A helper object used for managing large batch write operations on Cloud Firestore.

firestore-multi-batch.ts

/*! firestore-multi-batch.ts | Samuel Jones 2021 | MIT License | gist.github.com/samthecodingman */

import { firestore } from "firebase-admin";

/**
 * Helper class to compile an expanding `firestore.WriteBatch`.
 *
 * Using an internal operations counter, this class will automatically start a
 * new `firestore.WriteBatch` instance when it detects it has hit the operations
 * limit of 500. Once prepared, you can commit the batches together.
 *
 * Note: `FieldValue` transform operations such as `serverTimestamp`, 
 * `arrayUnion`, `arrayRemove`, `increment` are counted as two operations. If
 * your written data makes use of one of these, you should use the appropriate
 * `transformCreate`, `transformSet` or `transformUpdate` method so that the 
 * internal counter is correctly increased by 2 (the normal versions only 
 * increase the counter by 1).
 * 
 * If not sure, just use `delete`, `transformCreate`, `transformSet`, or
 * `transformUpdate` functions for every operation as this will make sure you
 * don't exceed the limit.
 *
 * @author Samuel Jones <@samthecodingman> [MIT License]
 * @see https://stackoverflow.com/a/66692467/3068190 
 * @see https://firebase.google.com/docs/firestore/manage-data/transactions
 * @see https://firebase.google.com/docs/reference/js/firebase.firestore.FieldValue
 * @see https://firebase.google.com/docs/reference/js/firebase.firestore.WriteBatch
 */
export class MultiBatch {
  dbRef: firestore.Firestore;
  batches: firestore.WriteBatch[];
  currentBatch: firestore.WriteBatch;
  currentBatchOpCount: number;
  committed: boolean;

  /** Initializes a new managed group of batch operations */
  constructor(dbRef: firestore.Firestore) {
    this.dbRef = dbRef;
    this.batches = [this.dbRef.batch()];
    this.currentBatch = this.batches[0];
    this.currentBatchOpCount = 0;
    this.committed = false;
  }

  /** INTERNAL: Increments operation counter and manages WriteBatch instances */
  _getCurrentBatch(count: number) {
    if (this.committed) throw new Error("MultiBatch already committed.");
    if (this.currentBatchOpCount + count > 500) {
      // operation limit exceeded, start a new batch
      this.currentBatch = this.dbRef.batch();
      this.currentBatchOpCount = 0;
      this.batches.push(this.currentBatch);
    }
    this.currentBatchOpCount += count;
    return this.currentBatch;
  }

  /** Creates the document, fails if it exists */
  create(ref: firestore.DocumentReference, data: firestore.DocumentData) {
    this._getCurrentBatch(1).create(ref, data);
    return this;
  }

  /**
   * Creates the document, fails if it exists. Used for commands that contain
   * serverTimestamp, arrayUnion, etc
   */
  transformCreate(
    ref: firestore.DocumentReference,
    data: firestore.DocumentData
  ) {
    this._getCurrentBatch(2).create(ref, data);
    return this;
  }

  /** Writes the document, creating/overwriting/etc as applicable. */
  set(
    ref: firestore.DocumentReference,
    data: firestore.DocumentData,
    options?: FirebaseFirestore.SetOptions
  ) {
    this._getCurrentBatch(1).set(ref, data, options);
    return this;
  }

  /**
   * Writes the document, creating/overwriting/etc as applicable. Used for
   * commands that contain serverTimestamp, arrayUnion, etc
   */
  transformSet(
    ref: firestore.DocumentReference,
    data: firestore.DocumentData,
    options?: FirebaseFirestore.SetOptions
  ) {
    this._getCurrentBatch(2).set(ref, data, options);
    return this;
  }

  /** Merges data into the document, failing if the document doesn't exist. */
  update(
    ref: firestore.DocumentReference,
    data: firestore.DocumentData,
    ...fieldsOrPrecondition: any[]
  ) {
    this._getCurrentBatch(1).update(ref, data, ...fieldsOrPrecondition);
    return this;
  }

  /**
   * Merges data into the document, failing if the document doesn't exist. Used
   * for commands that contain serverTimestamp, arrayUnion, etc
   */
  transformUpdate(
    ref: firestore.DocumentReference,
    data: firestore.DocumentData,
    ...fieldsOrPrecondition: any[]
  ) {
    this._getCurrentBatch(2).update(ref, data, ...fieldsOrPrecondition);
    return this;
  }

  /** Used when for basic update operations */
  delete(ref: firestore.DocumentReference) {
    this._getCurrentBatch(1).delete(ref);
    return this;
  }

  /**
   * Commits all of the batches to Firestore.
   *
   * Note: Unlike normal batch operations, this may cause one or more atomic
   * writes. One batch may succeed where others fail. By default, if any batch
   * fails, it will fail the whole promise. This can be suppressed by passing in
   * a truthy value as the first argument and checking the results returned by
   * this method.
   *
   * @param {boolean} [suppressErrors=false] Whether to suppress errors on a
   * per-batch basis.
   * @return {firestore.WriteResult[][]} array containing an array of
   * `WriteResult` objects for each batch.
   */
  commit(suppressErrors?: false): Promise<firestore.WriteResult[][]>;
  /**
   *
   * Commits all of the batches to Firestore.
   *
   * Note: Unlike normal batch operations, this may cause one or more atomic
   * writes. One batch may succeed where others fail. By default, if any batch
   * fails, it will fail the whole promise. This can be suppressed by passing in
   * a truthy value as the first argument and checking the results returned by
   * this method.
   *
   * @param {boolean} [suppressErrors=false] Whether to suppress errors on a
   * per-batch basis.
   * @return {firestore.WriteResult[]} array containing an array of
   * `WriteResult` objects and error-batch pairs, for each batch.
   */
  commit(
    suppressErrors: true
  ): Promise<
    (firestore.WriteResult[] | { error: any; batch: firestore.WriteBatch })[]
  >;
  commit(suppressErrors = false) {
    this.committed = true;
    const mapCallback = suppressErrors
      ? (batch: firestore.WriteBatch) =>
          batch.commit().catch((error) => ({ error, batch }))
      : (batch: firestore.WriteBatch) => batch.commit();

    return Promise.all(this.batches.map(mapCallback));
  }
}