(drill-mcp) 28/44: fix: never let unexpected exceptions escape run_query as raw tracebacks

[cg...@apache.org]

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 <cg...@gmail.com>
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()