import { DrizzleClient } from "./DrizzleClient";
import type { Transaction } from "./types";
import { isShuttingDown, incrementTransactions, decrementTransactions } from "./shutdown";

export const withTransaction = async <T>(
    callback: (tx: Transaction) => Promise<T>,
    tx?: Transaction
): Promise<T> => {
    if (tx) {
        return await callback(tx);
    } else {
        if (isShuttingDown()) {
            throw new Error("System is shutting down, no new transactions allowed.");
        }

        incrementTransactions();
        try {
            return await DrizzleClient.transaction(async (newTx) => {
                return await callback(newTx);
            });
        } finally {
            decrementTransactions();
        }
    }
};