This is an automated email from the ASF dual-hosted git repository. cgivre pushed a commit to branch feat/drill-mcp-server in repository https://gitbox.apache.org/repos/asf/drill-mcp.git
commit 8f413e7b861d0e7593819846a2b1886b30bc166f Author: cgivre <[email protected]> AuthorDate: Wed Aug 12 01:34:25 2026 -0400 fix: never let unexpected exceptions escape run_query as raw tracebacks Broaden guard._check's parse-failure handling from SqlglotError to Exception so a RecursionError from adversarial input (deeply nested parens) becomes a PolicyError instead of an uncaught crash. Add a catch-all in DrillTools.run_query around the policy check as a second line of defense, deliberately without embedding the exception text (which could carry a path). Validate sql/max_rows argument types at the tool boundary. Strengthen list_tables/ describe_table tests to assert forwarded arguments, fix _visible's fail-open truthiness, and split the tools section comment for Task 9. --- drill_mcp/guard.py | 11 +++++++---- drill_mcp/server.py | 17 +++++++++++++++-- tests/test_guard.py | 8 ++++++++ tests/test_server.py | 25 +++++++++++++++++++++++++ 4 files changed, 55 insertions(+), 6 deletions(-) diff --git a/drill_mcp/guard.py b/drill_mcp/guard.py index 9eb0066..ce7b8a3 100644 --- a/drill_mcp/guard.py +++ b/drill_mcp/guard.py @@ -101,10 +101,13 @@ def _check(sql: str, policy: Policy, depth: int) -> None: try: statements = [s for s in sqlglot.parse(sql, read=DIALECT) if s is not None] - except sqlglot.errors.SqlglotError as exc: - # Catches both ParseError and TokenError (siblings under SqlglotError, - # not parent/child) — malformed input such as an unterminated string - # literal must be a PolicyError, never an uncaught crash. + except Exception as exc: + # Catches ParseError and TokenError (siblings under SqlglotError, not + # parent/child) as well as anything else the parser can throw on + # adversarial input, such as a RecursionError from deeply nested + # parentheses. A parse failure of any kind must be a PolicyError, + # never an uncaught crash — scoped to this call only, so a PolicyError + # raised by the policy logic further down is never caught here. raise PolicyError( f"could not parse SQL, so it cannot be checked against policy and is rejected: {exc}" ) from exc diff --git a/drill_mcp/server.py b/drill_mcp/server.py index 0f9af37..5576ebf 100644 --- a/drill_mcp/server.py +++ b/drill_mcp/server.py @@ -56,16 +56,29 @@ class DrillTools: raise ToolError(f"schema '{schema}' is hidden by configuration") def _visible(self, schema: str | None) -> bool: - return not (schema and matches_prefix(schema, self._policy.hidden_schemas)) + return not matches_prefix(schema or "", self._policy.hidden_schemas) - # -- tools ------------------------------------------------------------- + # -- query and metadata tools ------------------------------------------- def run_query(self, sql: str, max_rows: int | None = None) -> dict[str, Any]: """Run a single SQL statement against Drill and return its rows.""" + if not isinstance(sql, str): + raise ToolError("sql must be a string") + if max_rows is not None and not isinstance(max_rows, int): + raise ToolError("max_rows must be an integer") + try: check(sql, self._policy) except PolicyError as exc: raise ToolError(str(exc)) from exc + except Exception as exc: + # A parse failure the guard itself did not convert (e.g. an + # exception type it does not anticipate) must still never reach + # the caller as a raw traceback. Deliberately no str(exc) here: + # an unexpected exception's text is exactly what might carry a + # path or internal detail; `from exc` keeps it for a developer + # without surfacing it to the model. + raise ToolError("could not check this SQL against policy; rejecting") from exc limit = self._effective_max_rows(max_rows) try: diff --git a/tests/test_guard.py b/tests/test_guard.py index 705ea90..73fa613 100644 --- a/tests/test_guard.py +++ b/tests/test_guard.py @@ -176,6 +176,14 @@ class TestInjectionAttempts: with pytest.raises(PolicyError, match="parse"): check("SELECT * FROM t WHERE note = 'it''s", CLOSED) + def test_recursion_error_during_parsing_is_a_policy_error_not_a_crash(self): + # Deeply nested parentheses can blow the recursive-descent parser's + # stack with a RecursionError, which is not a SqlglotError. A parse + # failure of any kind must be rejected, never propagate raw. + deeply_nested = "SELECT " + "(" * 400 + "1" + ")" * 400 + with pytest.raises(PolicyError, match="parse"): + check(deeply_nested, CLOSED) + class TestMatchesPrefix: def test_exact_match(self): diff --git a/tests/test_server.py b/tests/test_server.py index 4637d1c..ca2dd89 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -92,6 +92,29 @@ class TestRunQuery: with pytest.raises(ToolError, match="no such table"): make_tools(client).run_query("SELECT * FROM nope") + def test_pathological_sql_that_crashes_the_parser_is_rejected_without_a_traceback(self): + client = MagicMock() + deeply_nested = "SELECT " + "(" * 400 + "1" + ")" * 400 + with pytest.raises(ToolError) as excinfo: + make_tools(client).run_query(deeply_nested) + client.query.assert_not_called() + message = str(excinfo.value) + assert "/" not in message + assert "\\" not in message + assert ".py" not in message + + def test_non_string_sql_is_rejected_as_a_tool_error(self): + client = MagicMock() + with pytest.raises(ToolError, match="sql must be a string"): + make_tools(client).run_query(5) + client.query.assert_not_called() + + def test_non_integer_max_rows_is_rejected_as_a_tool_error(self): + client = MagicMock() + with pytest.raises(ToolError, match="max_rows must be an integer"): + make_tools(client).run_query("SELECT 1", max_rows="10") + client.query.assert_not_called() + class TestListSchemas: def test_returns_all_schemas_by_default(self): @@ -127,6 +150,7 @@ class TestListTables: client = MagicMock() client.tables.return_value = [{"name": "t", "type": "TABLE"}] assert make_tools(client).list_tables("dfs.tmp") == [{"name": "t", "type": "TABLE"}] + client.tables.assert_called_once_with("dfs.tmp") def test_hidden_schema_is_refused(self): client = MagicMock() @@ -147,6 +171,7 @@ class TestDescribeTable: client = MagicMock() client.columns.return_value = [{"name": "id", "data_type": "INTEGER", "nullable": True}] assert make_tools(client).describe_table("dfs.tmp", "t")[0]["name"] == "id" + client.columns.assert_called_once_with("dfs.tmp", "t") def test_hidden_schema_is_refused(self): client = MagicMock()
