fix: use resource type to identify type of error

google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseNotFoundException.java

@@ -16,19 +16,24 @@
 
 package com.google.cloud.spanner;
 
+import com.google.cloud.spanner.SpannerException.ResourceNotFoundException;
+import com.google.rpc.ResourceInfo;
 import javax.annotation.Nullable;
 
 /**
  * Exception thrown by Cloud Spanner when an operation detects that the database that is being used
  * no longer exists. This type of error has its own subclass as it is a condition that should cause
  * the client library to stop trying to send RPCs to the backend until the user has taken action.
  */
-public class DatabaseNotFoundException extends SpannerException {
+public class DatabaseNotFoundException extends ResourceNotFoundException {
   private static final long serialVersionUID = -6395746612598975751L;
 
   /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */
   DatabaseNotFoundException(
-      DoNotConstructDirectly token, @Nullable String message, @Nullable Throwable cause) {
-    super(token, ErrorCode.NOT_FOUND, false, message, cause);
+      DoNotConstructDirectly token,
+      @Nullable String message,
+      ResourceInfo resourceInfo,
+      @Nullable Throwable cause) {
+    super(token, message, resourceInfo, cause);
   }
 }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionNotFoundException.java

@@ -16,19 +16,24 @@
 
 package com.google.cloud.spanner;
 
+import com.google.cloud.spanner.SpannerException.ResourceNotFoundException;
+import com.google.rpc.ResourceInfo;
 import javax.annotation.Nullable;
 
 /**
  * Exception thrown by Cloud Spanner when an operation detects that the session that is being used
  * is no longer valid. This type of error has its own subclass as it is a condition that should
  * normally be hidden from the user, and the client library should try to fix this internally.
  */
-public class SessionNotFoundException extends SpannerException {
+public class SessionNotFoundException extends ResourceNotFoundException {
   private static final long serialVersionUID = -6395746612598975751L;
 
   /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */
   SessionNotFoundException(
-      DoNotConstructDirectly token, @Nullable String message, @Nullable Throwable cause) {
-    super(token, ErrorCode.NOT_FOUND, false, message, cause);
+      DoNotConstructDirectly token,
+      @Nullable String message,
+      ResourceInfo resourceInfo,
+      @Nullable Throwable cause) {
+    super(token, message, resourceInfo, cause);
   }
 }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerException.java

@@ -19,6 +19,7 @@
 import com.google.cloud.grpc.BaseGrpcServiceException;
 import com.google.common.base.Preconditions;
 import com.google.protobuf.util.Durations;
+import com.google.rpc.ResourceInfo;
 import com.google.rpc.RetryInfo;
 import io.grpc.Metadata;
 import io.grpc.Status;
@@ -27,6 +28,24 @@
 
 /** Base exception type for all exceptions produced by the Cloud Spanner service. */
 public class SpannerException extends BaseGrpcServiceException {
+  /** Base exception type for NOT_FOUND exceptions for known resource types. */
+  public abstract static class ResourceNotFoundException extends SpannerException {
+    private final ResourceInfo resourceInfo;
+
+    ResourceNotFoundException(
+        DoNotConstructDirectly token,
+        @Nullable String message,
+        ResourceInfo resourceInfo,
+        @Nullable Throwable cause) {
+      super(token, ErrorCode.NOT_FOUND, /* retryable */ false, message, cause);
+      this.resourceInfo = resourceInfo;
+    }
+
+    public String getResourceName() {
+      return resourceInfo.getResourceName();
+    }
+  }
+
   private static final long serialVersionUID = 20150916L;
   private static final Metadata.Key<RetryInfo> KEY_RETRY_INFO =
       ProtoUtils.keyForProto(RetryInfo.getDefaultInstance());

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java

@@ -21,12 +21,14 @@
 import com.google.cloud.spanner.SpannerException.DoNotConstructDirectly;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Predicate;
+import com.google.rpc.ResourceInfo;
 import io.grpc.Context;
+import io.grpc.Metadata;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
+import io.grpc.protobuf.ProtoUtils;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.TimeoutException;
-import java.util.regex.Pattern;
 import javax.annotation.Nullable;
 import javax.net.ssl.SSLHandshakeException;
 
@@ -37,13 +39,11 @@
  * ErrorCode#ABORTED} are always represented by {@link AbortedException}.
  */
 public final class SpannerExceptionFactory {
-  static final String DATABASE_NOT_FOUND_MSG =
-      "Database not found: projects/.*/instances/.*/databases/.*\n"
-          + "resource_type: \"type.googleapis.com/google.spanner.admin.database.v1.Database\"\n"
-          + "resource_name: \"projects/.*/instances/.*/databases/.*\"\n"
-          + "description: \"Database does not exist.\"\n";
-  private static final Pattern DATABASE_NOT_FOUND_MSG_PATTERN =
-      Pattern.compile(".*" + DATABASE_NOT_FOUND_MSG + ".*");
+  static final String SESSION_RESOURCE_TYPE = "type.googleapis.com/google.spanner.v1.Session";
+  static final String DATABASE_RESOURCE_TYPE =
+      "type.googleapis.com/google.spanner.admin.database.v1.Database";
+  private static final Metadata.Key<ResourceInfo> KEY_RESOURCE_INFO =
+      ProtoUtils.keyForProto(ResourceInfo.getDefaultInstance());
 
   public static SpannerException newSpannerException(ErrorCode code, @Nullable String message) {
     return newSpannerException(code, message, null);
@@ -175,6 +175,16 @@ private static String formatMessage(ErrorCode code, @Nullable String message) {
     return message.startsWith(code.toString()) ? message : code + ": " + message;
   }
 
+  private static ResourceInfo extractResourceInfo(Throwable cause) {
+    if (cause != null) {
+      Metadata trailers = Status.trailersFromThrowable(cause);
+      if (trailers != null) {
+        return trailers.get(KEY_RESOURCE_INFO);
+      }
+    }
+    return null;
+  }
+
   private static SpannerException newSpannerExceptionPreformatted(
       ErrorCode code, @Nullable String message, @Nullable Throwable cause) {
     // This is the one place in the codebase that is allowed to call constructors directly.
@@ -183,10 +193,13 @@ private static SpannerException newSpannerExceptionPreformatted(
       case ABORTED:
         return new AbortedException(token, message, cause);
       case NOT_FOUND:
-        if (message != null && message.contains("Session not found")) {
-          return new SessionNotFoundException(token, message, cause);
-        } else if (message != null && DATABASE_NOT_FOUND_MSG_PATTERN.matcher(message).matches()) {
-          return new DatabaseNotFoundException(token, message, cause);
+        ResourceInfo resourceInfo = extractResourceInfo(cause);
+        if (resourceInfo != null) {
+          if (resourceInfo.getResourceType().equals(SESSION_RESOURCE_TYPE)) {
+            return new SessionNotFoundException(token, message, resourceInfo, cause);
+          } else if (resourceInfo.getResourceType().equals(DATABASE_RESOURCE_TYPE)) {
+            return new DatabaseNotFoundException(token, message, resourceInfo, cause);
+          }
         }
         // Fall through to the default.
       default:

google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java

@@ -51,24 +51,12 @@
 
 @RunWith(JUnit4.class)
 public class DatabaseClientImplTest {
-  private static final String DATABASE_NOT_FOUND_FORMAT =
-      "Database not found: projects/%s/instances/%s/databases/%s\n"
-          + "resource_type: \"type.googleapis.com/google.spanner.admin.database.v1.Database\"\n"
-          + "resource_name: \"projects/%s/instances/%s/databases/%s\"\n"
-          + "description: \"Database does not exist.\"\n";
   private static final String TEST_PROJECT = "my-project";
   private static final String TEST_INSTANCE = "my-instance";
   private static final String TEST_DATABASE = "my-database";
-  private static final String DATABASE_NOT_FOUND_MSG =
+  private static final String DATABASE_NAME =
       String.format(
-          "com.google.cloud.spanner.SpannerException: NOT_FOUND: io.grpc.StatusRuntimeException: NOT_FOUND: "
-              + DATABASE_NOT_FOUND_FORMAT,
-          TEST_PROJECT,
-          TEST_INSTANCE,
-          TEST_DATABASE,
-          TEST_PROJECT,
-          TEST_INSTANCE,
-          TEST_DATABASE);
+          "projects/%s/instances/%s/databases/%s", TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE);
   private static MockSpannerServiceImpl mockSpanner;
   private static Server server;
   private static LocalChannelProvider channelProvider;
@@ -283,7 +271,8 @@ public Long run(TransactionContext transaction) throws Exception {
   public void testDatabaseDoesNotExistOnPrepareSession() throws Exception {
     mockSpanner.setBeginTransactionExecutionTime(
         SimulatedExecutionTime.ofStickyException(
-            Status.NOT_FOUND.withDescription(DATABASE_NOT_FOUND_MSG).asRuntimeException()));
+            SpannerExceptionFactoryTest.newStatusResourceNotFoundException(
+                "Database", SpannerExceptionFactory.DATABASE_RESOURCE_TYPE, TEST_DATABASE)));
     DatabaseClientImpl dbClient =
         (DatabaseClientImpl)
             spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
@@ -323,7 +312,8 @@ public Void run(TransactionContext transaction) throws Exception {
   public void testDatabaseDoesNotExistOnInitialization() throws Exception {
     mockSpanner.setBatchCreateSessionsExecutionTime(
         SimulatedExecutionTime.ofStickyException(
-            Status.NOT_FOUND.withDescription(DATABASE_NOT_FOUND_MSG).asRuntimeException()));
+            SpannerExceptionFactoryTest.newStatusResourceNotFoundException(
+                "Database", SpannerExceptionFactory.DATABASE_RESOURCE_TYPE, DATABASE_NAME)));
     DatabaseClientImpl dbClient =
         (DatabaseClientImpl)
             spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
@@ -342,7 +332,8 @@ public void testDatabaseDoesNotExistOnInitialization() throws Exception {
   public void testDatabaseDoesNotExistOnCreate() throws Exception {
     mockSpanner.setBatchCreateSessionsExecutionTime(
         SimulatedExecutionTime.ofStickyException(
-            Status.NOT_FOUND.withDescription(DATABASE_NOT_FOUND_MSG).asRuntimeException()));
+            SpannerExceptionFactoryTest.newStatusResourceNotFoundException(
+                "Database", SpannerExceptionFactory.DATABASE_RESOURCE_TYPE, DATABASE_NAME)));
     // Ensure there are no sessions in the pool by default.
     try (Spanner spanner =
         SpannerOptions.newBuilder()
@@ -376,7 +367,8 @@ public void testDatabaseDoesNotExistOnCreate() throws Exception {
   public void testDatabaseDoesNotExistOnReplenish() throws Exception {
     mockSpanner.setBatchCreateSessionsExecutionTime(
         SimulatedExecutionTime.ofStickyException(
-            Status.NOT_FOUND.withDescription(DATABASE_NOT_FOUND_MSG).asRuntimeException()));
+            SpannerExceptionFactoryTest.newStatusResourceNotFoundException(
+                "Database", SpannerExceptionFactory.DATABASE_RESOURCE_TYPE, DATABASE_NAME)));
     DatabaseClientImpl dbClient =
         (DatabaseClientImpl)
             spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
@@ -471,7 +463,8 @@ public void testDatabaseIsDeletedAndThenRecreated() throws Exception {
       // Simulate that the database has been deleted.
       mockSpanner.setStickyGlobalExceptions(true);
       mockSpanner.addException(
-          Status.NOT_FOUND.withDescription(DATABASE_NOT_FOUND_MSG).asRuntimeException());
+          SpannerExceptionFactoryTest.newStatusResourceNotFoundException(
+              "Database", SpannerExceptionFactory.DATABASE_RESOURCE_TYPE, DATABASE_NAME));
 
       // All subsequent calls should fail with a DatabaseNotFoundException.
       try (ResultSet rs = dbClient.singleUse().executeQuery(SELECT1)) {

google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java

@@ -31,6 +31,7 @@
 import com.google.protobuf.Timestamp;
 import com.google.protobuf.Value.KindCase;
 import com.google.rpc.Code;
+import com.google.rpc.ResourceInfo;
 import com.google.rpc.RetryInfo;
 import com.google.spanner.v1.BatchCreateSessionsRequest;
 import com.google.spanner.v1.BatchCreateSessionsResponse;
@@ -727,10 +728,21 @@ public void getSession(GetSessionRequest request, StreamObserver<Session> respon
   }
 
   private <T> void setSessionNotFound(String name, StreamObserver<T> responseObserver) {
+    ResourceInfo resourceInfo =
+        ResourceInfo.newBuilder()
+            .setResourceType(SpannerExceptionFactory.SESSION_RESOURCE_TYPE)
+            .setResourceName(name)
+            .build();
+    Metadata.Key<ResourceInfo> key =
+        Metadata.Key.of(
+            resourceInfo.getDescriptorForType().getFullName() + Metadata.BINARY_HEADER_SUFFIX,
+            ProtoLiteUtils.metadataMarshaller(resourceInfo));
+    Metadata trailers = new Metadata();
+    trailers.put(key, resourceInfo);
     responseObserver.onError(
         Status.NOT_FOUND
             .withDescription(String.format("Session not found: Session with id %s not found", name))
-            .asRuntimeException());
+            .asRuntimeException(trailers));
   }
 
   @Override

google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolStressTest.java

@@ -162,8 +162,7 @@ public ApiFuture<Empty> answer(InvocationOnMock invocation) throws Throwable {
                 synchronized (lock) {
                   if (expiredSessions.contains(session.getName())) {
                     return ApiFutures.immediateFailedFuture(
-                        SpannerExceptionFactory.newSpannerException(
-                            ErrorCode.NOT_FOUND, "Session not found"));
+                        SpannerExceptionFactoryTest.newSessionNotFoundException(session.getName()));
                   }
                   if (sessions.remove(session.getName()) == null) {
                     setFailed(closedSessions.get(session.getName()));
@@ -185,8 +184,7 @@ public ApiFuture<Empty> answer(InvocationOnMock invocation) throws Throwable {
               public Void answer(InvocationOnMock invocation) throws Throwable {
                 if (random.nextInt(100) < 10) {
                   expireSession(session);
-                  throw SpannerExceptionFactory.newSpannerException(
-                      ErrorCode.NOT_FOUND, "Session not found");
+                  throw SpannerExceptionFactoryTest.newSessionNotFoundException(session.getName());
                 }
                 synchronized (lock) {
                   if (sessions.put(session.getName(), true)) {

google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java

@@ -94,6 +94,7 @@ public class SessionPoolTest extends BaseSessionPoolTest {
   DatabaseId db = DatabaseId.of("projects/p/instances/i/databases/unused");
   SessionPool pool;
   SessionPoolOptions options;
+  private String sessionName = String.format("%s/sessions/s", db.getName());
 
   @Parameters(name = "min sessions = {0}")
   public static Collection<Object[]> data() {
@@ -818,7 +819,7 @@ public void poolWorksWhenSessionNotFound() {
     SessionImpl mockSession2 = mockSession();
     final LinkedList<SessionImpl> sessions =
         new LinkedList<>(Arrays.asList(mockSession1, mockSession2));
-    doThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.NOT_FOUND, "Session not found"))
+    doThrow(SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName))
         .when(mockSession1)
         .prepareReadWriteTransaction();
     doAnswer(
@@ -1029,8 +1030,7 @@ public void testSessionNotFoundSingleUse() {
     ReadContext closedContext = mock(ReadContext.class);
     ResultSet closedResultSet = mock(ResultSet.class);
     when(closedResultSet.next())
-        .thenThrow(
-            SpannerExceptionFactory.newSpannerException(ErrorCode.NOT_FOUND, "Session not found"));
+        .thenThrow(SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName));
     when(closedContext.executeQuery(statement)).thenReturn(closedResultSet);
     when(closedSession.singleUse()).thenReturn(closedContext);
 
@@ -1088,8 +1088,7 @@ public void testSessionNotFoundReadOnlyTransaction() {
     Statement statement = Statement.of("SELECT 1");
     final SessionImpl closedSession = mockSession();
     when(closedSession.readOnlyTransaction())
-        .thenThrow(
-            SpannerExceptionFactory.newSpannerException(ErrorCode.NOT_FOUND, "Session not found"));
+        .thenThrow(SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName));
 
     final SessionImpl openSession = mockSession();
     ReadOnlyTransaction openTransaction = mock(ReadOnlyTransaction.class);
@@ -1155,7 +1154,7 @@ public void testSessionNotFoundReadWriteTransaction() {
     final Statement queryStatement = Statement.of("SELECT 1");
     final Statement updateStatement = Statement.of("UPDATE FOO SET BAR=1 WHERE ID=2");
     final SpannerException sessionNotFound =
-        SpannerExceptionFactory.newSpannerException(ErrorCode.NOT_FOUND, "Session not found");
+        SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName);
     for (ReadWriteTransactionTestStatementType statementType :
         ReadWriteTransactionTestStatementType.values()) {
       final ReadWriteTransactionTestStatementType executeStatementType = statementType;
@@ -1340,7 +1339,7 @@ public Integer run(TransactionContext transaction) throws Exception {
   @Test
   public void testSessionNotFoundOnPrepareTransaction() {
     final SpannerException sessionNotFound =
-        SpannerExceptionFactory.newSpannerException(ErrorCode.NOT_FOUND, "Session not found");
+        SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName);
     final SessionImpl closedSession = mock(SessionImpl.class);
     when(closedSession.getName())
         .thenReturn("projects/dummy/instances/dummy/database/dummy/sessions/session-closed");
@@ -1394,7 +1393,7 @@ public void run() {
   @Test
   public void testSessionNotFoundWrite() {
     SpannerException sessionNotFound =
-        SpannerExceptionFactory.newSpannerException(ErrorCode.NOT_FOUND, "Session not found");
+        SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName);
     List<Mutation> mutations = Arrays.asList(Mutation.newInsertBuilder("FOO").build());
     final SessionImpl closedSession = mockSession();
     when(closedSession.write(mutations)).thenThrow(sessionNotFound);
@@ -1446,7 +1445,7 @@ public void run() {
   @Test
   public void testSessionNotFoundWriteAtLeastOnce() {
     SpannerException sessionNotFound =
-        SpannerExceptionFactory.newSpannerException(ErrorCode.NOT_FOUND, "Session not found");
+        SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName);
     List<Mutation> mutations = Arrays.asList(Mutation.newInsertBuilder("FOO").build());
     final SessionImpl closedSession = mockSession();
     when(closedSession.writeAtLeastOnce(mutations)).thenThrow(sessionNotFound);
@@ -1497,7 +1496,7 @@ public void run() {
   @Test
   public void testSessionNotFoundPartitionedUpdate() {
     SpannerException sessionNotFound =
-        SpannerExceptionFactory.newSpannerException(ErrorCode.NOT_FOUND, "Session not found");
+        SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName);
     Statement statement = Statement.of("UPDATE FOO SET BAR=1 WHERE 1=1");
     final SessionImpl closedSession = mockSession();
     when(closedSession.executePartitionedUpdate(statement)).thenThrow(sessionNotFound);