• Home
• Toolsets
• SQL databases: SQLToolset

↑↓ Navigate ⏎ Select Esc Close

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 a table’s columns as JSON, with a name_contains filter and a bounded summary on very wide tables (see Bounded schema results)
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).

Restricting the agent

Every limit can be set on one task: the toolset holds most of them, and the operator holds the tool-call limit. This agent can query two tables, cannot write, gets bounded results, and has a budget for both refused and successful calls:

airflow/providers/common/ai/example_dags/example_sql_toolset.py

    @dag(tags=["example"])
    def example_sql_toolset_restricted():
        AgentOperator(
            task_id="revenue_by_region",
            prompt="What was last week's revenue by region?",
            llm_conn_id="pydanticai_default",
            toolsets=[
                SQLToolset(
                    # Log in as a role granted SELECT on orders and customers only.
                    db_conn_id="warehouse_agent_reader",
                    # Refuse any query that reaches another table or a source the parser cannot check.
                    allowed_tables=["orders", "customers"],
                    # Accept this function, which the SQL parser does not recognize.
                    allowed_functions=["json_build_object"],
                    # The default: allow only SELECT-family, DESCRIBE and SHOW statements.
                    allow_writes=False,
                    # Return at most 100 rows and 32 KiB from one query.
                    max_rows=100,
                    max_result_bytes=32 * 1024,
                    # Summarize a table wider than 50 columns instead of listing every column.
                    max_columns=50,
                    # Let the model correct a refused or failed call up to 3 times.
                    max_retries=3,
                )
            ],
            # Fail the run before a 21st successful tool call, counted across retries on Airflow 3.3+.
            usage_limits=UsageLimits(tool_calls_limit=20),
        )

A refused query never reaches the database. The model gets the reason back as an error it can correct, and the run carries on. Run against a Postgres warehouse that also holds a secrets table, these queries were refused with these messages:

SELECT token FROM secrets

The query tool failed: Query references tables that are not in the allowed tables list: secrets. Use list_tables to see the allowed tables.

SELECT pg_read_file('/etc/passwd')

The query tool failed: Query uses a data source that cannot be checked against allowed_tables: function(s) the parser cannot verify against allowed_tables (pg_read_file); if these functions are trusted, permit them via allowed_functions. Query the allowed tables directly: use list_tables to see them.

DELETE FROM orders

The query tool failed: Statement type 'Delete' is not allowed. Allowed types: Select, Union, Intersect, Except, Describe, Show

Each message ends with the same two lines:

Use the list_tables and get_schema tools to inspect the database, then fix the query and try again.

Fix the errors and try again.

This query runs, because the example lists json_build_object in allowed_functions:

SELECT c.region, json_build_object('revenue', sum(o.amount)) AS revenue
FROM orders o JOIN customers c ON c.id = o.customer_id
GROUP BY c.region ORDER BY c.region

It returns:

{"columns":["region","revenue"],"rows":[["AMER",{"revenue":50}],["EMEA",{"revenue":200}]],"row_count":2}

Without allowed_functions, the same query is refused with the same message as pg_read_file, naming json_build_object.

The two budgets count different calls, and running out of either fails the task. max_retries counts refused and failed calls, and a fourth refused query call in a row ends the run with UnexpectedModelBehavior:

Tool 'query' exceeded max retries count of 3. Consider raising the retry limit, or see the docs on tool retries: https://pydantic.dev/docs/ai/tools-toolsets/tools-advanced/#tool-retries

tool_calls_limit counts successful calls only, and fails the run with UsageLimitExceeded before any call that would take the count past 20:

The next tool call(s) would exceed the tool_calls_limit of 20 (tool_calls=21). Consider raising the limit, or see the docs on usage limits for budget-aware patterns: https://pydantic.dev/docs/ai/core-concepts/agent/#usage-limits

The task’s own retries then decide whether it runs again, and the budgets differ there too. max_retries starts again on each attempt. On Airflow 3.3 and later, usage_limits counts across every attempt, so a retry after tool_calls_limit ran out fails at its first tool call; see the usage budget. UsageLimits also keeps pydantic-ai’s default request_limit of 50 model requests, which refused calls use up too.

allowed_tables works by parsing the SQL, so a query the parser reads differently from the database, or a function listed in allowed_functions, can get past it. The connection’s role is the limit that holds regardless. The example’s warehouse_agent_reader connection logs in as a role created with:

CREATE ROLE warehouse_agent_reader LOGIN PASSWORD '...';
GRANT SELECT ON orders, customers TO warehouse_agent_reader;

With that role and no allowed_tables, SELECT token FROM secrets reaches Postgres, which refuses it, and the model gets The query tool failed: permission denied for table secrets. How allowed_tables is enforced lists what the parser checks and where it stops.

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. To have a person approve a query call before it runs, wrap the toolset with .approval_required() and return tool_def.name == "query" from its function. Reads and writes both go through query, and a task instance can pause for approval once per Dag run, so this fits an agent that runs one query, such as a single write; see Approve an agent’s tool calls.

• 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, and the byte backstop that also triggers the get_schema summary. Default 64 KiB. See Bounded query results and Bounded schema results.

• max_columns: Maximum columns get_schema returns in full. Default 100. Above it the result becomes a bounded summary. See Bounded schema results.

• max_retries: How many times the model may correct a failed call to these tools. Default None, the agent’s retries. See How often the model may correct a failed call.

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.

Bounded schema results

get_schema has the same context problem as query but cannot be solved the same way. Column names are what the agent needs to write SQL, so truncating to the first N columns would leave it unable to reference or discover the rest. The tool filters and summarizes instead.

Filter with name_contains. Pass a case-insensitive substring to get back only the columns whose name contains it, so on a very wide table the agent asks for the columns relevant to its question rather than all of them.

A summary replaces a very wide list. Above max_columns (default 100), or when the serialized columns exceed max_result_bytes, the full list is replaced by a summary that reports the shape and points at the filter:

{"column_count": 3200, "truncated": true, "truncated_by": "max_columns",
 "hint": "...", "type_histogram": {"VARCHAR": 2000, "NUMBER": 1200},
 "sample_columns": [{"name": "id", "type": "NUMBER"}]}

truncated_by is max_columns or max_result_bytes. type_histogram counts columns per type (capped to the most common, the tail folded into one entry), and sample_columns previews the first columns – both only while they fit the budget, so a pathologically small budget still returns the column_count and hint. A filtered call echoes name_contains and adds total_columns so a subset is never mistaken for the whole table.

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 (How often the model may correct a failed call) 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

@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.

Previous Next

Was this entry helpful?