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

flowfile_core/flowfile/connection_manager/_connection_manager.py

@@ -1,9 +1,10 @@
-from typing import Dict, List, Any
+from typing import Any
+
 from flowfile_core.flowfile.connection_manager.models import Connection
 
 
 class ConnectionManager:
-    connections: Dict[str, Dict[str, Connection]]
+    connections: dict[str, dict[str, Connection]]
 
     def add_connection(self, connection_group: str, connection_name: str, connection: Connection):
         existing_connections_in_group = self.connections.get(connection_group)
@@ -34,15 +35,15 @@ class ConnectionManager:
         self.raise_if_connection_does_not_exist(connection_group, connection_name)
         self.connections[connection_group][connection_name] = connection
 
-    def insert_settings_raw(self, connection_group: str, connection_name: str, settings: Dict[str, Any]):
+    def insert_settings_raw(self, connection_group: str, connection_name: str, settings: dict[str, Any]):
         connection = Connection(group=connection_group, name=connection_name, config_setting=settings)
         self.add_connection(connection_group, connection_name, connection)
 
-    def connection_groups(self) -> List[str]:
+    def connection_groups(self) -> list[str]:
         return list(self.connections.keys())
 
     def get_available_connections_in_group(self, group_name: str):
         connection_group = self.connections.get(group_name)
         if connection_group is None:
             return []
-        return list(connection_group.keys())
+        return list(connection_group.keys())

flowfile_core/flowfile/connection_manager/models.py

@@ -1,5 +1,5 @@
-from typing import Any
 from dataclasses import dataclass
+from typing import Any
 
 
 @dataclass

flowfile_core/flowfile/database_connection_manager/db_connections.py

@@ -1,10 +1,12 @@
-from flowfile_core.schemas.input_schema import FullDatabaseConnection, FullDatabaseConnectionInterface
-from flowfile_core.schemas.cloud_storage_schemas import FullCloudStorageConnection, FullCloudStorageConnectionInterface
 from sqlalchemy.orm import Session
-from flowfile_core.database.models import (DatabaseConnection as DBConnectionModel, Secret,
-                                           CloudStorageConnection as DBCloudStorageConnection)
-from flowfile_core.secret_manager.secret_manager import store_secret, SecretInput, decrypt_secret
+
 from flowfile_core.database.connection import get_db_context
+from flowfile_core.database.models import CloudStorageConnection as DBCloudStorageConnection
+from flowfile_core.database.models import DatabaseConnection as DBConnectionModel
+from flowfile_core.database.models import Secret
+from flowfile_core.schemas.cloud_storage_schemas import FullCloudStorageConnection, FullCloudStorageConnectionInterface
+from flowfile_core.schemas.input_schema import FullDatabaseConnection, FullDatabaseConnectionInterface
+from flowfile_core.secret_manager.secret_manager import SecretInput, decrypt_secret, store_secret
 
 
 def store_database_connection(db: Session, connection: FullDatabaseConnection, user_id: int) -> DBConnectionModel:
@@ -32,7 +34,7 @@ def store_database_connection(db: Session, connection: FullDatabaseConnection, u
         username=connection.username,
         password_id=password_id,
         ssl_enabled=connection.ssl_enabled,
-        user_id=user_id
+        user_id=user_id,
     )
 
     # Add and commit the new connection to the database
@@ -47,10 +49,11 @@ def get_database_connection(db: Session, connection_name: str, user_id: int) ->
     """
     Get a database connection by its name and user ID.
     """
-    db_connection = db.query(DBConnectionModel).filter(
-        DBConnectionModel.connection_name == connection_name,
-        DBConnectionModel.user_id == user_id
-    ).first()
+    db_connection = (
+        db.query(DBConnectionModel)
+        .filter(DBConnectionModel.connection_name == connection_name, DBConnectionModel.user_id == user_id)
+        .first()
+    )
 
     return db_connection
 
@@ -59,10 +62,13 @@ def get_cloud_connection(db: Session, connection_name: str, user_id: int) -> DBC
     """
     Get a cloud storage connection by its name and user ID.
     """
-    db_connection = db.query(DBCloudStorageConnection).filter(
-        DBCloudStorageConnection.connection_name == connection_name,
-        DBCloudStorageConnection.user_id == user_id
-    ).first()
+    db_connection = (
+        db.query(DBCloudStorageConnection)
+        .filter(
+            DBCloudStorageConnection.connection_name == connection_name, DBCloudStorageConnection.user_id == user_id
+        )
+        .first()
+    )
 
     return db_connection
 
@@ -87,7 +93,7 @@ def get_database_connection_schema(db: Session, connection_name: str, user_id: i
             database_type=db_connection.database_type,
             username=db_connection.username,
             password=password_secret.encrypted_value,
-            ssl_enabled=db_connection.ssl_enabled
+            ssl_enabled=db_connection.ssl_enabled,
         )
 
     return None
@@ -116,10 +122,11 @@ def delete_database_connection(db: Session, connection_name: str, user_id: int)
     """
     Delete a database connection by its name and user ID.
     """
-    db_connection = db.query(DBConnectionModel).filter(
-        DBConnectionModel.connection_name == connection_name,
-        DBConnectionModel.user_id == user_id
-    ).first()
+    db_connection = (
+        db.query(DBConnectionModel)
+        .filter(DBConnectionModel.connection_name == connection_name, DBConnectionModel.user_id == user_id)
+        .first()
+    )
 
     if db_connection:
         db.delete(db_connection)
@@ -131,7 +138,8 @@ def delete_database_connection(db: Session, connection_name: str, user_id: int)
 
 
 def database_connection_interface_from_db_connection(
-        db_connection: DBConnectionModel) -> FullDatabaseConnectionInterface:
+    db_connection: DBConnectionModel,
+) -> FullDatabaseConnectionInterface:
     """
     Convert a database connection from the database model to the interface model.
     """
@@ -142,7 +150,7 @@ def database_connection_interface_from_db_connection(
         host=db_connection.host,
         port=db_connection.port,
         database=db_connection.database,
-        ssl_enabled=db_connection.ssl_enabled
+        ssl_enabled=db_connection.ssl_enabled,
     )
 
 
@@ -151,9 +159,7 @@ def get_all_database_connections_interface(db: Session, user_id: int) -> list[Fu
     Get all database connections for a user.
     """
     # Get the raw query results
-    query_results = db.query(DBConnectionModel).filter(
-        DBConnectionModel.user_id == user_id
-    ).all()
+    query_results = db.query(DBConnectionModel).filter(DBConnectionModel.user_id == user_id).all()
 
     # Convert with explicit type assertion
     result = []
@@ -168,7 +174,9 @@ def get_all_database_connections_interface(db: Session, user_id: int) -> list[Fu
     return result
 
 
-def store_cloud_connection(db: Session, connection: FullCloudStorageConnection, user_id: int) -> DBCloudStorageConnection:
+def store_cloud_connection(
+    db: Session, connection: FullCloudStorageConnection, user_id: int
+) -> DBCloudStorageConnection:
     """
     Placeholder function to store a cloud database connection.
     This function should be implemented based on specific cloud provider requirements.
@@ -180,20 +188,29 @@ def store_cloud_connection(db: Session, connection: FullCloudStorageConnection,
             f" Please use a unique connection name or delete the existing connection first."
         )
     if connection.aws_secret_access_key is not None:
-        aws_secret_access_key_ref_id = store_secret(db,
-                                                    SecretInput(name=connection.connection_name + "_aws_secret_access_key",
-                                                                value=connection.aws_secret_access_key), user_id).id
+        aws_secret_access_key_ref_id = store_secret(
+            db,
+            SecretInput(
+                name=connection.connection_name + "_aws_secret_access_key", value=connection.aws_secret_access_key
+            ),
+            user_id,
+        ).id
     else:
         aws_secret_access_key_ref_id = None
     if connection.azure_client_secret is not None:
-        azure_client_secret_ref_id = store_secret(db,
-                                                  SecretInput(name=connection.connection_name + "azure_client_secret",
-                                                              value=connection.azure_client_secret), user_id).id
+        azure_client_secret_ref_id = store_secret(
+            db,
+            SecretInput(name=connection.connection_name + "azure_client_secret", value=connection.azure_client_secret),
+            user_id,
+        ).id
     else:
         azure_client_secret_ref_id = None
     if connection.azure_account_key is not None:
-        azure_account_key_ref_id = store_secret(db, SecretInput(name=connection.connection_name + "azure_account_key",
-                                                                value=connection.azure_account_key), user_id).id
+        azure_account_key_ref_id = store_secret(
+            db,
+            SecretInput(name=connection.connection_name + "azure_account_key", value=connection.azure_account_key),
+            user_id,
+        ).id
     else:
         azure_account_key_ref_id = None
 
@@ -202,24 +219,21 @@ def store_cloud_connection(db: Session, connection: FullCloudStorageConnection,
         storage_type=connection.storage_type,
         auth_method=connection.auth_method,
         user_id=user_id,
-
-            # AWS S3 fields
+        # AWS S3 fields
         aws_region=connection.aws_region,
         aws_access_key_id=connection.aws_access_key_id,
         aws_role_arn=connection.aws_role_arn,
         aws_secret_access_key_id=aws_secret_access_key_ref_id,
         aws_allow_unsafe_html=connection.aws_allow_unsafe_html,
-
         # Azure ADLS fields
         azure_account_name=connection.azure_account_name,
         azure_tenant_id=connection.azure_tenant_id,
         azure_client_id=connection.azure_client_id,
         azure_account_key_id=azure_account_key_ref_id,
         azure_client_secret_id=azure_client_secret_ref_id,
-
         # Common fields
         endpoint_url=connection.endpoint_url,
-        verify_ssl=connection.verify_ssl
+        verify_ssl=connection.verify_ssl,
     )
     db.add(db_cloud_connection)
     db.commit()
@@ -228,7 +242,8 @@ def store_cloud_connection(db: Session, connection: FullCloudStorageConnection,
 
 
 def get_full_cloud_storage_interface_from_db(
-        db_cloud_connection: DBCloudStorageConnection) -> FullCloudStorageConnectionInterface:
+    db_cloud_connection: DBCloudStorageConnection,
+) -> FullCloudStorageConnectionInterface:
     """
     Convert a cloud storage connection from the database model to the interface model.
     """
@@ -244,7 +259,7 @@ def get_full_cloud_storage_interface_from_db(
         azure_tenant_id=db_cloud_connection.azure_tenant_id,
         azure_client_id=db_cloud_connection.azure_client_id,
         endpoint_url=db_cloud_connection.endpoint_url,
-        verify_ssl=db_cloud_connection.verify_ssl
+        verify_ssl=db_cloud_connection.verify_ssl,
     )
 
 
@@ -291,12 +306,13 @@ def get_cloud_connection_schema(db: Session, connection_name: str, user_id: int)
         azure_client_id=db_connection.azure_client_id,
         azure_client_secret=azure_client_secret,
         endpoint_url=db_connection.endpoint_url,
-        verify_ssl=db_connection.verify_ssl
+        verify_ssl=db_connection.verify_ssl,
     )
 
 
 def cloud_connection_interface_from_db_connection(
-        db_connection: DBCloudStorageConnection) -> FullCloudStorageConnectionInterface:
+    db_connection: DBCloudStorageConnection,
+) -> FullCloudStorageConnectionInterface:
     """
     Converts a DBCloudStorageConnection model to a FullCloudStorageConnectionInterface model,
     which safely exposes non-sensitive data.
@@ -313,7 +329,7 @@ def cloud_connection_interface_from_db_connection(
         azure_tenant_id=db_connection.azure_tenant_id,
         azure_client_id=db_connection.azure_client_id,
         endpoint_url=db_connection.endpoint_url,
-        verify_ssl=db_connection.verify_ssl
+        verify_ssl=db_connection.verify_ssl,
     )
 
 
@@ -339,7 +355,7 @@ def delete_cloud_connection(db: Session, connection_name: str, user_id: int) ->
             db_connection.aws_session_token_id,
             db_connection.azure_account_key_id,
             db_connection.azure_client_secret_id,
-            db_connection.azure_sas_token_id
+            db_connection.azure_sas_token_id,
         ]
         # Filter out None values
         secret_ids_to_delete = [id for id in secret_ids_to_delete if id is not None]

flowfile_core/flowfile/database_connection_manager/models.py

@@ -12,4 +12,4 @@ class DatabaseConnectionOutput(BaseModel):
     password: str | None = None  # Password can be None if not stored in the database
     ssl_mode: str | None = None  # SSL mode can be None if not applicable
     created_at: str
-    updated_at: str
+    updated_at: str

flowfile_core/flowfile/extensions.py

@@ -1,10 +1,11 @@
+from functools import lru_cache
+
+import polars as pl
+
 from flowfile_core.flowfile._extensions.real_time_interface import get_realtime_func_results
 from flowfile_core.flowfile.flow_node.flow_node import FlowNode
 from flowfile_core.schemas.output_model import InstantFuncResult
 from flowfile_core.utils.arrow_reader import read_top_n
-import pyarrow as pa
-import polars as pl
-from functools import lru_cache
 
 
 @lru_cache(maxsize=16)
@@ -14,23 +15,27 @@ def get_first_row(arrow_path: str) -> pl.DataFrame:
 
 def get_instant_func_results(node_step: FlowNode, func_string: str) -> InstantFuncResult:
     if len(node_step.main_input) == 0:
-        return InstantFuncResult(result='No input data connected, so cannot evaluate the result', success=None)
+        return InstantFuncResult(result="No input data connected, so cannot evaluate the result", success=None)
     node_input = node_step.main_input[0]
     try:
-        if node_input.node_stats.has_run_with_current_setup and node_input.is_setup and node_input.results.example_data_path:
+        if (
+            node_input.node_stats.has_run_with_current_setup
+            and node_input.is_setup
+            and node_input.results.example_data_path
+        ):
             df = get_first_row(node_input.results.example_data_path)
         else:
             df = node_input.get_predicted_resulting_data().data_frame.collect()
     except:
-        return InstantFuncResult(result='Could not get data from previous step', success=None)
+        return InstantFuncResult(result="Could not get data from previous step", success=None)
     try:
-        real_time_result = get_realtime_func_results(df=df,
-                                                     func_string=func_string)
-        if node_step.name == 'filter' and not real_time_result.is_filterable_result():
-            return InstantFuncResult(result='Result is not filterable,'
-                                            ' make sure the function results in a true or false output', success=False)
+        real_time_result = get_realtime_func_results(df=df, func_string=func_string)
+        if node_step.name == "filter" and not real_time_result.is_filterable_result():
+            return InstantFuncResult(
+                result="Result is not filterable," " make sure the function results in a true or false output",
+                success=False,
+            )
         r = InstantFuncResult(result=real_time_result.readable_result, success=real_time_result.success)
     except Exception as e:
         r = InstantFuncResult(result=str(e), success=False)
     return r
-

flowfile_core/flowfile/flow_data_engine/cloud_storage_reader.py

@@ -1,12 +1,15 @@
+from collections.abc import Callable
+from typing import Any, Literal
+
 import boto3
 from botocore.exceptions import ClientError
-from typing import Optional, Dict, Any, Callable, Literal
 
 from flowfile_core.schemas.cloud_storage_schemas import FullCloudStorageConnection
 
 
-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
@@ -45,7 +48,7 @@ class CloudStorageReader:
     """Helper class to handle different cloud storage authentication methods and read operations."""
 
     @staticmethod
-    def get_storage_options(connection: FullCloudStorageConnection) -> Dict[str, Any]:
+    def get_storage_options(connection: FullCloudStorageConnection) -> dict[str, Any]:
         """
         Build storage options dict based on the connection type and auth method.
 
@@ -65,13 +68,12 @@ class CloudStorageReader:
             raise ValueError(f"Unsupported storage type: {connection.storage_type}")
 
     @staticmethod
-    def _get_s3_storage_options(connection: 'FullCloudStorageConnection') -> Dict[str, Any]:
+    def _get_s3_storage_options(connection: "FullCloudStorageConnection") -> dict[str, Any]:
         """Build S3-specific storage options."""
         auth_method = connection.auth_method
         if auth_method == "aws-cli":
             return create_storage_options_from_boto_credentials(
-                profile_name=connection.connection_name,
-                region_name=connection.aws_region
+                profile_name=connection.connection_name, region_name=connection.aws_region
             )
 
         storage_options = {}
@@ -81,7 +83,7 @@ class CloudStorageReader:
             storage_options["endpoint_url"] = connection.endpoint_url
         if not connection.verify_ssl:
             storage_options["verify"] = "False"
-        if connection.aws_allow_unsafe_html: # Note: Polars uses aws_allow_http
+        if connection.aws_allow_unsafe_html:  # Note: Polars uses aws_allow_http
             storage_options["aws_allow_http"] = "true"
 
         if auth_method == "access_key":
@@ -92,20 +94,20 @@ class CloudStorageReader:
 
         elif auth_method == "iam_role":
             # Correctly implement IAM role assumption using boto3 STS client.
-            sts_client = boto3.client('sts', region_name=connection.aws_region)
+            sts_client = boto3.client("sts", region_name=connection.aws_region)
             assumed_role_object = sts_client.assume_role(
                 RoleArn=connection.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"] = credentials['SecretAccessKey']
-            storage_options["aws_session_token"] = credentials['SessionToken']
+            credentials = assumed_role_object["Credentials"]
+            storage_options["aws_access_key_id"] = credentials["AccessKeyId"]
+            storage_options["aws_secret_access_key"] = credentials["SecretAccessKey"]
+            storage_options["aws_session_token"] = credentials["SessionToken"]
 
         return storage_options
 
     @staticmethod
-    def _get_adls_storage_options(connection: 'FullCloudStorageConnection') -> Dict[str, Any]:
+    def _get_adls_storage_options(connection: "FullCloudStorageConnection") -> dict[str, Any]:
         """Build Azure ADLS-specific storage options."""
         storage_options = {}
 
@@ -133,14 +135,14 @@ class CloudStorageReader:
         return storage_options
 
     @staticmethod
-    def _get_gcs_storage_options(connection: 'FullCloudStorageConnection') -> Dict[str, Any]:
+    def _get_gcs_storage_options(connection: "FullCloudStorageConnection") -> dict[str, Any]:
         """Build GCS-specific storage options."""
         # GCS typically uses service account authentication
         # Implementation would depend on how credentials are stored
         return {}
 
     @staticmethod
-    def get_credential_provider(connection: 'FullCloudStorageConnection') -> Optional[Callable]:
+    def get_credential_provider(connection: "FullCloudStorageConnection") -> Callable | None:
         """
         Get a credential provider function if needed for the authentication method.
 
@@ -165,7 +167,7 @@ class CloudStorageReader:
         return None
 
 
-def get_first_file_from_s3_dir(source: str, storage_options: Dict[str, Any] = None) -> str:
+def get_first_file_from_s3_dir(source: str, storage_options: dict[str, Any] = None) -> str:
     """
     Get the first parquet file from an S3 directory path.
 
@@ -188,7 +190,7 @@ def get_first_file_from_s3_dir(source: str, storage_options: Dict[str, Any] = No
     ClientError
         If S3 access fails
     """
-    if not source.startswith('s3://'):
+    if not source.startswith("s3://"):
         raise ValueError("Source must be a valid S3 URI starting with 's3://'")
     bucket_name, prefix = _parse_s3_path(source)
     file_extension = _get_file_extension(source)
@@ -211,39 +213,39 @@ def _get_file_extension(source: str) -> str:
 
 def _parse_s3_path(source: str) -> tuple[str, str]:
     """Parse S3 URI into bucket name and prefix."""
-    path_parts = source[5:].split('/', 1)  # Remove 's3://'
+    path_parts = source[5:].split("/", 1)  # Remove 's3://'
     bucket_name = path_parts[0]
-    prefix = path_parts[1] if len(path_parts) > 1 else ''
+    prefix = path_parts[1] if len(path_parts) > 1 else ""
     return bucket_name, prefix
 
 
 def _remove_wildcards_from_prefix(prefix: str) -> str:
     """Remove wildcard patterns from S3 prefix."""
-    return prefix.split('*')[0]
+    return prefix.split("*")[0]
 
 
-def _create_s3_client(storage_options: Optional[Dict[str, Any]]):
+def _create_s3_client(storage_options: dict[str, Any] | None):
     """Create boto3 S3 client with optional credentials."""
     if storage_options is None:
-        return boto3.client('s3')
+        return boto3.client("s3")
 
     # Handle both 'aws_region' and 'region_name' keys
     client_options = storage_options.copy()
-    if 'aws_region' in client_options:
-        client_options['region_name'] = client_options.pop('aws_region')
+    if "aws_region" in client_options:
+        client_options["region_name"] = client_options.pop("aws_region")
 
-    return boto3.client('s3', **{k: v for k, v in client_options.items() if k != "aws_allow_http"})
+    return boto3.client("s3", **{k: v for k, v in client_options.items() if k != "aws_allow_http"})
 
 
-def _get_first_file(s3_client, bucket_name: str, base_prefix: str, file_extension: str) -> Dict[Any, Any]:
+def _get_first_file(s3_client, bucket_name: str, base_prefix: str, file_extension: str) -> dict[Any, Any]:
     """List all parquet files in S3 bucket with given prefix."""
     try:
-        paginator = s3_client.get_paginator('list_objects_v2')
+        paginator = s3_client.get_paginator("list_objects_v2")
         pages = paginator.paginate(Bucket=bucket_name, Prefix=base_prefix)
         for page in pages:
-            if 'Contents' in page:
-                for obj in page['Contents']:
-                    if obj['Key'].endswith(f".{file_extension}"):
+            if "Contents" in page:
+                for obj in page["Contents"]:
+                    if obj["Key"].endswith(f".{file_extension}"):
                         return obj
             else:
                 raise ValueError(f"No objects found in s3://{bucket_name}/{base_prefix}")