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

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,
+    )