Prefect SQLAlchemy Integration

Warnings

• breaking SQLAlchemy 2.0 introduced significant API changes compared to previous 1.x versions. If you are upgrading your `prefect-sqlalchemy` integration or `SQLAlchemy` itself, ensure your database interaction code is compatible with SQLAlchemy 2.x patterns. `prefect-sqlalchemy` version 0.4.0 and later added explicit support for SQLAlchemy 2.x.
  Fix: Upgrade `prefect-sqlalchemy` to 0.4.0 or higher. Review SQLAlchemy 2.0 migration guides and update database interaction code accordingly, especially for connection and query execution.
• gotcha Database drivers are not included: `prefect-sqlalchemy` provides the integration layer but does not bundle specific database drivers (e.g., `psycopg2-binary` for PostgreSQL, `aiosqlite` for async SQLite, `pymysql` for MySQL). You must install the appropriate driver(s) for your target database(s) separately.
  Fix: Install the necessary database driver using `pip install <driver_package>` (e.g., `pip install psycopg2-binary`) alongside `prefect-sqlalchemy`.
• gotcha Block state management: It is generally recommended to load and consume an `SqlAlchemyConnector` (or `AsyncSqlAlchemyConnector`) within the scope of a single task or flow. Passing the connector instance across separate tasks or flows might lead to loss of connection/cursor state and unexpected behavior.
  Fix: Load the `SqlAlchemyConnector` block inside the task or flow where it is being used, or ensure its lifecycle is properly managed if shared. Using it as a context manager (`with connector: ...`) is the recommended pattern to ensure proper resource closure.
• gotcha Using `SqlAlchemyConnector.load()` requires a pre-existing block document. If you attempt to load a block that hasn't been saved yet (either programmatically or via the Prefect UI), it will result in an error.
  Fix: Before calling `SqlAlchemyConnector.load('your-block-name')`, ensure you have saved a block with that name using `connector_instance.save('your-block-name')` or by configuring it through the Prefect UI.

Install

• pip install prefect-sqlalchemy Basic Installation
• pip install "prefect[sqlalchemy]" Install with Prefect (includes core Prefect)

Imports

• SqlAlchemyConnector

  from prefect_sqlalchemy import SqlAlchemyConnector

• AsyncSqlAlchemyConnector

  from prefect_sqlalchemy import AsyncSqlAlchemyConnector

• ConnectionComponents

  from prefect_sqlalchemy import ConnectionComponents

• SyncDriver

  from prefect_sqlalchemy import SyncDriver

• AsyncDriver

  from prefect_sqlalchemy import AsyncDriver

• DatabaseCredentials

  from prefect_sqlalchemy import SqlAlchemyConnector

  The `DatabaseCredentials` block has been replaced by `SqlAlchemyConnector` in recent versions.

Quickstart

This quickstart demonstrates how to set up an `SqlAlchemyConnector` block, save it to Prefect, and then use it within Prefect tasks and a flow to interact with a SQLite database. It includes creating a table, inserting data, and fetching results. Remember to run the block saving part first, or configure your block via the Prefect UI, before executing the flow.

import os
from prefect import flow, task
from prefect_sqlalchemy import SqlAlchemyConnector, ConnectionComponents, SyncDriver

@task
def setup_table(block_name: str) -> None:
    """Sets up a table and inserts data using the SQLAlchemy connector."""
    with SqlAlchemyConnector.load(block_name) as connector:
        connector.execute(
            "CREATE TABLE IF NOT EXISTS customers (name varchar, address varchar);"
        )
        connector.execute(
            "INSERT INTO customers (name, address) VALUES (:name, :address);",
            parameters={"name": "Marvin", "address": "Highway 42"},
        )
        connector.execute_many(
            "INSERT INTO customers (name, address) VALUES (:name, :address);",
            seq_of_parameters=[
                {"name": "Ford", "address": "Highway 42"},
                {"name": "Unknown", "address": "Highway 42"},
            ],
        )

@task
def fetch_data(block_name: str) -> list:
    """Fetches all data from the customers table."""
    all_rows = []
    with SqlAlchemyConnector.load(block_name) as connector:
        while True:
            new_rows = connector.fetch_many("SELECT * FROM customers", size=2)
            if len(new_rows) == 0:
                break
            all_rows.extend(new_rows)
    return all_rows

@flow(name="SQLAlchemy Flow Example")
def sqlalchemy_flow(block_name: str) -> list:
    """Orchestrates database setup and data fetching."""
    setup_table(block_name)
    all_rows = fetch_data(block_name)
    return all_rows

if __name__ == "__main__":
    # --- Configuration and Block Saving (Run this once to create the block) ---
    # In a real scenario, use environment variables for sensitive info.
    # Replace 'my-sqlite-block' with your desired block name.
    # For SQLite, a file-based database is sufficient.
    sqlite_block_name = os.environ.get('PREFECT_SQL_BLOCK_NAME', 'my-sqlite-block')
    sqlite_db_path = os.environ.get('SQLITE_DB_PATH', 'prefect.db')

    # Create and save the connector block programmatically
    # (Alternatively, create it via the Prefect UI)
    connector = SqlAlchemyConnector(
        connection_info=ConnectionComponents(
            driver=SyncDriver.SQLITE_PYSQLITE,
            database=sqlite_db_path
        )
    )
    connector.save(sqlite_block_name)
    print(f"Saved SqlAlchemyConnector block as '{sqlite_block_name}' pointing to '{sqlite_db_path}'")
    print("You can now run the flow using this block.")

    # --- Running the Flow ---
    # Ensure the block 'my-sqlite-block' exists in your Prefect server/Cloud.
    results = sqlalchemy_flow(sqlite_block_name)
    print("Fetched Data:", results)