@clickup/pg-microsharding/src/internal/cleanUpPubSub.ts

Microshards support for PostgreSQL

import { pubName, subName } from "./names";
import { ident, sql } from "./quote";
import { runSql } from "./runSql";

/**
 * Removes all traces of logical replication from the schema.
 */
export async function cleanUpPubSub({
  fromDsn,
  toDsn,
  schema,
}: {
  fromDsn: string;
  toDsn: string;
  schema: string;
}): Promise<void> {
  const subExists = await runSql.one(
    toDsn,
    sql`
      SELECT 1
      FROM pg_subscription
      JOIN pg_database ON pg_database.oid = subdbid AND datname = current_database()
      WHERE subname = ${subName(schema)}
    `,
  );
  if (subExists) {
    await runSql(
      toDsn,
      sql`ALTER SUBSCRIPTION ${ident(subName(schema))} DISABLE`,
      "Disabling the destination subscription",
    );
    await runSql(
      toDsn,
      sql`ALTER SUBSCRIPTION ${ident(subName(schema))} SET (slot_name = NONE)`,
      "Disjoining destination subscription from the replication slot",
    );
    await runSql(
      toDsn,
      sql`DROP SUBSCRIPTION ${ident(subName(schema))} CASCADE`,
      "Dropping destination subscription",
    );
  }

  await runSql(
    fromDsn,
    sql`
      SELECT pg_terminate_backend(active_pid)
        FROM pg_replication_slots
        WHERE database = current_database() AND slot_name = ${subName(schema)} AND active_pid IS NOT NULL;
      SELECT pg_drop_replication_slot(slot_name)
        FROM pg_replication_slots
        WHERE database = current_database() AND slot_name = ${subName(schema)}
    `,
  );

  const pubExists = await runSql.one(
    fromDsn,
    sql`SELECT 1 FROM pg_publication WHERE pubname = ${pubName(schema)}`,
  );
  if (pubExists) {
    await runSql(
      fromDsn,
      sql`DROP PUBLICATION ${ident(pubName(schema))}`,
      "Dropping source publication",
    );
  }
}