Transaction cache layer to simplify its use?

[pepoospina]

Transactions have a significant limitation which is that you cannot read after you have already made one write to the DB.

On our app, we are using a Transaction manager that keeps all the writes in a memory cache, use it to return reads (if there is a cached write apply it to the document being read), and then, after everything is done, execute all writes at once.

I was not able to find an official library for this. Do you know of anything like it?

This is our code. You can also switch between a Transaction and a Batch operation. I thought this was a cool tool and so I wanted to share the code here.

Its not professional code, and there are probably many rough corners, but it could help inspire others or maybe an official tool by Firebase.

TransactionManager

import {
  DocumentReference,
  DocumentSnapshot,
  PartialWithFieldValue,
  Precondition,
  Query,
  QuerySnapshot,
  SetOptions,
  Timestamp,
  Transaction,
  UpdateData,
  WithFieldValue,
  WriteBatch,
} from 'firebase-admin/firestore';
import { merge } from 'lodash';

export interface TransactionCreate<T = any> {
  documentRef: DocumentReference<T>;
  data: WithFieldValue<T>;
}

export interface TransactionSet<T = any> {
  documentRef: DocumentReference<T>;
  data: PartialWithFieldValue<T>;
  options?: SetOptions;
}

export interface TransactionUpdate<T = any> {
  documentRef: DocumentReference<T>;
  data: UpdateData<T>;
  precondition?: Precondition;
}

export interface TransactionDelete<T = any> {
  documentRef: DocumentReference<T>;
  precondition?: Precondition;
}
enum MutationType {
  create = 'CREATE',
  set = 'SET',
  update = 'UPDATE',
  delete = 'DELETE',
}

export interface Mutation<T = any, D = any> {
  type: MutationType;
  documentRef: DocumentReference<T>;
  data?: D;
  options?: SetOptions;
  precondition?: Precondition;
}

/**
 * An interface that is common to the Transaction and WriteBatch types
 * in firestore
 */
export interface Mutator {
  get<T>(documentRef: DocumentReference<T>): Promise<DocumentSnapshot<T>>;
  create<T>(documentRef: DocumentReference<T>, data: WithFieldValue<T>): void;
  set<T>(
    documentRef: DocumentReference<T>,
    data: PartialWithFieldValue<T>,
    options: SetOptions
  ): void;
  set<T>(documentRef: DocumentReference<T>, data: WithFieldValue<T>): void;
  update<T>(
    documentRef: DocumentReference<T>,
    data: UpdateData<T>,
    precondition?: Precondition
  ): void;
  delete(
    documentRef: DocumentReference<any>,
    precondition?: Precondition
  ): void;
}

/** batches db writes to occur after all reads were done */
export class TransactionManager {
  /** an array of all the  */
  private mutations: Map<string, Mutation[]> = new Map();

  constructor(
    readonly transaction?: Transaction,
    readonly batch?: WriteBatch
  ) {}

  now() {
    return new Timestamp(Math.round(Date.now() / 1000), 0);
  }

  /**
   * Created and modified documents are hot cached by the manager to
   * allow for reads after they have been created/modified
   *
   * The cached object overwrites the properties of the object in firestore. The final result
   * of get will be an object with all the properties
   */
  private getCumMutationOf(id: string) {
    const mutations = this.mutations.get(id);
    if (!mutations) return {};

    /** merge all mutations data */
    const obj = mutations.reduce((_obj, mut) => {
      switch (mut.type) {
        case MutationType.create:
        case MutationType.set:
        case MutationType.update:
          merge(_obj, mut.data);
          break;

        case MutationType.delete:
          return {};
      }

      return _obj;
    }, {});

    return obj;
  }

  async mutate<D>(documentRef: DocumentReference<any>, mutation: Mutation<D>) {
    const current = this.mutations.get(documentRef.id) || [];
    current.push(mutation);
    return this.mutations.set(documentRef.id, current);
  }

  async get<T>(
    documentRef: DocumentReference<T>
  ): Promise<DocumentSnapshot<T>> {
    const cached = this.getCumMutationOf(documentRef.id);
    const onFirestore = this.transaction
      ? await this.transaction.get(documentRef)
      : await documentRef.get();

    const mergedData = merge(onFirestore.data(), cached);

    return {
      id: documentRef.id,
      ref: documentRef,
      exists: cached !== undefined || onFirestore.exists,
      readTime: this.now(),
      data: () => mergedData,
      get: () => {
        throw new Error('not implemented');
      },
      isEqual: () => {
        throw new Error('not implemented');
      },
    };
  }

  /** queries cannot be cached, so the always hit firestore. This
   * method 'query', different from 'get' must be used to query
   */
  query<T>(query: Query<T>): Promise<QuerySnapshot<T>> {
    return this.transaction ? this.transaction.get(query) : query.get();
  }

  create(documentRef: DocumentReference<any>, data: WithFieldValue<any>): void {
    this.mutate(documentRef, {
      type: MutationType.create,
      documentRef,
      data,
    });
  }

  set(
    documentRef: DocumentReference<any>,
    data: PartialWithFieldValue<any>,
    options?: SetOptions
  ): void {
    this.mutate(documentRef, {
      type: MutationType.set,
      documentRef,
      data,
      options,
    });
  }

  update(
    documentRef: DocumentReference<any>,
    data: UpdateData<any>,
    precondition?: Precondition
  ): void {
    this.mutate(documentRef, {
      type: MutationType.update,
      documentRef,
      data,
      precondition,
    });
  }

  delete(
    documentRef: DocumentReference<any>,
    precondition?: Precondition
  ): void {
    this.mutate(documentRef, {
      type: MutationType.delete,
      documentRef,
      precondition,
    });
  }

  get mutator(): Mutator {
    if (this.transaction) {
      return this.transaction;
    }
    if (this.batch) {
      return this.batch as unknown as Mutator;
    }

    throw new Error('Either transaction or bath need to be defined');
  }

  async applyWrites(): Promise<void> {
    /** apply all writes of different element concurrently */

    this.mutations.forEach((mutations) => {
      /** apply writes of one element sequentially */
      for (const mut of mutations) {
        switch (mut.type) {
          case MutationType.create:
            this.mutator.create(mut.documentRef, mut.data);
            break;

          case MutationType.set:
            mut.options
              ? this.mutator.set(mut.documentRef, mut.data, mut.options)
              : this.mutator.set(mut.documentRef, mut.data);
            break;

          case MutationType.update:
            mut.precondition
              ? this.mutator.update(mut.documentRef, mut.data, mut.precondition)
              : this.mutator.update(mut.documentRef, mut.data);
            break;

          case MutationType.delete:
            mut.precondition
              ? this.mutator.delete(mut.documentRef, mut.precondition)
              : this.mutator.delete(mut.documentRef);
            break;
        }
      }
    });

    if (this.batch) {
      await this.batch.commit();
    }
  }
}

export type HandleWithTxManager<R = void, P = void> = (
  manager: TransactionManager,
  payload?: P
) => Promise<R>;

export enum ManagerModes {
  /** batch all writes until after the function call */
  BATCH = 'BATCH',
  /** snapshot reads at function start and batch all writes until after the function call */
  TRANSACTION = 'TRANSACTION',
}

export interface ManagerConfig {
  mode: ManagerModes;
}

DBInstance

import { initializeApp } from 'firebase-admin/app';
import { Firestore, getFirestore } from 'firebase-admin/firestore';

import { CollectionNames } from '../@shared/utils/collectionNames';
import { IS_EMULATOR } from '../config/config.runtime';
import { logger } from '../instances/logger';
// import { SERVICE_ACCOUNT_ID } from '../config/config.runtime';
import {
  HandleWithTxManager,
  ManagerConfig,
  ManagerModes,
  TransactionManager,
} from './transaction.manager';

export const app = IS_EMULATOR
  ? initializeApp({
      projectId: 'demo-sensenets',
    })
  : initializeApp();

const DEBUG = false;

export class DBInstance {
  public firestore: Firestore;

  constructor() {
    this.firestore = getFirestore();
  }

  /** a wrapper of TransactionManager to instantiate and applyWrites automatically */
  async run<R, P>(
    func: HandleWithTxManager<R, P>,
    payload?: P,
    config: ManagerConfig = { mode: ManagerModes.TRANSACTION }
  ): Promise<R> {
    switch (config.mode) {
      case ManagerModes.TRANSACTION:
        return this.firestore.runTransaction(async (transaction) => {
          if (DEBUG) logger.debug('Transaction started');
          const manager = new TransactionManager(transaction);
          
          const result = await func(manager, payload);
          if (DEBUG) logger.debug('Transaction function executed');
          
          await manager.applyWrites();
          if (DEBUG) logger.debug('Transaction writes applied');
          return result;
        });

      case ManagerModes.BATCH: {
        const batch = this.firestore.batch();
        const manager = new TransactionManager(undefined, batch);
        const result = await func(manager, payload);
        await manager.applyWrites();
        return result;
      }
    }
  }
}

And an example of how we use it:

  db.run(async (manager) => {
      return postsManager.processing.createOrUpdatePostDrafts(
        payload.postId,
        manager
      );
    });

We need to send the manager as a parameter from the upper controller layers, down to our DB repository services. But it works just fine. Also, writes are synchronous operations (as they operate on the cache).

[pepoospina]

@Feiyang1

sorry for the ping. Is there a tool that do this already? I find it extremely useful