durable.ts

/*

Whats the difference between a workflow and a step?

- workflow is the high level things that gets run.
- steps are cached.

*/
import BetterSqlite3 from "better-sqlite3";
import mapValues from "lodash/mapValues";

export type Database = {
  get(key: string): any;
  set(key: string, value: any): void;
};

export class InMemoryKV implements Database {
  private data: Record<string, any> = {};
  get(key: string): any {
    return this.data[key];
  }
  set(key: string, value: any): void {
    this.data[key] = value;
  }
}

// https://github.com/tc39/proposal-async-context
export type DurableFn<I extends any[] = any[], O = any> = {
  name: string;
  fn: (...args: I) => Promise<O>;
};

export function durable<I extends any[], O>(db: Database, fn: DurableFn<I, O>): DurableFn<I, O>["fn"] {
  return async (...args: I) => {
    const key = `${fn.name}(${args.map((x) => JSON.stringify(x)).join(",")})`;
    const cached = db.get(key);
    if (cached) return cached;
    const result = await fn.fn(...args);
    db.set(key, result);
    return result;
  };
}

export function durables<F extends { [key: string]: DurableFn["fn"] }>(db: Database, fns: F): F {
  return mapValues(fns, (fn, name) => durable(db, { name, fn })) as F;
}

export class SQLiteKV implements Database {
  public db: BetterSqlite3.Database;
  private getStmt: BetterSqlite3.Statement;
  private upsertStmt: BetterSqlite3.Statement;

  constructor(filePath: string) {
    this.db = new BetterSqlite3(filePath);
    this.db.exec(
      `CREATE TABLE IF NOT EXISTS kv (
        key TEXT PRIMARY KEY,
        value TEXT NOT NULL
      )`
    );
    this.getStmt = this.db.prepare("SELECT value FROM kv WHERE key = ?");
    this.upsertStmt = this.db.prepare(
      "INSERT INTO kv (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value=excluded.value"
    );
  }

  get(key: string): any {
    const row = this.getStmt.get(key) as { value: string } | undefined;
    if (!row) return undefined;
    return JSON.parse(row.value);
  }

  set(key: string, value: any): void {
    const payload = JSON.stringify(value);
    this.upsertStmt.run(key, payload);
  }
}

export type Cache = Database & {
  cached<F extends { [key: string]: DurableFn["fn"] }>(fns: F): F;
};

export class SQLiteCache extends SQLiteKV implements Cache {
  cached<F extends { [key: string]: DurableFn["fn"] }>(fns: F): F {
    return durables(this, fns);
  }
}