Merge pull request #24 from vinted/fix/correctly-resolve-exception

fix: correctly resolve nested exception
vinted · Mar 19, 2024 · 79df424 · 79df424
2 parents 8065b79 + 9f6f8b8
commit 79df424
Showing 1 changed file with 18 additions and 1 deletion.
diff --git a/src/main/java/com/vinted/flink/bigquery/sink/async/AsyncBigQuerySinkWriter.java b/src/main/java/com/vinted/flink/bigquery/sink/async/AsyncBigQuerySinkWriter.java
@@ -4,11 +4,14 @@
 import com.google.cloud.bigquery.storage.v1.Exceptions;
 import com.google.cloud.bigquery.storage.v1.ProtoRows;
 import com.google.cloud.bigquery.storage.v1.StreamWriter;
+import com.google.common.base.Preconditions;
 import com.vinted.flink.bigquery.metric.AsyncBigQueryStreamMetrics;
 import com.vinted.flink.bigquery.metric.BigQueryStreamMetrics;
 import com.vinted.flink.bigquery.model.Rows;
 import com.vinted.flink.bigquery.sink.ExecutorProvider;
 import io.grpc.Status;
+import io.grpc.StatusException;
+import io.grpc.StatusRuntimeException;
 import org.apache.flink.api.connector.sink2.Sink;
 import org.apache.flink.connector.base.sink.writer.AsyncSinkWriter;
 import org.apache.flink.connector.base.sink.writer.BufferedRequestState;
@@ -152,7 +155,7 @@ protected void submitRequestEntries(List<StreamRequest> list, Consumer<List<Stre
                                         throw new AsyncWriterException(traceId, status.getCode(), t);
                                     }
                                 case UNKNOWN:
-                                    if (status.getCause() instanceof Exceptions.MaximumRequestCallbackWaitTimeExceededException) {
+                                    if (isMaximumRequestCallbackWaitTimeExceededException(t)) {
                                         logger.info("Trace-id {} request timed out: {}", traceId, t.getMessage());
                                         Optional.ofNullable(this.metrics.get(request.getStream()))
                                                 .ifPresent(AsyncBigQueryStreamMetrics::incrementTimeoutCount);
@@ -194,6 +197,20 @@ protected void submitRequestEntries(List<StreamRequest> list, Consumer<List<Stre
 
     }
 
+    private boolean isMaximumRequestCallbackWaitTimeExceededException(Throwable t) {
+        if (t == null) {
+            return false;
+        }
+
+        for (Throwable cause = t; cause != null; cause = cause.getCause()) {
+            if (cause instanceof Exceptions.MaximumRequestCallbackWaitTimeExceededException) {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
     private List<StreamRequest> retry(Throwable t, String traceId, StreamRequest request) {
         var status = Status.fromThrowable(t);
         request.setRetries(request.getRetries() - 1);