fix(observability): fix the spans for the abort transaction

observability-test/database.ts

@@ -93,7 +93,7 @@
 
 class FakeSession {
   calledWith_: IArguments;
   formattedName_: any;
   constructor() {
     this.calledWith_ = arguments;
   }
@@ -682,7 +682,7 @@
           'Expected that secondRetrySpan is the child to parentSpan'
         );
 
-        const expectedEventNames = ['No session available'];
+        const expectedEventNames = ['No session available', 'Using Session'];
         assert.deepStrictEqual(
           actualEventNames,
           expectedEventNames,
@@ -1558,7 +1558,7 @@
         );
 
         // We don't expect events.
-        const expectedEventNames = [];
+        const expectedEventNames = ['Using Session'];
         assert.deepStrictEqual(
           actualEventNames,
           expectedEventNames,

observability-test/helper.ts

@@ -36,6 +36,7 @@ export const cacheSessionEvents = [
   'Acquiring session',
   'Cache hit: has usable session',
   'Acquired session',
+  'Using Session',
 ];
 
 /**
@@ -82,14 +83,25 @@ export async function verifySpansAndEvents(
       actualEventNames.push(event.name);
     });
   });
+
+  assert.strictEqual(
+    actualSpanNames.length,
+    expectedSpans.length,
+    `Span count mismatch: Expected ${expectedSpans.length} spans, but received ${actualSpanNames.length} spans`
+  );
   assert.deepStrictEqual(
     actualSpanNames,
     expectedSpans,
     `span names mismatch:\n\tGot:  ${actualSpanNames}\n\tWant: ${expectedSpans}`
   );
+  assert.strictEqual(
+    actualEventNames.length,
+    expectedEvents.length,
+    `Event count mismatch: Expected ${expectedEvents.length} events, but received ${actualEventNames.length} events`
+  );
   assert.deepStrictEqual(
-    actualEventNames,
-    expectedEvents,
+    actualEventNames.sort(),
+    expectedEvents.sort(),
     `Unexpected events:\n\tGot:  ${actualEventNames}\n\tWant: ${expectedEvents}`
   );
 }

observability-test/spanner.ts

@@ -142,6 +142,11 @@ describe('EndToEnd', async () => {
     tracerProvider: tracerProvider,
     enableExtendedTracing: false,
   });
+  let dbCounter = 1;
+
+  function newTestDatabase(): Database {
+    return instance.database(`database-${dbCounter++}`);
+  }
 
   const server = setupResult.server;
   const spannerMock = setupResult.spannerMock;
@@ -195,7 +200,6 @@ describe('EndToEnd', async () => {
             'CloudSpanner.Snapshot.begin',
             'CloudSpanner.Database.getSnapshot',
             'CloudSpanner.Snapshot.runStream',
-            'CloudSpanner.Snapshot.run',
           ];
           const expectedEventNames = [
             'Begin Transaction',
@@ -221,7 +225,7 @@ describe('EndToEnd', async () => {
         transaction!.commit();
 
         const expectedSpanNames = ['CloudSpanner.Database.getTransaction'];
-        const expectedEventNames = [...cacheSessionEvents, 'Using Session'];
+        const expectedEventNames = [...cacheSessionEvents];
         await verifySpansAndEvents(
           traceExporter,
           expectedSpanNames,
@@ -241,11 +245,7 @@ describe('EndToEnd', async () => {
             'CloudSpanner.Snapshot.runStream',
             'CloudSpanner.Database.runStream',
           ];
-          const expectedEventNames = [
-            'Starting stream',
-            ...cacheSessionEvents,
-            'Using Session',
-          ];
+          const expectedEventNames = ['Starting stream', ...cacheSessionEvents];
           await verifySpansAndEvents(
             traceExporter,
             expectedSpanNames,
@@ -263,11 +263,7 @@ describe('EndToEnd', async () => {
         'CloudSpanner.Database.runStream',
         'CloudSpanner.Database.run',
       ];
-      const expectedEventNames = [
-        'Starting stream',
-        ...cacheSessionEvents,
-        'Using Session',
-      ];
+      const expectedEventNames = ['Starting stream', ...cacheSessionEvents];
       await verifySpansAndEvents(
         traceExporter,
         expectedSpanNames,
@@ -283,16 +279,15 @@ describe('EndToEnd', async () => {
         await transaction!.end();
         const expectedSpanNames = [
           'CloudSpanner.Snapshot.runStream',
-          'CloudSpanner.Snapshot.run',
           'CloudSpanner.Transaction.commit',
           'CloudSpanner.Database.runTransaction',
         ];
         const expectedEventNames = [
           'Starting stream',
-          'Transaction Creation Done',
           'Starting Commit',
           'Commit Done',
           ...cacheSessionEvents,
+          'Transaction Creation Done',
         ];
 
         await verifySpansAndEvents(
@@ -311,14 +306,12 @@ describe('EndToEnd', async () => {
 
       const expectedSpanNames = [
         'CloudSpanner.Snapshot.runStream',
-        'CloudSpanner.Snapshot.run',
         'CloudSpanner.Database.runTransactionAsync',
       ];
       const expectedEventNames = [
         'Starting stream',
-        'Transaction Creation Done',
         ...cacheSessionEvents,
-        'Using Session',
+        'Transaction Creation Done',
       ];
       await verifySpansAndEvents(
         traceExporter,
@@ -327,6 +320,101 @@ describe('EndToEnd', async () => {
       );
     });
 
+    it.skip('runTransaction with abort', done => {
+      let attempts = 0;
+      let rowCount = 0;
+      const database = newTestDatabase();
+      database.runTransaction(async (err, transaction) => {
+        assert.ifError(err);
+        if (!attempts) {
+          spannerMock.abortTransaction(transaction!);
+        }
+        attempts++;
+        transaction!.run(selectSql, (err, rows) => {
+          assert.ifError(err);
+          rows.forEach(() => rowCount++);
+          transaction!
+            .commit()
+            .catch(done)
+            .then(async () => {
+              const expectedSpanNames = [
+                'CloudSpanner.Database.batchCreateSessions',
+                'CloudSpanner.SessionPool.createSessions',
+                'CloudSpanner.Snapshot.runStream',
+                'CloudSpanner.Snapshot.begin',
+                'CloudSpanner.Snapshot.runStream',
+                'CloudSpanner.Transaction.commit',
+                'CloudSpanner.Database.runTransaction',
+              ];
+              const expectedEventNames = [
+                ...batchCreateSessionsEvents,
+                'Starting stream',
+                'Begin Transaction',
+                'Transaction Creation Done',
+                'Starting stream',
+                'Starting Commit',
+                'Commit Done',
+                ...waitingSessionsEvents,
+                'Retrying transaction',
+              ];
+              await verifySpansAndEvents(
+                traceExporter,
+                expectedSpanNames,
+                expectedEventNames
+              );
+              database
+                .close()
+                .catch(done)
+                .then(() => done());
+            });
+        });
+      });
+    });
+
+    it('runTransactionAsync with abort', async () => {
+      let attempts = 0;
+      const database = newTestDatabase();
+      await database.runTransactionAsync((transaction): Promise<number> => {
+        if (!attempts) {
+          spannerMock.abortTransaction(transaction);
+        }
+        attempts++;
+        return transaction.run(selectSql).then(([rows]) => {
+          let count = 0;
+          rows.forEach(() => count++);
+          return transaction.commit().then(() => count);
+        });
+      });
+      assert.strictEqual(attempts, 2);
+      const expectedSpanNames = [
+        'CloudSpanner.Database.batchCreateSessions',
+        'CloudSpanner.SessionPool.createSessions',
+        'CloudSpanner.Snapshot.runStream',
+        'CloudSpanner.Snapshot.begin',
+        'CloudSpanner.Snapshot.runStream',
+        'CloudSpanner.Transaction.commit',
+        'CloudSpanner.Database.runTransactionAsync',
+      ];
+      const expectedEventNames = [
+        ...batchCreateSessionsEvents,
+        'Starting stream',
+        'Stream broken. Not safe to retry',
+        'Begin Transaction',
+        'Transaction Creation Done',
+        'Starting stream',
+        'Starting Commit',
+        'Commit Done',
+        ...waitingSessionsEvents,
+        'Retrying transaction',
+      ];
+      await verifySpansAndEvents(
+        traceExporter,
+        expectedSpanNames,
+        expectedEventNames
+      );
+      await database.close();
+    });
+
     it('writeAtLeastOnce', done => {
       const blankMutations = new MutationSet();
       database.writeAtLeastOnce(blankMutations, async (err, response) => {
@@ -340,7 +428,6 @@ describe('EndToEnd', async () => {
           'Starting Commit',
           'Commit Done',
           ...cacheSessionEvents,
-          'Using Session',
         ];
         await verifySpansAndEvents(
           traceExporter,
@@ -371,7 +458,6 @@ describe('EndToEnd', async () => {
           const expectedSpanNames = [
             'CloudSpanner.Snapshot.begin',
             'CloudSpanner.Snapshot.runStream',
-            'CloudSpanner.Snapshot.run',
             'CloudSpanner.Dml.runUpdate',
             'CloudSpanner.PartitionedDml.runUpdate',
             'CloudSpanner.Database.runPartitionedUpdate',
@@ -530,7 +616,6 @@ describe('ObservabilityOptions injection and propagation', async () => {
           const expectedSpanNames = [
             'CloudSpanner.Database.getTransaction',
             'CloudSpanner.Snapshot.runStream',
-            'CloudSpanner.Snapshot.run',
           ];
           assert.deepStrictEqual(
             actualSpanNames,
@@ -540,7 +625,6 @@ describe('ObservabilityOptions injection and propagation', async () => {
 
           const expectedEventNames = [
             ...cacheSessionEvents,
-            'Using Session',
             'Starting stream',
             'Transaction Creation Done',
           ];
@@ -585,7 +669,6 @@ describe('ObservabilityOptions injection and propagation', async () => {
           const expectedSpanNames = [
             'CloudSpanner.Snapshot.begin',
             'CloudSpanner.Snapshot.runStream',
-            'CloudSpanner.Snapshot.run',
             'CloudSpanner.Dml.runUpdate',
           ];
           assert.deepStrictEqual(
@@ -649,7 +732,6 @@ describe('ObservabilityOptions injection and propagation', async () => {
 
             const expectedEventNames = [
               ...cacheSessionEvents,
-              'Using Session',
               'Starting stream',
             ];
             assert.deepStrictEqual(
@@ -694,7 +776,6 @@ describe('ObservabilityOptions injection and propagation', async () => {
             const expectedSpanNames = [
               'CloudSpanner.Snapshot.begin',
               'CloudSpanner.Snapshot.runStream',
-              'CloudSpanner.Snapshot.run',
               'CloudSpanner.Dml.runUpdate',
               'CloudSpanner.Transaction.rollback',
             ];
@@ -1247,7 +1328,6 @@ SELECT 1p
       'CloudSpanner.Database.batchCreateSessions',
       'CloudSpanner.SessionPool.createSessions',
       'CloudSpanner.Snapshot.runStream',
-      'CloudSpanner.Snapshot.run',
       'CloudSpanner.Snapshot.begin',
       'CloudSpanner.Snapshot.begin',
       'CloudSpanner.Transaction.commit',
@@ -1259,19 +1339,6 @@ SELECT 1p
       expectedSpanNames,
       `span names mismatch:\n\tGot:  ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
     );
-    const spanSnapshotRun = spans[3];
-    assert.strictEqual(spanSnapshotRun.name, 'CloudSpanner.Snapshot.run');
-    const wantSpanErr = '6 ALREADY_EXISTS: ' + messageBadInsertAlreadyExistent;
-    assert.deepStrictEqual(
-      spanSnapshotRun.status.code,
-      SpanStatusCode.ERROR,
-      'Unexpected status code'
-    );
-    assert.deepStrictEqual(
-      spanSnapshotRun.status.message,
-      wantSpanErr,
-      'Unexpexcted error message'
-    );
 
     const databaseBatchCreateSessionsSpan = spans[0];
     assert.strictEqual(
@@ -1300,8 +1367,7 @@ SELECT 1p
 
     // We need to ensure a strict relationship between the spans.
     // |-Database.runTransactionAsync |-------------------------------------|
-    //   |-Snapshot.run                |------------------------|
-    //      |-Snapshot.runStream           |---------------------|
+    //   |-Snapshot.runStream           |---------------------|
     //   |-Transaction.commit                                 |--------|
     //      |-Snapshot.begin                                   |------|
     //       |-Snapshot.commit                                  |-----|
@@ -1322,12 +1388,6 @@ SELECT 1p
       'Expected that Database.runTransaction is the parent to Transaction.commmit'
     );
 
-    assert.deepStrictEqual(
-      spanSnapshotRun.parentSpanId,
-      spanDatabaseRunTransactionAsync.spanContext().spanId,
-      'Expected that Database.runTransaction is the parent to Snapshot.run'
-    );
-
     // Assert that despite all being exported, SessionPool.createSessions
     // is not in the same trace as runStream, createSessions is invoked at
     // Spanner Client instantiation, thus before database.run is invoked.
@@ -1848,7 +1908,6 @@ describe('Traces for ExecuteStream broken stream retries', () => {
       'CloudSpanner.Database.batchCreateSessions',
       'CloudSpanner.SessionPool.createSessions',
       'CloudSpanner.Snapshot.runStream',
-      'CloudSpanner.Snapshot.run',
       'CloudSpanner.Dml.runUpdate',
       'CloudSpanner.Snapshot.begin',
       'CloudSpanner.Transaction.commit',