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 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 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")); + } }