Repository navigation
Allow canceling SQL streaming pipeline - #40666
ahmedabu98 wants to merge 2 commits into
Conversation
There was a problem hiding this comment.
📋 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
blockOnRunonDirectOptionsand delegating the blocking wait to a cancel-awarewaitUntilFinishhelper is well-reasoned and avoids hanging insideDirectRunner.run(). - Consider guarding against a potentially
nullcancelFlagfromDataContext.Variable.CANCEL_FLAG.get(root)and avoiding redundantwaitUntilFinishcalls whenblockOnRunis 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 returnnullif the providedDataContextimplementation does not populate thecancelFlagvariable. PassingnullintotoEnumerablewill cause aNullPointerExceptionwhencancelFlag.get()is called inwaitUntilFinishorlimitRun. Consider defaulting to a newAtomicBoolean()if it is null, or keeping thetoEnumerable(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: WhenDirectRunneris used withblockOnRun=true(the default),run(pipeline, cancelFlag)already callswaitUntilFinish(result, cancelFlag), and thenrunCollector(andcounton line 463) immediately callswaitUntilFinish(result, cancelFlag)a second time, spawning a secondBeamSqlPipelineCancellerthread for an already-finished pipeline. InrunCollector,run(...)doesn't need to wait ifwaitUntilFinishis 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 thefinallyblock executes before all in-flight bundles have terminated. Any bundle starting concurrently after removal will observevalues == nullinstartBundleand throw aNullPointerExceptioninprocessElement. Guardingif (values != null)inprocessElementavoids noisy exceptions during cancellation.
|
Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment |
|
assign set of reviewers |
|
Assigning reviewers: R: @Abacn for label java. Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
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.