HyukjinKwon commented on code in PR #59425:
URL: https://github.com/apache/spark/pull/59425#discussion_r4270893370
##########
sql/connect/server/src/test/scala/org/apache/spark/sql/connect/pipelines/SparkDeclarativePipelinesServerSuite.scala:
##########
@@ -1164,4 +1164,127 @@ class SparkDeclarativePipelinesServerSuite
relation.getCommand.getSqlCommand.getInput)
}
}
+
+ // Helper to build a SqlCommand ExecutePlanRequest with
PipelineAnalysisContext.
+ private def buildSqlCommandRequestWithPipelineContext(
+ graphId: String,
+ input: proto.Relation): proto.ExecutePlanRequest = {
+ val context = proto.PipelineAnalysisContext
+ .newBuilder()
+ .setDataflowGraphId(graphId)
+ .build()
+ val userContext = proto.UserContext
+ .newBuilder()
+ .addExtensions(com.google.protobuf.Any.pack(context))
+ .setUserId("test_user")
+ .build()
+ val plan = proto.Plan
+ .newBuilder()
+ .setCommand(
+ proto.Command
+ .newBuilder()
+ .setSqlCommand(proto.SqlCommand.newBuilder().setInput(input).build())
+ .build())
+ .build()
+ proto.ExecutePlanRequest
+ .newBuilder()
+ .setUserContext(userContext)
+ .setPlan(plan)
+ .setSessionId(UUID.randomUUID().toString)
+ .build()
+ }
+
+ // Helper to build a SQL Relation.
+ private def sqlRelation(query: String): proto.Relation =
+ proto.Relation
+ .newBuilder()
+ .setSql(proto.SQL.newBuilder().setQuery(query))
+ .build()
+
+ // Helper to build a WITH_RELATIONS Relation with the given root SQL and
reference SQLs.
+ private def withRelationsRelation(
+ rootSql: String,
+ refs: Seq[(String, String)]): proto.Relation = {
+ val withRelationsBuilder =
proto.WithRelations.newBuilder().setRoot(sqlRelation(rootSql))
+ refs.foreach { case (alias, refSql) =>
+ withRelationsBuilder.addReferences(
+ proto.Relation
+ .newBuilder()
+ .setSubqueryAlias(
+ proto.SubqueryAlias
+ .newBuilder()
+ .setAlias(alias)
+ .setInput(sqlRelation(refSql))
+ .build())
+ .build())
+ }
+
proto.Relation.newBuilder().setWithRelations(withRelationsBuilder.build()).build()
+ }
+
+ test(
+ "SPARK-57189: blockUnsupportedSqlCommand catches Command at root of
WITH_RELATIONS " +
+ "without executing it") {
+ withRawBlockingStub { implicit stub =>
+ val graphId = createDataflowGraph
+ val tableName = "spark_catalog.default.test_with_relations_root_table"
+ try {
+ sql(s"DROP TABLE IF EXISTS $tableName")
+ val request = buildSqlCommandRequestWithPipelineContext(
+ graphId,
+ withRelationsRelation(
+ rootSql = s"CREATE TABLE $tableName (id INT) USING parquet",
+ refs = Seq("ref0" -> "SELECT 1 AS id")))
+ val ex = intercept[RuntimeException] {
+ stub.executePlan(request).next()
+ }
+
assert(ex.getMessage.contains("UNSUPPORTED_PIPELINE_SPARK_SQL_COMMAND"))
+ // Verify the guard caught the Command before any side effect could
materialize.
+ assert(!spark.catalog.tableExists(tableName))
+ } finally {
+ sql(s"DROP TABLE IF EXISTS $tableName")
+ }
+ }
+ }
+
+ test(
+ "SPARK-57189: blockUnsupportedSqlCommand catches Command inside reference
of " +
Review Comment:
**Nit (P3):** None of the new tests reaches the recursive arm for
references. `withRelationsRelation` always builds the reference input with
`sqlRelation(...)`, so this test goes through `case _` and `transformSql`. If
the recursion were replaced with
`PipelinesHandler.blockUnsupportedSqlCommand(transformRelation(ref.getSubqueryAlias.getInput))`,
all three tests would still pass. A reference whose input is itself a valid
SQL-with-references `WITH_RELATIONS` with a `CREATE TABLE` root would then run
inside the guard again.
That shape is reachable from PySpark: a DataFrame returned by
`spark.sql("...", df=...)` carries its `WITH_RELATIONS` relation, and it
becomes a reference input when you pass it to another `spark.sql` call. Could
you add a case where `ref0`'s input is a `WITH_RELATIONS` with root `CREATE
TABLE $tableName ...` and one plain reference, and assert the same error and
`!tableExists`? The helper would need a variant that takes a `proto.Relation`
for the reference input.
**Verification:**
- **Regression:** With the reference recursion replaced by a non-recursive
transformRelation check, the nested CREATE TABLE executes in the guard and the
new test fails on its error or table-existence assertion.
- **Behavior:** With the PR's recursive walker, the new test passes, and the
existing SPARK-57189 tests still pass.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala:
##########
@@ -3198,6 +3203,29 @@ class SparkConnectPlanner(
}
}
+ /**
+ * Walks the relation tree and applies
[[PipelinesHandler.blockUnsupportedSqlCommand]] to every
+ * component that would be executed by [[executeSQLWithRefs]] or
[[executeSQL]] at the normal
+ * path. For a valid SQL-with-references [[proto.WithRelations]], the root
SQL is parsed
+ * (without execution) and each reference's input is inspected recursively
so nested
+ * [[proto.WithRelations]] cannot re-trigger the executing transformation
chain. For other
+ * relation types, [[transformRelation]] produces an unresolved plan without
side effects and is
Review Comment:
**Nit (P3):** This says the `case _` branch is side-effect free, but
`transformRelation(rel)` still recurses into child relations. So a valid
SQL-with-references `WithRelations` nested under any other node still gets
executed here. For example, a reference whose input is
`Limit(WithRelations{root: SQL "CREATE TABLE t ...", refs: [...]})` goes
`transformLimit` -> `transformWithRelations` -> `transformSqlWithRefs` ->
`executeSQLWithRefs`. The CREATE TABLE runs inside the guard, the blocklist
only sees the `Limit`, and the normal path runs it again. Dataset-building
wrappers such as `SHOW_STRING` behave the same way.
The guard behaved this way before this PR too, so it is not a regression.
But this sentence, the "nested [[proto.WithRelations]] cannot re-trigger the
executing transformation chain" sentence above it, and the new call-site
comment in `handleSqlCommand` now promise more than the walker delivers. It
also matters for the planned follow-up: checking the whole parsed plan would
not catch this, because the execution happens before there is a plan to inspect.
Could you narrow the wording to what the two arms actually cover (a valid
SQL-with-references `WithRelations` at the top level or as a direct reference
input) and note the wrapped case as a known limitation?
**Verification:**
- **Inspection:** The revised comments make no unconditional
side-effect-free or no-re-trigger claim for the fallback branch, and they
describe coverage consistently with the code's two match arms.
- **Compile:** The connect server module still compiles, and the Scala style
checks accept the edited comments.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
