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

Microshards support for PostgreSQL

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

/**
 * Removes old and semi-migrated schemas.
 */
export async function cleanup({
  dsn,
  noOldShards,
  confirm,
}: {
  dsn: string;
  noOldShards?: (oldSchemaNameRe: string) => Promise<void>;
  confirm?: (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);
    return;
  }

  if (confirm && !(await confirm(oldSchemas))) {
    return;
  }

  // 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}`,
    );
  }
}