feat: Add native AbortSignal support

dev/conformance/runner.ts

@@ -272,7 +272,10 @@
         updateTime: {},
       });
     }
-    return response(res);
+    const promise = response(res) as any;
+    // Add cancel method to make it a CancellablePromise
+    promise.cancel = () => {};
+    return promise;
   };
 }
 
@@ -284,7 +287,9 @@
     );
     const expectedQuery = STRUCTURED_QUERY_TYPE.fromObject(spec.query);
     expect(actualQuery).to.deep.equal(expectedQuery);
-    const stream = through2.obj();
+    const stream = through2.obj() as any;
+    // Add cancel method to make it a CancellableStream
+    stream.cancel = () => {};
     setImmediate(() => {
       // Empty query always emits a readTime
       stream.push({readTime: {seconds: 0, nanos: 0}});
@@ -299,7 +304,9 @@
   return (request: api.IBatchGetDocumentsRequest) => {
     const getDocument = spec.request;
     expect(request.documents![0]).to.equal(getDocument.name);
-    const stream = through2.obj();
+    const stream = through2.obj() as any;
+    // Add cancel method to make it a CancellableStream
+    stream.cancel = () => {};
     setImmediate(() => {
       stream.push({
         missing: getDocument.name,
@@ -432,7 +439,12 @@
     const expectedSnapshots = spec.snapshots;
     const writeStream = through2.obj();
     const overrides: ApiOverride = {
-      listen: () => duplexify.obj(through2.obj(), writeStream),
+      listen: () => {
+        const stream = duplexify.obj(through2.obj(), writeStream) as any;
+        // Add cancel method to make it a CancellableStream
+        stream.cancel = () => {};
+        return stream;
+      },
     };
 
     return createInstance(overrides).then(() => {

dev/src/abort-util.ts

@@ -0,0 +1,79 @@
+/*!
+ * Copyright 2024 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/**
+ * Interface for objects that can be cancelled.
+ */
+export interface Cancellable {
+  cancel(): void;
+}
+
+/**
+ * Utility class for working with AbortSignal and cancellable operations.
+ */
+export class AbortUtil {
+  /**
+   * Throws an error if the AbortSignal is already aborted.
+   */
+  static throwIfAborted(signal: AbortSignal | null): void {
+    if (signal?.aborted) {
+      throw new Error('The operation was aborted');
+    }
+  }
+
+  /**
+   * Creates a Promise that rejects when the AbortSignal is aborted.
+   */
+  static createAbortPromise<T>(
+    signal: AbortSignal,
+    cancellable: Cancellable
+  ): Promise<T> {
+    return new Promise<T>((_, reject) => {
+      const onAbort = () => {
+        cancellable.cancel();
+        reject(new Error('The operation was aborted'));
+      };
+
+      if (signal.aborted) {
+        onAbort();
+      } else {
+        signal.addEventListener('abort', onAbort, { once: true });
+      }
+    });
+  }
+
+  /**
+   * Makes a Promise cancellable with an AbortSignal.
+   */
+  static async makeCancellable<T>(
+    promise: Promise<T>,
+    cancellable: Cancellable,
+    signal: AbortSignal | null
+  ): Promise<T> {
+    if (!signal) {
+      return promise;
+    }
+
+    // Check if already aborted
+    AbortUtil.throwIfAborted(signal);
+
+    // Race the original promise against the abort promise
+    return Promise.race([
+      promise,
+      AbortUtil.createAbortPromise<T>(signal, cancellable)
+    ]);
+  }
+}

dev/src/bulk-writer.ts

@@ -246,7 +246,7 @@
     return this.docPaths.has(documentRef.path);
   }
 
-  async bulkCommit(options: {requestTag?: string} = {}): Promise<void> {
+  async bulkCommit(options: {requestTag?: string, abortSignal?: AbortSignal} = {}): Promise<void> {
     return this._firestore._traceUtil.startActiveSpan(
       SPAN_NAME_BULK_WRITER_COMMIT,
       async () => {
@@ -266,7 +266,7 @@
           response = await this._commit<
             api.BatchWriteRequest,
             api.BatchWriteResponse
-          >({retryCodes, methodName: 'batchWrite', requestTag: tag});
+          >(options, {retryCodes, methodName: 'batchWrite', requestTag: tag}); //TODO: add public options
         } catch (err) {
           // Map the failure to each individual write's result.
           const ops = Array.from({length: this.pendingOps.length});

dev/src/collection-group.ts

@@ -80,7 +80,8 @@ export class CollectionGroup<
    * `QueryPartition`s.
    */
   async *getPartitions(
-    desiredPartitionCount: number
+    desiredPartitionCount: number,
+    options?: firestore.FirestoreRequestOptions
   ): AsyncIterable<QueryPartition<AppModelType, DbModelType>> {
     const partitions: Array<api.IValue>[] = [];
 
@@ -108,7 +109,8 @@ export class CollectionGroup<
             'partitionQueryStream',
             /* bidirectional= */ false,
             request,
-            tag
+            tag,
+            options
           );
           stream.resume();
 

dev/src/document-reader.ts

@@ -22,7 +22,7 @@
 import {logger} from './logger';
 import {Firestore} from './index';
 import {Timestamp} from './timestamp';
-import {DocumentData} from '@google-cloud/firestore';
+import {DocumentData, FirestoreRequestOptions} from '@google-cloud/firestore';
 import api = google.firestore.v1;
 
 interface BatchGetResponse<AppModelType, DbModelType extends DocumentData> {
@@ -79,9 +79,10 @@
    * @param requestTag A unique client-assigned identifier for this request.
    */
   async get(
-    requestTag: string
+    requestTag: string,
+    options?: FirestoreRequestOptions
   ): Promise<Array<DocumentSnapshot<AppModelType, DbModelType>>> {
-    const {result} = await this._get(requestTag);
+    const {result} = await this._get(requestTag, options); //TODO: add public options
     return result;
   }
 
@@ -92,9 +93,10 @@
    * @param requestTag A unique client-assigned identifier for this request.
    */
   async _get(
-    requestTag: string
+    requestTag: string,
+    options: FirestoreRequestOptions | undefined
   ): Promise<BatchGetResponse<AppModelType, DbModelType>> {
-    await this.fetchDocuments(requestTag);
+    await this.fetchDocuments(requestTag, options);
 
     // BatchGetDocuments doesn't preserve document order. We use the request
     // order to sort the resulting documents.
@@ -125,7 +127,7 @@
     };
   }
 
-  private async fetchDocuments(requestTag: string): Promise<void> {
+  private async fetchDocuments(requestTag: string, options: FirestoreRequestOptions | undefined): Promise<void> {
     if (!this.outstandingDocuments.size) {
       return;
     }
@@ -156,7 +158,8 @@
         'batchGetDocuments',
         /* bidirectional= */ false,
         request,
-        requestTag
+        requestTag,
+        options
       );
       stream.resume();
 
@@ -217,7 +220,7 @@
         shouldRetry
       );
       if (shouldRetry) {
-        return this.fetchDocuments(requestTag);
+        return this.fetchDocuments(requestTag, options);
       } else {
         throw error;
       }

dev/src/index.ts

@@ -74,6 +74,7 @@
   validateTimestamp,
 } from './validate';
 import {WriteBatch} from './write-batch';
+import {AbortUtil, Cancellable} from './abort-util';
 
 import {interfaces} from './v1/firestore_client_config.json';
 const serviceConfig = interfaces['google.firestore.v1.Firestore'];
@@ -1288,9 +1289,9 @@
    * });
    * ```
    */
-  listCollections(): Promise<CollectionReference[]> {
+  listCollections(options?: firestore.FirestoreRequestOptions): Promise<CollectionReference[]> {
     const rootDocument = new DocumentReference(this, ResourcePath.EMPTY);
-    return rootDocument.listCollections();
+    return rootDocument.listCollections(options);
   }
 
   /**
@@ -1331,9 +1332,10 @@
           1
         );
 
-        const {documents, fieldMask} = parseGetAllArguments(
+        const { parsedGetAllArguments, readOptions } = parseGetAllArguments(
           documentRefsOrReadOptions
         );
+        const { documents, fieldMask } = parsedGetAllArguments;
 
         this._traceUtil.currentSpan().setAttributes({
           [ATTRIBUTE_KEY_IS_TRANSACTIONAL]: false,
@@ -1348,7 +1350,7 @@
         return this.initializeIfNeeded(tag)
           .then(() => {
             const reader = new DocumentReader(this, documents, fieldMask);
-            return reader.get(tag);
+            return reader.get(tag, readOptions);
           })
           .catch(err => {
             throw wrapError(err, stack);
@@ -1793,11 +1795,17 @@
     methodName: FirestoreUnaryMethod,
     request: Req,
     requestTag: string,
+    options: firestore.FirestoreRequestOptions | undefined,
     retryCodes?: number[]
   ): Promise<Resp> {
+    const abortSignal = options?.abortSignal || null;
+
+    // Check if already aborted before starting
+    AbortUtil.throwIfAborted(abortSignal);
+
     const callOptions = this.createCallOptions(methodName, retryCodes);
 
-    return this._clientPool.run(
+    const requestPromise = this._clientPool.run(
       requestTag,
       /* requiresGrpc= */ false,
       async gapicClient => {
@@ -1808,9 +1816,21 @@
             'Sending request: %j',
             request
           );
-          const [result] = await (
-            gapicClient[methodName] as UnaryMethod<Req, Resp>
-          )(request, callOptions);
+
+          // Make the GAX call - this returns a CancellablePromise
+          const gaxCall = (gapicClient[methodName] as UnaryMethod<Req, Resp>)(request, callOptions);
+
+          // Create cancellable wrapper for the GAX call
+          const cancellable: Cancellable = {
+            cancel: () => {
+              logger('Firestore.request', requestTag, 'Cancelling request due to abort signal');
+              gaxCall.cancel();
+            }
+          };
+
+          // Make the call cancellable with AbortSignal
+          const [result] = await AbortUtil.makeCancellable(gaxCall, cancellable, abortSignal);
+
           logger(
             'Firestore.request',
             requestTag,
@@ -1824,6 +1844,8 @@
         }
       }
     );
+
+    return requestPromise;
   }
 
   /**
@@ -1847,8 +1869,14 @@
     methodName: FirestoreStreamingMethod,
     bidrectional: boolean,
     request: {},
-    requestTag: string
+    requestTag: string,
+    options: firestore.FirestoreRequestOptions | undefined
   ): Promise<Duplex> {
+    const abortSignal = options?.abortSignal || null;
+
+    // Check if already aborted before starting
+    AbortUtil.throwIfAborted(abortSignal);
+
     const callOptions = this.createCallOptions(methodName);
 
     const bidirectional = methodName === 'listen';
@@ -1874,6 +1902,15 @@
           const stream = bidirectional
             ? gapicClient[methodName](callOptions)
             : gapicClient[methodName](request, callOptions);
+
+          // Create a cancellable object for the stream
+          const cancellable: Cancellable = {
+            cancel: () => {
+              logger('Firestore.requestStream', requestTag, 'Cancelling stream due to abort signal');
+              (stream as any).cancel();
+            }
+          };
+
           const logStream = new Transform({
             objectMode: true,
             transform: (chunk, encoding, callback) => {
@@ -1901,12 +1938,20 @@
           stream.pipe(logStream);
 
           const lifetime = new Deferred<void>();
-          const resultStream = await this._initializeStream(
+          const streamPromise = this._initializeStream(
             stream,
             lifetime,
             requestTag,
             bidirectional ? request : undefined
           );
+
+          // Make the stream initialization cancellable
+          const resultStream = await AbortUtil.makeCancellable(
+            streamPromise,
+            cancellable,
+            abortSignal
+          );
+
           resultStream.on('end', () => {
             stream.end();
             this._traceUtil