This is an automated email from the ASF dual-hosted git repository.
ahmedabu98 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new a412e69657e enable canceling SQL streaming pipeline (#40666)
a412e69657e is described below
commit a412e69657e6528ef893003345dfba37909535fd
Author: Ahmed Abualsaud <[email protected]>
AuthorDate: Tue Oct 20 10:59:59 2026 -0700
enable canceling SQL streaming pipeline (#40666)
---
.../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());
