Flowfile 0.5.1__py3-none-any.whl → 0.5.3__py3-none-any.whl

flowfile_worker/create/utils.py

@@ -1,13 +1,14 @@
-from faker import Faker
+from collections.abc import Iterable
 from functools import partial
 from random import randint
+
 import polars as pl
-from typing import Iterable
+from faker import Faker
 
 
 def create_fake_data(n_records: int = 1000) -> pl.DataFrame:
     fake = Faker()
-    selector = partial(randint,0)
+    selector = partial(randint, 0)
     min_range = partial(min, n_records)
     # Pre-generation of static data
     cities = [fake.city() for _ in range(min_range(7000))]
@@ -35,18 +36,20 @@ def create_fake_data(n_records: int = 1000) -> pl.DataFrame:
     data = []
     for i in range(n_records):
         name = generate_name()
-        data.append(dict(
-            ID=randint(1, 1000000),
-            Name=name,
-            Address=generate_address(),
-            City=cities[selector(min_range(7000))-1],
-            Email=generate_email(name),
-            Phone=generate_phone_number(),
-            DOB=dob[selector(min_range(100_000))-1],
-            Work=companies[selector(min_range(100_000))-1],
-            Zipcode=zipcodes[selector(min_range(200_000))-1],
-            Country=countries[selector(min_range(50))-1]
-        ))
+        data.append(
+            dict(
+                ID=randint(1, 1000000),
+                Name=name,
+                Address=generate_address(),
+                City=cities[selector(min_range(7000)) - 1],
+                Email=generate_email(name),
+                Phone=generate_phone_number(),
+                DOB=dob[selector(min_range(100_000)) - 1],
+                Work=companies[selector(min_range(100_000)) - 1],
+                Zipcode=zipcodes[selector(min_range(200_000)) - 1],
+                Country=countries[selector(min_range(50)) - 1],
+            )
+        )
 
     return pl.DataFrame(data)
 
@@ -68,7 +71,7 @@ def standardize_col_dtype(vals):
         return [convert_to_string(v) for v in vals]
 
 
-def create_pl_df_type_save(raw_data: Iterable[Iterable], orient: str = 'row') -> pl.DataFrame:
+def create_pl_df_type_save(raw_data: Iterable[Iterable], orient: str = "row") -> pl.DataFrame:
     """
         orient : {'col', 'row'}, default None
         Whether to interpret two-dimensional data as columns or as rows. If None,
@@ -78,7 +81,7 @@ def create_pl_df_type_save(raw_data: Iterable[Iterable], orient: str = 'row') ->
     :param orient:
     :return: polars dataframe
     """
-    if orient == 'row':
-        raw_data = zip(*raw_data)
+    if orient == "row":
+        raw_data = zip(*raw_data, strict=False)
     raw_data = [standardize_col_dtype(values) for values in raw_data]
-    return pl.DataFrame(raw_data, orient='col')
+    return pl.DataFrame(raw_data, orient="col")

flowfile_worker/external_sources/s3_source/main.py

@@ -4,23 +4,17 @@ This module provides functionality to write Polars LazyFrames to various cloud s
 services (S3, Azure ADLS, Google Cloud Storage) in different file formats.
 """
 
-import polars as pl
-from typing import Dict, Any
 from logging import Logger
+from typing import Any
+
+import polars as pl
 
-from flowfile_worker.external_sources.s3_source.models import (
-    CloudStorageWriteSettings,
-    WriteSettings
-)
+from flowfile_worker.external_sources.s3_source.models import CloudStorageWriteSettings, WriteSettings
 from flowfile_worker.utils import collect_lazy_frame
 
 
 def _write_parquet_to_cloud(
-    df: pl.LazyFrame,
-    resource_path: str,
-    storage_options: Dict[str, Any],
-    write_settings: WriteSettings,
-    logger: Logger
+    df: pl.LazyFrame, resource_path: str, storage_options: dict[str, Any], write_settings: WriteSettings, logger: Logger
 ) -> None:
     """Write LazyFrame to a Parquet file in cloud storage.
 
@@ -49,7 +43,7 @@ def _write_parquet_to_cloud(
             # Fall back to collecting and writing if sink fails
             logger.warning(f"Failed to use sink_parquet, falling back to collect and write: {str(e)}")
             pl_df = collect_lazy_frame(df)
-            sink_kwargs['file'] = sink_kwargs.pop("path")
+            sink_kwargs["file"] = sink_kwargs.pop("path")
             pl_df.write_parquet(**sink_kwargs)
 
     except Exception as e:
@@ -58,11 +52,7 @@ def _write_parquet_to_cloud(
 
 
 def _write_delta_to_cloud(
-    df: pl.LazyFrame,
-    resource_path: str,
-    storage_options: Dict[str, Any],
-    write_settings: WriteSettings,
-    logger: Logger
+    df: pl.LazyFrame, resource_path: str, storage_options: dict[str, Any], write_settings: WriteSettings, logger: Logger
 ) -> None:
     """Write LazyFrame to Delta Lake format in cloud storage.
 
@@ -85,11 +75,7 @@ def _write_delta_to_cloud(
 
 
 def _write_csv_to_cloud(
-    df: pl.LazyFrame,
-    resource_path: str,
-    storage_options: Dict[str, Any],
-    write_settings: WriteSettings,
-    logger: Logger
+    df: pl.LazyFrame, resource_path: str, storage_options: dict[str, Any], write_settings: WriteSettings, logger: Logger
 ) -> None:
     """Write LazyFrame to a CSV file in cloud storage.
 
@@ -120,11 +106,7 @@ def _write_csv_to_cloud(
 
 
 def _write_json_to_cloud(
-    df: pl.LazyFrame,
-    resource_path: str,
-    storage_options: Dict[str, Any],
-    write_settings: WriteSettings,
-    logger: Logger
+    df: pl.LazyFrame, resource_path: str, storage_options: dict[str, Any], write_settings: WriteSettings, logger: Logger
 ) -> None:
     """Write LazyFrame to a line-delimited JSON (NDJSON) file in cloud storage.
 
@@ -149,7 +131,7 @@ def _write_json_to_cloud(
         except Exception as e:
             # Fall back to collecting and writing if sink fails
             pl_df = collect_lazy_frame(df)
-            sink_kwargs['file'] = sink_kwargs.pop("path")
+            sink_kwargs["file"] = sink_kwargs.pop("path")
             pl_df.write_ndjson(**sink_kwargs)
             logger.error(f"Failed to use sink_ndjson, falling back to collect and write: {str(e)}")
 
@@ -157,6 +139,7 @@ def _write_json_to_cloud(
         logger.error(f"Failed to write JSON to {resource_path}: {str(e)}")
         raise Exception(f"Failed to write JSON to cloud storage: {str(e)}")
 
+
 writers = {
     "parquet": _write_parquet_to_cloud,
     "delta": _write_delta_to_cloud,
@@ -165,11 +148,7 @@ writers = {
 }
 
 
-def write_df_to_cloud(
-    df: pl.LazyFrame,
-    settings: CloudStorageWriteSettings,
-    logger: Logger
-) -> None:
+def write_df_to_cloud(df: pl.LazyFrame, settings: CloudStorageWriteSettings, logger: Logger) -> None:
     """Write a Polars LazyFrame to an object in cloud storage.
 
     Supports writing to S3, Azure ADLS, and Google Cloud Storage. Currently supports
@@ -187,30 +166,18 @@ def write_df_to_cloud(
     """
     connection = settings.connection
     write_settings = settings.write_settings
-    logger.info(
-        f"Writing to {connection.storage_type} storage: {write_settings.resource_path}"
-    )
+    logger.info(f"Writing to {connection.storage_type} storage: {write_settings.resource_path}")
     # Validate write mode
-    if write_settings.write_mode == 'append' and write_settings.file_format != "delta":
-        raise NotImplementedError(
-            "The 'append' write mode is not yet supported for this destination."
-        )
+    if write_settings.write_mode == "append" and write_settings.file_format != "delta":
+        raise NotImplementedError("The 'append' write mode is not yet supported for this destination.")
 
     storage_options = connection.get_storage_options()
 
     # Dispatch to the appropriate writer
     writer_func = writers.get(write_settings.file_format)
     if not writer_func:
-        raise ValueError(
-            f"Unsupported file format for writing: {write_settings.file_format}"
-        )
-
-    writer_func(
-        df,
-        write_settings.resource_path,
-        storage_options,
-        write_settings,
-        logger
-    )
+        raise ValueError(f"Unsupported file format for writing: {write_settings.file_format}")
+
+    writer_func(df, write_settings.resource_path, storage_options, write_settings, logger)
 
     logger.info(f"Successfully wrote data to {write_settings.resource_path}")

flowfile_worker/external_sources/s3_source/models.py

@@ -1,16 +1,21 @@
 """Cloud storage connection schemas for S3, ADLS, and other cloud providers."""
 
-from typing import Optional, Literal, Dict, Any
+from typing import Any, Literal
+
 import boto3
 from pydantic import BaseModel, SecretStr
+
 from flowfile_worker.secrets import decrypt_secret
 
 CloudStorageType = Literal["s3", "adls", "gcs"]
-AuthMethod = Literal["access_key", "iam_role", "service_principal", "managed_identity", "sas_token", "aws-cli", "env_vars"]
+AuthMethod = Literal[
+    "access_key", "iam_role", "service_principal", "managed_identity", "sas_token", "aws-cli", "env_vars"
+]
 
 
-def create_storage_options_from_boto_credentials(profile_name: Optional[str],
-                                                 region_name: Optional[str] = None) -> Dict[str, Any]:
+def create_storage_options_from_boto_credentials(
+    profile_name: str | None, region_name: str | None = None
+) -> dict[str, Any]:
     """
     Create a storage options dictionary from AWS credentials using a boto3 profile.
     This is the most robust way to handle profile-based authentication as it
@@ -47,29 +52,30 @@ def create_storage_options_from_boto_credentials(profile_name: Optional[str],
 
 class FullCloudStorageConnection(BaseModel):
     """Internal model with decrypted secrets"""
+
     storage_type: CloudStorageType
     auth_method: AuthMethod
-    connection_name: Optional[str] = "None"  # This is the reference to the item we will fetch that contains the data
+    connection_name: str | None = "None"  # This is the reference to the item we will fetch that contains the data
 
     # AWS S3
-    aws_region: Optional[str] = None
-    aws_access_key_id: Optional[str] = None
-    aws_secret_access_key: Optional[SecretStr] = None
-    aws_role_arn: Optional[str] = None
-    aws_allow_unsafe_html: Optional[bool] = None
+    aws_region: str | None = None
+    aws_access_key_id: str | None = None
+    aws_secret_access_key: SecretStr | None = None
+    aws_role_arn: str | None = None
+    aws_allow_unsafe_html: bool | None = None
 
     # Azure ADLS
-    azure_account_name: Optional[str] = None
-    azure_account_key: Optional[SecretStr] = None
-    azure_tenant_id: Optional[str] = None
-    azure_client_id: Optional[str] = None
-    azure_client_secret: Optional[SecretStr] = None
+    azure_account_name: str | None = None
+    azure_account_key: SecretStr | None = None
+    azure_tenant_id: str | None = None
+    azure_client_id: str | None = None
+    azure_client_secret: SecretStr | None = None
 
     # Common
-    endpoint_url: Optional[str] = None
+    endpoint_url: str | None = None
     verify_ssl: bool = True
 
-    def get_storage_options(self) -> Dict[str, Any]:
+    def get_storage_options(self) -> dict[str, Any]:
         """
         Build storage options dict based on the connection type and auth method.
 
@@ -79,15 +85,14 @@ class FullCloudStorageConnection(BaseModel):
         if self.storage_type == "s3":
             return self._get_s3_storage_options()
 
-    def _get_s3_storage_options(self) -> Dict[str, Any]:
+    def _get_s3_storage_options(self) -> dict[str, Any]:
         """Build S3-specific storage options."""
         auth_method = self.auth_method
         print(f"Building S3 storage options for auth_method: '{auth_method}'")
 
         if auth_method == "aws-cli":
             return create_storage_options_from_boto_credentials(
-                profile_name=self.connection_name,
-                region_name=self.aws_region
+                profile_name=self.connection_name, region_name=self.aws_region
             )
 
         storage_options = {}
@@ -103,27 +108,29 @@ class FullCloudStorageConnection(BaseModel):
         if auth_method == "access_key":
             storage_options["aws_access_key_id"] = self.aws_access_key_id
             storage_options["aws_secret_access_key"] = decrypt_secret(
-                self.aws_secret_access_key.get_secret_value()).get_secret_value()
+                self.aws_secret_access_key.get_secret_value()
+            ).get_secret_value()
             # Explicitly clear any session token from the environment
             storage_options["aws_session_token"] = ""
 
         elif auth_method == "iam_role":
             # Correctly implement IAM role assumption using boto3 STS client.
-            sts_client = boto3.client('sts', region_name=self.aws_region)
+            sts_client = boto3.client("sts", region_name=self.aws_region)
             assumed_role_object = sts_client.assume_role(
                 RoleArn=self.aws_role_arn,
-                RoleSessionName="PolarsCloudStorageReaderSession"  # A descriptive session name
+                RoleSessionName="PolarsCloudStorageReaderSession",  # A descriptive session name
             )
-            credentials = assumed_role_object['Credentials']
-            storage_options["aws_access_key_id"] = credentials['AccessKeyId']
-            storage_options["aws_secret_access_key"] = decrypt_secret(credentials['SecretAccessKey']).get_secret_value()
-            storage_options["aws_session_token"] = decrypt_secret(credentials['SessionToken']).get_secret_value()
+            credentials = assumed_role_object["Credentials"]
+            storage_options["aws_access_key_id"] = credentials["AccessKeyId"]
+            storage_options["aws_secret_access_key"] = decrypt_secret(credentials["SecretAccessKey"]).get_secret_value()
+            storage_options["aws_session_token"] = decrypt_secret(credentials["SessionToken"]).get_secret_value()
 
         return storage_options
 
 
 class WriteSettings(BaseModel):
     """Settings for writing to cloud storage"""
+
     resource_path: str  # s3://bucket/path/to/file.csv
 
     write_mode: Literal["overwrite", "append"] = "overwrite"

flowfile_worker/external_sources/sql_source/main.py

@@ -1,6 +1,8 @@
+from io import BytesIO
+
 import polars as pl
+
 from flowfile_worker.external_sources.sql_source.models import DatabaseReadSettings, DatabaseWriteSettings
-from io import BytesIO
 
 
 def write_df_to_database(df: pl.DataFrame, database_write_settings: DatabaseWriteSettings):
@@ -11,9 +13,11 @@ def write_df_to_database(df: pl.DataFrame, database_write_settings: DatabaseWrit
         database_write_settings (DatabaseWriteSettings): The settings for the database connection and table.
     """
     # Write the DataFrame to the database
-    df.write_database(table_name=database_write_settings.table_name,
-                      connection=database_write_settings.connection.create_uri(),
-                      if_table_exists=database_write_settings.if_exists)
+    df.write_database(
+        table_name=database_write_settings.table_name,
+        connection=database_write_settings.connection.create_uri(),
+        if_table_exists=database_write_settings.if_exists,
+    )
     return True
 
 
@@ -53,4 +57,3 @@ def read_sql_source(database_read_settings: DatabaseReadSettings):
     # Read the query into a DataFrame
     df = read_query_as_pd_df(database_read_settings.query, database_read_settings.connection.create_uri())
     return df
-

flowfile_worker/external_sources/sql_source/models.py

@@ -1,17 +1,20 @@
-from typing import Optional, Literal
+from typing import Literal
+
 from pydantic import BaseModel, SecretStr
+
 from flowfile_worker.secrets import decrypt_secret
 
 
 class DataBaseConnection(BaseModel):
     """Database connection configuration with secure password handling."""
-    username: Optional[str] = None
-    password: Optional[SecretStr] = None  # Encrypted password
-    host: Optional[str] = None
-    port: Optional[int] = None
-    database: Optional[str] = None  # The database name
+
+    username: str | None = None
+    password: SecretStr | None = None  # Encrypted password
+    host: str | None = None
+    port: int | None = None
+    database: str | None = None  # The database name
     database_type: str = "postgresql"  # Database type (postgresql, mysql, etc.)
-    url: Optional[str] = None
+    url: str | None = None
 
     def get_decrypted_secret(self) -> SecretStr:
         return decrypt_secret(self.password.get_secret_value())
@@ -48,7 +51,6 @@ class DataBaseConnection(BaseModel):
         if self.port:
             port_section = f":{self.port}"
         if self.database:
-
             base_uri = f"{self.database_type}://{credentials}{self.host}{port_section}/{self.database}"
         else:
             base_uri = f"{self.database_type}://{credentials}{self.host}{port_section}"
@@ -57,6 +59,7 @@ class DataBaseConnection(BaseModel):
 
 class DatabaseReadSettings(BaseModel):
     """Settings for SQL source."""
+
     connection: DataBaseConnection
     query: str
     flowfile_flow_id: int = 1
@@ -65,8 +68,9 @@ class DatabaseReadSettings(BaseModel):
 
 class DatabaseWriteSettings(BaseModel):
     """Settings for SQL sink."""
+
     connection: DataBaseConnection
     table_name: str
-    if_exists: Literal['append', 'replace', 'fail'] = 'append'
+    if_exists: Literal["append", "replace", "fail"] = "append"
     flowfile_flow_id: int = 1
     flowfile_node_id: int | str = -1

flowfile_worker/flow_logger.py

@@ -1,7 +1,9 @@
 import logging
+
 import requests
-from flowfile_worker.models import RawLogInput
+
 from flowfile_worker.configs import FLOWFILE_CORE_URI
+from flowfile_worker.models import RawLogInput
 
 LOGGING_URL = FLOWFILE_CORE_URI + "/raw_logs"
 
@@ -23,12 +25,12 @@ class FlowfileLogHandler(logging.Handler):
                 flowfile_flow_id=self.flowfile_flow_id,
                 log_message=log_message,
                 log_type=record.levelname.upper(),
-                extra={
-                }
+                extra={},
             )
             if self.flowfile_flow_id != -1 and self.flowfile_node_id != -1:
-                response = requests.post(LOGGING_URL, json=raw_log_input.__dict__,
-                                         headers={"Content-Type": "application/json"})
+                response = requests.post(
+                    LOGGING_URL, json=raw_log_input.__dict__, headers={"Content-Type": "application/json"}
+                )
                 if response.status_code != 200:
                     raise Exception(f"Failed to send log: {response.text}")
         except Exception as e:
@@ -45,13 +47,13 @@ def get_worker_logger(flowfile_flow_id: int, flowfile_node_id: int | str) -> log
     if not logger.handlers:
         stream_handler = logging.StreamHandler()
         stream_handler.setLevel(logging.DEBUG)
-        stream_formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+        stream_formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
         stream_handler.setFormatter(stream_formatter)
         logger.addHandler(stream_handler)
 
-        http_handler = FlowfileLogHandler(flowfile_flow_id=flowfile_flow_id, flowfile_node_id = flowfile_node_id)
+        http_handler = FlowfileLogHandler(flowfile_flow_id=flowfile_flow_id, flowfile_node_id=flowfile_node_id)
         http_handler.setLevel(logging.INFO)
-        http_formatter = logging.Formatter('%(message)s')
+        http_formatter = logging.Formatter("%(message)s")
         http_handler.setFormatter(http_formatter)
         logger.addHandler(http_handler)