feat: Add overloaded executePartitionedDmlStatement to support upda…

…te options (#2025) Expose `UpdateOption` parameters for `executePartitionedDmlStatement` in Spanner. Ref: Support required in Spring Lib to set priority for partitioned DML b/287076754
GoogleCloudPlatform · Jul 18, 2023 · df65e54 · df65e54
1 parent dcbf202
commit df65e54
Show file tree

Hide file tree

Showing 5 changed files with 100 additions and 1 deletion.
diff --git a/...in/java/com/google/cloud/spring/data/spanner/core/ReadOnlyTransactionSpannerTemplate.java b/...in/java/com/google/cloud/spring/data/spanner/core/ReadOnlyTransactionSpannerTemplate.java
@@ -18,6 +18,7 @@
 
 import com.google.cloud.spanner.DatabaseClient;
 import com.google.cloud.spanner.Mutation;
+import com.google.cloud.spanner.Options.UpdateOption;
 import com.google.cloud.spanner.ReadContext;
 import com.google.cloud.spanner.ReadOnlyTransaction;
 import com.google.cloud.spanner.Statement;
@@ -72,6 +73,12 @@ public long executePartitionedDmlStatement(Statement statement) {
         "A read-only transaction template cannot execute partitioned DML.");
   }
 
+  @Override
+  public long executePartitionedDmlStatement(Statement statement, UpdateOption... options) {
+    throw new SpannerDataException(
+        "A read-only transaction template cannot execute partitioned DML.");
+  }
+
   @Override
   protected ReadContext getReadContext() {
     return this.readOnlyTransaction;

diff --git a/...n/java/com/google/cloud/spring/data/spanner/core/ReadWriteTransactionSpannerTemplate.java b/...n/java/com/google/cloud/spring/data/spanner/core/ReadWriteTransactionSpannerTemplate.java
@@ -18,6 +18,7 @@
 
 import com.google.cloud.spanner.DatabaseClient;
 import com.google.cloud.spanner.Mutation;
+import com.google.cloud.spanner.Options.UpdateOption;
 import com.google.cloud.spanner.ReadContext;
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.TimestampBound;
@@ -77,6 +78,12 @@ public long executePartitionedDmlStatement(Statement statement) {
         "A read-write transaction template cannot execute partitioned DML.");
   }
 
+  @Override
+  public long executePartitionedDmlStatement(Statement statement, UpdateOption... options) {
+    throw new SpannerDataException(
+        "A read-write transaction template cannot execute partitioned DML.");
+  }
+
   @Override
   protected ReadContext getReadContext(TimestampBound timestampBound) {
     throw new SpannerDataException(

diff --git a/...ta-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/SpannerOperations.java b/...ta-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/SpannerOperations.java
@@ -18,6 +18,7 @@
 
 import com.google.cloud.spanner.Key;
 import com.google.cloud.spanner.KeySet;
+import com.google.cloud.spanner.Options.UpdateOption;
 import com.google.cloud.spanner.ReadContext;
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.Struct;
@@ -50,6 +51,15 @@ public interface SpannerOperations {
    */
   long executePartitionedDmlStatement(Statement statement);
 
+  /**
+   * Execute a DML statement in partitioned mode. This is not available inside of transactions.
+   *
+   * @param statement the DML statement to execute.
+   * @param options marks options applicable to update operation.
+   * @return the lower-bound of number of rows affected.
+   */
+  long executePartitionedDmlStatement(Statement statement, UpdateOption... options);
+
   /**
    * Finds a single stored object using a key.
    *

diff --git a/...data-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/SpannerTemplate.java b/...data-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/SpannerTemplate.java
@@ -22,6 +22,7 @@
 import com.google.cloud.spanner.Mutation;
 import com.google.cloud.spanner.Options.QueryOption;
 import com.google.cloud.spanner.Options.ReadOption;
+import com.google.cloud.spanner.Options.UpdateOption;
 import com.google.cloud.spanner.ReadContext;
 import com.google.cloud.spanner.ReadOnlyTransaction;
 import com.google.cloud.spanner.ResultSet;
@@ -146,14 +147,20 @@ public long executeDmlStatement(Statement statement) {
 
   @Override
   public long executePartitionedDmlStatement(Statement statement) {
+    return executePartitionedDmlStatement(statement, new UpdateOption[] {});
+  }
+
+  @Override
+  public long executePartitionedDmlStatement(Statement statement, UpdateOption... options) {
     Assert.notNull(statement, "A non-null statement is required.");
+    Assert.notNull(options, "A non-null UpdateOption is required.");
     maybeEmitEvent(new BeforeExecuteDmlEvent(statement));
     long rowsAffected =
         doWithOrWithoutTransactionContext(
             x -> {
               throw new SpannerDataException("Cannot execute partitioned DML in a transaction.");
             },
-            () -> this.databaseClientProvider.get().executePartitionedUpdate(statement));
+            () -> this.databaseClientProvider.get().executePartitionedUpdate(statement, options));
     maybeEmitEvent(new AfterExecuteDmlEvent(statement, rowsAffected));
     return rowsAffected;
   }

diff --git a/...spanner/src/test/java/com/google/cloud/spring/data/spanner/core/SpannerTemplateTests.java b/...spanner/src/test/java/com/google/cloud/spring/data/spanner/core/SpannerTemplateTests.java
@@ -38,6 +38,8 @@
 import com.google.cloud.spanner.Mutation;
 import com.google.cloud.spanner.Options;
 import com.google.cloud.spanner.Options.ReadOption;
+import com.google.cloud.spanner.Options.ReadQueryUpdateTransactionOption;
+import com.google.cloud.spanner.Options.RpcPriority;
 import com.google.cloud.spanner.ReadContext;
 import com.google.cloud.spanner.ReadOnlyTransaction;
 import com.google.cloud.spanner.ResultSet;
@@ -89,6 +91,9 @@ class SpannerTemplateTests {
 
   private static final Statement DML = Statement.of("update statement");
 
+  private static final ReadQueryUpdateTransactionOption OPTION_RPC_PRIORITY_HIGH =
+      Options.priority(RpcPriority.HIGH);
+
   private DatabaseClient databaseClient;
 
   private SpannerMappingContext mappingContext;
@@ -196,6 +201,69 @@ void executePartitionedDmlTest() {
         x -> x.verify(this.databaseClient, times(1)).executePartitionedUpdate(DML));
   }
 
+  @Test
+  void executePartitionedDmlWithOptionsTest() {
+    when(this.databaseClient.executePartitionedUpdate(DML, OPTION_RPC_PRIORITY_HIGH))
+        .thenReturn(333L);
+    verifyBeforeAndAfterEvents(
+        new BeforeExecuteDmlEvent(DML),
+        new AfterExecuteDmlEvent(DML, 333L),
+        () -> this.spannerTemplate.executePartitionedDmlStatement(DML, OPTION_RPC_PRIORITY_HIGH),
+        x ->
+            x.verify(this.databaseClient, times(1))
+                .executePartitionedUpdate(DML, OPTION_RPC_PRIORITY_HIGH));
+  }
+
+  @Test
+  void readOnlyTransactionPartitionedDmlWithOptionsTest() {
+
+    ReadOnlyTransaction readOnlyTransaction = mock(ReadOnlyTransaction.class);
+    when(this.databaseClient.readOnlyTransaction(
+            TimestampBound.ofReadTimestamp(Timestamp.ofTimeMicroseconds(333))))
+        .thenReturn(readOnlyTransaction);
+
+    SpannerReadOptions testSpannerReadOptions =
+        new SpannerReadOptions().setTimestamp(Timestamp.ofTimeMicroseconds(333));
+    Function<SpannerTemplate, Void> testSpannerOperations =
+        spannerOperations -> {
+          spannerOperations.executePartitionedDmlStatement(
+              Statement.of("fail"), OPTION_RPC_PRIORITY_HIGH);
+          return null;
+        };
+
+    assertThatThrownBy(
+            () ->
+                this.spannerTemplate.performReadOnlyTransaction(
+                    testSpannerOperations, testSpannerReadOptions))
+        .hasMessage("A read-only transaction template cannot execute partitioned DML.");
+  }
+
+  @Test
+  void readWriteTransactionPartitionedDmlWithOptionsTest() {
+
+    TransactionRunner transactionRunner = mock(TransactionRunner.class);
+    when(this.databaseClient.readWriteTransaction()).thenReturn(transactionRunner);
+
+    TransactionContext transactionContext = mock(TransactionContext.class);
+
+    when(transactionRunner.run(any()))
+        .thenAnswer(
+            invocation -> {
+              TransactionCallable transactionCallable = invocation.getArgument(0);
+              return transactionCallable.run(transactionContext);
+            });
+    Function<SpannerTemplate, String> testSpannerOperations =
+        spannerTemplate -> {
+          spannerTemplate.executePartitionedDmlStatement(
+              Statement.of("DML statement here"), OPTION_RPC_PRIORITY_HIGH);
+          return "all done";
+        };
+
+    assertThatThrownBy(
+            () -> this.spannerTemplate.performReadWriteTransaction(testSpannerOperations))
+        .hasMessage("A read-write transaction template cannot execute partitioned DML.");
+  }
+
   @Test
   void readWriteTransactionTest() {