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
The following commit(s) were added to refs/heads/main by this push:
new 1bef3fb Check every DROP target under sqlglot 30.20+
1bef3fb is described below
commit 1bef3fb62b611b6cbebaf68aa4561fdf735f434b
Author: cgivre <[email protected]>
AuthorDate: Tue Oct 6 09:34:45 2026 -0400
Check every DROP target under sqlglot 30.20+
sqlglot 30.20 moved DROP targets into a `tables` list, so the write
guard could not find the target and rejected permitted DROPs. Read the
list (falling back to `this`) and check each table against
writable_plugins.
---
drill_mcp/guard.py | 40 +++++++++++++++++++++-------------------
tests/test_guard.py | 6 +++++-
2 files changed, 26 insertions(+), 20 deletions(-)
diff --git a/drill_mcp/guard.py b/drill_mcp/guard.py
index a20c32a..a5919d5 100644
--- a/drill_mcp/guard.py
+++ b/drill_mcp/guard.py
@@ -212,23 +212,24 @@ def _check_write(statement: exp.Expression, policy:
Policy, depth: int) -> None:
raise PolicyError(
f"{statement.key.upper()} {kind or 'UNKNOWN'} is not permitted"
)
- target = _write_target(statement)
- if target is None:
+ targets = _write_targets(statement)
+ if not targets:
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'})"
- )
+ for target in targets:
+ 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")
@@ -250,11 +251,12 @@ def _explain_body(statement: exp.Command) -> str:
return re.sub(r"^\s*PLAN\s+FOR\s+", "", text, flags=re.IGNORECASE)
-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 _write_targets(statement: exp.Expression) -> list[exp.Table]:
+ # sqlglot >= 30.20 puts DROP targets in a `tables` list (DROP TABLE a, b);
+ # older releases and CREATE use `this`. Any non-Table target fails closed.
+ targets = statement.args.get("tables") or [statement.this]
+ targets = [t.this if isinstance(t, exp.Schema) else t for t in targets]
+ return targets if all(isinstance(t, exp.Table) for t in targets) else []
def _schema_prefix(table: exp.Table) -> str:
diff --git a/tests/test_guard.py b/tests/test_guard.py
index 36795ef..3b78d21 100644
--- a/tests/test_guard.py
+++ b/tests/test_guard.py
@@ -124,6 +124,10 @@ class TestWritesWithAnAllowlist:
def test_drop_in_allowed_plugin_permitted(self):
check("DROP TABLE dfs.tmp.old", OPEN)
+ def test_multi_table_drop_checks_every_target(self):
+ with pytest.raises(PolicyError):
+ check("DROP TABLE dfs.tmp.old, s3.bucket.prod", 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)
@@ -440,7 +444,7 @@ class TestCoverageGaps:
def test_ctas_with_explicit_column_list_permitted(self):
# The parenthesized column list wraps the target table in an
- # exp.Schema node; _write_target must unwrap it to find the table.
+ # exp.Schema node; _write_targets must unwrap it to find the table.
check("CREATE TABLE dfs.tmp.out (a INT) AS SELECT 1", OPEN)
def test_create_with_no_resolvable_target_rejected(self):
