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 665a7e258910ef5b738ab7fcdfdb2c1d29f37869 Author: cgivre <[email protected]> AuthorDate: Wed Aug 12 16:34:19 2026 -0400 Fix test env leakage and enforce row cap in the tool layer Add an autouse conftest fixture that strips DRILL_* env vars before each test, making the suite hermetic regardless of the ambient shell environment (test_cli_flags_reach_the_config goes through main(), which legitimately reads os.environ, so per-test env={} alone cannot fix it). Also slice DrillTools.run_query's rows to the effective max_rows limit after hidden-schema filtering. Both backend clients already cap rows, but run_query is the last chokepoint before the model, so it should not simply trust what the client returns. --- drill_mcp/server.py | 6 ++++++ tests/conftest.py | 42 ++++++++++++++++++++++++++++++++++++++++++ tests/test_server.py | 12 ++++++++++++ 3 files changed, 60 insertions(+) diff --git a/drill_mcp/server.py b/drill_mcp/server.py index c4c4bf3..9383553 100644 --- a/drill_mcp/server.py +++ b/drill_mcp/server.py @@ -150,6 +150,12 @@ class DrillTools: if self._policy.hidden_schemas and is_show_command(sql): rows = [row for row in rows if self._visible(_first_value(row))] + # Belt-and-suspenders: both backend clients already cap `rows` to + # `limit`, but this tool method is the last chokepoint before the + # model sees the data, so the cap is enforced here too rather than + # trusted from below. + rows = rows[:limit] + payload: dict[str, Any] = { "columns": result.columns, "rows": rows, diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..a2ae2b9 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,42 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, +# either express or implied. See the License for the specific +# language governing permissions and limitations under the +# License. +# + +"""Shared pytest fixtures. + +The suite must be hermetic: nothing in a developer's or CI runner's real +shell environment should be able to change test outcomes. `main()` and +`load_config()` read `os.environ` directly when no explicit `env=` is +passed, so a stray `DRILL_*` variable in the ambient environment (e.g. a +developer's local `.env` sourced into their shell) can otherwise leak into +a test that assumes defaults. Strip every `DRILL_*` variable before each +test runs so the whole suite behaves the same regardless of the ambient +environment. +""" + +import os + +import pytest + + [email protected](autouse=True) +def _clean_drill_env(monkeypatch): + """Remove all DRILL_* environment variables for the duration of a test.""" + for key in list(os.environ): + if key.startswith("DRILL_"): + monkeypatch.delenv(key, raising=False) diff --git a/tests/test_server.py b/tests/test_server.py index c323598..f8cb08e 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -48,6 +48,18 @@ class TestRunQuery: make_tools(client, max_rows=100).run_query("SELECT 1") assert client.query.call_args.kwargs["max_rows"] == 100 + def test_enforces_the_row_cap_even_if_the_client_returns_more(self): + # RestClient and JdbcClient both cap rows themselves, but run_query + # is the last chokepoint before the model, so it must not simply + # trust whatever the client hands back. + client = MagicMock() + client.query.return_value = QueryResult( + ["a"], [{"a": i} for i in range(10)], "q1", False + ) + result = make_tools(client, max_rows=3).run_query("SELECT 1") + assert len(result["rows"]) == 3 + assert result["rows"] == [{"a": 0}, {"a": 1}, {"a": 2}] + def test_caller_may_lower_the_cap(self): client = MagicMock() client.query.return_value = QueryResult()
