From 9f35bd3fd415576cc351e1f201e19c876296838c Mon Sep 17 00:00:00 2001
From: Emmanuel T Odeke <emmanuel@orijtech.com>
Date: Mon, 7 Apr 2025 20:43:59 +0300
Subject: [PATCH] chore: add x_goog_spanner_request_id as attribute to spans

This change adds the x-goog-spanner-request-id value that is
sent as a gRPC header, but as a span attribute with the
key `x_goog_spanner_request_id` to aid in better debugging
and correlation.

Updates #2200
---
 src/index.ts             |  3 +++
 src/request_id_header.ts | 20 +++++++++++++++++
 test/spanner.ts          | 46 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 69 insertions(+)

diff --git a/src/index.ts b/src/index.ts
index 912e77395..3be6d3688 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -90,6 +90,7 @@ import {
   ensureInitialContextManagerSet,
 } from './instrument';
 import {
+  attributeXGoogSpannerRequestIdToActiveSpan,
   injectRequestIDIntoError,
   nextSpannerClientId,
 } from './request_id_header';
@@ -1562,6 +1563,8 @@ class Spanner extends GrpcService {
           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,
diff --git a/src/request_id_header.ts b/src/request_id_header.ts
index 34ca0ca21..d3f103f3f 100644
--- a/src/request_id_header.ts
+++ b/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,
diff --git a/test/spanner.ts b/test/spanner.ts
index 1b02f6b29..9838e28ec 100644
--- a/test/spanner.ts
+++ b/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.
   });