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

[rdtr]

What changes were proposed in this pull request?

In SparkConnectPlanner.handleSqlCommand, the SDP guard previously received a queryPlan built by transformRelation(relation). For WITH_RELATIONS matching isValidSQLWithRefs, that chain routes through executeSQLWithRefs → session.sql(...) and eagerly executes
the root SQL (and any Commands resolved from references), producing a CommandResult that bypasses the blocklist. The normal execution path then runs the same query again.

This PR introduces a recursive helper blockUnsupportedSqlCommandIn(rel) that walks the relation tree and inspects each component the normal handler would execute, without triggering execution. For a valid SQL-with-references WITH_RELATIONS, it parses the root SQL with
transformSql to get an unresolved plan and checks each reference's SubqueryAlias.input the same way, so a valid SQL-with-references WITH_RELATIONS used directly as a reference input is also parsed instead of executed. For other relation types, the behavior is unchanged.

Known limitation: a SQL-with-references WITH_RELATIONS nested under another relation (for example a Limit), or wrapped in a Dataset-building relation such as SHOW_STRING, still goes through transformRelation and is executed inside the guard, as before this PR.

Why are the changes needed?

Two bugs defeat the SDP guard:

1. Bypassed guard. The blocklist rejects any Command, plus a few side-effecting plans that are not Command subclasses (CreateTableAsSelect, InsertIntoStatement, etc.). But after execution the plan is wrapped as CommandResult, which is neither, so it passes. Commands have already changed state by the time the guard runs.
   Commands inside reference SubqueryAlias inputs also execute, because executeSQLWithRefs runs each reference as its own query (Dataset.ofRows(session, transformRelation(ref.input))) to register its temp view.

2. Double execution. After the guard, handleSqlCommand falls through to the normal execution path and runs executeSQLWithRefs again, causing duplicate side effects.

The guard should match the runtime's execution surface: inspect both the root SQL and each reference's input, without triggering execution itself.

Does this PR introduce any user-facing change?

Yes. SDP requests containing blocked Commands inside WITH_RELATIONS (at the root or in a reference's input) now fail with UNSUPPORTED_PIPELINE_SPARK_SQL_COMMAND before any side effects occur.

How was this patch tested?

Added tests in SparkDeclarativePipelinesServerSuite. The three command tests run both outside and inside a pipeline flow function (with a flow name in the pipeline analysis context), because the guard runs before the flow-function early return:

• Command at the root of WITH_RELATIONS is blocked, and the target table is asserted not to exist (verifying no side effect).
• Command embedded in a SubqueryAlias reference input is blocked, and the target table is asserted not to exist.
• Command at the root of a WITH_RELATIONS that is itself a reference input is blocked, and the target table is asserted not to exist. This covers the recursive check of reference inputs.
• A valid SELECT WITH_RELATIONS in pipeline analysis context is not blocked. This test passes with and without the fix. It keeps the new recursive check from blocking a valid query in the future.

Without the fix, the three command tests fail in both variants because the target table is created as a side effect of the guard.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Claude Opus 5.5)

[rdtr]

@SCHJonathan @sryza @sarutak Could you take a look when you have time? This is about the SDP guard in handleSqlCommand from SPARK-54020: for WITH_RELATIONS, the guard currently runs the SQL before checking it.

[sarutak]

Thank you @rdtr, I've left one comment.

Other than the inline comment, one pre-existing limitation, not a blocker for this PR and out of its scope:
PipelinesHandler.blockUnsupportedSqlCommand inspects only the top plan node, so a command nested below a non-Command top node (for example WITH cte AS (...) INSERT INTO t ..., which parses to UnresolvedWith) passes the guard while eagerlyExecuteCommands still runs it at execution. This is a property of the blocklist, not of this change, and it affects the plain SQL path identically, so this PR neither introduces nor worsens it.

[sarutak]

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]

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.

[sarutak]

Also, a few small things on the description before merge, since it lands in git log as the commit message:

1. Accuracy: the first "Why are the changes needed?" bullet says the blocklist "checks for Command subclasses (CreateTableAsSelect, InsertIntoStatement, etc.)". Those two examples are actually the plans that are NOT Command subclasses: CreateTableAsSelect extends V2CreateTableAsSelectPlan and InsertIntoStatement extends UnaryParsedStatement, and blockUnsupportedSqlCommand checks them separately for exactly that reason (its own comment reads "Following commands are not subclasses of Command but have side effects"). Suggest rewording to distinguish the two, for example: "The blocklist rejects any Command, plus a few side-effecting plans that are not Command subclasses (CreateTableAsSelect, InsertIntoStatement, etc.)".

2. Minor mechanism accuracy (optional): the same bullet says commands inside reference SubqueryAlias inputs execute "because eagerlyExecuteCommands walks the resolved plan tree". For a reference input the command actually runs because executeSQLWithRefs eagerly executes each reference as its own query via Dataset.ofRows(session, transformRelation(ref.input)) to register the temp view. eagerlyExecuteCommands is the tree-walk mechanism for a command nested within a single query plan, which is a different path. The conclusion (reference commands run as a side effect) is correct; only the stated mechanism is slightly off.

3. Template headings: the section headings use ## but the Spark PR template uses ### (### What changes were proposed in this pull request?, etc.). Minor, but worth matching the template.

4. Could you drop the opening line "This replaces [SPARK-57189][CONNECT] Fix double execution and guard bypass in handleSqlCommand for WITH_RELATIONS #56241, which was closed as stale without a review. It is rebased on the current master with no other changes."? That context is useful in the PR conversation, but it becomes noise in the permanent commit message. The PR link history already records the lineage.

[sarutak]

The linter failure seems relevant to this change.
https://github.com/rdtr/spark/actions/runs/38358754437/job/115168705636

Could you fix it?

[sarutak]

cc: @szehon-ho, and @dongjoon-hyun from #53156

[rdtr]

@sarutak Thank you for the detailed review!

• Lint: fixed in 3921af2. scalafmt only rewrapped one doc comment in SparkConnectPlanner. dev/lint-scala passes locally.
• Description: I updated all four points. The blocklist sentence now says it rejects any Command plus a few side-effecting plans that are not Command subclasses, the reference case now points to executeSQLWithRefs running each reference with Dataset.ofRows, the headings use ###, and the opening line about [SPARK-57189][CONNECT] Fix double execution and guard bypass in handleSqlCommand for WITH_RELATIONS #56241 is removed.
• The top-node limitation: agreed that it is out of scope here. I'd like to fix it as a follow-up after this PR, by checking the whole parsed plan. I'll open a JIRA for it.

[HyukjinKwon]

Review summary

The fix works for the shapes it targets. A valid SQL-with-references WITH_RELATIONS, at the top level or used directly as a reference input, is now parsed and checked instead of executed, and the normal path runs it once. The two command tests fail without the change and pass with it.

Two small, non-blocking points are inline. First, the new scaladoc and call-site comment say the case _ fallback never executes anything and that nested WithRelations cannot re-trigger execution. That is not true for a WithRelations nested under another relation node, or for Dataset-building wrappers, which still run inside the guard just as they did before this PR. The wording should match what the walker covers. Second, the recursive reference arm has no test, so a regression there would go unnoticed.

Findings

2 total: 0 P0, 0 P1, 0 P2, 2 P3.

Nit (P3)

• Guard scaladoc claims the fallback branch is side-effect free — sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala:3212 — see inline.
• No test reaches the recursive reference branch — sql/connect/server/src/test/scala/org/apache/spark/sql/connect/pipelines/SparkDeclarativePipelinesServerSuite.scala:1250 — see inline.

PR description suggestions

• The first section says the helper "recursively inspects each reference's SubqueryAlias.input so nested WITH_RELATIONS cannot re-trigger the executing chain". That holds only when the nested WITH_RELATIONS is a reference's direct input. One nested under another relation node still goes through transformRelation and is executed in the guard, as it was before this PR. Consider scoping the sentence to direct reference inputs, or mentioning the remaining case as a known limitation.

Generated by Omnigent on Databricks.

[HyukjinKwon]

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.

[rdtr]

Thanks, fixed in d4ed446. The comments now cover only the top-level and direct-reference cases, and the scaladoc notes the wrapped case as a known limitation. I updated the PR description too.

[HyukjinKwon]

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]

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

[sarutak]

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]

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.

[sarutak]

@rdtr Could you address @HyukjinKwon's two P3 comments when you get a chance? Both look correct to me:

1. SparkConnectPlanner.scala (the case _ fallback): transformRelation recurses into child relations, so a valid SQL-with-references WITH_RELATIONS wrapped under another node (e.g. Limit(WithRelations{root: CREATE TABLE ...})) still runs inside the guard, going through transformLimit, then transformWithRelations, then executeSQLWithRefs. The "side-effect free" / "nested WITH_RELATIONS cannot re-trigger" wording promises more than the two match arms deliver; narrowing it to what the arms cover (top-level or direct reference input) and noting the wrapped case as a known limitation would be accurate.

2. The three new tests never reach the recursive reference arm: withRelationsRelation always builds reference inputs with sqlRelation(...), so replacing the recursion with a non-recursive transformRelation check would keep them green. A reference whose input is itself a WITH_RELATIONS with a CREATE TABLE root would pin that arm.

I also left one related inline nit on the test helper (flow-function coverage) that fits naturally with HyukjinKwon's helper-variant suggestion.

[rdtr]

@sarutak Thanks! Both points from @HyukjinKwon and your flow-function nit are addressed, and the PR description is updated.