Adjust tests to ensure transaction.end() where necessary

Author: odeke-em

observability-test/database.ts

@@ -774,6 +774,7 @@ describe('Database', () => {
           callback(null, RESPONSE);
         },
         once() {},
+        end() {},
       };
 
       database.batchTransaction = (identifier, options) => {
@@ -782,10 +783,14 @@ describe('Database', () => {
         return fakeTransaction;
       };
 
-      database.createBatchTransaction(opts, (err, transaction, resp) => {
+      database.createBatchTransaction(opts, async (err, transaction, resp) => {
         assert.strictEqual(err, null);
         assert.strictEqual(transaction, fakeTransaction);
         assert.strictEqual(resp, RESPONSE);
+        transaction!.end();
+
+        await provider.forceFlush();
+        traceExporter.forceFlush();
         const spans = traceExporter.getFinishedSpans();
         assert.strictEqual(spans.length, 1, 'Exactly 1 span expected');
         withAllSpansHaveDBName(spans);
@@ -839,8 +844,8 @@ describe('Database', () => {
         begin(callback) {
           callback(error, RESPONSE);
         },
-
         once() {},
+        end() {},
       };
 
       database.batchTransaction = () => {
@@ -926,9 +931,11 @@ describe('Database', () => {
 
       getSessionStub.callsFake(callback => callback(fakeError));
 
-      database.getTransaction(err => {
+      database.getTransaction(async err => {
         assert.strictEqual(err, fakeError);
 
+        await provider.forceFlush();
+        traceExporter.forceFlush();
         const spans = traceExporter.getFinishedSpans();
         assert.strictEqual(spans.length, 1, 'Exactly 1 span expected');
         withAllSpansHaveDBName(spans);
@@ -975,9 +982,10 @@ describe('Database', () => {
     });
 
     it('with no errors', done => {
-      database.getTransaction((err, transaction) => {
+      database.getTransaction(async (err, transaction) => {
         assert.ifError(err);
         assert.strictEqual(transaction, fakeTransaction);
+        transaction!.end();
 
         const spans = traceExporter.getFinishedSpans();
         withAllSpansHaveDBName(spans);

observability-test/spanner.ts

@@ -273,13 +273,17 @@ describe('EndToEnd', () => {
       const withAllSpansHaveDBName = generateWithAllSpansHaveDBName(
         database.formattedName_
       );
-      database.getTransaction((err, transaction) => {
+      database.getTransaction(async (err, transaction) => {
         assert.ifError(err);
         assert.ok(transaction);
+        transaction!.end();
+        transaction!.commit();
 
+        await tracerProvider.forceFlush();
         traceExporter.forceFlush();
         const spans = traceExporter.getFinishedSpans();
         withAllSpansHaveDBName(spans);
+        console.log(`flushed spans: ${spans.toString()}`);
 
         const actualEventNames: string[] = [];
         const actualSpanNames: string[] = [];
@@ -807,7 +811,10 @@ describe('ObservabilityOptions injection and propagation', async () => {
       database.getTransaction((err, tx) => {
         assert.ifError(err);
 
-        tx!.run('SELECT 1', (err, rows) => {
+        tx!.run('SELECT 1', async (err, rows) => {
+          tx!.end();
+
+          await tracerProvider.forceFlush();
           traceExporter.forceFlush();
 
           const spans = traceExporter.getFinishedSpans();
@@ -860,9 +867,11 @@ describe('ObservabilityOptions injection and propagation', async () => {
         traceExporter.reset();
 
         tx!.begin();
-        tx!.runUpdate(updateSql, (err, rowCount) => {
+        tx!.runUpdate(updateSql, async (err, rowCount) => {
           assert.ifError(err);
+          tx!.end();
 
+          await tracerProvider.forceFlush();
           traceExporter.forceFlush();
 
           const spans = traceExporter.getFinishedSpans();
@@ -914,9 +923,10 @@ describe('ObservabilityOptions injection and propagation', async () => {
           .on('data', () => rowCount++)
           .on('error', assert.ifError)
           .on('stats', _stats => {})
-          .on('end', () => {
+          .on('end', async () => {
             tx!.end();
 
+            await tracerProvider.forceFlush();
             traceExporter.forceFlush();
 
             const spans = traceExporter.getFinishedSpans();
@@ -970,7 +980,9 @@ describe('ObservabilityOptions injection and propagation', async () => {
 
         tx!.runUpdate(updateSql, async (err, rowCount) => {
           assert.ifError(err);
-          tx!.rollback(err => {
+          tx!.rollback(async err => {
+            tx!.end();
+            await tracerProvider.forceFlush();
             traceExporter.forceFlush();
 
             const spans = traceExporter.getFinishedSpans();
@@ -1557,9 +1569,9 @@ SELECT 1p
     const expectedSpanNames = [
       'CloudSpanner.Database.batchCreateSessions',
       'CloudSpanner.SessionPool.createSessions',
+      'CloudSpanner.Snapshot.runStream',
       'CloudSpanner.Snapshot.run',
       'CloudSpanner.Snapshot.begin',
-      'CloudSpanner.Snapshot.runStream',
       'CloudSpanner.Snapshot.begin',
       'CloudSpanner.Transaction.commit',
       'CloudSpanner.Transaction.commit',
@@ -1570,7 +1582,7 @@ SELECT 1p
       expectedSpanNames,
       `span names mismatch:\n\tGot:  ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
     );
-    const spanSnapshotRun = spans[2];
+    const spanSnapshotRun = spans[3];
     assert.strictEqual(spanSnapshotRun.name, 'CloudSpanner.Snapshot.run');
     const wantSpanErr = '6 ALREADY_EXISTS: ' + messageBadInsertAlreadyExistent;
     assert.deepStrictEqual(
@@ -1653,12 +1665,12 @@ SELECT 1p
       'Requesting 25 sessions',
       'Creating 25 sessions',
       'Requested for 25 sessions returned 25',
-      'Begin Transaction',
-      'Transaction Creation Done',
       'Starting stream',
       'Stream broken. Safe to retry',
       'Begin Transaction',
       'Transaction Creation Done',
+      'Begin Transaction',
+      'Transaction Creation Done',
       'Starting Commit',
       'Commit Done',
       'Acquiring session',

src/database.ts

@@ -856,6 +856,7 @@ class Database extends common.GrpcServiceObject {
             callback!(err as ServiceError, null, undefined);
             return;
           }
+
           const transaction = this.batchTransaction(
             {session: session!},
             options
@@ -874,6 +875,7 @@ class Database extends common.GrpcServiceObject {
               return;
             }
             span.addEvent('Using Session', {'session.id': session?.id});
+            span.end();
             callback!(null, transaction, resp!);
           });
         });
@@ -1130,7 +1132,9 @@ class Database extends common.GrpcServiceObject {
         setSpanErrorAndException(span, e as Error);
         this.emit('error', e);
       } finally {
-        span.end();
+        if (span.isRecording()) {
+          span.end();
+        }
       }
     });
   }
@@ -2213,11 +2217,10 @@ class Database extends common.GrpcServiceObject {
             'session.id': session?.id,
           });
           setSpanError(span, err);
-          span.end();
         } else {
           setSpanError(span, err);
-          span.end();
         }
+        span.end();
         cb!(err as grpc.ServiceError | null, transaction);
       });
     });

src/transaction.ts

@@ -1370,12 +1370,14 @@ export class Snapshot extends EventEmitter {
           span.end();
         });
 
-      finished(resultStream, err => {
-        if (err) {
-          setSpanError(span, err);
-        }
-        span.end();
-      });
+      if (resultStream instanceof Stream) {
+        finished(resultStream, err => {
+          if (err) {
+            setSpanError(span, err);
+          }
+          span.end();
+        });
+      }
       return resultStream;
     });
   }
@@ -2133,7 +2135,7 @@ export class Transaction extends Dml {
       opts: this._observabilityOptions,
       dbName: this._dbName!,
     };
-    startTrace('Transaction.commit', traceConfig, async span => {
+    return startTrace('Transaction.commit', traceConfig, span => {
       if (this.id) {
         reqOpts.transactionId = this.id as Uint8Array;
       } else if (!this._useInRunner) {