feat: run and runStream can return query stats (#857)

* feat: run and runStream can return query stats

* fix: use const enum

* fix: add missing semicolon

Co-authored-by: skuruppu <skuruppu@google.com>

Author: olavloite

handwritten/spanner/src/database.ts

@@ -53,6 +53,7 @@ import {
   ExecuteSqlRequest,
   RunUpdateCallback,
   RunResponse,
+  RunCallback,
 } from './transaction';
 import {
   AsyncRunTransactionCallback,
@@ -112,7 +113,6 @@ export type UpdateSchemaResponse = [
 
 type PoolRequestCallback = RequestCallback<Session>;
 
-type RunCallback = RequestCallback<Row[]>;
 type ResultSetStats = spannerClient.spanner.v1.ResultSetStats;
 
 type GetSessionsOptions = PagedRequest<google.spanner.v1.IListSessionsRequest>;
@@ -1553,6 +1553,7 @@ class Database extends GrpcServiceObject {
     optionsOrCallback?: TimestampBounds | RunCallback,
     cb?: RunCallback
   ): void | Promise<RunResponse> {
+    let stats: ResultSetStats;
     const rows: Row[] = [];
     const callback =
       typeof optionsOrCallback === 'function'
@@ -1565,11 +1566,12 @@ class Database extends GrpcServiceObject {
 
     this.runStream(query, options)
       .on('error', callback!)
+      .on('stats', _stats => (stats = _stats))
       .on('data', row => {
         rows.push(row);
       })
       .on('end', () => {
-        callback!(null, rows);
+        callback!(null, rows, stats);
       });
   }
   runPartitionedUpdate(query: string | ExecuteSqlRequest): Promise<[number]>;
@@ -1783,6 +1785,7 @@ class Database extends GrpcServiceObject {
             snapshot.end();
           }
         })
+        .on('stats', stats => proxyStream.emit('stats', stats))
         .once('end', endListener)
         .pipe(proxyStream);
     });

handwritten/spanner/src/v1/spanner_client.d.ts

@@ -185,7 +185,7 @@ declare namespace SpannerClient {
     (error: null | ServiceError, response: ExecuteBatchDmlResponse): void;
   }
 
-  enum QueryMode {
+  const enum QueryMode {
     NORMAL,
     PLAN,
     PROFILE

handwritten/spanner/test/mockserver/mockspanner.ts

@@ -26,6 +26,7 @@ import ExecuteBatchDmlResponse = google.spanner.v1.ExecuteBatchDmlResponse;
 import ResultSet = google.spanner.v1.ResultSet;
 import Status = google.rpc.Status;
 import Any = google.protobuf.Any;
+import QueryMode = google.spanner.v1.ExecuteSqlRequest.QueryMode;
 
 const PROTO_PATH = 'spanner.proto';
 const IMPORT_PATH = __dirname + '/../../../protos';
@@ -474,7 +475,8 @@ export class MockSpanner {
           switch (res.type) {
             case StatementResultType.RESULT_SET:
               const partialResultSets = MockSpanner.toPartialResultSets(
-                res.resultSet
+                res.resultSet,
+                call.request.queryMode
               );
               // Resume on the next index after the last one seen by the client.
               const resumeIndex =
@@ -527,9 +529,13 @@ export class MockSpanner {
   /**
    * Splits a ResultSet into one PartialResultSet per row. This ensure that we can also test returning multiple partial results sets from a streaming method.
    * @param resultSet The ResultSet to split.
+   * @param queryMode The query mode that was used to execute the query.
    */
   private static toPartialResultSets(
-    resultSet: protobuf.ResultSet
+    resultSet: protobuf.ResultSet,
+    queryMode:
+      | google.spanner.v1.ExecuteSqlRequest.QueryMode
+      | keyof typeof google.spanner.v1.ExecuteSqlRequest.QueryMode
   ): protobuf.PartialResultSet[] {
     const res: protobuf.PartialResultSet[] = [];
     let first = true;
@@ -545,6 +551,13 @@ export class MockSpanner {
       }
       res.push(partial);
     }
+    if (queryMode === QueryMode.PROFILE || queryMode === 'PROFILE') {
+      // Include an empty query plan and statistics.
+      res[res.length - 1].stats = {
+        queryStats: {fields: {}},
+        queryPlan: {planNodes: []},
+      };
+    }
     return res;
   }
 

handwritten/spanner/test/spanner.ts

@@ -38,12 +38,14 @@ import {
   SessionPoolExhaustedError,
   SessionPoolOptions,
 } from '../src/session-pool';
-import CreateInstanceMetadata = google.spanner.admin.instance.v1.CreateInstanceMetadata;
 import {Json} from '../src/codec';
+import CreateInstanceMetadata = google.spanner.admin.instance.v1.CreateInstanceMetadata;
 import Done = Mocha.Done;
 import QueryOptions = google.spanner.v1.ExecuteSqlRequest.QueryOptions;
 import v1 = google.spanner.v1;
 import IQueryOptions = google.spanner.v1.ExecuteSqlRequest.IQueryOptions;
+import ResultSetStats = google.spanner.v1.ResultSetStats;
+import {SpannerClient as s} from '../src/v1';
 
 function numberToEnglishWord(num: number): string {
   switch (num) {
@@ -172,6 +174,97 @@ describe('Spanner with mock server', () => {
       }
     });
 
+    it('should return statistics', async () => {
+      const database = newTestDatabase();
+      try {
+        const [rows, stats] = await database.run({
+          sql: selectSql,
+          queryMode: s.QueryMode.PROFILE,
+        });
+        assert.strictEqual(rows.length, 3);
+        assert.ok(stats);
+        assert.ok(stats.queryPlan);
+      } finally {
+        await database.close();
+      }
+    });
+
+    it('should return statistics from snapshot', async () => {
+      const database = newTestDatabase();
+      try {
+        const [snapshot] = await database.getSnapshot();
+        const [rows, stats] = await snapshot.run({
+          sql: selectSql,
+          queryMode: s.QueryMode.PROFILE,
+        });
+        assert.strictEqual(rows.length, 3);
+        assert.ok(stats);
+        assert.ok(stats.queryPlan);
+        snapshot.end();
+      } finally {
+        await database.close();
+      }
+    });
+
+    it('should emit query statistics', done => {
+      const database = newTestDatabase();
+      let rowCount = 0;
+      let stats: ResultSetStats;
+      database
+        .runStream({
+          sql: selectSql,
+          queryMode: s.QueryMode.PROFILE,
+        })
+        .on('data', () => rowCount++)
+        .on('stats', _stats => (stats = _stats))
+        .on('end', () => {
+          assert.strictEqual(rowCount, 3);
+          assert.ok(stats);
+          assert.ok(stats.queryPlan);
+          database.close().then(() => done());
+        });
+    });
+
+    it('should emit query statistics from snapshot', done => {
+      const database = newTestDatabase();
+      let rowCount = 0;
+      let stats: ResultSetStats;
+      database.getSnapshot().then(response => {
+        const [snapshot] = response;
+        snapshot
+          .runStream({
+            sql: selectSql,
+            queryMode: s.QueryMode.PROFILE,
+          })
+          .on('data', () => rowCount++)
+          .on('stats', _stats => (stats = _stats))
+          .on('end', () => {
+            assert.strictEqual(rowCount, 3);
+            assert.ok(stats);
+            assert.ok(stats.queryPlan);
+            snapshot.end();
+            database.close().then(() => done());
+          });
+      });
+    });
+
+    it('should call callback with statistics', done => {
+      const database = newTestDatabase();
+      database.run(
+        {
+          sql: selectSql,
+          queryMode: s.QueryMode.PROFILE,
+        },
+        (err, rows, stats) => {
+          assert.ifError(err);
+          assert.strictEqual(rows.length, 3);
+          assert.ok(stats);
+          assert.ok(stats.queryPlan);
+          database.close().then(() => done());
+        }
+      );
+    });
+
     it('should execute update', async () => {
       const update = {
         sql: insertSql,