feat: add PartitionedUpdate support to executor

Author: rajatbhatta

google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/CloudClientExecutor.java

@@ -22,13 +22,9 @@
 import com.google.api.gax.paging.Page;
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.api.gax.rpc.DeadlineExceededException;
-import com.google.api.gax.rpc.TransportChannelProvider;
 import com.google.api.gax.rpc.UnavailableException;
-import com.google.auth.Credentials;
-import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.ByteArray;
 import com.google.cloud.Date;
-import com.google.cloud.NoCredentials;
 import com.google.cloud.Timestamp;
 import com.google.cloud.spanner.Backup;
 import com.google.cloud.spanner.BatchClient;
@@ -52,6 +48,7 @@
 import com.google.cloud.spanner.Mutation;
 import com.google.cloud.spanner.Mutation.WriteBuilder;
 import com.google.cloud.spanner.Options;
+import com.google.cloud.spanner.Options.RpcPriority;
 import com.google.cloud.spanner.Partition;
 import com.google.cloud.spanner.PartitionOptions;
 import com.google.cloud.spanner.ReadContext;
@@ -128,6 +125,8 @@
 import com.google.spanner.executor.v1.MutationAction.Mod;
 import com.google.spanner.executor.v1.MutationAction.UpdateArgs;
 import com.google.spanner.executor.v1.OperationResponse;
+import com.google.spanner.executor.v1.PartitionedUpdateAction;
+import com.google.spanner.executor.v1.PartitionedUpdateAction.ExecutePartitionedUpdateOptions;
 import com.google.spanner.executor.v1.QueryAction;
 import com.google.spanner.executor.v1.ReadAction;
 import com.google.spanner.executor.v1.RestoreCloudDatabaseAction;
@@ -145,8 +144,6 @@
 import com.google.spanner.v1.TypeCode;
 import io.grpc.Status;
 import io.grpc.stub.StreamObserver;
-import java.io.ByteArrayInputStream;
-import java.io.File;
 import java.io.IOException;
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
@@ -165,7 +162,6 @@
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
-import org.apache.commons.io.FileUtils;
 import org.jetbrains.annotations.NotNull;
 import org.threeten.bp.Duration;
 import org.threeten.bp.LocalDate;
@@ -747,20 +743,6 @@ private synchronized Spanner getClient() throws IOException {
 
   // Return the spanner client, create one if not exists.
   private synchronized Spanner getClient(long timeoutSeconds) throws IOException {
-    // Create a cloud spanner client
-    Credentials credentials;
-    if (WorkerProxy.serviceKeyFile.isEmpty()) {
-      credentials = NoCredentials.getInstance();
-    } else {
-      credentials =
-          GoogleCredentials.fromStream(
-              new ByteArrayInputStream(
-                  FileUtils.readFileToByteArray(new File(WorkerProxy.serviceKeyFile))),
-              HTTP_TRANSPORT_FACTORY);
-    }
-
-    TransportChannelProvider channelProvider =
-        CloudUtil.newChannelProviderHelper(WorkerProxy.spannerPort);
 
     Duration rpcTimeout = Duration.ofHours(1L);
     if (timeoutSeconds > 0) {
@@ -779,12 +761,7 @@ private synchronized Spanner getClient(long timeoutSeconds) throws IOException {
 
     // Cloud Spanner Client does not support global retry settings,
     // Thus, we need to add retry settings to each individual stub.
-    SpannerOptions.Builder optionsBuilder =
-        SpannerOptions.newBuilder()
-            .setProjectId(PROJECT_ID)
-            .setHost(HOST_PREFIX + WorkerProxy.spannerPort)
-            .setCredentials(credentials)
-            .setChannelProvider(channelProvider);
+    SpannerOptions.Builder optionsBuilder = SpannerOptions.newBuilder().setProjectId(PROJECT_ID);
 
     SpannerStubSettings.Builder stubSettingsBuilder =
         optionsBuilder.getSpannerStubSettingsBuilder();
@@ -886,6 +863,13 @@ private Status executeAction(
       } else if (action.hasExecutePartition()) {
         return executeExecutePartition(
             action.getExecutePartition(), outcomeSender, executionContext);
+      } else if (action.hasPartitionedUpdate()) {
+        if (dbPath == null) {
+          throw SpannerExceptionFactory.newSpannerException(
+              ErrorCode.INVALID_ARGUMENT, "Database path must be set for this action");
+        }
+        DatabaseClient dbClient = getClient().getDatabaseClient(DatabaseId.of(dbPath));
+        return executePartitionedUpdate(action.getPartitionedUpdate(), dbClient, outcomeSender);
       } else if (action.hasCloseBatchTxn()) {
         return executeCloseBatchTxn(action.getCloseBatchTxn(), outcomeSender, executionContext);
       } else if (action.hasExecuteChangeStreamQuery()) {
@@ -1974,6 +1958,33 @@ private Status executeExecutePartition(
     }
   }
 
+  /** Execute a partitioned update which runs different partitions in parallel. */
+  private Status executePartitionedUpdate(
+      PartitionedUpdateAction action, DatabaseClient dbClient, OutcomeSender sender) {
+    try {
+      ExecutePartitionedUpdateOptions options = action.getOptions();
+      Long count =
+          dbClient.executePartitionedUpdate(
+              Statement.of(action.getUpdate().getSql()),
+              Options.tag(options.getTag()),
+              Options.priority(RpcPriority.getEnumFromProto(options.getRpcPriority())));
+      SpannerActionOutcome outcome =
+          SpannerActionOutcome.newBuilder()
+              .setStatus(toProto(Status.OK))
+              .addDmlRowsModified(count)
+              .build();
+      sender.sendOutcome(outcome);
+      return sender.finishWithOK();
+    } catch (SpannerException e) {
+      return sender.finishWithError(toStatus(e));
+    } catch (Exception e) {
+      return sender.finishWithError(
+          toStatus(
+              SpannerExceptionFactory.newSpannerException(
+                  ErrorCode.INVALID_ARGUMENT, "Unexpected error: " + e.getMessage())));
+    }
+  }
+
   /** Build a child partition record proto out of childPartitionRecord returned by client. */
   private ChildPartitionsRecord buildChildPartitionRecord(Struct childPartitionRecord)
       throws Exception {