Add an event for logical termination and add events to tests.

googleapis · Oct 29, 2024 · 89e63b4 · 89e63b4
1 parent 4b66844
commit 89e63b4
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 2 deletions.
diff --git a/dev/src/reference/query-util.ts b/dev/src/reference/query-util.ts
@@ -180,6 +180,7 @@ export class QueryUtil<
     const tag = requestTag();
     const startTime = Date.now();
     const isExplain = explainOptions !== undefined;
+    const methodName = 'runQuery';
 
     let numDocumentsReceived = 0;
     let lastReceivedDocument: QueryDocumentSnapshot<
@@ -245,6 +246,11 @@ export class QueryUtil<
 
         if (proto.done) {
           logger('QueryUtil._stream', tag, 'Trigger Logical Termination.');
+          this._firestore._traceUtil
+            .currentSpan()
+            .addEvent(
+              `Firestore.${methodName}: Received RunQueryResponse.Done.`
+            );
           backendStream.unpipe(stream);
           backendStream.resume();
           backendStream.end();
@@ -265,7 +271,6 @@ export class QueryUtil<
         let streamActive: Deferred<boolean>;
         do {
           streamActive = new Deferred<boolean>();
-          const methodName = 'runQuery';
 
           this._firestore._traceUtil
             .currentSpan()

diff --git a/dev/system-test/tracing.ts b/dev/system-test/tracing.ts
@@ -652,7 +652,7 @@ describe('Tracing Tests', () => {
 
     // Expect that the span exists first.
     const span = getSpanByName(spanName);
-    expect(span).to.not.be.null;
+    expect(span, `Could not find the span named ${spanName}`).to.not.be.null;
 
     // Assert that the expected attributes are present in the span attributes.
     // Note that the span attributes may be a superset of the attributes passed
@@ -664,6 +664,29 @@ describe('Tracing Tests', () => {
     }
   }
 
+  // Ensures that the given span exists and has the given attributes.
+  function expectSpanHasEvents(spanName: string, eventNames: string[]): void {
+    // The Cloud Trace API does not return span attributes and events.
+    if (testConfig.e2e) {
+      return;
+    }
+
+    // Expect that the span exists first.
+    const span = getSpanByName(spanName);
+    expect(span, `Could not find the span named ${spanName}`).to.not.be.null;
+
+    // Assert that the expected attributes are present in the span attributes.
+    // Note that the span attributes may be a superset of the attributes passed
+    // to this function.
+    if (span?.events) {
+      const numEvents = eventNames.length;
+      expect(numEvents).to.equal(span.events.length);
+      for (let i = 0; i < numEvents; ++i) {
+        expect(span.events[i].name).to.equal(eventNames[i]);
+      }
+    }
+  }
+
   describe(IN_MEMORY_TEST_SUITE_TITLE, () => {
     describe(NON_GLOBAL_OTEL_TEST_SUITE_TITLE, () => {
       describe(GRPC_TEST_SUITE_TITLE, () => {
@@ -745,6 +768,11 @@ describe('Tracing Tests', () => {
         SPAN_NAME_DOC_REF_GET,
         SPAN_NAME_BATCH_GET_DOCUMENTS
       );
+      expectSpanHasEvents(SPAN_NAME_BATCH_GET_DOCUMENTS, [
+        'Firestore.batchGetDocuments: Start',
+        'Firestore.batchGetDocuments: First response received',
+        'Firestore.batchGetDocuments: Completed',
+      ]);
     });
 
     it('document reference create()', async () => {
@@ -812,6 +840,11 @@ describe('Tracing Tests', () => {
       );
       await waitForCompletedSpans(2);
       expectSpanHierarchy(SPAN_NAME_TEST_ROOT, SPAN_NAME_AGGREGATION_QUERY_GET);
+      expectSpanHasEvents(SPAN_NAME_AGGREGATION_QUERY_GET, [
+        'Firestore.runAggregationQuery: Start',
+        'Firestore.runAggregationQuery: First response received',
+        'Firestore.runAggregationQuery: Completed',
+      ]);
     });
 
     it('collection reference add()', async () => {
@@ -844,6 +877,13 @@ describe('Tracing Tests', () => {
       );
       await waitForCompletedSpans(2);
       expectSpanHierarchy(SPAN_NAME_TEST_ROOT, SPAN_NAME_QUERY_GET);
+      expectSpanHasEvents(SPAN_NAME_QUERY_GET, [
+        'RunQuery',
+        'Firestore.runQuery: Start',
+        'Firestore.runQuery: First response received',
+        'Firestore.runQuery: Received RunQueryResponse.Done.',
+        'Firestore.runQuery: Completed',
+      ]);
     });
 
     it('firestore getAll()', async () => {
@@ -854,6 +894,11 @@ describe('Tracing Tests', () => {
       );
       await waitForCompletedSpans(2);
       expectSpanHierarchy(SPAN_NAME_TEST_ROOT, SPAN_NAME_BATCH_GET_DOCUMENTS);
+      expectSpanHasEvents(SPAN_NAME_BATCH_GET_DOCUMENTS, [
+        'Firestore.batchGetDocuments: Start',
+        'Firestore.batchGetDocuments: First response received',
+        'Firestore.batchGetDocuments: Completed',
+      ]);
     });
 
     it('transaction', async () => {