mcp/core/database.py

from sqlalchemy import create_engine
from sqlalchemy.engine import Engine
import logging
from typing import Dict

from . import config

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# A dictionary to hold the initialized SQLAlchemy engines
_db_engines: Dict[str, Engine] = {}

def get_db_connections() -> Dict[str, Engine]:
    """
    Initializes and returns a dictionary of SQLAlchemy engines for all configured databases.
    This function is idempotent.
    """
    global _db_engines
    if not _db_engines:
        logger.info("Initializing database connections...")
        for db_name, conn_str in config.DB_CONNECTIONS.items():
            try:
                engine = create_engine(conn_str)
                # Test the connection
                with engine.connect():
                    logger.info(f"Successfully connected to {db_name}")
                    _db_engines[db_name] = engine
            except Exception as e:
                logger.error(f"Failed to connect to {db_name}: {e}")
    return _db_engines

def close_db_connections():
    """Closes all active database connections."""
    global _db_engines
    logger.info("Closing database connections...")
    for engine in _db_engines.values():
        engine.dispose()
    _db_engines = {}