nsivabalan commented on code in PR #20329:
URL: https://github.com/apache/hudi/pull/20329#discussion_r4395355033
##########
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:
⚠️ **IMPORTANT.** This closes path-named tables, but Hudi's own table-valued
functions still accept raw paths as string arguments. `HoodieSparkBaseAnalysis`
treats a first argument containing `/` as the table base path for
`hudi_table_changes`, `hudi_query_timeline`, `hudi_filesystem_view` and the
others. These all pass the guardrail as written, with or without
`runSQLOnFiles=false`:
```sql
SELECT * FROM hudi_table_changes('s3a://other-bucket/tbl', 'latest_state',
'earliest')
SELECT * FROM hudi_query_timeline('s3a://other-bucket/tbl', 'false')
```
Some options:
- Reject any table function (`exp.Table` whose `this` is a function) with a
string-literal argument containing `/`, `\` or `:`.
- Allow-list table functions in the Spark dialect: `range`, `explode`, and
the Hudi TVFs only with identifier arguments.
Relatedly, `IDENTIFIER('text.`/etc/passwd`')` passes through unchanged. The
chart's `runSQLOnFiles=false` covers it, but it's cheap to reject
`IDENTIFIER(...)` in the Spark dialect too, since the gateway can target other
Connect servers. Please add both to
`test_spark_script_and_path_reads_are_rejected`.
--
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]
