Skip to content

Allow canceling SQL streaming pipeline #40666

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ahmedabu98 merged 2 commits into from
Oct 20, 2026
Merged

Allow canceling SQL streaming pipeline #40666

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9a56e45
enable canceling SQL streaming pipeline
ahmedabu98 Oct 19, 2026
c61c5cd
Merge branch 'master' of https://github.com/ahmedabu98/beam into sql-…
ahmedabu98 Oct 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
123 changes: 99 additions & 24 deletions ...ql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamEnumerableConverter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,11 +28,14 @@
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;
Expand All @@ -58,6 +61,7 @@
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;
Expand Down Expand Up @@ -111,16 +115,18 @@ public Result implement(EnumerableRelImplementor implementor, Prefer prefer) {
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);
}
Expand Down Expand Up @@ -161,24 +167,30 @@ static List<Row> toRowList(PipelineOptions options, BeamRelNode node) {
} 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(
PipelineOptions options,
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);
Expand All @@ -198,7 +210,7 @@ private static PipelineResult limitRun(
}

try {
if (values.size() >= limitCount) {
if (values.size() >= limitCount || cancelFlag.get()) {
result.cancel();
break;
}
Expand All @@ -207,21 +219,78 @@ private static PipelineResult limitRun(
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<>();

Expand All @@ -233,14 +302,16 @@ private static Queue<Row> collectRows(PipelineOptions options, BeamRelNode node)
"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<>();

Expand All @@ -254,8 +325,11 @@ private static Enumerable<Object> limitCollect(PipelineOptions options, BeamRelN
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) {
Expand Down Expand Up @@ -378,14 +452,15 @@ private static Object fieldToAvatica(Schema.FieldType type, Object beamValue) {
}
}

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 =
Expand Down
51 changes: 51 additions & 0 deletions .../extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/JdbcDriverTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,11 @@
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;
Expand All @@ -36,13 +38,21 @@
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;
Expand Down Expand Up @@ -449,6 +459,47 @@ public void testInternalConnect_unbounded_limit() throws Exception {
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<>();

Expand Down
13 changes: 9 additions & 4 deletions ...rc/test/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamEnumerableConverterTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@

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;
Expand Down Expand Up @@ -77,7 +78,8 @@ public void testToEnumerable_collectSingle() {
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());
Expand All @@ -97,7 +99,8 @@ public void testToEnumerable_collectMultiple() {
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());
Expand Down Expand Up @@ -173,7 +176,8 @@ public void testToEnumerable_count() {
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());
Expand Down Expand Up @@ -225,7 +229,8 @@ public void testToEnumerable_collectNullValue() {
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());
Expand Down
Loading