@clickup/pg-microsharding/src/api/cleanup.ts

Microshards support for PostgreSQL

import { getPublisherDsn } from "../internal/getPublisherDsn";
import { getSchemaName } from "../internal/getSchemaName";
import { listActiveSchemas } from "../internal/listActiveSchemas";
import { schemaCleanupRe } from "../internal/names";
import { ident, sql } from "../internal/quote";
import { resultAbort } from "../internal/resultAbort";
import { runSql } from "../internal/runSql";

/**
 * Removes old and semi-migrated schemas, and optionally orphan microshard
 * schemas (those matching the microshard pattern but absent from the active
 * shards list).
 */
export async function cleanup({
  dsn,
  noOldShards,
  confirm,
  noOrphanShards,
  confirmOrphanShards,
}: {
  dsn: string;
  noOldShards?: (oldSchemaNameRe: string) => Promise<void>;
  confirm?: (schemas: string[]) => Promise<boolean>;
  noOrphanShards?: () => Promise<void>;
  confirmOrphanShards?: (schemas: string[]) => Promise<boolean>;
}): Promise<void> {
  const dummyNo = 101;
  const dummySchemaName = await getSchemaName(dsn, dummyNo);
  const schemaNameRe = dummySchemaName!.replace(
    new RegExp(`0*${dummyNo}0*`),
    "\\d+",
  );
  const oldSchemaNameRe = schemaCleanupRe(schemaNameRe);

  const allSchemas = await runSql.column(
    dsn,
    sql`SELECT nspname FROM pg_namespace`,
  );
  const oldSchemas = allSchemas.filter((schema) =>
    schema.match(oldSchemaNameRe),
  );
  if (oldSchemas.length === 0) {
    await noOldShards?.(oldSchemaNameRe);
  } else if (!confirm || (await confirm(oldSchemas))) {
    // Remove sequentially, otherwise there is a high chance of having
    // shared_buffers OOM.
    for (const schema of oldSchemas) {
      await runSql(
        dsn,
        sql`DROP SCHEMA ${ident(schema)} CASCADE`,
        `Dropping redundant schema ${schema}`,
      );
    }
  }

  if (noOrphanShards || confirmOrphanShards) {
    const activeSet = new Set(await listActiveSchemas({ dsn }));
    const orphanSchemas = allSchemas.filter(
      (schema) =>
        schema.match(new RegExp(`^${schemaNameRe}$`)) && !activeSet.has(schema),
    );
    if (orphanSchemas.length === 0) {
      await noOrphanShards?.();
    } else if (
      !confirmOrphanShards ||
      (await confirmOrphanShards(orphanSchemas))
    ) {
      for (const schema of orphanSchemas) {
        // Discover the publisher's connection string from the subscription, so
        // we can clean up both ends of any lingering replication. Falls back to
        // the current DSN if no subscription exists (e.g. leftover source shard).
        const fromDsn =
          (await getPublisherDsn({ subscriberDsn: dsn, schema })) || dsn;
        await resultAbort({ fromDsn, toDsn: dsn, schema });
      }
    }
  }
}