scripts/loaders/miami/atomic-load.ts

// Keep candidate loads off the served tables until all their writes can commit.
// This module has no connection/env side effects so failure paths can be tested.
import crypto from "node:crypto";
import type { MiamiSubgroup, Prisma, PrismaClient } from "@prisma/client";

export type Fact = {
  schoolNumber: string;
  year: string;
  metricKey: string;
  subgroup?: MiamiSubgroup;
  value: number | null;
  suppressed?: boolean;
  denominator?: number | null;
};
export type LoadOptions = {
  sourceId: string;
  sourceUrl?: string;
  yearCovered?: string;
  scriptPath: string;
  sha256?: string;
  bytes?: number;
  notes?: string;
  gitSha?: string | null;
};
type Counts = { inserted: number; updated: number; suppressed: number };
type Operation = (tx: Prisma.TransactionClient) => Promise<unknown>;
type Pending = { opts: LoadOptions; previousId: string | null; facts: Fact[]; operations: Operation[]; hasFacts: boolean };

export function createDataLoader(db: Pick<PrismaClient, "$transaction" | "miamiDataLoad">) {
  const pending = new Map<string, Pending>();
  const latest = (client: Pick<Prisma.TransactionClient, "miamiDataLoad">, sourceId: string) => client.miamiDataLoad.findFirst({
    where: { sourceId, OR: [{ rowsInserted: { gt: 0 } }, { rowsUpdated: { gt: 0 } }] },
    orderBy: [{ loadedAt: "desc" }, { id: "desc" }], select: { id: true },
  });
  const get = (id: string) => {
    const load = pending.get(id);
    if (!load) throw new Error(`[load] unknown or already finished candidate ${id}`);
    return load;
  };

  async function startDataLoad(opts: LoadOptions) {
    const id = crypto.randomUUID();
    const previous = await latest(db, opts.sourceId);
    pending.set(id, { opts: { ...opts }, previousId: previous?.id ?? null, facts: [], operations: [], hasFacts: false });
    return { id };
  }

  async function upsertFacts(facts: Fact[], id: string, _concurrency = 8) {
    const load = get(id);
    for (const f of facts) {
      if (!f.schoolNumber || !f.year || !f.metricKey || (f.value !== null && !Number.isFinite(f.value)) ||
        (f.denominator != null && !Number.isFinite(f.denominator))) {
        throw new Error(`[load] invalid candidate fact ${JSON.stringify(f)}`);
      }
    }
    load.hasFacts = true;
    load.facts.push(...facts.map((f) => ({ ...f })));
    // Counts describe staged candidates. They become served only at finishDataLoad.
    return { written: facts.length, suppressed: facts.filter((f) => f.suppressed).length, fkFailures: new Set<string>() };
  }

  function stageLoadOperation(id: string, operation: Operation) {
    get(id).operations.push(operation);
  }

  async function finishDataLoad(id: string, counts: Counts, notes?: string) {
    const load = get(id);
    try {
      if (counts.inserted + counts.updated === 0 || (load.hasFacts && load.facts.length === 0) ||
        (!load.hasFacts && load.operations.length === 0)) {
        throw new Error(`[load] ${load.opts.sourceId}: no candidate writes; previous served load preserved`);
      }
      if (load.hasFacts && counts.inserted !== load.facts.length) {
        throw new Error(`[load] ${load.opts.sourceId}: staged fact count disagrees with completion count`);
      }
      return await db.$transaction(async (tx) => {
        // The lock prevents interleaved replacements. Rechecking the generation
        // also rejects a candidate staged against a source that changed meanwhile.
        await tx.$executeRaw`SELECT pg_advisory_xact_lock(hashtext(${'miami-load:' + load.opts.sourceId}))`;
        const previous = await latest(tx, load.opts.sourceId);
        if ((previous?.id ?? null) !== load.previousId) {
          throw new Error(`[load] ${load.opts.sourceId}: another load committed while this candidate was staged; rerun from current inputs`);
        }
        const keys = [...new Set(load.facts.map((f) => f.metricKey))];
        if (keys.length) {
          const registered = new Set((await tx.miamiMetricDefinition.findMany({ where: { key: { in: keys } }, select: { key: true } })).map((r) => r.key));
          const absent = keys.filter((k) => !registered.has(k));
          if (absent.length) throw new Error(`[facts] unregistered metric keys (run metric_definitions.ts first): ${absent.join(", ")}`);
        }
        await tx.miamiDataLoad.create({ data: { id, ...load.opts } });
        for (const f of load.facts) {
          const subgroup = f.subgroup ?? "ALL";
          const data = { value: f.value, suppressed: f.suppressed ?? false, denominator: f.denominator ?? null, sourceLoadId: id };
          await tx.miamiSchoolYearMetric.upsert({
            where: { schoolNumber_year_metricKey_subgroup: { schoolNumber: f.schoolNumber, year: f.year, metricKey: f.metricKey, subgroup } },
            update: data, create: { schoolNumber: f.schoolNumber, year: f.year, metricKey: f.metricKey, subgroup, ...data },
          });
        }
        for (const operation of load.operations) await operation(tx);
        const stale = load.hasFacts ? await tx.$executeRaw`
          DELETE FROM miami_school_year_metrics
           WHERE source_load_id IN (SELECT id FROM miami_data_loads WHERE id <> ${id} AND source_id = ${load.opts.sourceId})` : 0;
        // Timestamp precision is milliseconds in this schema. Give committed
        // generations a strict order even if two short transactions finish in
        // the same millisecond, or this transaction began while waiting for a lock.
        await tx.$executeRaw`
          UPDATE miami_data_loads
             SET loaded_at = GREATEST(clock_timestamp(),
               (SELECT MAX(loaded_at) + INTERVAL '1 millisecond' FROM miami_data_loads
                 WHERE source_id = ${load.opts.sourceId} AND id <> ${id}))
           WHERE id = ${id}`;
        return tx.miamiDataLoad.update({
          where: { id }, data: { rowsInserted: counts.inserted, rowsUpdated: counts.updated, rowsSuppressed: counts.suppressed,
            notes: [load.opts.notes, notes, load.hasFacts ? `pruned ${stale} stale cells from earlier loads` : null].filter(Boolean).join("; ") || null },
        });
      }, { maxWait: 60_000, timeout: 300_000 });
    } finally {
      pending.delete(id);
    }
  }

  return { startDataLoad, upsertFacts, stageLoadOperation, finishDataLoad };
}