Skip to content

Instantly share code, notes, and snippets.

@katowulf

katowulf/DbBatch.ts

Last active Sep 13, 2019
Embed
What would you like to do?
Batch utility for seeding data and deleting collections larger than 500 docs.
export default {
databaseURL: 'https://YOUR_PROJECT_ID.firebaseio.com',
basePath: [], // root
//basePath: 'organizations/companyName', // use a doc as an ad hoc namespace
}
/**********************************
* This is a convenience utility to
* abstract a bit of the Firestore
* boilerplate and complexity dealing
* with docs vs collections.
*********************************/
import conf from './conf';
const TRANSACTION_LIMIT = 500;
import * as admin from "firebase-admin";
import {DocumentReference, FirebaseFirestore} from "@firebase/firestore-types";
if( admin.apps.length === 0 ) {
admin.initializeApp();
}
// this is not a privileged ref
const db = admin.firestore();
function buildPath(url: string|Array<string>) {
let ref: any = db;
splitUrl(conf.basePath, url).forEach(p => {
if( typeof ref.collection === 'function' ) {
ref = ref.collection(p);
}
else {
ref = ref.doc(p);
}
});
return ref;
}
function getData(url: string|Array<string>) {
return buildPath(url).get().then(snap => {
if( snap instanceof admin.firestore.DocumentSnapshot ) {
if( snap.exists ) {
return snap.data();
}
else {
return null;
}
}
else {
const data = {};
snap.forEach(doc => data[doc.id] = doc.data());
return data;
}
});
}
function splitUrl(...url: Array<string|Array<string>>): Array<string> {
let parts = [];
url.forEach(u => {
if( typeof u === 'string' ) {
u.replace(/^\//, '').replace(/\/$/, '').split('/').forEach(uu =>{
if( uu !== '' ) { parts.push(uu); }
});
}
else {
u.forEach(ubit => parts = [...parts, ...splitUrl(ubit)]);
}
});
return parts;
}
export default {
path: buildPath,
get: getData,
relativePath(ref: DocumentReference) {
const re = new RegExp(`^${conf.basePath}/?`);
return ref.path.replace(re, '');
},
newId: function() {
return buildPath('foo/bar').id;
},
addToArray: function(url: string|Array<string>, field: string, val) {
const data = {};
data[field] = admin.firestore.FieldValue.arrayUnion(val);
return buildPath(url).update(data);
},
removeFromArray: function(url: string|Array<string>, field: string, val) {
const data = {};
data[field] = admin.firestore.FieldValue.arrayRemove(val);
return buildPath(url).update(data);
},
batch: function() { return db.batch(); },
root: function(): FirebaseFirestore.DocumentReference { return db.doc(conf.basePath); },
transaction: db.runTransaction.bind(db),
FieldValue: admin.firestore.FieldValue,
TRANSACTION_LIMIT: TRANSACTION_LIMIT
};
import db from './db';
// Set this to the max ops per batch listed in the Firestore docs
// https://firebase.google.com/docs/firestore/manage-data/transactions
export const MAX_BATCH_OPS = 500;
// A batch processor that will accept any number of requests.
// It processes batches in the maximum number of events allowed.
// Does not guarantee atomicity across batches.
//
export default class DbBatch {
private opsCount: number;
private totalOpsCount: number;
private batchCount: number;
private currentBatch;
private promises: Array<Promise<any>>;
constructor() {
this.promises = [];
this.opsCount = 0;
this.totalOpsCount = 0;
this.batchCount = 0;
}
// Add a new document at path
create(path: string | Array<string>, data: any) {
const batch = this.getBatch();
const id = db.newId();
batch.set(db.path([path, id]), data);
return id;
}
// Add or replace a document at path
set(path: string | Array<string>, data: any) {
const batch = this.getBatch();
batch.set(db.path(path), data);
}
// Update properties on a document at path
update(path: string | Array<string>, data: any) {
const batch = this.getBatch();
batch.update(db.path(path), data);
}
// Delete a single document at path
delete(path: string | Array<string>) {
const batch = this.getBatch();
batch.delete(db.path(path));
}
// Delete all docs in a collection (but not any child collections)
deleteCollection(collectionName: string | Array<string>) {
const query = db.path(collectionName).orderBy('__name__');
this.deleteDocsMatchingQuery(query);
}
// Delete all docs matching the query provided, do not include
// limit attributes on the query. Those are added internally.
deleteDocsMatchingQuery(query: FirebaseFirestore.Query) {
// Unfortunately, there's no synchronous way to update the totalOpsCount
// based on number of docs affected, so we count this as a single DB op
// We do return the number of docs deleted in the promise.
this.totalOpsCount += 1;
this.batchCount += 1;
// Process batches of MAX_BATCH_OPS until we run out of docs
const nextBatch = (resolve, reject, limitQuery, totalDocsDeleted = 0) => {
const currentBatch = db.batch();
limitQuery.get().then(snapshot => {
totalDocsDeleted += snapshot.size;
if (snapshot.size > 0) {
snapshot.docs.forEach(doc => {
currentBatch.delete(doc.ref);
});
currentBatch.commit().then(() => {
process.nextTick(() => nextBatch(resolve, reject, limitQuery, totalDocsDeleted));
});
}
else {
resolve(totalDocsDeleted);
}
})
};
// For query-based deletes, there are going to be an unknown
// number of batches (each making asynchronous calls). This means that
// we can't depend on all the promises being collected in this.promises
// before commitAll() is invoked. We solve this by using a single
// promise and passing the resolve method indefinitely until we run out
// of batches.
this.promises.push(new Promise((resolve, reject) => {
const limitQuery = query.limit(MAX_BATCH_OPS);
nextBatch(resolve, reject, limitQuery);
}));
}
// Get the current number of batch operations that have been sent.
count() {
return this.totalOpsCount;
}
// Resolves when all batch operations have been completed.
commitAll() {
if( this.currentBatch ) {
this.promises.push(this.currentBatch.commit());
}
this.currentBatch = null;
return Promise.all(this.promises).then(() => this);
}
private getBatch() {
if( this.opsCount === MAX_BATCH_OPS ) {
this.promises.push(this.currentBatch.commit());
this.opsCount = 0;
}
if( this.opsCount === 0 ) {
this.batchCount++;
this.currentBatch = db.batch();
}
this.opsCount++;
this.totalOpsCount++;
return this.currentBatch;
}
static deleteCollections(list: Array<string>) {
const batch = new DbBatch();
list.forEach(url => batch.deleteCollection(url));
return batch.commitAll();
}
static loadSeedData(data) {
const batch = new DbBatch();
Object.keys(data).forEach(docPath => {
batch.set(docPath, data[docPath]);
});
return batch.commitAll();
}
static deleteRecursively(docPath: string|Array<string>) {
const batch = new DbBatch();
return this.recurseAndDelete(batch, db.path(docPath)).then(() => batch.commitAll());
}
private static recurseAndDelete(batch: DbBatch, docRef: FirebaseFirestore.DocumentReference) {
// Delete the current document
batch.getBatch().delete(docRef);
// Delete subcollections of the doc
return docRef.listCollections().then(subcollections => {
return Promise.all(subcollections.map(coll => {
return coll.listDocuments().then(docs => {
// For each document, rinse and repeat
return Promise.all(docs.map(doc => {
return this.recurseAndDelete(batch, doc);
}));
});
}));
})
}
}
//
// When calling DbData.loadSeedData(...), pass in an object similar to this
//
export default {
'users/U1TR3ljGptbEyihYUqiVhxOsoQn1': { name: 'Kato' },
'users/Md4mdB6cWOTNF2o7lXZf0s0BHdq2': { name: 'Puf' },
'users/e1sRcmQ7PudVx4gDsr8wuyosGOk1': { name: 'Randi' },
'users/3873ll6llkZqpxkx6v1ECq6TSOg2': { name: 'Jeff' },
'users/EFVDAJizX3fuD1o8FhFmJ28k5bt2': { name: 'Kiana' },
'users/jyb8sSRn1yawd3aTzNVo3u7CbOm2': { name: 'Mike' },
"groups/group1": { name: 'Group 1' },
'groups/group1/members/U1TR3ljGptbEyihYUqiVhxOsoQn1': {},
'groups/group1/members/Md4mdB6cWOTNF2o7lXZf0s0BHdq2': {},
// circular relationship
"groups/group3": { name: 'Group 3' },
'groups/group3/members/3873ll6llkZqpxkx6v1ECq6TSOg2': {},
'groups/group3/members/U1TR3ljGptbEyihYUqiVhxOsoQn1': {},
"groups/group4": { name: 'Group 4' },
'groups/group4/members/EFVDAJizX3fuD1o8FhFmJ28k5bt2': {},
'groups/group4/members/3873ll6llkZqpxkx6v1ECq6TSOg2': {},
"groups/group5": { name: 'Group 5' },
'groups/group5/members/jyb8sSRn1yawd3aTzNVo3u7CbOm2': {},
// linear relationship
"groups/group2": { name: 'Group 2' },
'groups/group2/members/U1TR3ljGptbEyihYUqiVhxOsoQn1': {},
'groups/group2/members/e1sRcmQ7PudVx4gDsr8wuyosGOk1': {},
"groups/group6": { name: 'Group 6' },
'groups/group6/members/EFVDAJizX3fuD1o8FhFmJ28k5bt2': {},
"groups/group7": { name: 'Group 7' },
"groups/group2/inherits/group1": {},
"groups/group3/inherits/group1": {},
"groups/group3/inherits/group4": {},
"groups/group4/inherits/group5": {},
"groups/group5/inherits/group3": {},
"groups/group6/inherits/group2": {},
"groups/group7/inherits/group6": {},
"docs/doc1": {
title: 'Doc 1',
owner: 'U1TR3ljGptbEyihYUqiVhxOsoQn1',
groups: ['group1'],
},
"docs/doc2": {
title: 'Doc 2',
owner: 'Md4mdB6cWOTNF2o7lXZf0s0BHdq2',
groups: ['group2']
},
"docs/doc3": {
title: 'Doc 3',
owner: '3873ll6llkZqpxkx6v1ECq6TSOg2',
groups: ['group3', 'group5']
},
"docs/doc4": {
title: 'Doc 4',
owner: 'U1TR3ljGptbEyihYUqiVhxOsoQn1',
groups: ['group4']
},
"docs/doc5": {
title: 'Doc 5',
owner: 'e1sRcmQ7PudVx4gDsr8wuyosGOk1',
},
"docs/doc6": {
title: 'Doc 6',
owner: 'EFVDAJizX3fuD1o8FhFmJ28k5bt2',
groups: ['group1']
}
};
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.