import threading import queue import time import logging from typing import Optional, Dict, Any from .connectorDbJson import DatabaseConnector logger = logging.getLogger(__name__) class DatabaseConnectorPool: """ A connection pool for DatabaseConnector instances to manage resources efficiently and ensure proper isolation between users. """ def __init__(self, max_connections: int = 100, max_idle_time: int = 300): """ Initialize the connection pool. Args: max_connections: Maximum number of connections in the pool max_idle_time: Maximum idle time in seconds before connection is considered stale """ self.max_connections = max_connections self.max_idle_time = max_idle_time self._pool = queue.Queue(maxsize=max_connections) self._created_connections = 0 self._lock = threading.Lock() self._connection_times = {} # Track when connections were created def _create_connector(self, dbHost: str, dbDatabase: str, dbUser: str = None, dbPassword: str = None, userId: str = None) -> DatabaseConnector: """Create a new DatabaseConnector instance.""" with self._lock: if self._created_connections >= self.max_connections: raise RuntimeError(f"Maximum connections ({self.max_connections}) exceeded") self._created_connections += 1 logger.debug(f"Creating new database connector (total: {self._created_connections})") connector = DatabaseConnector( dbHost=dbHost, dbDatabase=dbDatabase, dbUser=dbUser, dbPassword=dbPassword, userId=userId ) # Track creation time connector_id = id(connector) self._connection_times[connector_id] = time.time() return connector def get_connector(self, dbHost: str, dbDatabase: str, dbUser: str = None, dbPassword: str = None, userId: str = None) -> DatabaseConnector: """ Get a database connector from the pool or create a new one. Args: dbHost: Database host path dbDatabase: Database name dbUser: Database user (optional) dbPassword: Database password (optional) userId: User ID for context (optional) Returns: DatabaseConnector instance """ try: # Try to get an existing connector from the pool connector = self._pool.get_nowait() # Check if connector is stale connector_id = id(connector) if connector_id in self._connection_times: idle_time = time.time() - self._connection_times[connector_id] if idle_time > self.max_idle_time: logger.debug(f"Connector {connector_id} is stale (idle: {idle_time}s), creating new one") # Remove stale connector from tracking if connector_id in self._connection_times: del self._connection_times[connector_id] # Create new connector return self._create_connector(dbHost, dbDatabase, dbUser, dbPassword, userId) # Update user context if provided if userId is not None: connector.updateContext(userId) logger.debug(f"Reusing existing connector {connector_id}") return connector except queue.Empty: # Pool is empty, create new connector return self._create_connector(dbHost, dbDatabase, dbUser, dbPassword, userId) def return_connector(self, connector: DatabaseConnector) -> None: """ Return a connector to the pool for reuse. Args: connector: DatabaseConnector instance to return """ try: # Update connection time connector_id = id(connector) self._connection_times[connector_id] = time.time() # Try to return to pool self._pool.put_nowait(connector) logger.debug(f"Returned connector {connector_id} to pool") except queue.Full: # Pool is full, discard connector logger.debug(f"Pool full, discarding connector {id(connector)}") with self._lock: self._created_connections -= 1 if id(connector) in self._connection_times: del self._connection_times[id(connector)] def cleanup_stale_connections(self) -> int: """ Clean up stale connections from the pool. Returns: Number of connections cleaned up """ cleaned = 0 current_time = time.time() # Check all tracked connections stale_connectors = [] for connector_id, creation_time in list(self._connection_times.items()): if current_time - creation_time > self.max_idle_time: stale_connectors.append(connector_id) # Remove stale connections from tracking for connector_id in stale_connectors: if connector_id in self._connection_times: del self._connection_times[connector_id] cleaned += 1 logger.debug(f"Cleaned up {cleaned} stale connections") return cleaned def get_stats(self) -> Dict[str, Any]: """Get pool statistics.""" with self._lock: return { "max_connections": self.max_connections, "created_connections": self._created_connections, "available_connections": self._pool.qsize(), "tracked_connections": len(self._connection_times) } # Global pool instance _connector_pool = None _pool_lock = threading.Lock() def get_connector_pool() -> DatabaseConnectorPool: """Get the global connector pool instance.""" global _connector_pool if _connector_pool is None: with _pool_lock: if _connector_pool is None: _connector_pool = DatabaseConnectorPool() return _connector_pool def get_connector(dbHost: str, dbDatabase: str, dbUser: str = None, dbPassword: str = None, userId: str = None) -> DatabaseConnector: """Get a database connector from the global pool.""" pool = get_connector_pool() return pool.get_connector(dbHost, dbDatabase, dbUser, dbPassword, userId) def return_connector(connector: DatabaseConnector) -> None: """Return a database connector to the global pool.""" pool = get_connector_pool() pool.return_connector(connector)