This is an automated email from the ASF dual-hosted git repository. ahmedabu98 pushed a commit to branch sql-shell-cancel in repository https://gitbox.apache.org/repos/asf/beam.git
commit 9a56e458381f3d5000b5e2356aef0c024a3988cc Author: Ahmed Abualsaud <[email protected]> AuthorDate: Mon Oct 19 12:10:18 2026 -0700 enable canceling SQL streaming pipeline --- .../sql/impl/rel/BeamEnumerableConverter.java | 123 +++++++++++++++++---- .../sdk/extensions/sql/impl/JdbcDriverTest.java | 51 +++++++++ .../sql/impl/rel/BeamEnumerableConverterTest.java | 13 ++- 3 files changed, 159 insertions(+), 28 deletions(-) diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamEnumerableConverter.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamEnumerableConverter.java index 69eb0a0a8d5..12a93975460 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamEnumerableConverter.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamEnumerableConverter.java @@ -28,11 +28,14 @@ import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Queue; +import java.util.concurrent.CancellationException; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.stream.Collectors; import java.util.stream.StreamSupport; import org.apache.beam.runners.direct.DirectOptions; +import org.apache.beam.runners.direct.DirectRunner; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.Pipeline.PipelineVisitor; import org.apache.beam.sdk.PipelineResult; @@ -58,6 +61,7 @@ import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollection.IsBounded; import org.apache.beam.sdk.values.PValue; import org.apache.beam.sdk.values.Row; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.DataContext; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.adapter.enumerable.EnumerableRel; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.adapter.enumerable.EnumerableRelImplementor; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.adapter.enumerable.PhysType; @@ -111,16 +115,18 @@ public class BeamEnumerableConverter extends ConverterImpl implements Enumerable final PhysType physType = PhysTypeImpl.of(implementor.getTypeFactory(), rowType, prefer.preferArray()); final Expression node = implementor.stash((BeamRelNode) getInput(), BeamRelNode.class); - list.add(Expressions.call(BeamEnumerableConverter.class, "toEnumerable", node)); + list.add( + Expressions.call( + BeamEnumerableConverter.class, "toEnumerable", node, implementor.getRootExpression())); return implementor.result(physType, list.toBlock()); } - public static Enumerable<Object> toEnumerable(BeamRelNode node) { + public static Enumerable<Object> toEnumerable(BeamRelNode node, DataContext root) { final ClassLoader originalClassLoader = Thread.currentThread().getContextClassLoader(); try { Thread.currentThread().setContextClassLoader(BeamEnumerableConverter.class.getClassLoader()); final PipelineOptions options = createPipelineOptions(node.getPipelineOptions()); - return toEnumerable(options, node); + return toEnumerable(options, node, DataContext.Variable.CANCEL_FLAG.get(root)); } finally { Thread.currentThread().setContextClassLoader(originalClassLoader); } @@ -161,16 +167,21 @@ public class BeamEnumerableConverter extends ConverterImpl implements Enumerable } else if (isLimitQuery(node)) { throw new UnsupportedOperationException("Does not support queries with LIMIT in toRowList."); } - return collectRows(options, node).stream().collect(Collectors.toList()); + return collectRows(options, node, new AtomicBoolean()).stream().collect(Collectors.toList()); } - static Enumerable<Object> toEnumerable(PipelineOptions options, BeamRelNode node) { + /** + * Runs the pipeline for {@code node}. Calcite sets {@code cancelFlag} when the JDBC statement is + * cancelled (e.g. with Ctrl+C in the SQL shell), which cancels the pipeline. + */ + static Enumerable<Object> toEnumerable( + PipelineOptions options, BeamRelNode node, AtomicBoolean cancelFlag) { if (node instanceof BeamIOSinkRel) { - return count(options, node); + return count(options, node, cancelFlag); } else if (isLimitQuery(node)) { - return limitCollect(options, node); + return limitCollect(options, node, cancelFlag); } - return Linq4j.asEnumerable(rowToAvaticaAndUnboxValues(collectRows(options, node))); + return Linq4j.asEnumerable(rowToAvaticaAndUnboxValues(collectRows(options, node, cancelFlag))); } private static PipelineResult limitRun( @@ -178,7 +189,8 @@ public class BeamEnumerableConverter extends ConverterImpl implements Enumerable BeamRelNode node, DoFn<Row, Void> doFn, Queue<Row> values, - int limitCount) { + int limitCount, + AtomicBoolean cancelFlag) { options.as(DirectOptions.class).setBlockOnRun(false); Pipeline pipeline = Pipeline.create(options); PCollection<Row> resultCollection = BeamSqlRelUtils.toPCollection(pipeline, node); @@ -198,7 +210,7 @@ public class BeamEnumerableConverter extends ConverterImpl implements Enumerable } try { - if (values.size() >= limitCount) { + if (values.size() >= limitCount || cancelFlag.get()) { result.cancel(); break; } @@ -207,21 +219,78 @@ public class BeamEnumerableConverter extends ConverterImpl implements Enumerable break; } } + checkNotCancelled(cancelFlag); return result; } - private static void runCollector(PipelineOptions options, BeamRelNode node) { + private static void runCollector( + PipelineOptions options, BeamRelNode node, AtomicBoolean cancelFlag) { Pipeline pipeline = Pipeline.create(options); PCollection<Row> resultCollection = BeamSqlRelUtils.toPCollection(pipeline, node); resultCollection.apply(ParDo.of(new Collector())); - PipelineResult result = pipeline.run(); - if (PipelineResult.State.FAILED.equals(result.waitUntilFinish())) { + PipelineResult result = run(pipeline, cancelFlag); + if (PipelineResult.State.FAILED.equals(waitUntilFinish(result, cancelFlag))) { throw new RuntimeException("Pipeline failed for unknown reason"); } } - private static Queue<Row> collectRows(PipelineOptions options, BeamRelNode node) { + /** + * Runs the pipeline. If the DirectRunner would block in {@code run()} until the pipeline + * finishes, waits here instead, where the wait can be cancelled. + */ + private static PipelineResult run(Pipeline pipeline, AtomicBoolean cancelFlag) { + DirectOptions options = pipeline.getOptions().as(DirectOptions.class); + boolean blockOnRun = + options.isBlockOnRun() + && options.getRunner().getCanonicalName().equals(DirectRunner.class.getName()); + options.setBlockOnRun(false); + PipelineResult result = pipeline.run(); + if (blockOnRun) { + waitUntilFinish(result, cancelFlag); + } + return result; + } + + /** + * Waits for the pipeline to finish. If {@code cancelFlag} is set while waiting, cancels the + * pipeline so that the wait returns, and throws. + */ + private static State waitUntilFinish(PipelineResult result, AtomicBoolean cancelFlag) { + Thread canceller = + new Thread( + () -> { + try { + while (!cancelFlag.get()) { + Thread.sleep(100); + } + result.cancel(); + } catch (InterruptedException e) { + // The pipeline finished. + } catch (IOException | RuntimeException e) { + LOG.warn("Failed to cancel the pipeline", e); + } + }, + "BeamSqlPipelineCanceller"); + canceller.setDaemon(true); + canceller.start(); + try { + State state = result.waitUntilFinish(); + checkNotCancelled(cancelFlag); + return state; + } finally { + canceller.interrupt(); + } + } + + private static void checkNotCancelled(AtomicBoolean cancelFlag) { + if (cancelFlag.get()) { + throw new CancellationException("Statement canceled"); + } + } + + private static Queue<Row> collectRows( + PipelineOptions options, BeamRelNode node, AtomicBoolean cancelFlag) { long id = options.getOptionsId(); Queue<Row> values = new ConcurrentLinkedQueue<>(); @@ -233,14 +302,16 @@ public class BeamEnumerableConverter extends ConverterImpl implements Enumerable "collectRowList is only available in direct runner."); Collector.globalValues.put(id, values); - - runCollector(options, node); - - Collector.globalValues.remove(id); + try { + runCollector(options, node, cancelFlag); + } finally { + Collector.globalValues.remove(id); + } return values; } - private static Enumerable<Object> limitCollect(PipelineOptions options, BeamRelNode node) { + private static Enumerable<Object> limitCollect( + PipelineOptions options, BeamRelNode node, AtomicBoolean cancelFlag) { long id = options.getOptionsId(); ConcurrentLinkedQueue<Row> values = new ConcurrentLinkedQueue<>(); @@ -254,8 +325,11 @@ public class BeamEnumerableConverter extends ConverterImpl implements Enumerable int limitCount = getLimitCount(node); Collector.globalValues.put(id, values); - limitRun(options, node, new Collector(), values, limitCount); - Collector.globalValues.remove(id); + try { + limitRun(options, node, new Collector(), values, limitCount, cancelFlag); + } finally { + Collector.globalValues.remove(id); + } // remove extra retrieved values while (values.size() > limitCount) { @@ -378,14 +452,15 @@ public class BeamEnumerableConverter extends ConverterImpl implements Enumerable } } - private static Enumerable<Object> count(PipelineOptions options, BeamRelNode node) { + private static Enumerable<Object> count( + PipelineOptions options, BeamRelNode node, AtomicBoolean cancelFlag) { Pipeline pipeline = Pipeline.create(options); BeamSqlRelUtils.toPCollection(pipeline, node).apply(ParDo.of(new RowCounter())); - PipelineResult result = pipeline.run(); + PipelineResult result = run(pipeline, cancelFlag); long count = 0; if (!containsUnboundedPCollection(pipeline)) { - if (PipelineResult.State.FAILED.equals(result.waitUntilFinish())) { + if (PipelineResult.State.FAILED.equals(waitUntilFinish(result, cancelFlag))) { throw new RuntimeException("Pipeline failed for unknown reason"); } MetricQueryResults metrics = diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/JdbcDriverTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/JdbcDriverTest.java index 4133fb8b070..b022d94777b 100644 --- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/JdbcDriverTest.java +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/JdbcDriverTest.java @@ -22,9 +22,11 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.containsInAnyOrder; import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.instanceOf; import static org.hamcrest.Matchers.nullValue; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import java.sql.Connection; @@ -36,13 +38,21 @@ import java.sql.SQLException; import java.sql.Statement; import java.sql.Timestamp; import java.util.ArrayList; +import java.util.Arrays; import java.util.Calendar; import java.util.List; import java.util.Locale; import java.util.Map; import java.util.TimeZone; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; +import org.apache.beam.sdk.extensions.sql.meta.catalog.InMemoryCatalogManager; import org.apache.beam.sdk.extensions.sql.meta.provider.ReadOnlyTableProvider; +import org.apache.beam.sdk.extensions.sql.meta.provider.seqgen.GenerateSequenceTableProvider; import org.apache.beam.sdk.extensions.sql.meta.provider.test.TestBoundedTable; import org.apache.beam.sdk.extensions.sql.meta.provider.test.TestTableProvider; import org.apache.beam.sdk.extensions.sql.meta.provider.test.TestUnboundedTable; @@ -449,6 +459,47 @@ public class JdbcDriverTest { assertFalse(resultSet2.next()); } + /** Statement.cancel() (e.g. Ctrl+C in the SQL shell) must stop a never-ending pipeline. */ + @Test + public void testCancelUnboundedQuery() throws Exception { + TestTableProvider testTables = new TestTableProvider(); + InMemoryCatalogManager catalogManager = new InMemoryCatalogManager(); + catalogManager.registerTableProvider(testTables); + catalogManager.registerTableProvider(new GenerateSequenceTableProvider()); + Connection connection = JdbcDriver.connect(catalogManager, PipelineOptionsFactory.create()); + Statement ddl = connection.createStatement(); + ddl.execute( + "CREATE EXTERNAL TABLE seq (sequence BIGINT, event_time TIMESTAMP) TYPE 'sequence'"); + ddl.execute("CREATE EXTERNAL TABLE sink (sequence BIGINT) TYPE 'test'"); + + ExecutorService executor = Executors.newCachedThreadPool(); + try { + Statement insert = connection.createStatement(); + Future<?> insertRun = + executor.submit(() -> insert.execute("INSERT INTO sink SELECT sequence FROM seq")); + while (testTables.tableRows("sink").isEmpty() && !insertRun.isDone()) { + Thread.sleep(100); + } + assertCancels(insert, insertRun); + + for (String select : + Arrays.asList("SELECT sequence FROM seq", "SELECT sequence FROM seq LIMIT 1000000")) { + Statement statement = connection.createStatement(); + assertCancels(statement, executor.submit(() -> statement.executeQuery(select))); + } + } finally { + executor.shutdownNow(); + } + } + + private static void assertCancels(Statement statement, Future<?> run) throws Exception { + statement.cancel(); + ExecutionException e = + assertThrows(ExecutionException.class, () -> run.get(30, TimeUnit.SECONDS)); + assertThat(e.getCause(), instanceOf(SQLException.class)); + assertThat(e.getCause().getMessage(), containsString("Statement canceled")); + } + private List<List<Object>> readResultSet(ResultSet result) throws Exception { List<List<Object>> results = new ArrayList<>(); diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamEnumerableConverterTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamEnumerableConverterTest.java index 697643d4715..42d138e2f38 100644 --- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamEnumerableConverterTest.java +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamEnumerableConverterTest.java @@ -23,6 +23,7 @@ import static org.junit.Assert.assertTrue; import java.math.BigDecimal; import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils; import org.apache.beam.sdk.extensions.sql.meta.SchemaBaseBeamTable; import org.apache.beam.sdk.options.PipelineOptions; @@ -77,7 +78,8 @@ public class BeamEnumerableConverterTest { ImmutableList.of(ImmutableList.of(rexBuilder.makeBigintLiteral(BigDecimal.ZERO))); BeamRelNode node = new BeamValuesRel(cluster, type, tuples, RelTraitSet.createEmpty()); - Enumerable<Object> enumerable = BeamEnumerableConverter.toEnumerable(options, node); + Enumerable<Object> enumerable = + BeamEnumerableConverter.toEnumerable(options, node, new AtomicBoolean()); Enumerator<Object> enumerator = enumerable.enumerator(); assertTrue(enumerator.moveNext()); @@ -97,7 +99,8 @@ public class BeamEnumerableConverterTest { rexBuilder.makeBigintLiteral(BigDecimal.ONE))); BeamRelNode node = new BeamValuesRel(cluster, type, tuples, RelTraitSet.createEmpty()); - Enumerable<Object> enumerable = BeamEnumerableConverter.toEnumerable(options, node); + Enumerable<Object> enumerable = + BeamEnumerableConverter.toEnumerable(options, node, new AtomicBoolean()); Enumerator<Object> enumerator = enumerable.enumerator(); assertTrue(enumerator.moveNext()); @@ -173,7 +176,8 @@ public class BeamEnumerableConverterTest { new FakeTable(), null); - Enumerable<Object> enumerable = BeamEnumerableConverter.toEnumerable(options, node); + Enumerable<Object> enumerable = + BeamEnumerableConverter.toEnumerable(options, node, new AtomicBoolean()); Enumerator<Object> enumerator = enumerable.enumerator(); assertTrue(enumerator.moveNext()); @@ -225,7 +229,8 @@ public class BeamEnumerableConverterTest { rexBuilder.makeNullLiteral(CalciteUtils.toRelDataType(TYPE_FACTORY, fieldType)))); BeamRelNode node = new BeamValuesRel(cluster, type, tuples, RelTraitSet.createEmpty()); - Enumerable<Object> enumerable = BeamEnumerableConverter.toEnumerable(options, node); + Enumerable<Object> enumerable = + BeamEnumerableConverter.toEnumerable(options, node, new AtomicBoolean()); Enumerator<Object> enumerator = enumerable.enumerator(); assertTrue(enumerator.moveNext());
