sibi-dst 2025.8.6__py3-none-any.whl → 2025.8.8__py3-none-any.whl

sibi_dst/df_helper/backends/sqlalchemy/_db_connection.py

@@ -29,6 +29,7 @@ class SqlAlchemyConnectionConfig(BaseModel):
     connection_url: str
     table: Optional[str] = None
     debug: bool = False
+    logger_extra: Optional[Dict[str, Any]] = {"sibi_dst_component": __name__}
 
     # --- Pool Configuration ---
     pool_size: int = int(os.environ.get("DB_POOL_SIZE", 5))
@@ -99,10 +100,10 @@ class SqlAlchemyConnectionConfig(BaseModel):
                 self.engine = wrapper["engine"]
                 wrapper["ref_count"] += 1
                 if self.debug:
-                    self.logger.debug(f"Reusing engine. Ref count: {wrapper['ref_count']}.")
+                    self.logger.debug(f"Reusing engine. Ref count: {wrapper['ref_count']}.", extra=self.logger_extra)
             else:
                 if self.debug:
-                    self.logger.debug(f"Creating new engine for key: {self._engine_key_instance}")
+                    self.logger.debug(f"Creating new engine for key: {self._engine_key_instance}", extra=self.logger_extra)
                 try:
                     new_engine = create_engine(
                         self.connection_url,
@@ -121,7 +122,7 @@ class SqlAlchemyConnectionConfig(BaseModel):
                         "active_connections": 0,
                     }
                 except Exception as e:
-                    self.logger.error(f"Failed to create engine: {e}")
+                    self.logger.error(f"Failed to create engine: {e}", extra=self.logger_extra)
                     raise SQLAlchemyError(f"Engine creation failed: {e}") from e
 
     def close(self) -> None:
@@ -134,14 +135,14 @@ class SqlAlchemyConnectionConfig(BaseModel):
             key = self._engine_key_instance
             wrapper = _ENGINE_REGISTRY.get(key)
             if not wrapper:
-                self.logger.warning("Attempted to close a config whose engine is not in the registry.")
+                self.logger.warning("Attempted to close a config whose engine is not in the registry.", extra=self.logger_extra)
             else:
                 wrapper["ref_count"] -= 1
                 if self.debug:
-                    self.logger.debug(f"Closing connection. Ref count now {wrapper['ref_count']}.")
+                    self.logger.debug(f"Closing connection. Ref count now {wrapper['ref_count']}.", extra=self.logger_extra)
                 if wrapper["ref_count"] <= 0:
                     if self.debug:
-                        self.logger.debug(f"Disposing engine as reference count is zero. Key: {key}")
+                        self.logger.debug(f"Disposing engine as reference count is zero. Key: {key}", extra=self.logger_extra)
                     try:
                         wrapper["engine"].dispose()
                     finally:
@@ -177,9 +178,9 @@ class SqlAlchemyConnectionConfig(BaseModel):
             with self.managed_connection() as conn:
                 conn.execute(text("SELECT 1"))
             if self.debug:
-                self.logger.debug("Database connection validated successfully.")
+                self.logger.debug("Database connection validated successfully.", extra=self.logger_extra)
         except OperationalError as e:
-            self.logger.error(f"Database connection failed: {e}")
+            self.logger.error(f"Database connection failed: {e}", extra=self.logger_extra)
             raise ValueError(f"DB connection failed: {e}") from e
 
     @contextmanager
@@ -204,8 +205,8 @@ class SqlAlchemyConnectionConfig(BaseModel):
             builder = SqlAlchemyModelBuilder(self.engine, self.table)
             self.model = builder.build_model()
             if self.debug:
-                self.logger.debug(f"Successfully built ORM model for table: {self.table}")
+                self.logger.debug(f"Successfully built ORM model for table: {self.table}", extra=self.logger_extra)
         except Exception as e:
-            self.logger.error(f"Failed to build ORM model for table '{self.table}': {e}")
+            self.logger.error(f"Failed to build ORM model for table '{self.table}': {e}", extra=self.logger_extra)
             raise ValueError(f"Model construction failed for table '{self.table}': {e}") from e
 

sibi_dst/df_helper/backends/sqlalchemy/_io_dask.py

@@ -38,6 +38,7 @@ class SQLAlchemyDask(ManagedResource):
         "TIME": "object",
         "UUID": "object",
     }
+    logger_extra: Dict[str, Any] = {"sibi_dst_component": __name__}
 
     def __init__(
         self,
@@ -97,7 +98,7 @@ class SQLAlchemyDask(ManagedResource):
         max_overflow = _to_int(max_overflow_attr, 10)
 
         cap = max(1, pool_size + max_overflow - 1)
-        self.logger.debug(f"Using a Cap of {cap} from pool size of {pool_size} and max overflow of {max_overflow}.")
+        self.logger.debug(f"Using a Cap of {cap} from pool size of {pool_size} and max overflow of {max_overflow}.", extra=self.logger_extra)
         return max(1, cap)
 
     # ---------- meta ----------
@@ -140,25 +141,25 @@ class SQLAlchemyDask(ManagedResource):
                     break
             except SASQLTimeoutError:
                 if attempt < retry_attempts - 1:
-                    self.logger.warning(f"Connection pool limit reached. Retrying in {backoff} seconds...")
+                    self.logger.warning(f"Connection pool limit reached. Retrying in {backoff} seconds...", extra=self.logger_extra)
                     time.sleep(backoff)
                     backoff *= 2
                 else:
                     self.total_records = -1
-                    self.logger.error("Failed to get a connection from the pool after retries.", exc_info=True)
+                    self.logger.error("Failed to get a connection from the pool after retries.", exc_info=True, extra=self.logger_extra)
                     return self.total_records, dd.from_pandas(meta_df, npartitions=1)
             except OperationalError as oe:
                 if "timeout" in str(oe).lower() and attempt < retry_attempts - 1:
-                    self.logger.warning("Operational timeout, retrying…", exc_info=self.debug)
+                    self.logger.warning("Operational timeout, retrying…", exc_info=self.debug, extra=self.logger_extra)
                     time.sleep(backoff)
                     backoff *= 2
                     continue
                 self.total_records = -1
-                self.logger.error("OperationalError during count.", exc_info=True)
+                self.logger.error("OperationalError during count.", exc_info=True, extra=self.logger_extra)
                 return self.total_records, dd.from_pandas(meta_df, npartitions=1)
             except Exception as e:
                 self.total_records = -1
-                self.logger.error(f"Unexpected error during count: {e}", exc_info=True)
+                self.logger.error(f"Unexpected error during count: {e}", exc_info=True, extra=self.logger_extra)
                 return self.total_records, dd.from_pandas(meta_df, npartitions=1)
 
         self.total_records = int(total)
@@ -167,7 +168,7 @@ class SQLAlchemyDask(ManagedResource):
             super().close()
             return self.total_records, dd.from_pandas(meta_df, npartitions=1)
 
-        self.logger.debug(f"Total records to fetch: {total}. Chunk size: {self.chunk_size}.")
+        self.logger.debug(f"Total records to fetch: {total}. Chunk size: {self.chunk_size}.", extra=self.logger_extra)
 
         @dask.delayed
         def get_chunk(sql_query, chunk_offset):
@@ -181,6 +182,6 @@ class SQLAlchemyDask(ManagedResource):
         offsets = range(0, total, self.chunk_size)
         delayed_chunks = [get_chunk(query, off) for off in offsets]
         ddf = dd.from_delayed(delayed_chunks, meta=meta_df)
-        self.logger.debug(f"Created Dask DataFrame with {ddf.npartitions} partitions.")
+        self.logger.debug(f"{self.model.__name__} created Dask DataFrame with {ddf.npartitions} partitions.", extra=self.logger_extra)
         return self.total_records, ddf
 

sibi_dst/df_helper/backends/sqlalchemy/_load_from_db.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Tuple
+from typing import Any, Tuple, Dict
 
 import dask.dataframe as dd
 import pandas as pd
@@ -15,6 +15,7 @@ class SqlAlchemyLoadFromDb(ManagedResource):
     """
     Orchestrates loading data from a database using SQLAlchemy into a Dask DataFrame.
     """
+    logger_extra: Dict[str, Any] = {"sibi_dst_component": __name__}
 
     def __init__(
         self,
@@ -43,86 +44,13 @@ class SqlAlchemyLoadFromDb(ManagedResource):
                 verbose=self.verbose,
                 debug=self.debug,
             ) as loader:
-                self.logger.debug(f"SQLAlchemyDask loader initialized for model: {self.model.__name__}")
+                self.logger.debug(f"SQLAlchemyDask loader initialized for model: {self.model.__name__}", extra=self.logger_extra)
                 self.total_records, dask_df = loader.read_frame()
                 return self.total_records, dask_df
         except Exception as e:
             self.total_records = -1
-            self.logger.error(f"{self.model.__name__} Failed to build and load data: {e}", exc_info=True)
+            self.logger.error(f"{self.model.__name__} Failed to build and load data: {e}", exc_info=True, extra=self.logger_extra)
             # empty df with correct columns
             columns = [c.name for c in self.model.__table__.columns]
             return self.total_records, dd.from_pandas(pd.DataFrame(columns=columns), npartitions=1)
 
-# from __future__ import annotations
-#
-# from typing import Any
-#
-# import dask.dataframe as dd
-# import pandas as pd
-#
-# from sibi_dst.utils import ManagedResource
-# from sibi_dst.df_helper.core import ParamsConfig, QueryConfig
-# from ._db_connection import SqlAlchemyConnectionConfig
-# from ._io_dask import SQLAlchemyDask
-#
-# class SqlAlchemyLoadFromDb(ManagedResource):
-#     """
-#     Orchestrates loading data from a database using SQLAlchemy into a Dask
-#     DataFrame by configuring and delegating to the SQLAlchemyDask loader.
-#     """
-#
-#     def __init__(
-#             self,
-#             plugin_sqlalchemy: SqlAlchemyConnectionConfig,
-#             plugin_query: QueryConfig = None,
-#             plugin_params: ParamsConfig = None,
-#             **kwargs,
-#     ):
-#         """
-#         Initializes the loader with all necessary configurations.
-#
-#         Args:
-#             plugin_sqlalchemy: The database connection configuration object.
-#             plugin_query: The query configuration object.
-#             plugin_params: The parameters and filters configuration object.
-#             logger: An optional logger instance.
-#             **kwargs: Must contain 'index_column' for Dask partitioning.
-#         """
-#         super().__init__(**kwargs)
-#         self.db_connection = plugin_sqlalchemy
-#         self.model = self.db_connection.model
-#         self.engine = self.db_connection.engine
-#         self.query_config = plugin_query
-#         self.params_config = plugin_params
-#         self.chunk_size = kwargs.get("chunk_size", self.params_config.df_params.get("chunk_size", 1000))
-#         self.total_records = -1 # Initialize total_records to -1 to indicate no records loaded yet
-#
-#     def build_and_load(self) -> tuple[int | Any, Any] | dd.DataFrame:
-#         """
-#         Builds and loads a Dask DataFrame from a SQLAlchemy source.
-#
-#         This method is stateless and returns the DataFrame directly.
-#
-#         Returns:
-#             A Dask DataFrame containing the queried data or an empty,
-#             correctly structured DataFrame if the query fails or returns no results.
-#         """
-#         try:
-#             # Instantiate and use the low-level Dask loader
-#             with SQLAlchemyDask(model=self.model,filters=self.params_config.filters if self.params_config else {},
-#                 engine=self.engine,
-#                 chunk_size=self.chunk_size,
-#                 logger=self.logger,
-#                 verbose=self.verbose,
-#                 debug=self.debug) as sqlalchemy_dask_loader:
-#                 self.logger.debug(f"SQLAlchemyDask loader initialized for model: {self.model.__name__}")
-#                 # Create the lazy DataFrame and read a record count
-#                 # if total_records less than 0, it means an error occurred during the loading process
-#                 self.total_records, dask_df = sqlalchemy_dask_loader.read_frame()
-#                 return self.total_records, dask_df
-#         except Exception as e:
-#             self.total_records = -1
-#             self.logger.error(f"{self.model.__name__} Failed to build and load data: {e}", exc_info=True)
-#             # Return an empty dataframe with the correct schema on failure
-#             columns = [c.name for c in self.model.__table__.columns]
-#             return self.total_records, dd.from_pandas(pd.DataFrame(columns=columns), npartitions=1)

sibi_dst/df_helper/backends/sqlalchemy/_sql_model_builder.py

@@ -48,107 +48,3 @@ class SqlAlchemyModelBuilder:
             return f"{sane_name}_field"
         return sane_name
 
-# import re
-# import keyword
-# import threading
-# from sqlalchemy import MetaData, Engine
-# from sqlalchemy.orm import DeclarativeBase
-#
-#
-# class Base(DeclarativeBase):
-#     """Shared declarative base for all ORM models."""
-#     pass
-#
-#
-# apps_label = "datacubes.models"
-#
-#
-# class SqlAlchemyModelBuilder:
-#     """
-#     Builds a single SQLAlchemy ORM model from a specific database table.
-#     This class is thread-safe and caches reflected table metadata to
-#     improve performance across multiple instantiations.
-#     """
-#     _lock = threading.Lock()
-#     _metadata_cache: dict[str, MetaData] = {}
-#
-#     def __init__(self, engine: Engine, table_name: str):
-#         """
-#         Initializes the model builder for a specific table.
-#
-#         Args:
-#             engine: The SQLAlchemy engine connected to the database.
-#             table_name: The name of the table to generate the model for.
-#         """
-#         self.engine = engine
-#         self.table_name = table_name
-#         self.class_name = self._normalize_class_name(self.table_name)
-#
-#         engine_key = str(engine.url)
-#
-#         # ✅ REFACTOR: Acquire lock to make cache access and creation atomic,
-#         # preventing a race condition between multiple threads.
-#         with self._lock:
-#             if engine_key not in self._metadata_cache:
-#                 self._metadata_cache[engine_key] = MetaData()
-#             self.metadata = self._metadata_cache[engine_key]
-#
-#     def build_model(self) -> type:
-#         """
-#         Builds and returns a database model class for the specified table.
-#         This process is atomic and thread-safe.
-#
-#         Raises:
-#             ValueError: If the specified table does not exist in the database.
-#         Returns:
-#             The dynamically created ORM model class.
-#         """
-#         with self._lock:
-#             # NOTE: Using a private SQLAlchemy API. This is a performance
-#             # optimization but may break in future versions of the library.
-#             registered_model = Base.registry._class_registry.get(self.class_name)
-#             if registered_model:
-#                 return registered_model
-#
-#             # Check if the table's schema is in our metadata cache
-#             table = self.metadata.tables.get(self.table_name)
-#
-#             # If not cached, reflect it from the database
-#             if table is None:
-#                 self.metadata.reflect(bind=self.engine, only=[self.table_name])
-#                 table = self.metadata.tables.get(self.table_name)
-#
-#             if table is None:
-#                 raise ValueError(
-#                     f"Table '{self.table_name}' does not exist in the database."
-#                 )
-#
-#             # Create the model class dynamically.
-#             attrs = {
-#                 "__tablename__": table.name,
-#                 "__table__": table,
-#                 "__module__": apps_label,
-#             }
-#             model = type(self.class_name, (Base,), attrs)
-#
-#             return model
-#
-#     @staticmethod
-#     def _normalize_class_name(table_name: str) -> str:
-#         """Converts a snake_case table_name to a CamelCase class name."""
-#         return "".join(word.capitalize() for word in table_name.split("_"))
-#
-#     @staticmethod
-#     def _normalize_column_name(column_name: str) -> str:
-#         """
-#         Sanitizes a column name to be a valid Python identifier.
-#         (Kept for utility, though not used in the final model creation).
-#         """
-#         sane_name = re.sub(r"\W", "_", column_name)
-#         sane_name = re.sub(r"^\d", r"_\g<0>", sane_name)
-#
-#         if keyword.iskeyword(sane_name):
-#             return f"{sane_name}_field"
-#         return sane_name
-#
-#

sibi_dst/utils/async_utils.py

@@ -0,0 +1,12 @@
+import asyncio
+import dask.dataframe as dd
+
+
+def is_dask_dataframe(df):
+    """Check if the given object is a Dask DataFrame."""
+    return isinstance(df, dd.DataFrame)
+
+async def to_thread(func, *args, **kwargs):
+    """Explicit helper to keep code clear where we hop off the event loop."""
+    return await asyncio.to_thread(func, *args, **kwargs)
+

sibi_dst/utils/boilerplate/__init__.py

@@ -0,0 +1,6 @@
+from .base_data_artifact import BaseDataArtifact
+from .base_data_cube import BaseDataCube
+
+__all__ = ["BaseDataCube",
+           "BaseDataArtifact"
+           ]

sibi_dst/utils/boilerplate/base_data_artifact.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import asyncio
+from typing import Any, Dict, Mapping, Optional, Type, Union
+from datetime import date, datetime
+
+import pandas as pd
+import dask.dataframe as dd
+from sibi_dst.df_helper import ParquetArtifact
+
+
+DateLike = Union[str, date, datetime, None]
+
+
+def _validate_and_format_date(name: str, value: DateLike) -> Optional[str]:
+    """
+    Normalize date-like input into a canonical string '%Y-%m-%d'.
+
+    - None -> None
+    - str/date/datetime -> parse with pandas.to_datetime, take .date(), return '%Y-%m-%d'
+    - else -> TypeError
+    """
+    if value is None:
+        return None
+    if isinstance(value, (str, date, datetime)):
+        try:
+            return pd.to_datetime(value).date().strftime("%Y-%m-%d")
+        except Exception as e:
+            raise ValueError(f"{name} must be a valid date, got {value!r}") from e
+    raise TypeError(f"{name} must be str, date, datetime, or None; got {type(value)}")
+
+
+class BaseDataArtifact(ParquetArtifact):
+    """
+    Base class for Parquet artifacts with optional date window.
+
+    Dates are always stored as strings in '%Y-%m-%d' format.
+    """
+
+    config: Mapping[str, Any] = {}
+
+    parquet_start_date: Optional[str]
+    parquet_end_date: Optional[str]
+    data_wrapper_class: Optional[Type[Any]]
+    class_params: Dict[str, Any]
+    df: Union[pd.DataFrame | dd.DataFrame] = None
+
+    def __init__(
+        self,
+        **kwargs: Any,
+    ) -> None:
+        merged = {**self.config, **kwargs}
+        super().__init__(**merged)
+
+        # Normalize and store as canonical strings
+        self.parquet_start_date = _validate_and_format_date("parquet_start_date", merged.get("parquet_start_date", None))
+        self.parquet_end_date   = _validate_and_format_date("parquet_end_date", merged.get("parquet_end_date", None))
+
+        self.data_wrapper_class = merged.get("data_wrapper_class", None)
+        self.class_params = merged.get("class_params", None) or {
+            "debug": self.debug,
+            "logger": self.logger,
+            "fs": self.fs,
+            "verbose": getattr(self, "verbose", False),
+        }
+
+        # Ordering check
+        if self.parquet_start_date and self.parquet_end_date:
+            if self.parquet_start_date > self.parquet_end_date:
+                raise ValueError(
+                    f"parquet_start_date {self.parquet_start_date} "
+                    f"cannot be after parquet_end_date {self.parquet_end_date}"
+                )
+
+    # -------- Optional hooks --------
+
+    def before_load(self, **kwargs: Any) -> None: return None
+    def after_load(self, **kwargs: Any) -> None: return None
+    async def abefore_load(self, **kwargs: Any) -> None: return None
+    async def aafter_load(self, **kwargs: Any) -> None: return None
+
+    # -------- Public API --------
+
+    def load(self, **kwargs: Any):
+        self.before_load(**kwargs)
+        self.df = super().load(**kwargs)
+        self.after_load(**kwargs)
+        return self.df
+
+    async def aload(self, **kwargs: Any):
+        await self.abefore_load(**kwargs)
+        df = await asyncio.to_thread(super().load, **kwargs)
+        self.df = df
+        await self.aafter_load(**kwargs)
+        return self.df
+
+    def has_date_window(self) -> bool:
+        return bool(self.parquet_start_date or self.parquet_end_date)
+
+    def date_window(self) -> tuple[Optional[str], Optional[str]]:
+        return self.parquet_start_date, self.parquet_end_date
+
+    def to_params(self) -> Dict[str, Any]:
+        return {
+            "parquet_start_date": self.parquet_start_date,
+            "parquet_end_date": self.parquet_end_date,
+            "data_wrapper_class": self.data_wrapper_class,
+            "class_params": dict(self.class_params),
+        }
+

sibi_dst/utils/boilerplate/base_data_cube.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from typing import Union
+import dask.dataframe as dd
+import pandas as pd
+
+from sibi_dst.df_helper import DfHelper
+
+
+class BaseDataCube(DfHelper):
+    """
+    Base cube with sync/async load hooks.
+
+    Subclasses *may* override:
+      - fix_data(self, **kwargs): synchronous, local transforms
+      - async afix_data(self, **kwargs): asynchronous transforms (I/O, awaits)
+
+    Semantics:
+      - load() -> runs fix_data() if defined
+      - aload() -> runs afix_data() if subclass overrides it, else fix_data()
+    """
+    df: Union[dd.DataFrame, pd.DataFrame, None] = None
+    config: dict = {}
+
+    def __init__(self, **kwargs):
+        # kwargs override class config
+        kwargs = {**self.config, **kwargs}
+        super().__init__(**kwargs)
+
+    # -------------------- optional hooks --------------------
+
+    def fix_data(self, **kwargs) -> None:
+        """Optional sync transform hook. Override in subclasses if needed."""
+        return None
+
+    async def afix_data(self, **kwargs) -> None:
+        """Optional async transform hook. Override in subclasses if needed."""
+        return None
+
+    # -------------------- internals --------------------
+
+    def _has_data(self) -> bool:
+        """Check if dataframe has rows; avoids hidden heavy ops where possible."""
+        if self.df is None:
+            return False
+        if isinstance(self.df, dd.DataFrame):
+            return bool(self.df.shape[0].compute() > 0)
+        return not self.df.empty
+
+    def _afix_data_is_overridden(self) -> bool:
+        """Check if subclass provided its own afix_data."""
+        return self.__class__.afix_data is not BaseDataCube.afix_data
+
+    def _fix_data_is_overridden(self) -> bool:
+        """Check if subclass provided its own fix_data."""
+        return self.__class__.fix_data is not BaseDataCube.fix_data
+
+    # -------------------- public API --------------------
+
+    def load(self, **kwargs):
+        """Sync load path with optional fix_data hook."""
+        self.df = super().load(**kwargs)
+        if self._has_data() and self._fix_data_is_overridden():
+            self.fix_data()
+        elif not self._has_data():
+            self.logger.debug(f"No data was found by {self.__class__.__name__} loader")
+        return self.df
+
+    async def aload(self, **kwargs):
+        """Async load path with optional afix_data/fix_data hook."""
+        self.df = await super().aload(**kwargs)
+        if self._has_data():
+            if self._afix_data_is_overridden():
+                await self.afix_data()
+            elif self._fix_data_is_overridden():
+                self.fix_data()
+        else:
+            self.logger.debug(f"No data was found by {self.__class__.__name__} loader")
+        return self.df