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 6fd9f1f1b2589ed170b0e631be76784c8c6a2775 Author: cgivre <[email protected]> AuthorDate: Wed Aug 12 13:14:08 2026 -0400 feat: follow the dialect methodology for dynamic-schema plugins fetch_columns now branches on plugin TYPE: file/mongo/splunk probe with SELECT ... LIMIT 1 (schema discovered at read time, DESCRIBE cannot answer), http raises an explanatory error, everything else uses DESCRIBE. QueryResult gains metadata (per-column type strings) on both REST and JDBC backends. fetch_plugin_type resolves a bare plugin name by matching in Python instead of LIKE, avoiding a wildcard-injection surface. --- drill_mcp/client_jdbc.py | 13 ++- drill_mcp/client_rest.py | 220 +++++++++++++++++++++++++++++++++++-------- tests/test_client_jdbc.py | 15 +++ tests/test_client_rest.py | 231 +++++++++++++++++++++++++++++++++++++++++----- 4 files changed, 417 insertions(+), 62 deletions(-) diff --git a/drill_mcp/client_jdbc.py b/drill_mcp/client_jdbc.py index 2b103c6..ca550b7 100644 --- a/drill_mcp/client_jdbc.py +++ b/drill_mcp/client_jdbc.py @@ -132,13 +132,24 @@ class JdbcClient: with closing(connection.cursor()) as cursor: cursor.execute(sql) rows = cursor.fetchmany(max_rows) - columns = [description[0] for description in cursor.description or []] + 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 + ] except Exception as exc: raise DrillError(self._scrub(str(exc))) from exc return QueryResult( columns=columns, rows=[dict(zip(columns, row)) for row in rows], truncated=max_rows > 0 and len(rows) >= max_rows, + metadata=metadata, ) # -- metadata -------------------------------------------------------------- diff --git a/drill_mcp/client_rest.py b/drill_mcp/client_rest.py index b890f13..91245a8 100644 --- a/drill_mcp/client_rest.py +++ b/drill_mcp/client_rest.py @@ -93,6 +93,11 @@ class QueryResult: rows: list[dict[str, Any]] = field(default_factory=list) query_id: str | None = None truncated: bool = False + # Per-column type strings aligned with `columns`, e.g. "VARCHAR(10)". + # 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) def quote_literal(value: str) -> str: @@ -216,19 +221,48 @@ def _error_text(response: httpx.Response) -> str: Query = Callable[[str, int], QueryResult] +# Plugin types whose schema is discovered at read time rather than registered +# in INFORMATION_SCHEMA. `DESCRIBE` cannot answer for these -- there is +# nothing durable to describe -- so `fetch_columns` must probe with a +# `SELECT ... LIMIT 1` instead. Matches sqlalchemy-drill's `get_columns` +# (base.py:405-470), which is the methodology this module follows rather +# than inventing its own. +DYNAMIC_SCHEMA_TYPES = ("file", "mongo", "splunk") + +# Strips size/precision info from a Drill type string, e.g. "VARCHAR(10)" -> +# "VARCHAR", "DECIMAL(10, 2)" -> "DECIMAL". Same approach drilldbapi's +# `Cursor.execute` uses on the `metadata` array (sad.py:278). +_TYPE_PRECISION = re.compile(r"\(.*\)") + + def fetch_plugin_type(query: Query, schema: str) -> str | None: """Return the storage plugin TYPE backing `schema`, or None if unknown. File-based plugins (`dfs`, `s3`) do not register their contents in INFORMATION_SCHEMA, so `fetch_tables` and `fetch_columns` must branch on this. + + A *bare* plugin name (`dfs`) has no exact SCHEMATA row of its own -- + only its workspaces do (`dfs.tmp`, `dfs.root`). SCHEMATA is fetched once, + unfiltered, and matched in Python: an exact match wins; otherwise the + first row whose name's leading dotted component equals `schema` is used. + This is deliberately NOT a `WHERE SCHEMA_NAME LIKE '%schema%'` clause + (which is what sqlalchemy-drill's `get_plugin_type` does, base.py:474) -- + `schema` is model-supplied, and `%`/`_` are wildcards in a LIKE pattern. """ + quote_literal_path(schema) # validate the identifier before any query fires result = query( - "SELECT SCHEMA_NAME, TYPE FROM INFORMATION_SCHEMA.`SCHEMATA` " - f"WHERE SCHEMA_NAME = {quote_literal_path(schema)}", - 1, + "SELECT SCHEMA_NAME, TYPE FROM INFORMATION_SCHEMA.`SCHEMATA`", + 10_000, ) - return result.rows[0].get("TYPE") if result.rows else None + prefix_match: str | None = None + for row in result.rows: + name = row.get("SCHEMA_NAME") + if name == schema: + return row.get("TYPE") + if prefix_match is None and name and name.split(".", 1)[0] == schema: + prefix_match = row.get("TYPE") + return prefix_match def fetch_schemas(query: Query) -> list[dict[str, Any]]: @@ -274,41 +308,120 @@ def fetch_tables(query: Query, schema: str) -> list[dict[str, Any]]: ] -def fetch_columns(query: Query, 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_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 _is_valid_file_identifier(table): - raise DrillError(f"invalid identifier: {table!r}") +def fetch_view_names(query: Query, schema: str) -> list[str]: + """Return the names of views registered directly under `schema`. - # Same split: file plugins have dynamic schemas and no - # INFORMATION_SCHEMA.`COLUMNS` rows. DESCRIBE is metadata-only -- - # deliberately NOT a `SELECT * ... LIMIT 1` probe, which would read user - # data to answer a metadata question. - if fetch_plugin_type(query, schema) == "file": - # `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 = query(f"DESCRIBE {target}", 10_000) - return [ - { - "name": row.get("COLUMN_NAME"), - "data_type": row.get("DATA_TYPE"), - "nullable": str(row.get("IS_NULLABLE", "")).upper() == "YES", - } - for row in result.rows - ] + 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. + + `schema` is a dotted path (each segment quoted individually via + `quote_identifier_path`); `table` stays inside ONE backtick pair via + `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. + """ + return f"{quote_identifier_path(schema)}.{quote_identifier(table)}" - result = query( - "SELECT COLUMN_NAME, DATA_TYPE, IS_NULLABLE FROM INFORMATION_SCHEMA.`COLUMNS` " - f"WHERE TABLE_SCHEMA = {quote_literal_path(schema)} " - f"AND TABLE_NAME = {quote_literal(table)} ORDER BY ORDINAL_POSITION", - 10_000, - ) + +def _columns_from_metadata(columns: list[str], metadata: list[str]) -> list[dict[str, Any]]: + """Build `describe_table` rows from a probe's `columns`/`metadata`. + + Never reads `result.rows` -- the caller passes only `columns` and + `metadata`, so a sampled row value cannot reach this function let alone + its return value. Nullability is unknowable from a single probed row + (a `NULL` here says nothing about whether the column CAN be null, and a + non-`NULL` value says nothing about whether it must); reporting `None` + is honest, guessing `True` or `False` is not. + """ + # `zip` silently truncates to the shorter list; pad rather than let a + # metadata array that is present but short (a malformed or unexpected + # response) drop trailing columns. + types = list(metadata) + [None] * max(0, len(columns) - len(metadata)) + result = [] + for name, type_str in zip(columns, types): + data_type = _TYPE_PRECISION.sub("", type_str) if type_str else None + result.append({"name": name, "data_type": data_type, "nullable": None}) + return result + + +def _probe_columns(query: Query, schema: str, table: str, plugin_type: str) -> list[dict[str, Any]]: + """Discover columns for a dynamic-schema plugin by probing one row. + + `DESCRIBE` cannot answer for `file`/`mongo`/`splunk`: their schema is + discovered at read time, not registered anywhere `DESCRIBE` can consult. + Mirrors sqlalchemy-drill's `get_columns` (base.py:405-451). + + Privacy: this reads one row from the underlying data, but only + `result.columns` and `result.metadata` are used below -- `result.rows` + is discarded unread. That is what makes the probe acceptable here: the + caller (`describe_table`) gets column names and types, never sampled + values. + """ + 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). + 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) + return _columns_from_metadata(result.columns, result.metadata) + + +def _describe_columns(query: Query, schema: str, table: str) -> list[dict[str, Any]]: + """Discover columns via `DESCRIBE`, for any plugin with a registered schema. + + Mirrors sqlalchemy-drill's `get_columns` `else` branch (base.py:453-472): + `DESCRIBE` is metadata-only and never reads user data, so it is + preferred whenever it can answer -- i.e. for anything NOT in + `DYNAMIC_SCHEMA_TYPES`. + """ + target = _probe_target(schema, table) + result = query(f"DESCRIBE {target}", 10_000) return [ { "name": row.get("COLUMN_NAME"), @@ -319,6 +432,38 @@ def fetch_columns(query: Query, schema: str, table: str) -> list[dict[str, Any]] ] +def fetch_columns(query: Query, 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_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 _is_valid_file_identifier(table): + raise DrillError(f"invalid identifier: {table!r}") + + plugin_type = fetch_plugin_type(query, schema) + + # An HTTP plugin has no column metadata until a query has actually been + # run against the endpoint -- there is no schema to DESCRIBE and no + # table to probe. Returning [] here would read as "this table has no + # columns"; fail loudly and explain instead. `schema`/`table` are + # already known to be valid identifiers at this point (checked above and + # by `fetch_plugin_type`'s `quote_literal_path` call), so they are safe + # to embed directly in the message. + if plugin_type == "http": + raise DrillError( + f"Drill cannot report columns for the HTTP plugin schema '{schema}' until a query " + "has been run against it. Run a query such as\n" + f" SELECT * FROM `{schema}`.`{table}` LIMIT 10\n" + "and read the column names from the result." + ) + + if plugin_type in DYNAMIC_SCHEMA_TYPES: + return _probe_columns(query, schema, table, plugin_type) + + return _describe_columns(query, schema, table) + + # -- client -------------------------------------------------------------- @@ -428,6 +573,7 @@ class RestClient: rows=rows, query_id=payload.get("queryId"), truncated=max_rows > 0 and len(rows) >= max_rows, + metadata=payload.get("metadata") or [], ) # -- metadata ---------------------------------------------------------- diff --git a/tests/test_client_jdbc.py b/tests/test_client_jdbc.py index 8e64500..202d1a6 100644 --- a/tests/test_client_jdbc.py +++ b/tests/test_client_jdbc.py @@ -80,6 +80,21 @@ def test_query_returns_columns_and_rows(fake_jaydebeapi): assert result.truncated is False +def test_query_populates_metadata_from_cursor_description(fake_jaydebeapi): + cursor = fake_jaydebeapi.connect.return_value.cursor.return_value + cursor.description = [("id", "INTEGER"), ("name", "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_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. + result = make_client().query("SELECT 1", max_rows=10) + assert result.metadata == [None, None] + + 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 fcf2b4f..5325723 100644 --- a/tests/test_client_rest.py +++ b/tests/test_client_rest.py @@ -401,8 +401,11 @@ class TestKerberosAuth: assert client._http.auth is sentinel -def query_response(columns, rows): - return httpx.Response(200, json={"columns": columns, "rows": rows, "queryId": "q"}) +def query_response(columns, rows, metadata=None): + payload = {"columns": columns, "rows": rows, "queryId": "q"} + if metadata is not None: + payload["metadata"] = metadata + return httpx.Response(200, json=payload) class TestMetadata: @@ -451,7 +454,53 @@ class TestMetadata: ) ) assert make_client().plugin_type("dfs.tmp") == "file" - assert b"'dfs.tmp'" in route.calls.last.request.read() + # SCHEMATA is fetched unfiltered (not `WHERE SCHEMA_NAME = ...`) and + # matched in Python -- a bare plugin name has no exact SCHEMATA row, + # only its workspaces do. See test_plugin_type_resolves_a_bare_plugin_name. + body = route.calls.last.request.read() + assert b"SCHEMATA" in body + assert b"WHERE" not in body + + @respx.mock + def test_plugin_type_resolves_a_bare_plugin_name(self): + # `WHERE SCHEMA_NAME = 'dfs'` finds nothing when only `dfs.tmp` and + # `dfs.root` exist as SCHEMATA rows. + respx.post(f"{BASE}/query.json").mock( + return_value=query_response( + ["SCHEMA_NAME", "TYPE"], + [ + {"SCHEMA_NAME": "dfs.tmp", "TYPE": "file"}, + {"SCHEMA_NAME": "dfs.root", "TYPE": "file"}, + ], + ) + ) + assert make_client().plugin_type("dfs") == "file" + + @respx.mock + def test_plugin_type_prefers_an_exact_match_over_a_prefix_match(self): + respx.post(f"{BASE}/query.json").mock( + return_value=query_response( + ["SCHEMA_NAME", "TYPE"], + [ + {"SCHEMA_NAME": "dfs.tmp", "TYPE": "file"}, + {"SCHEMA_NAME": "dfs", "TYPE": "exact"}, + ], + ) + ) + assert make_client().plugin_type("dfs") == "exact" + + @respx.mock + def test_plugin_type_bare_name_resolution_rejects_injection(self): + # The Python-side prefix match never interpolates `schema` into SQL, + # but the identifier is still validated up front -- fail fast, no + # network call, and no chance of the malicious string leaking into + # a later query built from the resolved plugin type. + route = respx.post(f"{BASE}/query.json").mock( + return_value=query_response(["SCHEMA_NAME", "TYPE"], []) + ) + with pytest.raises(DrillError, match="invalid identifier"): + make_client().plugin_type("dfs' OR '1'='1") + assert not route.called class TestFilePluginMetadata: @@ -504,41 +553,134 @@ class TestFilePluginMetadata: assert b"INFORMATION_SCHEMA" in route.calls[1].request.read() @respx.mock - def test_columns_uses_describe_for_a_file_plugin(self): + def test_columns_probes_a_file_plugin_instead_of_describe(self): + # DESCRIBE cannot answer for a file plugin: its schema is discovered + # at read time, not registered anywhere DESCRIBE can consult. Follows + # sqlalchemy-drill's get_columns (base.py:405-451): probe with + # SELECT ... LIMIT 1, read `columns`/`metadata`, strip precision. 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"}], - ), + query_response(["id"], [{"id": 12345}], metadata=["BIGINT"]), ] ) assert make_client().columns("dfs.tmp", "sales.csv") == [ - {"name": "id", "data_type": "BIGINT", "nullable": True} + {"name": "id", "data_type": "BIGINT", "nullable": None} ] body = route.calls[1].request.read() - assert b"DESCRIBE" in body + assert b"DESCRIBE" not in body + assert b"SELECT * FROM" in body + assert b"LIMIT 1" 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( + def test_columns_probe_strips_precision_from_the_type_string(self): + respx.post(f"{BASE}/query.json").mock( + side_effect=[ + self._schemata("file"), + query_response(["name"], [{"name": "Alice"}], metadata=["VARCHAR(10)"]), + ] + ) + assert make_client().columns("dfs.tmp", "people.csv") == [ + {"name": "name", "data_type": "VARCHAR", "nullable": None} + ] + + @respx.mock + def test_columns_probe_falls_back_to_none_type_when_metadata_is_absent(self): + # Older Drill (< 1.19) omits the `metadata` array entirely; that is + # not an error. + respx.post(f"{BASE}/query.json").mock( + side_effect=[ + self._schemata("file"), + query_response(["id"], [{"id": 1}]), # no metadata=... + ] + ) + assert make_client().columns("dfs.tmp", "sales.csv") == [ + {"name": "id", "data_type": None, "nullable": None} + ] + + @respx.mock + def test_columns_probe_never_leaks_the_sampled_row_value(self): + # Privacy constraint: the probe reads one row, but describe_table + # must return ONLY column names and types -- never the sampled row. + 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"}], + ["ssn"], [{"ssn": "078-05-1120-SENTINEL"}], metadata=["VARCHAR(11)"] ), ] ) + result = make_client().columns("dfs.tmp", "people.csv") + assert "078-05-1120-SENTINEL" not in repr(result) + + @respx.mock + def test_columns_probes_a_mongo_plugin_with_double_star(self): + route = respx.post(f"{BASE}/query.json").mock( + side_effect=[ + self._schemata("mongo"), + query_response(["id"], [{"id": 1}], metadata=["BIGINT"]), + ] + ) + assert make_client().columns("dfs.tmp", "mycollection") == [ + {"name": "id", "data_type": "BIGINT", "nullable": None} + ] + assert len(route.calls) == 2 + 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. + assert b"`dfs`.`tmp`.`mycollection`" in body + + @respx.mock + def test_columns_probes_a_splunk_plugin(self): + route = respx.post(f"{BASE}/query.json").mock( + side_effect=[ + self._schemata("splunk"), + query_response(["host"], [{"host": "web1"}], metadata=["VARCHAR"]), + ] + ) + assert make_client().columns("dfs.tmp", "main") == [ + {"name": "host", "data_type": "VARCHAR", "nullable": None} + ] + body = route.calls[1].request.read() + 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. + route = respx.post(f"{BASE}/query.json").mock( + side_effect=[ + self._schemata("file"), + query_response(["id"], [{"id": 1}], metadata=["BIGINT"]), + ] + ) + assert make_client().columns("dfs.tmp", "top_sales") == [ + {"name": "id", "data_type": "BIGINT", "nullable": None} + ] + body = route.calls[1].request.read() + assert b"SELECT * FROM" in body + assert b"`top_sales`" 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(["id"], [{"id": 1}], metadata=["BIGINT"]), + ] + ) make_client().columns("dfs.tmp", "archive.2024.json") body = route.calls[1].request.read() assert b"`archive.2024.json`" in body @@ -548,14 +690,11 @@ class TestFilePluginMetadata: 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"}], - ), + query_response(["id"], [{"id": 1}], metadata=["BIGINT"]), ] ) assert make_client().columns("dfs.tmp", "README") == [ - {"name": "id", "data_type": "BIGINT", "nullable": True} + {"name": "id", "data_type": "BIGINT", "nullable": None} ] assert b"`README`" in route.calls[1].request.read() @@ -576,7 +715,9 @@ class TestFilePluginMetadata: assert not route.called @respx.mock - def test_columns_uses_information_schema_for_a_non_file_plugin(self): + def test_columns_uses_describe_for_a_non_dynamic_plugin(self): + # "jdbc" is not in DYNAMIC_SCHEMA_TYPES, so DESCRIBE (metadata-only, + # never reads user data) answers directly -- no probe needed. route = respx.post(f"{BASE}/query.json").mock( side_effect=[ self._schemata("jdbc"), @@ -588,7 +729,49 @@ class TestFilePluginMetadata: ) result = make_client().columns("mysql.app", "t") assert result == [{"name": "id", "data_type": "INTEGER", "nullable": False}] - assert b"INFORMATION_SCHEMA" in route.calls[1].request.read() + body = route.calls[1].request.read() + assert b"DESCRIBE" in body + assert b"LIMIT 1" not in body + assert b"SELECT *" not in body + + @respx.mock + def test_columns_on_an_http_plugin_raises_an_explanatory_error(self): + respx.post(f"{BASE}/query.json").mock(return_value=self._schemata("http")) + with pytest.raises(DrillError) as exc_info: + make_client().columns("dfs.tmp", "results") + message = str(exc_info.value) + assert "dfs.tmp" in message + 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):
