You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@drill.apache.org by cg...@apache.org on 2026/08/12 20:25:10 UTC
(drill-mcp) 33/44: feat: follow the dialect methodology for dynamic-schema plugins
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 <cg...@gmail.com>
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):