Propagate requestId into some more SpannerException values

Author: odeke-em

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

@@ -274,6 +274,7 @@ private List<Partition> partitionReadUsingIndex(
           return partitionReadUsingIndex(
               partitionOptions, table, index, keys, columns, true, option);
         }
+        e.setRequestId(reqId);
         throw e;
       }
     }
@@ -332,6 +333,7 @@ private List<Partition> partitionQuery(
         if (!isFallback && maybeMarkUnimplementedForPartitionedOps(e)) {
           return partitionQuery(partitionOptions, statement, true, option);
         }
+        e.setRequestId(reqId);
         throw e;
       }
     }

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

@@ -127,6 +127,9 @@ long executeStreamingPartitionedUpdate(
           request = newTransactionRequestFrom(statement, options);
           // Create a new xGoogSpannerRequestId.
           reqId = session.getRequestIdCreator().nextRequestId(1 /*TODO: infer channelId*/, 0);
+        } catch (SpannerException e) {
+          e.setRequestId(reqId);
+          throw e;
         }
       }
       if (!foundStats) {
@@ -226,7 +229,8 @@ private ByteString initTransaction(final Options options) {
     if (tx.getId().isEmpty()) {
       throw SpannerExceptionFactory.newSpannerException(
           ErrorCode.INTERNAL,
-          "Failed to init transaction, missing transaction id\n" + session.getName());
+          "Failed to init transaction, missing transaction id\n" + session.getName(),
+          reqId);
     }
     return tx.getId();
   }

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

@@ -187,10 +187,10 @@ private void backoffSleep(Context context, long backoffMillis) throws SpannerExc
       }
       if (latch.await(backoffMillis, TimeUnit.MILLISECONDS)) {
         // Woken by context cancellation.
-        throw newSpannerExceptionForCancellation(context, null, null /*TODO: requestId*/);
+        throw newSpannerExceptionForCancellation(context, null, this.xGoogRequestId);
       }
     } catch (InterruptedException interruptExcept) {
-      throw newSpannerExceptionForCancellation(context, interruptExcept, null /*TODO: requestId*/);
+      throw newSpannerExceptionForCancellation(context, interruptExcept, this.xGoogRequestId);
     } finally {
       context.removeListener(listener);
     }
@@ -324,6 +324,7 @@ && prepareIteratorForRetryOnDifferentGrpcChannel()) {
         }
         span.addAnnotation("Stream broken. Not safe to retry", spannerException);
         span.setStatus(spannerException);
+        spannerException.setRequestId(this.xGoogRequestId);
         throw spannerException;
       } catch (RuntimeException e) {
         span.addAnnotation("Stream broken. Not safe to retry", e);

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

@@ -57,7 +57,7 @@ public String getResourceName() {
 
   private final ErrorCode code;
   private final ApiException apiException;
-  private final XGoogSpannerRequestId requestId;
+  private XGoogSpannerRequestId requestId;
 
   /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */
   SpannerException(
@@ -197,4 +197,9 @@ public ErrorDetails getErrorDetails() {
     }
     return null;
   }
+
+  /** Sets the requestId. This method is meant to be used internally and not by customers. */
+  public void setRequestId(XGoogSpannerRequestId reqId) {
+    this.requestId = reqId;
+  }
 }