sibi-dst 0.3.62__py3-none-any.whl → 0.3.64__py3-none-any.whl

sibi_dst/df_helper/_parquet_artifact.py

@@ -158,6 +158,7 @@ class ParquetArtifact(DfHelper):
         if not manifest_exists:
             self.logger.info(f"Creating new manifest at {self.missing_manifest_path}")
             self.mmanifest.save()
+            self.mmanifest.cleanup_temp_manifests()
         else:
             self.logger.info(f"Manifest already exists at {self.missing_manifest_path}")
 
@@ -206,6 +207,7 @@ class ParquetArtifact(DfHelper):
         try:
             if self.mmanifest and self.mmanifest._new_records:
                 self.mmanifest.save()
+                self.mmanifest.cleanup_temp_manifests()
             if getattr(self, "_entered", False) and self.fs and self._own_fs:
                 self.fs.close()
         except Exception as e:

sibi_dst/df_helper/backends/sqlalchemy/__init__.py

@@ -1,4 +1,3 @@
-from ._filter_handler import SqlAlchemyFilterHandler
 from ._db_connection import SqlAlchemyConnectionConfig
 from ._load_from_db import SqlAlchemyLoadFromDb
 from ._sql_model_builder import SqlAlchemyModelBuilder
@@ -7,5 +6,4 @@ __all__ = [
     'SqlAlchemyConnectionConfig',
     'SqlAlchemyModelBuilder',
     'SqlAlchemyLoadFromDb',
-    'SqlAlchemyFilterHandler'
 ]

sibi_dst/df_helper/backends/sqlalchemy/_db_connection.py

@@ -1,51 +1,85 @@
 from __future__ import annotations
-from typing import Any, Optional, ClassVar, Generator, Type
+
 import threading
 from contextlib import contextmanager
-from pydantic import BaseModel, field_validator, ValidationError, model_validator
+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.pool import QueuePool, NullPool, StaticPool
+from sqlalchemy.exc import OperationalError, SQLAlchemyError
 from sqlalchemy.orm import sessionmaker, Session
+from sqlalchemy.pool import QueuePool, NullPool, StaticPool
+
+# Assuming these are your project's internal modules
 from sibi_dst.utils import Logger
 from ._sql_model_builder import SqlAlchemyModelBuilder
 
 
 class SqlAlchemyConnectionConfig(BaseModel):
     """
-    Thread-safe, registry-backed SQLAlchemy connection manager with:
-      - Shared engine reuse
-      - Active connection tracking
-      - Idle-pool and database-level cleanup
-      - Dynamic ORM model building via SqlAlchemyModelBuilder
-      - Optional session factory
+    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.
     """
+    # --- Public Configuration ---
     connection_url: str
     table: Optional[str] = None
-    model: Optional[Any] = None
-    engine: Optional[Engine] = None
-    logger: Logger = None
     debug: bool = False
 
+    # --- Pool Configuration ---
     pool_size: int = 5
     max_overflow: int = 10
     pool_timeout: int = 30
-    pool_recycle: int = 300
+    pool_recycle: int = 1800
     pool_pre_ping: bool = True
-    poolclass: Type = QueuePool
+    poolclass: Type[QueuePool] = QueuePool
 
+    # --- Internal & Runtime State ---
+    model: Optional[Type[Any]] = None
+    engine: Optional[Engine] = None
+    logger: Optional[Logger] = None
     session_factory: Optional[sessionmaker] = None
-    _owns_engine: bool = False
 
-    _engine_registry: ClassVar[dict[tuple, Engine]] = {}
+    # --- Private State ---
+    _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()
-    _active_connections: ClassVar[int] = 0
 
-    class Config:
-        arbitrary_types_allowed = True
-        underscore_attrs_are_private = True
+    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."""
+        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")
     @classmethod
@@ -56,78 +90,139 @@ class SqlAlchemyConnectionConfig(BaseModel):
 
     @model_validator(mode="after")
     def _init_all(self) -> SqlAlchemyConnectionConfig:
+        """Orchestrates the initialization process after Pydantic validation."""
         self._init_logger()
+        self._engine_key_instance = self._get_engine_key()
         self._init_engine()
         self._validate_conn()
         self._build_model()
-        self.session_factory = sessionmaker(bind=self.engine, expire_on_commit=False)
+        if self.engine:
+            self.session_factory = sessionmaker(bind=self.engine, expire_on_commit=False)
         return self
 
     def _init_logger(self) -> None:
-        self.logger = self.logger or Logger.default_logger(logger_name=self.__class__.__name__)
-        self.logger.set_level(Logger.DEBUG if self.debug else Logger.INFO)
-
-    def _engine_key(self) -> tuple:
+        """Initializes the logger for this instance."""
+        if self.logger is None:
+            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)
+
+    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 = parsed.set(query=query)
-        key = [str(normalized)]
+        normalized_url = parsed.set(query=query)
+        key_parts = [str(normalized_url)]
         if self.poolclass not in (NullPool, StaticPool):
-            key += [self.pool_size, self.max_overflow, self.pool_timeout, self.pool_recycle, self.pool_pre_ping, self.table]
-        return tuple(key)
+            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:
-        key = self._engine_key()
+        """Initializes or reuses a shared SQLAlchemy Engine."""
         with self._registry_lock:
-            existing = self._engine_registry.get(key)
-            if existing:
-                self.engine = existing
-                self._owns_engine = False
-                self.logger.debug(f"Reusing engine {key}")
+            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']}.")
             else:
-                self.logger.debug(f"Creating engine {key}")
-                self.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._attach_events()
-                self._engine_registry[key] = self.engine
-                self._owns_engine = True
+                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()
+                    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
+
+    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.")
+            return
+
+        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']}.")
+
+            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.
+        self._closed = True
+
+    # ... (the rest of your methods like _attach_events, _on_checkout, get_session, etc. remain unchanged)
+    # They are omitted here for brevity but should be included in your final file.
 
     def _attach_events(self) -> None:
-        event.listen(self.engine, "checkout", self._on_checkout)
-        event.listen(self.engine, "checkin", self._on_checkin)
+        """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)
 
     def _on_checkout(self, *args) -> None:
+        """Event listener for when a connection is checked out from the pool."""
         with self._registry_lock:
-            type(self)._active_connections += 1
-        self.logger.debug(f"Checked out, active: {self.active_connections}")
+            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}")
 
     def _on_checkin(self, *args) -> None:
+        """Event listener for when a connection is returned to the pool."""
         with self._registry_lock:
-            type(self)._active_connections = max(type(self)._active_connections - 1, 0)
-        self.logger.debug(f"Checked in, active: {self.active_connections}")
+            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:
-        return type(self)._active_connections
+        """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:
             with self.managed_connection() as conn:
                 conn.execute(text("SELECT 1"))
-            self.logger.debug("Connection OK")
+            self.logger.debug("Database connection validated successfully.")
         except OperationalError as e:
-            self.logger.error(f"Connection failed: {e}")
-            raise ValidationError(f"DB connection failed: {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, Any]:
+    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
@@ -135,68 +230,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")
+            raise RuntimeError("Session factory not initialized. Cannot get a session.")
         return self.session_factory()
 
     def _build_model(self) -> None:
-        """Dynamically build and assign the ORM model if table is set"""
+        """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"Model built for table: {self.table}")
+            self.logger.debug(f"Successfully built ORM model for table: {self.table}")
         except Exception as e:
-            self.logger.error(f"Model build failed: {e}")
-            raise ValidationError(f"Model construction error: {e}") from e
-
-    def dispose_idle_connections(self) -> int:
-        key = self._engine_key()
-        with self._registry_lock:
-            if self._engine_registry.get(key) is not self.engine:
-                self.logger.debug("Engine changed")
-                return 0
-            pool = self.engine.pool
-            if isinstance(pool, QueuePool):
-                count = pool.checkedin()
-                pool.dispose()
-                self.logger.debug(f"Disposed {count}")
-                return count
-            self.logger.warning(f"No idle dispose for {type(pool).__name__}")
-            return 0
-
-    def terminate_idle_connections(self, idle_seconds: int = 300) -> int:
-        terminated = 0
-        dialect = self.engine.dialect.name
-        with self.managed_connection() as conn:
-            if dialect == 'postgresql':
-                res = conn.execute(text(
-                    f"SELECT pg_terminate_backend(pid) FROM pg_stat_activity "
-                    f"WHERE state='idle' AND (now() - query_start) > interval '{idle_seconds} seconds' "
-                    f"AND pid<>pg_backend_pid()"
-                ))
-                terminated = res.rowcount
-            elif dialect == 'mysql':
-                for row in conn.execute(text("SHOW PROCESSLIST")):
-                    if row.Command == 'Sleep' and row.Time > idle_seconds:
-                        conn.execute(text(f"KILL {row.Id}"))
-                        terminated += 1
-            else:
-                self.logger.warning(f"Idle termination not supported: {dialect}")
-        self.logger.debug(f"Terminated {terminated}")
-        return terminated
-
-    def close(self) -> None:
-        with self._registry_lock:
-            key = self._engine_key()
-            if not self._owns_engine:
-                self.logger.debug("Not owner, skipping close")
-                return
-            if self._engine_registry.get(key) != self.engine:
-                self.logger.debug("Engine not in registry")
-                return
-            self.engine.dispose()
-            del self._engine_registry[key]
-            type(self)._active_connections = 0
-            self.logger.debug(f"Engine closed {key}")
+            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