Spanner: migrate all unary call methods to gapic and inject headers (#3112)

* inject headers to grpcCallContext before making calls

* Suppress retry in gapic

Author: yihanzhen

google-cloud-bom/pom.xml

@@ -170,9 +170,9 @@
     <testing.version>0.43.1-alpha-SNAPSHOT</testing.version><!-- {x-version-update:google-cloud-testing:current} -->
 
     <api-common.version>1.5.0</api-common.version>
-    <gax.version>1.23.0</gax.version>
-    <gax-grpc.version>1.23.0</gax-grpc.version>
-    <gax-httpjson.version>0.40.0</gax-httpjson.version>
+    <gax.version>1.24.0</gax.version>
+    <gax-grpc.version>1.24.0</gax-grpc.version>
+    <gax-httpjson.version>0.41.0</gax-httpjson.version>
     <generated-proto-beta.version>0.8.0</generated-proto-beta.version>
     <generated-proto-ga.version>1.7.0</generated-proto-ga.version>
   </properties>

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

@@ -266,7 +266,7 @@ Session createSession(final DatabaseId db) throws SpannerException {
               new Callable<com.google.spanner.v1.Session>() {
                 @Override
                 public com.google.spanner.v1.Session call() throws Exception {
-                  return rawGrpcRpc.createSession(
+                  return gapicRpc.createSession(
                       db.getName(), getOptions().getSessionLabels(), options);
                 }
               });
@@ -806,7 +806,7 @@ public Timestamp writeAtLeastOnce(Iterable<Mutation> mutations) throws SpannerEx
                 new Callable<CommitResponse>() {
                   @Override
                   public CommitResponse call() throws Exception {
-                    return rawGrpcRpc.commit(request, options);
+                    return gapicRpc.commit(request, options);
                   }
                 });
         Timestamp t = Timestamp.fromProto(response.getCommitTimestamp());
@@ -872,7 +872,7 @@ public void close() {
             new Callable<Void>() {
               @Override
               public Void call() throws Exception {
-                rawGrpcRpc.deleteSession(name, options);
+                gapicRpc.deleteSession(name, options);
                 return null;
               }
             });
@@ -898,7 +898,7 @@ ByteString beginTransaction() {
                 new Callable<Transaction>() {
                   @Override
                   public Transaction call() throws Exception {
-                    return rawGrpcRpc.beginTransaction(request, options);
+                    return gapicRpc.beginTransaction(request, options);
                   }
                 });
         if (txn.getId().isEmpty()) {

google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java

@@ -18,14 +18,18 @@
 
 import static com.google.cloud.spanner.SpannerExceptionFactory.newSpannerException;
 
+import com.google.api.core.ApiFunction;
 import com.google.api.gax.core.CredentialsProvider;
 import com.google.api.gax.core.GaxProperties;
 import com.google.api.gax.grpc.GaxGrpcProperties;
+import com.google.api.gax.grpc.GrpcCallContext;
 import com.google.api.gax.grpc.GrpcTransportChannel;
 import com.google.api.gax.rpc.ApiClientHeaderProvider;
 import com.google.api.gax.rpc.FixedTransportChannelProvider;
 import com.google.api.gax.rpc.HeaderProvider;
+import com.google.api.gax.rpc.StatusCode;
 import com.google.api.gax.rpc.TransportChannelProvider;
+import com.google.api.gax.rpc.UnaryCallSettings;
 import com.google.api.pathtemplate.PathTemplate;
 import com.google.cloud.ServiceOptions;
 import com.google.cloud.grpc.GrpcTransportOptions;
@@ -43,6 +47,7 @@
 import com.google.cloud.spanner.v1.stub.SpannerStub;
 import com.google.cloud.spanner.v1.stub.SpannerStubSettings;
 import com.google.common.base.MoreObjects;
+import com.google.common.collect.ImmutableSet;
 import com.google.longrunning.GetOperationRequest;
 import com.google.longrunning.Operation;
 import com.google.protobuf.FieldMask;
@@ -68,6 +73,8 @@
 import com.google.spanner.v1.BeginTransactionRequest;
 import com.google.spanner.v1.CommitRequest;
 import com.google.spanner.v1.CommitResponse;
+import com.google.spanner.v1.CreateSessionRequest;
+import com.google.spanner.v1.DeleteSessionRequest;
 import com.google.spanner.v1.ExecuteSqlRequest;
 import com.google.spanner.v1.PartitionQueryRequest;
 import com.google.spanner.v1.PartitionReadRequest;
@@ -140,25 +147,59 @@ public GapicSpannerRpc(SpannerOptions options) throws IOException {
                 .build());
     CredentialsProvider credentialsProvider =
         GrpcTransportOptions.setUpCredentialsProvider(options);
-
-    this.stub =
+    
+    // Disabling retry for now because spanner handles retry in SpannerImpl.
+    // We will finally want to improve gax but for smooth transitioning we
+    // preserve the retry in SpannerImpl
+    try {
+      // TODO: bump the version of gax and remove this try-catch block
+      // applyToAllUnaryMethods does not throw exception in the latest version
+      this.stub =
         GrpcSpannerStub.create(
             SpannerStubSettings.newBuilder()
                 .setTransportChannelProvider(channelProvider)
                 .setCredentialsProvider(credentialsProvider)
+                .applyToAllUnaryMethods(
+                    new ApiFunction<UnaryCallSettings.Builder<?, ?>, Void>() {
+                      @Override
+                      public Void apply(UnaryCallSettings.Builder<?,?> builder) {
+                        builder.setRetryableCodes(ImmutableSet.<StatusCode.Code>of());
+                        return null;
+                      }
+                    })
                 .build());
+    
     this.instanceStub =
         GrpcInstanceAdminStub.create(
             InstanceAdminStubSettings.newBuilder()
                 .setTransportChannelProvider(channelProvider)
                 .setCredentialsProvider(credentialsProvider)
+                .applyToAllUnaryMethods(
+                    new ApiFunction<UnaryCallSettings.Builder<?, ?>, Void>() {
+                      @Override
+                      public Void apply(UnaryCallSettings.Builder<?,?> builder) {
+                        builder.setRetryableCodes(ImmutableSet.<StatusCode.Code>of());
+                        return null;
+                      }
+                    })
                 .build());
     this.databaseStub =
         GrpcDatabaseAdminStub.create(
             DatabaseAdminStubSettings.newBuilder()
                 .setTransportChannelProvider(channelProvider)
                 .setCredentialsProvider(credentialsProvider)
+                .applyToAllUnaryMethods(
+                    new ApiFunction<UnaryCallSettings.Builder<?, ?>, Void>() {
+                      @Override
+                      public Void apply(UnaryCallSettings.Builder<?,?> builder) {
+                        builder.setRetryableCodes(ImmutableSet.<StatusCode.Code>of());
+                        return null;
+                      }
+                    })
                 .build());
+    } catch (Exception e) {
+      throw SpannerExceptionFactory.newSpannerException(e);
+    }
   }
 
   @Override
@@ -171,9 +212,9 @@ public Paginated<InstanceConfig> listInstanceConfigs(int pageSize, @Nullable Str
     }
     ListInstanceConfigsRequest request = requestBuilder.build();
 
-    // TODO: put projectName in metadata
+    GrpcCallContext context = newCallContext(null, projectName);
     ListInstanceConfigsResponse response = 
-        get(instanceStub.listInstanceConfigsCallable().futureCall(request));
+        get(instanceStub.listInstanceConfigsCallable().futureCall(request, context));
     return new Paginated<>(response.getInstanceConfigsList(), response.getNextPageToken());
   }
 
@@ -182,8 +223,8 @@ public InstanceConfig getInstanceConfig(String instanceConfigName) throws Spanne
     GetInstanceConfigRequest request = 
         GetInstanceConfigRequest.newBuilder().setName(instanceConfigName).build();
 
-    // TODO: put projectName in metadata
-    return get(instanceStub.getInstanceConfigCallable().futureCall(request));
+    GrpcCallContext context = newCallContext(null, projectName);
+    return get(instanceStub.getInstanceConfigCallable().futureCall(request, context));
   }
 
   @Override
@@ -199,9 +240,9 @@ public Paginated<Instance> listInstances(
     }
     ListInstancesRequest request = requestBuilder.build();
 
-    // TODO: put projectName in metadata
+    GrpcCallContext context = newCallContext(null, projectName);
     ListInstancesResponse response = 
-        get(instanceStub.listInstancesCallable().futureCall(request));
+        get(instanceStub.listInstancesCallable().futureCall(request, context));
     return new Paginated<>(response.getInstancesList(), response.getNextPageToken());
   }
 
@@ -214,34 +255,36 @@ public Operation createInstance(String parent, String instanceId, Instance insta
             .setInstanceId(instanceId)
             .setInstance(instance)
             .build();
-    // TODO: put parent in metadata
-    return get(instanceStub.createInstanceCallable().futureCall(request));
+
+    GrpcCallContext context = newCallContext(null, parent);
+    return get(instanceStub.createInstanceCallable().futureCall(request, context));
   }
 
   @Override
   public Operation updateInstance(Instance instance, FieldMask fieldMask) throws SpannerException {
     UpdateInstanceRequest request =
         UpdateInstanceRequest.newBuilder().setInstance(instance).setFieldMask(fieldMask).build();
-    // TODO: put instance.getName() in metadata
-    return get(instanceStub.updateInstanceCallable().futureCall(request));
+    
+    GrpcCallContext context = newCallContext(null, instance.getName());
+    return get(instanceStub.updateInstanceCallable().futureCall(request, context));
   }
 
   @Override
   public Instance getInstance(String instanceName) throws SpannerException {
     GetInstanceRequest request = 
         GetInstanceRequest.newBuilder().setName(instanceName).build();
 
-    // TODO: put instanceName in metadata
-    return get(instanceStub.getInstanceCallable().futureCall(request));
+    GrpcCallContext context = newCallContext(null, instanceName);
+    return get(instanceStub.getInstanceCallable().futureCall(request, context));
   }
 
   @Override
   public void deleteInstance(String instanceName) throws SpannerException {
     DeleteInstanceRequest request = 
         DeleteInstanceRequest.newBuilder().setName(instanceName).build();
 
-    // TODO: put instanceName in metadata
-    get(instanceStub.deleteInstanceCallable().futureCall(request));
+    GrpcCallContext context = newCallContext(null, instanceName);
+    get(instanceStub.deleteInstanceCallable().futureCall(request, context));
   }
 
   @Override
@@ -254,8 +297,9 @@ public Paginated<Database> listDatabases(
     }
     ListDatabasesRequest request = requestBuilder.build();
 
-    // TODO: put instanceName in metadata
-    ListDatabasesResponse response = get(databaseStub.listDatabasesCallable().futureCall(request));
+    GrpcCallContext context = newCallContext(null, instanceName);
+    ListDatabasesResponse response = get(databaseStub.listDatabasesCallable()
+        .futureCall(request, context));
     return new Paginated<>(response.getDatabasesList(), response.getNextPageToken());
   }
 
@@ -268,8 +312,8 @@ public Operation createDatabase(String instanceName, String createDatabaseStatem
             .setCreateStatement(createDatabaseStatement)
             .addAllExtraStatements(additionalStatements)
             .build();
-    // TODO: put instanceName in metadata
-    return get(databaseStub.createDatabaseCallable().futureCall(request));
+    GrpcCallContext context = newCallContext(null, instanceName);
+    return get(databaseStub.createDatabaseCallable().futureCall(request, context));
   }
 
   @Override
@@ -281,17 +325,17 @@ public Operation updateDatabaseDdl(String databaseName, Iterable<String> updateD
             .addAllStatements(updateDatabaseStatements)
             .setOperationId(MoreObjects.firstNonNull(updateId, ""))
             .build();
-    // TODO: put databaseName in metadata
-    return get(databaseStub.updateDatabaseDdlCallable().futureCall(request));
+    GrpcCallContext context = newCallContext(null, databaseName);
+    return get(databaseStub.updateDatabaseDdlCallable().futureCall(request, context));
   }
 
   @Override
   public void dropDatabase(String databaseName) throws SpannerException {
     DropDatabaseRequest request =
         DropDatabaseRequest.newBuilder().setDatabase(databaseName).build();
 
-    // TODO: put databaseName in metadata
-    get(databaseStub.dropDatabaseCallable().futureCall(request));
+    GrpcCallContext context = newCallContext(null, databaseName);
+    get(databaseStub.dropDatabaseCallable().futureCall(request, context));
   }
 
   @Override
@@ -301,43 +345,56 @@ public Database getDatabase(String databaseName) throws SpannerException {
         .setName(databaseName)
         .build();
 
-    // TODO: put databaseName in metadata
-    return get(databaseStub.getDatabaseCallable().futureCall(request));
+    GrpcCallContext context = newCallContext(null, databaseName);
+    return get(databaseStub.getDatabaseCallable().futureCall(request, context));
   }
 
   @Override
   public List<String> getDatabaseDdl(String databaseName) throws SpannerException {
     GetDatabaseDdlRequest request = 
         GetDatabaseDdlRequest.newBuilder().setDatabase(databaseName).build();
 
-    // TODO: put databaseName in metadata
-    return get(databaseStub.getDatabaseDdlCallable().futureCall(request))
+    GrpcCallContext context = newCallContext(null, databaseName);
+    return get(databaseStub.getDatabaseDdlCallable().futureCall(request, context))
                .getStatementsList();
   }
 
   @Override
   public Operation getOperation(String name) throws SpannerException {
     GetOperationRequest request = GetOperationRequest.newBuilder().setName(name).build();
-    // TODO: put name in metadata
-    return get(databaseStub.getOperationsStub().getOperationCallable().futureCall(request));
+    GrpcCallContext context = newCallContext(null, name);
+    return get(databaseStub.getOperationsStub().getOperationCallable()
+        .futureCall(request, context));
   }
 
   @Override
   public Session createSession(String databaseName, @Nullable Map<String, String> labels,
       @Nullable Map<Option, ?> options) throws SpannerException {
-    throw new UnsupportedOperationException("Not implemented yet.");
+    CreateSessionRequest.Builder requestBuilder =
+        CreateSessionRequest.newBuilder().setDatabase(databaseName);
+    if (labels != null && !labels.isEmpty()) {
+      Session.Builder session = Session.newBuilder().putAllLabels(labels);
+      requestBuilder.setSession(session);
+    }
+    CreateSessionRequest request = requestBuilder.build();
+    GrpcCallContext context = newCallContext(options, databaseName);
+    return get(stub.createSessionCallable().futureCall(request, context));
   }
 
   @Override
   public void deleteSession(String sessionName, @Nullable Map<Option, ?> options)
       throws SpannerException {
-    throw new UnsupportedOperationException("Not implemented yet.");
+    DeleteSessionRequest request =
+        DeleteSessionRequest.newBuilder().setName(sessionName).build();
+    GrpcCallContext context = newCallContext(options, sessionName);
+    get(stub.deleteSessionCallable().futureCall(request, context));
   }
 
   @Override
   public StreamingCall read(
       ReadRequest request, ResultStreamConsumer consumer, @Nullable Map<Option, ?> options) {
-    throw new UnsupportedOperationException("Not implemented yet.");
+    GrpcCallContext context = newCallContext(options, request.getSession());
+    throw new UnsupportedOperationException("not implemented yet");
   }
 
   @Override
@@ -349,33 +406,36 @@ public StreamingCall executeQuery(
   @Override
   public Transaction beginTransaction(
       BeginTransactionRequest request, @Nullable Map<Option, ?> options) throws SpannerException {
-    throw new UnsupportedOperationException("Not implemented yet.");
+    GrpcCallContext context = newCallContext(options, request.getSession());
+    return get(stub.beginTransactionCallable().futureCall(request, context));
   }
 
   @Override
   public CommitResponse commit(CommitRequest commitRequest, @Nullable Map<Option, ?> options)
       throws SpannerException {
-    throw new UnsupportedOperationException("Not implemented yet.");
+    GrpcCallContext context = newCallContext(options, commitRequest.getSession());
+    return get(stub.commitCallable().futureCall(commitRequest, context));
   }
 
   @Override
   public void rollback(RollbackRequest request, @Nullable Map<Option, ?> options)
       throws SpannerException {
-    throw new UnsupportedOperationException("Not implemented yet.");
+    GrpcCallContext context = newCallContext(options, request.getSession());
+    get(stub.rollbackCallable().futureCall(request, context));
   }
 
   @Override
   public PartitionResponse partitionQuery(
       PartitionQueryRequest request, @Nullable Map<Option, ?> options) throws SpannerException {
-    throw new UnsupportedOperationException("Not implemented yet.");
+    GrpcCallContext context = newCallContext(options, request.getSession());
+    return get(stub.partitionQueryCallable().futureCall(request, context));
   }
 
   @Override
   public PartitionResponse partitionRead(
       PartitionReadRequest request, @Nullable Map<Option, ?> options) throws SpannerException {
-    // TODO(pongad): Figure out metadata
-    // TODO(pongad): Figure out channel affinity
-    return get(stub.partitionReadCallable().futureCall(request));
+    GrpcCallContext context = newCallContext(options, request.getSession());
+    return get(stub.partitionReadCallable().futureCall(request, context));
   }
 
   /** Gets the result of an async RPC call, handling any exceptions encountered. */
@@ -391,4 +451,14 @@ private static <T> T get(final Future<T> future) throws SpannerException {
       throw newSpannerException(context, e);
     }
   }
+
+  private GrpcCallContext newCallContext(@Nullable Map<Option, ?> options, String resource) {
+    GrpcCallContext context = GrpcCallContext.createDefault();
+    if (options != null) {
+      context = context.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue());
+    }
+    context = context.withExtraHeaders(
+        metadataProvider.newExtraHeaders(resource, projectName));
+    return context;
+  }
 }