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 8f72443e940a348a07341e141dfb4bc7f90f84ce Author: cgivre <[email protected]> AuthorDate: Wed Aug 12 00:26:57 2026 -0400 fix: quote file-plugin table names as a single identifier in DESCRIBE Concatenating schema+table before quote_identifier_path split a dotted filename like sales.csv into separate schema/table segments, producing an invalid DESCRIBE. Quote the table with a new quote_identifier() that keeps a dotted filename inside one backtick pair. --- drill_mcp/client_rest.py | 39 ++++++++++++++++++++++++++++++++------- tests/test_client_rest.py | 43 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 75 insertions(+), 7 deletions(-) diff --git a/drill_mcp/client_rest.py b/drill_mcp/client_rest.py index 6e7142d..43dd341 100644 --- a/drill_mcp/client_rest.py +++ b/drill_mcp/client_rest.py @@ -127,6 +127,28 @@ def quote_identifier_path(value: str) -> str: return ".".join(f"`{part}`" for part in parts) +# `_IDENTIFIER` plus a literal ".", for filenames like "sales.csv" that are ONE +# identifier, not a dotted path -- see `quote_identifier`. Still excludes +# backticks: a backtick in the value would break out of the quoting below, +# which is the entire trust boundary this regex exists to enforce. +_FILE_IDENTIFIER = re.compile(r"[A-Za-z0-9_$.-]+") + + +def quote_identifier(value: str) -> str: + """Quote a single identifier that may itself contain dots, e.g. a filename. + + `sales.csv` is one identifier, not a two-part path: a dotted filename must + stay inside a single backtick pair, or Drill reads the extension as the + table name and the stem as part of the schema (see + sqlalchemy_drill.base.DrillIdentifierPreparer.format_drill_table, which + quotes plugin.`workspace`.`file.ext` the same way). Reject rather than + escape -- same trust boundary as `quote_identifier_path`. + """ + if not _FILE_IDENTIFIER.fullmatch(value): + raise DrillError(f"invalid identifier: {value!r}") + return f"`{value}`" + + _QUERY_ID = re.compile(r"[A-Za-z0-9-]+") @@ -330,9 +352,10 @@ class RestClient: def columns(self, schema: str, table: str) -> list[dict[str, Any]]: # Validate the table name up front, before the plugin_type lookup fires # a query: an invalid table name should never make it to the network. - # File-plugin table names are filenames and may contain a literal "." - # (e.g. "sales.csv"), so validate segment-by-segment like a dotted path. - if any(not _IDENTIFIER.fullmatch(part) for part in table.split(".")): + # `_FILE_IDENTIFIER` (not `_IDENTIFIER`) because file-plugin table + # names are filenames and may contain a literal "." (e.g. "sales.csv") + # as ONE identifier -- see `quote_identifier`. + if not _FILE_IDENTIFIER.fullmatch(table): raise DrillError(f"invalid identifier: {table!r}") # Same split: file plugins have dynamic schemas and no @@ -340,10 +363,12 @@ class RestClient: # deliberately NOT a `SELECT * ... LIMIT 1` probe, which would read user # data to answer a metadata question. if self.plugin_type(schema) == "file": - result = self.query( - f"DESCRIBE {quote_identifier_path(schema + '.' + table)}", - max_rows=10_000, - ) + # `table` is ONE identifier (a filename), not a further dotted + # path -- quote it with `quote_identifier`, not + # `quote_identifier_path`, or "sales.csv" would be split into a + # schema segment "sales" and a table segment "csv". + target = f"{quote_identifier_path(schema)}.{quote_identifier(table)}" + result = self.query(f"DESCRIBE {target}", max_rows=10_000) return [ { "name": row.get("COLUMN_NAME"), diff --git a/tests/test_client_rest.py b/tests/test_client_rest.py index fb804d0..3d32e92 100644 --- a/tests/test_client_rest.py +++ b/tests/test_client_rest.py @@ -486,10 +486,53 @@ class TestFilePluginMetadata: ] body = route.calls[1].request.read() assert b"DESCRIBE" in body + # The filename is ONE identifier, not a further dotted path: it must + # stay inside a single backtick pair, or Drill reads the extension as + # the table name and the stem as part of the schema. + assert b"`sales.csv`" in body + assert b"`sales`.`csv`" not in body # Metadata-only: never read user rows to answer a metadata question. assert b"LIMIT 1" not in body assert b"SELECT *" not in body + @respx.mock + def test_columns_keeps_a_multi_dot_filename_in_one_backtick_pair(self): + route = respx.post(f"{BASE}/query.json").mock( + side_effect=[ + self._schemata("file"), + query_response( + ["COLUMN_NAME", "DATA_TYPE", "IS_NULLABLE"], + [{"COLUMN_NAME": "id", "DATA_TYPE": "BIGINT", "IS_NULLABLE": "YES"}], + ), + ] + ) + make_client().columns("dfs.tmp", "archive.2024.json") + body = route.calls[1].request.read() + assert b"`archive.2024.json`" in body + + @respx.mock + def test_columns_works_for_a_file_with_no_extension(self): + route = respx.post(f"{BASE}/query.json").mock( + side_effect=[ + self._schemata("file"), + query_response( + ["COLUMN_NAME", "DATA_TYPE", "IS_NULLABLE"], + [{"COLUMN_NAME": "id", "DATA_TYPE": "BIGINT", "IS_NULLABLE": "YES"}], + ), + ] + ) + assert make_client().columns("dfs.tmp", "README") == [ + {"name": "id", "data_type": "BIGINT", "nullable": True} + ] + assert b"`README`" in route.calls[1].request.read() + + @respx.mock + def test_columns_rejects_a_backtick_in_the_table_name(self): + route = respx.post(f"{BASE}/query.json").mock(return_value=self._schemata("file")) + with pytest.raises(DrillError, match="invalid identifier"): + make_client().columns("dfs.tmp", "sales`; DROP TABLE x --.csv") + assert not route.called + @respx.mock def test_columns_uses_information_schema_for_a_non_file_plugin(self): route = respx.post(f"{BASE}/query.json").mock(
