feat: Query profiling feature (#1221)

* Add the right files from the preview branch

Copy paste over the files from the preview branch

* Thread Query mode through the library

Thread the query mode through the client library. Allow calls to get made using the query mode.

* Create a basic workflow for testing profiling grpc

This code allows basic query profiling to work and creates a basic code snippet that can be used to make it work.

* This test makes a call to the mock server

This will be useful for debugging grpc issues.

* Now the test works with the mock server in repo

This mock server will be super useful for ensuring that the right data gets transmitted in private preview so that we can confirm the client library is working correctly.

* Delete js file

We use the ts file instead.

* Try explain

See if Explain gets passed along to the mock server.

* Update protos.json

Protos.json is needed to pass the plan along.

* Try an aggregate query

* Add test to catch breaking change to return type

Make sure changes don’t break the use of this callback by adding this test.

* Try redefining RunQueryCallback

Define it with a more general data structure.

* Revert Gapic changes to match head branch

The merge caused generated files to have a diff. The diff should not be in the PR.

* Remove only

only should not be used here

* Add data structures for return values

Data structures need to be added for the return values that contain stats.

* Add plumbing to send stats info back

Add stats to info if it is available. If there are no results then end the stream and send the info back. This way, stats will always be sent in info if they are available and the program won ’t break if there are no results.

* Set test to only that we intend to explore

* Add a comment about stats

Explain what happens when the result set is empty. Just send the stats back.

* Delete the mock server code

The mock server code was used for debugging and now we don’t need it since the service is working.

* Remove calls to nightly

Calls to nightly don’t have a second database to work with. Regular calls work now so nightly calls are not necessary.

* Introduce profiling tests again

Bring the query profiling tests back.

* Revert "Remove calls to nightly"

This reverts commit 040d0a5.

* Stats are optional

Stats do not necessarily come from the server.

* Write some tests to test each mode

Each query profiling mode needs to be explored.

* Add code for parsing the stats returned

Stats returned needs to be parsed by a special library that will remove the complexities of the Struct object.

* Add dependencies necessary for parsing

A library is needed for parsing the Struct values. Add the libraries and use them.

* Add assertions for the expected plan/stats

Expected plan and expected stats should be used in the tests. This ensures the tests check for the proper stats values.

* Refactor info with stats build and add info to cb

Add a specific type for the runAggregationQuery callback so that it can now support the info parameter. In order to allow runAggregationQuery to make use of creating info, we also refactor info construction into a separate function.

* Modify the parser for runAggregationQuery

The parser for runAggregationQuery should have a deterministic process for computing results.

* Add asserts for the return values of runAggregate

Make sure that the entities and the plan are correct.

* Complete tests for runQuery and aggregation query

The assertion checks for runQuery and runAggregationQuery should be done and they test each mode.

* Add tests for Query and AggregateQuery

Tests should ensure run functions work properly for the run function on both of these objects.

* Add initial transaction tests for runQuery

The runQuery tests have been added which get the right results. Next the right assert statements and info collection will be added to the tests.

* Add checks on info to the transaction tests

Checks against info are needed to be sure that stats are passed back to the caller properly.

* Fix tests for aggregate queries in transactions

Add tests for running aggregate queries inside transactions.

* Ran linter, added all tests for runQueryStream

Added a temporary test for runQueryStream. Also ran the linter.

* Change parsing of return values

Return values are going to look different for users. Change the code so that the parsing is done differently.

* Reformat the info function

This function is more readable if we eliminate some of the tertiary arguments and complex logic for building the info object.

* Change tests as a result of structure changes

The structure of the QueryInfo object is changed. Modify the tests to work with the new structure.

* Use import and not require

import is better for catching compile time errors and is more customary.

* Better spacing for imports

Change the spacing so that the imports are all in one place.

* Introduce a single function for checking the execution stats. Make sure all the tests use this function. Pull out the run query plan and the run aggregation query plan.

* Fix the tests so that they call the right fns

Add assertion checks to check the query plan against some expected value and make sure the right assertion checks are done for the right tests.

* Finish the tests for the streaming call

Finish the tests for specifying no mode, specifying normal mode, EXPLAIN mode and EXPLAIN_ANALYZE mode. Make sure the tests  pass.

* Delete code that will not be used anymore

There is a lot of boilerplate code that was needed for the streaming call. Get rid of it here.

* Make changes to match new proto

Code change to use new proto was made so that code will compile.

* Add Explain Metrics interface

Make slight change to withBeginTransaction so that code compiles under new structure. Also group plan and statistics under the new explainMetrics interface.

* Remove bytesReturned from test

Proto should not be passing along bytesReturned anymore.

* Fix system tests to use values matching new struct

* Remove calls to nightly

Feature is now fully ready so nightly tests should not be done because the feature is expected to work in production.

* Query profiling

Add test for runQuery. Send back plan summary and execution stats.

* Add a test for runAggregationQuery

runAggregationQuery needs a unit test to validate request/return data.

* Parameterize the query profiling tests

* run the linter

* Export Query Mode in index.ts

Query mode needs to be exported so that it can be accessed by the user.

* Change data structure types to match return values

* Remove TODO

* remove import

* delete the query profiling samples

* Remove abstraction for RunQueryCallback

* Change the comment to describe new data types

* Remove TODO

* linting fixes

* Update type to include runAggregationQuery

* Put else back in

This change is actually simpler because it doesn’t introduce a let. It is also a much smaller diff.

* mode is not needed in sharedQueryOptions

* Revert "mode is not needed in sharedQueryOptions"

This reverts commit b8d0c63.

* Rearrange imports

rearrange the imports to simplify the diff.

* Revert imports to simplify diff

* Don’t change Entities position

Simplify diff

* Move timestamp import back

* This interface is only needed once

Define the interface inline

* Remove QueryMode and replace with explainOptions

* A few system tests

Add a few tests for the different explain options cases.

* Add a few tests for the false case

* Add more specific types to introduced function

* mode parameter is no longer required

* This signature change is no longer required

* Update the comment for getInfoFromStats

* GapicExplainOptions are no longer needed.

* Set analyze to false to match description

* Add a test for runQueryStream and analyze set fals

* Add a test for analyze set to false

* Import ExplainOptions

* Remove bytesReturned from the interface

* Make types in test more specific

* name as string

* Rely on 2 dependencies from google-gax instead

* Change expectations in the test to reflect new val

* Ran linter

* Remove extraneous import

* Modify stats return type

src/query.ts

@@ -23,6 +23,7 @@ import {Transaction} from './transaction';
 import {CallOptions} from 'google-gax';
 import {RunQueryStreamOptions} from '../src/request';
 import * as gaxInstance from 'google-gax';
+import {google} from '../protos/protos';
 
 export type Operator =
   | '='
@@ -597,10 +598,15 @@ export interface IntegerTypeCastOptions {
   properties?: string | string[];
 }
 
+export interface ExplainOptions {
+  analyze?: boolean;
+}
+
 export interface RunQueryOptions {
   consistency?: 'strong' | 'eventual';
   readTime?: number;
   gaxOptions?: CallOptions;
+  explainOptions?: ExplainOptions;
   wrapNumbers?: boolean | IntegerTypeCastOptions;
 }
 
@@ -620,4 +626,24 @@ export interface RunQueryInfo {
     | 'MORE_RESULTS_AFTER_LIMIT'
     | 'MORE_RESULTS_AFTER_CURSOR'
     | 'NO_MORE_RESULTS';
+  explainMetrics?: ExplainMetrics;
+}
+
+export interface ExplainMetrics {
+  planSummary?: PlanSummary;
+  executionStats?: ExecutionStats;
+}
+export interface ExecutionStats {
+  resultsReturned?: number;
+  executionDuration?: google.protobuf.IDuration;
+  readOperations?: number;
+  debugStats?: {
+    [key: string]: any;
+  };
+}
+
+export interface PlanSummary {
+  indexesUsed: {
+    [key: string]: any;
+  }[];
 }

src/request.ts

@@ -44,9 +44,10 @@ import {
   KeyProto,
   ResponseResult,
   Entities,
-  ValueProto,
 } from './entity';
 import {
+  ExplainMetrics,
+  ExplainOptions,
   Query,
   QueryProto,
   RunQueryInfo,
@@ -57,6 +58,100 @@ import {
 import {Datastore} from '.';
 import ITimestamp = google.protobuf.ITimestamp;
 import {AggregateQuery} from './aggregate';
+import * as protos from '../protos/protos';
+import {serializer} from 'google-gax';
+import * as gax from 'google-gax';
+type JSONValue =
+  | string
+  | number
+  | boolean
+  | null
+  | JSONValue[]
+  | {
+      [key: string]: JSONValue;
+    };
+
+const root = gax.protobuf.loadSync('google/protobuf/struct.proto');
+const Struct = root.lookupType('Struct');
+
+// This function decodes Struct proto values
+function decodeStruct(structValue: google.protobuf.IStruct): JSONValue {
+  return serializer.toProto3JSON(Struct.fromObject(structValue));
+}
+
+// This function gets a RunQueryInfo object that contains explain metrics that
+// were returned from the server.
+function getInfoFromStats(
+  resp:
+    | protos.google.datastore.v1.IRunQueryResponse
+    | protos.google.datastore.v1.IRunAggregationQueryResponse
+): RunQueryInfo {
+  // Decode struct values stored in planSummary and executionStats
+  const explainMetrics: ExplainMetrics = {};
+  if (
+    resp &&
+    resp.explainMetrics &&
+    resp.explainMetrics.planSummary &&
+    resp.explainMetrics.planSummary.indexesUsed
+  ) {
+    Object.assign(explainMetrics, {
+      planSummary: {
+        indexesUsed: resp.explainMetrics.planSummary.indexesUsed.map(
+          (index: google.protobuf.IStruct) => decodeStruct(index)
+        ),
+      },
+    });
+  }
+  if (resp && resp.explainMetrics && resp.explainMetrics.executionStats) {
+    const executionStats = {};
+    {
+      const resultsReturned =
+        resp.explainMetrics.executionStats.resultsReturned;
+      if (resultsReturned) {
+        Object.assign(executionStats, {
+          resultsReturned:
+            typeof resultsReturned === 'string'
+              ? parseInt(resultsReturned)
+              : resultsReturned,
+        });
+      }
+    }
+    {
+      const executionDuration =
+        resp.explainMetrics.executionStats.executionDuration;
+      if (executionDuration) {
+        Object.assign(executionStats, {
+          executionDuration:
+            typeof executionDuration === 'string'
+              ? parseInt(executionDuration)
+              : executionDuration,
+        });
+      }
+    }
+    {
+      const readOperations = resp.explainMetrics.executionStats.readOperations;
+      if (readOperations) {
+        Object.assign(executionStats, {
+          readOperations:
+            typeof readOperations === 'string'
+              ? parseInt(readOperations)
+              : readOperations,
+        });
+      }
+    }
+    {
+      const debugStats = resp.explainMetrics.executionStats.debugStats;
+      if (debugStats) {
+        Object.assign(executionStats, {debugStats: decodeStruct(debugStats)});
+      }
+    }
+    Object.assign(explainMetrics, {executionStats});
+  }
+  if (explainMetrics.planSummary || explainMetrics.executionStats) {
+    return {explainMetrics};
+  }
+  return {};
+}
 
 /**
  * A map of read consistency values to proto codes.
@@ -566,12 +661,15 @@ class DatastoreRequest {
   runAggregationQuery(
     query: AggregateQuery,
     options: RunQueryOptions,
-    callback: RequestCallback
+    callback: RunAggregationQueryCallback
+  ): void;
+  runAggregationQuery(
+    query: AggregateQuery,
+    callback: RunAggregationQueryCallback
   ): void;
-  runAggregationQuery(query: AggregateQuery, callback: RequestCallback): void;
   runAggregationQuery(
     query: AggregateQuery,
-    optionsOrCallback?: RunQueryOptions | RequestCallback,
+    optionsOrCallback?: RunQueryOptions | RunAggregationQueryCallback,
     cb?: RequestCallback
   ): void | Promise<RunQueryResponse> {
     const options =
@@ -605,6 +703,7 @@ class DatastoreRequest {
         gaxOpts: options.gaxOptions,
       },
       (err, res) => {
+        const info = getInfoFromStats(res);
         if (res && res.batch) {
           const results = res.batch.aggregationResults;
           const finalResults = results
@@ -621,9 +720,9 @@ class DatastoreRequest {
                 )
               )
             );
-          callback(err, finalResults);
+          callback(err, finalResults, info);
         } else {
-          callback(err, res);
+          callback(err, [], info);
         }
       }
     );
@@ -825,9 +924,16 @@ class DatastoreRequest {
         return;
       }
 
-      const info: RunQueryInfo = {
+      if (!resp.batch) {
+        // If there are no results then send any stats back and end the stream.
+        stream.emit('info', getInfoFromStats(resp));
+        stream.push(null);
+        return;
+      }
+
+      const info = Object.assign(getInfoFromStats(resp), {
         moreResults: resp.batch.moreResults,
-      };
+      });
 
       if (resp.batch.endCursor) {
         info.endCursor = resp.batch.endCursor.toString('base64');
@@ -908,6 +1014,9 @@ class DatastoreRequest {
     options: RunQueryStreamOptions = {}
   ): SharedQueryOptions {
     const sharedQueryOpts = this.getRequestOptions(options);
+    if (options.explainOptions) {
+      sharedQueryOpts.explainOptions = options.explainOptions;
+    }
     if (query.namespace) {
       sharedQueryOpts.partitionId = {
         namespaceId: query.namespace,
@@ -1161,6 +1270,14 @@ export interface RequestCallback {
     b?: any
   ): void;
 }
+export interface RunAggregationQueryCallback {
+  (
+    a?: Error | null,
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    b?: any,
+    c?: RunQueryInfo
+  ): void;
+}
 export interface RequestConfig {
   client: string;
   gaxOpts?: CallOptions;
@@ -1170,6 +1287,7 @@ export interface RequestConfig {
 }
 export interface SharedQueryOptions {
   databaseId?: string;
+  explainOptions?: ExplainOptions;
   projectId?: string;
   partitionId?: google.datastore.v1.IPartitionId | null;
   readOptions?: {