Merge branch 'partitioned-query' of github.com:pratickchokhani/java-spanner into partitioned-query

Author: pratickchokhani

google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java

@@ -79,9 +79,7 @@ public BatchReadOnlyTransaction batchReadOnlyTransaction(TimestampBound bound) {
   @Override
   public BatchReadOnlyTransaction batchReadOnlyTransaction(BatchTransactionId batchTransactionId) {
     SessionImpl session =
-        sessionClient.sessionWithId(
-            checkNotNull(batchTransactionId).getSessionId(),
-            batchTransactionId.isMultiplexedSession());
+        sessionClient.sessionWithId(checkNotNull(batchTransactionId).getSessionId());
     return new BatchReadOnlyTransactionImpl(
         MultiUseReadOnlyTransaction.newBuilder()
             .setSession(session)

google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchTransactionId.java

@@ -32,7 +32,6 @@ public class BatchTransactionId implements Serializable {
   private final ByteString transactionId;
   private final String sessionId;
   private final Timestamp timestamp;
-  private final boolean isMultiplexedSession;
   private static final long serialVersionUID = 8067099123096783939L;
 
   BatchTransactionId(
@@ -43,7 +42,6 @@ public class BatchTransactionId implements Serializable {
     this.transactionId = Preconditions.checkNotNull(transactionId);
     this.sessionId = Preconditions.checkNotNull(sessionId);
     this.timestamp = Preconditions.checkNotNull(timestamp);
-    this.isMultiplexedSession = isMultiplexedSession;
   }
 
   ByteString getTransactionId() {
@@ -58,15 +56,11 @@ Timestamp getTimestamp() {
     return timestamp;
   }
 
-  public boolean isMultiplexedSession() {
-    return isMultiplexedSession;
-  }
-
   @Override
   public String toString() {
     return String.format(
-        "transactionId: %s, sessionId: %s, timestamp: %s, isMultiplexedSession: %s",
-        transactionId.toStringUtf8(), sessionId, timestamp, isMultiplexedSession);
+        "transactionId: %s, sessionId: %s, timestamp: %s",
+        transactionId.toStringUtf8(), sessionId, timestamp);
   }
 
   @Override

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

@@ -317,7 +317,7 @@ public long executePartitionedUpdate(final Statement stmt, final UpdateOption...
     if (useMultiplexedSessionPartitionedOps) {
       return getMultiplexedSession().executePartitionedUpdate(stmt, options);
     }
-    return executePartitionedUpdateSession(stmt, options);
+    return executePartitionedUpdateWithPooledSession(stmt, options);
   }
 
   private long executePartitionedUpdateWithPooledSession(

google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedMultiplexedSessionTransaction.java

@@ -232,15 +232,9 @@ private SessionReference getSessionReference() {
    */
   @Override
   public long executePartitionedUpdate(Statement stmt, UpdateOption... options) {
-    try {
-      SessionReference sessionReference = getSessionReference();
-      return new MultiplexedSessionTransaction(
-              client, span, sessionReference, NO_CHANNEL_HINT, /* singleUse = */ true)
-          .executePartitionedUpdate(stmt, options);
-    } catch (InterruptedException e) {
-      throw new RuntimeException(e);
-    } catch (ExecutionException e) {
-      throw new RuntimeException(e);
-    }
+    SessionReference sessionReference = getSessionReference();
+    return new MultiplexedSessionTransaction(
+            client, span, sessionReference, NO_CHANNEL_HINT, /* singleUse = */ true)
+        .executePartitionedUpdate(stmt, options);
   }
 }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java

@@ -556,7 +556,8 @@ public AsyncTransactionManager transactionManagerAsync(TransactionOption... opti
 
   @Override
   public long executePartitionedUpdate(Statement stmt, UpdateOption... options) {
-    return createMultiplexedSessionTransaction(/* singleUse = */ true).executePartitionedUpdate(stmt, options);
+    return createMultiplexedSessionTransaction(/* singleUse = */ true)
+        .executePartitionedUpdate(stmt, options);
   }
 
   /**

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java

@@ -414,16 +414,10 @@ private List<SessionImpl> internalBatchCreateSessions(
 
   /** Returns a {@link SessionImpl} that references the existing session with the given name. */
   SessionImpl sessionWithId(String name) {
-    return sessionWithId(name, /*isMultiplexedSession= */ false);
-  }
-
-  /** Returns a {@link SessionImpl} that references the existing session with the given name. */
-  SessionImpl sessionWithId(String name, boolean isMultiplexedSession) {
     final Map<SpannerRpc.Option, ?> options;
     synchronized (this) {
       options = optionMap(SessionOption.channelHint(sessionChannelCounter++));
     }
-    return new SessionImpl(
-        spanner, new SessionReference(name, /*createTime= */ null, isMultiplexedSession, options));
+    return new SessionImpl(spanner, new SessionReference(name, options));
   }
 }

google-cloud-spanner/src/test/java/com/google/cloud/spanner/BatchClientImplTest.java

@@ -132,14 +132,12 @@ public void testBatchReadOnlyTxnWithBound() throws Exception {
     assertThat(batchTxn.getReadTimestamp()).isEqualTo(t);
     assertThat(batchTxn.getReadTimestamp())
         .isEqualTo(batchTxn.getBatchTransactionId().getTimestamp());
-    assertEquals(batchTxn.getBatchTransactionId().isMultiplexedSession(), isMultiplexedSession);
   }
 
   @Test
   public void testBatchReadOnlyTxnWithTxnId() {
     when(txnID.getSessionId()).thenReturn(SESSION_NAME);
     when(txnID.getTransactionId()).thenReturn(TXN_ID);
-    when(txnID.isMultiplexedSession()).thenReturn(isMultiplexedSession);
     Timestamp t = Timestamp.parseTimestamp(TIMESTAMP);
     when(txnID.getTimestamp()).thenReturn(t);
 
@@ -149,8 +147,6 @@ public void testBatchReadOnlyTxnWithTxnId() {
     assertThat(batchTxn.getReadTimestamp()).isEqualTo(t);
     assertThat(batchTxn.getReadTimestamp())
         .isEqualTo(batchTxn.getBatchTransactionId().getTimestamp());
-    assertThat(batchTxn.getBatchTransactionId().isMultiplexedSession())
-        .isEqualTo(isMultiplexedSession);
   }
 
   @Test