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

Microshards support for PostgreSQL

import { libSchema } from "../../internal/names";
import * as names from "../../internal/names";
import { ident, join, sql } from "../../internal/quote";
import { runSql } from "../../internal/runSql";
import { dropTestDatabase } from "./dropTestDatabase";
import { getTestDsn } from "./getTestDsn";

const now = Date.now();
const testDBs: Array<{ dsn: string; dbName: string }> = [];

// Clean the test DBs even when the tests fail. This allows to not accumulate
// too many of them (although it is still not guaranteed that this code will
// execute: the process may crash for instance).
afterAll(async () => {
  // Drop the DBs sequentially and in the reverse order of creation, assuming
  // that there may be cross-DB dependencies (e.g. subscriptions).
  for (const { dsn, dbName } of testDBs.reverse()) {
    await dropTestDatabase(dsn, dbName).catch((e) =>
      // eslint-disable-next-line no-console
      console.error("In afterAll cleanup:", e),
    );
  }
});

/**
 * Creates a brand new test database with the name equals to the current DB name
 * with "~" + JEST_WORKER_ID suffix. The database is cleaned before returning.
 */
export async function ensureTestDBExists({
  from,
  to,
  state,
  dbNameSuffix,
  skipInstall,
}: {
  from?: number;
  to?: number;
  state?: "active" | "inactive";
  dbNameSuffix?: string;
  skipInstall?: boolean;
}): Promise<string> {
  const dsn = getTestDsn();

  // Subscription and slot names are globally unique across all databases, so we
  // have to suffix them to avoid collisions in concurrent tests.
  jest
    .spyOn(names, "subName")
    .mockImplementation(
      (schema: string) => `pg_microsharding_move_${schema}_sub_test${now}`,
    );

  // DB name must be unique per the running Jest worker.
  const testDB =
    new URL(dsn).pathname.substring(1) +
    `~pg_micro_${process.env["JEST_WORKER_ID"]}` +
    (dbNameSuffix ?? "");
  const testDsnUrl = new URL(dsn);
  testDsnUrl.pathname = `/${testDB}`;
  const testDsn = testDsnUrl.toString();

  testDBs.push({ dsn, dbName: testDB });
  await dropTestDatabase(dsn, testDB);
  await runSql(dsn, sql`CREATE DATABASE ${ident(testDB)}`);

  if (!skipInstall) {
    await runSql(
      testDsn,
      join(
        [
          sql`CREATE SCHEMA IF NOT EXISTS ${libSchema()};`,
          sql`SET search_path TO ${libSchema()};`,
          sql`\\ir ${__dirname}/../../../sql/pg-microsharding-up.sql`,
          ...(from && to
            ? [
                sql`SELECT ${libSchema()}.microsharding_ensure_exist(${from}, ${to});`,
                state === "active"
                  ? sql`SELECT ${libSchema()}.microsharding_ensure_active(${from}, ${to});`
                  : sql`SELECT ${libSchema()}.microsharding_ensure_inactive(${from}, ${to});`,
              ]
            : []),
        ],
        "\n",
      ),
    );
  }

  return testDsn;
}