SQL databases: `SQLToolset`¶

Curated toolset wrapping DbApiHook with four tools:

Tool	Description
`list_tables`	Lists available table names (filtered by `allowed_tables` if set)
`get_schema`	Returns column names and types for a table
`query`	Executes a SQL query and returns bounded, columnar JSON (see Bounded query results)
`check_query`	Validates SQL syntax without executing it

from airflow.providers.common.ai.toolsets.sql import SQLToolset

toolset = SQLToolset(
    db_conn_id="postgres_default",
    allowed_tables=["customers", "orders"],
    max_rows=20,
)

The DbApiHook is resolved lazily from db_conn_id on first tool call via BaseHook.get_connection(conn_id).get_hook().

In read-only mode (allow_writes=False, the default) the query tool also accepts read-only metadata statements – DESCRIBE/DESC and SHOW – in addition to SELECT-family queries. Agents commonly open with DESCRIBE to learn a table’s columns, so permitting it keeps runs deterministic instead of hard-failing on schema discovery. The toolset passes the connection’s dialect to the validator, so SHOW is recognized on databases that support it (Snowflake, MySQL, etc.); on databases without SHOW it stays rejected. Data-modifying statements remain blocked – including ones hidden behind DESCRIBE/EXPLAIN (e.g. EXPLAIN DELETE ..., DESCRIBE DROP TABLE ...), which the validator rejects by scanning the parsed statement for write operations. When allowed_tables is set it scopes these statements too: a DESCRIBE names a table, so its target must be on the list, while SHOW enumerates objects beyond any single table and is rejected outright (see How allowed_tables is enforced).

Multi-schema warehouses¶

When an agent’s tables live in several schemas of one database – common on Snowflake – list them with schema-qualified allowed_tables entries:

SQLToolset(
    db_conn_id="snowflake_hq",
    allowed_tables=["MODEL_ASTRO.DEPLOYMENT_IMAGE_DETAILS", "MODEL_CRM.SF_ASTRO_ORGS"],
)

list_tables then introspects each referenced schema and returns the matching tables fully qualified (e.g. MODEL_ASTRO.DEPLOYMENT_IMAGE_DETAILS), and get_schema routes each qualified name to its own schema. Without this, a single schema only covers one namespace, and leaving schema unset made introspection query a literal "None" schema and fail. Unqualified entries fall back to schema, and table-name matching is case-insensitive (databases reflect identifiers in their own case). For tables in a different database, use a separate toolset whose connection points at that database.

Templated connection IDs¶

db_conn_id is a Jinja template, rendered for each task instance just before it runs, so one toolset definition can reach a different database depending on where and for what the task runs:

Per environment. The same Dag reads the staging warehouse in staging and the production one in production, with the environment name kept in a Variable: SQLToolset(db_conn_id="warehouse_{{ var.value.environment }}").
Per unit of work. A mapped task gives each map index its own connection – one per customer, region, or shard – as in the example below.

Each task instance renders its own copy of the toolset, so the object in the Dag file keeps its template and no rendered connection carries over to another task instance. The task log records which connection each instance got, as a Rendered toolset sql-warehouse_prod line. A toolset wrapped with .prefixed() or .filtered(), passed as a Toolset capability, or passed in agent_params["toolsets"] is rendered the same way. A Toolset capability built from a callable is resolved when the run starts and is not rendered. MCPToolset.mcp_conn_id and HookToolset’s hook connection ID are templated the same way (see Templated connection IDs).

Warning

Build the connection ID from values the Dag controls – a Variable, upstream task output – not from params or dag_run.conf. Whoever triggers the Dag controls those, and a task can read any connection it names, so a templated db_conn_id taken from trigger input lets the trigger pick the database.

Only the connection ID is templated. allowed_tables is validated when the toolset is created, so a template in it stays a literal table name. DataFusionToolset takes data source configs and is not templated.

One connection per customer¶

Customer-facing analytics must only ever read one customer’s rows. The boundary that holds is the database’s own: a role or database per customer, reached through its own Airflow connection. A mapped agent task can give each customer’s task instance that customer’s connection:

from airflow.providers.common.ai.toolsets.sql import SQLToolset
from airflow.sdk import dag, task


@dag
def customer_reports():
    @task
    def customers() -> list[str]:
        return ["acme", "globex"]

    @task.agent(
        llm_conn_id="pydanticai_default",
        toolsets=[SQLToolset(db_conn_id="analytics_{{ task.op_kwargs.customer }}")],
    )
    def report(customer: str) -> str:
        return f"Summarize this month's orders for {customer}."

    report.expand(customer=customers())


customer_reports()

The acme task instance queries through analytics_acme and the globex one through analytics_globex. Write {{ task.op_kwargs.customer }}, not {{ customer }}: the task’s arguments are not template variables, and the undefined name fails the task.

AgentOperator mapped over prompts has no customer argument to read, so the connection has to come from the map index. That works when the prompts are built from the same list, in the same order, as the one the template indexes:

names = customers()
AgentOperator.partial(
    task_id="report",
    llm_conn_id="pydanticai_default",
    toolsets=[SQLToolset(db_conn_id="analytics_{{ ti.xcom_pull(task_ids='customers')[ti.map_index] }}")],
).expand(prompt=names.map(lambda name: f"Summarize this month's orders for {name}."))

Prefer @task.agent where you can: task.op_kwargs names the customer directly instead of relying on the two lists lining up.

Parameters¶

db_conn_id: Airflow connection ID for the database. Templated (see Templated connection IDs).
allowed_tables: Restrict the agent to a fixed set of tables. Omit the argument (the default) to expose all tables in schema. No value means allow-all: None and an empty list both raise ValueError, so an allow-list built at runtime that resolves to nothing fails at import instead of silently exposing every table. Entries may be schema-qualified ("SCHEMA.TABLE") to span multiple schemas; see above. Matching is case-insensitive. When set, the list is enforced on query and check_query as well as discovery – every table a query references must be on it. See How allowed_tables is enforced for what this does and does not guarantee.
allowed_functions: Names of functions that sqlglot does not recognize as builtins but are safe to run while allowed_tables is active (e.g. ["json_build_object"] or a project UDF). None (default) rejects every unrecognized function. Matching is case-insensitive. Only consulted when allowed_tables is set.
schema: Default schema/namespace for unqualified table listing and introspection. Schema-qualified allowed_tables entries override it per table.
allow_writes: Allow data-modifying SQL (INSERT, UPDATE, DELETE, etc.). Default False – only SELECT-family and read-only metadata (DESCRIBE/SHOW) statements are permitted.
max_rows: Maximum rows returned from the query tool. Default 50. Rows beyond it are not read out of a DBAPI cursor; what the driver has already transferred is its own call. See Bounded query results.
max_result_bytes: Budget for the serialized query result. Default 64 KiB. See Bounded query results.

Bounded query results¶

A tool result stays in the model’s message history for the rest of the run, so its cost is re-paid on every subsequent model request. The query tool of both SQLToolset and DataFusionToolset bounds that in three ways.

The result is columnar. Column names appear once, not once per row:

{"columns": ["id", "name"], "rows": [[1, "Alice"], [2, "Bob"]], "row_count": 2}

On a table with thousands of columns the repeated names, not the values, are the bulk of a row-of-dicts payload. Positional rows also keep columns that share a name – SELECT o.id, c.id – which a dict per row silently collapsed to one.

Rows are fetched, not filtered. max_rows bounds what leaves the cursor, so a query matching a whole table costs the worker roughly what one matching max_rows costs. How much is saved depends on the driver: with a server-side cursor the remaining rows are never sent, while a client-buffering driver (psycopg2’s default cursor, MySQLdb) has already received them and only the per-row conversion is skipped. Hooks whose cursor is not DBAPI 2.0 (ExasolHook passes a pyexasol statement) fall back to a full fetch, where the payload is bounded but the transfer is not. DataFusionToolset pushes the bound into the query instead: it runs the statement with a DataFusion LIMIT of max_rows + 1, so the engine never materializes more than that and the extra row only signals truncation.

A byte budget bounds the payload. max_rows caps rows, which says nothing about size – one row of a 3000-column table is larger than a thousand rows of a narrow one. max_result_bytes is what actually bounds context. Rows are returned as a contiguous prefix: the result stops at the first row that does not fit the remaining budget rather than skipping it and packing later ones, so a single wide row early in the result ends it. The result says which limit it hit:

{"columns": ["..."], "rows": ["..."], "row_count": 3,
 "truncated": true, "truncated_by": "max_result_bytes"}

truncated_by is max_rows or max_result_bytes. When not even one row fits, or the column names alone exceed the budget, the result carries a hint telling the agent to narrow its projection – the only move that helps. total_rows is present when the driver reports a row count for the query; several (SQLite, some warehouse drivers) do not, and it is then omitted rather than guessed. DataFusionToolset never reports it, because it reads only max_rows + 1 rows and so has no total to report; an agent that needs one runs COUNT(*).

The default budget is deliberately generous: the columnar shape alone shrinks a wide result several-fold, so results that fit before still fit. Lower max_result_bytes when an agent makes many queries in one run, since every result is re-paid on every later request.

When to choose it¶

Choose it when the question is a query and the data is in a DBAPI database. SQLToolset gives the agent four tools: list tables, get schema, query, check query. Set allowed_tables and that allow-list is enforced by parsing the SQL rather than by matching strings; see How allowed_tables is enforced for how the walk handles CTEs, subqueries and joins.

What it cannot do

allowed_tables is an application-level guardrail, not a replacement for database permissions. Its own docstring says so, and names the residual gap: an engine or query the parser reads differently. Point db_conn_id at a least-privilege role whose grants match the allow-list.
It cannot bound the fetch for every driver. Hooks that hand their handler something other than a DBAPI cursor (ExasolHook and its pyexasol statement, for instance) fall back to a full fetch. The payload handed to the model is still bounded; the transfer is not. See Bounded query results.
Its parser-level closure is opt-in, not the default. The table walk returns immediately while allowed_tables is unset, so out of the box the agent reaches every table the connection can see. Only omitting the argument grants that: an explicit None or empty list is rejected at construction. DESCRIBE and SHOW both pass, on dialects that parse them, while allowed_tables stays unset. Set allowed_tables and the walk turns fail-closed for SHOW, but not for DESCRIBE: it instead becomes an ordinary table reference, allowed only when the table it names is on the list. See How allowed_tables is enforced for what else the walk rejects once allowed_tables is active. Statements that modify data, COPY among them, are rejected either way while allow_writes is False.
It does not classify failures. A connection error or a typo in a column name reaching list_tables, get_schema or query becomes one ModelRetry, so the two are treated the same way until the retry budget runs out and the task fails for Airflow to retry. Two paths do not raise: check_query catches its own errors and reports them back as a normal {"valid": false, ...} result, and get_schema returns a normal {"error": ...} result instead of raising when the requested table is outside allowed_tables; other get_schema failures still raise and still become a ModelRetry.

A real example. example_pydantic_ai_hook.py builds an agent around SQLToolset inside a plain @task function, with no operator involved:

airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py[source]

@dag(schedule=None, tags=["example"])
def example_task_with_toolsets():
    """Use toolsets directly in a @task function without AgentOperator."""

    @task
    def analyze_revenue() -> str:
        from airflow.providers.common.ai.toolsets.sql import SQLToolset

        hook = PydanticAIHook(llm_conn_id="pydanticai_default")
        agent = hook.create_agent(
            output_type=str,
            instructions=(
                "You are a sales analytics assistant. "
                "Use the SQL tools to explore the database schema and answer questions."
            ),
            toolsets=[
                SQLToolset(
                    db_conn_id="my_database",
                    allowed_tables=["customers", "orders"],
                    max_rows=20,
                ),
            ],
        )
        result = agent.run_sync("Which customers have spent the most? Show the top 5.")
        return result.output

    analyze_revenue()

Credentials and where it runs. db_conn_id is resolved through BaseHook.get_connection, and the connection must supply a DbApiHook. Queries run from the Airflow worker process against the database. Its tool calls act as barriers, as they do for the other routes that build their own tools; see Tool calls as barriers.

SQL databases: SQLToolset¶