Bound SQL toolset query results by size, not just row count (#71317)

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

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 request. max_rows bounded rows, which
says nothing about size: one row of a 3000-column table dwarfs a thousand rows
of a narrow one, and the truncation happened after fetching the whole result,
so the worker paid the full transfer cost for rows it then discarded.

Three changes to the query tool of SQLToolset and DataFusionToolset:

- The result is columnar. Column names are serialized once rather than once
  per row, and positional rows keep same-named columns that a dict per row
  silently collapsed.
- max_result_bytes bounds the serialized payload. Rows are dropped from the
  end until it fits, and the result names the limit it hit so the agent can
  narrow its projection instead of paging.
- SQLToolset fetches through DbApiHook.run's handler protocol with fetchmany
  rather than get_records, so rows past max_rows never leave the cursor.

* Address review: trustworthy total_rows, byte-accurate budget, honest docs

- total_rows: rowcount is only a query total on drivers that buffer the whole
  result. python-oracledb reports rows fetched so far, so after a capped fetch
  it equals the cap -- a 10M-row query would report total_rows: 51. Discard any
  count no larger than what was fetched; row_count is already the total when the
  result was not truncated.
- Byte budget: ensure_ascii escaped each CJK character to six bytes instead of
  three, truncating a non-ASCII result several times earlier than the equivalent
  English one and charging the model for the escapes. Measure encoded bytes so
  max_result_bytes means what it says.
- Truncation hint: "No row fits" was false whenever one wide row preceded narrow
  ones, and the partial case -- the common one -- carried no guidance at all.
  Name the row that stopped it, and hint on every byte-capped result.
- A cursor that can neither describe nor fetch now raises instead of rendering
  as an empty table; the Exasol full-fetch path reports its exact total.
- max_rows docs no longer imply the fetch bound reaches the database. No hook
  opens a server-side cursor here, so a client-buffering driver has already
  transferred the rows; only the Python conversion is skipped.
- Changelog note for the output-shape change.

* Correct docs that described the pre-fix truncation behaviour

Rows are returned as a contiguous prefix, stopping at the first that does not fit
the remaining budget. Three places still said "dropped from the end until the
payload fits", which reads as though a wide row would be skipped and later ones
packed in.

Also: the max_rows parameter entry claimed rows beyond it are "never fetched from
the cursor", contradicting the qualification the same page carries further down,
and the query tool description told agents that `truncated` means more rows
matched, when it also fires when the result was simply too large.

Author: kaxil

providers/common/ai/docs/changelog.rst

@@ -25,6 +25,18 @@
 Changelog
 ---------
 
+.. note::
+    The ``query`` tool of ``SQLToolset`` and ``DataFusionToolset`` returns a different
+    shape. Rows were a dict per row alongside a ``count`` of every matching row:
+    ``{"rows": [{"id": 1, "name": "a"}], "count": 900}``. They are now columnar, and
+    ``row_count`` counts the rows actually returned:
+    ``{"columns": ["id", "name"], "rows": [[1, "a"]], "row_count": 1}``. A truncated
+    result also carries ``truncated_by`` -- ``max_rows`` or the new ``max_result_bytes``
+    -- and ``total_rows`` appears only when the driver reports a trustworthy query
+    total. The tool's own description states the new shape, so agents adapt without
+    changes; update any system prompt that describes the old shape, and any code
+    calling ``toolset.call_tool("query", ...)`` directly.
+
 0.7.0
 .....
 

providers/common/ai/docs/toolsets.rst

@@ -128,7 +128,8 @@ Curated toolset wrapping
    * - ``get_schema``
      - Returns column names and types for a table
    * - ``query``
-     - Executes a SQL query and returns rows as JSON
+     - Executes a SQL query and returns bounded, columnar JSON (see
+       :ref:`bounded-query-results`)
    * - ``check_query``
      - Validates SQL syntax without executing it
 
@@ -203,6 +204,62 @@ Parameters
   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 :ref:`bounded-query-results`.
+- ``max_result_bytes``: Budget for the serialized ``query`` result. Default 64 KiB.
+  See :ref:`bounded-query-results`.
+
+.. _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:
+
+.. code-block:: json
+
+    {"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, and ``DataFusionToolset`` materializes the full result in the
+engine before the toolset sees it; in both the payload is bounded but the transfer is
+not.
+
+**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:
+
+.. code-block:: json
+
+    {"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.
+
+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.
 
 ``DataFusionToolset``
 ---------------------
@@ -223,7 +280,8 @@ querying files on object stores (S3, local filesystem, Iceberg) via Apache DataF
    * - ``get_schema``
      - Returns column names and types for a table (Arrow schema)
    * - ``query``
-     - Executes a SQL query and returns rows as JSON
+     - Executes a SQL query and returns bounded, columnar JSON (see
+       :ref:`bounded-query-results`)
 
 Each :class:`~airflow.providers.common.sql.config.DataSourceConfig` entry
 registers a table backed by Parquet, CSV, Avro, or Iceberg data. Multiple
@@ -267,6 +325,8 @@ Parameters
   permitted. DataFusion on object stores is mostly read-only, but it does
   support DDL for in-memory tables; this guard blocks those by default.
 - ``max_rows``: Maximum rows returned from the ``query`` tool. Default ``50``.
+- ``max_result_bytes``: Budget for the serialized ``query`` result. Default 64 KiB.
+  See :ref:`bounded-query-results`.
 
 ``LoggingToolset``
 ------------------
@@ -616,10 +676,14 @@ No single layer is sufficient — they work together.
        ``allowed_functions``. Fail-closed, but only as exact as the SQL parser. Not a
        security boundary -- always pair it with least-privilege database grants. See
        :ref:`allowed-tables-enforcement` below.
-   * - **SQLToolset: max_rows**
-     - Truncates query results to ``max_rows`` (default 50), preventing the
-       agent from pulling entire tables into context.
-     - Does not limit the number of queries the agent can make.
+   * - **SQLToolset: max_rows / max_result_bytes**
+     - Bounds a query result by rows (default 50) and by serialized size
+       (default 64 KiB), preventing the agent from pulling entire tables into
+       context.
+     - Does not limit the number of queries the agent can make, and each result
+       stays in message history for the rest of the run. Rows past ``max_rows``
+       are not read out of the cursor, but a client-buffering driver has already
+       transferred them -- this bounds context, not database or network load.
    * - **MCPToolset: external server**
      - Connects the agent to tools exposed by an MCP server, authenticated
        through an Airflow connection.
@@ -735,7 +799,8 @@ Recommended Configuration
         db_conn_id="analytics_readonly",  # Connection with SELECT-only grants
         allowed_tables=["orders", "customers"],  # Hide other tables from agent
         allow_writes=False,  # Default — validates SQL
-        max_rows=50,  # Default — truncate large results
+        max_rows=50,  # Default — cap rows
+        max_result_bytes=65536,  # Default — cap bytes; lower it for wide tables
     )
 
 **Agents that need to modify data** (use with caution):
@@ -763,8 +828,10 @@ Before deploying an agent task to production:
    agent can call any exposed tool with any arguments.
 4. **Read-only default**: Keep ``allow_writes=False`` unless the task
    specifically requires writes.
-5. **Row limits**: Set ``max_rows`` appropriate to the use case. Large
-   result sets consume LLM context and increase cost.
+5. **Result limits**: Set ``max_rows`` and ``max_result_bytes`` appropriate to
+   the use case. ``max_rows`` alone does not bound size -- on wide tables it is
+   ``max_result_bytes`` that keeps a result from dominating the context window for
+   the rest of the run.
 6. **Model budget**: Configure pydantic-ai's ``model_settings`` (e.g.
    ``max_tokens``) and ``retries`` to bound cost and prevent runaway loops.
 7. **System prompt**: Include safety instructions in ``system_prompt`` (e.g.

providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py

@@ -36,6 +36,11 @@
 from pydantic_ai.tools import ToolDefinition
 from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool
 
+from airflow.providers.common.ai.utils.query_results import (
+    DEFAULT_MAX_RESULT_BYTES,
+    QUERY_TOOL_DESCRIPTION as _QUERY_DESCRIPTION,
+    build_query_result,
+)
 from airflow.providers.common.ai.utils.tool_definition import build_args_validator
 
 if TYPE_CHECKING:
@@ -97,6 +102,15 @@ class DataFusionToolset(AbstractToolset[Any]):
         are permitted.
     :param max_rows: Maximum number of rows returned from the ``query`` tool.
         Default ``50``.
+    :param max_result_bytes: Budget for the serialized ``query`` result, in bytes.
+        Default 64 KiB. ``max_rows`` bounds rows, which says nothing about size: one
+        row of a 3000-column table is larger than a thousand rows of a narrow one, and
+        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 request. Rows are returned as a
+        contiguous prefix, stopping at the first that does not fit the remaining budget
+        rather than skipping it and packing later ones, so one wide row early in the
+        result ends it. The result reports which limit it hit so the agent can narrow
+        its projection rather than page through the table.
     """
 
     def __init__(
@@ -105,12 +119,14 @@ def __init__(
         *,
         allow_writes: bool = False,
         max_rows: int = 50,
+        max_result_bytes: int = DEFAULT_MAX_RESULT_BYTES,
     ) -> None:
         if not datasource_configs:
             raise ValueError("datasource_configs must contain at least one DataSourceConfig")
         self._datasource_configs = datasource_configs
         self._allow_writes = allow_writes
         self._max_rows = max_rows
+        self._max_result_bytes = max_result_bytes
         self._engine: DataFusionEngine | None = None
 
     @property
@@ -133,7 +149,7 @@ async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]:
         for name, description, schema in (
             ("list_tables", "List available table names.", _LIST_TABLES_SCHEMA),
             ("get_schema", "Get column names and types for a table.", _GET_SCHEMA_SCHEMA),
-            ("query", "Execute a SQL query and return rows as JSON.", _QUERY_SCHEMA),
+            ("query", _QUERY_DESCRIPTION, _QUERY_SCHEMA),
         ):
             tool_def = ToolDefinition(
                 name=name,
@@ -199,16 +215,17 @@ def _query(self, sql: str) -> str:
             col_names = list(pydict.keys())
             num_rows = len(next(iter(pydict.values()), []))
 
-            result: list[dict[str, Any]] = [
-                {col: pydict[col][i] for col in col_names} for i in range(min(num_rows, self._max_rows))
-            ]
-
-            truncated = num_rows > self._max_rows
-            output: dict[str, Any] = {"rows": result, "count": num_rows}
-            if truncated:
-                output["truncated"] = True
-                output["max_rows"] = self._max_rows
-            return json.dumps(output, default=str)
+            # DataFusion has already materialised the full result, so unlike SQLToolset
+            # there is nothing left to avoid fetching -- only the payload is bounded.
+            rows = [[pydict[col][i] for col in col_names] for i in range(min(num_rows, self._max_rows))]
+            return build_query_result(
+                col_names,
+                rows,
+                max_rows=self._max_rows,
+                max_result_bytes=self._max_result_bytes,
+                more_rows_available=num_rows > self._max_rows,
+                total_rows=num_rows,
+            )
         except SQLSafetyError as ex:
             log.warning("query failed SQL safety validation: %s", ex)
             raise