Bump version to 0.9.6.3; update type hints, improve MQTT client configuration, and enhance storage options handling

Author: legout

pyproject.toml

@@ -1,19 +1,11 @@
 [project]
 name = "FlowerPower"
 description = "A simple workflow framework. Hamilton + APScheduler = FlowerPower"
-authors = [{name = "Volker L.", email = "[email protected]"}]
+authors = [{ name = "Volker L.", email = "[email protected]" }]
 readme = "README.md"
 requires-python = ">= 3.11"
-version = "0.9.6.2"
-keywords = [
-  "hamilton",
-  "workflow",
-  "pipeline",
-  "scheduler",
-  "apscheduler",
-  "dask",
-  "ray",
-]
+version = "0.9.6.3"
+keywords = ["hamilton", "workflow", "pipeline", "scheduler", "apscheduler", "dask", "ray"]
 dependencies = [
   "sf-hamilton[visualization]>=1.69.0",
   "python-dotenv>=1.0.1",
@@ -32,21 +24,14 @@ dependencies = [
 
 
 [project.urls]
-"Homepage" = "https://github.com/legout/flowerpower"
+"Homepage"    = "https://github.com/legout/flowerpower"
 "Bug Tracker" = "https://github.com/legout/flowerpower/issues"
 
 [project.scripts]
 flowerpower = "flowerpower.cli:app"
 
 [project.optional-dependencies]
-io = [
-  "datafusion>=43.1.0",
-  "duckdb>=1.1.3",
-  "orjson>=3.10.12",
-  "polars>=1.15.0",
-  "pyarrow>=18.1.0",
-  "pydala2>=0.9.4.5",
-]
+io = ["datafusion>=43.1.0", "duckdb>=1.1.3", "orjson>=3.10.12", "polars>=1.15.0", "pyarrow>=18.1.0", "pydala2>=0.9.4.5"]
 mongodb = ["pymongo>=4.7.2"]
 mqtt = ["paho-mqtt>=2.1.0", "orjson>=3.10.11"]
 opentelemetry = [
@@ -65,11 +50,7 @@ scheduler = [
 ]
 tui = ["textual>=0.85.2"]
 ui = ["sf-hamilton-ui>=0.0.11"]
-webserver = [
-  "sanic>=24.6.0",
-  "sanic-ext>=23.12.0",
-  "orjson>=3.10.11",
-]
+webserver = ["sanic>=24.6.0", "sanic-ext>=23.12.0", "orjson>=3.10.11"]
 
 #[build-system]
 #build-backend = "hatchling.build"

src/flowerpower/io/base.py

@@ -1,4 +1,4 @@
-import os
+import posixpath
 from typing import Any, Generator
 
 import datafusion
@@ -14,11 +14,11 @@
 from ..utils.polars import pl
 from ..utils.sql import sql2polars_filter, sql2pyarrow_filter
 from ..utils.storage_options import (
-    # AwsStorageOptions,
-    # AzureStorageOptions,
-    # GcsStorageOptions,
-    # GitHubStorageOptions,
-    # GitLabStorageOptions,
+    AwsStorageOptions,
+    AzureStorageOptions,
+    GcsStorageOptions,
+    GitHubStorageOptions,
+    GitLabStorageOptions,
     StorageOptions,
 )
 import importlib
@@ -62,7 +62,16 @@ class BaseFileIO(BaseModel):
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
     path: str | list[str]
-    storage_options: StorageOptions | dict[str, Any] | None = None
+    storage_options: (
+        StorageOptions
+        | AwsStorageOptions
+        | AzureStorageOptions
+        | GcsStorageOptions
+        | GitLabStorageOptions
+        | GitHubStorageOptions
+        | dict[str, Any]
+        | None
+    ) = None
     fs: AbstractFileSystem | None = None
     format: str | None = None
 
@@ -74,6 +83,8 @@ def model_post_init(self, __context):
             self.storage_options = StorageOptions(
                 **self.storage_options
             ).storage_options
+        if isinstance(self.storage_options, StorageOptions):
+            self.storage_options = self.storage_options.storage_options
 
         if self.fs is None:
             self.fs = get_filesystem(
@@ -326,7 +337,6 @@ def iter_polars(
         batch_size: int = 1,
         **kwargs,
     ) -> Generator[pl.DataFrame | pl.LazyFrame, None, None]:
-
         if lazy:
             yield from self._iter_polars_lazyframe(batch_size=batch_size, **kwargs)
         yield from self._iter_polars_dataframe(batch_size=batch_size, **kwargs)
@@ -754,13 +764,13 @@ def write(
             | list[
                 pl.DataFrame | pl.LazyFrame | pa.Table | pd.DataFrame | dict[str, Any]
             ]
-        ) | None = None,
+        )
+        | None = None,
         basename: str | None = None,
         concat: bool = False,
         mode: str = "append",
         **kwargs,
     ):
-
         self.fs.write_files(
             data=data or self.data,
             basename=basename or self.basename,
@@ -819,7 +829,8 @@ def write(
                 | pd.DataFrame
                 | dict[str, Any]
             ]
-        ) | None = None,
+        )
+        | None = None,
         unique: bool | list[str] | str = False,
         delta_subset: str | None = None,
         alter_schema: bool = False,

src/flowerpower/io/loader/deltatable.py

@@ -11,19 +11,21 @@
 # from hamilton.function_modifiers import dataloader
 
 from ...utils.polars import pl
-from ..base import BaseFileIO
+from ..base import BaseDatasetLoader
 
 # from ..utils import get_dataframe_metadata, get_delta_metadata
 
 
-class DeltaTableLoader(BaseFileIO):
+class DeltaTableLoader(BaseDatasetLoader):
     """Delta table loader.
 
     This class is responsible for loading Delta tables into several dataframe formats,
     duckdb and datafusion.
 
     """
 
+    delta_table: DeltaTable | None = None
+
     def model_post_init(self, __context):
         super().model_post_init(__context)
         self.format = "delta"
@@ -39,63 +41,63 @@ def dt(self) -> DeltaTable:
 
     def to_pyarrow_dataset(self) -> pds.Dataset:
         """Converts the DeltaTable to a PyArrow Dataset."""
-        self.delta_table.to_pyarrow_dataset()
+        return self.delta_table.to_pyarrow_dataset()
 
     def to_pyrrow_table(self) -> pa.Table:
         """Converts the DeltaTable to a PyArrow Table."""
-        self.delta_table.to_pyarrow_table()
-
-    def to_polars(self, lazy: bool = True) -> pl.DataFrame | pl.LazyFrame:
-        """Converts the DeltaTable to a Polars DataFrame."""
-        if lazy:
-            return pl.scan_pyarrow_dataset(self.to_pyarrow_dataset())
-
-        return pl.from_arrow(self.to_pyarrow_table())
-
-    def to_pandas(self) -> pd.DataFrame:
-        """Converts the DeltaTable to a Pandas DataFrame."""
-        return self.to_pyrrow_table().to_pandas()
-
-    def to_duckdb_relation(
-        self, conn: duckdb.DuckDBPyConnection | None = None, lazy: bool = True
-    ) -> duckdb.DuckDBPyRelation:
-        """Converts the DeltaTable to a DuckDB relation."""
-        if conn is None:
-            conn = duckdb.connect()
-        if lazy:
-            return conn.from_arrow(self.to_pyarrow_dataset())
-        return conn.from_arrow(self.to_pyarrow_table())
-
-    def register_in_duckdb(
-        self,
-        conn: duckdb.DuckDBPyConnection | None = None,
-        name: str | None = None,
-        lazy: bool = True,
-    ) -> duckdb.DuckDBPyConnection:
-        """Registers the DeltaTable in a DuckDB connection."""
-        if name is None:
-            name = f"{self.format}:{self.path}"
-        if conn is None:
-            conn = duckdb.connect()
-        table = self.to_duckdb_relation(conn=conn, lazy=lazy)
-        conn.register(name, table)
-        return conn
-
-    def register_in_datafusion(
-        self,
-        ctx: dtf.SessionContext | None = None,
-        name: str | None = None,
-        lazy: bool = True,
-    ) -> dtf.SessionContext:
-        """Registers the DeltaTable in a DataFusion context."""
-        if name is None:
-            name = f"{self.format}:{self.path}"
-        if ctx is None:
-            ctx = dtf.SessionContext()
-        if lazy:
-            ctx.register_dataset(name, self.to_pyarrow_dataset())
-        else:
-            ctx.register_record_batches(name, [self.to_pyarrow_table().to_batches()])
-            # table = table.to_table()
-        # ctx.register_dataset(name, table)
-        return ctx
+        return self.delta_table.to_pyarrow_table()
+
+    # def to_polars(self, lazy: bool = True) -> pl.DataFrame | pl.LazyFrame:
+    #     """Converts the DeltaTable to a Polars DataFrame."""
+    #     if lazy:
+    #         return pl.scan_pyarrow_dataset(self.to_pyarrow_dataset())
+
+    #     return pl.from_arrow(self.to_pyarrow_table())
+
+    # def to_pandas(self) -> pd.DataFrame:
+    #     """Converts the DeltaTable to a Pandas DataFrame."""
+    #     return self.to_pyrrow_table().to_pandas()
+
+    # def to_duckdb_relation(
+    #     self, conn: duckdb.DuckDBPyConnection | None = None, lazy: bool = True
+    # ) -> duckdb.DuckDBPyRelation:
+    #     """Converts the DeltaTable to a DuckDB relation."""
+    #     if conn is None:
+    #         conn = duckdb.connect()
+    #     if lazy:
+    #         return conn.from_arrow(self.to_pyarrow_dataset())
+    #     return conn.from_arrow(self.to_pyarrow_table())
+
+    # def register_in_duckdb(
+    #     self,
+    #     conn: duckdb.DuckDBPyConnection | None = None,
+    #     name: str | None = None,
+    #     lazy: bool = True,
+    # ) -> duckdb.DuckDBPyConnection:
+    #     """Registers the DeltaTable in a DuckDB connection."""
+    #     if name is None:
+    #         name = f"{self.format}:{self.path}"
+    #     if conn is None:
+    #         conn = duckdb.connect()
+    #     table = self.to_duckdb_relation(conn=conn, lazy=lazy)
+    #     conn.register(name, table)
+    #     return conn
+
+    # def register_in_datafusion(
+    #     self,
+    #     ctx: dtf.SessionContext | None = None,
+    #     name: str | None = None,
+    #     lazy: bool = True,
+    # ) -> dtf.SessionContext:
+    #     """Registers the DeltaTable in a DataFusion context."""
+    #     if name is None:
+    #         name = f"{self.format}:{self.path}"
+    #     if ctx is None:
+    #         ctx = dtf.SessionContext()
+    #     if lazy:
+    #         ctx.register_dataset(name, self.to_pyarrow_dataset())
+    #     else:
+    #         ctx.register_record_batches(name, [self.to_pyarrow_table().to_batches()])
+    #         # table = table.to_table()
+    #     # ctx.register_dataset(name, table)
+    #     return ctx

src/flowerpower/io/saver/deltatable.py

@@ -20,9 +20,10 @@ class DeltaTableWriter(BaseDatasetWriter):
         ```
     """
 
-    description = None
+    description: str | None = None
 
     def model_post_init(self, __context):
+        super().model_post_init(__context)
         self.format = "delta"
 
     def write(
@@ -44,7 +45,8 @@ def write(
                 | pd.DataFrame
                 | dict[str, Any]
             ]
-        ) | None = None,
+        )
+        | None = None,
         mode: str = "append",  # "overwrite" | "append" | "error | "ignore"
         schema: pa.Schema | None = None,
         schema_mode: str | None = None,  # "merge" | "overwrite"
@@ -76,7 +78,7 @@ def write(
             data = [data]
         if isinstance(data[0], dict):
             data = [_dict_to_dataframe(d) for d in data]
-        if isinstance(data[0], pa.LazyFrame):
+        if isinstance(data[0], pl.LazyFrame):
             data = [d.collect() for d in data]
         if isinstance(data[0], pl.DataFrame):
             data = pl.concat(data, how="diagonal_relaxed").to_arrow()

src/flowerpower/mqtt.py

@@ -230,8 +230,8 @@ def from_event_broker(cls, base_dir: str):
         if event_broker_cfg is not None:
             if event_broker_cfg.get("type", None) == "mqtt":
                 return cls(
-                    user=event_broker_cfg.get("user", None),
-                    pw=event_broker_cfg.get("pw", None),
+                    user=event_broker_cfg.get("username", None),
+                    pw=event_broker_cfg.get("password", None),
                     host=event_broker_cfg.get("host", "localhost"),
                     port=event_broker_cfg.get("port", 1883),
                     transport=event_broker_cfg.get("transport", "tcp"),
@@ -387,6 +387,10 @@ def start_pipeline_listener(
     storage_options: dict = {},
     fs: AbstractFileSystem | None = None,
     background: bool = False,
+    host: str | None = None,
+    port: int | None = None,
+    username: str | None = None,
+    password: str | None = None,
     **kwargs,
 ):
     """
@@ -410,6 +414,13 @@ def start_pipeline_listener(
         background: Run the listener in the background
         **kwargs: Additional keyword arguments
     """
+    if host and port:
+        client = MQTTClient(
+            user=username,
+            pw=password,
+            host=host,
+            port=port,
+        )
     client = MQTTClient.from_event_broker(base_dir=base_dir)
     client.start_pipeline_listener(
         name=name,