chore: add x_goog_spanner_request_id as attribute to spans

src/index.ts

@@ -90,6 +90,7 @@
   ensureInitialContextManagerSet,
 } from './instrument';
 import {
+  attributeXGoogSpannerRequestIdToActiveSpan,
   injectRequestIDIntoError,
   nextSpannerClientId,
 } from './request_id_header';
@@ -154,7 +155,7 @@
   directedReadOptions?: google.spanner.v1.IDirectedReadOptions | null;
   defaultTransactionOptions?: Pick<RunTransactionOptions, 'isolationLevel'>;
   observabilityOptions?: ObservabilityOptions;
   interceptors?: any[];
 }
 export interface RequestConfig {
   client: string;
@@ -1562,6 +1563,8 @@
           carrier[key] = value; // Set the span context (trace and span ID)
         },
       });
+      // Attach the x-goog-spanner-request-id to the currently active span.
+      attributeXGoogSpannerRequestIdToActiveSpan(config);
       const requestFn = gaxClient[config.method].bind(
         gaxClient,
         reqOpts,

src/request_id_header.ts

@@ -17,6 +17,7 @@
 import {randomBytes} from 'crypto';
 // eslint-disable-next-line n/no-extraneous-import
 import * as grpc from '@grpc/grpc-js';
+import {getActiveOrNoopSpan} from './instrument';
 const randIdForProcess = randomBytes(8)
   .readUint32LE(0)
   .toString(16)
@@ -187,12 +188,31 @@ export interface RequestIDError extends grpc.ServiceError {
   requestID: string;
 }
 
+const X_GOOG_SPANNER_REQUEST_ID_SPAN_ATTR = 'x_goog_spanner_request_id';
+
+/*
+ * attributeXGoogSpannerRequestIdToActiveSpan extracts the x-goog-spanner-request-id
+ * from config, if possible and then adds it as an attribute to the current/active span.
+ * Since x-goog-spanner-request-id is associated with RPC invoking methods, it is invoked
+ * long after tracing has been performed.
+ */
+function attributeXGoogSpannerRequestIdToActiveSpan(config: any) {
+  const reqId = extractRequestID(config);
+  if (!(reqId && reqId.length > 0)) {
+    return;
+  }
+  const span = getActiveOrNoopSpan();
+  span.setAttribute(X_GOOG_SPANNER_REQUEST_ID_SPAN_ATTR, reqId);
+}
+
 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,
+  X_GOOG_SPANNER_REQUEST_ID_SPAN_ATTR,
+  attributeXGoogSpannerRequestIdToActiveSpan,
   craftRequestId,
   injectRequestIDIntoError,
   injectRequestIDIntoHeaders,

test/spanner.ts

@@ -61,6 +61,7 @@ import {
   RequestIDError,
   X_GOOG_REQ_ID_REGEX,
   X_GOOG_SPANNER_REQUEST_ID_HEADER,
+  X_GOOG_SPANNER_REQUEST_ID_SPAN_ATTR,
   randIdForProcess,
   resetNthClientId,
 } from '../src/request_id_header';
@@ -5780,6 +5781,23 @@ describe('Spanner with mock server', () => {
   });
 
   describe('XGoogRequestId', () => {
+    const exporter = new InMemorySpanExporter();
+    const provider = new NodeTracerProvider({
+      sampler: new AlwaysOnSampler(),
+      exporter: exporter,
+    });
+    provider.addSpanProcessor(new SimpleSpanProcessor(exporter));
+    provider.register();
+
+    beforeEach(async () => {
+      await exporter.forceFlush();
+      await exporter.reset();
+    });
+
+    after(async () => {
+      await provider.shutdown();
+    });
+
     it('with retry on aborted query', async () => {
       let attempts = 0;
       const database = newTestDatabase();
@@ -5850,6 +5868,34 @@ describe('Spanner with mock server', () => {
       await database.close();
     });
 
+    it('check span attributes for x-goog-spanner-request-id', async () => {
+      const database = newTestDatabase();
+      await database.runTransactionAsync(async transaction => {
+        await transaction!.run(selectSql);
+        await transaction!.commit();
+      });
+
+      await exporter.forceFlush();
+      const spans = exporter.getFinishedSpans();
+
+      // The RPC invoking spans that we expect to have our value.
+      const rpcMakingSpans = [
+        'CloudSpanner.Database.batchCreateSessions',
+        'CloudSpanner.Snapshot.run',
+        'CloudSpanner.Transaction.commit',
+      ];
+
+      spans.forEach(span => {
+        if (rpcMakingSpans.includes(span.name)) {
+          assert.strictEqual(
+            X_GOOG_SPANNER_REQUEST_ID_SPAN_ATTR in span.attributes,
+            true,
+            `Missing ${X_GOOG_SPANNER_REQUEST_ID_SPAN_ATTR} for ${span.name}`
+          );
+        }
+      });
+    });
+
     // TODO(@odeke-em): introduce tests for incremented attempts to verify
     // that retries from GAX produce the required results.
   });