From 115b54ad5802e31e869eab161bf874af87df2aa1 Mon Sep 17 00:00:00 2001 From: YuehuCao Date: Fri, 19 Sep 2025 12:11:51 +0800 Subject: [PATCH] refactor(name): rename the tenant DB connection middleware --- ....py => tenant_DB_connection_middleware.py} | 75 ++++++++++++++++--- 1 file changed, 63 insertions(+), 12 deletions(-) rename apps/notification/webapi/middleware/{tenant_middleware.py => tenant_DB_connection_middleware.py} (68%) diff --git a/apps/notification/webapi/middleware/tenant_middleware.py b/apps/notification/webapi/middleware/tenant_DB_connection_middleware.py similarity index 68% rename from apps/notification/webapi/middleware/tenant_middleware.py rename to apps/notification/webapi/middleware/tenant_DB_connection_middleware.py index 89a587e..c251b93 100644 --- a/apps/notification/webapi/middleware/tenant_middleware.py +++ b/apps/notification/webapi/middleware/tenant_DB_connection_middleware.py @@ -3,16 +3,46 @@ from urllib.parse import urlparse from motor.motor_asyncio import AsyncIOMotorClient from beanie import init_beanie from contextlib import asynccontextmanager +import os from common.config.app_settings import app_settings from common.log.module_logger import ModuleLogger -from webapi.providers.database import client from webapi.middleware.freeleaps_auth_middleware import request_context_var -from backend.models.models import MessageTemplateDoc +from backend.models.models import ( + MessageTemplateDoc, + EmailSenderDoc, + EmailSendStatusDoc, + EmailTrackingDoc, + EmailBounceDoc, + UsageLogDoc +) + +# MongoDB config (moved from providers/database.py) +MONGODB_URI = os.getenv('MONGODB_URI') +MONGODB_NAME = os.getenv('MONGODB_NAME') + +# Create MongoDB client (moved from providers/database.py) +client = AsyncIOMotorClient( + MONGODB_URI, + serverSelectionTimeoutMS=60000, + minPoolSize=5, + maxPoolSize=20, + heartbeatFrequencyMS=20000, +) + +# Define all document models (moved from providers/database.py) +document_models = [ + MessageTemplateDoc, + EmailSenderDoc, + EmailSendStatusDoc, + EmailTrackingDoc, + EmailBounceDoc, + UsageLogDoc +] -class TenantMiddleware: - """TenantMiddleware: Request-level tenant database isolation middleware +class TenantDBConnectionMiddleware: + """TenantDBConnectionMiddleware: Request-level tenant database isolation middleware Depends on Auth middleware to get the product_id after API Key validation. Must be executed after Auth middleware. @@ -20,7 +50,7 @@ class TenantMiddleware: def __init__(self, app): self.app = app - self.module_logger = ModuleLogger(sender_id=TenantMiddleware) + self.module_logger = ModuleLogger(sender_id=TenantDBConnectionMiddleware) async def __call__(self, scope, receive, send): """Process request, automatically switch tenant database""" @@ -30,6 +60,15 @@ class TenantMiddleware: request = Request(scope, receive) + # Check if request has API Key - if not, use default database (compatibility mode) + api_key = request.headers.get("X-API-KEY") + if not api_key or api_key == "": + await self.module_logger.log_info(f"No API Key provided - using default database for path: {request.url.path}") + # Ensure default database connection is available for business logic + await self._ensure_default_database_initialized() + await self.app(scope, receive, send) + return + # Get product_id from Auth middleware context try: request_context = request_context_var.get() @@ -82,8 +121,9 @@ class TenantMiddleware: @asynccontextmanager async def _get_tenant_database_context(self, product_id: str): """Get tenant database context manager, ensure request-level isolation""" + tenant_client = None try: - # Get tenant information directly from main database + # Use existing main database client to get tenant information main_db = client[app_settings.MONGODB_NAME] tenant_collection = main_db.tenant_doc @@ -94,8 +134,8 @@ class TenantMiddleware: if tenant_doc.get("status") != "active": raise ValueError(f"Tenant {product_id} is not active") - # Create new database connection (each request has its own connection) - client = AsyncIOMotorClient( + # Create new database connection for tenant (each request has its own connection) + tenant_client = AsyncIOMotorClient( tenant_doc["mongodb_uri"], serverSelectionTimeoutMS=10000, minPoolSize=5, @@ -104,11 +144,11 @@ class TenantMiddleware: # Extract database name from URI db_name = self._extract_db_name_from_uri(tenant_doc["mongodb_uri"]) - database = client[db_name] + database = tenant_client[db_name] # Initialize Beanie for this request (using independent database instance) # TODO: check whether only MessageTemplateDoc is needed - db_models = [MessageTemplateDoc] + db_models = [MessageTemplateDoc, EmailSenderDoc, EmailSendStatusDoc] await init_beanie(database=database, document_models=db_models) await self.module_logger.log_info(f"Only initialized base model for tenant {product_id}") @@ -119,8 +159,8 @@ class TenantMiddleware: finally: # Close connection safely after request try: - if client: - client.close() + if tenant_client: + tenant_client.close() await self.module_logger.log_info(f"Closed database connection for tenant {product_id}") except Exception as close_error: logger.warning(f"Error closing database connection for tenant {product_id}: {str(close_error)}") @@ -137,4 +177,15 @@ class TenantMiddleware: return parsed.path.lstrip('/') or app_settings.MONGODB_NAME except Exception: return app_settings.MONGODB_NAME + + async def _ensure_default_database_initialized(self): + """Ensure default database connection is properly initialized for requests without API Key""" + try: + # Initialize default database with all models for compatibility mode + default_db = client[MONGODB_NAME] + await init_beanie(database=default_db, document_models=document_models) + await self.module_logger.log_info("Default database initialized for compatibility mode") + except Exception as e: + await self.module_logger.log_error(f"Failed to initialize default database: {str(e)}") + \ No newline at end of file