botversion-sdk 2.2.3 → 2.3.1

package/undo/lib/capture-sql.js

@@ -0,0 +1,632 @@
+// lib/capture-sql.js
+"use strict";
+// PG CAPTURE (Phase 2): saves before/after rows of tagged UPDATE and INSERT statements on the host's own connection.
+// Never throws into the host. Logs: table names, counts and reasons only.
+
+const log = require("./log");
+const tagger = require("./tagger");
+const store = require("./snapshot-store");
+const guard = require("./capture-guard");
+const gate = require("./gate");
+
+// Statements that start with one of these words are writes, even when the parser cannot read them.
+const WRITE_WORDS = new Set([
+  "INSERT",
+  "UPDATE",
+  "DELETE",
+  "TRUNCATE",
+  "DROP",
+  "ALTER",
+  "MERGE",
+]);
+
+const { classify, parseSingle } = require("./sql-classify");
+
+let toSql = null;
+try {
+  toSql = require("pgsql-ast-parser").toSql;
+} catch (_) {
+  /* parser missing: nothing is capturable */
+}
+
+const txState = new WeakMap(); // client -> is the host inside a transaction?
+const pkCache = new Map();
+
+const MAX_DEPTH = 5; // how many levels of linked rows a delete may reach
+class Skip extends Error {}
+
+// Finds the real schema and name of the table a statement points to.
+const RESOLVE_SQL =
+  "SELECT n.nspname::text AS s, c.relname::text AS n FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE c.oid = $1::text::regclass";
+
+// Foreign keys that point AT a table. rule: c = cascade, n = set null, d = set default, a/r = the database refuses the delete.
+const FK_SQL =
+  "SELECT cn.nspname::text AS child_schema, cl.relname::text AS child_table, c.confdeltype::text AS rule, " +
+  "(SELECT array_agg(a.attname::text ORDER BY k.ord) FROM unnest(c.conkey) WITH ORDINALITY k(attnum, ord) JOIN pg_attribute a ON a.attrelid = c.conrelid AND a.attnum = k.attnum) AS child_cols, " +
+  "(SELECT array_agg(a.attname::text ORDER BY k.ord) FROM unnest(c.confkey) WITH ORDINALITY k(attnum, ord) JOIN pg_attribute a ON a.attrelid = c.confrelid AND a.attnum = k.attnum) AS parent_cols " +
+  "FROM pg_constraint c JOIN pg_class cl ON cl.oid = c.conrelid JOIN pg_namespace cn ON cn.oid = cl.relnamespace " +
+  "WHERE c.contype = 'f' AND c.confrelid = $1::text::regclass";
+
+const q = (s) => '"' + String(s).replace(/"/g, '""') + '"';
+const tableSql = (t) => (t.schema ? q(t.schema) + "." : "") + q(t.name);
+const label = (t) => (t.schema ? t.schema + "." : "") + t.name;
+const keyOf = (row, cols) =>
+  cols
+    .map((c) => {
+      const v = row[c];
+      return v instanceof Date ? v.toISOString() : String(v);
+    })
+    .join("\u0001");
+const pickPk = (row, pk) => {
+  const o = {};
+  for (const c of pk) o[c] = row[c];
+  return o;
+};
+
+// Gives each $n parameter a new number (1, 2, 3...) so the read only gets the values it uses.
+function remap(node, map, order) {
+  if (Array.isArray(node)) return node.map((n) => remap(n, map, order));
+  if (!node || typeof node !== "object") return node;
+  if (node.type === "parameter" && /^\$\d+$/.test(node.name)) {
+    const idx = parseInt(node.name.slice(1), 10);
+    if (!map.has(idx)) {
+      map.set(idx, map.size + 1);
+      order.push(idx);
+    }
+    return Object.assign({}, node, { name: "$" + map.get(idx) });
+  }
+  const out = {};
+  for (const k of Object.keys(node))
+    if (k !== "_location") out[k] = remap(node[k], map, order);
+  return out;
+}
+
+// Remembers whether the host's own statements opened a transaction on this connection.
+function trackTransaction(client, text) {
+  if (/^\s*ROLLBACK(\s+(WORK|TRANSACTION))?\s+TO\b/i.test(text)) return;
+  const m = /^\s*([A-Za-z]+)/.exec(text);
+  const w = m ? m[1].toUpperCase() : "";
+  if (w === "BEGIN" || w === "START") txState.set(client, true);
+  else if (w === "COMMIT" || w === "ROLLBACK" || w === "END" || w === "ABORT")
+    txState.set(client, false);
+}
+
+async function primaryKeyOf(run, table, inTx) {
+  const key = tableSql(table);
+  if (pkCache.has(key)) return pkCache.get(key);
+  const res = await store.guarded(run, inTx, () =>
+    run(
+      "SELECT a.attname AS name FROM pg_index i JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey) " +
+        "WHERE i.indrelid = $1::text::regclass AND i.indisprimary ORDER BY array_position(i.indkey::int2[], a.attnum)",
+      [key],
+    ),
+  );
+  const pk = res.rows.map((r) => r.name);
+  if (pk.length) pkCache.set(key, pk);
+  return pk;
+}
+
+async function readByPk(run, table, pk, pkObjs, inTx) {
+  if (!pkObjs.length) return [];
+  const params = [];
+  const tuples = pkObjs.map(
+    (o) =>
+      "(" +
+      pk
+        .map((c) => {
+          params.push(o[c]);
+          return "$" + params.length;
+        })
+        .join(", ") +
+      ")",
+  );
+  const sql = `SELECT * FROM ${tableSql(table)} WHERE (${pk.map(q).join(", ")}) IN (${tuples.join(", ")})`;
+  return (await store.guarded(run, inTx, () => run(sql, params))).rows;
+}
+
+async function beforeUpdate(run, stmt, values, st, inTx, notSaved) {
+  const pk = await primaryKeyOf(run, stmt.table, inTx);
+  let c = store.checkCapturable({ rowCount: 0, primaryKeyColumns: pk });
+  if (!c.ok) {
+    notSaved(c.reason);
+    return null;
+  }
+  const map = new Map();
+  const order = [];
+  const sel = {
+    type: "select",
+    columns: [{ expr: { type: "ref", name: "*" } }],
+    from: [{ type: "table", name: stmt.table }],
+    limit: { limit: { type: "integer", value: store.ROW_LIMIT + 1 } },
+  };
+  if (stmt.where) sel.where = remap(stmt.where, map, order);
+  const sql = toSql.statement(sel);
+  const rows = (
+    await store.guarded(run, inTx, () =>
+      run(
+        sql,
+        order.map((i) => values[i - 1]),
+      ),
+    )
+  ).rows;
+  c = store.checkCapturable({ rowCount: rows.length, primaryKeyColumns: pk });
+  if (!c.ok) {
+    notSaved(c.reason);
+    return null;
+  }
+  const saved = [];
+  for (const row of rows) {
+    const pkObj = pickPk(row, pk);
+    const r = await store.saveSnapshot(
+      run,
+      st,
+      {
+        table: label(stmt.table),
+        operation: "update",
+        primaryKey: pkObj,
+        before: row,
+        after: null,
+        isLinked: false,
+      },
+      { inTransaction: inTx },
+    );
+    if (!r.ok) {
+      notSaved(r.reason);
+      break;
+    }
+    saved.push({ id: r.id, key: keyOf(row, pk), pkObj });
+  }
+  return { pk, saved };
+}
+
+async function afterUpdate(run, stmt, before, st, inTx) {
+  const rows = await readByPk(
+    run,
+    stmt.table,
+    before.pk,
+    before.saved.map((s) => s.pkObj),
+    inTx,
+  );
+  const byKey = new Map(before.saved.map((s) => [s.key, s]));
+  let n = 0;
+  for (const row of rows) {
+    const s = byKey.get(keyOf(row, before.pk));
+    if (!s) continue;
+    const r = await store.saveAfter(run, s.id, row, { inTransaction: inTx });
+    if (r.ok) n++;
+  }
+}
+
+async function afterInsert(run, stmt, result, st, inTx, notSaved) {
+  const pk = await primaryKeyOf(run, stmt.into, inTx);
+  const c = store.checkCapturable({
+    rowCount: result && result.rows ? result.rows.length : 0,
+    primaryKeyColumns: pk,
+  });
+  if (!c.ok) return notSaved(c.reason);
+  const rows = result && result.rows;
+  if (
+    !stmt.returning ||
+    !Array.isArray(rows) ||
+    !rows.length ||
+    !pk.every((k) => k in rows[0])
+  ) {
+    return notSaved(
+      "new row primary key not returned - add RETURNING <primary key> in your own code",
+    );
+  }
+  const full = await readByPk(
+    run,
+    stmt.into,
+    pk,
+    rows.map((r) => pickPk(r, pk)),
+    inTx,
+  );
+  let n = 0;
+  for (const row of full) {
+    const r = await store.saveSnapshot(
+      run,
+      st,
+      {
+        table: label(stmt.into),
+        operation: "insert",
+        primaryKey: pickPk(row, pk),
+        before: null,
+        after: row,
+        isLinked: false,
+      },
+      { inTransaction: inTx },
+    );
+    if (!r.ok) return notSaved(r.reason);
+    n++;
+  }
+}
+
+// Rows of a child table whose link columns match the given parent rows.
+async function readByCols(run, table, childCols, parentCols, parentRows, inTx) {
+  const params = [];
+  const tuples = parentRows.map(
+    (r) =>
+      "(" +
+      parentCols
+        .map((c) => {
+          params.push(r[c]);
+          return "$" + params.length;
+        })
+        .join(", ") +
+      ")",
+  );
+  const sql = `SELECT * FROM ${tableSql(table)} WHERE (${childCols.map(q).join(", ")}) IN (${tuples.join(", ")}) LIMIT ${store.ROW_LIMIT + 1}`;
+  return (await store.guarded(run, inTx, () => run(sql, params))).rows;
+}
+
+// Works out every row a DELETE will remove or change (through foreign keys). Saves nothing. Throws Skip when unsafe.
+async function planDelete(run, stmt, values, inTx) {
+  const res = await store.guarded(run, inTx, () =>
+    run(RESOLVE_SQL, [tableSql(stmt.from)]),
+  );
+  if (!res.rows.length) throw new Skip("table not found");
+  const root = { schema: res.rows[0].s, name: res.rows[0].n };
+  const rootPk = await primaryKeyOf(run, root, inTx);
+  let c = store.checkCapturable({ rowCount: 0, primaryKeyColumns: rootPk });
+  if (!c.ok) throw new Skip(c.reason);
+
+  // 1. The rows the host is deleting.
+  const map = new Map();
+  const order = [];
+  const sel = {
+    type: "select",
+    columns: [{ expr: { type: "ref", name: "*" } }],
+    from: [{ type: "table", name: stmt.from }],
+    limit: { limit: { type: "integer", value: store.ROW_LIMIT + 1 } },
+  };
+  if (stmt.where) sel.where = remap(stmt.where, map, order);
+  const sql = toSql.statement(sel);
+  const rootRows = (
+    await store.guarded(run, inTx, () =>
+      run(
+        sql,
+        order.map((i) => values[i - 1]),
+      ),
+    )
+  ).rows;
+  c = store.checkCapturable({
+    rowCount: rootRows.length,
+    primaryKeyColumns: rootPk,
+  });
+  if (!c.ok) throw new Skip(c.reason);
+
+  // 2. Rows reached through foreign keys, level by level.
+  //    kind "delete" = removed by the database (cascade); kind "update" = link cleared (set null / default).
+  const groups = [
+    { table: root, pk: rootPk, rows: rootRows, depth: 0, kind: "delete" },
+  ];
+  const tag = (t, row, pk) => label(t) + "\u0002" + keyOf(row, pk);
+  const deletedAt = new Map();
+  const cleared = new Set();
+  for (const r of rootRows) deletedAt.set(tag(root, r, rootPk), 0);
+  let frontier = rootRows.length ? [groups[0]] : [];
+  let total = rootRows.length;
+  for (let depth = 1; frontier.length; depth++) {
+    const next = [];
+    for (const parent of frontier) {
+      const links = (
+        await store.guarded(run, inTx, () =>
+          run(FK_SQL, [tableSql(parent.table)]),
+        )
+      ).rows;
+      for (const link of links) {
+        const cascades = link.rule === "c";
+        if (!cascades && link.rule !== "n" && link.rule !== "d") continue; // restrict / no action
+        const child = { schema: link.child_schema, name: link.child_table };
+        const found = await readByCols(
+          run,
+          child,
+          link.child_cols,
+          link.parent_cols,
+          parent.rows,
+          inTx,
+        );
+        if (!found.length) continue;
+        if (depth > MAX_DEPTH) throw new Skip("links go too deep");
+        const childPk = await primaryKeyOf(run, child, inTx);
+        const cc = store.checkCapturable({
+          rowCount: found.length,
+          primaryKeyColumns: childPk,
+        });
+        if (!cc.ok) throw new Skip(cc.reason);
+        const fresh = [];
+        for (const r of found) {
+          const k = tag(child, r, childPk);
+          if (cascades) {
+            if (deletedAt.has(k)) {
+              // Same row reached at another level would be restored in the wrong order: do not guess.
+              if (deletedAt.get(k) !== depth)
+                throw new Skip(
+                  "a row is reached through links of different depth",
+                );
+              continue;
+            }
+            deletedAt.set(k, depth);
+          } else {
+            if (cleared.has(k)) continue;
+            cleared.add(k);
+          }
+          fresh.push(r);
+        }
+        total += fresh.length;
+        if (total > store.ROW_LIMIT) throw new Skip("too-many-rows");
+        if (!fresh.length) continue;
+        const grp = {
+          table: child,
+          pk: childPk,
+          rows: fresh,
+          depth,
+          kind: cascades ? "delete" : "update",
+        };
+        groups.push(grp);
+        if (cascades) next.push(grp);
+      }
+    }
+    frontier = next;
+  }
+  // A row that is deleted does not also need its link cleared.
+  for (const g of groups) {
+    if (g.kind === "update")
+      g.rows = g.rows.filter((r) => !deletedAt.has(tag(g.table, r, g.pk)));
+  }
+  return groups.filter((g) => g.rows.length);
+}
+
+// Saves the rows a DELETE will remove or change. Deepest rows first, deleted rows last (Undo goes last to first).
+async function beforeDelete(run, stmt, values, st, inTx, notSaved) {
+  let groups;
+  try {
+    groups = await planDelete(run, stmt, values, inTx);
+  } catch (e) {
+    if (e instanceof Skip) {
+      notSaved(e.message);
+      return null;
+    }
+    throw e;
+  }
+  const ordered = groups.slice().sort((a, b) => b.depth - a.depth);
+  const saved = [];
+  const fills = [];
+  for (const g of ordered) {
+    const ids = new Map();
+    for (const row of g.rows) {
+      const r = await store.saveSnapshot(
+        run,
+        st,
+        {
+          table: label(g.table),
+          operation: g.kind,
+          primaryKey: pickPk(row, g.pk),
+          before: row,
+          after: null,
+          isLinked: g.depth > 0,
+        },
+        { inTransaction: inTx },
+      );
+      if (!r.ok) {
+        // Do not keep a half-saved set: it could restore only part of the data.
+        if (saved.length) {
+          try {
+            await store.guarded(run, inTx, () =>
+              run(`DELETE FROM ${store.TABLE} WHERE id = ANY($1::uuid[])`, [
+                saved.map((s) => s.id),
+              ]),
+            );
+          } catch (_) {
+            /* ignore */
+          }
+        }
+        notSaved(r.reason);
+        return null;
+      }
+      saved.push({ id: r.id });
+      ids.set(keyOf(row, g.pk), r.id);
+    }
+    if (g.kind === "update") fills.push({ g, ids });
+  }
+  return { saved, fills };
+}
+
+// Rows whose link was cleared: save how they look now (needed for the safety check in Undo).
+async function afterDelete(run, before, st, inTx, name) {
+  for (const { g, ids } of before.fills) {
+    const rows = await readByPk(
+      run,
+      g.table,
+      g.pk,
+      g.rows.map((r) => pickPk(r, g.pk)),
+      inTx,
+    );
+    for (const row of rows) {
+      const id = ids.get(keyOf(row, g.pk));
+      if (!id) continue;
+      await store.saveAfter(run, id, row, { inTransaction: inTx });
+    }
+  }
+}
+
+// Before an INSERT runs: can its new row be saved? (It needs a primary key and RETURNING that includes it.)
+async function beforeInsert(run, stmt, inTx, notSaved) {
+  const pk = await primaryKeyOf(run, stmt.into, inTx);
+  const c = store.checkCapturable({ rowCount: 0, primaryKeyColumns: pk });
+  if (!c.ok) return notSaved(c.reason);
+  const cols = Array.isArray(stmt.returning) ? stmt.returning : [];
+  const refs = cols
+    .filter((x) => x && x.expr && x.expr.type === "ref" && !x.alias)
+    .map((x) => x.expr.name);
+  if (!(refs.includes("*") || pk.every((k) => refs.includes(k)))) {
+    notSaved(
+      "new row primary key not returned - add RETURNING <primary key> in your own code",
+    );
+  }
+}
+
+async function runCapture(client, rawRun, stmt, values, st, hostCall) {
+  const inTx = txState.get(client) === true;
+  const name = label(
+    stmt.type === "update"
+      ? stmt.table
+      : stmt.type === "delete"
+        ? stmt.from
+        : stmt.into,
+  );
+  let phase = "before"; // "before" = the write has not run yet, so it can still be blocked
+  let failedWhy = null;
+  const notSaved = (why) => {
+    if (phase === "before") failedWhy = failedWhy || why;
+    if (phase !== "before" || st.confirmed) st.unsaved = true; // the write runs (or already ran) without a saved snapshot
+  };
+
+  let before = null;
+  try {
+    if (stmt.type === "update")
+      before = await beforeUpdate(rawRun, stmt, values, st, inTx, notSaved);
+    else if (stmt.type === "delete")
+      before = await beforeDelete(rawRun, stmt, values, st, inTx, notSaved);
+    else if (stmt.type === "insert")
+      await beforeInsert(rawRun, stmt, inTx, notSaved);
+  } catch (e) {
+    notSaved("capture error " + ((e && e.code) || "unknown"));
+  }
+
+  // Cannot be saved: remove what this statement saved (a half-saved set could restore only part of the data), then block unless confirmed.
+  if (failedWhy) {
+    const ids = before && before.saved ? before.saved.map((s) => s.id) : [];
+    if (ids.length) {
+      try {
+        await store.guarded(rawRun, inTx, () =>
+          rawRun(`DELETE FROM ${store.TABLE} WHERE id = ANY($1::uuid[])`, [
+            ids,
+          ]),
+        );
+      } catch (_) {
+        /* ignore */
+      }
+    }
+    before = null; // nothing from this statement stays half-saved
+    if (!st.confirmed) await gate.block(st, name, failedWhy);
+  }
+
+  phase = "after";
+  let result;
+  try {
+    result = await hostCall();
+  } catch (e) {
+    // The host's write failed: remove the snapshots saved for it (inside a transaction they roll back by themselves).
+    if (before && before.saved && before.saved.length && !inTx) {
+      try {
+        await rawRun(`DELETE FROM ${store.TABLE} WHERE id = ANY($1::uuid[])`, [
+          before.saved.map((s) => s.id),
+        ]);
+      } catch (_) {
+        /* ignore */
+      }
+    }
+    throw e;
+  }
+  try {
+    if (stmt.type === "update" && before && before.saved.length)
+      await afterUpdate(rawRun, stmt, before, st, inTx);
+    if (stmt.type === "insert")
+      await afterInsert(rawRun, stmt, result, st, inTx, notSaved);
+    if (stmt.type === "delete" && before && before.saved.length)
+      await afterDelete(rawRun, before, st, inTx, name);
+  } catch (e) {
+    notSaved("capture error " + ((e && e.code) || "unknown"));
+  }
+  return result;
+}
+
+// Returns { value } when capture took over (it runs the host's query itself, or blocks it), otherwise null.
+function maybeCapture(client, args, original) {
+  const st = tagger.getCurrent();
+  if (!st || guard.isUnhooked() || !toSql) return null;
+  const first = args[0];
+  const text =
+    typeof first === "string"
+      ? first
+      : first && typeof first.text === "string"
+        ? first.text
+        : null;
+  if (text === null) return null;
+  trackTransaction(client, text);
+  if (
+    first &&
+    typeof first === "object" &&
+    (first.rowMode === "array" || typeof first.submit === "function")
+  )
+    return null;
+
+  const last = args[args.length - 1];
+  const cb = typeof last === "function" ? last : null;
+
+  // A write that cannot be saved: blocked until the user confirms (or runs unsaved if they already did).
+  const refuse = (table, why) => {
+    if (st.confirmed) {
+      st.unsaved = true; // runs without a saved snapshot, so Undo will not cover it
+      return null;
+    }
+    const p = gate.block(st, table, why);
+    if (cb) {
+      p.then(null, (e) => process.nextTick(cb, e));
+      return { value: undefined };
+    }
+    return { value: p };
+  };
+
+  const c = classify(text);
+  if (c.parseFailed)
+    return WRITE_WORDS.has(c.firstWord)
+      ? refuse("?", "statement could not be read")
+      : null;
+  if (c.writes.length === 0) return null; // a read
+  const tbl = c.writes[0].tables.filter(Boolean)[0] || "?";
+  if (c.multi || c.writes.length !== 1)
+    return refuse(tbl, "several writes in one statement");
+  const w = c.writes[0];
+  if (w.kind !== "update" && w.kind !== "insert" && w.kind !== "delete")
+    return refuse(tbl, w.kind + " cannot be saved");
+  const bad = w.notes.find((n) =>
+    ["has-subquery", "upsert", "insert-from-select"].includes(n),
+  );
+  if (bad) return refuse(tbl, bad + " cannot be saved");
+  const stmt = parseSingle(text);
+  if (
+    !stmt ||
+    stmt.type !== w.kind ||
+    (stmt.type === "update" && stmt.from) ||
+    (stmt.type === "delete" && stmt.using)
+  )
+    return refuse(tbl, "this statement shape cannot be saved");
+
+  const hostArgs = Array.prototype.slice.call(args, 0, cb ? -1 : args.length);
+  let values =
+    typeof first === "string"
+      ? hostArgs[1]
+      : first.values !== undefined
+        ? first.values
+        : hostArgs[1];
+  if (!Array.isArray(values)) values = [];
+
+  // original = pg's own query function, so our reads and saves are never seen by the hook again.
+  const rawRun = (sql, params) => original.call(client, sql, params);
+  const p = runCapture(client, rawRun, stmt, values, st, () =>
+    Promise.resolve(original.apply(client, hostArgs)),
+  );
+  if (cb) {
+    p.then(
+      (r) => process.nextTick(cb, null, r),
+      (e) => process.nextTick(cb, e),
+    );
+    return { value: undefined };
+  }
+  return { value: p };
+}
+
+module.exports = { maybeCapture };