import { PG_UNDEFINED_TABLE, PG_UNDEFINED_COLUMN } from "@drdgvhbh/postgres-error-codes";
import postgres from "postgres";
import { equals } from "../../../../../shared/src/set-utils.js";
import { parse } from "../../../../../shared/src/valita.js";
import { READONLY } from "../../../db/mode-enum.js";
import { TsvParser } from "../../../db/pg-copy.js";
import { getTypeParsers } from "../../../db/pg-type-parser.js";
import { TransactionPool } from "../../../db/transaction-pool.js";
import { pgClient } from "../../../types/pg.js";
import { SchemaIncompatibilityError } from "../common/backfill-manager.js";
import { tableMetadataSchema, columnMetadataSchema } from "./backfill-metadata.js";
import { makeSelectPublishedStmt, createReplicationSlot } from "./initial-sync.js";
import { toStateVersionString } from "./lsn.js";
import { getPublicationInfo } from "./schema/published.js";
const POSTGRES_COPY_CHUNK_SIZE = 64 * 1024;
async function* streamBackfill(lc, upstreamURI, { slot, publications }, bf, opts = {}) {
  lc = lc.withContext("component", "backfill").withContext("table", bf.table.name);
  const { flushThresholdBytes = POSTGRES_COPY_CHUNK_SIZE } = opts;
  const db = pgClient(lc, upstreamURI, {
    connection: { ["application_name"]: "backfill-stream" }
  });
  const tx = new TransactionPool(lc, READONLY).run(db);
  try {
    const watermark = await setSnapshot(lc, upstreamURI, tx, slot);
    const { tableSpec, backfill } = await validateSchema(
      tx,
      publications,
      bf,
      watermark
    );
    const types = await getTypeParsers(db, { returnJsonAsString: true });
    const { relation, columns } = backfill;
    const cols = [...relation.rowKey.columns, ...columns];
    yield* stream(
      lc,
      tx,
      backfill,
      makeSelectPublishedStmt(tableSpec, cols),
      cols.map((col) => types.getTypeParser(tableSpec.columns[col].typeOID)),
      flushThresholdBytes
    );
  } catch (e) {
    if (e instanceof postgres.PostgresError && (e.code === PG_UNDEFINED_TABLE || e.code === PG_UNDEFINED_COLUMN)) {
      throw new SchemaIncompatibilityError(bf, String(e), { cause: e });
    }
    throw e;
  } finally {
    tx.setDone();
    void tx.done().catch(() => {
    });
    void db.end().catch((e) => lc.warn?.(`error closing backfill connection`, e));
  }
}
async function* stream(lc, tx, backfill, selectStmt, colParsers, flushThresholdBytes) {
  const start = performance.now();
  lc.info?.(`Starting backfill copy stream:`, selectStmt);
  const copyStream = await tx.processReadTask(
    (sql) => sql.unsafe(`COPY (${selectStmt}) TO STDOUT`).readable()
  );
  const tsvParser = new TsvParser();
  let totalRows = 0;
  let totalBytes = 0;
  let totalMsgs = 0;
  let rowValues = [];
  let bufferedBytes = 0;
  const logFlushed = () => {
    lc.debug?.(
      `Flushed ${rowValues.length} rows, ${bufferedBytes} bytes (total: rows=${totalRows}, msgs=${totalMsgs}, bytes=${totalBytes})`
    );
  };
  let row = Array.from({ length: colParsers.length });
  let col = 0;
  for await (const data of copyStream) {
    const chunk = data;
    for (const text of tsvParser.parse(chunk)) {
      row[col] = text === null ? null : colParsers[col](text);
      if (++col === colParsers.length) {
        rowValues.push(row);
        totalRows++;
        row = Array.from({ length: colParsers.length });
        col = 0;
      }
    }
    bufferedBytes += chunk.byteLength;
    totalBytes += chunk.byteLength;
    if (bufferedBytes >= flushThresholdBytes) {
      yield { tag: "backfill", ...backfill, rowValues };
      totalMsgs++;
      logFlushed();
      rowValues = [];
      bufferedBytes = 0;
    }
  }
  if (rowValues.length > 0) {
    yield { tag: "backfill", ...backfill, rowValues };
    totalMsgs++;
    logFlushed();
  }
  yield { tag: "backfill-completed", ...backfill };
  const elapsed = performance.now() - start;
  lc.info?.(
    `Finished streaming ${totalRows} rows, ${totalMsgs} msgs, ${totalBytes} bytes (${elapsed.toFixed(3)} ms)`
  );
}
async function setSnapshot(lc, upstreamURI, tx, slotNamePrefix) {
  const replicationSession = pgClient(lc, upstreamURI, {
    ["fetch_types"]: false,
    // Necessary for the streaming protocol
    connection: { replication: "database" }
    // https://www.postgresql.org/docs/current/protocol-replication.html
  });
  const tempSlot = `${slotNamePrefix}_bf_${Date.now()}`;
  try {
    const { snapshot_name: snapshot, consistent_point: lsn } = await createReplicationSlot(lc, replicationSession, tempSlot);
    await tx.processReadTask(
      (sql) => sql.unsafe(`SET TRANSACTION SNAPSHOT '${snapshot}'`)
    );
    await replicationSession.unsafe(`DROP_REPLICATION_SLOT "${tempSlot}"`);
    const watermark = toStateVersionString(lsn);
    lc.info?.(`Opened snapshot transaction at LSN ${lsn} (${watermark})`);
    return watermark;
  } catch (e) {
    await replicationSession.unsafe(
      /*sql*/
      `SELECT pg_drop_replication_slot(slot_name) FROM pg_replication_slots
         WHERE slot_name = '${tempSlot}'`
    );
    lc.error?.(`Failed to create backfill snapshot`, e);
    throw e;
  } finally {
    await replicationSession.end();
  }
}
function validateSchema(tx, publications, bf, watermark) {
  return tx.processReadTask(async (sql) => {
    const { tables } = await getPublicationInfo(sql, publications);
    const spec = tables.find(
      (spec2) => spec2.schema === bf.table.schema && spec2.name === bf.table.name
    );
    if (!spec) {
      throw new SchemaIncompatibilityError(
        bf,
        `Table has been renamed or dropped`
      );
    }
    const tableMeta = parse(bf.table.metadata, tableMetadataSchema);
    if (spec.schemaOID !== tableMeta.schemaOID) {
      throw new SchemaIncompatibilityError(
        bf,
        `Schema no longer corresponds to the original schema`
      );
    }
    if (spec.oid !== tableMeta.relationOID) {
      throw new SchemaIncompatibilityError(
        bf,
        `Table no longer corresponds to the original table`
      );
    }
    if (!equals(
      new Set(Object.keys(tableMeta.rowKey)),
      new Set(spec.replicaIdentityColumns)
    )) {
      throw new SchemaIncompatibilityError(
        bf,
        "Row key (e.g. PRIMARY KEY or INDEX) has changed"
      );
    }
    const allCols = [
      ...Object.entries(tableMeta.rowKey),
      ...Object.entries(bf.columns)
    ];
    for (const [col, val] of allCols) {
      const colSpec = spec.columns[col];
      if (!colSpec) {
        throw new SchemaIncompatibilityError(
          bf,
          `Column ${col} has been renamed or dropped`
        );
      }
      const colMeta = parse(val, columnMetadataSchema);
      if (colMeta.attNum !== colSpec.pos) {
        throw new SchemaIncompatibilityError(
          bf,
          `Column ${col} no longer corresponds to the original column`
        );
      }
    }
    const backfill = {
      relation: {
        schema: bf.table.schema,
        name: bf.table.name,
        rowKey: { columns: Object.keys(tableMeta.rowKey) }
      },
      columns: Object.keys(bf.columns).filter(
        (col) => !(col in tableMeta.rowKey)
      ),
      watermark
    };
    return { tableSpec: spec, backfill };
  });
}
export {
  streamBackfill
};
//# sourceMappingURL=backfill-stream.js.map
