diff_app_middleware_tenant_context.py

diff --git a/app/middleware/tenant_context.py b/app/middleware/tenant_context.py
index a24544f3..fbc62d9d 100644
--- a/app/middleware/tenant_context.py
+++ b/app/middleware/tenant_context.py
@@ -4,12 +4,15 @@ Provides tenant_id injection and multi-tenant isolation enforcement
 """
 import logging
 import uuid
+from datetime import datetime, timezone
 from typing import Optional
 from fastapi import Request, HTTPException, status
-from sqlalchemy.orm import Session
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
 
-from app.db import get_db
+from app.db import get_async_session
 from app.models import ApiKey, Tenant
+from app.utils.auth import derive_api_key_role
 
 logger = logging.getLogger(__name__)
 
@@ -39,19 +42,21 @@ class TenantContextMiddleware:
             return
 
         try:
-            # Extract and validate tenant context
             tenant_info = await self._extract_tenant_context(request)
 
-            # Inject tenant info into request state
             if tenant_info:
                 request.state.tenant_id = tenant_info["tenant_id"]
                 request.state.tenant_role = tenant_info["role"]
                 request.state.tenant_environment = tenant_info["environment"]
                 request.state.tenant_active = tenant_info["active"]
 
-                logger.debug(f"Tenant context set: {tenant_info['tenant_id']} ({tenant_info['role']}, {tenant_info['environment']})")
+                logger.debug(
+                    "Tenant context set: %s (%s, %s)",
+                    tenant_info['tenant_id'],
+                    tenant_info['role'],
+                    tenant_info['environment']
+                )
 
-                # Check tenant is active
                 if not tenant_info["active"]:
                     raise HTTPException(
                         status_code=status.HTTP_403_FORBIDDEN,
@@ -61,7 +66,6 @@ class TenantContextMiddleware:
             await self.app(scope, receive, send)
 
         except HTTPException as e:
-            # Return authentication/authorization error
             response_body = f'{{"detail": "{e.detail}"}}'.encode()
             await send({
                 "type": "http.response.start",
@@ -80,9 +84,9 @@ class TenantContextMiddleware:
             "/docs",
             "/openapi.json",
             "/redoc",
-            "/signup",  # Public signup endpoint
-            "/admin/",  # Admin endpoints (handled separately)
-            "/metrics",  # Monitoring endpoints
+            "/signup",
+            "/admin/",
+            "/metrics",
             "/favicon.ico",
         ]
         return any(path.startswith(skip_path) for skip_path in skip_paths)
@@ -91,7 +95,6 @@ class TenantContextMiddleware:
         """Extract tenant context from API key"""
         auth_header = request.headers.get("Authorization")
         if not auth_header:
-            # No authorization header - allow for public endpoints
             return None
 
         if not auth_header.startswith("Bearer "):
@@ -102,13 +105,8 @@ class TenantContextMiddleware:
 
         api_key = auth_header.replace("Bearer ", "")
 
-        db = next(get_db())
-        try:
-            # Look up API key and associated tenant
-            api_key_record = db.query(ApiKey).filter(
-                ApiKey.key_hash == self._hash_api_key(api_key),
-                ApiKey.active == True
-            ).first()
+        async with get_async_session() as db:
+            api_key_record, tenant = await self._fetch_api_key_and_tenant(db, api_key)
 
             if not api_key_record:
                 raise HTTPException(
@@ -116,27 +114,17 @@ class TenantContextMiddleware:
                     detail="Invalid or inactive API key"
                 )
 
-            # Check if API key has expired
-            if api_key_record.expires_at and api_key_record.expires_at < db.func.now():
-                raise HTTPException(
-                    status_code=status.HTTP_401_UNAUTHORIZED,
-                    detail="API key has expired"
-                )
-
-            # Update last used timestamp
-            api_key_record.last_used_at = db.func.now()
-            db.commit()
-
-            # Get tenant information
-            tenant = db.query(Tenant).filter(Tenant.id == api_key_record.tenant_id).first()
             if not tenant:
                 raise HTTPException(
                     status_code=status.HTTP_401_UNAUTHORIZED,
                     detail="Tenant not found"
                 )
 
-            # Determine role (admin keys can access admin endpoints)
-            role = "admin" if self._is_admin_key(api_key_record) else "client"
+            # Update last used timestamp
+            api_key_record.last_used_at = datetime.now(timezone.utc)
+            await db.commit()
+
+            role = derive_api_key_role(api_key_record, tenant)
 
             return {
                 "tenant_id": tenant.id,
@@ -147,8 +135,34 @@ class TenantContextMiddleware:
                 "plan": tenant.plan
             }
 
-        finally:
-            db.close()
+    async def _fetch_api_key_and_tenant(
+        self,
+        db: AsyncSession,
+        api_key: str
+    ) -> tuple[Optional[ApiKey], Optional[Tenant]]:
+        key_hash = self._hash_api_key(api_key)
+
+        api_key_result = await db.execute(
+            select(ApiKey).where(
+                ApiKey.key_hash == key_hash,
+                ApiKey.active == True
+            )
+        )
+        api_key_record = api_key_result.scalar_one_or_none()
+
+        if not api_key_record:
+            return None, None
+
+        now = datetime.now(timezone.utc)
+        if api_key_record.expires_at and api_key_record.expires_at <= now:
+            return None, None
+
+        tenant_result = await db.execute(
+            select(Tenant).where(Tenant.id == api_key_record.tenant_id)
+        )
+        tenant = tenant_result.scalar_one_or_none()
+
+        return api_key_record, tenant
 
     def _hash_api_key(self, api_key: str) -> str:
         """Hash API key for lookup (same method used in API key manager)"""
@@ -157,7 +171,6 @@ class TenantContextMiddleware:
 
     def _is_admin_key(self, api_key_record: ApiKey) -> bool:
         """Determine if this is an admin API key"""
-        # Admin keys can be identified by name pattern or special tenant
         return (
             api_key_record.name.lower().startswith("admin") or
             api_key_record.tenant_id == "admin" or
@@ -182,11 +195,10 @@ def get_current_tenant_environment(request: Request) -> str:
 
 def require_tenant_context(request: Request) -> str:
     """Helper function that requires tenant context and returns tenant_id"""
-    # Check for demo mode
     import os
     demo_mode = os.getenv("DEMO_MODE", "false").lower() == "true"
     if demo_mode:
-        return "demo"  # Return dummy tenant ID for demo mode
+        return "demo"
 
     tenant_id = get_current_tenant_id(request)
     if not tenant_id:
@@ -218,7 +230,6 @@ def log_tenant_access(request: Request, resource: str, action: str = "access"):
     """Log tenant access for compliance audit"""
     tenant_id = get_current_tenant_id(request)
     if tenant_id:
-        # Generate request ID for tracing
         request_id = getattr(request.state, "request_id", str(uuid.uuid4()))
 
         logger.info(
@@ -233,4 +244,4 @@ def log_tenant_access(request: Request, resource: str, action: str = "access"):
                 "user_agent": request.headers.get("user-agent", ""),
                 "ip_address": request.client.host if request.client else "unknown"
             }
-        )
\ No newline at end of file
+        )