sibi-flux 2025.12.0__py3-none-any.whl

sibi_flux/utils/boilerplate/base_cube_router.py

@@ -0,0 +1,283 @@
+import asyncio
+import json
+from contextlib import suppress
+from collections.abc import Mapping
+from typing import Any, Optional, Type, Sequence, List
+
+import pandas as pd
+from fastapi import APIRouter, Body, HTTPException, Query, Request
+from fastapi.encoders import jsonable_encoder
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel, ConfigDict, Field, create_model
+
+from sibi_flux.dask_cluster import safe_compute
+from sibi_flux.dask_cluster.client_manager import get_persistent_client
+from sibi_flux.utils.boilerplate.base_data_cube import BaseDatacube
+from sibi_flux.datacube._data_cube import Datacube
+
+
+class BaseCubeRouter:
+    """
+    A reusable router class for Datacubes that provides:
+    - POST /: Data retrieval with dynamic filters (optional pagination).
+    - POST /stream: SSE streaming with dynamic filters.
+
+    Attributes:
+        cube_cls: The Datacube class to wrap.
+        router: The APIRouter instance containing the endpoints.
+        FiltersModel: Dynamic Pydantic model for validation.
+        pagination: Whether to enable pagination logic on the main endpoint.
+    """
+
+    def __init__(
+        self,
+        cube_cls: Type[BaseDatacube] | Type[Datacube],
+        prefix: str = "",
+        tags: Sequence[str] | None = None,
+        pagination: bool = True,
+    ):
+        self.cube_cls = cube_cls
+        self.pagination = pagination
+        self.router = APIRouter(prefix=prefix, tags=tags or [])
+        self.FiltersModel = self._create_filters_model()
+        self._setup_routes()
+
+    def _generate_docstring_and_example(self, field_defs: Mapping[str, Any]) -> str:
+        """
+        Generates a JSON-like representation of the model for the docstring.
+        """
+        example_dict = {}
+        for name, (type_, field) in field_defs.items():
+            # Simplistic type mapping for display
+            type_name = "string"
+            if type_ is Optional[bool]:
+                type_name = True
+            elif type_ is Optional[int]:
+                type_name = 0
+            elif type_ is Optional[float]:
+                type_name = 0.0
+
+            example_dict[name] = type_name
+
+        pretty_json = json.dumps(example_dict, indent=2)
+        return f"Filter options:\n\n```json\n{pretty_json}\n```"
+
+    def _create_filters_model(self) -> Type[BaseModel]:
+        # Default behavior: if exclude_columns is empty/missing, exclude 'id'.
+        _exclude_columns = self.cube_cls.config.get("exclude_columns", [])
+        if not _exclude_columns:
+            _exclude_columns = ["id"]
+
+        try:
+            # Introspect dtypes by loading an empty slice/schema
+            # We use n_records=1 at init to hint the backend to limit query size
+            # Note: This requires the DB/Backend to be accessible at startup
+            cube = self.cube_cls(n_records=1)
+
+            # Try to load just metadata/empty df
+            # If backend doesn't support limit, this might raise Error or load full data.
+            # We catch exceptions to fallback.
+            df = cube.load()
+
+            if df is None:
+                return self._fallback_create_model(_exclude_columns)
+
+            _field_definitions = {}
+
+            # Identify boolean fields from config as override/hint
+            _boolean_fields = getattr(self.cube_cls, "boolean_fields", [])
+            _field_map = self.cube_cls.config.get("field_map", {})
+            _target_columns = set(_field_map.values()) if _field_map else None
+
+            for col in df.columns:
+                if col in _exclude_columns:
+                    continue
+
+                # If field_map is provided, use only targeted columns
+                if _target_columns is not None and col not in _target_columns:
+                    continue
+
+                if col not in df.columns:
+                    continue
+
+                dtype = df.dtypes[col]
+                desc = f"Filter by {col}"
+                field_type: Any = Optional[str]  # Default
+
+                if col in _boolean_fields:
+                    field_type = Optional[bool]
+                    desc += " (bool)"
+                elif pd.api.types.is_bool_dtype(dtype):
+                    field_type = Optional[bool]
+                    desc += " (bool)"
+                elif pd.api.types.is_integer_dtype(dtype):
+                    field_type = Optional[int]
+                    desc += " (int)"
+                elif pd.api.types.is_float_dtype(dtype):
+                    field_type = Optional[float]
+                    desc += " (float)"
+
+                _field_definitions[col] = (
+                    field_type,
+                    Field(default=None, description=desc),
+                )
+
+            model = create_model(
+                f"{self.cube_cls.__name__}Filters",
+                **_field_definitions,
+                __config__=ConfigDict(extra="ignore"),
+            )
+            model.__doc__ = self._generate_docstring_and_example(_field_definitions)
+            return model
+
+        except Exception:
+            # Fallback to config-based generation
+            # print(f"DTO INTROSPECTION ERROR for {self.cube_cls.__name__}: {e}")
+            return self._fallback_create_model(_exclude_columns)
+
+    def _fallback_create_model(self, exclude_columns: Sequence[str]) -> Type[BaseModel]:
+        # Retrieve fields from the Cube configuration
+        _field_map = self.cube_cls.config.get("field_map", {})
+        _output_columns = list(_field_map.values()) if _field_map else []
+        _boolean_fields = getattr(self.cube_cls, "boolean_fields", [])
+
+        _field_definitions = {}
+        for col in _output_columns:
+            if col in exclude_columns:
+                continue
+
+            if col in _boolean_fields:
+                _field_definitions[col] = (
+                    Optional[bool],
+                    Field(default=None, description=f"Filter by {col} (bool)"),
+                )
+            else:
+                _field_definitions[col] = (
+                    Optional[str | int],
+                    Field(default=None, description=f"Filter by {col}"),
+                )
+
+        FiltersModel: Any = create_model(
+            f"{self.cube_cls.__name__}Filters",
+            **_field_definitions,
+            __config__=ConfigDict(extra="ignore"),
+        )
+        FiltersModel.__doc__ = self._generate_docstring_and_example(_field_definitions)
+        return FiltersModel
+
+    async def _materialize_df(self, df: Any) -> pd.DataFrame:
+        if hasattr(df, "compute"):
+            # Use the shared persistent client if available
+            client = get_persistent_client()
+            return await asyncio.to_thread(safe_compute, df, dask_client=client)
+        return df
+
+    async def _stream_data_task(
+        self, cube: BaseDatacube, filters: Mapping[str, Any]
+    ) -> None:
+        async with cube:
+            await cube.emit("status", msg="loading")
+            try:
+                df = await cube.aload(**filters)
+                df = await self._materialize_df(df)
+
+                if df is not None and not df.empty:
+                    records = df.to_dict(orient="records")
+                    await cube.emit("data", data=jsonable_encoder(records))
+
+                await cube.emit("status", msg="complete")
+
+            except Exception as e:
+                await cube.emit("error", detail=str(e))
+
+    def _setup_routes(self):
+        # We define the functions inside closure to capture 'self' reliably
+        # and use the dynamic FiltersModel for type annotation and FastAPI dependency.
+
+        FiltersModel = self.FiltersModel
+
+        async def _core_logic(cube, filters):
+            if isinstance(filters, dict):
+                filters = FiltersModel(**filters)
+            filter_dict = filters.model_dump(exclude_unset=True)
+            try:
+                df = await cube.aload(**filter_dict)
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=str(e))
+            df = await self._materialize_df(df)
+            return df
+
+        if self.pagination:
+
+            @self.router.post(
+                "",
+                summary=f"Get Paged {self.cube_cls.__name__} Data",
+                description="Retrieve data with filtering and pagination.",
+            )
+            async def get_cube_data(
+                filters: Any = Body(
+                    default_factory=lambda: FiltersModel(),
+                    description="Filters to apply.",
+                    openapi_examples={},
+                ),
+                page: int = Query(1, ge=1, description="Page number."),
+                page_size: int = Query(
+                    50, ge=1, le=1000, description="Records per page."
+                ),
+            ) -> List[dict[str, Any]]:
+                cube = self.cube_cls()
+                df = await _core_logic(cube, filters)
+
+                if df is None or df.empty:
+                    return []
+
+                start = (page - 1) * page_size
+                end = start + page_size
+                paged_df = df.iloc[start:end]
+                return paged_df.to_dict(orient="records")
+
+        else:
+
+            @self.router.post(
+                "",
+                summary=f"Get {self.cube_cls.__name__} Data",
+                description="Retrieve all data with filtering (pagination disabled).",
+            )
+            async def get_cube_data_all(
+                filters: Any = Body(
+                    default_factory=lambda: FiltersModel(),
+                    description="Filters to apply.",
+                    openapi_examples={},
+                ),
+            ) -> List[dict[str, Any]]:
+                cube = self.cube_cls()
+                df = await _core_logic(cube, filters)
+
+                if df is None or df.empty:
+                    return []
+
+                return df.to_dict(orient="records")
+
+        @self.router.post(
+            "/stream",
+            summary=f"Stream {self.cube_cls.__name__} Data",
+            description="Stream data updates via SSE.",
+        )
+        async def stream_cube_data(
+            request: Request,
+            filters: Any = Body(
+                default_factory=lambda: FiltersModel(),
+                description="Filters to apply.",
+                openapi_examples={},
+            ),
+        ) -> StreamingResponse:
+            cube = self.cube_cls(auto_sse=True)
+            if isinstance(filters, dict):
+                filters = FiltersModel(**filters)
+            filter_dict = filters.model_dump(exclude_unset=True)
+
+            asyncio.create_task(self._stream_data_task(cube, filter_dict))
+
+            return StreamingResponse(
+                cube.get_sse().aiter_sse(), media_type="text/event-stream"
+            )

sibi_flux/utils/boilerplate/base_data_cube.py

@@ -0,0 +1,132 @@
+"""
+Base DataCube module.
+
+Provides the foundational classes for DataCube definitions and loading logic.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Optional, Any, Dict
+
+import dask.dataframe as dd
+import pandas as pd
+
+from sibi_flux.df_helper import DfHelper
+from sibi_flux.dask_cluster import dask_is_empty
+
+
+DataFrameType = dd.DataFrame | pd.DataFrame
+
+
+class BaseDatacube(DfHelper):
+    """
+    Base cube with sync/async load hooks and resilient Dask handling.
+
+    Lifecycle:
+      1. load() / aload() calls parent DfHelper to fetch raw data.
+      2. If data exists, calls fix_data() (sync) or afix_data() (async).
+      3. Result is stored in self.df and returned.
+
+    Subclasses should override:
+      - fix_data(self, df): for CPU-bound transformations (Pandas/Dask map_partitions)
+      - afix_data(self, df): for I/O-bound transformations (DB lookups, etc.)
+    """
+
+    # Class-level config overrides (e.g. backend='parquet')
+    config: Dict[str, Any] = {}
+
+    def __init__(self, **kwargs):
+
+        combined_config = {**self.config, **kwargs}
+        super().__init__(**combined_config)
+
+        # State container
+        self.df: Optional[DataFrameType] = None
+
+    # -----------------------------------------------------------------------
+    # Hooks (Override these)
+    # -----------------------------------------------------------------------
+    def fix_data(self, df: DataFrameType, **kwargs) -> DataFrameType:
+        """
+        Synchronous transformation hook.
+
+        Args:
+            df: The loaded dataframe (Pandas or Dask)
+            **kwargs: The options passed to load()
+
+        Returns:
+            The transformed dataframe.
+        """
+        return df
+
+    async def afix_data(self, df: DataFrameType, **kwargs) -> DataFrameType:
+        """
+        Asynchronous transformation hook.
+
+        Defaults to calling fix_data(). Override this ONLY if you need
+        to perform async operations (like awaiting other DB calls).
+        """
+        return self.fix_data(df, **kwargs)
+
+    # -----------------------------------------------------------------------
+    # Public API
+    # -----------------------------------------------------------------------
+    def load(self, **kwargs) -> DataFrameType:
+        """
+        Synchronous load pipeline.
+        """
+        # 1. Load Raw Data (Delegate to DfHelper)
+        df = super().load(**kwargs)
+
+        # 2. Check Emptiness (Resilient)
+        if not self._is_empty(df):
+            # 3. Apply Transform Hook
+            df = self.fix_data(df, **kwargs)
+        else:
+            self.logger.debug(f"No data loaded by {self.__class__.__name__}")
+
+        # 4. Update State & Return
+        self.df = df
+        return df
+
+    async def aload(self, **kwargs) -> DataFrameType:
+        """
+        Asynchronous load pipeline.
+        """
+        # 1. Load Raw Data
+        df = await super().aload(**kwargs)
+
+        # 2. Check Emptiness (Non-blocking)
+        # _is_empty triggers dask.compute(), so we must offload to thread
+        is_empty = await asyncio.to_thread(self._is_empty, df)
+
+        if not is_empty:
+            # 3. Apply Async Transform Hook
+            # (Note: afix_data calls fix_data by default, so this covers both cases)
+            df = await self.afix_data(df, **kwargs)
+        else:
+            self.logger.debug(f"No data loaded by {self.__class__.__name__}")
+
+        # 4. Update State & Return
+        self.df = df
+        return df
+
+    # -----------------------------------------------------------------------
+    # Internals
+    # -----------------------------------------------------------------------
+    def _is_empty(self, df: Optional[DataFrameType]) -> bool:
+        """
+        Robust emptiness check using dask_cluster.
+        """
+        if df is None:
+            return True
+
+        if isinstance(df, pd.DataFrame):
+            return df.empty
+
+        # Use our new resilient checker
+        # This prevents the "compute whole graph to check length" disaster
+        # Pass the client if we have one (from DfHelper/DaskClientMixin)
+        client = getattr(self, "dask_client", None)
+        return dask_is_empty(df, dask_client=client, logger=self.logger)

sibi_flux/utils/boilerplate/base_pipeline_template.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import pandas as pd
+
+from sibi_flux.pipelines import BasePipeline
+
+
+class PipelineTemplate:
+    """
+    A reusable base class for executing product-related pipelines end-to-end.
+    """
+
+    def __init__(
+        self,
+        start_date: str,
+        end_date: str,
+        fs_instance,
+        storage_path: str,
+        dataset_cls,
+        filename: str,
+        date_field: str = "last_activity_dt",
+        **kwargs,
+    ):
+        self.start_date = start_date
+        self.end_date = end_date
+        self.max_workers = kwargs.pop("max_workers", 4)
+        self.fs = fs_instance
+        self.storage_path = storage_path
+
+        self.pipeline = BasePipeline(
+            start_date=self.start_date,
+            end_date=self.end_date,
+            dataset_cls=dataset_cls,
+            parquet_storage_path=self.storage_path,
+            fs=self.fs,
+            filename=filename,
+            date_field=date_field,
+            max_workers=self.max_workers,
+        )
+
+    async def to_parquet(self, **kwargs) -> pd.DataFrame:
+        await self.pipeline.to_parquet(**kwargs)
+        df = await self.pipeline.from_parquet(**kwargs)
+        return df
+
+    async def from_parquet(self, **kwargs) -> pd.DataFrame:
+        df = await self.pipeline.from_parquet(**kwargs)
+        return df
+
+    async def to_clickhouse(self, clickhouse_conf, **kwargs) -> None:
+        cnf = clickhouse_conf.copy()
+        cnf["table"] = self.pipeline.filename
+        cnf["overwrite"] = True
+        await self.pipeline.to_clickhouse(cnf, **kwargs)

sibi_flux/utils/boilerplate/hybrid_data_loader.py

@@ -0,0 +1,193 @@
+import datetime
+from typing import Optional
+import pandas as pd
+import dask.dataframe as dd
+from sibi_flux.logger import Logger
+
+TODAY = datetime.date.today()
+YESTERDAY = TODAY - datetime.timedelta(days=1)
+TODAY_STR = TODAY.strftime("%Y-%m-%d")
+YESTERDAY_STR = YESTERDAY.strftime("%Y-%m-%d")
+
+
+class HybridDataLoader:
+    """
+    Hybrid loader that merges historical (Parquet) and live (API/DB) data
+    in a consistent, schema-safe, timezone-normalized way.
+    """
+
+    def __init__(
+        self,
+        start_date: str,
+        end_date: str,
+        historical_reader,
+        live_reader,
+        date_field: str,
+        **kwargs,
+    ):
+        self.start_date = self._validate_date_format(start_date)
+        self.end_date = self._validate_date_format(end_date)
+        self.historical_reader = historical_reader
+        self.live_reader = live_reader
+        self.date_field = date_field
+
+        self.logger = kwargs.get("logger", Logger.default_logger(logger_name=__name__))
+        self.debug = kwargs.get("debug", False)
+        self.logger.set_level(Logger.DEBUG if self.debug else Logger.INFO)
+
+        self._validate_date_range()
+
+        self._read_live_flag = self.end_date == TODAY_STR
+        self._is_single_today = self.start_date == self.end_date == TODAY_STR
+        self._is_single_historical = self.start_date == self.end_date != TODAY_STR
+
+    # ------------------------------------------------------------------ #
+    # Validation
+    # ------------------------------------------------------------------ #
+    @staticmethod
+    def _validate_date_format(date_str: str) -> str:
+        try:
+            datetime.datetime.strptime(date_str, "%Y-%m-%d")
+            return date_str
+        except ValueError:
+            raise ValueError(f"Date '{date_str}' is not in valid YYYY-MM-DD format")
+
+    def _validate_date_range(self):
+        start = datetime.datetime.strptime(self.start_date, "%Y-%m-%d").date()
+        end = datetime.datetime.strptime(self.end_date, "%Y-%m-%d").date()
+        if end < start:
+            raise ValueError(
+                f"End date ({self.end_date}) cannot be before start date ({self.start_date})"
+            )
+
+    @staticmethod
+    def _normalize_datetimes(df: dd.DataFrame, cols: list[str]) -> dd.DataFrame:
+        """Normalize datetime columns to UTC safely."""
+        if not cols:
+            return df
+
+        def _to_ts(pdf: pd.DataFrame) -> pd.DataFrame:
+            for c in cols:
+                if c in pdf.columns:
+                    pdf[c] = pd.to_datetime(pdf[c], errors="coerce", utc=True)
+            return pdf
+
+        return df.map_partitions(_to_ts)
+
+    @staticmethod
+    def _create_empty_dataframe(meta: Optional[pd.DataFrame] = None) -> dd.DataFrame:
+        if meta is not None and not meta.empty:
+            return dd.from_pandas(meta.iloc[0:0], npartitions=1)
+        return dd.from_pandas(pd.DataFrame(), npartitions=1)
+
+    # ------------------------------------------------------------------ #
+    # Data loading methods
+    # ------------------------------------------------------------------ #
+    async def _load_today_data(self, **kwargs) -> Optional[dd.DataFrame]:
+        """Load today's live data."""
+        self.logger.debug("Loading today's live data...")
+        date_filter = {f"{self.date_field}__date": TODAY_STR}
+        filters = {**kwargs, **date_filter}
+
+        try:
+            reader_obj = self.live_reader(logger=self.logger, debug=self.debug)
+            if not hasattr(reader_obj, "aload"):
+                raise TypeError("live_reader must expose an async aload() method")
+            today_df = await reader_obj.aload(**filters)
+            if today_df is not None:
+                today_df = self._normalize_datetimes(today_df, [self.date_field])
+            return today_df
+        except Exception as e:
+            self.logger.error(f"Failed to load today's data: {e}")
+            return None if not self.debug else (_ for _ in ()).throw(e)
+
+    async def _load_historical_data(
+        self, start_date: str, end_date: str, **kwargs
+    ) -> dd.DataFrame:
+        """Load historical data."""
+        self.logger.debug(f"Loading historical data from {start_date} to {end_date}...")
+        try:
+            reader_obj = self.historical_reader(
+                parquet_start_date=start_date,
+                parquet_end_date=end_date,
+                logger=self.logger,
+                debug=self.debug,
+            )
+            if not hasattr(reader_obj, "aload"):
+                raise TypeError("historical_reader must expose an async aload() method")
+            df = await reader_obj.aload(**kwargs)
+            df = self._normalize_datetimes(df, [self.date_field])
+            return df
+        except Exception as e:
+            self.logger.error(f"Failed to load historical data: {e}")
+            if self.debug:
+                raise
+            return self._create_empty_dataframe()
+
+    # ------------------------------------------------------------------ #
+    # Orchestrator
+    # ------------------------------------------------------------------ #
+    # ------------------------------------------------------------------ #
+    # Orchestrator
+    # ------------------------------------------------------------------ #
+    async def aload(self, **kwargs) -> dd.DataFrame:
+        """Load and concatenate data from historical and live sources."""
+        self.logger.debug(
+            f"[HybridLoader] start={self.start_date}, end={self.end_date}, "
+            f"read_live={self._read_live_flag}"
+        )
+
+        # Case 1: only today
+        if self._is_single_today:
+            today_df = await self._load_today_data(**kwargs)
+            return today_df if today_df is not None else self._create_empty_dataframe()
+
+        # Case 2: purely historical
+        if not self._read_live_flag:
+            return await self._load_historical_data(
+                self.start_date, self.end_date, **kwargs
+            )
+
+        # Case 3: mixed historical + live
+        hist_df = await self._load_historical_data(
+            self.start_date, YESTERDAY_STR, **kwargs
+        )
+        live_df = await self._load_today_data(**kwargs)
+
+        if hist_df is None:
+            hist_df = self._create_empty_dataframe()
+        if live_df is None:
+            live_df = self._create_empty_dataframe()
+
+        # Standardize / Validate both before concat if schema is enforced
+        # (Assuming the caller might want to check schemas here)
+        # For now, we rely on the readers returning relatively clean data,
+        # but we can do a naive concat.
+
+        # Note: Previous implementation did heavy schema alignment.
+        # Dask concat handles most alignment if columns differ but dtypes are compatible.
+        # If columns are completely missing, Dask might warn or error.
+
+        try:
+            return dd.concat(
+                [hist_df, live_df],
+                ignore_unknown_divisions=True,
+                interleave_partitions=True,
+            )
+        except Exception as e:
+            self.logger.warning(
+                f"Simple concat failed: {e}. Attempting robust alignment..."
+            )
+            # Robust alignment fallback can be implemented here if strictly needed,
+            # possibly offloaded to a thread.
+            self.logger.error(
+                "Robust alignment not yet reimplemented with DfValidator."
+            )
+            raise e
+
+    # ------------------------------------------------------------------ #
+    def __repr__(self):
+        return (
+            f"HybridDataLoader(start='{self.start_date}', end='{self.end_date}', "
+            f"read_live={self._read_live_flag})"
+        )

sibi_flux/utils/clickhouse_writer/__init__.py

@@ -0,0 +1,6 @@
+from ._clickhouse_writer import ClickHouseWriter, _to_bool
+
+__all__ = [
+    "ClickHouseWriter",
+    "_to_bool",
+]