feat: add recursiveDelete()

dev/src/index.ts

@@ -16,14 +16,14 @@
 
 import * as firestore from '@google-cloud/firestore';
 
-import {CallOptions, grpc, RetryOptions} from 'google-gax';
+import {CallOptions, GoogleError, grpc, RetryOptions, Status} from 'google-gax';
 import {Duplex, PassThrough, Transform} from 'stream';
 
 import {URL} from 'url';
 
 import {google} from '../protos/firestore_v1_proto_api';
 import {ExponentialBackoff, ExponentialBackoffSetting} from './backoff';
-import {BulkWriter} from './bulk-writer';
+import {BulkWriter, BulkWriterError} from './bulk-writer';
 import {BundleBuilder} from './bundle';
 import {fieldsFromJson, timestampFromJson} from './convert';
 import {
@@ -76,6 +76,7 @@ const serviceConfig = interfaces['google.firestore.v1.Firestore'];
 
 import api = google.firestore.v1;
 import {CollectionGroup} from './collection-group';
+import {DocumentData} from '@google-cloud/firestore';
 
 export {
   CollectionReference,
@@ -141,7 +142,7 @@ const CLOUD_RESOURCE_HEADER = 'google-cloud-resource-prefix';
 /*!
  * The maximum number of times to retry idempotent requests.
  */
-const MAX_REQUEST_RETRIES = 5;
+export const MAX_REQUEST_RETRIES = 5;
 
 /*!
  * The default number of idle GRPC channel to keep.
@@ -166,7 +167,7 @@ const MAX_CONCURRENT_REQUESTS_PER_CLIENT = 100;
  *
  * @private
  */
-const REFERENCE_NAME_MIN_ID = '__id-9223372036854775808__';
+export const REFERENCE_NAME_MIN_ID = '__id-9223372036854775808__';
 
 /**
  * Document data (e.g. for use with
@@ -399,6 +400,26 @@ export class Firestore implements firestore.Firestore {
    */
   private registeredListenersCount = 0;
 
+  /**
+   * A lazy-loaded BulkWriter instance to be used with recursiveDelete() if no
+   * BulkWriter instance is provided.
+   *
+   * @private
+   */
+  private _bulkWriter: BulkWriter | undefined;
+
+  /**
+   * Lazy-load the Firestore's default BulkWriter.
+   *
+   * @private
+   */
+  private getBulkWriter(): BulkWriter {
+    if (!this._bulkWriter) {
+      this._bulkWriter = this.bulkWriter();
+    }
+    return this._bulkWriter;
+  }
+
   /**
    * Number of pending operations on the client.
    *
@@ -1198,15 +1219,100 @@ export class Firestore implements firestore.Firestore {
     this.bulkWritersCount -= 1;
   }
 
+  /**
+   * Recursively deletes all documents and subcollections at and under the
+   * specified level.
+   *
+   * If any deletes fail, the promise is rejected with an error message
+   * containing the number of failed writes and the stack trace of the last
+   * failed write. The provided reference is deleted regardless of whether
+   * all deletes succeeded, except when Firestore fails to fetch the provided
+   * reference's descendants.
+   *
+   * Firestore uses a BulkWriter instance with default settings to perform the
+   * deletes. To customize throttling rates or add success/error callbacks,
+   * pass in a custom BulkWriter instance.
+   *
+   * @param ref The reference of a document or collection to delete.
+   * @param bulkWriter Custom BulkWriter instance with which to perform the
+   * deletes with.
+   * @return A promise that resolves when all deletes have been performed.
+   * The promise is rejected if any of the deletes fail.
+   */
+  recursiveDelete<T = DocumentData>(
+    ref: CollectionReference<T> | DocumentReference<T>,
+    bulkWriter?: BulkWriter
+  ): Promise<void> {
+    const docStream = this.getAllDescendants(ref);
+    const writer = bulkWriter ?? this.getBulkWriter();
+    const deleteCompleted = new Deferred<void>();
+    let errorCount = 0;
+    let lastError: Error | undefined;
+
+    // Capture the error stack to preserve stack tracing across async calls.
+    const stack = Error().stack!;
+
+    docStream
+      .on('error', err => {
+        err.code = Status.FAILED_PRECONDITION;
+        err.message =
+          'Failed to fetch children documents. ' +
+          'The provided reference was not deleted.';
+        deleteCompleted.reject(wrapError(err, stack));
+      })
+      .on('data', (snap: QueryDocumentSnapshot) => {
+        const docRef = new DocumentReference(this, snap.ref._path);
+        writer.delete(docRef).catch(err => {
+          errorCount++;
+          lastError = err;
+        });
+      })
+      .on('end', () => {
+        writer.flush().then(async () => {
+          if (ref instanceof DocumentReference) {
+            try {
+              await ref.delete();
+            } catch (err) {
+              logger(
+                'Firestore.recursiveDelete',
+                null,
+                'failed to delete the provided reference',
+                ref.path
+              );
+              lastError = err;
+            }
+          }
+
+          if (lastError === undefined) {
+            deleteCompleted.resolve();
+          } else {
+            let error = new GoogleError(
+              `${errorCount} ` +
+                `${errorCount > 1 ? 'deletes' : 'delete'} ` +
+                'failed. The last delete failed with: '
+            );
+            if (lastError instanceof BulkWriterError) {
+              error.code = (lastError.code as number) as Status;
+            }
+            error = wrapError(error, stack);
+
+            // Wrap the BulkWriter error last to provide the full stack trace.
+            deleteCompleted.reject(wrapError(error, lastError.stack ?? ''));
+          }
+        });
+      });
+
+    return deleteCompleted.promise;
+  }
+
   /**
    * Retrieves all descendant documents nested under the provided reference.
    *
    * @private
    * @return {Stream<QueryDocumentSnapshot>} Stream of descendant documents.
    */
-  // TODO(chenbrian): Make this a private method after adding recursive delete.
-  _getAllDescendants(
-    ref: CollectionReference | DocumentReference
+  private getAllDescendants<T = DocumentData>(
+    ref: CollectionReference<T> | DocumentReference<T>
   ): NodeJS.ReadableStream {
     // The parent is the closest ancestor document to the location we're
     // deleting. If we are deleting a document, the parent is the path of that

dev/system-test/firestore.ts

@@ -2629,69 +2629,95 @@ describe('BulkWriter class', () => {
     return firestore.terminate();
   });
 
-  // TODO(chenbrian): This is a temporary test used to validate that the
-  //  StructuredQuery calls work properly. Remove these tests after adding
-  //  recursive delete tests.
-  it('finds nested documents and collection', async () => {
-    // ROOT-DB
-    // └── randomCol
-    //     ├── anna
-    //     └── bob
-    //         └── parentsCol
-    //             ├── charlie
-    //             └── daniel
-    //                 └── childCol
-    //                     ├── ernie
-    //                     └── francis
-    const batch = firestore.batch();
-    batch.set(randomCol.doc('anna'), {name: 'anna'});
-    batch.set(randomCol.doc('bob'), {name: 'bob'});
-    batch.set(randomCol.doc('bob/parentsCol/charlie'), {name: 'charlie'});
-    batch.set(randomCol.doc('bob/parentsCol/daniel'), {name: 'daniel'});
-    batch.set(randomCol.doc('bob/parentsCol/daniel/childCol/ernie'), {
-      name: 'ernie',
-    });
-    batch.set(randomCol.doc('bob/parentsCol/daniel/childCol/francis'), {
-      name: 'francis',
-    });
-    await batch.commit();
+  describe('recursiveDelete()', () => {
+    async function countDocumentChildren(
+      ref: DocumentReference
+    ): Promise<number> {
+      let count = 0;
+      const collections = await ref.listCollections();
+      for (const collection of collections) {
+        count += await countCollectionChildren(collection);
+      }
+      return count;
+    }
 
-    const numStreamItems = async (
-      stream: NodeJS.ReadableStream
-    ): Promise<number> => {
+    async function countCollectionChildren(
+      ref: CollectionReference
+    ): Promise<number> {
       let count = 0;
-      // eslint-disable-next-line @typescript-eslint/no-unused-vars
-      for await (const _ of stream) {
-        ++count;
+      const docs = await ref.listDocuments();
+      for (const doc of docs) {
+        count += (await countDocumentChildren(doc)) + 1;
       }
       return count;
-    };
+    }
 
-    // Query all descendants of collections.
-    let descendantsStream = await firestore._getAllDescendants(randomCol);
-    expect(await numStreamItems(descendantsStream)).to.equal(6);
-    descendantsStream = await firestore._getAllDescendants(
-      randomCol.doc('bob').collection('parentsCol')
-    );
-    expect(await numStreamItems(descendantsStream)).to.equal(4);
-    descendantsStream = await firestore._getAllDescendants(
-      randomCol.doc('bob').collection('parentsCol/daniel/childCol')
-    );
-    expect(await numStreamItems(descendantsStream)).to.equal(2);
+    beforeEach(async () => {
+      // ROOT-DB
+      // └── randomCol
+      //     ├── anna
+      //     └── bob
+      //         └── parentsCol
+      //             ├── charlie
+      //             └── daniel
+      //                 └── childCol
+      //                     ├── ernie
+      //                     └── francis
+      const batch = firestore.batch();
+      batch.set(randomCol.doc('anna'), {name: 'anna'});
+      batch.set(randomCol.doc('bob'), {name: 'bob'});
+      batch.set(randomCol.doc('bob/parentsCol/charlie'), {name: 'charlie'});
+      batch.set(randomCol.doc('bob/parentsCol/daniel'), {name: 'daniel'});
+      batch.set(randomCol.doc('bob/parentsCol/daniel/childCol/ernie'), {
+        name: 'ernie',
+      });
+      batch.set(randomCol.doc('bob/parentsCol/daniel/childCol/francis'), {
+        name: 'francis',
+      });
+      await batch.commit();
+    });
 
-    // Query all descendants of documents.
-    descendantsStream = await firestore._getAllDescendants(
-      randomCol.doc('bob')
-    );
-    expect(await numStreamItems(descendantsStream)).to.equal(4);
-    descendantsStream = await firestore._getAllDescendants(
-      randomCol.doc('bob/parentsCol/daniel')
-    );
-    expect(await numStreamItems(descendantsStream)).to.equal(2);
-    descendantsStream = await firestore._getAllDescendants(
-      randomCol.doc('anna')
-    );
-    expect(await numStreamItems(descendantsStream)).to.equal(0);
+    it('recursiveDelete on top-level collection', async () => {
+      await firestore.recursiveDelete(randomCol);
+      expect(await countCollectionChildren(randomCol)).to.equal(0);
+    });
+
+    it('recursiveDelete on nested collection', async () => {
+      const coll = randomCol.doc('bob').collection('parentsCol');
+      await firestore.recursiveDelete(coll);
+
+      expect(await countCollectionChildren(coll)).to.equal(0);
+      expect(await countCollectionChildren(randomCol)).to.equal(2);
+    });
+
+    it('recursiveDelete on nested document', async () => {
+      const doc = randomCol.doc('bob/parentsCol/daniel');
+      await firestore.recursiveDelete(doc);
+
+      const docSnap = await doc.get();
+      expect(docSnap.exists).to.be.false;
+      expect(await countDocumentChildren(randomCol.doc('bob'))).to.equal(1);
+      expect(await countCollectionChildren(randomCol)).to.equal(3);
+    });
+
+    it('recursiveDelete on leaf document', async () => {
+      const doc = randomCol.doc('bob/parentsCol/daniel/childCol/ernie');
+      await firestore.recursiveDelete(doc);
+
+      const docSnap = await doc.get();
+      expect(docSnap.exists).to.be.false;
+      expect(await countCollectionChildren(randomCol)).to.equal(5);
+    });
+
+    it('recursiveDelete with custom BulkWriter instance', async () => {
+      const bulkWriter = firestore.bulkWriter();
+      let callbackCount = 0;
+      bulkWriter.onWriteResult(() => {
+        callbackCount++;
+      });
+      await firestore.recursiveDelete(randomCol, bulkWriter);
+      expect(callbackCount).to.equal(6);
+    });
   });
 
   it('can retry failed writes with a provided callback', async () => {