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

Microshards support for PostgreSQL

import { ident, sql } from "./quote";
import { runSql } from "./runSql";

/**
 * Returns an estimate for the number of rows in the table.
 * - Method "explain" is good for the source tables (fromDsn). It relies on the
 *   ANALYZE work and is precise enough for the table which exists for a while.
 *   It's not good for the COPYing table, since EXPLAIN severely under-estimates
 *   the number of rows copied so far.
 * - Method "pg_stat_progress_copy" is good while a COPY query is running. It is
 *   very precise for the destination tables (toDsn).
 */
export async function getRowCount({
  dsn,
  schema,
  table,
  method,
}: {
  dsn: string;
  schema: string;
  table: string;
  method: "explain" | "pg_stat_progress_copy";
}): Promise<number | null> {
  if (method === "explain") {
    const firstLine = await runSql.one(
      dsn,
      sql`EXPLAIN SELECT 1 FROM ${ident(schema)}.${ident(table)}`,
    );
    return firstLine.match(/rows=(\d+)/) ? parseInt(RegExp.$1) : null;
  } else {
    const firstLine = await runSql.one(
      dsn,
      sql`
        SELECT tuples_processed
        FROM pg_stat_progress_copy
        JOIN pg_class ON pg_class.oid = relid
        JOIN pg_namespace ON pg_namespace.oid = relnamespace
        WHERE
          datname = current_database()
          AND nspname = ${schema}
          AND relname = ${table}
      `,
    );
    return firstLine.match(/^\d+$/s) ? parseInt(firstLine) : null;
  }
}