Skip to content

Allow canceling SQL streaming pipeline - #40666

Open
ahmedabu98 wants to merge 2 commits into
masterfrom
sql-shell-cancel
Open

ahmedabu98 wants to merge 2 commits into
masterfrom
sql-shell-cancel

Conversation

@ahmedabu98

Copy link
Copy Markdown
Contributor

Allow users to cancel a running pipeline in Beam SQL shell with Ctrl+C. Unbounded or long-running queries stop and return to the prompt instead of hanging.

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

📋 Review Summary

This PR enables canceling running Beam SQL pipelines via JDBC Statement.cancel() (such as pressing Ctrl+C in the SQL shell) by propagating Calcite's DataContext.Variable.CANCEL_FLAG into BeamEnumerableConverter. The implementation cleanly handles DirectRunner's blocking behavior, ensures global collector state is cleaned up in finally blocks, and includes comprehensive tests covering INSERT, SELECT, and SELECT ... LIMIT on unbounded sources.

🔍 General Feedback

  • The approach of disabling blockOnRun on DirectOptions and delegating the blocking wait to a cancel-aware waitUntilFinish helper is well-reasoned and avoids hanging inside DirectRunner.run().
  • Consider guarding against a potentially null cancelFlag from DataContext.Variable.CANCEL_FLAG.get(root) and avoiding redundant waitUntilFinish calls when blockOnRun is enabled.

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

💬 Inline Findings

  • sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamEnumerableConverter.java (line 129): 🟡 Medium: DataContext.Variable.CANCEL_FLAG.get(root) can return null if the provided DataContext implementation does not populate the cancelFlag variable. Passing null into toEnumerable will cause a NullPointerException when cancelFlag.get() is called in waitUntilFinish or limitRun. Consider defaulting to a new AtomicBoolean() if it is null, or keeping the toEnumerable(BeamRelNode) overload if external callers might invoke it.
      AtomicBoolean cancelFlag = DataContext.Variable.CANCEL_FLAG.get(root);
      return toEnumerable(options, node, cancelFlag != null ? cancelFlag : new AtomicBoolean());
  • sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamEnumerableConverter.java (line 233): 🟢 Low: When DirectRunner is used with blockOnRun=true (the default), run(pipeline, cancelFlag) already calls waitUntilFinish(result, cancelFlag), and then runCollector (and count on line 463) immediately calls waitUntilFinish(result, cancelFlag) a second time, spawning a second BeamSqlPipelineCanceller thread for an already-finished pipeline. In runCollector, run(...) doesn't need to wait if waitUntilFinish is called right after.

  • sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamEnumerableConverter.java (line 357): 🟡 Medium: When a pipeline is cancelled while bundles are still actively processing in background threads, Collector.globalValues.remove(id) in the finally block executes before all in-flight bundles have terminated. Any bundle starting concurrently after removal will observe values == null in startBundle and throw a NullPointerException in processElement. Guarding if (values != null) in processElement avoids noisy exceptions during cancellation.

@github-actions

Copy link
Copy Markdown
Contributor

Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment assign set of reviewers

@ahmedabu98

Copy link
Copy Markdown
Contributor Author

assign set of reviewers

@github-actions

Copy link
Copy Markdown
Contributor

Assigning reviewers:

R: @Abacn for label java.
Backup reviewer: chamikaramj (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

  • stop reviewer notifications - opt out of the automated review tooling
  • remind me after tests pass - tag the comment author after tests pass
  • waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant