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 47d99a40b5158aea0780be1f4725e108b2db8de8 Author: cgivre <[email protected]> AuthorDate: Tue Aug 11 15:34:29 2026 -0400 feat: SQL guard with deny-by-default write policy --- drill_mcp/guard.py | 146 +++++++++++++++++++++++++++++++++++++++++++++ tests/test_guard.py | 168 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 314 insertions(+) diff --git a/drill_mcp/guard.py b/drill_mcp/guard.py new file mode 100644 index 0000000..f2cdeeb --- /dev/null +++ b/drill_mcp/guard.py @@ -0,0 +1,146 @@ +"""SQL policy enforcement. Pure: string in, decision out, no I/O. + +Deny by default. Reads are permitted; a write is permitted only when its target +resolves into a plugin the operator explicitly listed in `writable_plugins`. + +A real parser rather than a regex, deliberately: a regex guard is defeated by +`-- CREATE TABLE` in a comment, by `'DROP TABLE'` inside a string literal, and +by statement stacking. This module is the only thing standing between a model +and the user's data. +""" + +from __future__ import annotations + +import re +from collections.abc import Iterable +from dataclasses import dataclass + +import sqlglot +from sqlglot import exp + +DIALECT = "postgres" # closest available fit for Drill's Calcite SQL + +# Commands sqlglot does not model as expressions, but which cannot write. +_SAFE_COMMANDS = {"SHOW", "DESCRIBE", "DESC", "EXPLAIN"} + +_READ_TYPES = (exp.Select, exp.Union, exp.Intersect, exp.Except, exp.Subquery, exp.Describe) + + +class PolicyError(Exception): + """Raised when a statement is not permitted. The message is shown to the caller.""" + + +@dataclass(frozen=True) +class Policy: + writable_plugins: tuple[str, ...] = () + hidden_schemas: tuple[str, ...] = () + + @classmethod + def from_config(cls, cfg) -> "Policy": + return cls( + writable_plugins=tuple(cfg.writable_plugins), + hidden_schemas=tuple(cfg.hidden_schemas), + ) + + +def matches_prefix(qualified: str, entries: Iterable[str]) -> bool: + """True if any entry is a dotted-component prefix of `qualified`. + + Component-wise so that `dfs` matches `dfs.tmp` but not `dfsx.tmp`. + """ + parts = [p.lower() for p in qualified.split(".") if p] + if not parts: + return False + for entry in entries: + entry_parts = [p.lower() for p in entry.split(".") if p] + if entry_parts and parts[: len(entry_parts)] == entry_parts: + return True + return False + + +def check(sql: str, policy: Policy) -> None: + """Return None if `sql` is permitted under `policy`; raise PolicyError otherwise.""" + if not sql or not sql.strip(): + raise PolicyError("empty SQL statement") + + try: + statements = [s for s in sqlglot.parse(sql, read=DIALECT) if s is not None] + except sqlglot.ParseError as exc: + raise PolicyError( + f"could not parse SQL, so it cannot be checked against policy and is rejected: {exc}" + ) from exc + + if len(statements) != 1: + raise PolicyError( + f"exactly one statement per call is permitted, got {len(statements)}" + ) + + statement = statements[0] + _check_hidden(statement, policy) + _check_write(statement, policy) + + +def _check_write(statement: exp.Expression, policy: Policy) -> None: + if isinstance(statement, _READ_TYPES): + return + + if isinstance(statement, exp.Command): + keyword = str(statement.this or "").upper() + if keyword in _SAFE_COMMANDS: + return + raise PolicyError(f"statement type {keyword or 'UNKNOWN'} is not permitted") + + if isinstance(statement, (exp.Create, exp.Drop)): + kind = (statement.args.get("kind") or "").upper() + if kind not in {"TABLE", "VIEW"}: + raise PolicyError(f"{statement.key.upper()} {kind or 'UNKNOWN'} is not permitted") + target = _write_target(statement) + if target is None: + raise PolicyError("could not determine the target of this statement, rejecting") + qualified = _schema_prefix(target) + if not qualified: + raise PolicyError( + f"write target '{target.name}' is not schema-qualified; " + "qualify it with a plugin listed in writable_plugins" + ) + if not matches_prefix(qualified, policy.writable_plugins): + raise PolicyError( + f"writes to '{qualified}' are not permitted; " + f"add it to writable_plugins to allow this (currently: " + f"{list(policy.writable_plugins) or 'none'})" + ) + return + + raise PolicyError(f"statement type {statement.key.upper()} is not permitted") + + +def _write_target(statement: exp.Expression) -> exp.Table | None: + target = statement.this + if isinstance(target, exp.Schema): + target = target.this + return target if isinstance(target, exp.Table) else None + + +def _schema_prefix(table: exp.Table) -> str: + return ".".join(part for part in (table.catalog, table.db) if part) + + +def _check_hidden(statement: exp.Expression, policy: Policy) -> None: + if not policy.hidden_schemas: + return + + for table in statement.find_all(exp.Table): + prefix = _schema_prefix(table) + if prefix and matches_prefix(prefix, policy.hidden_schemas): + raise PolicyError(f"schema '{prefix}' is hidden by configuration") + + if isinstance(statement, exp.Command): + # ponytail: sqlglot leaves SHOW's remainder as raw text, so scan it for + # hidden names. Safe because the tokenizer already stripped comments and + # string literals cannot appear in a SHOW target. Upgrade to a real parse + # if Drill's SHOW grammar ever gains an expression argument. + remainder = str(statement.args.get("expression") or "") + for entry in policy.hidden_schemas: + head = entry.split(".")[0] + if re.search(rf"\b{re.escape(head)}\b", remainder, re.IGNORECASE): + raise PolicyError(f"schema '{entry}' is hidden by configuration") diff --git a/tests/test_guard.py b/tests/test_guard.py new file mode 100644 index 0000000..42d4ff5 --- /dev/null +++ b/tests/test_guard.py @@ -0,0 +1,168 @@ +import pytest +import sqlglot +from sqlglot import exp + +from drill_mcp.guard import Policy, PolicyError, check, matches_prefix + + +class TestSqlglotAssumptions: + """Characterization tests: what the guard relies on sqlglot doing.""" + + def test_parse_returns_one_statement_per_semicolon(self): + assert len(sqlglot.parse("SELECT 1; SELECT 2", read="postgres")) == 2 + + def test_select_parses_to_select(self): + stmt = sqlglot.parse_one("SELECT * FROM dfs.tmp.foo", read="postgres") + assert isinstance(stmt, exp.Select) + + def test_table_exposes_catalog_db_name(self): + table = sqlglot.parse_one("SELECT * FROM dfs.tmp.foo", read="postgres").find(exp.Table) + assert table.catalog == "dfs" + assert table.db == "tmp" + assert table.name == "foo" + + def test_two_part_name_populates_db_not_catalog(self): + table = sqlglot.parse_one("SELECT * FROM sys.options", read="postgres").find(exp.Table) + assert table.catalog == "" + assert table.db == "sys" + assert table.name == "options" + + def test_comments_are_stripped_by_the_tokenizer(self): + stmt = sqlglot.parse_one("-- CREATE TABLE evil\nSELECT 1", read="postgres") + assert isinstance(stmt, exp.Select) + + def test_ctas_target_is_reachable_from_this(self): + stmt = sqlglot.parse_one( + "CREATE TABLE dfs.tmp.out AS SELECT * FROM dfs.raw.src", read="postgres" + ) + assert isinstance(stmt, exp.Create) + target = stmt.this.this if isinstance(stmt.this, exp.Schema) else stmt.this + assert isinstance(target, exp.Table) + assert target.catalog == "dfs" + assert target.db == "tmp" + + +OPEN = Policy(writable_plugins=("dfs.tmp",)) +CLOSED = Policy() + + +class TestReadsAreAllowed: + @pytest.mark.parametrize( + "sql", + [ + "SELECT 1", + "SELECT * FROM dfs.tmp.foo", + "SELECT a, b FROM dfs.tmp.foo WHERE a > 1 ORDER BY b", + "WITH x AS (SELECT * FROM dfs.tmp.foo) SELECT * FROM x", + "SELECT * FROM dfs.tmp.a UNION SELECT * FROM dfs.tmp.b", + "SELECT * FROM dfs.tmp.a JOIN dfs.tmp.b ON a.id = b.id", + "SHOW SCHEMAS", + "SHOW TABLES", + "DESCRIBE dfs.tmp.foo", + ], + ) + def test_permitted(self, sql): + check(sql, CLOSED) # does not raise + + +class TestWritesAreDeniedByDefault: + @pytest.mark.parametrize( + "sql", + [ + "CREATE TABLE dfs.tmp.out AS SELECT * FROM dfs.tmp.src", + "CREATE VIEW dfs.tmp.v AS SELECT 1", + "DROP TABLE dfs.tmp.foo", + "DROP VIEW dfs.tmp.v", + "INSERT INTO dfs.tmp.foo VALUES (1)", + "ALTER SYSTEM SET `planner.width.max_per_node` = 4", + "ALTER SESSION SET `store.format` = 'json'", + "USE dfs.tmp", + "REFRESH TABLE METADATA dfs.tmp.foo", + ], + ) + def test_rejected_with_no_writable_plugins(self, sql): + with pytest.raises(PolicyError): + check(sql, CLOSED) + + +class TestWritesWithAnAllowlist: + def test_ctas_into_allowed_plugin_permitted(self): + check("CREATE TABLE dfs.tmp.out AS SELECT * FROM dfs.raw.src", OPEN) + + def test_create_view_into_allowed_plugin_permitted(self): + check("CREATE VIEW dfs.tmp.v AS SELECT 1", OPEN) + + def test_drop_in_allowed_plugin_permitted(self): + check("DROP TABLE dfs.tmp.old", OPEN) + + def test_ctas_into_other_plugin_rejected(self): + with pytest.raises(PolicyError, match="writable_plugins"): + check("CREATE TABLE s3.bucket.out AS SELECT 1", OPEN) + + def test_sibling_workspace_rejected(self): + with pytest.raises(PolicyError): + check("CREATE TABLE dfs.raw.out AS SELECT 1", OPEN) + + def test_plugin_level_entry_permits_any_workspace(self): + check("CREATE TABLE dfs.raw.out AS SELECT 1", Policy(writable_plugins=("dfs",))) + + def test_matching_is_case_insensitive(self): + check("CREATE TABLE DFS.TMP.OUT AS SELECT 1", OPEN) + + def test_insert_still_rejected_even_when_plugin_writable(self): + with pytest.raises(PolicyError): + check("INSERT INTO dfs.tmp.foo VALUES (1)", OPEN) + + def test_unqualified_target_rejected(self): + with pytest.raises(PolicyError): + check("CREATE TABLE bare AS SELECT 1", OPEN) + + +class TestInjectionAttempts: + def test_write_hidden_in_a_comment_is_not_a_write(self): + check("-- CREATE TABLE dfs.tmp.x AS SELECT 1\nSELECT 1", CLOSED) + + def test_write_inside_a_string_literal_is_not_a_write(self): + check("SELECT 'DROP TABLE dfs.tmp.foo' AS note", CLOSED) + + def test_stacked_statements_rejected(self): + with pytest.raises(PolicyError, match="one statement"): + check("SELECT 1; DROP TABLE dfs.tmp.foo", CLOSED) + + def test_stacked_statements_rejected_even_when_both_are_reads(self): + with pytest.raises(PolicyError, match="one statement"): + check("SELECT 1; SELECT 2", CLOSED) + + def test_trailing_semicolon_is_fine(self): + check("SELECT 1;", CLOSED) + + def test_empty_sql_rejected(self): + with pytest.raises(PolicyError): + check(" ", CLOSED) + + def test_unparseable_sql_rejected(self): + with pytest.raises(PolicyError, match="parse"): + check("SELECT FROM WHERE ((", CLOSED) + + +class TestMatchesPrefix: + def test_exact_match(self): + assert matches_prefix("dfs.tmp", ["dfs.tmp"]) + + def test_parent_entry_matches_child(self): + assert matches_prefix("dfs.tmp", ["dfs"]) + + def test_child_entry_does_not_match_parent(self): + assert not matches_prefix("dfs", ["dfs.tmp"]) + + def test_sibling_does_not_match(self): + assert not matches_prefix("dfs.raw", ["dfs.tmp"]) + + def test_case_insensitive(self): + assert matches_prefix("DFS.TMP", ["dfs.tmp"]) + + def test_no_partial_component_match(self): + assert not matches_prefix("dfsx.tmp", ["dfs"]) + + def test_empty_entries_never_match(self): + assert not matches_prefix("dfs.tmp", [])
