Update modules/database/database_init.py

modules/database/database_init.py

@@ -1,80 +1,143 @@
-import os
-import logging
-from azure.cosmos import CosmosClient
-from pymongo import MongoClient
-import certifi
-
-logging.basicConfig(level=logging.DEBUG)
-logger = logging.getLogger(__name__)
-
-# Variables globales para Cosmos DB SQL API
-cosmos_client = None
-user_database = None
-user_container = None
-application_requests_container = None
-user_feedback_container = None
-
-# Variables globales para Cosmos DB MongoDB API
-mongo_client = None
-mongo_db = None
-
-def initialize_cosmos_sql_connection():
-    global cosmos_client, user_database, user_container, application_requests_container, user_feedback_container
-    try:
-        cosmos_endpoint = os.environ.get("COSMOS_ENDPOINT")
-        cosmos_key = os.environ.get("COSMOS_KEY")
-
-        if not cosmos_endpoint or not cosmos_key:
-            raise ValueError("COSMOS_ENDPOINT and COSMOS_KEY environment variables must be set")
-
-        cosmos_client = CosmosClient(cosmos_endpoint, cosmos_key)
-        user_database = cosmos_client.get_database_client("user_database")
-        user_container = user_database.get_container_client("users")
-        application_requests_container = user_database.get_container_client("application_requests")
-        user_feedback_container = user_database.get_container_client("user_feedback")
-
-        logger.info("Conexión a Cosmos DB SQL API exitosa")
-        return True
-    except Exception as e:
-        logger.error(f"Error al conectar con Cosmos DB SQL API: {str(e)}", exc_info=True)
-        return False
-
-def initialize_mongodb_connection():
-    global mongo_client, mongo_db
-    try:
-        cosmos_mongodb_connection_string = os.getenv("MONGODB_CONNECTION_STRING")
-        if not cosmos_mongodb_connection_string:
-            raise ValueError("MONGODB_CONNECTION_STRING environment variable is not set")
-
-        mongo_client = MongoClient(cosmos_mongodb_connection_string,
-                                   tls=True,
-                                   tlsCAFile=certifi.where(),
-                                   retryWrites=False,
-                                   serverSelectionTimeoutMS=5000,
-                                   connectTimeoutMS=10000,
-                                   socketTimeoutMS=10000)
-
-        mongo_client.admin.command('ping')
-
-        mongo_db = mongo_client['aideatext_db']
-
-        logger.info("Conexión a Cosmos DB MongoDB API exitosa")
-        return True
-    except Exception as e:
-        logger.error(f"Error al conectar con Cosmos DB MongoDB API: {str(e)}", exc_info=True)
-        return False
-
-def initialize_database_connections():
-    sql_success = initialize_cosmos_sql_connection()
-    mongodb_success = initialize_mongodb_connection()
-    return sql_success and mongodb_success
-
-def get_sql_containers():
-    if user_container is None or application_requests_container is None or user_feedback_container is None:
-        initialize_cosmos_sql_connection()
-    return user_container, application_requests_container, user_feedback_container
-
-def get_mongodb():
-    if mongo_db is None:
-        initialize_mongodb_connection()
+import os
+import logging
+from azure.cosmos import CosmosClient
+from pymongo import MongoClient
+import certifi
+
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
+# Variables globales para Cosmos DB SQL API
+cosmos_client = None
+user_database = None
+user_container = None
+application_requests_container = None
+user_feedback_container = None
+user_sessions_container = None  # Agregado para sessions
+
+# Variables globales para Cosmos DB MongoDB API
+mongo_client = None
+mongo_db = None
+
+def verify_container_partition_key(container, expected_path):
+    """Verifica la configuración de partition key de un contenedor"""
+    try:
+        container_props = container.read()
+        partition_key_paths = container_props['partitionKey']['paths']
+        
+        logger.info(f"Container: {container.id}, Partition Key Paths: {partition_key_paths}")
+        
+        if expected_path not in partition_key_paths:
+            logger.warning(
+                f"Partition key incorrecta en {container.id}. "
+                f"Actual: {partition_key_paths}, Esperada: {expected_path}"
+            )
+            return False
+        return True
+    except Exception as e:
+        logger.error(f"Error verificando partition key en {container.id}: {str(e)}")
+        return False
+
+def initialize_cosmos_sql_connection():
+    """Inicializa y verifica la conexión a Cosmos DB SQL API"""
+    global cosmos_client, user_database, user_container, \
+           application_requests_container, user_feedback_container, \
+           user_sessions_container
+
+    try:
+        cosmos_endpoint = os.environ.get("COSMOS_ENDPOINT")
+        cosmos_key = os.environ.get("COSMOS_KEY")
+        
+        if not cosmos_endpoint or not cosmos_key:
+            raise ValueError("COSMOS_ENDPOINT and COSMOS_KEY environment variables must be set")
+
+        # Inicializar cliente y obtener contenedores
+        cosmos_client = CosmosClient(cosmos_endpoint, cosmos_key)
+        user_database = cosmos_client.get_database_client("user_database")
+        
+        # Configuración de contenedores y sus partition keys esperadas
+        containers_config = {
+            "users": ("/id", user_container),
+            "application_requests": ("/email", application_requests_container),
+            "user_feedback": ("/username", user_feedback_container),
+            "users_sessions": ("/username", user_sessions_container)
+        }
+
+        # Inicializar y verificar cada contenedor
+        for container_id, (expected_path, container_ref) in containers_config.items():
+            container = user_database.get_container_client(container_id)
+            if not verify_container_partition_key(container, expected_path):
+                logger.error(f"Verificación de partition key falló para {container_id}")
+                return False
+                
+            # Actualizar la referencia global
+            if container_id == "users":
+                user_container = container
+            elif container_id == "application_requests":
+                application_requests_container = container
+            elif container_id == "user_feedback":
+                user_feedback_container = container
+            elif container_id == "users_sessions":
+                user_sessions_container = container
+
+        logger.info("Conexión a Cosmos DB SQL API exitosa y particiones verificadas")
+        return True
+
+    except Exception as e:
+        logger.error(f"Error al conectar con Cosmos DB SQL API: {str(e)}", exc_info=True)
+        return False
+
+def initialize_mongodb_connection():
+    """Inicializa la conexión a Cosmos DB MongoDB API"""
+    global mongo_client, mongo_db
+    try:
+        cosmos_mongodb_connection_string = os.getenv("MONGODB_CONNECTION_STRING")
+        if not cosmos_mongodb_connection_string:
+            raise ValueError("MONGODB_CONNECTION_STRING environment variable is not set")
+
+        mongo_client = MongoClient(
+            cosmos_mongodb_connection_string,
+            tls=True,
+            tlsCAFile=certifi.where(),
+            retryWrites=False,
+            serverSelectionTimeoutMS=5000,
+            connectTimeoutMS=10000,
+            socketTimeoutMS=10000
+        )
+        
+        # Verificar conexión
+        mongo_client.admin.command('ping')
+        mongo_db = mongo_client['aideatext_db']
+        
+        logger.info("Conexión a Cosmos DB MongoDB API exitosa")
+        return True
+        
+    except Exception as e:
+        logger.error(f"Error al conectar con Cosmos DB MongoDB API: {str(e)}", exc_info=True)
+        return False
+
+def initialize_database_connections():
+    """Inicializa todas las conexiones a bases de datos"""
+    sql_success = initialize_cosmos_sql_connection()
+    mongodb_success = initialize_mongodb_connection()
+    
+    if sql_success and mongodb_success:
+        logger.info("Todas las conexiones a bases de datos inicializadas correctamente")
+    else:
+        logger.error("Falló la inicialización de una o más conexiones a bases de datos")
+        
+    return sql_success and mongodb_success
+
+def get_sql_containers():
+    """Obtiene los contenedores SQL, reinicializando si es necesario"""
+    if any(container is None for container in 
+           [user_container, application_requests_container, 
+            user_feedback_container, user_sessions_container]):
+        initialize_cosmos_sql_connection()
+    return user_container, application_requests_container, user_feedback_container, user_sessions_container
+
+def get_mongodb():
+    """Obtiene la conexión MongoDB, reinicializando si es necesario"""
+    if mongo_db is None:
+        initialize_mongodb_connection()
     return mongo_db