gateway/modules/connectors/connectorPool.py

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)