feat: allow DML batches in transactions to execute analyzeUpdate (#3114)

Executing analyzeUpdate in a DML Batch inside a transaction should
be possible, as the parent transaction can be used for that.

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java

@@ -176,8 +176,6 @@ private Void endUnitOfWorkSpan() {
     return null;
   }
 
-  abstract boolean isSingleUse();
-
   /**
    * Returns a descriptive name for the type of transaction / unit of work. This is used in error
    * messages.

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractMultiUseTransaction.java

@@ -96,7 +96,7 @@ public String toString() {
   }
 
   @Override
-  boolean isSingleUse() {
+  public boolean isSingleUse() {
     return false;
   }
 

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java

@@ -97,7 +97,7 @@ private DdlBatch(Builder builder) {
   }
 
   @Override
-  boolean isSingleUse() {
+  public boolean isSingleUse() {
     return false;
   }
 

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DmlBatch.java

@@ -78,12 +78,12 @@ static Builder newBuilder() {
 
   private DmlBatch(Builder builder) {
     super(builder);
-    this.transaction = builder.transaction;
+    this.transaction = Preconditions.checkNotNull(builder.transaction);
     this.statementTag = builder.statementTag;
   }
 
   @Override
-  boolean isSingleUse() {
+  public boolean isSingleUse() {
     return false;
   }
 
@@ -174,8 +174,11 @@ public ApiFuture<Long> executeUpdateAsync(
   @Override
   public ApiFuture<ResultSet> analyzeUpdateAsync(
       CallType callType, ParsedStatement update, AnalyzeMode analyzeMode, UpdateOption... options) {
-    throw SpannerExceptionFactory.newSpannerException(
-        ErrorCode.FAILED_PRECONDITION, "Analyzing updates is not allowed for DML batches.");
+    if (transaction.isSingleUse()) {
+      throw SpannerExceptionFactory.newSpannerException(
+          ErrorCode.FAILED_PRECONDITION, "Analyzing updates is not allowed for DML batches.");
+    }
+    return transaction.analyzeUpdateAsync(callType, update, analyzeMode, options);
   }
 
   @Override

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java

@@ -180,7 +180,7 @@ private SingleUseTransaction(Builder builder) {
   }
 
   @Override
-  boolean isSingleUse() {
+  public boolean isSingleUse() {
     return true;
   }
 

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/UnitOfWork.java

@@ -77,6 +77,9 @@ public boolean isActive() {
   /** @return <code>true</code> if this unit of work is still active. */
   boolean isActive();
 
+  /** Returns true if this transaction can only be used for a single statement. */
+  boolean isSingleUse();
+
   /**
    * Commits the changes in this unit of work to the database. For read-only transactions, this only
    * closes the {@link ReadContext}. This method will throw a {@link SpannerException} if called for

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

@@ -355,8 +355,9 @@ public void testAnalyzeUpdateStatementDdlBatch() {
   }
 
   @Test
-  public void testAnalyzeUpdateDmlBatch() {
+  public void testAnalyzeUpdateDmlBatch_AutoCommit() {
     try (Connection connection = createConnection()) {
+      connection.setAutocommit(true);
       connection.startBatchDml();
 
       SpannerException exception =
@@ -371,8 +372,23 @@ public void testAnalyzeUpdateDmlBatch() {
   }
 
   @Test
-  public void testAnalyzeUpdateStatementDmlBatch() {
+  public void testAnalyzeUpdateDmlBatch_Transactional() {
     try (Connection connection = createConnection()) {
+      connection.setAutocommit(false);
+      connection.startBatchDml();
+
+      assertNotNull(connection.analyzeUpdate(PLAN_UPDATE, QueryAnalyzeMode.PLAN));
+      assertEquals(-1L, connection.executeUpdate(INSERT_STATEMENT));
+      connection.runBatch();
+
+      assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    }
+  }
+
+  @Test
+  public void testAnalyzeUpdateStatementDmlBatch_AutoCommit() {
+    try (Connection connection = createConnection()) {
+      connection.setAutocommit(true);
       connection.startBatchDml();
 
       SpannerException exception =
@@ -385,4 +401,18 @@ public void testAnalyzeUpdateStatementDmlBatch() {
     assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
     assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));
   }
+
+  @Test
+  public void testAnalyzeUpdateStatementDmlBatch_Transactional() {
+    try (Connection connection = createConnection()) {
+      connection.setAutocommit(false);
+      connection.startBatchDml();
+
+      connection.analyzeUpdateStatement(PLAN_UPDATE, QueryAnalyzeMode.PLAN);
+      assertEquals(-1L, connection.executeUpdate(INSERT_STATEMENT));
+      connection.runBatch();
+
+      assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    }
+  }
 }