Remove dataStream.once span nesting

Author: odeke-em

observability-test/database.ts

@@ -1878,11 +1878,14 @@ describe('Database', () => {
         .on('error', err => {
           assert.fail(err);
         })
-        .on('end', () => {
+        .on('end', async () => {
           assert.strictEqual(endStub.callCount, 1);
           assert.strictEqual(endStub2.callCount, 1);
           assert.strictEqual(rows, 1);
 
+          await provider.forceFlush();
+          await traceExporter.forceFlush();
+
           const spans = traceExporter.getFinishedSpans();
           assert.strictEqual(spans.length, 2, 'Exactly 1 span expected');
           withAllSpansHaveDBName(spans);
@@ -1907,35 +1910,35 @@ describe('Database', () => {
           );
 
           // Ensure that the span actually produced an error that was recorded.
-          const secondSpan = spans[1];
-          assert.strictEqual(
+          const lastSpan = spans[0];
+          assert.deepStrictEqual(
             SpanStatusCode.ERROR,
-            secondSpan.status.code,
+            lastSpan.status.code,
             'Expected an ERROR span status'
           );
-          assert.strictEqual(
+          assert.deepStrictEqual(
             'Session not found',
-            secondSpan.status.message,
+            lastSpan.status.message,
             'Mismatched span status message'
           );
 
           // Ensure that the final span that got retries did not error.
-          const firstSpan = spans[0];
-          assert.strictEqual(
+          const firstSpan = spans[1];
+          assert.deepStrictEqual(
             SpanStatusCode.UNSET,
             firstSpan.status.code,
-            'Unexpected an span status code'
+            'Unexpected span status code'
           );
-          assert.strictEqual(
+          assert.deepStrictEqual(
             undefined,
             firstSpan.status.message,
             'Unexpected span status message'
           );
 
           const expectedEventNames = [
-            'Using Session',
             'Using Session',
             'No session available',
+            'Using Session',
           ];
           assert.deepStrictEqual(
             actualEventNames,

src/database.ts

@@ -856,7 +856,6 @@ class Database extends common.GrpcServiceObject {
             callback!(err as ServiceError, null, undefined);
             return;
           }
-
           const transaction = this.batchTransaction(
             {session: session!},
             options
@@ -3088,11 +3087,9 @@ class Database extends common.GrpcServiceObject {
               // Create a new data stream and add it to the end user stream.
               dataStream = this.runStream(query, options);
               dataStream.pipe(proxyStream);
-              dataStream.on('end', () => span.end());
-              dataStream.on('error', err => {
-                setSpanError(span, err);
-                span.end();
-              });
+              // Explicitly invoking span.end() here,
+              // instead of inside dataStream.on('error').
+              span.end();
             } else {
               proxyStream.destroy(err);
               snapshot.end();