chore(spanner): support AbortedException in begin

Author: harshachinta

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

@@ -17,6 +17,7 @@
 package com.google.cloud.spanner;
 
 import com.google.api.gax.rpc.ApiException;
+import com.google.protobuf.ByteString;
 import javax.annotation.Nullable;
 
 /**
@@ -31,6 +32,7 @@ public class AbortedException extends SpannerException {
    * new transaction attempt) before a retry can succeed.
    */
   private static final boolean IS_RETRYABLE = false;
+  private ByteString transactionID;
 
   /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */
   AbortedException(
@@ -46,6 +48,9 @@ public class AbortedException extends SpannerException {
       @Nullable ApiException apiException,
       @Nullable XGoogSpannerRequestId reqId) {
     super(token, ErrorCode.ABORTED, IS_RETRYABLE, message, cause, apiException, reqId);
+    if (cause instanceof AbortedException) {
+      this.transactionID = ((AbortedException) cause).getTransactionID();
+    }
   }
 
   /**
@@ -55,4 +60,12 @@ public class AbortedException extends SpannerException {
   public boolean isEmulatorOnlySupportsOneTransactionException() {
     return getMessage().endsWith("The emulator only supports one transaction at a time.");
   }
+
+  void setTransactionID(ByteString transactionID) {
+    this.transactionID = transactionID;
+  }
+
+  ByteString getTransactionID() {
+    return this.transactionID;
+  }
 }

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

@@ -900,6 +900,11 @@ public TransactionContext begin() {
       return internalBegin();
     }
 
+    @Override
+    public TransactionContext begin(AbortedException exception) {
+      return begin();
+    }
+
     private TransactionContext internalBegin() {
       TransactionContext res = new SessionPoolTransactionContext(this, delegate.begin());
       session.get().markUsed();

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

@@ -793,6 +793,7 @@ public SpannerException onError(SpannerException e, boolean withBeginTransaction
         long delay = -1L;
         if (exceptionToThrow instanceof AbortedException) {
           delay = exceptionToThrow.getRetryDelayInMillis();
+          ((AbortedException) exceptionToThrow).setTransactionID(this.transactionId);
         }
         if (delay == -1L) {
           txnLogger.log(

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java

@@ -97,6 +97,11 @@ public TransactionContext begin() {
       return txContext;
     }
 
+    @Override
+    public TransactionContext begin(AbortedException exception) {
+      return begin();
+    }
+
     @Override
     public void commit() {
       switch (commitBehavior) {

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java

@@ -36,6 +36,7 @@
 import com.google.api.core.ApiFuture;
 import com.google.api.gax.longrunning.OperationFuture;
 import com.google.cloud.Timestamp;
+import com.google.cloud.spanner.AbortedException;
 import com.google.cloud.spanner.AsyncResultSet;
 import com.google.cloud.spanner.BatchClient;
 import com.google.cloud.spanner.CommitResponse;
@@ -130,6 +131,11 @@ public TransactionContext begin() {
       return txContext;
     }
 
+    @Override
+    public TransactionContext begin(AbortedException exception) {
+      return begin();
+    }
+
     @Override
     public void commit() {
       switch (commitBehavior) {