Handle error on source split in BigQueryStorageStreamSource (#33540)

apache · Jan 22, 2025 · 8e75484 · 8e75484
1 parent 0d3cfa9
commit 8e75484
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/.../java/org/apache/beam/runners/dataflow/worker/util/common/worker/WorkProgressUpdater.java b/.../java/org/apache/beam/runners/dataflow/worker/util/common/worker/WorkProgressUpdater.java
@@ -309,7 +309,7 @@ private void reportProgress() {
       LOG.info("Cancelling workitem execution: {}", workString(), e);
       worker.abort();
     } catch (Throwable e) {
-      LOG.warn("Error reporting workitem progress update to Dataflow service: ", e);
+      LOG.error("Error reporting workitem progress update to Dataflow service: ", e);
     }
   }
 

diff --git a/...atform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageStreamSource.java b/...atform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageStreamSource.java
@@ -368,8 +368,13 @@ public synchronized BigQueryStorageStreamSource<T> getCurrentSource() {
               .setName(source.readStream.getName())
               .setFraction((float) fraction)
               .build();
-
-      SplitReadStreamResponse splitResponse = storageClient.splitReadStream(splitRequest);
+      SplitReadStreamResponse splitResponse;
+      try {
+        splitResponse = storageClient.splitReadStream(splitRequest);
+      } catch (Exception e) {
+        LOG.warn("Skip split read stream due to failed request: ", e);
+        return null;
+      }
       if (!splitResponse.hasPrimaryStream() || !splitResponse.hasRemainderStream()) {
         // No more splits are possible!
         impossibleSplitPointCalls.inc();