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 = {
+        const headers = {
           [CLOUD_RESOURCE_HEADER]: (this.session.parent as Database)
             .formattedName_,
         };
+        config.headers = injectRequestIDIntoHeaders(headers, this.session);
         delete query.partitionOptions;
         this.session.request(config, (err, resp) => {
           if (err) {
@@ -293,7 +295,7 @@ class BatchTransaction extends Snapshot {
             method: 'partitionRead',
             reqOpts,
             gaxOpts: options.gaxOptions,
-            headers: headers,
+            headers: injectRequestIDIntoHeaders(headers, this.session),
           },
           (err, partitions, resp) => {
             if (err) {

src/database.ts

@@ -113,6 +113,13 @@ import {
   setSpanErrorAndException,
   traceConfig,
 } from './instrument';
+import {
+  AtomicCounter,
+  X_GOOG_SPANNER_REQUEST_ID_HEADER,
+  craftRequestId,
+  newAtomicCounter,
+} from './request_id_header';
+
 export type GetDatabaseRolesCallback = RequestCallback<
   IDatabaseRole,
   databaseAdmin.spanner.admin.database.v1.IListDatabaseRolesResponse
@@ -355,6 +362,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,
@@ -469,6 +478,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_,
@@ -489,6 +504,11 @@ class Database extends common.GrpcServiceObject {
       Database.getEnvironmentQueryOptions()
     );
   }
+
+  _nextNthRequest(): number {
+    return this._nthRequest.increment();
+  }
+
   /**
    * @typedef {array} SetDatabaseMetadataResponse
    * @property {object} 0 The {@link Database} metadata.
@@ -697,14 +717,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) {
@@ -728,6 +754,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.
    *
@@ -993,7 +1039,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);
     }
@@ -1915,6 +1965,12 @@ class Database extends common.GrpcServiceObject {
       delete (gaxOpts as GetSessionsOptions).pageToken;
     }
 
+    const headers = this._metadataWithRequestId(
+      this._nextNthRequest(),
+      1,
+      this.commonHeaders_
+    );
+
     return startTrace('Database.getSessions', this._traceConfig, span => {
       this.request<
         google.spanner.v1.ISession,
@@ -1925,7 +1981,7 @@ class Database extends common.GrpcServiceObject {
           method: 'listSessions',
           reqOpts,
           gaxOpts,
-          headers: this.commonHeaders_,
+          headers: headers,
         },
         (err, sessions, nextPageRequest, ...args) => {
           if (err) {

src/index.ts

@@ -23,6 +23,7 @@
 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 @@
   ObservabilityOptions,
   ensureInitialContextManagerSet,
 } 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');
@@ -148,6 +153,7 @@
   directedReadOptions?: google.spanner.v1.IDirectedReadOptions | null;
   defaultTransactionOptions?: Pick<RunTransactionOptions, 'isolationLevel'>;
   observabilityOptions?: ObservabilityOptions;
+  interceptors?: any[];
 }
 export interface RequestConfig {
   client: string;
@@ -251,6 +257,7 @@
   directedReadOptions: google.spanner.v1.IDirectedReadOptions | null;
   defaultTransactionOptions: RunTransactionOptions;
   _observabilityOptions: ObservabilityOptions | undefined;
+  readonly _nthClientId: number;
 
   /**
    * Placeholder used to auto populate a column with the commit timestamp.
@@ -312,6 +319,7 @@
         }
       }
     }
+
     options = Object.assign(
       {
         libName: 'gccl',
@@ -388,6 +396,7 @@
       this._observabilityOptions?.enableEndToEndTracing
     );
     ensureInitialContextManagerSet();
+    this._nthClientId = nextSpannerClientId();
   }
 
   /**
@@ -1562,7 +1571,56 @@
           },
         })
       );
-      callback(null, requestFn);
+
+      // Wrap requestFn to inject the spanner request id into every returned error.
+      const wrappedRequestFn = (...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, wrappedRequestFn);
     });
   }
 

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!)!;

src/request_id_header.ts

@@ -187,8 +187,11 @@ export interface RequestIDError extends grpc.ServiceError {
   requestID: string;
 }
 
+const X_GOOG_REQ_ID_REGEX = /^1\.[0-9A-Fa-f]{8}(\.\d+){3}\.\d+/;
+
 export {
   AtomicCounter,
+  X_GOOG_REQ_ID_REGEX,
   X_GOOG_SPANNER_REQUEST_ID_HEADER,
   craftRequestId,
   injectRequestIDIntoError,

src/session.ts

@@ -44,7 +44,7 @@ import {
 import {grpc, CallOptions} from 'google-gax';
 import IRequestOptions = google.spanner.v1.IRequestOptions;
 import {Spanner} from '.';
-
+import {injectRequestIDIntoHeaders, nextNthRequest} from './request_id_header';
 export type GetSessionResponse = [Session, r.Response];
 
 /**
@@ -317,13 +317,19 @@ export class Session extends common.GrpcServiceObject {
     const reqOpts = {
       name: this.formattedName_,
     };
+    const database = this.parent as Database;
     return this.request(
       {
         client: 'SpannerClient',
         method: 'deleteSession',
         reqOpts,
         gaxOpts,
-        headers: this.commonHeaders_,
+        headers: injectRequestIDIntoHeaders(
+          this.commonHeaders_,
+          this,
+          nextNthRequest(database),
+          1
+        ),
       },
       callback!
     );
@@ -389,13 +395,19 @@ export class Session extends common.GrpcServiceObject {
     if (this._getSpanner().routeToLeaderEnabled) {
       addLeaderAwareRoutingHeader(headers);
     }
+    const database = this.parent as Database;
     return this.request(
       {
         client: 'SpannerClient',
         method: 'getSession',
         reqOpts,
         gaxOpts,
-        headers: headers,
+        headers: injectRequestIDIntoHeaders(
+          headers,
+          this.session,
+          nextNthRequest(database),
+          1
+        ),
       },
       (err, resp) => {
         if (resp) {
@@ -441,17 +453,25 @@ export class Session extends common.GrpcServiceObject {
       session: this.formattedName_,
       sql: 'SELECT 1',
     };
+
+    const database = this.parent as Database;
     return this.request(
       {
         client: 'SpannerClient',
         method: 'executeSql',
         reqOpts,
         gaxOpts,
-        headers: this.commonHeaders_,
+        headers: injectRequestIDIntoHeaders(
+          this.commonHeaders_,
+          this,
+          nextNthRequest(database),
+          1
+        ),
       },
       callback!
     );
   }
+
   /**
    * Create a PartitionedDml transaction.
    *
@@ -534,6 +554,11 @@ export class Session extends common.GrpcServiceObject {
   private _getSpanner(): Spanner {
     return this.parent.parent.parent as Spanner;
   }
+
+  private channelId(): number {
+    // The Node.js client does not use a gRPC channel pool, so this always returns 1.
+    return 1;
+  }
 }
 
 /*! Developer Documentation