You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by ah...@apache.org on 2026/10/19 19:10:33 UTC
(beam) branch sql-shell-cancel created (now 9a56e458381)
This is an automated email from the ASF dual-hosted git repository.
ahmedabu98 pushed a change to branch sql-shell-cancel
in repository https://gitbox.apache.org/repos/asf/beam.git
at 9a56e458381 enable canceling SQL streaming pipeline
This branch includes the following new commits:
new 9a56e458381 enable canceling SQL streaming pipeline
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(beam) 01/01: enable canceling SQL streaming pipeline
Posted by ah...@apache.org.
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 <a....@gmail.com>
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());