cshuo commented on code in PR #20329:
URL: https://github.com/apache/hudi/pull/20329#discussion_r4396088226
##########
hudi-agent-gateway/src/hudi_agent_gateway/tools/guardrails.py:
##########
@@ -42,42 +43,58 @@
exp.Command, # also covers CALL and EXPLAIN, which parse as generic
commands
exp.Use,
)
+_SPARK_REFLECTION_FUNCTIONS = frozenset({"java_method", "reflect",
"try_reflect"})
-_HINT = "Provide exactly one read-only SELECT statement (Trino SQL)."
-
-def enforce_guardrails(sql: str, row_cap: int) -> str:
+def enforce_guardrails(
+ sql: str, row_cap: int, *, dialect: str = "trino", spark_catalog: str =
"spark_catalog"
+) -> str:
"""Validate ``sql`` as a single read-only SELECT and cap its LIMIT.
Returns the (possibly rewritten) SQL to execute. Raises
:class:`ToolInputError` on any violation.
"""
+ hint = f"Provide exactly one read-only SELECT statement ({dialect} SQL)."
try:
- statements = sqlglot.parse(sql, read="trino")
+ statements = sqlglot.parse(sql, read=dialect)
except sqlglot.errors.SqlglotError as e:
- raise ToolInputError(f"could not parse SQL: {e}", hint=_HINT) from e
+ raise ToolInputError(f"could not parse SQL: {e}", hint=hint) from e
statements = [s for s in statements if s is not None]
if len(statements) != 1:
raise ToolInputError(
- f"exactly one statement is allowed, got {len(statements)}",
hint=_HINT
+ f"exactly one statement is allowed, got {len(statements)}",
hint=hint
)
stmt = statements[0]
if not isinstance(stmt, (exp.Select, exp.SetOperation)):
raise ToolInputError(
- f"read-only: only SELECT is allowed, got {type(stmt).__name__}",
hint=_HINT
+ f"read-only: only SELECT is allowed, got {type(stmt).__name__}",
hint=hint
)
# Defense in depth: no write/DDL/command node anywhere in the tree
# (catches e.g. a CTE wrapping an INSERT).
for node in stmt.walk():
if isinstance(node, _FORBIDDEN_NODES):
raise ToolInputError(
- f"read-only: {type(node).__name__} is not allowed", hint=_HINT
+ f"read-only: {type(node).__name__} is not allowed", hint=hint
)
-
- return _cap_limit(stmt, row_cap).sql(dialect="trino")
+ if dialect == "spark" and isinstance(node, exp.QueryTransform):
+ raise ToolInputError("read-only: Spark script transforms are not
allowed", hint=hint)
+ if dialect == "spark" and isinstance(node, exp.Table):
Review Comment:
Fixed: Hudi TVFs now require literal catalog table names; path/computed
arguments, filesystem-view subpaths, and `IDENTIFIER(...)` are rejected. Added
regression tests for nested and LATERAL calls.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
