feat: support analyzeUpdate (#1867)

* feat: support analyzeUpdate Adds support for analyzeUpdate for DML statements, similarly to analyzeQuery for DQL statements. Executing a DML statement in PLAN mode only returns the query plan, but does not actually execute the statement. Executing a DML statement in PROFILE mode executes the statement and returns the query plan and execution statistics. Fixes #1866 * test: fix integration test * build: add ignored differences to clirr * fix: plan returns zero row count * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: olavloite

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) {