feat: support analyzeUpdate

README.md

@@ -49,7 +49,7 @@ If you are using Maven without BOM, add this to your dependencies:
 If you are using Gradle 5.x or later, add this to your dependencies
 
 ```Groovy
-implementation platform('com.google.cloud:libraries-bom:25.2.0')
+implementation platform('com.google.cloud:libraries-bom:25.3.0')
 
 implementation 'com.google.cloud:google-cloud-spanner'
 ```

google-cloud-spanner/clirr-ignored-differences.xml

@@ -70,4 +70,15 @@
     <className>com/google/cloud/spanner/DatabaseAdminClient</className>
     <method>com.google.api.gax.longrunning.OperationFuture createDatabase(java.lang.String, java.lang.String, com.google.cloud.spanner.Dialect, java.lang.Iterable)</method>
   </difference>
+  <!-- Support analyzeUpdate -->
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/TransactionContext</className>
+    <method>com.google.spanner.v1.ResultSetStats analyzeUpdate(com.google.cloud.spanner.Statement, com.google.cloud.spanner.ReadContext$QueryAnalyzeMode, com.google.cloud.spanner.Options$UpdateOption[])</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/connection/Connection</className>
+    <method>com.google.spanner.v1.ResultSetStats analyzeUpdate(com.google.cloud.spanner.Statement, com.google.cloud.spanner.ReadContext$QueryAnalyzeMode)</method>
+  </difference>
 </differences>

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

@@ -95,7 +95,7 @@ long executeStreamingPartitionedUpdate(
               resumeToken = rs.getResumeToken();
             }
             if (rs.hasStats()) {
-              foundStats = true;
+              foundStats = rs.getStats().hasRowCountLowerBound();
               updateCount += rs.getStats().getRowCountLowerBound();
             }
           }

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

@@ -67,6 +67,7 @@
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.common.util.concurrent.SettableFuture;
 import com.google.protobuf.Empty;
+import com.google.spanner.v1.ResultSetStats;
 import io.opencensus.common.Scope;
 import io.opencensus.metrics.DerivedLongCumulative;
 import io.opencensus.metrics.DerivedLongGauge;
@@ -713,6 +714,16 @@ public ApiFuture<Void> bufferAsync(Iterable<Mutation> mutations) {
       return delegate.bufferAsync(mutations);
     }
 
+    @Override
+    public ResultSetStats analyzeUpdate(
+        Statement statement, QueryAnalyzeMode analyzeMode, UpdateOption... options) {
+      try {
+        return delegate.analyzeUpdate(statement, analyzeMode, options);
+      } catch (SessionNotFoundException e) {
+        throw handler.handleSessionNotFound(e);
+      }
+    }
+
     @Override
     public long executeUpdate(Statement statement, UpdateOption... options) {
       try {

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

@@ -18,6 +18,7 @@
 
 import com.google.api.core.ApiFuture;
 import com.google.cloud.spanner.Options.UpdateOption;
+import com.google.spanner.v1.ResultSetStats;
 
 /**
  * Context for a single attempt of a locking read-write transaction. This type of transaction is the
@@ -126,6 +127,19 @@ default ApiFuture<Void> bufferAsync(Iterable<Mutation> mutations) {
    */
   ApiFuture<Long> executeUpdateAsync(Statement statement, UpdateOption... options);
 
+  /**
+   * Analyzes a DML statement and returns query plan and/or execution statistics information.
+   *
+   * <p>{@link com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#PLAN} only returns the plan for
+   * the statement. {@link com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#PROFILE} executes
+   * the DML statement, returns the modified row count and execution statistics, and the effects of
+   * the DML statement will be visible to subsequent operations in the transaction.
+   */
+  default ResultSetStats analyzeUpdate(
+      Statement statement, QueryAnalyzeMode analyzeMode, UpdateOption... options) {
+    throw new UnsupportedOperationException("method should be overwritten");
+  }
+
   /**
    * Executes a list of DML statements in a single request. The statements will be executed in order
    * and the semantics is the same as if each statement is executed by {@code executeUpdate} in a

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

@@ -43,6 +43,7 @@
 import com.google.spanner.v1.ExecuteSqlRequest;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
 import com.google.spanner.v1.RequestOptions;
+import com.google.spanner.v1.ResultSetStats;
 import com.google.spanner.v1.RollbackRequest;
 import com.google.spanner.v1.Transaction;
 import com.google.spanner.v1.TransactionOptions;
@@ -669,13 +670,39 @@ public ApiFuture<Void> bufferAsync(Iterable<Mutation> mutations) {
       return ApiFutures.immediateFuture(null);
     }
 
+    @Override
+    public ResultSetStats analyzeUpdate(
+        Statement statement, QueryAnalyzeMode analyzeMode, UpdateOption... options) {
+      Preconditions.checkNotNull(analyzeMode);
+      QueryMode queryMode;
+      switch (analyzeMode) {
+        case PLAN:
+          queryMode = QueryMode.PLAN;
+          break;
+        case PROFILE:
+          queryMode = QueryMode.PROFILE;
+          break;
+        default:
+          throw SpannerExceptionFactory.newSpannerException(
+              ErrorCode.INVALID_ARGUMENT, "Unknown analyze mode: " + analyzeMode);
+      }
+      return internalExecuteUpdate(statement, queryMode, options);
+    }
+
     @Override
     public long executeUpdate(Statement statement, UpdateOption... options) {
+      ResultSetStats resultSetStats = internalExecuteUpdate(statement, QueryMode.NORMAL, options);
+      // For standard DML, using the exact row count.
+      return resultSetStats.getRowCountExact();
+    }
+
+    private ResultSetStats internalExecuteUpdate(
+        Statement statement, QueryMode queryMode, UpdateOption... options) {
       beforeReadOrQuery();
       final ExecuteSqlRequest.Builder builder =
           getExecuteSqlRequestBuilder(
               statement,
-              QueryMode.NORMAL,
+              queryMode,
               Options.fromUpdateOptions(options),
               /* withTransactionSelector = */ true);
       try {
@@ -689,8 +716,7 @@ public long executeUpdate(Statement statement, UpdateOption... options) {
           throw new IllegalArgumentException(
               "DML response missing stats possibly due to non-DML statement as input");
         }
-        // For standard DML, using the exact row count.
-        return resultSet.getStats().getRowCountExact();
+        return resultSet.getStats();
       } catch (Throwable t) {
         throw onError(
             SpannerExceptionFactory.asSpannerException(t), builder.getTransaction().hasBegin());

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

@@ -37,6 +37,7 @@
 import com.google.cloud.spanner.TimestampBound;
 import com.google.cloud.spanner.connection.StatementResult.ResultType;
 import com.google.spanner.v1.ExecuteBatchDmlRequest;
+import com.google.spanner.v1.ResultSetStats;
 import java.util.Iterator;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
@@ -922,7 +923,8 @@ default RpcPriority getRPCPriority() {
   AsyncResultSet executeQueryAsync(Statement query, QueryOption... options);
 
   /**
-   * Analyzes a query and returns query plan and/or query execution statistics information.
+   * Analyzes a query or a DML statement and returns query plan and/or query execution statistics
+   * information.
    *
    * <p>The query plan and query statistics information is contained in {@link
    * com.google.spanner.v1.ResultSetStats} that can be accessed by calling {@link
@@ -957,6 +959,18 @@ default RpcPriority getRPCPriority() {
    */
   long executeUpdate(Statement update);
 
+  /**
+   * Analyzes a DML statement and returns query plan and/or execution statistics information.
+   *
+   * <p>{@link com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#PLAN} only returns the plan for
+   * the statement. {@link com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#PROFILE} executes
+   * the DML statement, returns the modified row count and execution statistics, and the effects of
+   * the DML statement will be visible to subsequent operations in the transaction.
+   */
+  default ResultSetStats analyzeUpdate(Statement update, QueryAnalyzeMode analyzeMode) {
+    throw new UnsupportedOperationException("Not implemented");
+  }
+
   /**
    * Executes the given statement asynchronously as a DML statement. If the statement does not
    * contain a valid DML statement, the method will throw a {@link SpannerException}.

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

@@ -48,6 +48,7 @@
 import com.google.common.base.Preconditions;
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
+import com.google.spanner.v1.ResultSetStats;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -293,6 +294,9 @@ public void close() {
   public ApiFuture<Void> closeAsync() {
     if (!isClosed()) {
       List<ApiFuture<Void>> futures = new ArrayList<>();
+      if (isBatchActive()) {
+        abortBatch();
+      }
       if (isTransactionStarted()) {
         futures.add(rollbackAsync());
       }
@@ -970,6 +974,7 @@ public long executeUpdate(Statement update) {
         "Statement is not an update statement: " + parsedStatement.getSqlWithoutComments());
   }
 
+  @Override
   public ApiFuture<Long> executeUpdateAsync(Statement update) {
     Preconditions.checkNotNull(update);
     ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
@@ -990,6 +995,27 @@ public ApiFuture<Long> executeUpdateAsync(Statement update) {
         "Statement is not an update statement: " + parsedStatement.getSqlWithoutComments());
   }
 
+  @Override
+  public ResultSetStats analyzeUpdate(Statement update, QueryAnalyzeMode analyzeMode) {
+    Preconditions.checkNotNull(update);
+    ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
+    ParsedStatement parsedStatement = getStatementParser().parse(update);
+    if (parsedStatement.isUpdate()) {
+      switch (parsedStatement.getType()) {
+        case UPDATE:
+          return get(internalAnalyzeUpdateAsync(parsedStatement, AnalyzeMode.of(analyzeMode)));
+        case CLIENT_SIDE:
+        case QUERY:
+        case DDL:
+        case UNKNOWN:
+        default:
+      }
+    }
+    throw SpannerExceptionFactory.newSpannerException(
+        ErrorCode.INVALID_ARGUMENT,
+        "Statement is not an update statement: " + parsedStatement.getSqlWithoutComments());
+  }
+
   @Override
   public long[] executeBatchUpdate(Iterable<Statement> updates) {
     Preconditions.checkNotNull(updates);
@@ -1101,7 +1127,9 @@ private ResultSet internalExecuteQuery(
       final AnalyzeMode analyzeMode,
       final QueryOption... options) {
     Preconditions.checkArgument(
-        statement.getType() == StatementType.QUERY, "Statement must be a query");
+        statement.getType() == StatementType.QUERY
+            || (statement.getType() == StatementType.UPDATE && analyzeMode != AnalyzeMode.NONE),
+        "Statement must either be a query or a DML mode with analyzeMode!=NONE");
     UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork();
     return get(
         transaction.executeQueryAsync(
@@ -1131,6 +1159,15 @@ private ApiFuture<Long> internalExecuteUpdateAsync(
         update, mergeUpdateRequestOptions(mergeUpdateStatementTag(options)));
   }
 
+  private ApiFuture<ResultSetStats> internalAnalyzeUpdateAsync(
+      final ParsedStatement update, AnalyzeMode analyzeMode, UpdateOption... options) {
+    Preconditions.checkArgument(
+        update.getType() == StatementType.UPDATE, "Statement must be an update");
+    UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork();
+    return transaction.analyzeUpdateAsync(
+        update, analyzeMode, mergeUpdateRequestOptions(mergeUpdateStatementTag(options)));
+  }
+
   private ApiFuture<long[]> internalExecuteBatchUpdateAsync(
       List<ParsedStatement> updates, UpdateOption... options) {
     UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork();

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

@@ -38,6 +38,7 @@
 import com.google.common.base.Preconditions;
 import com.google.spanner.admin.database.v1.DatabaseAdminGrpc;
 import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata;
+import com.google.spanner.v1.ResultSetStats;
 import com.google.spanner.v1.SpannerGrpc;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -201,6 +202,13 @@ public ApiFuture<Long> executeUpdateAsync(ParsedStatement update, UpdateOption..
         ErrorCode.FAILED_PRECONDITION, "Executing updates is not allowed for DDL batches.");
   }
 
+  @Override
+  public ApiFuture<ResultSetStats> analyzeUpdateAsync(
+      ParsedStatement update, AnalyzeMode analyzeMode, UpdateOption... options) {
+    throw SpannerExceptionFactory.newSpannerException(
+        ErrorCode.FAILED_PRECONDITION, "Analyzing updates is not allowed for DDL batches.");
+  }
+
   @Override
   public ApiFuture<long[]> executeBatchUpdateAsync(
       Iterable<ParsedStatement> updates, UpdateOption... options) {

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

@@ -33,6 +33,7 @@
 import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
 import com.google.common.base.Preconditions;
 import com.google.common.util.concurrent.MoreExecutors;
+import com.google.spanner.v1.ResultSetStats;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -161,6 +162,13 @@ public ApiFuture<Long> executeUpdateAsync(ParsedStatement update, UpdateOption..
     return ApiFutures.immediateFuture(-1L);
   }
 
+  @Override
+  public ApiFuture<ResultSetStats> analyzeUpdateAsync(
+      ParsedStatement update, AnalyzeMode analyzeMode, UpdateOption... options) {
+    throw SpannerExceptionFactory.newSpannerException(
+        ErrorCode.FAILED_PRECONDITION, "Analyzing updates is not allowed for DML batches.");
+  }
+
   @Override
   public ApiFuture<long[]> executeBatchUpdateAsync(
       Iterable<ParsedStatement> updates, UpdateOption... options) {

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

@@ -30,6 +30,7 @@
 import com.google.cloud.spanner.TimestampBound;
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
 import com.google.common.base.Preconditions;
+import com.google.spanner.v1.ResultSetStats;
 
 /**
  * Transaction that is used when a {@link Connection} is in read-only mode or when the transaction
@@ -163,6 +164,14 @@ public ApiFuture<Long> executeUpdateAsync(ParsedStatement update, UpdateOption..
         "Update statements are not allowed for read-only transactions");
   }
 
+  @Override
+  public ApiFuture<ResultSetStats> analyzeUpdateAsync(
+      ParsedStatement update, AnalyzeMode analyzeMode, UpdateOption... options) {
+    throw SpannerExceptionFactory.newSpannerException(
+        ErrorCode.FAILED_PRECONDITION,
+        "Analyzing updates is not allowed for read-only transactions");
+  }
+
   @Override
   public ApiFuture<long[]> executeBatchUpdateAsync(
       Iterable<ParsedStatement> updates, UpdateOption... options) {