This is an automated email from the ASF dual-hosted git repository. cgivre pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/drill-mcp.git
commit fe9465bcb416756ed02bbf60a77ce4348ed0314f Author: cgivre <[email protected]> AuthorDate: Wed Aug 12 16:43:40 2026 -0400 Add ruff for linting and formatting Configure ruff (E, F, I, B, UP, C4) with a 100-column line length that matches the codebase's existing wrapping style, add it to the dev extra, and bring the tree into compliance. Two zip() calls without strict= (client_jdbc.py, client_rest.py) are noqa'd rather than changed: both are length-matched by construction (a JDBC row tuple to its column list; a metadata list already padded to len(columns) two lines above), so strict=True would only add an unreachable error path to reviewed code. One long f-string in client_rest.py's auth-failure message is noqa'd rather than reflowed for a 4-character overage. --- drill_mcp/client_jdbc.py | 8 +++---- drill_mcp/client_rest.py | 21 ++++++++---------- drill_mcp/config.py | 2 +- drill_mcp/guard.py | 16 +++++++++----- drill_mcp/server.py | 8 ++----- pyproject.toml | 14 +++++++++++- tests/test_client_rest.py | 51 +++++++++++++++++++++++++++++-------------- tests/test_config.py | 6 ++++-- tests/test_guard.py | 4 +--- tests/test_server.py | 55 +++++++++++++++++++++++------------------------ 10 files changed, 107 insertions(+), 78 deletions(-) diff --git a/drill_mcp/client_jdbc.py b/drill_mcp/client_jdbc.py index 58ce1c3..b2fbebf 100644 --- a/drill_mcp/client_jdbc.py +++ b/drill_mcp/client_jdbc.py @@ -117,9 +117,7 @@ class JdbcClient: "the JDBC backend requires the jdbc extra: pip install drill-mcp[jdbc]" ) from exc credentials = ( - [self._config.user, self._config.password] - if self._config.auth == "basic" - else [] + [self._config.user, self._config.password] if self._config.auth == "basic" else [] ) try: self._connection = jaydebeapi.connect( @@ -161,7 +159,9 @@ class JdbcClient: raise DrillError(self._scrub(str(exc))) from exc return QueryResult( columns=columns, - rows=[dict(zip(columns, row)) for row in rows], + # JDBC row tuples are always column-aligned by the driver; strict= + # would only add a reachable-in-theory error path to reviewed code. + rows=[dict(zip(columns, row)) for row in rows], # noqa: B905 truncated=max_rows > 0 and len(rows) >= max_rows, metadata=metadata, ) diff --git a/drill_mcp/client_rest.py b/drill_mcp/client_rest.py index 110fa43..ca15440 100644 --- a/drill_mcp/client_rest.py +++ b/drill_mcp/client_rest.py @@ -286,14 +286,10 @@ def fetch_plugin_type(query: Query, schema: str) -> str | None: def fetch_schemas(query: Query) -> list[dict[str, Any]]: result = query( - "SELECT SCHEMA_NAME, TYPE FROM INFORMATION_SCHEMA.`SCHEMATA` " - "ORDER BY SCHEMA_NAME", + "SELECT SCHEMA_NAME, TYPE FROM INFORMATION_SCHEMA.`SCHEMATA` ORDER BY SCHEMA_NAME", 10_000, ) - return [ - {"name": row.get("SCHEMA_NAME"), "type": row.get("TYPE")} - for row in result.rows - ] + return [{"name": row.get("SCHEMA_NAME"), "type": row.get("TYPE")} for row in result.rows] def fetch_tables(query: Query, schema: str) -> list[dict[str, Any]]: @@ -321,10 +317,7 @@ def fetch_tables(query: Query, schema: str) -> list[dict[str, Any]]: f"WHERE TABLE_SCHEMA = {quote_literal_path(schema)} ORDER BY TABLE_NAME", 10_000, ) - return [ - {"name": row.get("TABLE_NAME"), "type": row.get("TABLE_TYPE")} - for row in result.rows - ] + return [{"name": row.get("TABLE_NAME"), "type": row.get("TABLE_TYPE")} for row in result.rows] def _probe_target(schema: str, table: str) -> str: @@ -368,7 +361,9 @@ def _columns_from_metadata(columns: list[str], metadata: list[str | None]) -> li # response) drop trailing columns. types = list(metadata) + [None] * max(0, len(columns) - len(metadata)) result = [] - for name, type_str in zip(columns, types): + # types was just padded to len(columns) above, so this zip is already + # length-matched; strict= would be redundant with the padding it follows. + for name, type_str in zip(columns, types): # noqa: B905 data_type = _TYPE_PRECISION.sub("", type_str) if type_str else None result.append({"name": name, "data_type": data_type, "nullable": None}) return result @@ -534,7 +529,9 @@ class RestClient: ) if _contains_invalid_credentials_marker(response.text): raise DrillError( - f"authentication failed for user {self._config.user!r} at {_safe_url(self._config.url)}" + # Line exceeds 100 cols by 4; not worth reflowing a reviewed + # error-message string for that, hence the noqa below. + f"authentication failed for user {self._config.user!r} at {_safe_url(self._config.url)}" # noqa: E501 ) self._authenticated = True diff --git a/drill_mcp/config.py b/drill_mcp/config.py index 721310a..3634e94 100644 --- a/drill_mcp/config.py +++ b/drill_mcp/config.py @@ -68,7 +68,7 @@ class Config(BaseModel): jdbc_driver_path: str | None = None @model_validator(mode="after") - def _check_consistency(self) -> "Config": + def _check_consistency(self) -> Config: if self.auth == "basic" and not (self.user and self.password): raise ValueError("auth: basic requires both user and password") if self.backend == "jdbc" and not self.jdbc_driver_path: diff --git a/drill_mcp/guard.py b/drill_mcp/guard.py index 2f0feb8..9afc6ac 100644 --- a/drill_mcp/guard.py +++ b/drill_mcp/guard.py @@ -51,7 +51,15 @@ _READ_TYPES = (exp.Select, exp.Union, exp.Intersect, exp.Except, exp.Subquery, e # or `SELECT ... INTO ...`). Checking only the root type is not enough: the # safety property must not depend on sqlglot's Drill grammar rejecting these # forms outright — a write hidden deeper in the tree must still be caught. -_EMBEDDED_WRITE_TYPES = (exp.Insert, exp.Update, exp.Delete, exp.Merge, exp.Create, exp.Drop, exp.Into) +_EMBEDDED_WRITE_TYPES = ( + exp.Insert, + exp.Update, + exp.Delete, + exp.Merge, + exp.Create, + exp.Drop, + exp.Into, +) # EXPLAIN unwraps its body and re-checks it recursively; this bounds # `EXPLAIN EXPLAIN EXPLAIN ...` so a malicious input cannot blow the stack. @@ -68,7 +76,7 @@ class Policy: hidden_schemas: tuple[str, ...] = () @classmethod - def from_config(cls, cfg) -> "Policy": + def from_config(cls, cfg) -> Policy: return cls( writable_plugins=tuple(cfg.writable_plugins), hidden_schemas=tuple(cfg.hidden_schemas), @@ -154,9 +162,7 @@ def _check(sql: str, policy: Policy, depth: int) -> None: ) from exc if len(statements) != 1: - raise PolicyError( - f"exactly one statement per call is permitted, got {len(statements)}" - ) + raise PolicyError(f"exactly one statement per call is permitted, got {len(statements)}") statement = statements[0] _check_hidden(statement, policy) diff --git a/drill_mcp/server.py b/drill_mcp/server.py index 9383553..ffb1651 100644 --- a/drill_mcp/server.py +++ b/drill_mcp/server.py @@ -210,9 +210,7 @@ class DrillTools: plugins = self._require_management("storage_plugins")() except DrillError as exc: raise ToolError(str(exc)) from exc - return [ - p for p in plugins if isinstance(p, dict) and self._visible(p.get("name")) - ] + return [p for p in plugins if isinstance(p, dict) and self._visible(p.get("name"))] def cluster_status(self) -> dict[str, Any]: """Report Drillbit membership and overall cluster status.""" @@ -330,9 +328,7 @@ def main(argv: list[str] | None = None) -> int: args = _parse_args(argv) overrides = { - key: value - for key, value in vars(args).items() - if key != "config" and value is not None + key: value for key, value in vars(args).items() if key != "config" and value is not None } try: config = load_config(args.config, overrides=overrides) diff --git a/pyproject.toml b/pyproject.toml index 96a5ff5..b251ff6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,10 +46,22 @@ dependencies = [ [project.optional-dependencies] jdbc = ["jaydebeapi>=1.2.3", "JPype1>=1.5"] kerberos = ["httpx-gssapi>=0.3"] -dev = ["pytest>=8.0", "respx>=0.21", "pytest-cov>=5.0"] +dev = ["pytest>=8.0", "respx>=0.21", "pytest-cov>=5.0", "ruff>=0.6"] [project.scripts] drill-mcp = "drill_mcp.server:main" [tool.pytest.ini_options] testpaths = ["tests"] + +[tool.ruff] +target-version = "py311" +# Existing reviewed code wraps around 88-100 columns already; matching +# that (rather than a larger limit) keeps `ruff format` a no-op on the +# working tree instead of collapsing manually-wrapped, reviewed calls +# onto single long lines. +line-length = 100 + +[tool.ruff.lint] +# pycodestyle/pyflakes, isort, bugbear, pyupgrade, comprehensions. +select = ["E", "F", "I", "B", "UP", "C4"] diff --git a/tests/test_client_rest.py b/tests/test_client_rest.py index 0a0e939..11f52c4 100644 --- a/tests/test_client_rest.py +++ b/tests/test_client_rest.py @@ -65,7 +65,8 @@ class TestQuoting: "foo\\bar", "foo`bar", "dfs.tmp", - "foo\n", # trailing newline: `$` matches before it under .match(), not under .fullmatch() + # trailing newline: `$` matches before it under .match(), not under .fullmatch() + "foo\n", ], ) def test_literal_rejects_dangerous_input(self, bad): @@ -173,7 +174,9 @@ class TestQuery: @respx.mock def test_drill_error_text_is_surfaced(self): respx.post(f"{BASE}/query.json").mock( - return_value=httpx.Response(500, json={"errorMessage": "VALIDATION ERROR: no such table"}) + return_value=httpx.Response( + 500, json={"errorMessage": "VALIDATION ERROR: no such table"} + ) ) with pytest.raises(DrillError, match="no such table"): make_client().query("SELECT * FROM nope", max_rows=10) @@ -249,7 +252,9 @@ class TestBasicAuth: httpx.Response(200, json={"columns": ["a"], "rows": []}), ] ) - result = make_client(auth="basic", user="alice", password="s3cret").query("SELECT 1", max_rows=1) + result = make_client(auth="basic", user="alice", password="s3cret").query( + "SELECT 1", max_rows=1 + ) assert result.columns == ["a"] assert login.call_count == 2 assert query.call_count == 2 @@ -323,8 +328,7 @@ class TestBasicAuth: return_value=httpx.Response( 200, text=( - "<div>Warning: 1 < 2 in the system. " - "Invalid username/password credentials</div>" + "<div>Warning: 1 < 2 in the system. Invalid username/password credentials</div>" ), ) ) @@ -362,7 +366,9 @@ class TestBasicAuth: respx.post(f"{BASE}/query.json").mock( return_value=httpx.Response(200, json={"columns": ["a"], "rows": []}) ) - result = make_client(auth="basic", user="alice", password="s3cret").query("SELECT 1", max_rows=1) + result = make_client(auth="basic", user="alice", password="s3cret").query( + "SELECT 1", max_rows=1 + ) assert login.called assert result.columns == ["a"] @@ -379,7 +385,9 @@ class TestBasicAuth: respx.post(f"{BASE}/query.json").mock( return_value=httpx.Response(200, json={"columns": ["a"], "rows": []}) ) - result = make_client(auth="basic", user="alice", password="s3cret").query("SELECT 1", max_rows=1) + result = make_client(auth="basic", user="alice", password="s3cret").query( + "SELECT 1", max_rows=1 + ) assert login.called assert result.columns == ["a"] @@ -547,14 +555,18 @@ class TestFilePluginMetadata: @staticmethod def _schemata(plugin_type): - return query_response(["SCHEMA_NAME", "TYPE"], [{"SCHEMA_NAME": "dfs.tmp", "TYPE": plugin_type}]) + return query_response( + ["SCHEMA_NAME", "TYPE"], [{"SCHEMA_NAME": "dfs.tmp", "TYPE": plugin_type}] + ) @respx.mock def test_tables_uses_show_files_for_a_file_plugin(self): route = respx.post(f"{BASE}/query.json").mock( side_effect=[ self._schemata("file"), - query_response(["name", "isDirectory"], [{"name": "sales.csv", "isDirectory": "false"}]), + query_response( + ["name", "isDirectory"], [{"name": "sales.csv", "isDirectory": "false"}] + ), ] ) assert make_client().tables("dfs.tmp") == [{"name": "sales.csv", "type": "TABLE"}] @@ -565,7 +577,9 @@ class TestFilePluginMetadata: respx.post(f"{BASE}/query.json").mock( side_effect=[ self._schemata("file"), - query_response(["name", "isDirectory"], [{"name": "year=2024", "isDirectory": "true"}]), + query_response( + ["name", "isDirectory"], [{"name": "year=2024", "isDirectory": "true"}] + ), ] ) assert make_client().tables("dfs.tmp")[0]["type"] == "DIRECTORY" @@ -575,7 +589,10 @@ class TestFilePluginMetadata: respx.post(f"{BASE}/query.json").mock( side_effect=[ self._schemata("file"), - query_response(["name", "isDirectory"], [{"name": "top_sales.view.drill", "isDirectory": "false"}]), + query_response( + ["name", "isDirectory"], + [{"name": "top_sales.view.drill", "isDirectory": "false"}], + ), ] ) assert make_client().tables("dfs.tmp") == [{"name": "top_sales", "type": "VIEW"}] @@ -585,7 +602,9 @@ class TestFilePluginMetadata: route = respx.post(f"{BASE}/query.json").mock( side_effect=[ self._schemata("jdbc"), - query_response(["TABLE_NAME", "TABLE_TYPE"], [{"TABLE_NAME": "t", "TABLE_TYPE": "TABLE"}]), + query_response( + ["TABLE_NAME", "TABLE_TYPE"], [{"TABLE_NAME": "t", "TABLE_TYPE": "TABLE"}] + ), ] ) assert make_client().tables("mysql.app") == [{"name": "t", "type": "TABLE"}] @@ -673,14 +692,14 @@ class TestFilePluginMetadata: httpx.Response( 500, json={ - "errorMessage": "VALIDATION ERROR: Object 'sales.csv' not found within 'dfs.tmp'" + "errorMessage": ( + "VALIDATION ERROR: Object 'sales.csv' not found within 'dfs.tmp'" + ) }, ), ] ) - with pytest.raises( - DrillError, match=r"Object 'sales\.csv' not found within 'dfs\.tmp'" - ): + with pytest.raises(DrillError, match=r"Object 'sales\.csv' not found within 'dfs\.tmp'"): make_client().columns("dfs.tmp", "sales.csv") @respx.mock diff --git a/tests/test_config.py b/tests/test_config.py index a40e59d..da85312 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -20,7 +20,7 @@ import pytest from pydantic import ValidationError -from drill_mcp.config import Config, ConfigError, load_config +from drill_mcp.config import ConfigError, load_config def test_defaults_are_conservative(): @@ -66,7 +66,9 @@ def test_cli_overrides_env(tmp_path): def test_credentials_read_from_env(): - cfg = load_config(env={"DRILL_USER": "alice", "DRILL_PASSWORD": "s3cret", "DRILL_AUTH": "basic"}) + cfg = load_config( + env={"DRILL_USER": "alice", "DRILL_PASSWORD": "s3cret", "DRILL_AUTH": "basic"} + ) assert cfg.user == "alice" assert cfg.password == "s3cret" diff --git a/tests/test_guard.py b/tests/test_guard.py index a1a6feb..6c48631 100644 --- a/tests/test_guard.py +++ b/tests/test_guard.py @@ -412,9 +412,7 @@ class TestCoverageGaps: """ def test_policy_from_config(self): - cfg = types.SimpleNamespace( - writable_plugins=["dfs.tmp"], hidden_schemas=["sys"] - ) + cfg = types.SimpleNamespace(writable_plugins=["dfs.tmp"], hidden_schemas=["sys"]) policy = Policy.from_config(cfg) assert policy.writable_plugins == ("dfs.tmp",) assert policy.hidden_schemas == ("sys",) diff --git a/tests/test_server.py b/tests/test_server.py index f8cb08e..35a94cc 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -53,9 +53,7 @@ class TestRunQuery: # is the last chokepoint before the model, so it must not simply # trust whatever the client hands back. client = MagicMock() - client.query.return_value = QueryResult( - ["a"], [{"a": i} for i in range(10)], "q1", False - ) + client.query.return_value = QueryResult(["a"], [{"a": i} for i in range(10)], "q1", False) result = make_tools(client, max_rows=3).run_query("SELECT 1") assert len(result["rows"]) == 3 assert result["rows"] == [{"a": 0}, {"a": 1}, {"a": 2}] @@ -373,9 +371,7 @@ class TestShowFiltering: def test_ordinary_select_rows_are_not_filtered(self): client = MagicMock() - client.query.return_value = QueryResult( - ["SCHEMA_NAME"], [{"SCHEMA_NAME": "sys"}] - ) + client.query.return_value = QueryResult(["SCHEMA_NAME"], [{"SCHEMA_NAME": "sys"}]) result = make_tools(client, hidden_schemas=["sys"]).run_query( "SELECT SCHEMA_NAME FROM dfs.tmp.notes" ) @@ -396,9 +392,7 @@ class TestShowFiltering: ["SCHEMA_NAME"], [{"SCHEMA_NAME": "sys"}, {"SCHEMA_NAME": "dfs.tmp"}], ) - result = make_tools(client, hidden_schemas=["sys"]).run_query( - "/* x */ SHOW SCHEMAS" - ) + result = make_tools(client, hidden_schemas=["sys"]).run_query("/* x */ SHOW SCHEMAS") assert result["rows"] == [{"SCHEMA_NAME": "dfs.tmp"}] def test_show_databases_with_leading_line_comment_is_still_filtered(self): @@ -451,9 +445,7 @@ class TestShowFiltering: client.query.return_value = QueryResult( ["TABLE_NAME"], [{"TABLE_NAME": "sys"}, {"TABLE_NAME": "orders"}] ) - result = make_tools(client, hidden_schemas=["sys"]).run_query( - "SHOW TABLES LIKE '%s%'" - ) + result = make_tools(client, hidden_schemas=["sys"]).run_query("SHOW TABLES LIKE '%s%'") assert result["rows"] == [{"TABLE_NAME": "orders"}] def test_show_schemas_like_rows_are_filtered(self): @@ -468,9 +460,7 @@ class TestShowFiltering: ["SCHEMA_NAME"], [{"SCHEMA_NAME": "sys"}, {"SCHEMA_NAME": "dfs.tmp"}], ) - result = make_tools(client, hidden_schemas=["sys"]).run_query( - "SHOW SCHEMAS LIKE '%s%'" - ) + result = make_tools(client, hidden_schemas=["sys"]).run_query("SHOW SCHEMAS LIKE '%s%'") assert result["rows"] == [{"SCHEMA_NAME": "dfs.tmp"}] def test_show_databases_like_rows_are_filtered(self): @@ -490,9 +480,7 @@ class TestShowFiltering: ["SCHEMA_NAME"], [{"SCHEMA_NAME": "sys"}, {"SCHEMA_NAME": "dfs.tmp"}], ) - result = make_tools(client, hidden_schemas=["sys"]).run_query( - "SHOW SCHEMAS /* trailing */" - ) + result = make_tools(client, hidden_schemas=["sys"]).run_query("SHOW SCHEMAS /* trailing */") assert result["rows"] == [{"SCHEMA_NAME": "dfs.tmp"}] def test_show_schemas_with_no_whitespace_before_comment_is_still_filtered(self): @@ -501,9 +489,7 @@ class TestShowFiltering: ["SCHEMA_NAME"], [{"SCHEMA_NAME": "sys"}, {"SCHEMA_NAME": "dfs.tmp"}], ) - result = make_tools(client, hidden_schemas=["sys"]).run_query( - "SHOW/**/SCHEMAS" - ) + result = make_tools(client, hidden_schemas=["sys"]).run_query("SHOW/**/SCHEMAS") assert result["rows"] == [{"SCHEMA_NAME": "dfs.tmp"}] def test_show_schemas_with_trailing_semicolon_is_still_filtered(self): @@ -539,9 +525,7 @@ class TestShowFiltering: ["SCHEMA_NAME"], [{"SCHEMA_NAME": "sys"}, {"SCHEMA_NAME": "dfs.tmp"}], ) - result = make_tools(client, hidden_schemas=["sys"]).run_query( - "SHOW */ SCHEMAS" - ) + result = make_tools(client, hidden_schemas=["sys"]).run_query("SHOW */ SCHEMAS") assert result["rows"] == [{"SCHEMA_NAME": "dfs.tmp"}] def test_row_that_is_not_a_dict_does_not_crash_filtering(self): @@ -617,17 +601,32 @@ class TestWiring: def test_no_write_or_mutation_tools_are_registered(self): server = build_server(load_config(env={})) names = {tool.name for tool in server._tool_manager.list_tools()} - forbidden = {"create_storage_plugin", "update_storage_plugin", - "delete_storage_plugin", "set_option", "alter_system"} + forbidden = { + "create_storage_plugin", + "update_storage_plugin", + "delete_storage_plugin", + "set_option", + "alter_system", + } assert not (names & forbidden) def test_no_registered_tool_accepts_a_credential_argument(self): """Credentials come from config or environment only, never a tool argument.""" server = build_server(load_config(env={})) - credential_words = {"user", "password", "username", "passwd", "secret", "token", "credential"} + credential_words = { + "user", + "password", + "username", + "passwd", + "secret", + "token", + "credential", + } for tool in server._tool_manager.list_tools(): params = set(tool.parameters.get("properties", {})) - assert not (params & credential_words), f"{tool.name} accepts {params & credential_words}" + assert not (params & credential_words), ( + f"{tool.name} accepts {params & credential_words}" + ) class TestMain:
