refactor: test

observability-test/database.ts

@@ -503,7 +503,6 @@ describe('Database', () => {
   });
 
   describe('getSnapshot', () => {
-    let fakePool: FakeSessionPool;
     let fakeSessionFactory: FakeSessionFactory;
     let fakeSession: FakeSession;
     let fakeMultiplexedSession: FakeMultiplexedSession;
@@ -2117,12 +2116,11 @@ describe('Database', () => {
       a: 'b',
       c: 'd',
     };
-
-    let fakePool: FakeSessionPool;
     let fakeSessionFactory: FakeSessionFactory;
     let fakeSession: FakeSession;
     let fakeSession2: FakeSession;
     let fakeMultiplexedSession: FakeMultiplexedSession;
+    let fakeMultiplexedSession2: FakeMultiplexedSession;
     let fakeSnapshot: FakeTransaction;
     let fakeSnapshot2: FakeTransaction;
     let fakeStream: Transform;
@@ -2142,23 +2140,32 @@ describe('Database', () => {
       beforeEach(() => {
         fakeSessionFactory = database.sessionFactory_;
         fakeMultiplexedSession = new FakeMultiplexedSession();
+        fakeMultiplexedSession2 = new FakeMultiplexedSession();
         fakeSnapshot = new FakeTransaction(
           {} as google.spanner.v1.TransactionOptions.ReadOnly
         );
+        fakeSnapshot2 = new FakeTransaction(
+          {} as google.spanner.v1.TransactionOptions.ReadOnly
+        );
         fakeStream = through.obj();
+        fakeStream2 = through.obj();
 
         getSessionStub = (
           sandbox.stub(fakeSessionFactory, 'getSession') as sinon.SinonStub
         )
           .onFirstCall()
-          .callsFake(callback => callback(null, fakeMultiplexedSession));
+          .callsFake(callback => callback(null, fakeMultiplexedSession))
+          .onSecondCall()
+          .callsFake(callback => callback(null, fakeMultiplexedSession2));
 
         (
           sandbox.stub(fakeMultiplexedSession, 'snapshot') as sinon.SinonStub
         ).returns(fakeSnapshot);
 
         sandbox.stub(fakeSnapshot, 'runStream').returns(fakeStream);
 
+        sandbox.stub(fakeSnapshot2, 'runStream').returns(fakeStream2);
+
         sandbox.stub(fakeSessionFactory, 'isMultiplexedEnabled').returns(true);
       });
 
@@ -2274,13 +2281,17 @@ describe('Database', () => {
           code: grpc.status.NOT_FOUND,
           message: 'Session not found',
         } as grpc.ServiceError;
+        const endStub = sandbox.stub(fakeSnapshot, 'end');
+        const endStub2 = sandbox.stub(fakeSnapshot2, 'end');
         let rows = 0;
 
         database
           .runStream(QUERY)
           .on('data', () => rows++)
           .on('error', err => {
             assert.strictEqual(err, sessionNotFoundError);
+            assert.strictEqual(endStub.callCount, 1);
+            assert.strictEqual(endStub2.callCount, 0);
             assert.strictEqual(rows, 0);
 
             provider.forceFlush();
@@ -2330,12 +2341,13 @@ describe('Database', () => {
           });
 
         fakeStream.emit('error', sessionNotFoundError);
+        fakeStream2.push('row1');
+        fakeStream2.push(null);
       });
     });
 
     describe('when GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSION is disable', () => {
       beforeEach(() => {
-        fakePool = database.pool_;
         fakeSessionFactory = database.sessionFactory_;
         fakeSession = new FakeSession();
         fakeSession2 = new FakeSession();
@@ -2498,7 +2510,7 @@ describe('Database', () => {
             await traceExporter.forceFlush();
 
             const spans = traceExporter.getFinishedSpans();
-            assert.strictEqual(spans.length, 2, 'Exactly 1 span expected');
+            assert.strictEqual(spans.length, 2, 'Exactly 2 spans expected');
             withAllSpansHaveDBName(spans);
 
             const actualSpanNames: string[] = [];

package.json

@@ -1,7 +1,7 @@
 {
   "name": "@google-cloud/spanner",
   "description": "Cloud Spanner Client Library for Node.js",
-  "version": "7.17.1",
+  "version": "7.16.0",
   "license": "Apache-2.0",
   "author": "Google Inc.",
   "engines": {

samples/package.json

@@ -17,7 +17,7 @@
   "dependencies": {
     "@google-cloud/kms": "^4.0.0",
     "@google-cloud/precise-date": "^4.0.0",
-    "@google-cloud/spanner": "^7.17.1",
+    "@google-cloud/spanner": "^7.16.0",
     "protobufjs": "^7.0.0",
     "yargs": "^17.0.0"
   },

src/index.ts

@@ -85,6 +85,7 @@ import * as v1 from './v1';
 import {
   ObservabilityOptions,
   ensureInitialContextManagerSet,
+  ensureContextPropagation,
 } from './instrument';
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires
@@ -377,6 +378,7 @@ class Spanner extends GrpcService {
       this._observabilityOptions?.enableEndToEndTracing
     );
     ensureInitialContextManagerSet();
+    ensureContextPropagation();
   }
 
   /**

src/instrument.ts

@@ -27,6 +27,7 @@ import {
   Span,
   SpanStatusCode,
   context,
+  propagation,
   trace,
   INVALID_SPAN_CONTEXT,
   ROOT_CONTEXT,
@@ -99,6 +100,8 @@ const {
   AsyncHooksContextManager,
 } = require('@opentelemetry/context-async-hooks');
 
+const {W3CTraceContextPropagator} = require('@opentelemetry/core');
+
 /*
  * This function ensures that async/await works correctly by
  * checking if context.active() returns an invalid/unset context
@@ -118,8 +121,14 @@ function ensureInitialContextManagerSet() {
   }
 }
 
+function ensureContextPropagation() {
+  propagation.setGlobalPropagator(new W3CTraceContextPropagator());
+}
+
 export {ensureInitialContextManagerSet};
 
+export {ensureContextPropagation};
+
 /**
  * startTrace begins an active span in the current active context
  * and passes it back to the set callback function. Each span will