rebenitez1802 commented on PR #42244: URL: https://github.com/apache/superset/pull/42244#issuecomment-5038376506
Request changes: sound approach and the size-guard/fail-safe design is solid, but for the single-SELECT `execute_sql` case (the most common one) truncation is defeated by a duplicated row list, so the headline fix mostly doesn't fire there. 🔴 **High — `execute_sql` row list is serialized twice, so row truncation can't shrink it** `truncate_query_result` picks `row_field="rows"` and `_bisect_row_limit` only mutates `data["rows"]`, but in `execute_sql._convert_to_response` the top-level `rows = last_data_stmt.data.rows` is the *same* data also serialized under `statements[*].data.rows` (`superset/mcp_service/sql_lab/tool/execute_sql.py:359-361`). The bisection measures `json.dumps(data)`, which still contains the full `statements` copy. Result for an oversized single `SELECT`: either the untouched `statements` copy alone still exceeds the limit → middleware re-check fails → hard `ToolError` (feature never engages), or it just fits → a partial payload ships where `rows`/`row_count` say "1 of N" while `statements[0].data.rows` + `statements[0].row_count` still hold the full set (metadata lies). Fix: when `row_field == "rows"` and a `statements` array is present, trim `statements[*].data.rows` in lockstep and update each `statements[*].row_count` — or bisect on the last data-bearing statement and derive the top-level slice from it. 🟡 **Medium — `row_field is None` fallback truncates silently (no `_response_truncated`/`_truncation_notes`)** In `truncate_query_result` (`superset/mcp_service/utils/token_utils.py:896-898`) the no-row-field branch delegates to `truncate_oversized_response`, which never sets the truncation markers — the info path relies on the *middleware* to add them (`middleware.py:1326-1329`), but `_try_truncate_data_query_response` does not (`middleware.py:1398`). Scenario: a DML-final multi-statement `execute_sql` where top-level `rows is None` but `statements[]` is large → the client gets a truncated `statements` array with no marker/note and treats partial data as complete. Fix: mirror the info path — set `_response_truncated`/`_truncation_notes` on the dict after a successful truncation in `_try_truncate_data_query_response` (or inside the fallback branch of `truncate_query_result`). 🟡 **Medium — Test gaps around the partial-data signal and the excel/fallback safety nets** Three untested paths could each hide a real regression: (1) no fixture carries `total_rows`, yet the docstring says callers "rely on the true `total_rows`" to detect partial data — nothing asserts it's preserved while `row_count` is overwritten; (2) the `row_field is None` → `truncate_oversized_response` fallback is never exercised; (3) the excel-export → hard `ToolError` guarantee is only unit-tested, never through `on_call_tool`. Fix: add a middleware test with `row_count=200, total_rows=5000` asserting `total_rows` survives; a `truncate_query_result` test with only a huge scalar field; and an `on_call_tool` `get_chart_data` excel test asserting `pytest.raises(ToolError)`. 🟢 **Low — CSV truncation cuts mid-row** `_bisect_string_length` on `csv_data` (`superset/mcp_service/utils/token_utils.py:835`) slices at an arbitrary character offset, so the last kept line is a partial/unterminated record that a strict parser can reject. `excel_data` is correctly left alone (verified). Fix: after computing `kept_len`, back off to the last `\n` at or before it (keeping the header) so only whole rows remain. 🟢 **Low — Broad `except Exception` masks truncation regressions** `_try_truncate_data_query_response` catches bare `Exception` (`middleware.py:1349-1360`) where the info path catches only `(MemoryError, RecursionError)`. It's fail-*safe* (any error → hard `ToolError`, never an over-budget response), but a real `KeyError`/`TypeError` in `truncate_query_result` would be silently swallowed and every oversized query would degrade to a hard error at WARNING level, hiding the bug. Fix: narrow the catch or log at `logger.exception`. --- **Cleared (checked, no issue):** binary searches converge correctly with no off-by-one/infinite-loop; the fail-open path is closed (post-truncation re-check falls back to hard error); placeholder-headroom reservation is sound (real note ≤ placeholder); `estimated_tokens` is measured on the same extracted payload that gets re-wrapped; `DATA_QUERY_TOOLS` names match the registered tools and are disjoint from `INFO_TOOLS`; no schema-validation bypass; mutation operates on fresh copies. Full suite passes locally (**153 passed**, with real tiktoken). The **High** item is the one to resolve before merge — without it, the PR's core promise doesn't hold for the most common `execute_sql` path. <sub>🤖 Team review via Claude Code</sub> -- 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] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
