@clickup/pg-microsharding/src/__tests__/internal/dropTestDatabase.ts

Microshards support for PostgreSQL

import { promiseAllMap } from "../../internal/promiseAllMap";
import { ident, sql } from "../../internal/quote";
import { runSql } from "../../internal/runSql";

/**
 * "One does not simply drop a PostgreSQL test database..."
 * https://www.postgresql.org/docs/current/sql-dropdatabase.html
 */
export async function dropTestDatabase(
  dsn: string,
  dbName: string,
): Promise<void> {
  const testDsnUrl = new URL(dsn);
  testDsnUrl.pathname = `/${dbName}`;
  const testDsn = testDsnUrl.toString();

  const preparedXacts = await runSql.column(
    dsn,
    sql`SELECT gid FROM pg_prepared_xacts WHERE database = ${dbName}`,
  );
  await promiseAllMap(preparedXacts, async (gid) => {
    await runSql(testDsn, sql`ROLLBACK PREPARED ${ident(gid)}`);
  });

  const subNames = await runSql.column(
    dsn,
    sql`
      SELECT subname
      FROM pg_subscription
      JOIN pg_database ON pg_database.oid = subdbid
      WHERE datname = ${dbName}
    `,
  );
  await promiseAllMap(subNames, async (subName) => {
    await runSql(
      testDsn,
      sql`
        SET statement_timeout TO '20s';
        ALTER SUBSCRIPTION ${ident(subName)} DISABLE;
        ALTER SUBSCRIPTION ${ident(subName)} SET (slot_name = NONE);
        DROP SUBSCRIPTION ${ident(subName)} CASCADE;
      `,
    );
  });

  const slotNames = await runSql.column(
    dsn,
    sql`SELECT slot_name FROM pg_replication_slots WHERE database = ${dbName}`,
  );
  await promiseAllMap(slotNames, async (slotName) => {
    await runSql(
      testDsn,
      sql`
        SELECT pg_terminate_backend(active_pid)
          FROM pg_replication_slots
          WHERE slot_name = ${slotName} AND active_pid IS NOT NULL;
        SELECT pg_drop_replication_slot(${slotName});
      `,
    );
  });

  await runSql(dsn, sql`DROP DATABASE IF EXISTS ${ident(dbName)} WITH (FORCE)`);
}