This is an automated email from the ASF dual-hosted git repository. cgivre pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/drill-mcp.git
commit a843a7258ecb044aed25a308f60c70575deffa07 Author: cgivre <[email protected]> AuthorDate: Wed Aug 12 16:26:29 2026 -0400 fix: enforce max_rows client-side in RestClient.query RestClient.query sent autoLimit to Drill and returned payload["rows"] unsliced, so the row cap depended entirely on Drill honoring that field. JdbcClient.query already enforces the cap client-side via fetchmany(max_rows); slice REST results the same way as defense in depth, keeping autoLimit as the efficient server-side path. --- drill_mcp/client_rest.py | 9 ++++++++- tests/test_client_rest.py | 18 ++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/drill_mcp/client_rest.py b/drill_mcp/client_rest.py index 1ca1461..110fa43 100644 --- a/drill_mcp/client_rest.py +++ b/drill_mcp/client_rest.py @@ -580,11 +580,18 @@ class RestClient: ) payload = _json(response, _safe_url(self._config.url)) rows = payload.get("rows") or [] + truncated = max_rows > 0 and len(rows) >= max_rows + # Defense in depth: `autoLimit` asks Drill to cap rows server-side, + # but the cap must not depend entirely on Drill honoring that field. + # Slice client-side too, exactly like `JdbcClient.query`'s + # `fetchmany(max_rows)` -- the two backends must agree on this. + if max_rows > 0: + rows = rows[:max_rows] return QueryResult( columns=payload.get("columns") or [], rows=rows, query_id=payload.get("queryId"), - truncated=max_rows > 0 and len(rows) >= max_rows, + truncated=truncated, metadata=payload.get("metadata") or [], ) diff --git a/tests/test_client_rest.py b/tests/test_client_rest.py index 72c4299..2aec56f 100644 --- a/tests/test_client_rest.py +++ b/tests/test_client_rest.py @@ -134,6 +134,24 @@ class TestQuery: ) assert make_client().query("SELECT 1", max_rows=2).truncated is True + @respx.mock + def test_slices_rows_to_max_rows_even_if_drill_ignores_autolimit(self): + # `autoLimit` asks Drill to cap rows server-side, but the cap must + # not depend entirely on Drill honoring that field. Simulate Drill + # returning more rows than requested (e.g. an older Drill version, or + # autoLimit simply not being respected) and confirm the client still + # enforces the cap itself -- exactly like JdbcClient.query's + # fetchmany(max_rows). + respx.post(f"{BASE}/query.json").mock( + return_value=httpx.Response( + 200, + json={"columns": ["a"], "rows": [{"a": str(i)} for i in range(10)]}, + ) + ) + result = make_client().query("SELECT 1", max_rows=2) + assert len(result.rows) == 2 + assert result.truncated is True + @respx.mock def test_not_truncated_when_max_rows_is_zero(self): # 0 >= 0 would be a false "truncated" without the max_rows > 0 guard.
