kaxil opened a new pull request, #71317: URL: https://github.com/apache/airflow/pull/71317
`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)`](https://github.com/apache/airflow/blob/274011b437/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py#L373) pulls the entire result set into the worker and [`rows[: self._max_rows]`](https://github.com/apache/airflow/blob/274011b437/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py#L381) 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: ```json {"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 ```python 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 ) ``` <!-- SPDX-License-Identifier: Apache-2.0 https://www.apache.org/licenses/LICENSE-2.0 --> <!-- Thank you for contributing! Please provide above a brief description of the changes made in this pull request. Write a good git commit message following this guide: https://chris.beams.io/posts/git-commit/ Please make sure that your code changes are covered with tests. And in case of new features or big changes remember to adjust the documentation. For user-facing UI changes, please attach before/after screenshots (or a short screen recording) so reviewers can assess the visual impact. Feel free to ping (in general) for the review if you do not see reaction for a few days (72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications. In case of an existing issue, reference it using one of the following: * closes: #ISSUE * related: #ISSUE --> --- ##### Was generative AI tooling used to co-author this PR? <!-- If generative AI tooling has been used in the process of authoring this PR, please change below checkbox to `[X]` followed by the name of the tool, uncomment the "Generated-by". --> - [ ] Yes (please specify the tool below) <!-- Generated-by: [Tool Name] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions) --> --- * Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#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](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed. * When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x). * For significant user-facing changes create newsfragment: `{pr_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments). You can add this file in a follow-up commit after the PR is created so you know the PR number. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
