From 57eb6d0078a498e2f792d76b685c113a6a52ea8e Mon Sep 17 00:00:00 2001
From: Yiru Tang <yiru@google.com>
Date: Mon, 25 Jul 2022 13:16:06 -0700
Subject: [PATCH] feat: Return explicit StreamWriterClosedException (#1709)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* feat: Return explicit StreamWriterClosedException

* fix format

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
---
 .../google/cloud/bigquery/storage/v1/Exceptions.java |  6 ++++++
 .../cloud/bigquery/storage/v1/StreamWriter.java      |  5 +++--
 .../cloud/bigquery/storage/v1/StreamWriterTest.java  | 12 ++++++++++++
 3 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/Exceptions.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/Exceptions.java
index ba1051cd67..ee1958f827 100644
--- a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/Exceptions.java
+++ b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/Exceptions.java
@@ -83,6 +83,12 @@ protected StreamFinalizedException(Status grpcStatus, String name) {
     }
   }
 
+  public static final class StreamWriterClosedException extends StorageException {
+    protected StreamWriterClosedException(Status grpcStatus, String name) {
+      super(grpcStatus, name, null, null, ImmutableMap.of());
+    }
+  }
+
   /**
    * 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.
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java
index 2c4ecdd75f..45b8e2593f 100644
--- a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java
+++ b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java
@@ -309,9 +309,10 @@ private ApiFuture<AppendRowsResponse> appendInternal(AppendRowsRequest message)
     try {
       if (userClosed) {
         requestWrapper.appendResult.setException(
-            new StatusRuntimeException(
+            new Exceptions.StreamWriterClosedException(
                 Status.fromCode(Status.Code.FAILED_PRECONDITION)
-                    .withDescription("Connection is already closed")));
+                    .withDescription("Connection is already closed"),
+                streamName));
         return requestWrapper.appendResult;
       }
       // Check if queue is going to be full before adding the request.
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java
index 9078e581bd..a4330be263 100644
--- a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java
+++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java
@@ -653,4 +653,16 @@ public void testWriterClosedStream() throws Exception {
       TimeUnit.SECONDS.sleep(1);
     }
   }
+
+  @Test
+  public void testWriterException() throws Exception {
+    StreamWriter writer = getTestStreamWriter();
+    writer.close();
+    ApiFuture<AppendRowsResponse> appendFuture1 = sendTestMessage(writer, new String[] {"A"}, 0);
+    Exceptions.StreamWriterClosedException actualError =
+        assertFutureException(Exceptions.StreamWriterClosedException.class, appendFuture1);
+    // The basic StatusRuntimeException API is not changed.
+    assertEquals(Status.Code.FAILED_PRECONDITION, actualError.getStatus().getCode());
+    assertTrue(actualError.getStatus().getDescription().contains("Connection is already closed"));
+  }
 }