fix(Spanner): retry and resume BatchWrite streams on UNAVAILABLE

handwritten/spanner/observability-test/database.ts

@@ -95,7 +95,7 @@
 
 class FakeSession {
   calledWith_: IArguments;
   formattedName_: any;
   constructor() {
     this.calledWith_ = arguments;
   }
@@ -142,7 +142,7 @@
   setQueuedMutations(mutation) {
     this._queuedMutations = mutation;
   }
   setReadWriteTransactionOptions(options: RunTransactionOptions) {}
   commit(
     options?: CommitOptions,
     callback?: CommitCallback,
@@ -1285,6 +1285,66 @@
       fakeDataStream.emit('data', 'response');
       fakeDataStream.end('end');
     });
+
+    it('ends span and does not start RPC when stream is destroyed before getSession resolves', done => {
+      let sessionCallback:
+        ((err: Error | null, session?: FakeSession) => void) | null = null;
+      getSessionStub.callsFake(callback => {
+        sessionCallback = callback;
+      });
+
+      const stream = database.batchWriteAtLeastOnce(mutationGroups, options);
+      stream.destroy();
+
+      setImmediate(() => {
+        sessionCallback!(null, fakeSession);
+        const spans = traceExporter.getFinishedSpans();
+        assert.strictEqual(spans.length, 1, 'Exactly 1 span expected');
+        withAllSpansHaveDBName(spans);
+        assert.strictEqual(
+          spans[0].name,
+          'CloudSpanner.Database.batchWriteAtLeastOnce',
+        );
+        assert.strictEqual(spans[0].events.length, 0);
+        done();
+      });
+    });
+
+    it('records error on span when stream is destroyed with an error', done => {
+      const abortErr = new Error('abort');
+      const stream = database.batchWriteAtLeastOnce(mutationGroups, options);
+      stream.on('error', err => {
+        assert.strictEqual(err, abortErr);
+        const spans = traceExporter.getFinishedSpans();
+        assert.strictEqual(spans.length, 1, 'Exactly 1 span expected');
+        withAllSpansHaveDBName(spans);
+        assert.strictEqual(spans[0].status.code, SpanStatusCode.ERROR);
+        assert.strictEqual(spans[0].status.message, 'abort');
+        done();
+      });
+      stream.destroy(abortErr);
+    });
+
+    it('ends span immediately with no events when mutationGroups is empty', done => {
+      database
+        .batchWriteAtLeastOnce([], options)
+        .on('data', () => {
+          done(new Error('Should not emit data'));
+        })
+        .on('error', done)
+        .on('end', () => {
+          const spans = traceExporter.getFinishedSpans();
+          assert.strictEqual(spans.length, 1, 'Exactly 1 span expected');
+          withAllSpansHaveDBName(spans);
+          assert.strictEqual(
+            spans[0].name,
+            'CloudSpanner.Database.batchWriteAtLeastOnce',
+          );
+          assert.strictEqual(spans[0].events.length, 0);
+          assert.notStrictEqual(spans[0].status.code, SpanStatusCode.ERROR);
+          done();
+        });
+    });
   });
 
   describe('runTransaction', () => {