Handle moving 2nd getSnapshot() retry to itself

Author: odeke-em

observability-test/database.ts

@@ -605,7 +605,7 @@ describe('Database', () => {
       // pool, so that the pool can remove it from its inventory.
       const releaseStub = sandbox.stub(fakePool, 'release');
 
-      database.getSnapshot((err, snapshot) => {
+      database.getSnapshot(async (err, snapshot) => {
         assert.ifError(err);
         assert.strictEqual(snapshot, fakeSnapshot2);
         // The first session that error should already have been released back
@@ -616,8 +616,9 @@ describe('Database', () => {
         snapshot.emit('end');
         assert.strictEqual(releaseStub.callCount, 2);
 
+        await provider.forceFlush();
+        await traceExporter.forceFlush();
         const spans = traceExporter.getFinishedSpans();
-        assert.strictEqual(spans.length, 2, 'Exactly 2 spans expected');
         withAllSpansHaveDBName(spans);
 
         const actualSpanNames: string[] = [];

observability-test/spanner.ts

@@ -283,7 +283,6 @@ describe('EndToEnd', () => {
         traceExporter.forceFlush();
         const spans = traceExporter.getFinishedSpans();
         withAllSpansHaveDBName(spans);
-        console.log(`flushed spans: ${spans.toString()}`);
 
         const actualEventNames: string[] = [];
         const actualSpanNames: string[] = [];
@@ -442,10 +441,10 @@ describe('EndToEnd', () => {
 
       database.runTransaction(async (err, transaction) => {
         assert.ifError(err);
-        const [rows] = await transaction!.run('SELECT 1');
+        await transaction!.run('SELECT 1');
         await transaction!.commit();
+        await traceExporter.forceFlush();
 
-        traceExporter.forceFlush();
         const spans = traceExporter.getFinishedSpans();
         withAllSpansHaveDBName(spans);
 
@@ -872,7 +871,7 @@ describe('ObservabilityOptions injection and propagation', async () => {
           tx!.end();
 
           await tracerProvider.forceFlush();
-          traceExporter.forceFlush();
+          await traceExporter.forceFlush();
 
           const spans = traceExporter.getFinishedSpans();
           withAllSpansHaveDBName(spans);

src/database.ts

@@ -2105,19 +2105,11 @@ class Database extends common.GrpcServiceObject {
               });
               session!.lastError = err;
               this.pool_.release(session!);
-              this.getSnapshot(options, (err, snapshot) => {
-                if (err) {
-                  setSpanError(span, err);
-                  span.end();
-                } else {
-                  snapshot!.once('end', () => span.end());
-                  snapshot!.once('error', err => {
-                    setSpanError(span, err);
-                    span.end();
-                  });
-                }
-                callback!(err, snapshot);
-              });
+              this.getSnapshot(options, callback!);
+              // Explicitly requested in code review that this span.end() be
+              // moved out of this.getSnapshot, and that there will a later refactor,
+              // similar to https://github.com/googleapis/nodejs-spanner/issues/2159
+              span.end();
             } else {
               span.addEvent('Using Session', {'session.id': session?.id});
               this.pool_.release(session!);