feat: Add support for new QueryModes

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

@@ -618,6 +618,12 @@ public final ResultSet analyzeQuery(Statement statement, QueryAnalyzeMode readCo
       case PLAN:
         return executeQueryInternal(
             statement, com.google.spanner.v1.ExecuteSqlRequest.QueryMode.PLAN);
+      case WITH_STATS:
+        return executeQueryInternal(
+            statement, com.google.spanner.v1.ExecuteSqlRequest.QueryMode.WITH_STATS);
+      case WITH_PLAN_AND_STATS:
+        return executeQueryInternal(
+            statement, com.google.spanner.v1.ExecuteSqlRequest.QueryMode.WITH_PLAN_AND_STATS);
       default:
         throw new IllegalStateException(
             "Unknown value for QueryAnalyzeMode : " + readContextQueryMode);

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

@@ -33,8 +33,22 @@ public interface ReadContext extends AutoCloseable {
   enum QueryAnalyzeMode {
     /** Retrieves only the query plan information. No result data is returned. */
     PLAN,
-    /** Retrieves both query plan and query execution statistics along with the result data. */
-    PROFILE
+    /**
+     * Retrieves the query plan, overall execution statistics, operator level execution statistics
+     * along with the result data. This has a performance overhead compared to the other modes. It
+     * isn't recommended to use this mode for production traffic.
+     */
+    PROFILE,
+    /**
+     * Retrieves the overall (but not operator-level) execution statistics along with the result
+     * data.
+     */
+    WITH_STATS,
+    /**
+     * Retrieves the query plan, overall (but not operator-level) execution statistics along with
+     * the result data.
+     */
+    WITH_PLAN_AND_STATS
   }
   /**
    * Reads zero or more rows from a database.

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

@@ -63,8 +63,9 @@ public interface ResultSet extends AutoCloseable, StructReader {
   void close();
 
   /**
-   * Returns the {@link ResultSetStats} for the query only if the query was executed in either the
-   * {@code PLAN} or the {@code PROFILE} mode via the {@link ReadContext#analyzeQuery(Statement,
+   * Returns the {@link ResultSetStats} for the query only if the query was executed in {@code
+   * PLAN}, {@code PROFILE}, {@code WITH_STATS} or the {@code WITH_PLAN_AND_STATS} mode via the
+   * {@link ReadContext#analyzeQuery(Statement,
    * com.google.cloud.spanner.ReadContext.QueryAnalyzeMode)} method or for DML statements in {@link
    * ReadContext#executeQuery(Statement, QueryOption...)}. Attempts to call this method on a {@code
    * ResultSet} not obtained from {@code analyzeQuery} or {@code executeQuery} will return a {@code

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

@@ -878,6 +878,12 @@ private ResultSet internalAnalyzeStatement(
         case PROFILE:
           queryMode = QueryMode.PROFILE;
           break;
+        case WITH_STATS:
+          queryMode = QueryMode.WITH_STATS;
+          break;
+        case WITH_PLAN_AND_STATS:
+          queryMode = QueryMode.WITH_PLAN_AND_STATS;
+          break;
         default:
           throw SpannerExceptionFactory.newSpannerException(
               ErrorCode.INVALID_ARGUMENT, "Unknown analyze mode: " + analyzeMode);

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

@@ -19,14 +19,27 @@
 import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode;
 
 /**
- * {@link AnalyzeMode} indicates whether a query should be executed as a normal query (NONE),
- * whether only a query plan should be returned, or whether the query should be profiled while
- * executed.
+ * {@link AnalyzeMode} controls the execution and returned information for a query:
+ *
+ * <ul>
+ *   <li>{@code NONE}: The default mode. Only the statement results are returned.
+ *   <li>{@code PLAN}: Returns only the query plan, without any results or execution statistics
+ *       information.
+ *   <li>{@code PROFILE}: Returns the query plan, overall execution statistics, operator-level
+ *       execution statistics along with the results. This mode has a performance overhead and is
+ *       not recommended for production traffic.
+ *   <li>{@code WITH_STATS}: Returns the overall (but not operator-level) execution statistics along
+ *       with the results.
+ *   <li>{@code WITH_PLAN_AND_STATS}: Returns the query plan, overall (but not operator-level)
+ *       execution statistics along with the results.
+ * </ul>
  */
 enum AnalyzeMode {
   NONE(null),
   PLAN(QueryAnalyzeMode.PLAN),
-  PROFILE(QueryAnalyzeMode.PROFILE);
+  PROFILE(QueryAnalyzeMode.PROFILE),
+  WITH_STATS(QueryAnalyzeMode.WITH_STATS),
+  WITH_PLAN_AND_STATS(QueryAnalyzeMode.WITH_PLAN_AND_STATS);
 
   private final QueryAnalyzeMode mode;
 
@@ -45,6 +58,10 @@ static AnalyzeMode of(QueryAnalyzeMode mode) {
         return AnalyzeMode.PLAN;
       case PROFILE:
         return AnalyzeMode.PROFILE;
+      case WITH_STATS:
+        return AnalyzeMode.WITH_STATS;
+      case WITH_PLAN_AND_STATS:
+        return AnalyzeMode.WITH_PLAN_AND_STATS;
       default:
         throw new IllegalArgumentException(mode + " is unknown");
     }

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

@@ -1327,9 +1327,13 @@ PartitionedQueryResultSet runPartitionedQuery(
    * 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.
+   * the statement. {@link com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#WITH_STATS} returns
+   * the overall (but not operator-level) execution statistics. {@link
+   * com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#WITH_PLAN_AND_STATS} returns the query
+   * plan and overall (but not operator-level) execution statistics. {@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
@@ -1345,6 +1349,10 @@ default ResultSetStats analyzeUpdate(Statement update, QueryAnalyzeMode analyzeM
    *
    * <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#WITH_STATS} returns the overall (but not
+   * operator-level) execution statistics. {@link
+   * com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#WITH_PLAN_AND_STATS} returns the query
+   * plan and overall (but not operator-level) execution statistics. {@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.

google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java

@@ -3538,6 +3538,38 @@ public void testBackendQueryOptionsWithAnalyzeQuery() {
     }
   }
 
+  @Test
+  public void testWithStatsQueryModeWithAnalyzeQuery() {
+    // Use a Spanner instance with MinSession=0 to prevent background requests
+    // from the session pool interfering with the test case.
+    try (Spanner spanner =
+        SpannerOptions.newBuilder()
+            .setProjectId("[PROJECT]")
+            .setChannelProvider(channelProvider)
+            .setCredentials(NoCredentials.getInstance())
+            .setSessionPoolOption(SessionPoolOptions.newBuilder().setMinSessions(0).build())
+            .build()
+            .getService()) {
+      DatabaseClient client =
+          spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE"));
+      try (ReadOnlyTransaction tx = client.readOnlyTransaction()) {
+        try (ResultSet rs =
+            tx.analyzeQuery(
+                Statement.newBuilder(SELECT1.getSql()).build(), QueryAnalyzeMode.WITH_STATS)) {
+          // Just iterate over the results to execute the query.
+          consumeResults(rs);
+        }
+      }
+      // Check that the last query was executed using a custom optimizer version and statistics
+      // package.
+      List<AbstractMessage> requests = mockSpanner.getRequests();
+      assertThat(requests).isNotEmpty();
+      assertThat(requests.get(requests.size() - 1)).isInstanceOf(ExecuteSqlRequest.class);
+      ExecuteSqlRequest request = (ExecuteSqlRequest) requests.get(requests.size() - 1);
+      assertThat(request.getQueryMode()).isEqualTo(QueryMode.WITH_STATS);
+    }
+  }
+
   @Test
   public void testBackendPartitionQueryOptions() {
     // Use a Spanner instance with MinSession=0 to prevent background requests

google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java

@@ -470,6 +470,63 @@ public void planResult() {
     resultSet.close();
   }
 
+  @Test
+  public void withStatsResult() {
+    Map<String, com.google.protobuf.Value> statsMap =
+        ImmutableMap.of(
+            "f1", Value.string("").toProto(),
+            "f2", Value.string("").toProto());
+    ResultSetStats stats =
+        ResultSetStats.newBuilder()
+            .setQueryStats(com.google.protobuf.Struct.newBuilder().putAllFields(statsMap).build())
+            .build();
+    ArrayList<Type.StructField> dataType = new ArrayList<>();
+    dataType.add(Type.StructField.of("data", Type.string()));
+    consumer.onPartialResultSet(
+        PartialResultSet.newBuilder()
+            .setMetadata(makeMetadata(Type.struct(dataType)))
+            .addValues(Value.string("d1").toProto())
+            .setChunkedValue(false)
+            .setStats(stats)
+            .build());
+    resultSet = resultSetWithMode(QueryMode.WITH_STATS);
+    consumer.onCompleted();
+    assertThat(resultSet.next()).isTrue();
+    assertThat(resultSet.next()).isFalse();
+    ResultSetStats receivedStats = resultSet.getStats();
+    assertThat(receivedStats).isEqualTo(stats);
+    resultSet.close();
+  }
+
+  @Test
+  public void withPlanAndStatsResult() {
+    Map<String, com.google.protobuf.Value> statsMap =
+        ImmutableMap.of(
+            "f1", Value.string("").toProto(),
+            "f2", Value.string("").toProto());
+    ResultSetStats stats =
+        ResultSetStats.newBuilder()
+            .setQueryPlan(QueryPlan.newBuilder().build())
+            .setQueryStats(com.google.protobuf.Struct.newBuilder().putAllFields(statsMap).build())
+            .build();
+    ArrayList<Type.StructField> dataType = new ArrayList<>();
+    dataType.add(Type.StructField.of("data", Type.string()));
+    consumer.onPartialResultSet(
+        PartialResultSet.newBuilder()
+            .setMetadata(makeMetadata(Type.struct(dataType)))
+            .addValues(Value.string("d1").toProto())
+            .setChunkedValue(false)
+            .setStats(stats)
+            .build());
+    resultSet = resultSetWithMode(QueryMode.WITH_PLAN_AND_STATS);
+    consumer.onCompleted();
+    assertThat(resultSet.next()).isTrue();
+    assertThat(resultSet.next()).isFalse();
+    ResultSetStats receivedStats = resultSet.getStats();
+    assertThat(stats).isEqualTo(receivedStats);
+    resultSet.close();
+  }
+
   @Test
   public void statsUnavailable() {
     ResultSetStats stats = ResultSetStats.newBuilder().build();

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

@@ -272,6 +272,11 @@ public static ConnectionImpl createConnection(final ConnectionOptions options, D
         .thenReturn(select1ResultSetWithStats);
     when(singleUseReadOnlyTx.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.PROFILE))
         .thenReturn(select1ResultSetWithStats);
+    when(singleUseReadOnlyTx.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.WITH_STATS))
+        .thenReturn(select1ResultSetWithStats);
+    when(singleUseReadOnlyTx.analyzeQuery(
+            Statement.of(SELECT), QueryAnalyzeMode.WITH_PLAN_AND_STATS))
+        .thenReturn(select1ResultSetWithStats);
     when(singleUseReadOnlyTx.getReadTimestamp())
         .then(
             invocation -> {
@@ -307,6 +312,11 @@ public static ConnectionImpl createConnection(final ConnectionOptions options, D
                   .thenReturn(select1ResultSetWithStats);
               when(txContext.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.PROFILE))
                   .thenReturn(select1ResultSetWithStats);
+              when(txContext.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.WITH_STATS))
+                  .thenReturn(select1ResultSetWithStats);
+              when(txContext.analyzeQuery(
+                      Statement.of(SELECT), QueryAnalyzeMode.WITH_PLAN_AND_STATS))
+                  .thenReturn(select1ResultSetWithStats);
               when(txContext.executeUpdate(Statement.of(UPDATE))).thenReturn(1L);
               return new SimpleTransactionManager(txContext, options.isReturnCommitStats());
             });
@@ -328,6 +338,10 @@ public static ConnectionImpl createConnection(final ConnectionOptions options, D
                   .thenReturn(select1ResultSetWithStats);
               when(tx.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.PROFILE))
                   .thenReturn(select1ResultSetWithStats);
+              when(tx.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.WITH_STATS))
+                  .thenReturn(select1ResultSetWithStats);
+              when(tx.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.WITH_PLAN_AND_STATS))
+                  .thenReturn(select1ResultSetWithStats);
               when(tx.getReadTimestamp())
                   .then(
                       ignored -> {

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

@@ -358,6 +358,54 @@ public void testPlanQuery() {
     }
   }
 
+  @Test
+  public void testWithStatsQuery() {
+    for (TimestampBound staleness : getTestTimestampBounds()) {
+      ParsedStatement parsedStatement = mock(ParsedStatement.class);
+      when(parsedStatement.getType()).thenReturn(StatementType.QUERY);
+      when(parsedStatement.isQuery()).thenReturn(true);
+      Statement statement = Statement.of("SELECT * FROM FOO");
+      when(parsedStatement.getStatement()).thenReturn(statement);
+      when(parsedStatement.getSqlWithoutComments()).thenReturn(statement.getSql());
+
+      ReadOnlyTransaction transaction = createSubject(staleness);
+      ResultSet rs =
+          get(
+              transaction.executeQueryAsync(
+                  CallType.SYNC, parsedStatement, AnalyzeMode.WITH_STATS));
+      assertThat(rs, is(notNullValue()));
+      // get all results and then get the stats
+      while (rs.next()) {
+        // do nothing
+      }
+      assertThat(rs.getStats(), is(notNullValue()));
+    }
+  }
+
+  @Test
+  public void testWithPlanAndStatsQuery() {
+    for (TimestampBound staleness : getTestTimestampBounds()) {
+      ParsedStatement parsedStatement = mock(ParsedStatement.class);
+      when(parsedStatement.getType()).thenReturn(StatementType.QUERY);
+      when(parsedStatement.isQuery()).thenReturn(true);
+      Statement statement = Statement.of("SELECT * FROM FOO");
+      when(parsedStatement.getStatement()).thenReturn(statement);
+      when(parsedStatement.getSqlWithoutComments()).thenReturn(statement.getSql());
+
+      ReadOnlyTransaction transaction = createSubject(staleness);
+      ResultSet rs =
+          get(
+              transaction.executeQueryAsync(
+                  CallType.SYNC, parsedStatement, AnalyzeMode.WITH_PLAN_AND_STATS));
+      assertThat(rs, is(notNullValue()));
+      // get all results and then get the stats
+      while (rs.next()) {
+        // do nothing
+      }
+      assertThat(rs.getStats(), is(notNullValue()));
+    }
+  }
+
   @Test
   public void testProfileQuery() {
     for (TimestampBound staleness : getTestTimestampBounds()) {

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

@@ -268,6 +268,42 @@ public void testProfileQuery() {
     assertThat(rs.getStats(), is(notNullValue()));
   }
 
+  @Test
+  public void testWithStatsQuery() {
+    ParsedStatement parsedStatement = mock(ParsedStatement.class);
+    when(parsedStatement.getType()).thenReturn(StatementType.QUERY);
+    when(parsedStatement.isQuery()).thenReturn(true);
+    Statement statement = Statement.of("SELECT * FROM FOO");
+    when(parsedStatement.getStatement()).thenReturn(statement);
+
+    ReadWriteTransaction transaction = createSubject();
+    ResultSet rs =
+        get(transaction.executeQueryAsync(CallType.SYNC, parsedStatement, AnalyzeMode.WITH_STATS));
+    assertThat(rs, is(notNullValue()));
+    while (rs.next()) {
+      // do nothing
+    }
+    assertThat(rs.getStats(), is(notNullValue()));
+  }
+
+  @Test
+  public void testWithPlanAndStatsQuery() {
+    ParsedStatement parsedStatement = mock(ParsedStatement.class);
+    when(parsedStatement.getType()).thenReturn(StatementType.QUERY);
+    when(parsedStatement.isQuery()).thenReturn(true);
+    Statement statement = Statement.of("SELECT * FROM FOO");
+    when(parsedStatement.getStatement()).thenReturn(statement);
+
+    ReadWriteTransaction transaction = createSubject();
+    ResultSet rs =
+        get(transaction.executeQueryAsync(CallType.SYNC, parsedStatement, AnalyzeMode.WITH_STATS));
+    assertThat(rs, is(notNullValue()));
+    while (rs.next()) {
+      // do nothing
+    }
+    assertThat(rs.getStats(), is(notNullValue()));
+  }
+
   @Test
   public void testExecuteUpdate() {
     ParsedStatement parsedStatement = mock(ParsedStatement.class);