Spanner: PartitionedDML transaction does not need a prepared session

google-cloud-clients/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java

@@ -32,6 +32,11 @@ class DatabaseClientImpl implements DatabaseClient {
   private static final String PARTITION_DML_TRANSACTION = "CloudSpanner.PartitionDMLTransaction";
   private static final Tracer tracer = Tracing.getTracer();
 
+  private enum SessionMode {
+    READ,
+    READ_WRITE
+  }
+
   static {
     TraceUtil.exportSpans(READ_WRITE_TRANSACTION, READ_ONLY_TRANSACTION, PARTITION_DML_TRANSACTION);
   }
@@ -57,6 +62,7 @@ public Timestamp write(final Iterable<Mutation> mutations) throws SpannerExcepti
     Span span = tracer.spanBuilder(READ_WRITE_TRANSACTION).startSpan();
     try (Scope s = tracer.withSpan(span)) {
       return runWithSessionRetry(
+          SessionMode.READ_WRITE,
           new Function<Session, Timestamp>() {
             @Override
             public Timestamp apply(Session session) {
@@ -76,6 +82,7 @@ public Timestamp writeAtLeastOnce(final Iterable<Mutation> mutations) throws Spa
     Span span = tracer.spanBuilder(READ_WRITE_TRANSACTION).startSpan();
     try (Scope s = tracer.withSpan(span)) {
       return runWithSessionRetry(
+          SessionMode.READ_WRITE,
           new Function<Session, Timestamp>() {
             @Override
             public Timestamp apply(Session session) {
@@ -182,7 +189,10 @@ public TransactionManager transactionManager() {
   public long executePartitionedUpdate(final Statement stmt) {
     Span span = tracer.spanBuilder(PARTITION_DML_TRANSACTION).startSpan();
     try (Scope s = tracer.withSpan(span)) {
+      // A partitioned update transaction does not need a prepared write session, as the transaction
+      // object will start a new transaction with specific options anyway.
       return runWithSessionRetry(
+          SessionMode.READ,
           new Function<Session, Long>() {
             @Override
             public Long apply(Session session) {
@@ -195,13 +205,17 @@ public Long apply(Session session) {
     }
   }
 
-  private <T> T runWithSessionRetry(Function<Session, T> callable) {
-    PooledSession session = getReadWriteSession();
+  private <T> T runWithSessionRetry(SessionMode mode, Function<Session, T> callable) {
+    PooledSession session =
+        mode == SessionMode.READ_WRITE ? getReadWriteSession() : getReadSession();
     while (true) {
       try {
         return callable.apply(session);
       } catch (SessionNotFoundException e) {
-        session = pool.replaceReadWriteSession(e, session);
+        session =
+            mode == SessionMode.READ_WRITE
+                ? pool.replaceReadWriteSession(e, session)
+                : pool.replaceReadSession(e, session);
       }
     }
   }

google-cloud-clients/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java

@@ -0,0 +1,96 @@
+/*
+ * Copyright 2019 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.spanner;
+
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.CoreMatchers.is;
+import static org.junit.Assert.assertThat;
+
+import com.google.api.gax.grpc.testing.LocalChannelProvider;
+import com.google.cloud.NoCredentials;
+import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
+import io.grpc.Server;
+import io.grpc.inprocess.InProcessServerBuilder;
+import java.io.IOException;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class DatabaseClientImplTest {
+  private static MockSpannerServiceImpl mockSpanner;
+  private static Server server;
+  private static LocalChannelProvider channelProvider;
+  private static Spanner spanner;
+  private static final Statement UPDATE_STATEMENT =
+      Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=2");
+  private static final long UPDATE_COUNT = 1L;
+
+  @BeforeClass
+  public static void startStaticServer() throws IOException {
+    mockSpanner = new MockSpannerServiceImpl();
+    mockSpanner.setAbortProbability(0.0D); // We don't want any unpredictable aborted transactions.
+    mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT));
+
+    String uniqueName = InProcessServerBuilder.generateName();
+    server =
+        InProcessServerBuilder.forName(uniqueName)
+            .directExecutor()
+            .addService(mockSpanner)
+            .build()
+            .start();
+    channelProvider = LocalChannelProvider.create(uniqueName);
+  }
+
+  @AfterClass
+  public static void stopServer() {
+    server.shutdown();
+  }
+
+  @Before
+  public void setUp() throws IOException {
+    mockSpanner.reset();
+    spanner =
+        SpannerOptions.newBuilder()
+            .setProjectId("[PROJECT]")
+            .setChannelProvider(channelProvider)
+            .setCredentials(NoCredentials.getInstance())
+            .build()
+            .getService();
+  }
+
+  @After
+  public void tearDown() throws Exception {
+    spanner.close();
+  }
+
+  /**
+   * Test that the update statement can be executed as a partitioned transaction that returns a
+   * lower bound update count.
+   */
+  @Test
+  public void testExecutePartitionedDml() {
+    DatabaseClient client =
+        spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
+    long updateCount = client.executePartitionedUpdate(UPDATE_STATEMENT);
+    assertThat(updateCount, is(equalTo(UPDATE_COUNT)));
+  }
+}