import { SqliteError } from '@rocicorp/zero-sqlite3';
import { AbortError } from "../../../../shared/src/abort-error.js";
import { assert, unreachable } from "../../../../shared/src/asserts.js";
import { stringify } from "../../../../shared/src/bigint-json.js";
import { must } from "../../../../shared/src/must.js";
import { createLiteIndexStatement, createLiteTableStatement, liteColumnDef, } from "../../db/create.js";
import { computeZqlSpecs, listIndexes, listTables, } from "../../db/lite-tables.js";
import { mapPostgresToLite, mapPostgresToLiteColumn, mapPostgresToLiteIndex, } from "../../db/pg-to-lite.js";
import { JSON_PARSED, liteRow, } from "../../types/lite.js";
import { liteTableName } from "../../types/names.js";
import { id } from "../../types/sql.js";
import { logDeleteOp, logResetOp, logSetOp, logTruncateOp, } from "./schema/change-log.js";
import { ZERO_VERSION_COLUMN_NAME, updateReplicationWatermark, } from "./schema/replication-state.js";
/**
 * The ChangeProcessor partitions the stream of messages into transactions
 * by creating a {@link TransactionProcessor} when a transaction begins, and dispatching
 * messages to it until the commit is received.
 *
 * From https://www.postgresql.org/docs/current/protocol-logical-replication.html#PROTOCOL-LOGICAL-MESSAGES-FLOW :
 *
 * "The logical replication protocol sends individual transactions one by one.
 *  This means that all messages between a pair of Begin and Commit messages
 *  belong to the same transaction."
 */
export class ChangeProcessor {
    #db;
    #mode;
    #failService;
    // The TransactionProcessor lazily loads table specs into this Map,
    // and reloads them after a schema change. It is cached here to avoid
    // reading them from the DB on every transaction.
    #tableSpecs = new Map();
    #currentTx = null;
    #failure;
    constructor(db, mode, failService) {
        this.#db = db;
        this.#mode = mode;
        this.#failService = failService;
    }
    #fail(lc, err) {
        if (!this.#failure) {
            this.#currentTx?.abort(lc); // roll back any pending transaction.
            this.#failure = ensureError(err);
            if (!(err instanceof AbortError)) {
                // Propagate the failure up to the service.
                lc.error?.('Message Processing failed:', this.#failure);
                this.#failService(lc, this.#failure);
            }
        }
    }
    abort(lc) {
        this.#fail(lc, new AbortError());
    }
    /** @return If a transaction was committed. */
    processMessage(lc, downstream) {
        const [type, message] = downstream;
        if (this.#failure) {
            lc.debug?.(`Dropping ${message.tag}`);
            return null;
        }
        try {
            const watermark = type === 'begin'
                ? downstream[2].commitWatermark
                : type === 'commit'
                    ? downstream[2].watermark
                    : undefined;
            return this.#processMessage(lc, message, watermark);
        }
        catch (e) {
            this.#fail(lc, e);
        }
        return null;
    }
    #beginTransaction(lc, commitVersion, jsonFormat) {
        const start = Date.now();
        // litestream can technically hold the lock for an arbitrary amount of time
        // when checkpointing a large commit. Crashing on the busy-timeout in this
        // scenario will either produce a corrupt backup or otherwise prevent
        // replication from proceeding.
        //
        // Instead, retry the lock acquisition indefinitely. If this masks
        // an unknown deadlock situation, manual intervention will be necessary.
        for (let i = 0;; i++) {
            try {
                return new TransactionProcessor(lc, this.#db, this.#mode, this.#tableSpecs, commitVersion, jsonFormat);
            }
            catch (e) {
                if (e instanceof SqliteError && e.code === 'SQLITE_BUSY') {
                    lc.warn?.(`SQLITE_BUSY for ${Date.now() - start} ms (attempt ${i + 1}). ` +
                        `This is only expected if litestream is performing a large ` +
                        `checkpoint.`, e);
                    continue;
                }
                throw e;
            }
        }
    }
    /** @return If a transaction was committed. */
    #processMessage(lc, msg, watermark) {
        if (msg.tag === 'begin') {
            if (this.#currentTx) {
                throw new Error(`Already in a transaction ${stringify(msg)}`);
            }
            this.#currentTx = this.#beginTransaction(lc, must(watermark), msg.json ?? JSON_PARSED);
            return null;
        }
        // For non-begin messages, there should be a #currentTx set.
        const tx = this.#currentTx;
        if (!tx) {
            throw new Error(`Received message outside of transaction: ${stringify(msg)}`);
        }
        if (msg.tag === 'commit') {
            // Undef this.#currentTx to allow the assembly of the next transaction.
            this.#currentTx = null;
            assert(watermark);
            const schemaUpdated = tx.processCommit(msg, watermark);
            return { watermark, schemaUpdated };
        }
        if (msg.tag === 'rollback') {
            this.#currentTx?.abort(lc);
            this.#currentTx = null;
            return null;
        }
        switch (msg.tag) {
            case 'insert':
                tx.processInsert(msg);
                break;
            case 'update':
                tx.processUpdate(msg);
                break;
            case 'delete':
                tx.processDelete(msg);
                break;
            case 'truncate':
                tx.processTruncate(msg);
                break;
            case 'create-table':
                tx.processCreateTable(msg);
                break;
            case 'rename-table':
                tx.processRenameTable(msg);
                break;
            case 'add-column':
                tx.processAddColumn(msg);
                break;
            case 'update-column':
                tx.processUpdateColumn(msg);
                break;
            case 'drop-column':
                tx.processDropColumn(msg);
                break;
            case 'drop-table':
                tx.processDropTable(msg);
                break;
            case 'create-index':
                tx.processCreateIndex(msg);
                break;
            case 'drop-index':
                tx.processDropIndex(msg);
                break;
            default:
                unreachable(msg);
        }
        return null;
    }
}
/**
 * The {@link TransactionProcessor} handles the sequence of messages from
 * upstream, from `BEGIN` to `COMMIT` and executes the corresponding mutations
 * on the {@link postgres.TransactionSql} on the replica.
 *
 * When applying row contents to the replica, the `_0_version` column is added / updated,
 * and a corresponding entry in the `ChangeLog` is added. The version value is derived
 * from the watermark of the preceding transaction (stored as the `nextStateVersion` in the
 * `ReplicationState` table).
 *
 *   Side note: For non-streaming Postgres transactions, the commitEndLsn (and thus
 *   commit watermark) is available in the `begin` message, so it could theoretically
 *   be used for the row version of changes within the transaction. However, the
 *   commitEndLsn is not available in the streaming (in-progress) transaction
 *   protocol, and may not be available for CDC streams of other upstream types.
 *   Therefore, the zero replication protocol is designed to not require the commit
 *   watermark when a transaction begins.
 *
 * Also of interest is the fact that all INSERT Messages are logically applied as
 * UPSERTs. See {@link processInsert} for the underlying motivation.
 */
class TransactionProcessor {
    #lc;
    #startMs;
    #db;
    #mode;
    #version;
    #tableSpecs;
    #jsonFormat;
    #schemaChanged = false;
    constructor(lc, db, mode, tableSpecs, commitVersion, jsonFormat) {
        this.#startMs = Date.now();
        this.#mode = mode;
        this.#jsonFormat = jsonFormat;
        switch (mode) {
            case 'serving':
                // Although the Replicator / Incremental Syncer is the only writer of the replica,
                // a `BEGIN CONCURRENT` transaction is used to allow View Syncers to simulate
                // (i.e. and `ROLLBACK`) changes on historic snapshots of the database for the
                // purpose of IVM).
                //
                // This TransactionProcessor is the only logic that will actually
                // `COMMIT` any transactions to the replica.
                db.beginConcurrent();
                break;
            case 'backup':
                // For the backup-replicator (i.e. replication-manager), there are no View Syncers
                // and thus BEGIN CONCURRENT is not necessary. In fact, BEGIN CONCURRENT can cause
                // deadlocks with forced wal-checkpoints (which `litestream replicate` performs),
                // so it is important to use vanilla transactions in this configuration.
                db.beginImmediate();
                break;
            case 'initial-sync':
                // When the ChangeProcessor is used for initial-sync, the calling code
                // handles the transaction boundaries.
                break;
            default:
                unreachable();
        }
        this.#db = db;
        this.#version = commitVersion;
        this.#lc = lc.withContext('version', commitVersion);
        this.#tableSpecs = tableSpecs;
        if (this.#tableSpecs.size === 0) {
            this.#reloadTableSpecs();
        }
    }
    #reloadTableSpecs() {
        this.#tableSpecs.clear();
        // zqlSpecs include the primary key derived from unique indexes
        const zqlSpecs = computeZqlSpecs(this.#lc, this.#db.db);
        for (let spec of listTables(this.#db.db)) {
            if (!spec.primaryKey) {
                spec = {
                    ...spec,
                    primaryKey: [
                        ...(zqlSpecs.get(spec.name)?.tableSpec.primaryKey ?? []),
                    ],
                };
            }
            this.#tableSpecs.set(spec.name, spec);
        }
    }
    #tableSpec(name) {
        return must(this.#tableSpecs.get(name), `Unknown table ${name}`);
    }
    #getKey({ row, numCols }, { relation }) {
        const keyColumns = relation.replicaIdentity !== 'full'
            ? relation.keyColumns // already a suitable key
            : this.#tableSpec(liteTableName(relation)).primaryKey;
        if (!keyColumns?.length) {
            throw new Error(`Cannot replicate table "${relation.name}" without a PRIMARY KEY or UNIQUE INDEX`);
        }
        // For the common case (replica identity default), the row is already the
        // key for deletes and updates, in which case a new object can be avoided.
        if (numCols === keyColumns.length) {
            return row;
        }
        const key = {};
        for (const col of keyColumns) {
            key[col] = row[col];
        }
        return key;
    }
    processInsert(insert) {
        const table = liteTableName(insert.relation);
        const newRow = liteRow(insert.new, this.#tableSpec(table), this.#jsonFormat);
        this.#upsert(table, {
            ...newRow.row,
            [ZERO_VERSION_COLUMN_NAME]: this.#version,
        });
        if (insert.relation.keyColumns.length === 0) {
            // INSERTs can be replicated for rows without a PRIMARY KEY or a
            // UNIQUE INDEX. These are written to the replica but not recorded
            // in the changeLog, because these rows cannot participate in IVM.
            //
            // (Once the table schema has been corrected to include a key, the
            //  associated schema change will reset pipelines and data can be
            //  loaded via hydration.)
            return;
        }
        const key = this.#getKey(newRow, insert);
        this.#logSetOp(table, key);
    }
    #upsert(table, row) {
        const columns = Object.keys(row).map(c => id(c));
        this.#db.run(`
      INSERT OR REPLACE INTO ${id(table)} (${columns.join(',')})
        VALUES (${Array.from({ length: columns.length }).fill('?').join(',')})
      `, Object.values(row));
    }
    // Updates by default are applied as UPDATE commands to support partial
    // row specifications from the change source. In particular, this is needed
    // to handle updates for which unchanged TOASTed values are not sent:
    //
    // https://www.postgresql.org/docs/current/protocol-logicalrep-message-formats.html#PROTOCOL-LOGICALREP-MESSAGE-FORMATS-TUPLEDATA
    //
    // However, in certain cases an UPDATE may be received for a row that
    // was not initially synced, such as when:
    // (1) an existing table is added to the app's publication, or
    // (2) a new sharding key is added to a shard during resharding.
    //
    // In order to facilitate "resumptive" replication, the logic falls back to
    // an INSERT if the update did not change any rows.
    // TODO: Figure out a solution for resumptive replication of rows
    //       with TOASTed values.
    processUpdate(update) {
        const table = liteTableName(update.relation);
        const newRow = liteRow(update.new, this.#tableSpec(table), this.#jsonFormat);
        const row = { ...newRow.row, [ZERO_VERSION_COLUMN_NAME]: this.#version };
        // update.key is set with the old values if the key has changed.
        const oldKey = update.key
            ? this.#getKey(liteRow(update.key, this.#tableSpec(table), this.#jsonFormat), update)
            : null;
        const newKey = this.#getKey(newRow, update);
        if (oldKey) {
            this.#logDeleteOp(table, oldKey);
        }
        this.#logSetOp(table, newKey);
        const currKey = oldKey ?? newKey;
        const conds = Object.keys(currKey).map(col => `${id(col)}=?`);
        const setExprs = Object.keys(row).map(col => `${id(col)}=?`);
        const { changes } = this.#db.run(`
      UPDATE ${id(table)}
        SET ${setExprs.join(',')}
        WHERE ${conds.join(' AND ')}
      `, [...Object.values(row), ...Object.values(currKey)]);
        // If the UPDATE did not affect any rows, perform an UPSERT of the
        // new row for resumptive replication.
        if (changes === 0) {
            this.#upsert(table, row);
        }
    }
    processDelete(del) {
        const table = liteTableName(del.relation);
        const rowKey = this.#getKey(liteRow(del.key, this.#tableSpec(table), this.#jsonFormat), del);
        this.#delete(table, rowKey);
        if (this.#mode === 'serving') {
            this.#logDeleteOp(table, rowKey);
        }
    }
    #delete(table, rowKey) {
        const conds = Object.keys(rowKey).map(col => `${id(col)}=?`);
        this.#db.run(`DELETE FROM ${id(table)} WHERE ${conds.join(' AND ')}`, Object.values(rowKey));
    }
    processTruncate(truncate) {
        for (const relation of truncate.relations) {
            const table = liteTableName(relation);
            // Update replica data.
            this.#db.run(`DELETE FROM ${id(table)}`);
            // Update change log.
            this.#logTruncateOp(table);
        }
    }
    processCreateTable(create) {
        const table = mapPostgresToLite(create.spec);
        this.#db.db.exec(createLiteTableStatement(table));
        this.#logResetOp(table.name);
        this.#lc.info?.(create.tag, table.name);
    }
    processRenameTable(rename) {
        const oldName = liteTableName(rename.old);
        const newName = liteTableName(rename.new);
        this.#db.db.exec(`ALTER TABLE ${id(oldName)} RENAME TO ${id(newName)}`);
        this.#bumpVersions(newName);
        this.#logResetOp(oldName);
        this.#lc.info?.(rename.tag, oldName, newName);
    }
    processAddColumn(msg) {
        const table = liteTableName(msg.table);
        const { name } = msg.column;
        const spec = mapPostgresToLiteColumn(table, msg.column);
        this.#db.db.exec(`ALTER TABLE ${id(table)} ADD ${id(name)} ${liteColumnDef(spec)}`);
        this.#bumpVersions(table);
        this.#lc.info?.(msg.tag, table, msg.column);
    }
    processUpdateColumn(msg) {
        const table = liteTableName(msg.table);
        let oldName = msg.old.name;
        const newName = msg.new.name;
        // update-column can ignore defaults because it does not change the values
        // in existing rows.
        //
        // https://www.postgresql.org/docs/current/sql-altertable.html#SQL-ALTERTABLE-DESC-SET-DROP-DEFAULT
        //
        // "The new default value will only apply in subsequent INSERT or UPDATE
        //  commands; it does not cause rows already in the table to change."
        //
        // This allows support for _changing_ column defaults to any expression,
        // since it does not affect what the replica needs to do.
        const oldSpec = mapPostgresToLiteColumn(table, msg.old, 'ignore-default');
        const newSpec = mapPostgresToLiteColumn(table, msg.new, 'ignore-default');
        // The only updates that are relevant are the column name and the data type.
        if (oldName === newName && oldSpec.dataType === newSpec.dataType) {
            this.#lc.info?.(msg.tag, 'no thing to update', oldSpec, newSpec);
            return;
        }
        // If the data type changes, we have to make a new column with the new data type
        // and copy the values over.
        if (oldSpec.dataType !== newSpec.dataType) {
            // Remember (and drop) the indexes that reference the column.
            const indexes = listIndexes(this.#db.db).filter(idx => idx.tableName === table && oldName in idx.columns);
            const stmts = indexes.map(idx => `DROP INDEX IF EXISTS ${id(idx.name)};`);
            const tmpName = `tmp.${newName}`;
            stmts.push(`
        ALTER TABLE ${id(table)} ADD ${id(tmpName)} ${liteColumnDef(newSpec)};
        UPDATE ${id(table)} SET ${id(tmpName)} = ${id(oldName)};
        ALTER TABLE ${id(table)} DROP ${id(oldName)};
        `);
            for (const idx of indexes) {
                // Re-create the indexes to reference the new column.
                idx.columns[tmpName] = idx.columns[oldName];
                delete idx.columns[oldName];
                stmts.push(createLiteIndexStatement(idx));
            }
            this.#db.db.exec(stmts.join(''));
            oldName = tmpName;
        }
        if (oldName !== newName) {
            this.#db.db.exec(`ALTER TABLE ${id(table)} RENAME ${id(oldName)} TO ${id(newName)}`);
        }
        this.#bumpVersions(table);
        this.#lc.info?.(msg.tag, table, msg.new);
    }
    processDropColumn(msg) {
        const table = liteTableName(msg.table);
        const { column } = msg;
        this.#db.db.exec(`ALTER TABLE ${id(table)} DROP ${id(column)}`);
        this.#bumpVersions(table);
        this.#lc.info?.(msg.tag, table, column);
    }
    processDropTable(drop) {
        const name = liteTableName(drop.id);
        this.#db.db.exec(`DROP TABLE IF EXISTS ${id(name)}`);
        this.#logResetOp(name);
        this.#lc.info?.(drop.tag, name);
    }
    processCreateIndex(create) {
        const index = mapPostgresToLiteIndex(create.spec);
        this.#db.db.exec(createLiteIndexStatement(index));
        // indexes affect tables visibility (e.g. sync-ability is gated on
        // having a unique index), so reset pipelines to refresh table schemas.
        this.#logResetOp(index.tableName);
        this.#lc.info?.(create.tag, index.name);
    }
    processDropIndex(drop) {
        const name = liteTableName(drop.id);
        this.#db.db.exec(`DROP INDEX IF EXISTS ${id(name)}`);
        this.#lc.info?.(drop.tag, name);
    }
    #bumpVersions(table) {
        this.#db.run(`UPDATE ${id(table)} SET ${id(ZERO_VERSION_COLUMN_NAME)} = ?`, this.#version);
        this.#logResetOp(table);
    }
    #logSetOp(table, key) {
        if (this.#mode === 'serving') {
            logSetOp(this.#db, this.#version, table, key);
        }
    }
    #logDeleteOp(table, key) {
        if (this.#mode === 'serving') {
            logDeleteOp(this.#db, this.#version, table, key);
        }
    }
    #logTruncateOp(table) {
        if (this.#mode === 'serving') {
            logTruncateOp(this.#db, this.#version, table);
        }
    }
    #logResetOp(table) {
        this.#schemaChanged = true;
        if (this.#mode === 'serving') {
            logResetOp(this.#db, this.#version, table);
        }
        this.#reloadTableSpecs();
    }
    /** @returns `true` if the schema was updated. */
    processCommit(commit, watermark) {
        if (watermark !== this.#version) {
            throw new Error(`'commit' version ${watermark} does not match 'begin' version ${this.#version}: ${stringify(commit)}`);
        }
        updateReplicationWatermark(this.#db, watermark);
        if (this.#schemaChanged) {
            const start = Date.now();
            this.#db.db.pragma('optimize');
            this.#lc.info?.(`PRAGMA optimized after schema change (${Date.now() - start} ms)`);
        }
        if (this.#mode !== 'initial-sync') {
            this.#db.commit();
        }
        const elapsedMs = Date.now() - this.#startMs;
        this.#lc.debug?.(`Committed tx@${this.#version} (${elapsedMs} ms)`);
        return this.#schemaChanged;
    }
    abort(lc) {
        lc.info?.(`aborting transaction ${this.#version}`);
        this.#db.rollback();
    }
}
function ensureError(err) {
    if (err instanceof Error) {
        return err;
    }
    const error = new Error();
    error.cause = err;
    return error;
}
//# sourceMappingURL=change-processor.js.map