sibi-dst 0.3.63__py3-none-any.whl → 2025.1.1__py3-none-any.whl

sibi_dst/v2/df_helper/backends/sqlalchemy/_db_connection.py

@@ -1,82 +1,357 @@
-from typing import Any, Optional
+from __future__ import annotations
 
-from pydantic import BaseModel, model_validator, ConfigDict
-from sqlalchemy import create_engine, text
+import threading
+from contextlib import contextmanager
+from typing import Any, Optional, ClassVar, Generator, Type, Dict
+
+from pydantic import (
+    BaseModel,
+    field_validator,
+    model_validator,
+    ConfigDict,
+)
+from sqlalchemy import create_engine, event, text
+from sqlalchemy.engine import url as sqlalchemy_url
 from sqlalchemy.engine import Engine
-from sqlalchemy.exc import OperationalError
+from sqlalchemy.exc import OperationalError, SQLAlchemyError
+from sqlalchemy.orm import sessionmaker, Session
+from sqlalchemy.pool import QueuePool, NullPool, StaticPool
 
-from sibi_dst.v2.utils import Logger
+# Restoring the original intended imports
+from sibi_dst.utils import Logger
 from ._model_builder import SqlAlchemyModelBuilder
 
 
 class SqlAlchemyConnectionConfig(BaseModel):
     """
-    Configuration for establishing an SQLAlchemy database connection and dynamically building
-    an ORM model for a specific table.
+    A thread-safe, registry-backed SQLAlchemy connection manager.
+
+    This class encapsulates database connection configuration and provides robust,
+    shared resource management. It is designed to be instantiated multiple times
+    with the same connection parameters, where it will safely share a single
+    underlying database engine and connection pool.
+
+    Key Features:
+      - Shared Engine & Pool: Reuses a single SQLAlchemy Engine for identical
+        database URLs and pool settings, improving application performance.
+      - Reference Counting: Safely manages the lifecycle of the shared engine,
+        disposing of it only when the last user has closed its connection config.
+      - Per-Engine Connection Tracking: Monitors active connections for each
+        engine individually.
+      - Dynamic ORM Model Building: If a table name is provided, it can
+        dynamically generate a SQLAlchemy ORM model for that table.
+      - Integrated Session Management: Provides a pre-configured session factory
+        for convenient database interactions.
 
     Attributes:
-        connection_url (str): The URL used to connect to the database.
-        table_name (Optional[str]): The name of the table for which the model will be built.
-        model (Any): The dynamically built SQLAlchemy model.
-        engine (Optional[Engine]): The SQLAlchemy engine instance.
+        connection_url (str): The database URL (e.g., "postgresql://u:p@h/d").
+        table (Optional[str]): If provided, an ORM model will be built for this table.
+        model (Optional[Any]): The dynamically generated ORM model class.
+        engine (Optional[Engine]): The underlying SQLAlchemy engine.
+        logger (Logger): The logger instance.
+        debug (bool): If True, sets logger to DEBUG level.
+        session_factory (Optional[sessionmaker]): A factory for creating Session objects.
     """
-
-    model_config = ConfigDict(arbitrary_types_allowed=True)
+    # --- Public Configuration ---
     connection_url: str
     table: Optional[str] = None
-    model: Any = None
-    engine: Optional[Engine] = None
     debug: bool = False
+
+    # --- Pool Configuration ---
+    pool_size: int = 5
+    max_overflow: int = 10
+    pool_timeout: int = 30
+    pool_recycle: int = 300
+    pool_pre_ping: bool = True
+    poolclass: Type[QueuePool] = QueuePool
+
+    # --- Internal & Runtime State ---
+    model: Optional[Type[Any]] = None
+    engine: Optional[Engine] = None
     logger: Optional[Logger] = None
-    add_relationships: bool = False
-    export_models: bool = False
-    export_file_name: str = 'models.py'
+    session_factory: Optional[sessionmaker] = None
+
+    # --- Private State ---
+    _engine_key_instance: tuple = ()
+
+    # --- Class-level Shared Resources ---
+    # Registry stores engine wrappers with metadata like ref_count.
+    # Format: { engine_key: {'engine': Engine, 'ref_count': int, 'active_connections': int} }
+    _engine_registry: ClassVar[Dict[tuple, Dict[str, Any]]] = {}
+    _registry_lock: ClassVar[threading.Lock] = threading.Lock()
+
+    # Pydantic v2 configuration
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+    )
+
+    @field_validator("pool_size", "max_overflow", "pool_timeout", "pool_recycle")
+    @classmethod
+    def _validate_pool_params(cls, v: int) -> int:
+        if v < 0:
+            raise ValueError("Pool parameters must be non-negative")
+        return v
 
     @model_validator(mode="after")
-    def validate_and_initialize(self) -> "SqlAlchemyConnectionConfig":
+    def _init_all(self) -> SqlAlchemyConnectionConfig:
         """
-        Validate the configuration, initialize the engine, test the connection, and build the model.
+        Orchestrates the initialization process after Pydantic validation.
+        This method sets up the logger, engine, validates the connection,
+        builds the ORM model, and creates the session factory.
+        """
+        self._init_logger()
+
+        # The engine key is generated and stored on the instance to ensure
+        # that even if public attributes are changed later, this instance
+        # can always find its original engine in the registry.
+        self._engine_key_instance = self._get_engine_key()
+
+        self._init_engine()
+        self._validate_conn()
+        self._build_model()
 
-        Raises:
-            ValueError: If `connection_url` or `table_name` is missing, or if the connection or model
-                        building fails.
+        if self.engine:
+            self.session_factory = sessionmaker(bind=self.engine, expire_on_commit=False)
+
+        return self
+
+    def _init_logger(self) -> None:
+        """Initializes the logger for this instance."""
+        if self.logger is None:
+            self.logger = Logger.default_logger(logger_name=self.__class__.__name__)
+        # Assuming the Logger has a set_level method that accepts standard logging levels
+        log_level = Logger.DEBUG if self.debug else Logger.INFO
+        self.logger.set_level(log_level)
+
+    def _get_engine_key(self) -> tuple:
+        """
+        Generates a unique, normalized key for an engine configuration.
+        This key is used as the dictionary key in the shared _engine_registry.
+        It intentionally excludes instance-specific details like `table` to
+        ensure proper engine sharing.
+        """
+        parsed = sqlalchemy_url.make_url(self.connection_url)
+        # Exclude pooling from the query params as they are handled separately
+        query = {k: v for k, v in parsed.query.items() if not k.startswith("pool_")}
+        normalized_url = parsed.set(query=query)
+
+        key_parts = [str(normalized_url)]
+        if self.poolclass not in (NullPool, StaticPool):
+            key_parts += [
+                self.pool_size, self.max_overflow, self.pool_timeout,
+                self.pool_recycle, self.pool_pre_ping
+            ]
+        return tuple(key_parts)
+
+    def _init_engine(self) -> None:
+        """
+        Initializes the SQLAlchemy Engine.
+        If an identical engine already exists in the registry, it is reused
+        and its reference count is incremented. Otherwise, a new engine
+        is created and added to the registry.
         """
-        self.logger = self.logger or Logger.default_logger(logger_name="sqlalchemy_connection", debug=self.debug)
-        self.logger.debug("Validating and initializing SQLAlchemy connection configuration.")
-        if not self.connection_url:
-            raise ValueError("`connection_url` must be provided.")
+        with self._registry_lock:
+            engine_wrapper = self._engine_registry.get(self._engine_key_instance)
+
+            if engine_wrapper:
+                # --- Reuse existing engine ---
+                self.engine = engine_wrapper['engine']
+                engine_wrapper['ref_count'] += 1
+                self.logger.debug(
+                    f"Reusing engine. Ref count: {engine_wrapper['ref_count']}. Key: {self._engine_key_instance}")
+            else:
+                # --- Create new engine ---
+                self.logger.debug(f"Creating new engine for key: {self._engine_key_instance}")
+                try:
+                    new_engine = create_engine(
+                        self.connection_url,
+                        pool_size=self.pool_size,
+                        max_overflow=self.max_overflow,
+                        pool_timeout=self.pool_timeout,
+                        pool_recycle=self.pool_recycle,
+                        pool_pre_ping=self.pool_pre_ping,
+                        poolclass=self.poolclass,
+                    )
+                    self.engine = new_engine
+                    self._attach_events()
+
+                    # Store the new engine and its metadata in the registry
+                    self._engine_registry[self._engine_key_instance] = {
+                        'engine': new_engine,
+                        'ref_count': 1,
+                        'active_connections': 0
+                    }
+                except Exception as e:
+                    self.logger.error(f"Failed to create engine: {e}")
+                    raise SQLAlchemyError(f"Engine creation failed: {e}") from e
 
-        # Initialize the engine.
-        self.engine = create_engine(self.connection_url)
-        self.logger.debug(f"Engine created for URL")
+    def _attach_events(self) -> None:
+        """Attaches checkout/checkin events to the engine for connection tracking."""
+        if self.engine:
+            event.listen(self.engine, "checkout", self._on_checkout)
+            event.listen(self.engine, "checkin", self._on_checkin)
 
-        # Validate the connection.
-        self.validate_connection()
+    def _on_checkout(self, *args) -> None:
+        """Event listener for when a connection is checked out from the pool."""
+        with self._registry_lock:
+            wrapper = self._engine_registry.get(self._engine_key_instance)
+            if wrapper:
+                wrapper['active_connections'] += 1
+        self.logger.debug(f"Connection checked out. Active: {self.active_connections}")
 
-        if not self.table:
-            raise ValueError("`table` must be provided to build the model.")
+    def _on_checkin(self, *args) -> None:
+        """Event listener for when a connection is returned to the pool."""
+        with self._registry_lock:
+            wrapper = self._engine_registry.get(self._engine_key_instance)
+            if wrapper:
+                wrapper['active_connections'] = max(0, wrapper['active_connections'] - 1)
+        self.logger.debug(f"Connection checked in. Active: {self.active_connections}")
 
+    @property
+    def active_connections(self) -> int:
+        """Returns the number of active connections for this instance's engine."""
+        with self._registry_lock:
+            wrapper = self._engine_registry.get(self._engine_key_instance)
+            return wrapper['active_connections'] if wrapper else 0
+
+    def _validate_conn(self) -> None:
+        """Tests the database connection by executing a simple query."""
         try:
-            builder = SqlAlchemyModelBuilder(self.engine, self.table, self.add_relationships, self.debug, self.logger)
+            with self.managed_connection() as conn:
+                conn.execute(text("SELECT 1"))
+            self.logger.debug("Database connection validated successfully.")
+        except OperationalError as e:
+            self.logger.error(f"Database connection failed: {e}")
+            # This will be caught by Pydantic and raised as a ValidationError
+            raise ValueError(f"DB connection failed: {e}") from e
+
+    @contextmanager
+    def managed_connection(self) -> Generator[Any, None, None]:
+        """Provides a single database connection from the engine pool."""
+        if not self.engine:
+            raise RuntimeError("Engine not initialized. Cannot get a connection.")
+        conn = self.engine.connect()
+        try:
+            yield conn
+        finally:
+            conn.close()
+
+    def get_session(self) -> Session:
+        """Returns a new SQLAlchemy Session from the session factory."""
+        if not self.session_factory:
+            raise RuntimeError("Session factory not initialized. Cannot get a session.")
+        return self.session_factory()
+
+    def _build_model(self) -> None:
+        """
+        Dynamically builds and assigns an ORM model if `self.table` is set.
+        Uses SqlAlchemyModelBuilder to reflect the table schema from the database.
+        """
+        if not self.table or not self.engine:
+            return
+        try:
+            builder = SqlAlchemyModelBuilder(self.engine, self.table)
             self.model = builder.build_model()
-            if self.export_models:
-                builder.export_models_to_file(self.export_file_name)
-            self.logger.debug(f"Successfully built model for table: {self.table}")
+            self.logger.debug(f"Successfully built ORM model for table: {self.table}")
         except Exception as e:
-            raise ValueError(f"Failed to build model for table {self.table}: {e}")
+            self.logger.error(f"Failed to build ORM model for table '{self.table}': {e}")
+            # Propagate as a ValueError to be caught by Pydantic validation
+            raise ValueError(f"Model construction failed for table '{self.table}': {e}") from e
 
-        return self
+    def close(self) -> None:
+        """
+        Decrements the engine's reference count and disposes of the engine
+        if the count reaches zero. This is the primary method for releasing
+        resources managed by this configuration object.
+        """
+        with self._registry_lock:
+            key = self._engine_key_instance
+            engine_wrapper = self._engine_registry.get(key)
+
+            if not engine_wrapper:
+                self.logger.warning("Attempted to close a config whose engine is not in the registry.")
+                return
+
+            engine_wrapper['ref_count'] -= 1
+            self.logger.debug(f"Closing config. Ref count is now {engine_wrapper['ref_count']} for key {key}.")
 
-    def validate_connection(self) -> None:
+            if engine_wrapper['ref_count'] <= 0:
+                self.logger.debug(f"Disposing engine as reference count is zero. Key: {key}")
+                engine_wrapper['engine'].dispose()
+                del self._engine_registry[key]
+
+    def dispose_idle_connections(self) -> int:
         """
-        Test the database connection by executing a simple query.
+        Closes and discards idle connections in this engine's connection pool.
+
+        This is a pool-level operation that affects only the connections held
+        by the pool, not those active at the database level.
 
-        Raises:
-            ValueError: If the connection cannot be established.
+        Returns:
+            int: The number of connections disposed of.
         """
-        try:
-            with self.engine.connect() as connection:
-                connection.execute(text("SELECT 1"))
-            self.logger.debug("Database connection validated.")
-        except OperationalError as e:
-            raise ValueError(f"Failed to connect to the database: {e}")
+        if not self.engine:
+            self.logger.warning("Cannot dispose idle connections: engine not initialized.")
+            return 0
+
+        pool = self.engine.pool
+        if isinstance(pool, QueuePool):
+            # The pool.dispose() method checks in all connections and clears the pool.
+            # The number of checked-in connections is the number of idle connections.
+            count = pool.checkedin()
+            pool.dispose()
+            self.logger.debug(f"Disposed {count} idle connections from the pool.")
+            return count
+
+        self.logger.warning(f"Idle connection disposal not applicable for pool type: {type(pool).__name__}")
+        return 0
+
+    def terminate_idle_connections(self, idle_seconds: int = 300) -> int:
+        """
+        Terminates idle connections at the database server level.
+
+        This is an administrative database operation that forcibly closes
+        backend processes. Use with caution. Support is dialect-specific.
+
+        Args:
+            idle_seconds (int): The minimum idle time in seconds for a
+                                connection to be terminated.
+
+        Returns:
+            int: The number of connections terminated.
+        """
+        if not self.engine:
+            self.logger.warning("Cannot terminate connections: engine not initialized.")
+            return 0
+
+        terminated_count = 0
+        dialect = self.engine.dialect.name
+
+        with self.managed_connection() as conn:
+            try:
+                if dialect == 'postgresql':
+                    query = text(
+                        "SELECT pg_terminate_backend(pid) FROM pg_stat_activity "
+                        "WHERE state = 'idle' "
+                        "AND (now() - query_start) > INTERVAL ':idle_seconds seconds' "
+                        "AND pid <> pg_backend_pid()"
+                    )
+                    res = conn.execute(query, {"idle_seconds": idle_seconds})
+                    terminated_count = res.rowcount if res.rowcount is not None else 0
+
+                elif dialect == 'mysql':
+                    process_list = conn.execute(text("SHOW PROCESSLIST"))
+                    for row in process_list:
+                        # Pylint compatible attribute access
+                        if getattr(row, 'Command', '') == 'Sleep' and getattr(row, 'Time', 0) > idle_seconds:
+                            conn.execute(text(f"KILL {getattr(row, 'Id')}"))
+                            terminated_count += 1
+                else:
+                    self.logger.warning(f"Idle connection termination is not supported for dialect: {dialect}")
+            except SQLAlchemyError as e:
+                self.logger.error(f"Error terminating idle connections for dialect {dialect}: {e}")
+
+        if terminated_count > 0:
+            self.logger.debug(f"Terminated {terminated_count} idle connections on the database server.")
+        return terminated_count
+

sibi_dst/v2/df_helper/backends/sqlalchemy/_io_dask.py

@@ -10,7 +10,7 @@ from sibi_dst.v2.utils import Logger
 
 
 class SQLAlchemyDask:
-    def __init__(self, model, filters, engine_url, chunk_size=1000, logger=None, debug=False):
+    def __init__(self, model, filters, engine, chunk_size=1000, logger=None, debug=False):
         """
         Initialize with an SQLAlchemy query and database engine URL.
 
@@ -26,7 +26,7 @@ class SQLAlchemyDask:
         self.filters = filters
         self.chunk_size = chunk_size
         self.debug = debug
-        self.engine = create_engine(engine_url)
+        self.engine = engine
         self.Session = sessionmaker(bind=self.engine)
         self.logger = logger or Logger.default_logger(logger_name=self.__class__.__name__)
         self.logger.set_level(logger.DEBUG if debug else logger.INFO)