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 b4671dc76d025c5b800c1a908d3d36e1acca4f89 Author: cgivre <[email protected]> AuthorDate: Wed Aug 12 13:50:37 2026 -0400 fix: close probe error/repr leaks and drop the unused view lookup - JDBC metadata now reads a DBAPITypeObject's .values instead of stringifying it, matching sqlalchemy-drill's get_columns. - The dynamic-schema probe no longer surfaces Drill's raw error text on failure (it can embed sampled cell content); it also now raises instead of silently returning [] when a probe finds zero rows. - Removed fetch_view_names: the view/file quoting split it existed to support was already proven unnecessary, so it had no caller. - QueryResult.metadata is annotated list[str | None]. --- drill_mcp/client_jdbc.py | 32 ++++++++++---- drill_mcp/client_rest.py | 104 ++++++++++++++++++++++++---------------------- tests/test_client_jdbc.py | 45 +++++++++++++++++++- tests/test_client_rest.py | 99 ++++++++++++++++++++++++++----------------- 4 files changed, 181 insertions(+), 99 deletions(-) diff --git a/drill_mcp/client_jdbc.py b/drill_mcp/client_jdbc.py index ca550b7..58ce1c3 100644 --- a/drill_mcp/client_jdbc.py +++ b/drill_mcp/client_jdbc.py @@ -45,6 +45,28 @@ from .config import Config DRIVER_CLASS = "org.apache.drill.jdbc.Driver" +def _type_name(description_entry: tuple) -> str | None: + """Extract a per-column type name from one `cursor.description` entry. + + `description_entry[1]` is the DB-API `type_code`. jaydebeapi's is + typically a `DBAPITypeObject`-like value carrying a `.values` tuple of + type name strings, not a bare string -- stringifying it directly (`str + (type_code)`) yields an object repr (`<...DBAPITypeObject object at + 0x...>`), not a usable type name. Mirrors sqlalchemy-drill's + `get_columns` (base.py:433-438), which checks for `.values` the same + way before falling back to `str()`. + """ + if len(description_entry) <= 1: + return None + type_code = description_entry[1] + if type_code is None: + return None + values = getattr(type_code, "values", None) + if values: + return str(values[0]) + return str(type_code) + + class JdbcClient: def __init__(self, config: Config) -> None: self._config = config @@ -134,15 +156,7 @@ class JdbcClient: rows = cursor.fetchmany(max_rows) description = cursor.description or [] columns = [entry[0] for entry in description] - # entry[1] is the DB-API type_code. jaydebeapi's is not a - # bare string, so it is stringified the same way the REST - # path's `metadata` array is consumed (client_rest.py's - # QueryResult.metadata docstring): a per-column type name, - # `None` when unknown, never an error. - metadata = [ - str(entry[1]) if len(entry) > 1 and entry[1] is not None else None - for entry in description - ] + metadata = [_type_name(entry) for entry in description] except Exception as exc: raise DrillError(self._scrub(str(exc))) from exc return QueryResult( diff --git a/drill_mcp/client_rest.py b/drill_mcp/client_rest.py index 91245a8..bffb88a 100644 --- a/drill_mcp/client_rest.py +++ b/drill_mcp/client_rest.py @@ -97,7 +97,9 @@ class QueryResult: # Drill's REST API returns this in a `metadata` array (Drill >= 1.19); # older Drill omits it. Absent metadata is not an error -- callers that # need types (e.g. `_probe_columns`) must tolerate an empty list here. - metadata: list[str] = field(default_factory=list) + # `None` entries are deliberate (absent metadata, or padding for a + # shorter-than-`columns` array), not just an artifact of the default. + metadata: list[str | None] = field(default_factory=list) def quote_literal(value: str) -> str: @@ -308,31 +310,6 @@ def fetch_tables(query: Query, schema: str) -> list[dict[str, Any]]: ] -def fetch_view_names(query: Query, schema: str) -> list[str]: - """Return the names of views registered directly under `schema`. - - sqlalchemy-drill's dialect uses this to pick between two different probe - SQL shapes in `get_columns` (base.py:423-428), because its own - `format_drill_table` quoting mishandles a view name. `_probe_columns` - below has no such failure mode (see its comment), so it does not need - this to build a query -- but the lookup itself is still a plugin-neutral - piece of Drill metadata worth exposing directly, mirroring the dialect's - `get_view_names` (base.py:362-372): a query failure (e.g. a Drill - version without the `VIEWS` INFORMATION_SCHEMA relation) is tolerated - and yields an empty list rather than raising. An invalid `schema` still - raises -- only the query itself is allowed to fail silently. - """ - literal_schema = quote_literal_path(schema) # raises on an invalid identifier - try: - result = query( - f"SELECT `TABLE_NAME` FROM INFORMATION_SCHEMA.`VIEWS` WHERE TABLE_SCHEMA = {literal_schema}", - 10_000, - ) - except DrillError: - return [] - return [row["TABLE_NAME"] for row in result.rows if row.get("TABLE_NAME")] - - def _probe_target(schema: str, table: str) -> str: """Quote `schema`.`table` the same way `_describe_columns` does. @@ -341,20 +318,25 @@ def _probe_target(schema: str, table: str) -> str: `quote_identifier` because file-plugin table names are filenames that may themselves contain a "." (e.g. "sales.csv") -- see `quote_identifier`. - sqlalchemy-drill's dialect instead special-cases this by counting dots - in `schema + "." + table` to decide where the plugin/workspace/filename - boundaries fall (`format_drill_table`, base.py:164-193), and for a view - wraps the WHOLE schema string in a single backtick pair instead - (base.py:425). Both produce valid SQL, but neither is needed here: this - codebase's identifier helpers already quote every schema segment and the - table/filename correctly and uniformly, for both a view and a plain - file, so the same quoting is reused for both `_probe_columns` branches - rather than replicating the dialect's ad hoc dot-counting. + sqlalchemy-drill's dialect instead special-cases a view's target by + counting dots in `schema + "." + table` to decide where the + plugin/workspace/filename boundaries fall (`format_drill_table`, + base.py:164-193): that heuristic assumes the trailing dotted token is a + file extension, so for `schema="dfs"`, `table="a.b"` it emits + ``dfs.a.`b` `` -- splitting a view name that happens to contain a dot in + two. The dialect's separate view branch in `get_columns` + (base.py:423-428) exists only to dodge that bug by wrapping the whole + schema string in a single backtick pair instead. `_probe_target` has no + such failure mode -- it quotes every schema segment and the + table/filename correctly and uniformly regardless of whether `table` + names a view or a file -- so there is no input on which the two + approaches disagree, and no second quoting scheme or view lookup is + needed here. """ return f"{quote_identifier_path(schema)}.{quote_identifier(table)}" -def _columns_from_metadata(columns: list[str], metadata: list[str]) -> list[dict[str, Any]]: +def _columns_from_metadata(columns: list[str], metadata: list[str | None]) -> list[dict[str, Any]]: """Build `describe_table` rows from a probe's `columns`/`metadata`. Never reads `result.rows` -- the caller passes only `columns` and @@ -387,28 +369,50 @@ def _probe_columns(query: Query, schema: str, table: str, plugin_type: str) -> l is discarded unread. That is what makes the probe acceptable here: the caller (`describe_table`) gets column names and types, never sampled values. + + That same privacy property is why a probe FAILURE is handled specially + below, not just left to propagate: Drill's own error text for a query + that fails while reading a row (a type-coercion or malformed-record + error, for instance) can embed the offending cell's content -- + `_error_text` passes `errorMessage` through verbatim, and `server.py` + surfaces `DrillError`'s text to the caller unchanged. `DESCRIBE` could + never trigger this path; only the probe can, so only the probe needs to + guard against it. """ if plugin_type == "mongo": - # Collections carry no dots, so the combined schema.table path is - # quoted segment-wise like any other dotted path (base.py:420-422). + # MongoDB collection names CAN contain dots (e.g. "logs.2024"); this + # is not special-cased, so such a name is quoted the same as any + # other dotted path -- one backtick pair per "." segment, exactly + # like a schema path. sqlalchemy-drill's dialect does the same + # (base.py:420-422), so this is not a regression, just a limitation + # shared with the reference: a dotted collection name resolves to a + # nested path rather than one opaque identifier. target = quote_identifier_path(f"{schema}.{table}") sql = f"SELECT `**` FROM {target} LIMIT 1" else: - # sqlalchemy-drill's dialect branches here on `table in views` - # (base.py:423-428) because its OWN quoting -- `format_drill_table` - # counting dots to split plugin/workspace/filename -- mishandles a - # view name that doesn't fit that 2-or-3-dot shape, so it falls back - # to wrapping the whole schema in one backtick pair for views - # instead. `_probe_target` doesn't have that failure mode: it quotes - # every schema segment and the table/filename correctly and - # uniformly regardless of whether `table` names a view or a file, so - # there is no second quoting scheme to fall back to here. A table - # that happens to be a registered view is still queried by - # `_probe_target`, unchanged. target = _probe_target(schema, table) sql = f"SELECT * FROM {target} LIMIT 1" - result = query(sql, 1) + try: + result = query(sql, 1) + except DrillError as exc: + # Deliberately does NOT include `exc`'s text in the new message -- + # only chains it as the cause -- because that text may be Drill's + # own error message, which can embed sampled cell content. + raise DrillError( + f"could not determine columns for `{schema}`.`{table}`: the probe query failed" + ) from exc + + if not result.columns: + # A dynamic-schema plugin discovers columns only by reading data; + # zero rows means Drill never had anything to infer a schema from. + # Returning [] here would read exactly like the "no columns" failure + # mode Step 2 rejects for HTTP plugins -- fail loudly instead. + raise DrillError( + f"columns could not be determined for `{schema}`.`{table}` because " + "the probe returned no rows; the table may be empty." + ) + return _columns_from_metadata(result.columns, result.metadata) diff --git a/tests/test_client_jdbc.py b/tests/test_client_jdbc.py index 202d1a6..02ae59f 100644 --- a/tests/test_client_jdbc.py +++ b/tests/test_client_jdbc.py @@ -27,6 +27,19 @@ from drill_mcp.client_rest import DrillError from drill_mcp.config import load_config +class _FakeDBAPITypeObject: + """Mimics jaydebeapi's real `type_code`: an object with a `.values` + tuple of type name strings, NOT a bare string. Deliberately has no + `__str__`/`__repr__` override, so `str(instance)` falls back to the + default object repr (`<...>`) -- a regression that stringifies this + directly instead of reading `.values` would produce that repr as the + reported type, which is exactly the failure this class exists to catch. + """ + + def __init__(self, *values): + self.values = values + + @pytest.fixture def fake_jaydebeapi(monkeypatch): module = MagicMock() @@ -81,13 +94,31 @@ def test_query_returns_columns_and_rows(fake_jaydebeapi): def test_query_populates_metadata_from_cursor_description(fake_jaydebeapi): + # Shaped like jaydebeapi's actual output: type_code is a + # DBAPITypeObject-like value with `.values`, not a bare string. A + # fabricated string type_code (e.g. `("id", "INTEGER")`) would pass even + # if the implementation just did `str(type_code)` -- this fixture would + # not, since `str()` on `_FakeDBAPITypeObject` yields an object repr. cursor = fake_jaydebeapi.connect.return_value.cursor.return_value - cursor.description = [("id", "INTEGER"), ("name", "VARCHAR")] + cursor.description = [ + ("id", _FakeDBAPITypeObject("INTEGER")), + ("name", _FakeDBAPITypeObject("VARCHAR")), + ] cursor.fetchmany.return_value = [(1, "x")] result = make_client().query("SELECT 1", max_rows=10) assert result.metadata == ["INTEGER", "VARCHAR"] +def test_query_metadata_type_code_without_values_falls_back_to_str(fake_jaydebeapi): + # A type_code that is already a bare string (no `.values` attribute) + # must still work -- some DB-API drivers use plain strings. + cursor = fake_jaydebeapi.connect.return_value.cursor.return_value + cursor.description = [("id", "INTEGER")] + cursor.fetchmany.return_value = [(1,)] + result = make_client().query("SELECT 1", max_rows=10) + assert result.metadata == ["INTEGER"] + + def test_query_metadata_is_none_per_column_when_type_code_is_absent(fake_jaydebeapi): # The default fixture's description entries carry a `None` type_code; # this must not raise, and must not fabricate a type. @@ -95,6 +126,18 @@ def test_query_metadata_is_none_per_column_when_type_code_is_absent(fake_jaydebe assert result.metadata == [None, None] +def test_query_metadata_never_leaks_an_object_repr(fake_jaydebeapi): + # Regression test: stringifying a DBAPITypeObject-like type_code + # directly (instead of reading `.values`) produces a Python object repr + # like "<...DBAPITypeObject object at 0x...>" -- assert that never + # appears in the reported metadata. + cursor = fake_jaydebeapi.connect.return_value.cursor.return_value + cursor.description = [("id", _FakeDBAPITypeObject("BIGINT"))] + cursor.fetchmany.return_value = [(1,)] + result = make_client().query("SELECT 1", max_rows=10) + assert all("object at 0x" not in str(m) for m in result.metadata) + + def test_query_respects_max_rows(fake_jaydebeapi): cursor = fake_jaydebeapi.connect.return_value.cursor.return_value cursor.fetchmany.return_value = [(1, "x"), (2, "y")] diff --git a/tests/test_client_rest.py b/tests/test_client_rest.py index 5325723..59d894e 100644 --- a/tests/test_client_rest.py +++ b/tests/test_client_rest.py @@ -618,6 +618,58 @@ class TestFilePluginMetadata: result = make_client().columns("dfs.tmp", "people.csv") assert "078-05-1120-SENTINEL" not in repr(result) + @respx.mock + def test_columns_probe_failure_does_not_leak_drills_error_text(self): + # A probe reads a row -- if Drill fails WHILE reading it (a + # type-coercion or malformed-record error), Drill's own error text + # can embed the offending cell's content. DESCRIBE could never + # trigger this; only the probe can, so the probe's failure path must + # not surface Drill's raw error message. + respx.post(f"{BASE}/query.json").mock( + side_effect=[ + self._schemata("file"), + httpx.Response( + 500, json={"errorMessage": "conversion failed on value SENTINEL-CELL-9182"} + ), + ] + ) + with pytest.raises(DrillError) as exc_info: + make_client().columns("dfs.tmp", "sales.csv") + message = str(exc_info.value) + assert "SENTINEL-CELL-9182" not in message + assert "dfs.tmp" in message + assert "sales.csv" in message + # The original Drill error is preserved as the exception chain, just + # not folded into the new message text. + assert "SENTINEL-CELL-9182" in str(exc_info.value.__cause__) + + @respx.mock + def test_columns_probe_raises_when_the_table_is_empty(self): + # A dynamic-schema plugin discovers columns only by reading data; a + # probe that returns zero rows (and so no columns) means Drill never + # had anything to infer a schema from. Returning [] would read + # exactly like the "no columns" failure mode Step 2 rejects for HTTP + # plugins. + respx.post(f"{BASE}/query.json").mock( + side_effect=[ + self._schemata("file"), + query_response([], []), + ] + ) + with pytest.raises(DrillError, match="no rows"): + make_client().columns("dfs.tmp", "empty.csv") + + @respx.mock + def test_columns_probe_rejects_injection_in_schema_before_any_query_fires(self): + # A malicious schema must never reach the wire through the probe's + # `_probe_target`/mongo interpolation sites -- `fetch_plugin_type`'s + # own validation rejects it first, before any query (including the + # SCHEMATA lookup) fires. + route = respx.post(f"{BASE}/query.json").mock(return_value=self._schemata("file")) + with pytest.raises(DrillError, match="invalid identifier"): + make_client().columns("dfs'; DROP TABLE x --", "sales.csv") + assert not route.called + @respx.mock def test_columns_probes_a_mongo_plugin_with_double_star(self): route = respx.post(f"{BASE}/query.json").mock( @@ -633,8 +685,9 @@ class TestFilePluginMetadata: body = route.calls[1].request.read() assert b"SELECT `**` FROM" in body assert b"LIMIT 1" in body - # Collections carry no dots, so the combined path is quoted - # segment-wise like any other dotted schema path. + # A dotted collection name is quoted segment-wise like any other + # dotted path (see `_probe_columns`'s mongo comment) -- not a special + # case, just the same `quote_identifier_path` treatment as a schema. assert b"`dfs`.`tmp`.`mycollection`" in body @respx.mock @@ -652,14 +705,11 @@ class TestFilePluginMetadata: assert b"SELECT * FROM" in body @respx.mock - def test_columns_probe_handles_a_table_that_is_a_registered_view(self): - # sqlalchemy-drill's dialect special-cases a view name here because - # its OWN quoting scheme (dot-counting to split plugin/workspace/ - # filename) mishandles a name that doesn't fit that shape. This - # module's `_probe_target` quotes every schema segment and the - # table/filename uniformly, with no such failure mode, so a table - # that happens to be a view needs no special handling: the same - # `SELECT * FROM ... LIMIT 1` probe answers correctly either way. + def test_columns_probe_quotes_a_view_name_the_same_way_as_a_file(self): + # This does NOT register a view via INFORMATION_SCHEMA.VIEWS -- there + # is no such lookup any more (see `_probe_target`'s docstring for + # why). It only confirms `_probe_target` produces the same quoting + # for a table name shaped like a view as for an ordinary file name. route = respx.post(f"{BASE}/query.json").mock( side_effect=[ self._schemata("file"), @@ -744,35 +794,6 @@ class TestFilePluginMetadata: assert "run" in message.lower() or "query" in message.lower() assert "LIMIT" in message - @respx.mock - def test_view_names_returns_the_view_names_for_a_schema(self): - route = respx.post(f"{BASE}/query.json").mock( - return_value=query_response(["TABLE_NAME"], [{"TABLE_NAME": "top_sales"}]) - ) - client = make_client() - from drill_mcp.client_rest import fetch_view_names - - assert fetch_view_names(client.query, "dfs.tmp") == ["top_sales"] - assert b"VIEWS" in route.calls.last.request.read() - - @respx.mock - def test_view_names_returns_empty_list_when_the_query_fails(self): - # Matches the dialect's tolerance (base.py:362-372): a missing VIEWS - # relation (e.g. an older Drill) must not break column lookup. - respx.post(f"{BASE}/query.json").mock( - return_value=httpx.Response(500, json={"errorMessage": "no such relation"}) - ) - from drill_mcp.client_rest import fetch_view_names - - assert fetch_view_names(make_client().query, "dfs.tmp") == [] - - @respx.mock - def test_view_names_rejects_injection_in_schema_name(self): - from drill_mcp.client_rest import fetch_view_names - - with pytest.raises(DrillError, match="invalid identifier"): - fetch_view_names(make_client().query, "dfs'; DROP TABLE x --") - @respx.mock def test_unknown_plugin_type_falls_back_to_information_schema(self): respx.post(f"{BASE}/query.json").mock(
