fix: change customer StorageException from RuntimeException to StatusRuntimeException

google-cloud-bigquerystorage/clirr-ignored-differences.xml

@@ -0,0 +1,14 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!-- see http://www.mojohaus.org/clirr-maven-plugin/examples/ignored-differences.html -->
+<differences>
+  <difference>
+    <differenceType>7004</differenceType>
+    <className>com/google/cloud/bigquery/storage/v1/Exceptions$SchemaMismatchedException</className>
+    <method>Exceptions$SchemaMismatchedException(java.lang.String, java.lang.String, java.lang.Throwable)</method>
+  </difference>
+  <difference>
+    <differenceType>7004</differenceType>
+    <className>com/google/cloud/bigquery/storage/v1/Exceptions$StreamFinalizedException</className>
+    <method>Exceptions$StreamFinalizedException(java.lang.String, java.lang.String, java.lang.Throwable)</method>
+  </difference>
+</differences>

google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/Exceptions.java

@@ -19,6 +19,9 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.protobuf.Any;
 import com.google.protobuf.InvalidProtocolBufferException;
+import io.grpc.Metadata;
+import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
 import io.grpc.protobuf.StatusProto;
 import javax.annotation.Nullable;
 
@@ -30,21 +33,23 @@ public WriterClosedException(String streamName) {
     }
   }
   /** Main Storage Exception. Might contain map of streams to errors for that stream. */
-  public static class StorageException extends RuntimeException {
+  public static class StorageException extends StatusRuntimeException {
 
     private final ImmutableMap<String, GrpcStatusCode> errors;
     private final String streamName;
 
     private StorageException() {
-      this(null, null, null, ImmutableMap.of());
+      this(null, null, null, null, null, ImmutableMap.of());
     }
 
     private StorageException(
+        @Nullable Status grpcStatus,
+        @Nullable Metadata metadata,
         @Nullable String message,
         @Nullable Throwable cause,
         @Nullable String streamName,
         ImmutableMap<String, GrpcStatusCode> errors) {
-      super(message, cause);
+      super(grpcStatus, metadata);
       this.streamName = streamName;
       this.errors = errors;
     }
@@ -60,8 +65,9 @@ public String getStreamName() {
 
   /** Stream has already been finalized. */
   public static final class StreamFinalizedException extends StorageException {
-    protected StreamFinalizedException(String name, String message, Throwable cause) {
-      super(message, cause, name, ImmutableMap.of());
+    protected StreamFinalizedException(
+        Status invalidArgument, Metadata metadata, String name, String message, Throwable cause) {
+      super(invalidArgument, metadata, message, cause, name, ImmutableMap.of());
     }
   }
 
@@ -70,8 +76,9 @@ protected StreamFinalizedException(String name, String message, Throwable cause)
    * This can be resolved by updating the table's schema with the message schema.
    */
   public static final class SchemaMismatchedException extends StorageException {
-    protected SchemaMismatchedException(String name, String message, Throwable cause) {
-      super(message, cause, name, ImmutableMap.of());
+    protected SchemaMismatchedException(
+        Status invalidArgument, Metadata metadata, String name, String message, Throwable cause) {
+      super(invalidArgument, metadata, message, cause, name, ImmutableMap.of());
     }
   }
 
@@ -98,15 +105,19 @@ private static StorageError toStorageError(com.google.rpc.Status rpcStatus) {
   public static StorageException toStorageException(
       com.google.rpc.Status rpcStatus, Throwable exception) {
     StorageError error = toStorageError(rpcStatus);
+    Status grpcStatus =
+        Status.fromCodeValue(rpcStatus.getCode()).withDescription(rpcStatus.getMessage());
     if (error == null) {
       return null;
     }
     switch (error.getCode()) {
       case STREAM_FINALIZED:
-        return new StreamFinalizedException(error.getEntity(), error.getErrorMessage(), exception);
+        return new StreamFinalizedException(
+            grpcStatus, null, error.getEntity(), error.getErrorMessage(), exception);
 
       case SCHEMA_MISMATCH_EXTRA_FIELDS:
-        return new SchemaMismatchedException(error.getEntity(), error.getErrorMessage(), exception);
+        return new SchemaMismatchedException(
+            grpcStatus, null, error.getEntity(), error.getErrorMessage(), exception);
 
       default:
         return null;

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java

@@ -26,6 +26,7 @@
 import com.google.api.gax.grpc.testing.MockServiceHelper;
 import com.google.api.gax.rpc.ApiException;
 import com.google.api.gax.rpc.StatusCode.Code;
+import com.google.api.gax.rpc.UnknownException;
 import com.google.cloud.bigquery.storage.test.Test.FooType;
 import com.google.cloud.bigquery.storage.v1.StorageError.StorageErrorCode;
 import com.google.common.base.Strings;
@@ -343,6 +344,20 @@ public void testAppendFailedSchemaError() throws Exception {
     writer.close();
   }
 
+  @Test
+  public void testAppendFailRandomException() throws Exception {
+    StreamWriter writer = getTestStreamWriter();
+    // Trigger a non-StatusRuntimeException for append operation
+    IllegalArgumentException illegalArgumentException =
+        new IllegalArgumentException("Illegal argument");
+    testBigQueryWrite.addException(illegalArgumentException);
+    ApiFuture<AppendRowsResponse> appendFuture1 = sendTestMessage(writer, new String[] {"A"});
+    UnknownException actualError = assertFutureException(UnknownException.class, appendFuture1);
+    assertEquals(Code.UNKNOWN, actualError.getStatusCode().getCode());
+
+    writer.close();
+  }
+
   @Test
   public void longIdleBetweenAppends() throws Exception {
     StreamWriter writer = getTestStreamWriter();

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteManualClientTest.java

@@ -19,6 +19,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 
 import com.google.api.core.ApiFuture;
@@ -27,11 +28,15 @@
 import com.google.cloud.bigquery.Schema;
 import com.google.cloud.bigquery.storage.test.Test.*;
 import com.google.cloud.bigquery.storage.v1.*;
+import com.google.cloud.bigquery.storage.v1.Exceptions.SchemaMismatchedException;
+import com.google.cloud.bigquery.storage.v1.Exceptions.StreamFinalizedException;
 import com.google.cloud.bigquery.testing.RemoteBigQueryHelper;
 import com.google.common.collect.ImmutableList;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.Descriptors;
 import com.google.protobuf.Descriptors.DescriptorValidationException;
+import io.grpc.Status;
+import io.grpc.Status.Code;
 import java.io.IOException;
 import java.math.BigDecimal;
 import java.util.*;
@@ -835,8 +840,10 @@ public void testStreamSchemaMisMatchError() throws IOException, InterruptedExcep
         Assert.fail("Should fail");
       } catch (ExecutionException e) {
         assertEquals(Exceptions.SchemaMismatchedException.class, e.getCause().getClass());
-        assertThat(e.getCause().getMessage())
-            .contains("Schema mismatch due to extra fields in user schema");
+        Exceptions.SchemaMismatchedException actualError = (SchemaMismatchedException) e.getCause();
+        assertNotNull(actualError.getStreamName());
+        // This verifies that the Beam connector can consume this custom exception's grpc StatusCode
+        assertEquals(Code.INVALID_ARGUMENT, Status.fromThrowable(e.getCause()).getCode());
       }
     }
   }
@@ -866,7 +873,10 @@ public void testStreamFinalizedError()
         Assert.fail("Should fail");
       } catch (ExecutionException e) {
         assertEquals(Exceptions.StreamFinalizedException.class, e.getCause().getClass());
-        assertThat(e.getCause().getMessage()).contains("Stream is finalized");
+        Exceptions.StreamFinalizedException actualError = (StreamFinalizedException) e.getCause();
+        assertNotNull(actualError.getStreamName());
+        // This verifies that the Beam connector can consume this custom exception's grpc StatusCode
+        assertEquals(Code.INVALID_ARGUMENT, Status.fromThrowable(e.getCause()).getCode());
       }
     }
   }