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

flowfile_frame/database/__init__.py

@@ -0,0 +1,36 @@
+"""Database module for flowfile_frame.
+
+This module provides functions for:
+- Managing database connections (create, list, delete)
+- Reading from databases
+- Writing to databases
+"""
+
+from flowfile_frame.database.connection_manager import (
+    create_database_connection,
+    create_database_connection_if_not_exists,
+    del_database_connection,
+    get_all_available_database_connections,
+    get_database_connection_by_name,
+)
+from flowfile_frame.database.frame_helpers import (
+    add_read_from_database,
+    add_write_to_database,
+    read_database,
+    write_database,
+)
+
+__all__ = [
+    # Connection management
+    "create_database_connection",
+    "create_database_connection_if_not_exists",
+    "del_database_connection",
+    "get_all_available_database_connections",
+    "get_database_connection_by_name",
+    # FlowGraph helpers
+    "add_read_from_database",
+    "add_write_to_database",
+    # Direct read/write
+    "read_database",
+    "write_database",
+]

flowfile_frame/database/connection_manager.py

@@ -0,0 +1,205 @@
+"""Database connection management for flowfile_frame.
+
+This module provides functions for managing database connections,
+similar to how cloud_storage/secret_manager.py handles cloud storage connections.
+"""
+
+from typing import Literal
+
+from pydantic import SecretStr
+
+from flowfile_core.database.connection import get_db_context
+from flowfile_core.flowfile.database_connection_manager.db_connections import (
+    get_database_connection,
+    get_database_connection_schema,
+    store_database_connection,
+)
+from flowfile_core.schemas.input_schema import (
+    FullDatabaseConnection,
+    FullDatabaseConnectionInterface,
+)
+
+
+def get_current_user_id() -> int:
+    """Get the current user ID for database operations.
+
+    Returns:
+        int: The current user ID (defaults to 1 for single-user mode).
+    """
+    # In single-file mode, we use user_id = 1
+    return 1
+
+
+def create_database_connection(
+    connection_name: str,
+    *,
+    database_type: Literal["postgresql", "mysql", "sqlite", "mssql", "oracle"] = "postgresql",
+    host: str | None = None,
+    port: int | None = None,
+    database: str | None = None,
+    username: str | None = None,
+    password: str | SecretStr | None = None,
+    ssl_enabled: bool = False,
+    url: str | None = None,
+) -> FullDatabaseConnection:
+    """Create and store a new database connection.
+
+    Args:
+        connection_name: Unique name for this connection.
+        database_type: Type of database (postgresql, mysql, sqlite, mssql, oracle).
+        host: Database server hostname.
+        port: Database server port.
+        database: Database name.
+        username: Database username.
+        password: Database password.
+        ssl_enabled: Whether to use SSL for the connection.
+        url: Full database URL (overrides other connection parameters).
+
+    Returns:
+        FullDatabaseConnection: The created connection object.
+
+    Raises:
+        ValueError: If a connection with this name already exists.
+    """
+    user_id = get_current_user_id()
+
+    # Convert password to SecretStr if it's a plain string
+    if isinstance(password, str):
+        password = SecretStr(password)
+
+    connection = FullDatabaseConnection(
+        connection_name=connection_name,
+        database_type=database_type,
+        host=host,
+        port=port,
+        database=database,
+        username=username,
+        password=password,
+        ssl_enabled=ssl_enabled,
+        url=url,
+    )
+
+    with get_db_context() as db:
+        store_database_connection(db, connection, user_id)
+
+    return connection
+
+
+def create_database_connection_if_not_exists(
+    connection_name: str,
+    *,
+    database_type: Literal["postgresql", "mysql", "sqlite", "mssql", "oracle"] = "postgresql",
+    host: str | None = None,
+    port: int | None = None,
+    database: str | None = None,
+    username: str | None = None,
+    password: str | SecretStr | None = None,
+    ssl_enabled: bool = False,
+    url: str | None = None,
+) -> FullDatabaseConnection:
+    """Create a database connection if it doesn't already exist.
+
+    Args:
+        connection_name: Unique name for this connection.
+        database_type: Type of database (postgresql, mysql, sqlite, mssql, oracle).
+        host: Database server hostname.
+        port: Database server port.
+        database: Database name.
+        username: Database username.
+        password: Database password.
+        ssl_enabled: Whether to use SSL for the connection.
+        url: Full database URL (overrides other connection parameters).
+
+    Returns:
+        FullDatabaseConnection: The existing or newly created connection.
+    """
+    user_id = get_current_user_id()
+
+    # Check if connection already exists
+    existing = get_database_connection_by_name(connection_name)
+    if existing:
+        return existing
+
+    return create_database_connection(
+        connection_name,
+        database_type=database_type,
+        host=host,
+        port=port,
+        database=database,
+        username=username,
+        password=password,
+        ssl_enabled=ssl_enabled,
+        url=url,
+    )
+
+
+def get_database_connection_by_name(connection_name: str) -> FullDatabaseConnection | None:
+    """Get a database connection by its name.
+
+    Args:
+        connection_name: The name of the connection to retrieve.
+
+    Returns:
+        FullDatabaseConnection if found, None otherwise.
+    """
+    user_id = get_current_user_id()
+    with get_db_context() as db:
+        return get_database_connection_schema(db, connection_name, user_id)
+
+
+def get_all_available_database_connections() -> list[FullDatabaseConnectionInterface]:
+    """Get all available database connections for the current user.
+
+    Returns:
+        List of database connection interfaces (without passwords).
+    """
+    from flowfile_core.database.models import DatabaseConnection as DBConnectionModel
+
+    user_id = get_current_user_id()
+    with get_db_context() as db:
+        connections = (
+            db.query(DBConnectionModel)
+            .filter(DBConnectionModel.user_id == user_id)
+            .all()
+        )
+
+        return [
+            FullDatabaseConnectionInterface(
+                connection_name=conn.connection_name,
+                database_type=conn.database_type,
+                username=conn.username,
+                host=conn.host,
+                port=conn.port,
+                database=conn.database,
+                ssl_enabled=conn.ssl_enabled,
+            )
+            for conn in connections
+        ]
+
+
+def del_database_connection(connection_name: str) -> bool:
+    """Delete a database connection by its name.
+
+    Args:
+        connection_name: The name of the connection to delete.
+
+    Returns:
+        True if the connection was deleted, False if it didn't exist.
+    """
+    from flowfile_core.database.models import DatabaseConnection as DBConnectionModel
+    from flowfile_core.database.models import Secret
+
+    user_id = get_current_user_id()
+    with get_db_context() as db:
+        connection = get_database_connection(db, connection_name, user_id)
+        if connection:
+            # Delete the associated password secret
+            if connection.password_id:
+                secret = db.query(Secret).filter(Secret.id == connection.password_id).first()
+                if secret:
+                    db.delete(secret)
+
+            db.delete(connection)
+            db.commit()
+            return True
+        return False

flowfile_frame/database/frame_helpers.py

@@ -0,0 +1,249 @@
+"""Database helper functions for FlowFrame operations.
+
+This module provides functions for reading from and writing to databases,
+similar to how cloud_storage/frame_helpers.py handles cloud storage operations.
+"""
+
+from typing import Literal
+
+import polars as pl
+
+from flowfile_core.flowfile.flow_graph import FlowGraph
+from flowfile_core.schemas import input_schema
+from flowfile_frame.database.connection_manager import get_current_user_id
+from flowfile_frame.utils import generate_node_id
+
+
+def add_read_from_database(
+    flow_graph: FlowGraph,
+    *,
+    connection_name: str,
+    table_name: str | None = None,
+    schema_name: str | None = None,
+    query: str | None = None,
+    description: str | None = None,
+) -> int:
+    """Add a database reader node to the flow graph.
+
+    Either table_name or query must be provided. If both are provided,
+    query takes precedence.
+
+    Args:
+        flow_graph: The flow graph to add the node to.
+        connection_name: Name of the stored database connection to use.
+        table_name: Name of the table to read from.
+        schema_name: Database schema name (e.g., 'public' for PostgreSQL).
+        query: SQL query to execute instead of reading a table.
+        description: Optional description for the node.
+
+    Returns:
+        int: The node ID of the created database reader node.
+
+    Raises:
+        ValueError: If neither table_name nor query is provided.
+    """
+    if table_name is None and query is None:
+        raise ValueError("Either 'table_name' or 'query' must be provided")
+
+    node_id = generate_node_id()
+    flow_id = flow_graph.flow_id
+
+    # Determine query mode
+    query_mode: Literal["table", "query"] = "query" if query else "table"
+
+    settings = input_schema.NodeDatabaseReader(
+        flow_id=flow_id,
+        node_id=node_id,
+        user_id=get_current_user_id(),
+        description=description,
+        database_settings=input_schema.DatabaseSettings(
+            connection_mode="reference",
+            database_connection_name=connection_name,
+            query_mode=query_mode,
+            table_name=table_name,
+            schema_name=schema_name,
+            query=query,
+        ),
+    )
+
+    flow_graph.add_database_reader(settings)
+    return node_id
+
+
+def add_write_to_database(
+    flow_graph: FlowGraph,
+    depends_on_node_id: int,
+    *,
+    connection_name: str,
+    table_name: str,
+    schema_name: str | None = None,
+    if_exists: Literal["append", "replace", "fail"] = "append",
+    description: str | None = None,
+) -> int:
+    """Add a database writer node to the flow graph.
+
+    Args:
+        flow_graph: The flow graph to add the node to.
+        depends_on_node_id: The node ID that this writer depends on.
+        connection_name: Name of the stored database connection to use.
+        table_name: Name of the table to write to.
+        schema_name: Database schema name (e.g., 'public' for PostgreSQL).
+        if_exists: What to do if the table already exists:
+            - 'append': Add rows to existing table
+            - 'replace': Drop and recreate table
+            - 'fail': Raise an error
+        description: Optional description for the node.
+
+    Returns:
+        int: The node ID of the created database writer node.
+    """
+    node_id = generate_node_id()
+    flow_id = flow_graph.flow_id
+
+    settings = input_schema.NodeDatabaseWriter(
+        flow_id=flow_id,
+        node_id=node_id,
+        user_id=get_current_user_id(),
+        depending_on_id=depends_on_node_id,
+        description=description,
+        database_write_settings=input_schema.DatabaseWriteSettings(
+            connection_mode="reference",
+            database_connection_name=connection_name,
+            table_name=table_name,
+            schema_name=schema_name,
+            if_exists=if_exists,
+        ),
+    )
+
+    flow_graph.add_database_writer(settings)
+    return node_id
+
+
+def read_database(
+    connection_name: str,
+    *,
+    table_name: str | None = None,
+    schema_name: str | None = None,
+    query: str | None = None,
+) -> pl.LazyFrame:
+    """Read data from a database using a stored connection.
+
+    This is a convenience function for reading data directly without
+    needing to create a FlowGraph.
+
+    Either table_name or query must be provided. If both are provided,
+    query takes precedence.
+
+    Args:
+        connection_name: Name of the stored database connection to use.
+        table_name: Name of the table to read from.
+        schema_name: Database schema name (e.g., 'public' for PostgreSQL).
+        query: SQL query to execute instead of reading a table.
+
+    Returns:
+        pl.LazyFrame: The data read from the database.
+
+    Raises:
+        ValueError: If neither table_name nor query is provided.
+        ValueError: If the connection is not found.
+    """
+    from flowfile_core.flowfile.database_connection_manager.db_connections import get_local_database_connection
+    from flowfile_core.flowfile.sources.external_sources.sql_source.sql_source import SqlSource
+    from flowfile_core.flowfile.sources.external_sources.sql_source.utils import construct_sql_uri
+    from flowfile_core.secret_manager.secret_manager import decrypt_secret
+
+    if table_name is None and query is None:
+        raise ValueError("Either 'table_name' or 'query' must be provided")
+
+    user_id = get_current_user_id()
+    connection = get_local_database_connection(connection_name, user_id)
+
+    if connection is None:
+        raise ValueError(f"Database connection '{connection_name}' not found")
+
+    # Construct the connection string
+    connection_string = construct_sql_uri(
+        database_type=connection.database_type,
+        host=connection.host,
+        port=connection.port,
+        database=connection.database,
+        username=connection.username,
+        password=decrypt_secret(connection.password.get_secret_value()),
+        url=connection.url,
+    )
+
+    # Create SQL source and read data
+    sql_source = SqlSource(
+        connection_string=connection_string,
+        query=query,
+        table_name=table_name,
+        schema_name=schema_name,
+    )
+
+    return sql_source.get_data()
+
+
+def write_database(
+    df: pl.DataFrame | pl.LazyFrame,
+    connection_name: str,
+    table_name: str,
+    *,
+    schema_name: str | None = None,
+    if_exists: Literal["append", "replace", "fail"] = "append",
+) -> None:
+    """Write data to a database using a stored connection.
+
+    This is a convenience function for writing data directly without
+    needing to create a FlowGraph.
+
+    Args:
+        df: The DataFrame or LazyFrame to write.
+        connection_name: Name of the stored database connection to use.
+        table_name: Name of the table to write to.
+        schema_name: Database schema name (e.g., 'public' for PostgreSQL).
+        if_exists: What to do if the table already exists:
+            - 'append': Add rows to existing table
+            - 'replace': Drop and recreate table
+            - 'fail': Raise an error
+
+    Raises:
+        ValueError: If the connection is not found.
+    """
+    from sqlalchemy import create_engine
+
+    from flowfile_core.flowfile.database_connection_manager.db_connections import get_local_database_connection
+    from flowfile_core.flowfile.sources.external_sources.sql_source.utils import construct_sql_uri
+    from flowfile_core.secret_manager.secret_manager import decrypt_secret
+
+    user_id = get_current_user_id()
+    connection = get_local_database_connection(connection_name, user_id)
+
+    if connection is None:
+        raise ValueError(f"Database connection '{connection_name}' not found")
+
+    # Collect if LazyFrame
+    if isinstance(df, pl.LazyFrame):
+        df = df.collect()
+
+    # Construct the connection string
+    connection_string = construct_sql_uri(
+        database_type=connection.database_type,
+        host=connection.host,
+        port=connection.port,
+        database=connection.database,
+        username=connection.username,
+        password=decrypt_secret(connection.password.get_secret_value()),
+        url=connection.url,
+    )
+
+    # Write to database using pandas (polars doesn't have direct SQL write support)
+    engine = create_engine(connection_string)
+    pandas_df = df.to_pandas()
+
+    pandas_df.to_sql(
+        name=table_name,
+        con=engine,
+        schema=schema_name,
+        if_exists=if_exists,
+        index=False,
+    )

flowfile_worker/configs.py

@@ -2,6 +2,7 @@
 
 import argparse
 import logging
+import multiprocessing
 import os
 import platform
 
@@ -15,8 +16,9 @@ logger.setLevel(logging.INFO)
 # Constants for worker and core configuration
 DEFAULT_SERVICE_HOST = "0.0.0.0" if platform.system() != "Windows" else "127.0.0.1"
 DEFAULT_SERVICE_PORT = 63579
-DEFAULT_CORE_HOST = "0.0.0.0" if platform.system() != "Windows" else "127.0.0.1"
-DEFAULT_CORE_PORT = 63578
+# Check environment variable for core host (used in Docker mode)
+DEFAULT_CORE_HOST = os.environ.get("CORE_HOST", "0.0.0.0" if platform.system() != "Windows" else "127.0.0.1")
+DEFAULT_CORE_PORT = int(os.environ.get("CORE_PORT", 63578))
 TEST_MODE = True if "TEST_MODE" in os.environ else False
 
 
@@ -69,19 +71,33 @@ def get_core_url(host, port):
     return f"http://{host}:{port}"
 
 
-# Parse arguments - defaults are already set in the argument parser
-args = parse_args()
+def _is_main_process():
+    """Check if we're running in the main process (not a spawned child)"""
+    return multiprocessing.current_process().name == "MainProcess"
 
-# These variables will already use defaults from argparse if not provided
-SERVICE_HOST = args.host
-SERVICE_PORT = args.port
-CORE_HOST = args.core_host
-CORE_PORT = args.core_port
 
-# Generate the core URI
-FLOWFILE_CORE_URI = get_core_url(CORE_HOST, CORE_PORT)
+# Only parse args and log in the main process
+# Spawned child processes will use environment variables or defaults
+if _is_main_process():
+    # Parse arguments - defaults are already set in the argument parser
+    args = parse_args()
+
+    # These variables will already use defaults from argparse if not provided
+    SERVICE_HOST = args.host
+    SERVICE_PORT = args.port
+    CORE_HOST = args.core_host
+    CORE_PORT = args.core_port
 
-logger.info(f"ConnectorX version: {__version__}")
-# Log configuration
-logger.info(f"Worker configured at {SERVICE_HOST}:{SERVICE_PORT}")
-logger.info(f"Core service configured at {FLOWFILE_CORE_URI}")
+    logger.info(f"ConnectorX version: {__version__}")
+    # Log configuration
+    logger.info(f"Worker configured at {SERVICE_HOST}:{SERVICE_PORT}")
+    logger.info(f"Core service configured at {get_core_url(CORE_HOST, CORE_PORT)}")
+else:
+    # In spawned processes, use defaults from environment variables
+    SERVICE_HOST = DEFAULT_SERVICE_HOST
+    SERVICE_PORT = DEFAULT_SERVICE_PORT
+    CORE_HOST = DEFAULT_CORE_HOST
+    CORE_PORT = DEFAULT_CORE_PORT
+
+# Generate the core URI (used by both main and spawned processes)
+FLOWFILE_CORE_URI = get_core_url(CORE_HOST, CORE_PORT)