Skip to content

Instantly share code, notes, and snippets.

@samthecodingman
Last active September 19, 2022 16:11
Show Gist options
  • Save samthecodingman/faba163b31488ab885bacefa7f63d121 to your computer and use it in GitHub Desktop.
Save samthecodingman/faba163b31488ab885bacefa7f63d121 to your computer and use it in GitHub Desktop.
A helper object used for managing large batch write operations on Cloud Firestore.
/*! 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));
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment