From df65e547278b27cc8dff5e70f140f722ed81c8c3 Mon Sep 17 00:00:00 2001 From: Prashant Mishra <11733935+prash-mi@users.noreply.github.com> Date: Tue, 18 Jul 2023 14:01:48 +0000 Subject: [PATCH] feat: Add overloaded `executePartitionedDmlStatement` to support update options (#2025) Expose `UpdateOption` parameters for `executePartitionedDmlStatement` in Spanner. Ref: Support required in Spring Lib to set priority for partitioned DML b/287076754 --- .../ReadOnlyTransactionSpannerTemplate.java | 7 ++ .../ReadWriteTransactionSpannerTemplate.java | 7 ++ .../data/spanner/core/SpannerOperations.java | 10 +++ .../data/spanner/core/SpannerTemplate.java | 9 ++- .../spanner/core/SpannerTemplateTests.java | 68 +++++++++++++++++++ 5 files changed, 100 insertions(+), 1 deletion(-) diff --git a/spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/ReadOnlyTransactionSpannerTemplate.java b/spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/ReadOnlyTransactionSpannerTemplate.java index 6a4e265300..afc23707e1 100644 --- a/spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/ReadOnlyTransactionSpannerTemplate.java +++ b/spring-cloud-gcp-data-spanner/src/main/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/spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/ReadWriteTransactionSpannerTemplate.java b/spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/ReadWriteTransactionSpannerTemplate.java index f983620d9d..2b30705c39 100644 --- a/spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/ReadWriteTransactionSpannerTemplate.java +++ b/spring-cloud-gcp-data-spanner/src/main/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/spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/SpannerOperations.java b/spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/SpannerOperations.java index e9a4439ab3..386c81eb0e 100644 --- a/spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/SpannerOperations.java +++ b/spring-cloud-gcp-data-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/spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/SpannerTemplate.java b/spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/SpannerTemplate.java index 6637907e72..f8a63fbcf3 100644 --- a/spring-cloud-gcp-data-spanner/src/main/java/com/google/cloud/spring/data/spanner/core/SpannerTemplate.java +++ b/spring-cloud-gcp-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/spring-cloud-gcp-data-spanner/src/test/java/com/google/cloud/spring/data/spanner/core/SpannerTemplateTests.java b/spring-cloud-gcp-data-spanner/src/test/java/com/google/cloud/spring/data/spanner/core/SpannerTemplateTests.java index 373e012f9d..4ab8561d1a 100644 --- a/spring-cloud-gcp-data-spanner/src/test/java/com/google/cloud/spring/data/spanner/core/SpannerTemplateTests.java +++ b/spring-cloud-gcp-data-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 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 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() {