Add new QueryMode's support in client library

Author: nareshz

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

@@ -34,7 +34,17 @@ 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
+    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/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java

@@ -3538,6 +3538,40 @@ 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