Bound SQL toolset query results by size, not just row count

[kaxil]

SQLToolset's query tool returns a JSON dict per row and caps the result with max_rows. Neither bounds what actually costs money.

A tool result stays in the model's message history for the rest of the agent run, so its size is re-paid on every subsequent model request. On wide tables that gets expensive fast, and max_rows does not help:

• It caps rows, not bytes. One row of a 3000-column table is larger than a thousand rows of a narrow one. A 50-row cap on a wide table still produces a multi-megabyte tool result.
• A dict per row repeats every column name on every row. At 3000 columns and 50 rows, the column names are serialized 50 times over. On a wide result the repeated names, not the values, are the bulk of the payload.
• It truncates after the fetch. hook.get_records(sql) pulls the entire result set into the worker and rows[: self._max_rows] then discards most of it. The full transfer cost is already paid for rows nothing ever sees.

DataFusionToolset's query tool has the same payload shape and the same row-only cap.

Solution

Three changes to the query tool of both toolsets.

Columnar result. Column names are serialized once instead of once per row:

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

Positional rows also fix a correctness bug: SELECT o.id, c.id produced a dict with one id key, silently dropping a column.

max_result_bytes budget (default 64 KiB). Rows are dropped from the end until the serialized payload fits. The result names the limit it hit -- truncated_by is max_rows or max_result_bytes -- so the agent can narrow its projection rather than page through the table. When not even one row fits, or the column names alone exceed the budget, the result carries a hint saying so.

Bounded fetch. SQLToolset fetches through DbApiHook.run's handler protocol with fetchmany(max_rows + 1) instead of get_records. Rows past the cap never leave the cursor; the extra row is what makes "there is more" knowable without fetching the rest.

Measurements

Real SQLite database, 1200-column table, default settings:

	Before	After
Tool result entering message history	2,349,614 B	59,324 B
Worker peak memory, query matching 200k rows	56.6 MiB	0.1 MiB

The memory figure is SQLite, whose driver genuinely stops producing rows. On a
client-buffering driver the result has already been transferred by the time the first
row is read, and only the Python-object conversion is skipped -- a real saving, but a
smaller one. No Airflow hook opens a server-side cursor on this path, so treat the
row cap as a bound on what the agent is shown rather than on database or network load.

Asked to run SELECT * FROM wide and report what it received, a model given no
explanation of the format answered:

I got back only 2 rows, and yes -- there are more rows I did not see, since the result
was truncated by the max_result_bytes byte limit (each row is very wide at 1,200
columns), meaning the query matched more rows than were returned.

which is the property that matters for a shape change: the agent reads a bounded result
as bounded rather than as an empty or complete table, and aligns positional values to
columns without being told how.

Design decisions

Why a byte budget rather than offloading the result to XCom or object storage. Offloading is the better answer for keeping the full result available to downstream tasks, but it needs a result-locator contract and a backend, which belongs with the task-state/checkpoint work rather than in the toolset. Bounding the payload is orthogonal and useful either way -- an offloaded result still needs a bounded digest in the message history.

Why the default budget is generous. 64 KiB is roughly 16k tokens: large enough that ordinary queries are unaffected, small enough that no single result dominates a context window. The columnar shape alone shrinks a wide result several-fold, so results that fit before still fit -- the budget only bites where the payload was already pathological. Deployments whose agents make many queries per run should lower it.

Why fetchmany rather than pushing a LIMIT into the SQL. Rewriting user SQL is dialect-sensitive and changes the meaning of aggregate queries. run(handler=...) is the documented DbApiHook extension point that get_records is itself built on, and the same path SQLExecuteQueryOperator uses.

Why total_rows is often absent. rowcount is only a query total on drivers that buffer the whole result up front. Others report rows fetched so far -- python-oracledb documents exactly that for SELECT -- so after a capped fetch it equals the cap. A ten-million-row Oracle query would otherwise report total_rows: 51, which reads as authoritative and is wrong. A count no larger than what was fetched is indistinguishable from that case, so it is discarded; nothing is lost when the result was not truncated, because row_count is already the total there. Agents that need an exact total can SELECT COUNT(*).

Why the payload is measured in UTF-8 bytes with ensure_ascii=False. With the default escaping, one CJK character costs six bytes instead of three, so an identical result in Japanese would be truncated several times earlier than in English and the model would pay several times the tokens for it.

Tradeoffs and limitations

• The result shape changed. count (total matched) is replaced by row_count (rows returned) plus optional total_rows; rows holds positional lists instead of dicts. The old count was the more confusing of the two -- it reported the full match count next to a truncated rows array. The query tool description states the new shape, so agents get it in-band.
• A wide table fills the default budget with very few rows. At 1200 columns, 64 KiB holds 2 rows. That is the bound doing its job rather than a regression -- the alternative was a 2.3 MB result -- but on tables that wide an agent should be selecting specific columns, and the hint in the result tells it so.
• How much the bounded fetch saves depends on the driver. With a server-side cursor the remaining rows are never sent. A client-buffering driver (psycopg2's default cursor, MySQLdb) has already received them and only skips the per-row conversion. The payload is bounded either way.
• Hooks whose cursor is not DBAPI 2.0 fall back to a full fetch. ExasolHook hands its handler a pyexasol statement that signals "produced rows" through result_type rather than description. Bounding the fetch there needs driver-specific knowledge, so those keep the previous full-fetch behaviour; the payload is still bounded.
• DataFusionToolset bounds the payload only. The engine materializes the full result before the toolset sees it, so there is nothing left to avoid fetching.
• BigQueryHook.get_records' location guard no longer applies. It raises a clear "Need to specify 'location'" error that run() does not; a BigQuery connection without location now fails further down instead. BigQuery is otherwise unaffected by the switch.
• get_schema is still unbounded. On a 3000-column table it returns 3000 column entries, and an agent usually calls it before querying. Truncating it is not obviously right -- those column names are information the agent needs to write SQL at all -- so it wants a different answer (column search or filtering) and is left alone here.

Usage

SQLToolset(
    db_conn_id="analytics_readonly",
    allowed_tables=["orders", "customers"],
    max_rows=50,  # Default -- cap rows
    max_result_bytes=16384,  # Lower than the 64 KiB default for wide tables
)

---

Was generative AI tooling used to co-author this PR?

• Yes (please specify the tool below)

---

• Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
• For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

[Copilot]

Pull request overview

This PR updates the Common AI SQL toolsets to keep query tool results cheap to retain in LLM message history by switching to a columnar result format, adding a max_result_bytes payload budget, and (for SQLToolset) fetching through DbApiHook.run(handler=...) with fetchmany(max_rows + 1) to avoid pulling unneeded rows out of the cursor.

Changes:

• Introduce a shared build_query_result() helper that emits bounded, columnar JSON with truncation metadata (truncated_by, hint, optional total_rows).
• Update SQLToolset and DataFusionToolset query to use the new shape and budget; SQLToolset also switches to a bounded cursor fetch pattern.
• Add/adjust unit tests and provider docs/changelog to describe the new query output contract and configuration.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
providers/common/ai/src/airflow/providers/common/ai/utils/query_results.py	New shared helper for bounded, columnar query tool results and tool description text.
providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py	SQLToolset.query now uses capped fetch via run(handler=...) and emits bounded columnar results; adds max_result_bytes.
providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py	DataFusionToolset.query now emits bounded columnar results; adds max_result_bytes.
providers/common/ai/tests/unit/common/ai/utils/test_query_results.py	New unit tests for payload shape, truncation semantics, and byte accounting.
providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py	Updates/expands tests for new output shape, bounded fetch behavior, and truncation metadata.
providers/common/ai/tests/unit/common/ai/toolsets/test_datafusion.py	Updates tests for new output shape and truncation metadata.
providers/common/ai/docs/toolsets.rst	Documents the new bounded/columnar result contract and new configuration parameter.
providers/common/ai/docs/changelog.rst	Notes the breaking/behavioral change in query tool output shape and guidance for callers/prompts.

Suppressed comments (1)

providers/common/ai/docs/toolsets.rst:244

• This docs section says rows are “dropped from the end until the payload fits”, but the actual implementation (build_query_result) returns a contiguous prefix and stops at the first row that doesn’t fit the byte budget (it won’t skip an oversized row and include later ones). Adjusting the wording here will better set expectations for users reading the docs.

**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 dropped from the end
until the payload fits, and the result says which limit it hit: