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

Microshards support for PostgreSQL

import { Client as PgClient } from "pg";

/**
 * A connected PG client with extra methods.
 */
export class Client extends PgClient {
  async serverVersion(): Promise<number[]> {
    const { rows } = await this.query("SHOW server_version");
    return String(rows[0]["server_version"])
      .split(".")
      .map((p) => parseInt(p));
  }

  async currentWalInsertLsn(): Promise<string> {
    const { rows } = await this.query<{ lsn: string }>(
      "SELECT pg_current_wal_insert_lsn() AS lsn",
    );
    return rows[0].lsn;
  }

  async confirmedFlushLsn(slotName: string): Promise<string> {
    const { rows } = await this.query<{ lsn: string; gap: string | number }>(
      `SELECT confirmed_flush_lsn AS lsn
         FROM pg_replication_slots
         WHERE database = current_database() AND slot_name = $1`,
      [slotName],
    );
    if (rows.length === 0) {
      throw new Error(`Slot ${slotName} not found`);
    }

    return rows[0].lsn;
  }

  async isMaster(): Promise<boolean> {
    const { rows } = await this.query<{ pg_is_in_recovery: boolean }>(
      "SELECT pg_is_in_recovery()",
    );
    return rows[0]?.pg_is_in_recovery === false;
  }
}

/**
 * Node-postgres doesn't really support sslmode=prefer: it treats it as "ssl is
 * required". This is a different default than e.g. psql has when no PGSSLMODE
 * is set. Also, it is not always able to locate the local certs, so we have to
 * use rejectUnauthorized=false. Thus, we try to set up an SSL connection first,
 * and if it fails, then we try non-SSL.
 */
export async function createConnectedClient(dsn: string): Promise<Client> {
  const dsnUrl = new URL(dsn);
  dsnUrl.searchParams.delete("sslmode");

  let sslError: unknown;
  const config = {
    connectionString: dsnUrl.toString(),
    application_name: "pg-microsharding",
  };

  const clientSsl = new Client({
    ...config,
    ssl: { rejectUnauthorized: false },
  });
  try {
    await clientSsl.connect();
    clientSsl.on("error", () => {});
    return clientSsl;
  } catch (e: unknown) {
    sslError = e;
    if (
      !`${e}`.includes("The server does not support SSL connections") &&
      !`${e}`.includes("no pg_hba.conf")
    ) {
      throw e;
    }
  }

  const clientPlain = new Client(config);
  try {
    await clientPlain.connect();
    clientPlain.on("error", () => {});
    return clientPlain;
  } catch (e: unknown) {
    throw (
      `Error connecting to ${dsn}:\n` +
      ` - with SSL: ${sslError}\n` +
      ` - without SSL: ${e}`
    );
  }
}