Skip to content

Commit

Permalink
feat: Add overloaded executePartitionedDmlStatement to support upda…
Browse files Browse the repository at this point in the history
…te options (#2025)

Expose `UpdateOption` parameters for `executePartitionedDmlStatement` in Spanner.

Ref: Support required in Spring Lib to set priority for partitioned DML b/287076754
  • Loading branch information
prash-mi authored Jul 18, 2023
1 parent dcbf202 commit df65e54
Show file tree
Hide file tree
Showing 5 changed files with 100 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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.
*
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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() {

Expand Down

0 comments on commit df65e54

Please sign in to comment.