import _ from 'lodash'; import { AutoCastable, Prop, RPC_MARSHAL } from 'civkit/civ-rpc'; import { Firestore, FieldValue, DocumentReference, Query, Timestamp, SetOptions, DocumentSnapshot, } from '@google-cloud/firestore'; // Firestore doesn't support JavaScript objects with custom prototypes (i.e. objects that were created via the \"new\" operator) function patchFireStoreArrogance(func: Function) { return function (this: unknown) { const origObjectGetPrototype = Object.getPrototypeOf; Object.getPrototypeOf = function (x) { const r = origObjectGetPrototype.call(this, x); if (!r) { return r; } return Object.prototype; }; try { return func.call(this, ...arguments); } finally { Object.getPrototypeOf = origObjectGetPrototype; } }; } Reflect.set(DocumentReference.prototype, 'set', patchFireStoreArrogance(Reflect.get(DocumentReference.prototype, 'set'))); Reflect.set(DocumentSnapshot, 'fromObject', patchFireStoreArrogance(Reflect.get(DocumentSnapshot, 'fromObject'))); function mapValuesDeep(v: any, fn: (i: any) => any): any { if (_.isPlainObject(v)) { return _.mapValues(v, (i) => mapValuesDeep(i, fn)); } else if (_.isArray(v)) { return v.map((i) => mapValuesDeep(i, fn)); } else { return fn(v); } } export type Constructor = { new(...args: any[]): T; }; export type Constructed = T extends Partial ? U : T extends object ? T : object; export function fromFirestore( this: Constructor, id: string, overrideCollection?: string ): Promise; export async function fromFirestore( this: any, id: string, overrideCollection?: string ) { const collection = overrideCollection || this.collectionName; if (!collection) { throw new Error(`Missing collection name to construct ${this.name}`); } const ref = this.DB.collection(overrideCollection || this.collectionName).doc(id); const ptr = await ref.get(); if (!ptr.exists) { return undefined; } const doc = this.from( // Fixes non-native firebase types mapValuesDeep(ptr.data(), (i: any) => { if (i instanceof Timestamp) { return i.toDate(); } return i; }) ); Object.defineProperty(doc, '_ref', { value: ref, enumerable: false }); Object.defineProperty(doc, '_id', { value: ptr.id, enumerable: true }); return doc; } export function fromFirestoreQuery( this: Constructor, query: Query ): Promise; export async function fromFirestoreQuery(this: any, query: Query) { const ptr = await query.get(); if (ptr.docs.length) { return ptr.docs.map(doc => { const r = this.from( mapValuesDeep(doc.data(), (i: any) => { if (i instanceof Timestamp) { return i.toDate(); } return i; }) ); Object.defineProperty(r, '_ref', { value: doc.ref, enumerable: false }); Object.defineProperty(r, '_id', { value: doc.id, enumerable: true }); return r; }); } return []; } export function setToFirestore( this: Constructor, doc: T, overrideCollection?: string, setOptions?: SetOptions ): Promise; export async function setToFirestore( this: any, doc: any, overrideCollection?: string, setOptions?: SetOptions ) { let ref: DocumentReference = doc._ref; if (!ref) { const collection = overrideCollection || this.collectionName; if (!collection) { throw new Error(`Missing collection name to construct ${this.name}`); } const predefinedId = doc._id || undefined; const hdl = this.DB.collection(overrideCollection || this.collectionName); ref = predefinedId ? hdl.doc(predefinedId) : hdl.doc(); Object.defineProperty(doc, '_ref', { value: ref, enumerable: false }); Object.defineProperty(doc, '_id', { value: ref.id, enumerable: true }); } await ref.set(doc, { merge: true, ...setOptions }); return doc; } export function deleteQueryBatch( this: Constructor, query: Query ): Promise; export async function deleteQueryBatch(this: any, query: Query) { const snapshot = await query.get(); const batchSize = snapshot.size; if (batchSize === 0) { return; } // Delete documents in a batch const batch = this.DB.batch(); snapshot.docs.forEach((doc) => { batch.delete(doc.ref); }); await batch.commit(); process.nextTick(() => { this.deleteQueryBatch(query); }); }; export function fromFirestoreDoc( this: Constructor, snapshot: DocumentSnapshot, ): T | undefined; export function fromFirestoreDoc( this: any, snapshot: DocumentSnapshot, ) { const doc = this.from( // Fixes non-native firebase types mapValuesDeep(snapshot.data(), (i: any) => { if (i instanceof Timestamp) { return i.toDate(); } return i; }) ); Object.defineProperty(doc, '_ref', { value: snapshot.ref, enumerable: false }); Object.defineProperty(doc, '_id', { value: snapshot.id, enumerable: true }); return doc; } const defaultFireStore = new Firestore({ projectId: process.env.GCLOUD_PROJECT, }); export class FirestoreRecord extends AutoCastable { static collectionName?: string; static OPS = FieldValue; static DB = defaultFireStore; static get COLLECTION() { if (!this.collectionName) { throw new Error('Not implemented'); } return this.DB.collection(this.collectionName); } @Prop() _id?: string; _ref?: DocumentReference>>; static fromFirestore = fromFirestore; static fromFirestoreDoc = fromFirestoreDoc; static fromFirestoreQuery = fromFirestoreQuery; static save = setToFirestore; static deleteQueryBatch = deleteQueryBatch; [RPC_MARSHAL]() { return { ...this, _id: this._id, _ref: this._ref?.path }; } degradeForFireStore(): this { return JSON.parse(JSON.stringify(this, function (k, v) { if (k === '') { return v; } if (typeof v === 'object' && v && (typeof v.degradeForFireStore === 'function')) { return v.degradeForFireStore(); } return v; })); } }