feat: analyze update returns param types (#2156)

* feat: analyzeStatement for both DML and queries

Add support for `analyzeStatement` that works for both DML and queries.

* feat: analyze update returns param types

Add the method analyzeUpdateStatement that returns the undeclared
parameters in an update statement. This allows connection based APIs
to return more metadata for a statement than is currently possible:
1. JDBC should return the parameter data types when PreparedStatement#getMetaData()
   is called.
2. PGAdapter should return the parameter types when a DescribeStatement
   message is received.

* chore: add clirr differences

* chore: address review comments

Author: olavloite

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

@@ -207,4 +207,19 @@
     <className>com/google/cloud/spanner/connection/AbstractStatementParser</className>
     <method>boolean checkReturningClauseInternal(java.lang.String)</method>
   </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/ResultSet</className>
+    <method>com.google.spanner.v1.ResultSetMetadata getMetadata()</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/TransactionContext</className>
+    <method>com.google.cloud.spanner.ResultSet analyzeUpdateStatement(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.cloud.spanner.ResultSet analyzeUpdateStatement(com.google.cloud.spanner.Statement, com.google.cloud.spanner.ReadContext$QueryAnalyzeMode, com.google.cloud.spanner.Options$UpdateOption[])</method>
+  </difference>
 </differences>

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

@@ -93,6 +93,7 @@ void onTransactionMetadata(Transaction transaction, boolean shouldIncludeId)
   static class GrpcResultSet extends AbstractResultSet<List<Object>> {
     private final GrpcValueIterator iterator;
     private final Listener listener;
+    private ResultSetMetadata metadata;
     private GrpcStruct currRow;
     private SpannerException error;
     private ResultSetStats statistics;
@@ -117,7 +118,7 @@ public boolean next() throws SpannerException {
       }
       try {
         if (currRow == null) {
-          ResultSetMetadata metadata = iterator.getMetadata();
+          metadata = iterator.getMetadata();
           if (metadata.hasTransaction()) {
             listener.onTransactionMetadata(
                 metadata.getTransaction(), iterator.isWithBeginTransaction());
@@ -146,6 +147,12 @@ public ResultSetStats getStats() {
       return statistics;
     }
 
+    @Override
+    public ResultSetMetadata getMetadata() {
+      checkState(metadata != null, "next() call required");
+      return metadata;
+    }
+
     @Override
     public void close() {
       listener.onDone(iterator.isWithBeginTransaction());

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

@@ -29,6 +29,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.ListeningScheduledExecutorService;
 import com.google.common.util.concurrent.MoreExecutors;
+import com.google.spanner.v1.ResultSetMetadata;
 import com.google.spanner.v1.ResultSetStats;
 import java.util.Collection;
 import java.util.LinkedList;
@@ -572,6 +573,11 @@ public ResultSetStats getStats() {
     return delegateResultSet.get().getStats();
   }
 
+  @Override
+  public ResultSetMetadata getMetadata() {
+    return delegateResultSet.get().getMetadata();
+  }
+
   @Override
   protected void checkValidState() {
     synchronized (monitor) {

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

@@ -19,6 +19,7 @@
 import com.google.common.base.Preconditions;
 import com.google.common.base.Supplier;
 import com.google.common.base.Suppliers;
+import com.google.spanner.v1.ResultSetMetadata;
 import com.google.spanner.v1.ResultSetStats;
 
 /** Forwarding implementation of ResultSet that forwards all calls to a delegate. */
@@ -76,4 +77,9 @@ public void close() {
   public ResultSetStats getStats() {
     return delegate.get().getStats();
   }
+
+  @Override
+  public ResultSetMetadata getMetadata() {
+    return delegate.get().getMetadata();
+  }
 }

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

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.spanner;
+
+import com.google.spanner.v1.ResultSet;
+import com.google.spanner.v1.ResultSetMetadata;
+import com.google.spanner.v1.ResultSetStats;
+import java.util.List;
+import javax.annotation.Nullable;
+
+class NoRowsResultSet extends AbstractResultSet<List<Object>> {
+  private final ResultSetStats stats;
+  private final ResultSetMetadata metadata;
+
+  NoRowsResultSet(ResultSet resultSet) {
+    this.stats = resultSet.getStats();
+    this.metadata = resultSet.getMetadata();
+  }
+
+  @Override
+  protected GrpcStruct currRow() {
+    throw SpannerExceptionFactory.newSpannerException(
+        ErrorCode.FAILED_PRECONDITION, "This result set has no rows");
+  }
+
+  @Override
+  public boolean next() throws SpannerException {
+    return false;
+  }
+
+  @Override
+  public void close() {}
+
+  @Nullable
+  @Override
+  public ResultSetStats getStats() {
+    return stats;
+  }
+
+  @Override
+  public ResultSetMetadata getMetadata() {
+    return metadata;
+  }
+
+  @Override
+  public Type getType() {
+    return Type.struct();
+  }
+}

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

@@ -17,6 +17,7 @@
 package com.google.cloud.spanner;
 
 import com.google.cloud.spanner.Options.QueryOption;
+import com.google.spanner.v1.ResultSetMetadata;
 import com.google.spanner.v1.ResultSetStats;
 import javax.annotation.Nullable;
 
@@ -73,4 +74,12 @@ public interface ResultSet extends AutoCloseable, StructReader {
    */
   @Nullable
   ResultSetStats getStats();
+
+  /**
+   * Returns the {@link ResultSetMetadata} for this {@link ResultSet}. This is method may only be
+   * called after calling {@link ResultSet#next()} at least once.
+   */
+  default ResultSetMetadata getMetadata() {
+    throw new UnsupportedOperationException("Method should be overridden");
+  }
 }

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

@@ -29,6 +29,7 @@
 import com.google.common.base.Supplier;
 import com.google.common.collect.Lists;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
+import com.google.spanner.v1.ResultSetMetadata;
 import com.google.spanner.v1.ResultSetStats;
 import java.math.BigDecimal;
 import java.util.List;
@@ -158,6 +159,12 @@ public ResultSetStats getStats() {
           "ResultSetStats are available only for results returned from analyzeQuery() calls");
     }
 
+    @Override
+    public ResultSetMetadata getMetadata() {
+      throw new UnsupportedOperationException(
+          "ResultSetMetadata are available only for results that were returned from Cloud Spanner");
+    }
+
     @Override
     public int getColumnCount() {
       return getType().getStructFields().size();

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

@@ -717,8 +717,14 @@ public ApiFuture<Void> bufferAsync(Iterable<Mutation> mutations) {
     @Override
     public ResultSetStats analyzeUpdate(
         Statement statement, QueryAnalyzeMode analyzeMode, UpdateOption... options) {
+      return analyzeUpdateStatement(statement, analyzeMode, options).getStats();
+    }
+
+    @Override
+    public ResultSet analyzeUpdateStatement(
+        Statement statement, QueryAnalyzeMode analyzeMode, UpdateOption... options) {
       try {
-        return delegate.analyzeUpdate(statement, analyzeMode, options);
+        return delegate.analyzeUpdateStatement(statement, analyzeMode, options);
       } catch (SessionNotFoundException e) {
         throw handler.handleSessionNotFound(e);
       }

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

@@ -135,12 +135,31 @@ default ApiFuture<Void> bufferAsync(Iterable<Mutation> mutations) {
    * 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.
+   *
+   * @deprecated Use {@link #analyzeUpdateStatement(Statement, QueryAnalyzeMode, UpdateOption...)}
+   *     instead to get both statement plan and parameter metadata
    */
+  @Deprecated
   default ResultSetStats analyzeUpdate(
       Statement statement, QueryAnalyzeMode analyzeMode, UpdateOption... options) {
     throw new UnsupportedOperationException("method should be overwritten");
   }
 
+  /**
+   * Analyzes a DML statement and returns query plan and statement parameter metadata and optionally
+   * execution statistics information.
+   *
+   * <p>{@link com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#PLAN} only returns the plan and
+   * parameter metadata 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 ResultSet analyzeUpdateStatement(
+      Statement statement, QueryAnalyzeMode analyzeMode, UpdateOption... options) {
+    throw new UnsupportedOperationException("method should be overwritten");
+  }
+
   /**
    * Executes a list of DML statements (which can include simple DML statements or DML statements
    * with returning clause) in a single request. The statements will be executed in order and the

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.ResultSet;
 import com.google.spanner.v1.ResultSetStats;
 import com.google.spanner.v1.RollbackRequest;
 import com.google.spanner.v1.Transaction;
@@ -673,6 +674,17 @@ public ApiFuture<Void> bufferAsync(Iterable<Mutation> mutations) {
     @Override
     public ResultSetStats analyzeUpdate(
         Statement statement, QueryAnalyzeMode analyzeMode, UpdateOption... options) {
+      return internalAnalyzeStatement(statement, analyzeMode, options).getStats();
+    }
+
+    @Override
+    public com.google.cloud.spanner.ResultSet analyzeUpdateStatement(
+        Statement statement, QueryAnalyzeMode analyzeMode, UpdateOption... options) {
+      return new NoRowsResultSet(internalAnalyzeStatement(statement, analyzeMode, options));
+    }
+
+    private ResultSet internalAnalyzeStatement(
+        Statement statement, QueryAnalyzeMode analyzeMode, UpdateOption... options) {
       Preconditions.checkNotNull(analyzeMode);
       QueryMode queryMode;
       switch (analyzeMode) {
@@ -691,12 +703,12 @@ public ResultSetStats analyzeUpdate(
 
     @Override
     public long executeUpdate(Statement statement, UpdateOption... options) {
-      ResultSetStats resultSetStats = internalExecuteUpdate(statement, QueryMode.NORMAL, options);
+      ResultSet resultSet = internalExecuteUpdate(statement, QueryMode.NORMAL, options);
       // For standard DML, using the exact row count.
-      return resultSetStats.getRowCountExact();
+      return resultSet.getStats().getRowCountExact();
     }
 
-    private ResultSetStats internalExecuteUpdate(
+    private ResultSet internalExecuteUpdate(
         Statement statement, QueryMode queryMode, UpdateOption... options) {
       beforeReadOrQuery();
       final ExecuteSqlRequest.Builder builder =
@@ -716,7 +728,7 @@ private ResultSetStats internalExecuteUpdate(
           throw new IllegalArgumentException(
               "DML response missing stats possibly due to non-DML statement as input");
         }
-        return resultSet.getStats();
+        return resultSet;
       } catch (Throwable t) {
         throw onError(
             SpannerExceptionFactory.asSpannerException(t), builder.getTransaction().hasBegin());

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

@@ -29,6 +29,7 @@
 import com.google.cloud.spanner.Mutation;
 import com.google.cloud.spanner.Options.QueryOption;
 import com.google.cloud.spanner.Options.RpcPriority;
+import com.google.cloud.spanner.Options.UpdateOption;
 import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode;
 import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.SpannerBatchUpdateException;
@@ -968,11 +969,30 @@ default RpcPriority getRPCPriority() {
    * 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.
+   *
+   * @deprecated Use {@link #analyzeUpdateStatement(Statement, QueryAnalyzeMode, UpdateOption...)}
+   *     instead
    */
+  @Deprecated
   default ResultSetStats analyzeUpdate(Statement update, QueryAnalyzeMode analyzeMode) {
     throw new UnsupportedOperationException("Not implemented");
   }
 
+  /**
+   * Analyzes a DML statement and returns execution plan, undeclared parameters and optionally
+   * execution statistics information.
+   *
+   * <p>{@link com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#PLAN} only returns the plan and
+   * undeclared parameters for the statement. {@link
+   * com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#PROFILE} also 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 ResultSet analyzeUpdateStatement(
+      Statement statement, QueryAnalyzeMode analyzeMode, UpdateOption... options) {
+    throw new UnsupportedOperationException("Not implemented");
+  }
+
   /**
    * Executes the given statement asynchronously as a simple DML statement. If the statement does
    * not contain a simple DML statement, the method will throw a {@link SpannerException}. A DML

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

@@ -1065,7 +1065,8 @@ public ResultSetStats analyzeUpdate(Statement update, QueryAnalyzeMode analyzeMo
     if (parsedStatement.isUpdate()) {
       switch (parsedStatement.getType()) {
         case UPDATE:
-          return get(internalAnalyzeUpdateAsync(parsedStatement, AnalyzeMode.of(analyzeMode)));
+          return get(internalAnalyzeUpdateAsync(parsedStatement, AnalyzeMode.of(analyzeMode)))
+              .getStats();
         case CLIENT_SIDE:
         case QUERY:
         case DDL:
@@ -1078,6 +1079,27 @@ public ResultSetStats analyzeUpdate(Statement update, QueryAnalyzeMode analyzeMo
         "Statement is not an update statement: " + parsedStatement.getSqlWithoutComments());
   }
 
+  @Override
+  public ResultSet analyzeUpdateStatement(
+      Statement statement, QueryAnalyzeMode analyzeMode, UpdateOption... options) {
+    Preconditions.checkNotNull(statement);
+    ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
+    ParsedStatement parsedStatement = getStatementParser().parse(statement);
+    switch (parsedStatement.getType()) {
+      case UPDATE:
+        return get(
+            internalAnalyzeUpdateAsync(parsedStatement, AnalyzeMode.of(analyzeMode), options));
+      case QUERY:
+      case CLIENT_SIDE:
+      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);
@@ -1224,7 +1246,7 @@ private ApiFuture<Long> internalExecuteUpdateAsync(
         update, mergeUpdateRequestOptions(mergeUpdateStatementTag(options)));
   }
 
-  private ApiFuture<ResultSetStats> internalAnalyzeUpdateAsync(
+  private ApiFuture<ResultSet> internalAnalyzeUpdateAsync(
       final ParsedStatement update, AnalyzeMode analyzeMode, UpdateOption... options) {
     Preconditions.checkArgument(
         update.getType() == StatementType.UPDATE, "Statement must be an update");

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

@@ -38,7 +38,6 @@
 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;
@@ -203,7 +202,7 @@ public ApiFuture<Long> executeUpdateAsync(ParsedStatement update, UpdateOption..
   }
 
   @Override
-  public ApiFuture<ResultSetStats> analyzeUpdateAsync(
+  public ApiFuture<ResultSet> analyzeUpdateAsync(
       ParsedStatement update, AnalyzeMode analyzeMode, UpdateOption... options) {
     throw SpannerExceptionFactory.newSpannerException(
         ErrorCode.FAILED_PRECONDITION, "Analyzing updates is not allowed for DDL batches.");