Flowfile 0.2.2__py3-none-any.whl

flowfile_worker/secrets.py

@@ -0,0 +1,148 @@
+"""
+Simplified secure storage module for FlowFile worker to read credentials and secrets.
+"""
+from cryptography.fernet import Fernet
+import os
+from pathlib import Path
+import json
+import logging
+from pydantic import SecretStr
+from flowfile_worker.configs import TEST_MODE
+
+# Set up logging
+logger = logging.getLogger(__name__)
+
+
+class SecureStorage:
+    """A secure local storage mechanism for reading secrets using Fernet encryption."""
+
+    def __init__(self):
+        app_data = os.environ.get("APPDATA") or os.path.expanduser("~/.config")
+        self.storage_path = Path(app_data) / "flowfile"
+        logger.debug(f"Using storage path: {self.storage_path}")
+        self.key_path = self.storage_path / ".secret_key"
+
+    def _get_store_path(self, service_name):
+        """Get the path to the encrypted store file for a service."""
+        return self.storage_path / f"{service_name}.json.enc"
+
+    def _read_store(self, service_name):
+        """Read and decrypt the store file for a service."""
+        path = self._get_store_path(service_name)
+        if not path.exists():
+            return {}
+
+        try:
+            with open(self.key_path, "rb") as f:
+                key = f.read()
+            with open(path, "rb") as f:
+                data = f.read()
+
+            return json.loads(Fernet(key).decrypt(data).decode())
+        except Exception as e:
+            logger.debug(f"Error reading from encrypted store: {e}")
+            return {}
+
+    def get_password(self, service_name, username):
+        """Retrieve a password from secure storage."""
+        store = self._read_store(service_name)
+        return store.get(username)
+
+
+_storage = SecureStorage()
+
+
+def get_password(service_name, username):
+    """
+    Retrieve a password from secure storage.
+
+    Args:
+        service_name: The name of the service
+        username: The username or key
+
+    Returns:
+        The stored password or None if not found
+    """
+    return _storage.get_password(service_name, username)
+
+
+def get_docker_secret_key():
+    """
+    Get the master key from Docker secret.
+
+    Returns:
+        str: The master key if successfully read from Docker secret.
+
+    Raises:
+        RuntimeError: If running in Docker but unable to access the secret.
+    """
+    secret_path = "/run/secrets/flowfile_master_key"
+    if os.path.exists(secret_path):
+        try:
+            with open(secret_path, "r") as f:
+                return f.read().strip()
+        except Exception as e:
+            logger.error(f"Failed to read master key from Docker secret: {e}")
+            raise RuntimeError("Failed to read master key from Docker secret")
+    else:
+        logger.critical("Running in Docker but flowfile_master_key secret is not mounted!")
+        raise RuntimeError("Docker secret 'flowfile_master_key' is not mounted")
+
+
+def get_master_key() -> str:
+    """
+    Get the master encryption key.
+
+    If in TEST_MODE, returns a test key.
+    If running in Docker, retrieves the key from Docker secrets.
+    Otherwise, retrieves the key from secure storage.
+
+    Returns:
+        str: The master encryption key
+
+    Raises:
+        ValueError: If the master key is not found in storage.
+    """
+    # First check for test mode
+    if TEST_MODE:
+        return b'06t640eu3AG2FmglZS0n0zrEdqadoT7lYDwgSmKyxE4='.decode()
+
+    # Next check if running in Docker
+    if os.environ.get("RUNNING_IN_DOCKER") == "true":
+        return get_docker_secret_key()
+
+    # Otherwise read from local storage
+    key = get_password("flowfile", "master_key")
+    if not key:
+        raise ValueError("Master key not found in storage.")
+    return key
+
+
+def decrypt_secret(encrypted_value) -> SecretStr:
+    """
+    Decrypt an encrypted value using the master key.
+
+    Args:
+        encrypted_value: The encrypted value as a string
+
+    Returns:
+        SecretStr: The decrypted value as a SecretStr
+    """
+    key = get_master_key().encode()
+    f = Fernet(key)
+    return SecretStr(f.decrypt(encrypted_value.encode()).decode())
+
+
+def encrypt_secret(secret_value):
+    """
+    Encrypt a secret value using the master key.
+
+    Args:
+        secret_value: The secret value to encrypt
+
+    Returns:
+        str: The encrypted value as a string
+    """
+    key = get_master_key().encode()
+    f = Fernet(key)
+    return f.encrypt(secret_value.encode()).decode()

flowfile_worker/spawner.py

@@ -0,0 +1,187 @@
+from flowfile_worker import status_dict
+from time import sleep
+import gc
+from typing import List, Tuple
+from multiprocessing import Process, Queue
+from flowfile_worker.process_manager import ProcessManager
+from flowfile_worker import models, mp_context, funcs, status_dict_lock
+
+# Initialize ProcessManager
+process_manager = ProcessManager()
+
+flowfile_node_id_type = int|str
+
+
+def handle_task(task_id: str, p: Process, progress: mp_context.Value, error_message: mp_context.Array, q: Queue):
+    """
+    Monitors and manages a running process task, updating its status and handling completion/errors.
+
+    Args:
+        task_id (str): Unique identifier for the task
+        p (Process): The multiprocessing Process object being monitored
+        progress (mp_context.Value): Shared value object tracking task progress (0-100)
+        error_message (mp_context.Array): Shared array for storing error messages
+        q (Queue): Queue for storing task results
+
+    Notes:
+        - Updates task status in status_dict while process is running
+        - Handles task cancellation, completion, and error states
+        - Cleans up process resources after completion
+    """
+    try:
+        with status_dict_lock:
+            status_dict[task_id].status = "Processing"
+
+        while p.is_alive():
+            sleep(1)
+            with progress.get_lock():
+                current_progress = progress.value
+            with status_dict_lock:
+                status_dict[task_id].progress = current_progress
+
+                # Check if the task has been cancelled via status_dict
+                if status_dict[task_id].status == "Cancelled":
+                    p.terminate()
+                    break
+
+            if current_progress == -1:
+                with status_dict_lock:
+                    status_dict[task_id].status = "Error"
+                    with error_message.get_lock():
+                        status_dict[task_id].error_message = error_message.value.decode().rstrip('\x00')
+                break
+
+        p.join()
+
+        with status_dict_lock:
+            status = status_dict[task_id]
+            if status.status != "Cancelled":
+                if progress.value == 100:
+                    status.status = "Completed"
+                    if not q.empty():
+                        status.results = q.get()
+                elif progress.value != -1:
+                    status_dict[task_id].status = "Unknown Error"
+
+    finally:
+        if p.is_alive():
+            p.terminate()
+        p.join()
+        process_manager.remove_process(task_id)  # Remove from process manager
+        del p, progress, error_message
+        gc.collect()
+
+
+def start_process(polars_serializable_object: bytes, task_id: str,
+                  operation: models.OperationType,
+                  file_ref: str, flowfile_flow_id: int,
+                  flowfile_node_id: flowfile_node_id_type,
+                  kwargs: dict = None) -> None:
+    """
+    Starts a new process for handling Polars dataframe operations.
+
+    Args:
+        polars_serializable_object (bytes): Serialized Polars dataframe
+        task_id (str): Unique identifier for the task
+        operation (models.OperationType): Type of operation to perform
+        file_ref (str): Reference to the file being processed
+        kwargs (dict, optional): Additional arguments for the operation. Defaults to {}
+        flowfile_flow_id: id of the flow that started the process
+        flowfile_node_id: id of the node that started the process
+
+    Notes:
+        - Creates shared memory objects for progress tracking and error handling
+        - Initializes and starts a new process for the specified operation
+        - Delegates to handle_task for process monitoring
+    """
+    if kwargs is None:
+        kwargs = {}
+    process_task = getattr(funcs, operation)
+    kwargs['polars_serializable_object'] = polars_serializable_object
+    kwargs['progress'] = mp_context.Value('i', 0)
+    kwargs['error_message'] = mp_context.Array('c', 1024)
+    kwargs['queue'] = Queue(maxsize=1)
+    kwargs['file_path'] = file_ref
+    kwargs['flowfile_flow_id'] = flowfile_flow_id
+    kwargs['flowfile_node_id'] = flowfile_node_id
+
+    p: Process = mp_context.Process(target=process_task, kwargs=kwargs)
+    p.start()
+
+    process_manager.add_process(task_id, p)
+    handle_task(task_id=task_id, p=p, progress=kwargs['progress'], error_message=kwargs['error_message'], q=kwargs['queue'])
+
+
+def start_generic_process(func_ref: callable, task_id: str,
+                          file_ref: str, flowfile_flow_id: int,
+                          flowfile_node_id: flowfile_node_id_type, kwargs: dict = None) -> None:
+    """
+    Starts a new process for handling generic function execution.
+
+    Args:
+        func_ref (callable): Reference to the function to be executed
+        task_id (str): Unique identifier for the task
+        file_ref (str): Reference to the file being processed
+        flowfile_flow_id: id of the flow that started the process
+        flowfile_node_id: id of the node that started the process
+        kwargs (dict, optional): Additional arguments for the function. Defaults to None.
+
+    Notes:
+        - Creates shared memory objects for progress tracking and error handling
+        - Initializes and starts a new process for the generic function
+        - Delegates to handle_task for process monitoring
+    """
+    kwargs = {} if kwargs is None else kwargs
+    kwargs['func'] = func_ref
+    kwargs['progress'] = mp_context.Value('i', 0)
+    kwargs['error_message'] = mp_context.Array('c', 1024)
+    kwargs['queue'] = Queue(maxsize=1)
+    kwargs['file_path'] = file_ref
+    kwargs['flowfile_flow_id'] = flowfile_flow_id
+    kwargs['flowfile_node_id'] = flowfile_node_id
+
+    process_task = getattr(funcs, 'generic_task')
+    p: Process = mp_context.Process(target=process_task, kwargs=kwargs)
+    p.start()
+
+    process_manager.add_process(task_id, p)  # Add process to process manager
+    handle_task(task_id=task_id, p=p, progress=kwargs['progress'],
+                error_message=kwargs['error_message'], q=kwargs['queue'])
+
+
+def start_fuzzy_process(left_serializable_object: bytes,
+                        right_serializable_object: bytes,
+                        file_ref: str,
+                        fuzzy_maps: List[models.FuzzyMapping],
+                        task_id: str,
+                        flowfile_flow_id: int,
+                        flowfile_node_id: flowfile_node_id_type) -> None:
+    """
+    Starts a new process for performing fuzzy joining operations on two datasets.
+
+    Args:
+        left_serializable_object (bytes): Serialized left dataframe
+        right_serializable_object (bytes): Serialized right dataframe
+        file_ref (str): Reference to the file being processed
+        fuzzy_maps (List[models.FuzzyMapping]): List of fuzzy mapping configurations
+        task_id (str): Unique identifier for the task
+        flowfile_flow_id: id of the flow that started the process
+        flowfile_node_id: id of the node that started the process
+    Notes:
+        - Creates shared memory objects for progress tracking and error handling
+        - Initializes and starts a new process for fuzzy joining operation
+        - Delegates to handle_task for process monitoring
+    """
+    progress = mp_context.Value('i', 0)
+    error_message = mp_context.Array('c', 1024)
+    q = Queue(maxsize=1)
+
+    args: Tuple[bytes, bytes, List[models.FuzzyMapping], mp_context.Array, str, mp_context.Value, Queue, int, flowfile_node_id_type] = \
+        (left_serializable_object, right_serializable_object, fuzzy_maps, error_message, file_ref, progress, q,
+         flowfile_flow_id, flowfile_node_id)
+
+    p: Process = mp_context.Process(target=funcs.fuzzy_join_task, args=args)
+    p.start()
+
+    process_manager.add_process(task_id, p)  # Add process to process manager
+    handle_task(task_id=task_id, p=p, progress=progress, error_message=error_message, q=q)

flowfile_worker/utils.py

@@ -0,0 +1,25 @@
+import polars as pl
+from polars.exceptions import PanicException
+from dataclasses import dataclass
+
+
+def collect_lazy_frame(lf: pl.LazyFrame) -> pl.DataFrame:
+    try:
+        return lf.collect(engine="streaming")
+    except PanicException:
+        return lf.collect(engine="auto")
+
+
+@dataclass
+class CollectStreamingInfo:
+    __slots__ = 'df', 'streaming_collect_available'
+    df: pl.DataFrame
+    streaming_collect_available: bool
+
+
+def collect_lazy_frame_and_get_streaming_info(lf: pl.LazyFrame) -> CollectStreamingInfo:
+    try:
+        df = lf.collect(engine="streaming")
+        return CollectStreamingInfo(df, True)
+    except PanicException:
+        return CollectStreamingInfo(lf.collect(), False)

test_utils/__init__.py

@@ -0,0 +1,3 @@
+"""Test utilities for Flowfile project."""
+
+__version__ = "0.1.0"

test_utils/postgres/__init__.py

@@ -0,0 +1 @@
+"""PostgreSQL test utilities."""

test_utils/postgres/commands.py

@@ -0,0 +1,109 @@
+"""
+Command-line interface for PostgreSQL test database management.
+
+This module provides command-line functions that can be called via Poetry scripts
+to start and stop PostgreSQL containers with sample data.
+"""
+
+import argparse
+import sys
+import logging
+from . import fixtures
+
+# Set up logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+logger = logging.getLogger("postgres_commands")
+
+
+def start_postgres():
+    """
+    Start PostgreSQL container with sample data.
+
+    This function is the entry point for the 'start_postgres' Poetry script.
+    It parses command-line arguments and starts a PostgreSQL container
+    with the requested configuration.
+
+    Returns:
+        Exit code (0 for success, 1 for failure)
+    """
+    # Check if Docker is available first
+    if not fixtures.is_docker_available():
+        logger.warning("Docker is not available. Cannot start PostgreSQL container.")
+        print("\n" + "=" * 50)
+        print("SKIPPING: Docker is not available on this system")
+        print("Tests requiring Docker will need to be skipped")
+        print("=" * 50 + "\n")
+        return 0  # Return success to allow pipeline to continue
+
+    parser = argparse.ArgumentParser(description="Start PostgreSQL container with sample data")
+    parser.add_argument("--schema", choices=["movies", "stocks"], default="movies",
+                        help="Sample schema to use (movies or stocks)")
+    parser.add_argument("--port", type=int, default=5433,
+                        help="Port to expose PostgreSQL on (default: 5433)")
+    parser.add_argument("--user", default="testuser",
+                        help="PostgreSQL username (default: testuser)")
+    parser.add_argument("--password", default="testpass",
+                        help="PostgreSQL password (default: testpass)")
+    parser.add_argument("--db", default="testdb",
+                        help="PostgreSQL database name (default: testdb)")
+    parser.add_argument("--container-name", default="test-postgres-sample",
+                        help="Docker container name (default: test-postgres-sample)")
+    parser.add_argument("--image-tag", default="test-sample-db",
+                        help="Docker image tag (default: test-sample-db)")
+
+    args = parser.parse_args()
+
+    # Setup and start
+    if fixtures.setup_postgres_samples(args.schema, args.port, args.user, args.password, args.db, args.image_tag):
+        if fixtures.start_postgres_container(args.container_name, args.port, args.image_tag)[1]:
+            fixtures.print_connection_info(
+                "localhost", args.port, args.db, args.user, args.password, args.container_name
+            )
+            return 0
+    return 1
+
+
+def stop_postgres():
+    """
+    Stop PostgreSQL container.
+
+    This function is the entry point for the 'stop_postgres' Poetry script.
+    It parses command-line arguments and stops the specified PostgreSQL container.
+
+    Returns:
+        Exit code (0 for success, 1 for failure)
+    """
+    # Check if Docker is available first
+    if not fixtures.is_docker_available():
+        logger.warning("Docker is not available. No PostgreSQL container to stop.")
+        return 0  # Return success to allow pipeline to continue
+
+    parser = argparse.ArgumentParser(description="Stop PostgreSQL container")
+    parser.add_argument("--container-name", default="test-postgres-sample",
+                        help="Docker container name (default: test-postgres-sample)")
+    parser.add_argument("--timeout", type=int, default=15,
+                        help="Timeout for graceful container stop in seconds (default: 15)")
+
+    args = parser.parse_args()
+
+    if fixtures.stop_postgres_container(args.container_name, args.timeout):
+        print(f"Container {args.container_name} stopped successfully")
+        return 0
+    else:
+        print(f"Failed to stop container {args.container_name}")
+        return 1
+
+
+if __name__ == "__main__":
+    # Allow direct script execution for testing
+    if len(sys.argv) > 1 and sys.argv[1] == "start":
+        sys.exit(start_postgres())
+    elif len(sys.argv) > 1 and sys.argv[1] == "stop":
+        sys.exit(stop_postgres())
+    else:
+        print("Usage: python -m test_utils.postgres.commands [start|stop] [options]")
+        sys.exit(1)