feat: add support for StorageError

README.md

@@ -49,20 +49,20 @@ If you are using Maven without BOM, add this to your dependencies:
 If you are using Gradle 5.x or later, add this to your dependencies
 
 ```Groovy
-implementation platform('com.google.cloud:libraries-bom:23.1.0')
+implementation platform('com.google.cloud:libraries-bom:24.0.0')
 
 implementation 'com.google.cloud:google-cloud-bigquerystorage'
 ```
 If you are using Gradle without BOM, add this to your dependencies
 
 ```Groovy
-implementation 'com.google.cloud:google-cloud-bigquerystorage:2.4.0'
+implementation 'com.google.cloud:google-cloud-bigquerystorage:2.5.0'
 ```
 
 If you are using SBT, add this to your dependencies
 
 ```Scala
-libraryDependencies += "com.google.cloud" % "google-cloud-bigquerystorage" % "2.4.0"
+libraryDependencies += "com.google.cloud" % "google-cloud-bigquerystorage" % "2.5.0"
 ```
 
 ## Authentication

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

@@ -0,0 +1,145 @@
+/*
+ * Copyright 2021 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.cloud.bigquery.storage.v1;
+
+import com.google.api.gax.grpc.GrpcStatusCode;
+import com.google.common.collect.ImmutableMap;
+import com.google.protobuf.Any;
+import com.google.protobuf.InvalidProtocolBufferException;
+import javax.annotation.Nullable;
+
+/** Exceptions for Storage Client Libraries. */
+public final class Exceptions {
+  /** Main Storage Exception. Might contain map of streams to errors for that stream. */
+  public static class StorageException extends RuntimeException {
+
+    protected StorageException() {}
+
+    protected StorageException(String message) {
+      super(message);
+    }
+
+    protected StorageException(String message, Throwable cause) {
+      super(message, cause);
+    }
+
+    protected StorageException(Throwable cause) {
+      super(cause);
+    }
+
+    private ImmutableMap<String, GrpcStatusCode> errors;
+
+    public ImmutableMap<String, GrpcStatusCode> getErrors() {
+      return errors;
+    }
+
+    public void setErrors(ImmutableMap<String, GrpcStatusCode> value) {
+      this.errors = value;
+    }
+
+    private String streamName;
+
+    public String getStreamName() {
+      return streamName;
+    }
+
+    protected void setStreamName(String value) {
+      this.streamName = value;
+    }
+  }
+
+  /** Stream has already been finalized. */
+  public static class StreamFinalizedException extends StorageException {
+    protected StreamFinalizedException() {}
+
+    protected StreamFinalizedException(String name, String message, Throwable cause) {
+      super(message, cause);
+      setStreamName(name);
+    }
+  }
+
+  /**
+   * There was a schema mismatch due to bigquery table with fewer fields than the input message.
+   * This can be resolved by updating the table's schema with the message schema.
+   */
+  public static class SchemaMismatchedException extends StorageException {
+    protected SchemaMismatchedException() {}
+
+    protected SchemaMismatchedException(String name, String message, Throwable cause) {
+      super(message, cause);
+      setStreamName(name);
+    }
+  }
+
+  private static StorageError toStorageError(com.google.rpc.Status rpcStatus) {
+    for (Any detail : rpcStatus.getDetailsList()) {
+      if (detail.is(StorageError.class)) {
+        try {
+          return detail.unpack(StorageError.class);
+        } catch (InvalidProtocolBufferException protoException) {
+          throw new IllegalStateException(protoException);
+        }
+      }
+    }
+    return null;
+  }
+
+  /* Converts a c.g.rpc.Status into a StorageException, if possible.
+   * Examines the embedded StorageError, and potentially returns a StreamFinalizedException or
+   * SchemaMismatchedException (both derive from StorageException).
+   * If there is no StorageError, or the StorageError is a different error it will return NULL.
+   */
+  @Nullable
+  public static StorageException toStorageException(
+      com.google.rpc.Status rpcStatus, Throwable exception) {
+    StorageError error = toStorageError(rpcStatus);
+    if (error == null) {
+      return null;
+    }
+    switch (error.getCode()) {
+      case STREAM_FINALIZED:
+        return new StreamFinalizedException(error.getEntity(), error.getErrorMessage(), exception);
+
+      case SCHEMA_MISMATCH_EXTRA_FIELDS:
+        return new SchemaMismatchedException(error.getEntity(), error.getErrorMessage(), exception);
+
+      default:
+        return null;
+    }
+  }
+
+  /* Converts a Throwable into a StorageException, if possible.
+   * Examines the embedded error message, and potentially returns a StreamFinalizedException or
+   * SchemaMismatchedException (both derive from StorageException).
+   * If there is no StorageError, or the StorageError is a different error it will return NULL.
+   */
+  @Nullable
+  public static StorageException toStorageException(Throwable exception) {
+    String errorMsg = exception.getMessage();
+    if (errorMsg == null) {
+      return null;
+    }
+    if (errorMsg.toLowerCase().contains("finazlied")) {
+      return new StreamFinalizedException(null, errorMsg, exception);
+    }
+    if (errorMsg.toLowerCase().contains("mismatch")) {
+      return new SchemaMismatchedException(null, errorMsg, exception);
+    }
+    return null;
+  }
+
+  private Exceptions() {}
+}

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

@@ -21,6 +21,7 @@
 import com.google.api.gax.rpc.FixedHeaderProvider;
 import com.google.api.gax.rpc.TransportChannelProvider;
 import com.google.cloud.bigquery.storage.v1.AppendRowsRequest.ProtoData;
+import com.google.cloud.bigquery.storage.v1.Exceptions.StorageException;
 import com.google.cloud.bigquery.storage.v1.StreamConnection.DoneCallback;
 import com.google.cloud.bigquery.storage.v1.StreamConnection.RequestCallback;
 import com.google.common.base.Preconditions;
@@ -460,11 +461,17 @@ private void requestCallback(AppendRowsResponse response) {
       this.lock.unlock();
     }
     if (response.hasError()) {
-      StatusRuntimeException exception =
-          new StatusRuntimeException(
-              Status.fromCodeValue(response.getError().getCode())
-                  .withDescription(response.getError().getMessage()));
-      requestWrapper.appendResult.setException(exception);
+      Exceptions.StorageException storageException =
+          Exceptions.toStorageException(response.getError(), null);
+      if (storageException != null) {
+        requestWrapper.appendResult.setException(storageException);
+      } else {
+        StatusRuntimeException exception =
+            new StatusRuntimeException(
+                Status.fromCodeValue(response.getError().getCode())
+                    .withDescription(response.getError().getMessage()));
+        requestWrapper.appendResult.setException(exception);
+      }
     } else {
       requestWrapper.appendResult.set(response);
     }
@@ -482,6 +489,10 @@ private void doneCallback(Throwable finalStatus) {
     } finally {
       this.lock.unlock();
     }
+    Exceptions.StorageException storageException = Exceptions.toStorageException(finalStatus);
+    if (storageException != null) {
+      throw new StorageException(finalStatus);
+    }
   }
 
   @GuardedBy("lock")

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

@@ -27,7 +27,9 @@
 import com.google.api.gax.rpc.ApiException;
 import com.google.api.gax.rpc.StatusCode.Code;
 import com.google.cloud.bigquery.storage.test.Test.FooType;
+import com.google.cloud.bigquery.storage.v1.StorageError.StorageErrorCode;
 import com.google.common.base.Strings;
+import com.google.protobuf.Any;
 import com.google.protobuf.DescriptorProtos;
 import com.google.protobuf.Int64Value;
 import io.grpc.Status;
@@ -304,6 +306,38 @@ public void testAppendSuccessAndInStreamError() throws Exception {
     writer.close();
   }
 
+  @Test
+  public void testAppendFailedSchemaError() throws Exception {
+    StreamWriter writer = getTestStreamWriter();
+
+    StorageError storageError =
+        StorageError.newBuilder()
+            .setCode(StorageErrorCode.SCHEMA_MISMATCH_EXTRA_FIELDS)
+            .setEntity("foobar")
+            .build();
+    com.google.rpc.Status statusProto =
+        com.google.rpc.Status.newBuilder()
+            .setCode(Code.INVALID_ARGUMENT.getHttpStatusCode())
+            .addDetails(Any.pack(storageError))
+            .build();
+
+    testBigQueryWrite.addResponse(createAppendResponse(0));
+    testBigQueryWrite.addResponse(AppendRowsResponse.newBuilder().setError(statusProto).build());
+    testBigQueryWrite.addResponse(createAppendResponse(1));
+
+    ApiFuture<AppendRowsResponse> appendFuture1 = sendTestMessage(writer, new String[] {"A"});
+    ApiFuture<AppendRowsResponse> appendFuture2 = sendTestMessage(writer, new String[] {"B"});
+    ApiFuture<AppendRowsResponse> appendFuture3 = sendTestMessage(writer, new String[] {"C"});
+
+    assertEquals(0, appendFuture1.get().getAppendResult().getOffset().getValue());
+    Exceptions.SchemaMismatchedException actualError =
+        assertFutureException(Exceptions.SchemaMismatchedException.class, appendFuture2);
+    assertEquals("foobar", actualError.getStreamName());
+    assertEquals(1, appendFuture3.get().getAppendResult().getOffset().getValue());
+
+    writer.close();
+  }
+
   @Test
   public void longIdleBetweenAppends() throws Exception {
     StreamWriter writer = getTestStreamWriter();