sibi-dst 2025.1.12__py3-none-any.whl → 2025.8.1__py3-none-any.whl

sibi_dst/df_helper/backends/sqlalchemy/_db_connection.py

@@ -1,15 +1,11 @@
+
 from __future__ import annotations
 import os
 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 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
@@ -17,32 +13,18 @@ from sqlalchemy.exc import OperationalError, SQLAlchemyError
 from sqlalchemy.orm import sessionmaker, Session
 from sqlalchemy.pool import QueuePool, NullPool, StaticPool, Pool
 
-# Assuming these are your project's internal modules
 from sibi_dst.utils import Logger
 from ._sql_model_builder import SqlAlchemyModelBuilder
 
+_ENGINE_REGISTRY_LOCK = threading.RLock()
+_ENGINE_REGISTRY: Dict[tuple, Dict[str, Any]] = {}
+
 
 class SqlAlchemyConnectionConfig(BaseModel):
     """
-    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 used as a context manager
-    to ensure resources are always released correctly.
-
-    Recommended Usage is via the `with` statement.
-    with SqlAlchemyConnectionConfig(...) as config:
-        session = config.get_session()
-        # ... do work ...
-    # config.close() is called automatically upon exiting the block.
-
-    Key Features:
-      - Context Manager Support: Guarantees resource cleanup.
-      - 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.
+    Thread-safe, registry-backed SQLAlchemy connection manager.
     """
+
     # --- Public Configuration ---
     connection_url: str
     table: Optional[str] = None
@@ -50,36 +32,28 @@ class SqlAlchemyConnectionConfig(BaseModel):
 
     # --- Pool Configuration ---
     pool_size: int = int(os.environ.get("DB_POOL_SIZE", 5))
-    max_overflow: int = int(os.environ.get("DB_MAX_OVERFLOW",10))
+    max_overflow: int = int(os.environ.get("DB_MAX_OVERFLOW", 10))
     pool_timeout: int = int(os.environ.get("DB_POOL_TIMEOUT", 30))
     pool_recycle: int = int(os.environ.get("DB_POOL_RECYCLE", 1800))
     pool_pre_ping: bool = True
     poolclass: Type[Pool] = QueuePool
 
-    # --- Internal & Runtime State ---
+    # --- Internal & Runtime State (normal fields; Pydantic allowed) ---
     model: Optional[Type[Any]] = None
     engine: Optional[Engine] = None
     logger: Optional[Logger] = None
-    _own_logger: bool = False  # Indicates if this instance owns the logger.
+    _own_logger: bool = False
     session_factory: Optional[sessionmaker] = None
 
-    # --- Private State ---
+    # --- Private State (plain Python values only) ---
     _engine_key_instance: tuple = ()
-    _closed: bool = False  # Flag to prevent double-closing.
-
-    # --- Class-level Shared Resources ---
-    _engine_registry: ClassVar[Dict[tuple, Dict[str, Any]]] = {}
-    _registry_lock: ClassVar[threading.Lock] = threading.Lock()
-
+    _closed: bool = False  # prevent double-closing
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    # Add __enter__ and __exit__ for context manager protocol
-    def __enter__(self) -> SqlAlchemyConnectionConfig:
-        """Enter the runtime context, returning self."""
+    def __enter__(self) -> "SqlAlchemyConnectionConfig":
         return self
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        """Exit the runtime context, ensuring that close() is called."""
         self.close()
 
     @field_validator("pool_size", "max_overflow", "pool_timeout", "pool_recycle")
@@ -90,8 +64,7 @@ class SqlAlchemyConnectionConfig(BaseModel):
         return v
 
     @model_validator(mode="after")
-    def _init_all(self) -> SqlAlchemyConnectionConfig:
-        """Orchestrates the initialization process after Pydantic validation."""
+    def _init_all(self) -> "SqlAlchemyConnectionConfig":
         self._init_logger()
         self._engine_key_instance = self._get_engine_key()
         self._init_engine()
@@ -102,17 +75,12 @@ class SqlAlchemyConnectionConfig(BaseModel):
         return self
 
     def _init_logger(self) -> None:
-        """Initializes the logger for this instance."""
-        # This is not a ManagedResource subclass, so we handle logger initialization directly.
-        # unless a logger is provided, we create our own.
         if self.logger is None:
             self._own_logger = True
             self.logger = Logger.default_logger(logger_name=self.__class__.__name__)
-            log_level = Logger.DEBUG if self.debug else Logger.INFO
-            self.logger.set_level(log_level)
+            self.logger.set_level(Logger.DEBUG if self.debug else Logger.INFO)
 
     def _get_engine_key(self) -> tuple:
-        """Generates a unique, normalized key for an engine configuration."""
         parsed = sqlalchemy_url.make_url(self.connection_url)
         query = {k: v for k, v in parsed.query.items() if not k.startswith("pool_")}
         normalized_url = parsed.set(query=query)
@@ -125,104 +93,97 @@ class SqlAlchemyConnectionConfig(BaseModel):
         return tuple(key_parts)
 
     def _init_engine(self) -> None:
-        """Initializes or reuses a shared SQLAlchemy Engine."""
-        with self._registry_lock:
-            engine_wrapper = self._engine_registry.get(self._engine_key_instance)
-            if engine_wrapper:
-                self.engine = engine_wrapper['engine']
-                engine_wrapper['ref_count'] += 1
-                self.logger.debug(f"Reusing engine. Ref count: {engine_wrapper['ref_count']}.")
+        with _ENGINE_REGISTRY_LOCK:
+            wrapper = _ENGINE_REGISTRY.get(self._engine_key_instance)
+            if wrapper:
+                self.engine = wrapper["engine"]
+                wrapper["ref_count"] += 1
+                if self.debug:
+                    self.logger.debug(f"Reusing engine. Ref count: {wrapper['ref_count']}.")
             else:
-                self.logger.debug(f"Creating new engine for key: {self._engine_key_instance}")
+                if self.debug:
+                    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,
+                        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()
-                    self._engine_registry[self._engine_key_instance] = {
-                        'engine': new_engine, 'ref_count': 1, 'active_connections': 0
+                    _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
 
-            #self.logger.debug(f"Connections Active: {self.active_connections}")
-
     def close(self) -> None:
-        """
-        Decrements the engine's reference count and disposes of the engine
-        if the count reaches zero. This is now typically called automatically
-        when exiting a `with` block.
-        """
-        # Prevent the method from running more than once per instance.
         if self._closed:
-            self.logger.debug("Attempted to close an already-closed config instance.")
+            if self.debug:
+                self.logger.debug("Attempted to close an already-closed config instance.")
             return
 
-        with self._registry_lock:
+        with _ENGINE_REGISTRY_LOCK:
             key = self._engine_key_instance
-            engine_wrapper = self._engine_registry.get(key)
-
-            if not engine_wrapper:
+            wrapper = _ENGINE_REGISTRY.get(key)
+            if not 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 connection within engine wrapper. Ref count is now {engine_wrapper['ref_count']}.")
-
-            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]
-
-        # Mark this instance as closed to prevent subsequent calls.
+            else:
+                wrapper["ref_count"] -= 1
+                if self.debug:
+                    self.logger.debug(f"Closing connection. Ref count now {wrapper['ref_count']}.")
+                if wrapper["ref_count"] <= 0:
+                    if self.debug:
+                        self.logger.debug(f"Disposing engine as reference count is zero. Key: {key}")
+                    try:
+                        wrapper["engine"].dispose()
+                    finally:
+                        del _ENGINE_REGISTRY[key]
         self._closed = True
 
-
     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)
+        if not self.engine:
+            return
+        event.listen(self.engine, "checkout", self._on_checkout)
+        event.listen(self.engine, "checkin", self._on_checkin)
 
     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)
+        with _ENGINE_REGISTRY_LOCK:
+            wrapper = _ENGINE_REGISTRY.get(self._engine_key_instance)
             if wrapper:
-                wrapper['active_connections'] += 1
+                wrapper["active_connections"] += 1
 
     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)
+        with _ENGINE_REGISTRY_LOCK:
+            wrapper = _ENGINE_REGISTRY.get(self._engine_key_instance)
             if wrapper:
-                wrapper['active_connections'] = max(0, wrapper['active_connections'] - 1)
+                wrapper["active_connections"] = max(0, wrapper["active_connections"] - 1)
 
     @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
+        with _ENGINE_REGISTRY_LOCK:
+            wrapper = _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:
             with self.managed_connection() as conn:
                 conn.execute(text("SELECT 1"))
-            self.logger.debug("Database connection validated successfully.")
+            if self.debug:
+                self.logger.debug("Database connection validated successfully.")
         except OperationalError as e:
             self.logger.error(f"Database connection failed: {e}")
             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()
@@ -232,19 +193,19 @@ class SqlAlchemyConnectionConfig(BaseModel):
             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 an ORM model if `self.table` is set."""
         if not self.table or not self.engine:
             return
         try:
             builder = SqlAlchemyModelBuilder(self.engine, self.table)
             self.model = builder.build_model()
-            self.logger.debug(f"Successfully built ORM model for table: {self.table}")
+            if self.debug:
+                self.logger.debug(f"Successfully built ORM model for table: {self.table}")
         except Exception as e:
             self.logger.error(f"Failed to build ORM model for table '{self.table}': {e}")
             raise ValueError(f"Model construction failed for table '{self.table}': {e}") from e
+

sibi_dst/df_helper/backends/sqlalchemy/_db_gatekeeper.py

@@ -0,0 +1,15 @@
+import threading
+
+
+class DBGatekeeper:
+    _locks = {}
+    _global_lock = threading.Lock()
+
+    @classmethod
+    def get(cls, key: str, max_concurrency: int):
+        with cls._global_lock:
+            sem = cls._locks.get(key)
+            if sem is None:
+                sem = threading.BoundedSemaphore(max_concurrency)
+                cls._locks[key] = sem
+            return sem