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

Microshards support for PostgreSQL

import { spawnSync } from "child_process";
import pickBy from "lodash/pickBy";
import { discoverShards } from "../internal/discoverShards";
import { log } from "../internal/logging";
import { libSchema } from "../internal/names";
import { normalizeDsn } from "../internal/normalizeDsn";
import { sql } from "../internal/quote";
import { runSql } from "../internal/runSql";

/**
 * Ensures that all shards in the range exist on the DSN, then runs a shell
 * script (presumably DB migration), and then optionally activates the shards.
 */
export async function allocate({
  dsns,
  from,
  to,
  migrateCmd,
  activate,
}: {
  dsns: string[];
  from: number;
  to: number;
  migrateCmd: string;
  activate: boolean;
}): Promise<void> {
  const [dsn, ...otherDsns] = dsns;

  const existingShards = await discoverShards({ dsns: otherDsns });
  const errors: string[] = [];
  for (let shard = from; shard <= to; shard++) {
    const existingShard = existingShards.find((s) => s.shard === shard);
    if (
      existingShard &&
      normalizeDsn(existingShard.dsn) !== normalizeDsn(dsn)
    ) {
      errors.push(
        `- microshard ${shard} already exists on ${existingShard.dsn}`,
      );
    }
  }

  if (errors.length > 0) {
    throw "Errors encountered:\n" + errors.join("\n");
  }

  await runSql(
    dsn,
    sql`SELECT ${libSchema()}.microsharding_ensure_exist(${from}, ${to})`,
    `Ensuring that microshards ${from}-${to} exist on ${dsn}`,
  );

  log(
    `Running DB migration shell command "${migrateCmd}"` +
      (activate ? " (Microshards will only be activated if it succeeds.)" : ""),
  );
  const dsnUrl = new URL(dsn);
  const { status } = spawnSync(migrateCmd, {
    shell: true,
    stdio: "inherit",
    env: {
      ...process.env,
      ...pickBy(
        {
          PGUSER: decodeURIComponent(dsnUrl.username),
          PGPASSWORD: decodeURIComponent(dsnUrl.password),
          PGHOST: dsnUrl.hostname,
          PGPORT: dsnUrl.port,
          PGDATABASE: decodeURIComponent(dsnUrl.pathname.slice(1)),
          PGSSLMODE: new URLSearchParams(dsnUrl.search).get("sslmode") || "",
        },
        (v) => v,
      ),
    },
  });
  if (status !== 0) {
    throw "Error detected. NEW MICROSHARDS WERE NOT ACTIVATED!";
  }

  if (activate) {
    await runSql(
      dsn,
      sql`SELECT ${libSchema()}.microsharding_ensure_active(${from}, ${to})`,
      `Ensuring that all microshards in the range ${from}-${to} are active`,
    );
  }
}