sibi-flux 2025.12.0__py3-none-any.whl

sibi_flux/df_helper/backends/sqlalchemy/_load_from_db.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+from typing import Any, Tuple, Dict, Optional
+
+import dask.dataframe as dd
+import pandas as pd
+
+from sibi_flux.core import ManagedResource
+from sibi_flux.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.
+    """
+
+    logger_extra: Dict[str, Any] = {"sibi_flux_component": __name__}
+
+    def __init__(
+        self,
+        plugin_sqlalchemy: SqlAlchemyConnectionConfig,
+        plugin_query: Optional[QueryConfig] = None,
+        plugin_params: Optional[ParamsConfig] = None,
+        **kwargs: Any,
+    ):
+        super().__init__(**kwargs)
+        self.db_connection = plugin_sqlalchemy
+
+        # Access properties safely (triggers lazy load if needed)
+        self.model = self.db_connection.model
+        self.engine = self.db_connection.engine
+
+        self.query_config = plugin_query
+        self.params_config = plugin_params
+
+        # Safe extraction of tuning parameters
+        df_params = self.params_config.df_params if self.params_config else None
+
+        # 1. Chunk Size (Default: 50k)
+        # Priority: kwargs > df_params > default
+        self.chunk_size = int(
+            kwargs.get("chunk_size", self._safe_get(df_params, "chunk_size", 50_000))
+        )
+
+        # 2. Index Column (Critical for Range Pagination)
+        # Priority: kwargs > df_params > None
+        self.db_index = kwargs.get(
+            "db_index", self._safe_get(df_params, "db_index", None)
+        )
+
+        # 3. Safety Limit
+        # Priority: kwargs > df_params > None
+        self.limit = kwargs.get("limit", self._safe_get(df_params, "limit", None))
+        if self.limit is not None:
+            self.limit = int(self.limit)
+
+        # 4. Engine Options
+        self.execution_options = kwargs.get("execution_options", {})
+
+        self.total_records = -1
+
+    @staticmethod
+    def _safe_get(obj: Any, key: str, default: Any) -> Any:
+        """Helper to extract values from Dict, Pydantic, or None safely."""
+        if obj is None:
+            return default
+        if isinstance(obj, dict):
+            return obj.get(key, default)
+        return getattr(obj, key, default)
+
+    def build_and_load(self) -> Tuple[int, dd.DataFrame]:
+        try:
+            # Determine pagination strategy automatically
+            # If we have an index_col, try 'range' (fast). Otherwise 'offset' (safe).
+            pagination_mode = "range" if self.db_index else "offset"
+
+            # Extract filters safely
+            filters = self.params_config.filters if self.params_config else {}
+
+            # Use Context Manager to ensure SQLAlchemyDask resources are cleaned up
+            with SQLAlchemyDask(
+                model=self.model,
+                filters=filters,
+                engine=self.engine,
+                chunk_size=self.chunk_size,
+                pagination=pagination_mode,
+                db_index=self.db_index,
+                limit=self.limit,
+                execution_options=self.execution_options,
+                logger=self.logger,
+                verbose=self.verbose,
+                debug=self.debug,
+            ) as loader:
+
+                self.logger.debug(
+                    f"SQLAlchemyDask initialized for {self.model.__name__} "
+                    f"(strategy={pagination_mode}, chunk={self.chunk_size})",
+                    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 load: {e}",
+                exc_info=True,
+                extra=self.logger_extra,
+            )
+            # Return empty dataframe structure on failure to prevent crashes downstream
+            try:
+                # Try to inspect columns from model to return empty DF with correct schema
+                columns = [c.name for c in self.model.__table__.columns]
+            except Exception:
+                columns = []
+
+            return -1, dd.from_pandas(pd.DataFrame(columns=columns), npartitions=1)
+
+    def _cleanup(self) -> None:
+        """
+        Clean up instance references to prevent memory leaks.
+        Note: The engine is owned by DfHelper/ConnectionConfig, not us.
+        """
+        try:
+            self.logger.debug(f"Cleaning up {self.__class__.__name__} refs")
+            self.db_connection = None
+            self.engine = None
+            self.model = None
+        except Exception as e:
+            self.logger.warning(f"Error during cleanup: {e}", extra=self.logger_extra)

sibi_flux/df_helper/backends/sqlalchemy/_model_registry.py

@@ -0,0 +1,239 @@
+from __future__ import annotations
+
+import hashlib
+import sys
+import types
+import threading
+from typing import Dict, Optional, Tuple
+
+from sqlalchemy import MetaData, Table
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import DeclarativeBase
+
+
+class Base(DeclarativeBase):
+    """Shared declarative base for all ORM models."""
+
+    pass
+
+
+# Default module label for generated classes
+apps_label = "datacubes.models"
+
+
+class ModelRegistry:
+    """
+    Thread-safe registry that reflects tables once per (engine, schema) and
+    returns a single mapped class per (engine, schema, table).
+    """
+
+    def __init__(self) -> None:
+        self._metadata_cache: Dict[Tuple[str, Optional[str]], MetaData] = {}
+        self._model_cache: Dict[Tuple[str, Optional[str], str], type] = {}
+        self._lock = threading.RLock()
+        self._md_locks: Dict[Tuple[str, Optional[str]], threading.Lock] = {}
+
+    # ---------- key helpers ----------
+    @staticmethod
+    def _engine_key(engine: Engine) -> str:
+        # Use URL string as key. Note: This includes passwords if not masked.
+        # For internal cache keys, this is fine.
+        return str(engine.url)
+
+    @staticmethod
+    def _qualified_key(schema: Optional[str], table: str) -> str:
+        return f"{schema}.{table}" if schema else table
+
+    @staticmethod
+    def _split_schema_and_table(name: str) -> Tuple[Optional[str], str]:
+        if "." in name:
+            s, t = name.split(".", 1)
+            return (s or None), t
+        return None, name
+
+    # ---------- class name helpers ----------
+    @staticmethod
+    def _normalize_class_name(table_name: str) -> str:
+        return "".join(part.capitalize() for part in table_name.split("_"))
+
+    @staticmethod
+    def _short_hash(*parts: str, length: int = 8) -> str:
+        h = hashlib.sha1("|".join(parts).encode("utf-8")).hexdigest()
+        return h[:length]
+
+    def _is_class_name_taken(self, class_name: str, module_label: str) -> bool:
+        # Check Python's module system directly to see if the name is taken
+        if module_label in sys.modules:
+            return hasattr(sys.modules[module_label], class_name)
+        return False
+
+    def _find_existing_model_for_table(self, tbl: Table) -> Optional[type]:
+        # Scan SQLAlchemy's global mapper registry
+        for mapper in list(Base.registry.mappers):
+            try:
+                mapped_cls = mapper.class_
+                mapped_tbl = getattr(mapped_cls, "__table__", None)
+                if mapped_tbl is tbl:
+                    return mapped_cls
+                # Check for equivalent table (same name/schema)
+                if isinstance(mapped_tbl, Table):
+                    if (mapped_tbl.schema == tbl.schema) and (
+                        mapped_tbl.name == tbl.name
+                    ):
+                        return mapped_cls
+            except Exception:
+                continue
+        return None
+
+    # ---------- metadata helpers ----------
+    def _get_or_create_metadata(self, ekey: str, schema: Optional[str]) -> MetaData:
+        md_key = (ekey, schema)
+        with self._lock:
+            md = self._metadata_cache.get(md_key)
+            if md is None:
+                md = MetaData(schema=schema)
+                self._metadata_cache[md_key] = md
+            return md
+
+    def _get_or_create_md_lock(
+        self, md_key: Tuple[str, Optional[str]]
+    ) -> threading.Lock:
+        with self._lock:
+            lock = self._md_locks.get(md_key)
+            if lock is None:
+                lock = threading.Lock()
+                self._md_locks[md_key] = lock
+            return lock
+
+    def _register_as_module_attribute(
+        self, cls: type, module_name: str, class_name: str
+    ) -> None:
+        """
+        Injects the class into sys.modules so it can be imported/pickled.
+        Ensures parent packages exist (e.g., 'datacubes' for 'datacubes.models').
+        """
+        if module_name not in sys.modules:
+            # Create the module
+            mod = types.ModuleType(module_name)
+            sys.modules[module_name] = mod
+
+            # Ensure parent packages exist logic
+            if "." in module_name:
+                parent_name, child_name = module_name.rsplit(".", 1)
+                if parent_name not in sys.modules:
+                    sys.modules[parent_name] = types.ModuleType(parent_name)
+
+                # Link child to parent so `import parent; parent.child` works
+                setattr(sys.modules[parent_name], child_name, mod)
+
+        # Register the class on the module
+        setattr(sys.modules[module_name], class_name, cls)
+
+    # def _register_as_module_attribute(self, cls: type, module_name: str, class_name: str) -> None:
+    #     """
+    #     Injects the class into sys.modules so it can be imported/pickled.
+    #     """
+    #     if module_name not in sys.modules:
+    #         # Create a dummy module on the fly
+    #         sys.modules[module_name] = types.ModuleType(module_name)
+    #
+    #     # Register the class on the module
+    #     setattr(sys.modules[module_name], class_name, cls)
+
+    # ---------- public API ----------
+    def get_model(
+        self,
+        engine: Engine,
+        table_name: str,
+        *,
+        refresh: bool = False,
+        schema: Optional[str] = None,
+        module_label: Optional[str] = None,
+        prefer_stable_names: bool = True,
+    ) -> type:
+        s2, tname = self._split_schema_and_table(table_name)
+        schema = schema if schema is not None else s2
+        ekey = self._engine_key(engine)
+        model_key = (ekey, schema, tname)
+        md_key = (ekey, schema)
+        module_label = module_label or apps_label
+
+        if refresh:
+            with self._lock:
+                self._model_cache.pop(model_key, None)
+                # Note: We don't drop metadata easily as it might be used by other models
+
+        # fast path: already cached model
+        with self._lock:
+            m = self._model_cache.get(model_key)
+            if m is not None:
+                return m
+
+        # ensure metadata and reflection are serialized per (engine, schema)
+        md = self._get_or_create_metadata(ekey, schema)
+        md_lock = self._get_or_create_md_lock(md_key)
+        qname = self._qualified_key(schema, tname)
+
+        # Reflection logic
+        tbl = md.tables.get(qname)
+        if tbl is None:
+            with md_lock:
+                # double-checked locking
+                tbl = md.tables.get(qname)
+                if tbl is None:
+                    # Reflect only this table
+                    md.reflect(bind=engine, only=[tname], schema=schema)
+                tbl = md.tables.get(qname)
+
+        if tbl is None:
+            raise ValueError(f"Table '{qname}' does not exist in the database.")
+
+        # If a mapped model for this Table already exists (anywhere), reuse it
+        reused = self._find_existing_model_for_table(tbl)
+        if reused is not None:
+            with self._lock:
+                self._model_cache[model_key] = reused
+            return reused
+
+        # pick class name
+        base_name = self._normalize_class_name(tname)
+        final_name = base_name
+
+        # Conflict resolution
+        if self._is_class_name_taken(base_name, module_label):
+            suffix = self._short_hash(ekey, schema or "", tname)
+            final_name = f"{base_name}_{suffix}"
+        elif not prefer_stable_names:
+            suffix = self._short_hash(ekey, schema or "", tname)
+            final_name = f"{base_name}_{suffix}"
+
+        # build the model
+        attrs = {
+            "__tablename__": tbl.name,
+            "__table__": tbl,
+            "__module__": module_label,
+        }
+
+        # Create the dynamic class
+        model_cls = type(final_name, (Base,), attrs)
+
+        # CRITICAL FIX: Make it pickleable by registering it
+        self._register_as_module_attribute(model_cls, module_label, final_name)
+
+        with self._lock:
+            self._model_cache[model_key] = model_cls
+        return model_cls
+
+    def clear(self) -> None:
+        with self._lock:
+            self._metadata_cache.clear()
+            self._model_cache.clear()
+            self._md_locks.clear()
+
+
+# Singleton
+_global_registry = ModelRegistry()
+
+
+def get_global_registry() -> ModelRegistry:
+    return _global_registry

sibi_flux/df_helper/backends/sqlalchemy/_sql_model_builder.py

@@ -0,0 +1,42 @@
+import keyword
+import re
+from typing import Type
+from sqlalchemy.engine import Engine
+
+from ._model_registry import get_global_registry, apps_label
+
+
+class SqlAlchemyModelBuilder:
+    """
+    Builds a single SQLAlchemy ORM model from a specific database table.
+    Delegates to the process-wide ModelRegistry for thread-safety and caching.
+    """
+
+    def __init__(self, engine: Engine, table_name: str):
+        self.engine = engine
+        self.table_name = table_name
+
+    def build_model(self) -> Type:
+        """Reflects the table and returns the mapped ORM class."""
+        registry = get_global_registry()
+
+        # The registry handles locking internally, so we don't need a lock here
+        return registry.get_model(
+            engine=self.engine,
+            table_name=self.table_name,
+            module_label=apps_label,
+            prefer_stable_names=True,
+        )
+
+    @staticmethod
+    def _normalize_class_name(table_name: str) -> str:
+        return "".join(word.capitalize() for word in table_name.split("_"))
+
+    @staticmethod
+    def _normalize_column_name(column_name: str) -> str:
+        # Useful helper if you ever need to generate column aliases
+        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_flux/df_helper/backends/utils.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from datetime import date
+from typing import Any, Optional
+
+import dask.dataframe as dd
+
+
+from sibi_flux.dask_cluster.core import safe_compute, safe_persist
+
+
+def is_dask_df(x: Any) -> bool:
+    return isinstance(x, dd.DataFrame)
+
+
+def maybe_persist(df: Any, persist: bool):
+    return safe_persist(df) if persist and is_dask_df(df) else df
+
+
+def maybe_compute(df: Any, as_pandas: bool):
+    return safe_compute(df) if as_pandas and is_dask_df(df) else df
+
+
+def parse_iso_date(value: Optional[str], field_name: str) -> Optional[date]:
+    if value is None:
+        return None
+    if not isinstance(value, str):
+        raise TypeError(f"{field_name} must be a string in YYYY-MM-DD format.")
+    try:
+        return date.fromisoformat(value)
+    except ValueError as e:
+        raise ValueError(f"{field_name} must be in YYYY-MM-DD format.") from e

sibi_flux/df_helper/core/__init__.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from ._defaults import sqlalchemy_field_conversion_map_dask, normalize_sqlalchemy_type
+from ._filter_handler import FilterHandler
+from ._params_config import ParamsConfig, DataFrameParams
+from ._query_config import QueryConfig
+
+__all__ = [
+    "ParamsConfig",
+    "QueryConfig",
+    "sqlalchemy_field_conversion_map_dask",
+    "normalize_sqlalchemy_type",
+    "FilterHandler",
+    "DataFrameParams",
+]

sibi_flux/df_helper/core/_defaults.py

@@ -0,0 +1,104 @@
+#  Copyright (c) 2023. ISTMO Center S.A.  All Rights Reserved
+#
+import json
+from typing import Dict, Callable
+
+import pandas as pd
+from sqlalchemy import (
+    String,
+    Text,
+    Integer,
+    BigInteger,
+    SmallInteger,
+    Float,
+    Boolean,
+    DateTime,
+    Date,
+    Time,
+    JSON,
+    Numeric,
+    UUID,
+)
+from sqlalchemy.dialects.mysql import TINYINT, MEDIUMTEXT
+
+# This is the defaults configuration file for the df_helper module.
+
+# conversion_map is a dictionary that maps the field types to their corresponding data type conversion functions.
+# Each entry in the dictionary is a pair of a field type (as a string) and a callable function that performs the
+# conversion. This mapping is used to convert the values in a pandas DataFrame to the appropriate data types based on
+# the db field type.
+
+sqlalchemy_field_conversion_map_dask: Dict[str, Callable] = {
+    String.__name__: lambda x: x.astype(str).fillna(""),
+    Text.__name__: lambda x: x.fillna("").astype(str),
+    Integer.__name__: lambda x: x.fillna(0).astype(int),
+    BigInteger.__name__: lambda x: pd.to_numeric(x, errors="coerce"),
+    SmallInteger.__name__: lambda x: pd.to_numeric(x, errors="coerce"),
+    Float.__name__: lambda x: pd.to_numeric(x, errors="coerce"),
+    Numeric.__name__: lambda x: pd.to_numeric(x, errors="coerce"),
+    Boolean.__name__: lambda x: x.astype(bool),
+    DateTime.__name__: lambda x: pd.to_datetime(x, errors="coerce"),
+    Date.__name__: lambda x: pd.to_datetime(x, errors="coerce").map_partitions(
+        lambda x: x.dt.date, meta=("date", "object")
+    ),
+    Time.__name__: lambda x: pd.to_datetime(x, errors="coerce").map_partitions(
+        lambda x: x.dt.time, meta=("time", "object")
+    ),
+    JSON.__name__: lambda x: x.map_partitions(
+        lambda s: s.apply(json.loads), meta=("json", "object")
+    ),
+    UUID.__name__: lambda x: x.astype(str),
+}
+
+
+# Conversion map with normalized SQLAlchemy field types
+# sqlalchemy_field_conversion_map_dask: Dict[str, callable] = {
+#     "String": lambda x: x.map_partitions(lambda s: s.astype(str), meta=("string", "string")),
+#     "Text": lambda x: x.map_partitions(lambda s: s.astype(str), meta=("text", "string")),
+#     "Integer": lambda x: pd.to_numeric(x, errors="coerce"),
+#     "SmallInteger": lambda x: pd.to_numeric(x, errors="coerce"),
+#     "BigInteger": lambda x: pd.to_numeric(x, errors="coerce"),
+#     "Float": lambda x: pd.to_numeric(x, errors="coerce"),
+#     "Numeric": lambda x: pd.to_numeric(x, errors="coerce"),
+#     "Boolean": lambda x: x.map_partitions(lambda s: s.fillna(False).astype(bool), meta=("boolean", "bool")),
+#     "DateTime": lambda x: pd.to_datetime(x, errors="coerce"),
+#     "Date": lambda x: pd.to_datetime(x, errors="coerce").map_partitions(lambda s: s.dt.date, meta=("date", "object")),
+#     "Time": lambda x: pd.to_datetime(x, errors="coerce").map_partitions(lambda s: s.dt.time, meta=("time", "object")),
+#     "JSON": lambda x: x.map_partitions(lambda s: s.apply(json.loads), meta=("json", "object")),
+# }
+
+
+def normalize_sqlalchemy_type(field_type):
+    """
+    Normalize SQLAlchemy field types to generic type names.
+    Handles dialect-specific types (e.g., MySQL).
+    """
+    # Map of generic SQLAlchemy types
+    type_mapping = {
+        String: "String",
+        Text: "Text",
+        Integer: "Integer",
+        SmallInteger: "SmallInteger",
+        BigInteger: "BigInteger",
+        Float: "Float",
+        Numeric: "Numeric",
+        Boolean: "Boolean",
+        DateTime: "DateTime",
+        Date: "Date",
+        Time: "Time",
+        JSON: "JSON",
+    }
+
+    # Dialect-specific types
+    dialect_mapping = {
+        TINYINT: "SmallInteger",
+        MEDIUMTEXT: "Text",
+    }
+
+    # Check if the field matches a generic or dialect-specific type
+    for sql_type, name in {**type_mapping, **dialect_mapping}.items():
+        if isinstance(field_type, sql_type):
+            return name
+
+    # Fallback to raw class name
+    return field_type.__class__.__name__