Skip to content

[SPARK-57189][CONNECT][SDP] Fix double execution and guard bypass in handleSqlCommand for WITH_RELATIONS #59425

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rdtr wants to merge 4 commits into
base: master
Choose a base branch
from
Open

[SPARK-57189][CONNECT][SDP] Fix double execution and guard bypass in handleSqlCommand for WITH_RELATIONS #59425

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e50853b
[SPARK-57189][CONNECT] Fix double execution and guard bypass in handl…
rdtr Oct 13, 2026
3921af2
Format with scalafmt and clarify the non-regression test
rdtr Oct 13, 2026
d4ed446
Scope the guard comments to what it covers and test a WITH_RELATIONS …
rdtr Oct 16, 2026
64236ec
Run the command tests also inside a pipeline flow function
rdtr Oct 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 30 additions & 1 deletion ...nect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3094,8 +3094,13 @@ class SparkConnectPlanner(
}

// Block unsupported SQL commands if the request comes from Spark Declarative Pipelines.
// For a valid SQL-with-references WITH_RELATIONS, transformRelation would route
// through executeSQLWithRefs and eagerly run the root SQL as a side effect, producing
// a CommandResult that bypasses the blocklist. blockUnsupportedSqlCommandIn checks such
// a relation at the top level or as a direct reference input without executing it. See
// its scaladoc for a known limitation.
if (PipelineAnalysisContextUtils.hasPipelineAnalysisContext(userContext)) {
PipelinesHandler.blockUnsupportedSqlCommand(queryPlan = transformRelation(relation))
blockUnsupportedSqlCommandIn(relation)
}

// If the spark.sql() is called inside a pipeline flow function, we don't need to execute
Expand Down Expand Up @@ -3198,6 +3203,30 @@ class SparkConnectPlanner(
}
}

/**
* Applies [[PipelinesHandler.blockUnsupportedSqlCommand]] to a relation without executing a
* valid SQL-with-references [[proto.WithRelations]] at the top level or as a direct reference
* input. For such a relation, the root SQL is parsed (not executed) and each reference input is
* checked with this method again. Any other relation is passed to [[transformRelation]].
*
* Known limitation: [[transformRelation]] still executes a SQL-with-references
* [[proto.WithRelations]] that is nested under another relation (for example a `Limit`) or
* wrapped in a Dataset-building relation such as `SHOW_STRING`, as it did before this check.
*/
private def blockUnsupportedSqlCommandIn(rel: proto.Relation): Unit = {
rel.getRelTypeCase match {
case proto.Relation.RelTypeCase.WITH_RELATIONS
if isValidSQLWithRefs(rel.getWithRelations) =>
val withRel = rel.getWithRelations
PipelinesHandler.blockUnsupportedSqlCommand(transformSql(withRel.getRoot.getSql))
withRel.getReferencesList.asScala.foreach { ref =>
blockUnsupportedSqlCommandIn(ref.getSubqueryAlias.getInput)
}
case _ =>
PipelinesHandler.blockUnsupportedSqlCommand(transformRelation(rel))
}
}

private def isValidSQLWithRefs(query: proto.WithRelations): Boolean = {
query.getRoot.getRelTypeCase match {
case proto.Relation.RelTypeCase.SQL =>
Expand Down
175 changes: 175 additions & 0 deletions ...t/scala/org/apache/spark/sql/connect/pipelines/SparkDeclarativePipelinesServerSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1164,4 +1164,179 @@ class SparkDeclarativePipelinesServerSuite
relation.getCommand.getSqlCommand.getInput)
}
}

// Helper to build a SqlCommand ExecutePlanRequest with PipelineAnalysisContext. A flow name
// marks the request as coming from inside a pipeline flow function.
private def buildSqlCommandRequestWithPipelineContext(
graphId: String,
input: proto.Relation,
flowName: Option[String] = None): proto.ExecutePlanRequest = {
val contextBuilder = proto.PipelineAnalysisContext
.newBuilder()
.setDataflowGraphId(graphId)

@sarutak sarutak Oct 19, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit, not blocking. This helper only sets dataflowGraphId, never flowName, so all three new tests exercise the path outside a pipeline flow function. The guard runs at the hasPipelineAnalysisContext check, which is before the isInsidePipelineFlowFunction early return, so the bug existed inside a flow function too: on the old code the guard ran executeSQLWithRefs (e.g. a CREATE TABLE root) and the resulting CommandResult passed the blocklist. There was no second execution there (the handler then returns the relation), but the command still ran. That is also the most common real shape: inside a query function, spark.sql("... {df}", df=...) sends exactly this WITH_RELATIONS with SubqueryAlias references, and register_flow wraps flow.func() in add_pipeline_analysis_context(..., flow_name=flow.name).

Could the helper take an optional flowName and the two command tests run both with and without it (e.g. via gridTest)? On the old code the flow-function variant would also fail on both intercept[RuntimeException] and assert(!tableExists); on the current code it pins that the guard keeps running before the flow-function return. This dovetails with @HyukjinKwon's suggestion below to give the helper a reference-input variant, so both are about widening what the new tests actually cover.

@rdtr rdtr Oct 19, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, done in 64236ec. The helper takes an optional flowName, and the command tests run both inside and outside a flow function. All of them fail without the fix.

flowName.foreach(contextBuilder.setFlowName)
val context = contextBuilder.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 =
withRelationsRelationOf(
rootSql,
refs.map { case (alias, refSql) =>
alias -> sqlRelation(refSql)
})

// Helper to build a WITH_RELATIONS Relation with the given root SQL and reference inputs.
private def withRelationsRelationOf(
rootSql: String,
refs: Seq[(String, proto.Relation)]): proto.Relation = {
val withRelationsBuilder = proto.WithRelations.newBuilder().setRoot(sqlRelation(rootSql))
refs.foreach { case (alias, input) =>
withRelationsBuilder.addReferences(
proto.Relation
.newBuilder()
.setSubqueryAlias(
proto.SubqueryAlias
.newBuilder()
.setAlias(alias)
.setInput(input)
.build())
.build())
}
proto.Relation.newBuilder().setWithRelations(withRelationsBuilder.build()).build()
}

namedGridTest(
"SPARK-57189: blockUnsupportedSqlCommand catches Command at root of WITH_RELATIONS " +
"without executing it")(
Map("outside a flow function" -> None, "inside a flow function" -> Some("flow1"))) {
flowName =>
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,
flowName = flowName,
input = 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")
}
}
}

namedGridTest(
"SPARK-57189: blockUnsupportedSqlCommand catches Command inside reference of " +

@HyukjinKwon HyukjinKwon Oct 14, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@rdtr rdtr Oct 19, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, added in d4ed446. It fails if the recursion is replaced with a plain transformRelation check.

"WITH_RELATIONS without executing it")(
Map("outside a flow function" -> None, "inside a flow function" -> Some("flow1"))) {
flowName =>
withRawBlockingStub { implicit stub =>
val graphId = createDataflowGraph
val tableName = "spark_catalog.default.test_with_relations_ref_table"
try {
sql(s"DROP TABLE IF EXISTS $tableName")
val request = buildSqlCommandRequestWithPipelineContext(
graphId,
flowName = flowName,
input = withRelationsRelation(
rootSql = "SELECT * FROM ref0",
refs = Seq("ref0" -> s"CREATE TABLE $tableName (id INT) USING parquet")))
val ex = intercept[RuntimeException] {
stub.executePlan(request).next()
}
assert(ex.getMessage.contains("UNSUPPORTED_PIPELINE_SPARK_SQL_COMMAND"))
// Verify the Command in the reference did not execute as a side effect.
assert(!spark.catalog.tableExists(tableName))
} finally {
sql(s"DROP TABLE IF EXISTS $tableName")
}
}
}

namedGridTest(
"SPARK-57189: blockUnsupportedSqlCommand catches Command in a WITH_RELATIONS used as a " +
"reference input without executing it")(
Map("outside a flow function" -> None, "inside a flow function" -> Some("flow1"))) {
flowName =>
withRawBlockingStub { implicit stub =>
val graphId = createDataflowGraph
val tableName = "spark_catalog.default.test_with_relations_nested_ref_table"
try {
sql(s"DROP TABLE IF EXISTS $tableName")
// A DataFrame from spark.sql("...", df=...) carries a WITH_RELATIONS relation, and it
// becomes a reference input when passed to another spark.sql call.
val nested = withRelationsRelation(
rootSql = s"CREATE TABLE $tableName (id INT) USING parquet",
refs = Seq("inner0" -> "SELECT 1 AS id"))
val request = buildSqlCommandRequestWithPipelineContext(
graphId,
flowName = flowName,
input = withRelationsRelationOf(
rootSql = "SELECT * FROM ref0",
refs = Seq("ref0" -> nested)))
val ex = intercept[RuntimeException] {
stub.executePlan(request).next()
}
assert(ex.getMessage.contains("UNSUPPORTED_PIPELINE_SPARK_SQL_COMMAND"))
// Verify the Command in the nested WITH_RELATIONS did not execute as a side effect.
assert(!spark.catalog.tableExists(tableName))
} finally {
sql(s"DROP TABLE IF EXISTS $tableName")
}
}
}

test(
"SPARK-57189: blockUnsupportedSqlCommand does not over-block a valid SELECT " +
"WITH_RELATIONS in pipeline analysis context") {
// This passes with and without the fix. It keeps the new recursive check from blocking a
// valid SELECT in the future. The two tests above are the ones that verify the fix.
withRawBlockingStub { implicit stub =>
val graphId = createDataflowGraph
val request = buildSqlCommandRequestWithPipelineContext(
graphId,
withRelationsRelation(
rootSql = "SELECT * FROM ref0",
refs = Seq("ref0" -> "SELECT 1 AS id")))
val response = stub.executePlan(request).next()
assert(response.hasSqlCommandResult)

@sarutak sarutak Oct 13, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This positive test passes both before and after the fix: reverting the production change does not turn it red (the old guard did notblock a plain SELECT either, so there was no over-block to begin with). It is therefore a forward-looking guard that protects the new recursive blockUnsupportedSqlCommandIn from a future regression where it would wrongly block a valid SELECT WITH_RELATIONS. That is worth protecting, but it is worth a one line comment (or a name that reflects "does not over-block a valid SELECT") so a future reader does not mistake it for the test that proves the double-execution/bypass fix. The two negative tests above are the fix-pinning ones: without the fix, transformRelation on the WITH_RELATIONS executes the CREATE TABLE and returns a CommandResult that is not on the blocklist, so both intercept[RuntimeException] and assert(!tableExists) fail. Those are non-vacuous. Only this positive case is a forward-looking non-regression check.

@rdtr rdtr Oct 13, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point, thanks. I renamed it to "blockUnsupportedSqlCommand does not over-block a valid SELECT WITH_RELATIONS in pipeline analysis context" and added a comment that it passes with and without the fix, and that the two tests above are the ones that verify the fix.

}
}
}