sibi-dst 0.3.62__tar.gz → 0.3.64__tar.gz

{sibi_dst-0.3.62 → sibi_dst-0.3.64}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sibi-dst
-Version: 0.3.62
+Version: 0.3.64
 Summary: Data Science Toolkit
 Author: Luis Valverde
 Author-email: lvalverdeb@gmail.com

{sibi_dst-0.3.62 → sibi_dst-0.3.64}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sibi-dst"
-version = "0.3.62"
+version = "0.3.64"
 description = "Data Science Toolkit"
 authors = ["Luis Valverde <lvalverdeb@gmail.com>"]
 readme = "README.md"

sibi_dst-0.3.64/sibi_dst/df_helper/_df_helper.py

@@ -0,0 +1,230 @@
+from __future__ import annotations
+
+import warnings
+from typing import Any, Dict, Optional, Union, TypeVar
+
+import dask.dataframe as dd
+import fsspec
+import pandas as pd
+from pydantic import BaseModel
+
+from sibi_dst.df_helper.core import QueryConfig, ParamsConfig, FilterHandler
+from sibi_dst.utils import Logger, ParquetSaver, ClickHouseWriter
+from .backends.http import HttpConfig
+from .backends.parquet import ParquetConfig
+from .backends.sqlalchemy import SqlAlchemyConnectionConfig, SqlAlchemyLoadFromDb
+
+warnings.filterwarnings("ignore")
+T = TypeVar("T", bound=BaseModel)
+
+
+# --- Backend Strategy Pattern Implementation ---
+
+class BaseBackend:
+    """Abstract base class defining clear sync and async loading interfaces."""
+
+    def __init__(self, helper: DfHelper):
+        self.helper = helper
+        self.logger = helper.logger
+
+    def load(self, **options) -> dd.DataFrame | pd.DataFrame:
+        """Synchronous data loading method. Must be implemented by sync backends."""
+        raise NotImplementedError(f"Backend '{self.__class__.__name__}' does not support synchronous loading.")
+
+    async def aload(self, **options) -> dd.DataFrame | pd.DataFrame:
+        """Asynchronous data loading method. By default, it calls the sync version."""
+        return self.load(**options)
+
+
+class SqlAlchemyBackend(BaseBackend):
+    def load(self, **options) -> dd.DataFrame:
+        try:
+            # Process incoming filter options into the ParamsConfig object
+            if options and hasattr(self.helper._backend_params, 'parse_params'):
+                self.helper._backend_params.parse_params(options)
+
+            db_loader = SqlAlchemyLoadFromDb(
+                plugin_sqlalchemy=self.helper.backend_db_connection,
+                plugin_query=self.helper._backend_query,
+                plugin_params=self.helper._backend_params,
+                logger=self.logger
+            )
+            return db_loader.build_and_load()
+        except Exception as e:
+            self.logger.error(f"Failed to load data from sqlalchemy: {e}", exc_info=self.debug)
+            return dd.from_pandas(pd.DataFrame(), npartitions=1)
+
+
+class ParquetBackend(BaseBackend):
+    """This backend is also purely synchronous."""
+
+    def load(self, **options) -> dd.DataFrame | pd.DataFrame:
+        try:
+            df = self.helper.backend_parquet.load_files()
+            if options and df is not None:
+                df = FilterHandler('dask', self.logger).apply_filters(df, filters=options)
+            return df
+        except Exception as e:
+            self.logger.error(f"Failed to load data from parquet: {e}", exc_info=self.debug)
+            return dd.from_pandas(pd.DataFrame(), npartitions=1)
+
+
+class HttpBackend(BaseBackend):
+    """This backend is purely asynchronous."""
+
+    def load(self, **options) -> dd.DataFrame | pd.DataFrame:
+        # This will correctly fail by raising NotImplementedError from the base class.
+        return self.helper.backend_http.fetch_data(**options)
+
+    async def aload(self, **options) -> Union[pd.DataFrame, dd.DataFrame]:
+        if not self.helper.backend_http:
+            self.logger.warning("HTTP plugin not configured properly.")
+            return dd.from_pandas(pd.DataFrame(), npartitions=1)
+        return await self.helper.backend_http.fetch_data(**options)
+
+
+# --- Main DfHelper Facade Class ---
+
+class DfHelper:
+    """
+    A reusable utility for loading data. It provides both sync (`load`) and
+    async (`aload`) methods to accommodate different backends.
+    """
+    _BACKEND_STRATEGIES = {
+        'sqlalchemy': SqlAlchemyBackend,
+        'parquet': ParquetBackend,
+        'http': HttpBackend,
+    }
+
+    default_config: Dict = None
+
+    def __init__(self, backend='sqlalchemy', **kwargs):
+        self.default_config = self.default_config or {}
+        kwargs = {**self.default_config.copy(), **kwargs}
+        self.backend = backend
+        self.debug = kwargs.get("debug", False)
+        self.logger = kwargs.get("logger", Logger.default_logger(logger_name=self.__class__.__name__))
+        self.logger.set_level(Logger.DEBUG if self.debug else Logger.INFO)
+        self.fs = kwargs.get("fs", fsspec.filesystem('file'))
+        kwargs.setdefault("fs", self.fs)
+        kwargs.setdefault("logger", self.logger)
+        self._backend_query = self._get_config(QueryConfig, kwargs)
+        self._backend_params = self._get_config(ParamsConfig, kwargs)
+        self.backend_db_connection: Optional[SqlAlchemyConnectionConfig] = None
+        self.backend_parquet: Optional[ParquetConfig] = None
+        self.backend_http: Optional[HttpConfig] = None
+
+        if self.backend == 'sqlalchemy':
+            self.backend_db_connection = self._get_config(SqlAlchemyConnectionConfig, kwargs)
+        elif self.backend == 'parquet':
+            self.backend_parquet = self._get_config(ParquetConfig, kwargs)
+        elif self.backend == 'http':
+            self.backend_http = self._get_config(HttpConfig, kwargs)
+
+        strategy_class = self._BACKEND_STRATEGIES.get(self.backend)
+        if not strategy_class: raise ValueError(f"Unsupported backend: {self.backend}")
+        self.backend_strategy = strategy_class(self)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self._cleanup()
+
+    def _cleanup(self):
+        active_config = getattr(self, f"backend_{self.backend}", None)
+        if active_config and hasattr(active_config, "close"):
+            self.logger.debug(f"Closing resources for '{self.backend}' backend.")
+            active_config.close()
+
+    def _get_config(self, model: T, kwargs: Dict[str, Any]) -> T:
+        recognized_keys = set(model.model_fields.keys())
+        model_kwargs = {k: kwargs[k] for k in recognized_keys if k in kwargs}
+        return model(**model_kwargs)
+
+    def load(self, as_pandas=False, **options) -> Union[pd.DataFrame, dd.DataFrame]:
+        """Loads data synchronously. Fails if backend is async-only."""
+        self.logger.debug(f"Loading data from {self.backend} backend with options: {options}")
+        df = self.backend_strategy.load(**options)
+        df = self._process_loaded_data(df)
+        df = self._post_process_df(df)
+        return df.compute() if as_pandas else df
+
+    async def aload(self, as_pandas=False, **options) -> Union[pd.DataFrame, dd.DataFrame]:
+        """Loads data asynchronously from any backend."""
+        df = await self.backend_strategy.aload(**options)
+        df = self._process_loaded_data(df)
+        df = self._post_process_df(df)
+        return df.compute() if as_pandas else df
+
+    def _post_process_df(self, df: dd.DataFrame) -> dd.DataFrame:
+        df_params = self._backend_params.df_params
+        if not df_params: return df
+        fieldnames, column_names, index_col = (df_params.get("fieldnames"), df_params.get("column_names"),
+                                               df_params.get("index_col"))
+        if not any([fieldnames, column_names, index_col]): return df
+        self.logger.debug("Post-processing DataFrame.")
+        if fieldnames:
+            valid_fieldnames = [f for f in fieldnames if f in df.columns]
+            if len(valid_fieldnames) < len(fieldnames): self.logger.warning(
+                f"Missing columns for filtering: {set(fieldnames) - set(valid_fieldnames)}")
+            df = df[valid_fieldnames]
+        if column_names:
+            if len(df.columns) != len(column_names): raise ValueError(
+                f"Length mismatch: DataFrame has {len(df.columns)} columns, but {len(column_names)} names were provided.")
+            df = df.rename(columns=dict(zip(df.columns, column_names)))
+        if index_col:
+            if index_col not in df.columns: raise ValueError(f"Index column '{index_col}' not found in DataFrame.")
+            df = df.set_index(index_col)
+        return df
+
+    def _process_loaded_data(self, df: dd.DataFrame) -> dd.DataFrame:
+        field_map = self._backend_params.field_map or {}
+        if not isinstance(field_map, dict) or not field_map: return df
+        if hasattr(df, 'npartitions') and df.npartitions == 1 and not len(df.head(1)): return df
+        self.logger.debug("Processing loaded data...")
+        rename_mapping = {k: v for k, v in field_map.items() if k in df.columns}
+        if rename_mapping: df = df.rename(columns=rename_mapping)
+        return df
+
+    def save_to_parquet(self, df: dd.DataFrame, parquet_filename: str, **kwargs):
+        if hasattr(df, 'npartitions') and df.npartitions == 1 and not len(df.head(1)):
+            self.logger.warning("Cannot save to parquet; DataFrame is empty.")
+            return
+        fs = kwargs.pop('fs', self.fs)
+        path = kwargs.pop('parquet_storage_path', self.backend_parquet.parquet_storage_path)
+        ParquetSaver(df, path, self.logger, fs).save_to_parquet(parquet_filename)
+        self.logger.debug(f"Parquet saved to {parquet_filename} in path: {path}.")
+
+    def save_to_clickhouse(self, df: dd.DataFrame, **credentials):
+        if hasattr(df, 'npartitions') and df.npartitions == 1 and not len(df.head(1)):
+            self.logger.warning("Cannot write to ClickHouse; DataFrame is empty.")
+            return
+        ClickHouseWriter(self.logger, **credentials).save_to_clickhouse(df)
+        self.logger.debug("Save to ClickHouse completed.")
+
+    def load_period(self, dt_field: str, start: str, end: str, **kwargs) -> Union[pd.DataFrame, dd.DataFrame]:
+        """Synchronous convenience method for loading a date range."""
+        final_kwargs = self._prepare_period_filters(dt_field, start, end, **kwargs)
+        return self.load(**final_kwargs)
+
+    async def aload_period(self, dt_field: str, start: str, end: str, **kwargs) -> Union[pd.DataFrame, dd.DataFrame]:
+        """Asynchronous convenience method for loading a date range."""
+        final_kwargs = self._prepare_period_filters(dt_field, start, end, **kwargs)
+        return await self.aload(**final_kwargs)
+
+    def _prepare_period_filters(self, dt_field: str, start: str, end: str, **kwargs) -> dict:
+        start_date, end_date = pd.to_datetime(start).date(), pd.to_datetime(end).date()
+        if start_date > end_date: raise ValueError("'start' date cannot be later than 'end' date.")
+        field_map = self._backend_params.field_map or {}
+        reverse_map = {v: k for k, v in field_map.items()} if field_map else {}
+        if len(reverse_map) != len(field_map): self.logger.warning(
+            "field_map values are not unique; reverse mapping may be unreliable.")
+        mapped_field = reverse_map.get(dt_field, dt_field)
+        if start_date == end_date:
+            kwargs[f"{mapped_field}__date"] = start_date
+        else:
+            kwargs[f"{mapped_field}__date__range"] = [start_date, end_date]
+        self.logger.debug(f"Period load generated filters: {kwargs}")
+        return kwargs
+

{sibi_dst-0.3.62 → sibi_dst-0.3.64}/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-0.3.62 → sibi_dst-0.3.64}/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-0.3.64/sibi_dst/df_helper/backends/sqlalchemy/_db_connection.py

@@ -0,0 +1,248 @@
+from __future__ import annotations
+
+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, 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):
+    """
+    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
+    debug: bool = False
+
+    # --- Pool Configuration ---
+    pool_size: int = 5
+    max_overflow: int = 10
+    pool_timeout: int = 30
+    pool_recycle: int = 1800
+    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
+    session_factory: Optional[sessionmaker] = None
+
+    # --- 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()
+
+    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
+    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 _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()
+        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__)
+        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_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 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']}.")
+            else:
+                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:
+        """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:
+            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:
+            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:
+            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}")
+            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 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}")
+        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