sibi-dst 2025.9.9__py3-none-any.whl → 2025.9.11__py3-none-any.whl

sibi_dst/utils/filepath_generator.py

@@ -160,160 +160,7 @@ class FilePathGenerator:
         # For local file, return absolute-like path without scheme or keep 'file://'? Keep scheme for consistency.
         return f"{self._protocol}://{path}"
 
-# import datetime
-# import re
-#
-# import fsspec
-#
-# from .log_utils import Logger
-#
-#
-# class FilePathGenerator:
-#     """
-#     Dynamically generates file paths by scanning directories starting from the base path
-#     and determining the innermost directory structure.
-#
-#     Now supports generating appropriate paths for both pandas and Dask.
-#     """
-#
-#     def __init__(self, base_path='', fs=None, logger=None, **kwargs):
-#         """
-#         Initialize the FilePathGenerator.
-#
-#         Parameters:
-#             base_path (str): Base directory path where data files are stored.
-#             fs (fsspec.AbstractFileSystem, optional): Filesystem object to use for file operations.
-#             logger (Logger, optional): Logger instance for logging information.
-#             **kwargs: Additional keyword arguments.
-#                 - debug (bool): If True, enables debug logging.
-#                 - storage_options (dict): Options for the filesystem (e.g., credentials, tokens).
-#                 - exclude_patterns (list): List of regex patterns to exclude from file paths.
-#                 - file_extension (str): File extension to look for (default: 'parquet').
-#         """
-#         self.base_path = base_path.rstrip('/')
-#         self.fs = fs  # Filesystem object
-#         self.logger = logger or Logger.default_logger(logger_name=self.__class__.__name__)
-#         self.debug = kwargs.get('debug', False)
-#         self.storage_options = kwargs.get('storage_options', {})
-#         self.exclude_patterns = kwargs.get('exclude_patterns', [])
-#         self.file_extension = kwargs.get('file_extension', 'parquet').lstrip('.')
-#
-#         # If fs is not provided, initialize it based on base_path and storage_options
-#         if self.fs is None:
-#             self.fs, _ = fsspec.core.url_to_fs(self.base_path, **self.storage_options)
-#
-#     def generate_file_paths(self, start_date, end_date, engine='dask'):
-#         """
-#         Generate paths dynamically for files within the date range by scanning directories.
-#         Returns a list of file paths compatible with the specified engine.
-#
-#         Parameters:
-#             start_date (str or datetime): Start date in 'YYYY-MM-DD' format or datetime object.
-#             end_date (str or datetime): End date in 'YYYY-MM-DD' format or datetime object.
-#             engine (str): 'pandas' or 'dask' to specify which library the paths are intended for.
-#
-#         Returns:
-#             list: List of file paths.
-#         """
-#         start_date = self._convert_to_datetime(start_date)
-#         end_date = self._convert_to_datetime(end_date)
-#
-#         paths = []
-#         curr_date = start_date
-#
-#         while curr_date <= end_date:
-#             year, month, day = curr_date.year, curr_date.month, curr_date.day
-#             day_paths = self._collect_paths(year, month, day, engine)
-#             if day_paths:
-#                 paths.extend(day_paths)
-#             curr_date += datetime.timedelta(days=1)
-#
-#         return paths
-#
-#     def _collect_paths(self, year, month, day, engine):
-#         """
-#         Collect appropriate paths for a given date, depending on the engine.
-#
-#         Parameters:
-#             year (int): Year component of the date.
-#             month (int): Month component of the date.
-#             day (int): Day component of the date.
-#             engine (str): 'pandas' or 'dask'.
-#
-#         Returns:
-#             list: List of file or directory paths.
-#         """
-#         base_dir = f"{self.base_path}/{year}/{str(month).zfill(2)}/{str(day).zfill(2)}"
-#
-#         if not self.fs.exists(base_dir):
-#             if self.debug:
-#                 self.logger.debug(f"Directory does not exist: {base_dir}")
-#             return []
-#
-#         if engine == 'dask':
-#             # Collect individual file paths
-#             file_pattern = f"{base_dir}/**/*.{self.file_extension}"
-#             all_paths = self.fs.glob(file_pattern)
-#
-#             if not all_paths and self.debug:
-#                 self.logger.debug(f"No files found with pattern: {file_pattern}")
-#
-#             # Exclude unwanted files and directories
-#             filtered_paths = self._exclude_unwanted_paths(all_paths)
-#
-#             # Filter out directories
-#             file_paths = [path for path in filtered_paths if not self.fs.isdir(path)]
-#
-#         elif engine == 'pandas':
-#             # Collect dataset directories
-#             # Assume that the base_dir is a Parquet dataset
-#             if self.fs.isdir(base_dir):
-#                 file_paths = [base_dir]
-#             else:
-#                 file_paths = []
-#
-#         else:
-#             raise ValueError("Engine must be 'pandas' or 'dask'.")
-#
-#         protocol = self.fs.protocol if isinstance(self.fs.protocol, str) else self.fs.protocol[0]
-#
-#         # Ensure the protocol is included in the paths
-#         file_paths = [
-#             f"{protocol}://{path}" if not path.startswith(f"{protocol}://") else path
-#             for path in file_paths
-#         ]
-#
-#         if self.debug:
-#             self.logger.debug(f"Collected {len(file_paths)} paths from {base_dir} for engine '{engine}'")
-#
-#         return file_paths
-#
-#     def _exclude_unwanted_paths(self, paths):
-#         """
-#         Exclude paths that match any of the exclusion patterns.
-#         """
-#         # Combine default patterns with user-provided patterns
-#         exclude_patterns = self.exclude_patterns
-#
-#         # Compile regex patterns for efficiency
-#         compiled_patterns = [re.compile(pattern) for pattern in exclude_patterns]
-#
-#         # Filter out paths matching any of the exclude patterns
-#         filtered_paths = [
-#             path for path in paths
-#             if not any(pattern.match(path) for pattern in compiled_patterns)
-#         ]
-#
-#         return filtered_paths
-#
-#     @staticmethod
-#     def _convert_to_datetime(date):
-#         """Convert a date string or datetime object into a datetime object."""
-#         if isinstance(date, str):
-#             return datetime.datetime.strptime(date, '%Y-%m-%d')
-#         return date
-#
-#
+
 # """
 # Usage:
 # # Initialize the generator

sibi_dst/utils/parquet_saver.py

@@ -71,6 +71,7 @@ class ParquetSaver(ManagedResource):
             max_delete_workers: int = 8,
             write_gate_max: int = 2,
             write_gate_key: Optional[str] = None,
+            partition_on: Optional[list[str]] = None,
             **kwargs: Any,
     ):
         super().__init__(**kwargs)
@@ -93,6 +94,7 @@ class ParquetSaver(ManagedResource):
         self.max_delete_workers = max(1, int(max_delete_workers))
         self.write_gate_max = max(1, int(write_gate_max))
         self.write_gate_key = (write_gate_key or self.parquet_storage_path).rstrip("/")
+        self.partition_on = partition_on
 
         # Fix: Remove deprecated coerce_timestamps parameter
         self.pyarrow_args.setdefault("compression", "zstd")
@@ -103,7 +105,18 @@ class ParquetSaver(ManagedResource):
 
     # ---------- public API ----------
     def save_to_parquet(self, output_directory_name: str = "default_output", overwrite: bool = True) -> str:
-        target_path = f"{self.parquet_storage_path}/{output_directory_name}".rstrip("/")
+        """
+        Save the Dask DataFrame to Parquet. If partition_on is provided, write as a
+        partitioned dataset without overwriting earlier partitions.
+        """
+        # Always treat as a directory target
+        if self.partition_on:
+            overwrite = False
+            # we override the output_directory_name and overwrite setting to avoid confusion since dask will (re) create subdirs
+            # Partitioned dataset → write directly to root directory
+            target_path = self.parquet_storage_path.rstrip("/")
+        else:
+            target_path = f"{self.parquet_storage_path}/{output_directory_name}".rstrip("/")
 
         sem = get_write_sem(self.write_gate_key, self.write_gate_max)
         with sem:
@@ -111,7 +124,7 @@ class ParquetSaver(ManagedResource):
                 self._clear_directory_safely(target_path)
             self.fs.mkdirs(target_path, exist_ok=True)
 
-            # Define a pyarrow schema and coerce the Dask frame to match it.
+            # Enforce schema before write
             schema = self._define_schema()
             ddf = self._coerce_ddf_to_schema(self.df_result, schema)
 
@@ -128,25 +141,25 @@ class ParquetSaver(ManagedResource):
                 pa.set_cpu_count(self.arrow_cpu)
 
             try:
+                params = {
+                    "path": target_path,
+                    "engine": "pyarrow",
+                    "filesystem": self.fs,
+                    "write_index": self.write_index,
+                    "write_metadata_file": self.write_metadata_file,
+                    **self.pyarrow_args,
+                }
+                self.partition_on = self.partition_on if isinstance(self.partition_on, list) else None
+                if self.partition_on:
+                    params["partition_on"] = self.partition_on
+
                 with self._local_dask_pool():
-                    ddf.to_parquet(
-                        path=target_path,
-                        engine="pyarrow",
-                        schema=schema,
-                        overwrite=False,
-                        filesystem=self.fs,
-                        write_index=self.write_index,
-                        write_metadata_file=self.write_metadata_file,
-                        **self.pyarrow_args,
-                    )
+                    ddf.to_parquet(**params)
             finally:
                 if old_arrow_cpu is not None:
                     pa.set_cpu_count(old_arrow_cpu)
 
-        self.logger.info(
-            f"Parquet dataset written: {target_path}",
-            extra=self.logger_extra,
-        )
+        self.logger.info(f"Parquet dataset written: {target_path}", extra=self.logger_extra)
         return target_path
 
     @contextmanager

sibi_dst/utils/progress/sse_runner.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import asyncio, contextlib, inspect, json
-from typing import Any, Awaitable, Callable, Dict, Optional, Union
+from typing import Any, Awaitable, Callable, Dict, Optional, Union, Mapping, MutableMapping
 from fastapi import Request
 from sse_starlette.sse import EventSourceResponse
 from sibi_dst.utils import Logger
@@ -13,6 +13,16 @@ TaskFn = Union[Task2, Task3]
 def _as_sse_msg(event: str, data: Any) -> dict:
     return {"event": event, "data": json.dumps(data) if not isinstance(data, (str, bytes)) else data}
 
+def _merge_ctx(*parts: Optional[Mapping[str, Any]]) -> Dict[str, Any]:
+    """Right-most precedence; shallow merge is sufficient for our keys."""
+    out: Dict[str, Any] = {}
+    for p in parts:
+        if not p:
+            continue
+        for k, v in p.items():
+            out[k] = v
+    return out
+
 class SSERunner:
     def __init__(self, *, task: TaskFn, logger: Logger, ping: int = 15,
                  headers: Optional[dict] = None, auto_complete: bool = True) -> None:
@@ -31,39 +41,57 @@ class SSERunner:
     async def _worker(self, queue: asyncio.Queue, task_id: str, ctx: Dict[str, Any]) -> None:
         self.logger.info(f"SSE {task_id}: start")
         try:
-            await queue.put(_as_sse_msg("progress", {"message": "Task started"}))
+            await queue.put(_as_sse_msg("progress", {"message": "Task started", "task_id": task_id}))
             payload = await self._call_task(queue, task_id, ctx)
             if self.auto_complete:
                 final = payload if payload is not None else {"status": "complete"}
+                if isinstance(final, dict) and "task_id" not in final:
+                    final["task_id"] = task_id
                 await queue.put(_as_sse_msg("complete", final))
             self.logger.info(f"SSE {task_id}: complete")
         except asyncio.CancelledError:
             raise
         except Exception as e:
             self.logger.error(f"SSE {task_id} failed: {e}", exc_info=True)
-            await queue.put(_as_sse_msg("error", {"detail": str(e)}))
+            await queue.put(_as_sse_msg("error", {"detail": str(e), "task_id": task_id}))
         finally:
             await queue.put(None)
 
-    def endpoint(self):
-        async def handler(request: Request):  # <-- only Request
+    def endpoint(self, *, ctx: Optional[Dict[str, Any]] = None):
+        """
+        Create an SSE endpoint.
+        - ctx: optional explicit context dict provided by the caller.
+               This overrides request-derived context and request.state.ctx.
+        Precedence when merging: request-derived < request.state.ctx < ctx (explicit).
+        """
+        async def handler(request: Request):
             queue: asyncio.Queue = asyncio.Queue()
             task_id = str(asyncio.get_running_loop().time()).replace(".", "")
             self.logger.debug(
-                f"SSE {task_id}: new request client={request.client} path={request.url.path} q={dict(request.query_params)}")
+                f"SSE {task_id}: new request client={request.client} path={request.url.path} q={dict(request.query_params)}"
+            )
 
-            ctx: Dict[str, Any] = {
-                "path": dict(request.path_params),          # <-- pull path params here
+            # Base ctx from the HTTP request
+            base_ctx: Dict[str, Any] = {
+                "path": dict(request.path_params),
                 "query": dict(request.query_params),
                 "method": request.method,
+                "headers": dict(request.headers) if hasattr(request, "headers") else None,
             }
             if request.headers.get("content-type", "").startswith("application/json"):
                 try:
-                    ctx["body"] = await request.json()
+                    base_ctx["body"] = await request.json()
                 except Exception:
-                    ctx["body"] = None
+                    base_ctx["body"] = None
+
+            # Pull any pre-attached ctx from request.state
+            state_ctx: Optional[Dict[str, Any]] = getattr(request.state, "ctx", None)
+
+            # Merge with precedence: base_ctx < state_ctx < explicit ctx
+            merged_ctx = _merge_ctx(base_ctx, state_ctx, ctx)
 
-            worker = asyncio.create_task(self._worker(queue, task_id, ctx))
+            # Run worker
+            worker = asyncio.create_task(self._worker(queue, task_id, merged_ctx))
 
             async def gen():
                 try: