feat(debugging): implement x-goog-spanner-request-id propagation per request

src/batch-transaction.ts

@@ -26,6 +26,7 @@ import {
   addLeaderAwareRoutingHeader,
 } from '../src/common';
 import {startTrace, setSpanError, traceConfig} from './instrument';
+import {injectRequestIDIntoHeaders} from './request_id_header';
 
 export interface TransactionIdentifier {
   session: string | Session;
@@ -157,7 +158,7 @@ class BatchTransaction extends Snapshot {
             method: 'partitionQuery',
             reqOpts,
             gaxOpts: query.gaxOptions,
-            headers: headers,
+            headers: injectRequestIDIntoHeaders(headers, this.session),
           },
           (err, partitions, resp) => {
             if (err) {
@@ -201,10 +202,11 @@ class BatchTransaction extends Snapshot {
           transaction: {id: this.id},
         });
         config.reqOpts = extend({}, query);
-        config.headers = {
-          [CLOUD_RESOURCE_HEADER]: (this.session.parent as Database)
-            .formattedName_,
+        const database = this.session.parent as Database;
+        const headers = {
+          [CLOUD_RESOURCE_HEADER]: database.formattedName_,
         };
+        config.headers = injectRequestIDIntoHeaders(headers, this.session);
         delete query.partitionOptions;
         this.session.request(config, (err, resp) => {
           if (err) {
@@ -286,14 +288,13 @@ class BatchTransaction extends Snapshot {
         if (this._getSpanner().routeToLeaderEnabled) {
           addLeaderAwareRoutingHeader(headers);
         }
-
         this.createPartitions_(
           {
             client: 'SpannerClient',
             method: 'partitionRead',
             reqOpts,
             gaxOpts: options.gaxOptions,
-            headers: headers,
+            headers: injectRequestIDIntoHeaders(headers, this.session),
           },
           (err, partitions, resp) => {
             if (err) {

src/database.ts

@@ -85,7 +85,6 @@ import {
   LongRunningCallback,
   NormalCallback,
   PagedOptionsWithFilter,
-  CLOUD_RESOURCE_HEADER,
   PagedResponse,
   RequestCallback,
   ResourceCallback,
@@ -112,6 +111,12 @@ import {
   setSpanErrorAndException,
   traceConfig,
 } from './instrument';
+import {
+  AtomicCounter,
+  X_GOOG_SPANNER_REQUEST_ID_HEADER,
+  craftRequestId,
+  newAtomicCounter,
+} from './request_id_header';
 
 export type GetDatabaseRolesCallback = RequestCallback<
   IDatabaseRole,
@@ -350,6 +355,8 @@ class Database extends common.GrpcServiceObject {
   > | null;
   _observabilityOptions?: ObservabilityOptions; // TODO: exmaine if we can remove it
   private _traceConfig: traceConfig;
+  private _nthRequest: AtomicCounter;
+  public _clientId: number;
   constructor(
     instance: Instance,
     name: string,
@@ -473,6 +480,12 @@ class Database extends common.GrpcServiceObject {
     };
 
     this.request = instance.request;
+    this._nthRequest = newAtomicCounter(0);
+    if (this.parent && this.parent.parent) {
+      this._clientId = (this.parent.parent as Spanner)._nthClientId;
+    } else {
+      this._clientId = instance._nthClientId;
+    }
     this._observabilityOptions = instance._observabilityOptions;
     this.commonHeaders_ = getCommonHeaders(
       this.formattedName_,
@@ -488,6 +501,11 @@ class Database extends common.GrpcServiceObject {
       Database.getEnvironmentQueryOptions()
     );
   }
+
+  _nextNthRequest(): number {
+    return this._nthRequest.increment();
+  }
+
   /**
    * @typedef {array} SetDatabaseMetadataResponse
    * @property {object} 0 The {@link Database} metadata.
@@ -696,14 +714,20 @@ class Database extends common.GrpcServiceObject {
       addLeaderAwareRoutingHeader(headers);
     }
 
+    const allHeaders = this._metadataWithRequestId(
+      this._nextNthRequest(),
+      1,
+      headers
+    );
+
     startTrace('Database.batchCreateSessions', this._traceConfig, span => {
       this.request<google.spanner.v1.IBatchCreateSessionsResponse>(
         {
           client: 'SpannerClient',
           method: 'batchCreateSessions',
           reqOpts,
           gaxOpts: options.gaxOptions,
-          headers: headers,
+          headers: allHeaders,
         },
         (err, resp) => {
           if (err) {
@@ -727,6 +751,26 @@ class Database extends common.GrpcServiceObject {
     });
   }
 
+  public _metadataWithRequestId(
+    nthRequest: number,
+    attempt: number,
+    priorMetadata?: {[k: string]: string}
+  ): {[k: string]: string} {
+    if (!priorMetadata) {
+      priorMetadata = {};
+    }
+    const withReqId = {
+      ...priorMetadata,
+    };
+    withReqId[X_GOOG_SPANNER_REQUEST_ID_HEADER] = craftRequestId(
+      this._clientId || 1,
+      1, // TODO: Properly infer the channelId
+      nthRequest,
+      attempt
+    );
+    return withReqId;
+  }
+
   /**
    * Get a reference to a {@link BatchTransaction} object.
    *
@@ -992,7 +1036,11 @@ class Database extends common.GrpcServiceObject {
     reqOpts.session.creatorRole =
       options.databaseRole || this.databaseRole || null;
 
-    const headers = this.commonHeaders_;
+    const headers = this._metadataWithRequestId(
+      this._nextNthRequest(),
+      1,
+      this.commonHeaders_
+    );
     if (this._getSpanner().routeToLeaderEnabled) {
       addLeaderAwareRoutingHeader(headers);
     }
@@ -1913,6 +1961,12 @@ class Database extends common.GrpcServiceObject {
       delete (gaxOpts as GetSessionsOptions).pageToken;
     }
 
+    const headers = this._metadataWithRequestId(
+      this._nextNthRequest(),
+      1,
+      this.resourceHeader_
+    );
+
     return startTrace('Database.getSessions', this._traceConfig, span => {
       this.request<
         google.spanner.v1.ISession,

src/index.ts

@@ -23,6 +23,7 @@ import {GoogleAuth, GoogleAuthOptions} from 'google-auth-library';
 import * as path from 'path';
 import {common as p} from 'protobufjs';
 import * as streamEvents from 'stream-events';
+import {EventEmitter} from 'events';
 import * as through from 'through2';
 import {
   codec,
@@ -87,6 +88,10 @@ import {
   ensureInitialContextManagerSet,
   ensureContextPropagation,
 } from './instrument';
+import {
+  injectRequestIDIntoError,
+  nextSpannerClientId,
+} from './request_id_header';
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 const gcpApiConfig = require('./spanner_grpc_config.json');
@@ -147,6 +152,7 @@ export interface SpannerOptions extends GrpcClientOptions {
   routeToLeaderEnabled?: boolean;
   directedReadOptions?: google.spanner.v1.IDirectedReadOptions | null;
   observabilityOptions?: ObservabilityOptions;
+  interceptors?: any[];
 }
 export interface RequestConfig {
   client: string;
@@ -249,6 +255,7 @@ class Spanner extends GrpcService {
   routeToLeaderEnabled = true;
   directedReadOptions: google.spanner.v1.IDirectedReadOptions | null;
   _observabilityOptions: ObservabilityOptions | undefined;
+  _nthClientId: number;
 
   /**
    * Placeholder used to auto populate a column with the commit timestamp.
@@ -310,6 +317,12 @@ class Spanner extends GrpcService {
         }
       }
     }
+
+    let interceptors: any[] = [];
+    if (options) {
+      interceptors = options.interceptors || [];
+    }
+
     options = Object.assign(
       {
         libName: 'gccl',
@@ -322,6 +335,10 @@ class Spanner extends GrpcService {
         'grpc.callInvocationTransformer': grpcGcp.gcpCallInvocationTransformer,
         'grpc.channelFactoryOverride': grpcGcp.gcpChannelFactoryOverride,
         'grpc.gcpApiConfig': grpcGcp.createGcpApiConfig(gcpApiConfig),
+
+        // TODO: Negotiate with the Google team to plumb gRPC
+        // settings such as interceptors to the gRPC client.
+        // 'grpc.interceptors': interceptors,
         grpc,
       },
       options || {}
@@ -379,6 +396,7 @@ class Spanner extends GrpcService {
     );
     ensureInitialContextManagerSet();
     ensureContextPropagation();
+    this._nthClientId = nextSpannerClientId();
   }
 
   /**
@@ -1553,7 +1571,53 @@ class Spanner extends GrpcService {
           },
         })
       );
-      callback(null, requestFn);
+
+      const wrapped = (...args) => {
+        const hasCallback =
+          args &&
+          args.length > 0 &&
+          typeof args[args.length - 1] === 'function';
+
+        switch (hasCallback) {
+          case true:
+            const cb = args[args.length - 1];
+            const priorArgs = args.slice(0, args.length - 1);
+            requestFn(...priorArgs, (...results) => {
+              if (results && results.length > 0) {
+                const err = results[0] as Error;
+                injectRequestIDIntoError(config, err);
+              }
+
+              cb(...results);
+            });
+            return;
+
+          case false:
+            const res = requestFn(...args);
+            const stream = res as EventEmitter;
+            if (stream) {
+              stream.on('error', err => {
+                injectRequestIDIntoError(config, err as Error);
+              });
+            }
+
+            const originallyPromise = res instanceof Promise;
+            if (!originallyPromise) {
+              return res;
+            }
+
+            return new Promise((resolve, reject) => {
+              requestFn(...args)
+                .then(resolve)
+                .catch(err => {
+                  injectRequestIDIntoError(config, err as Error);
+                  reject(err);
+                });
+            });
+        }
+      };
+
+      callback(null, wrapped);
     });
   }
 

src/instance.ts

@@ -980,6 +980,10 @@ class Instance extends common.GrpcServiceObject {
     if (!this.databases_.has(key!)) {
       const db = new Database(this, name, poolOptions, queryOptions);
       db._observabilityOptions = this._observabilityOptions;
+      const parent = this.parent as Spanner;
+      if (parent && parent._nthClientId) {
+        db._clientId = parent._nthClientId;
+      }
       this.databases_.set(key!, db);
     }
     return this.databases_.get(key!)!;