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

Microshards support for PostgreSQL

import partition from "lodash/partition";
import { sql } from "./quote";
import { runSql } from "./runSql";

export interface TableInfo {
  /** Raw schema name (unquoted). */
  schema: string;
  /** Raw table name (unquoted). */
  name: string;
  /** Raw owner name (unquoted). */
  owner: string;
  /** Whether this is a base table for a partitioned table. Such tables have no
   * data and must not be mentioned in publications. */
  isPartitioned: boolean;
  /** Existing index that can be used on the subscriber end to identify the rows
   * for received updates and deletes. */
  identity: {
    /** Index name. */
    indexName: string;
    /** Whether the index is a replica identity. */
    isReplIdent: boolean;
    /** Whether the index is a primary key. */
    isPrimaryKey: boolean;
    /** Index definition. */
    def: string;
  } | null;
}

/**
 * Returns all tables in a given schema on fromDsn.
 *
 * For logical subscription purposes, we only need the `tables` that may
 * potentially contain data (e.g. regular tables, or partitions). But we also
 * return the base tables for partitioned tables in `partitionedTables`, since
 * we validate the primary key coverage in `validatePrimaryKeysFromDump()` for sanity
 * checking purposes.
 */
export async function getTablesInSchema({
  fromDsn,
  schema,
}: {
  fromDsn: string;
  schema: string;
}): Promise<{ tables: TableInfo[]; partitionedTables: TableInfo[] }> {
  const rows = await runSql(
    fromDsn,
    sql`
      SELECT DISTINCT ON (pg_namespace.nspname, pg_class.relname, pg_roles.rolname)
        pg_namespace.nspname,
        pg_class.relname,
        pg_roles.rolname,
        pg_class.relkind,
        idx_class.relname,
        pg_index.indisreplident,
        pg_index.indisprimary,
        pg_get_indexdef(idx_class.oid)
      FROM pg_class
      JOIN pg_namespace ON relnamespace = pg_namespace.oid
      JOIN pg_roles ON relowner = pg_roles.oid
      LEFT JOIN pg_index ON indrelid = pg_class.oid AND (indisreplident OR indisprimary)
      LEFT JOIN pg_class idx_class ON idx_class.oid = pg_index.indexrelid
      WHERE pg_class.relkind IN ('r', 'p') AND nspname = ${schema}
      ORDER BY
        pg_namespace.nspname,
        pg_class.relname,
        pg_roles.rolname,
        pg_index.indisreplident DESC,
        pg_index.indisprimary DESC
    `,
  );
  const res = rows.map((row) => ({
    schema: row[0],
    name: row[1],
    owner: row[2],
    isPartitioned: row[3] === "p",
    identity: row[4]
      ? {
          indexName: row[4],
          isReplIdent: row[5] === "t",
          isPrimaryKey: row[6] === "t",
          def: row[7],
        }
      : null,
  }));
  const [partitionedTables, tables] = partition(
    res,
    (table) => table.isPartitioned,
  );
  return { tables, partitionedTables };
}