v2.22.3: Fix partitioned tables support if they exist in microshards

Author: dimikot

package-lock.json

@@ -1,7 +1,7 @@
 {
   "name": "@clickup/pg-microsharding",
   "description": "Microshards support for PostgreSQL",
-  "version": "2.22.2",
+  "version": "2.22.3",
   "license": "MIT",
   "keywords": [
     "postgresql",

package.json

@@ -30,9 +30,28 @@ beforeEach(async () => {
       await runSql(
         fromDsn,
         sql`
-        CREATE TABLE sh0001.tbl(id bigserial PRIMARY KEY, s text);
-        INSERT INTO sh0001.tbl(s) SELECT s::text FROM generate_series(1, 1000) AS s;
-      `,
+          CREATE TABLE sh0001.tbl(id bigserial PRIMARY KEY, s text);
+          INSERT INTO sh0001.tbl(s)
+            SELECT s::text FROM generate_series(1, 1000) AS s;
+
+          CREATE TABLE sh0001."weirdly""named;tbl"(id bigserial PRIMARY KEY, s text);
+          INSERT INTO sh0001."weirdly""named;tbl"(s)
+            SELECT s::text FROM generate_series(1, 1000) AS s;
+
+          CREATE TABLE sh0001.partitioned_tbl(
+            grp bigint,
+            idx numeric(64, 32),
+            PRIMARY KEY (grp, idx)
+          ) PARTITION BY HASH (grp);
+          CREATE TABLE sh0001.partitioned_tbl_p0
+            PARTITION OF sh0001.partitioned_tbl
+            FOR VALUES WITH (MODULUS 2, REMAINDER 0);
+          CREATE TABLE sh0001.partitioned_tbl_p1
+            PARTITION OF sh0001.partitioned_tbl
+            FOR VALUES WITH (MODULUS 2, REMAINDER 1);
+          INSERT INTO sh0001.partitioned_tbl(grp, idx)
+            SELECT s, s FROM generate_series(1, 1000) AS s;
+        `,
       );
       await move({
         shard: 1,
@@ -50,6 +69,18 @@ beforeEach(async () => {
           await runSql.one(toDsn, sql`SELECT nextval('sh0001.tbl_id_seq')`),
         ),
       ).toBeGreaterThan(1500);
+      expect(
+        await runSql.one(
+          toDsn,
+          sql`SELECT count(1) FROM sh0001."weirdly""named;tbl"`,
+        ),
+      ).toEqual("1000");
+      expect(
+        await runSql.one(
+          toDsn,
+          sql`SELECT count(1) FROM sh0001.partitioned_tbl`,
+        ),
+      ).toEqual("1000");
       expect(out.stdout).toMatch(/ALTER SUBSCRIPTION.*ENABLE/m);
     } catch (e: unknown) {
       out.print();

src/api/__tests__/move.with-tables.test.ts

@@ -19,7 +19,7 @@ import {
   runSqlProgressTransactional,
 } from "../internal/runSqlProgress";
 import { startCopyingTables } from "../internal/startCopyingTables";
-import { validateTables } from "../internal/validateTables";
+import { validatePrimaryKeysFromDump } from "../internal/validatePrimaryKeysFromDump";
 import { waitUntilBackfillCompletes } from "../internal/waitUntilBackfillCompletes";
 import { waitUntilEverythingIsNotLaggingMuch } from "../internal/waitUntilEverythingIsNotLaggingMuch";
 import type { Unlock } from "../internal/waitUntilIncrementalCompletes";
@@ -53,12 +53,18 @@ export async function move({
       throw `Can't determine schema name for microshard number ${shard}`;
     }
 
-    const tables = await getTablesInSchema({ fromDsn, schema });
+    const { tables, partitionedTables } = await getTablesInSchema({
+      fromDsn,
+      schema,
+    });
     const { preData, primaryKeys, postData } = await getDump({
       fromDsn,
       schema,
     });
-    validateTables({ tables, primaryKeys });
+    validatePrimaryKeysFromDump({
+      primaryKeys,
+      allTables: [...tables, ...partitionedTables],
+    });
 
     await waitUntilEverythingIsNotLaggingMuch({
       when: "before starting the work",

src/api/move.ts

@@ -39,7 +39,7 @@ export async function getDump({
       `Reading post-data DDL (indexes, foreign keys) for ${schema} from ${fromDsn}`,
     ),
   );
-  // To logical replication to work, the tables on the subscription part must
+  // For logical replication to work, each table on the subscription part must
   // have an index that allows the logical worker to find the row it receives
   // from the publication end. PostgreSQL 16+ can use regular indexes for that,
   // but the earlier versions require the subscription end to have either a

src/internal/getDump.ts

@@ -1,3 +1,4 @@
+import partition from "lodash/partition";
 import { sql } from "./quote";
 import { runSql } from "./runSql";
 
@@ -8,6 +9,9 @@ export interface TableInfo {
   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: {
@@ -22,20 +26,30 @@ export interface TableInfo {
   } | 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<TableInfo[]> {
+}): 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,
@@ -45,7 +59,7 @@ export async function getTablesInSchema({
       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 = 'r' AND nspname = ${schema}
+      WHERE pg_class.relkind IN ('r', 'p') AND nspname = ${schema}
       ORDER BY
         pg_namespace.nspname,
         pg_class.relname,
@@ -54,17 +68,23 @@ export async function getTablesInSchema({
         pg_index.indisprimary DESC
     `,
   );
-  return rows.map((row) => ({
+  const res = rows.map((row) => ({
     schema: row[0],
     name: row[1],
     owner: row[2],
-    identity: row[3]
+    isPartitioned: row[3] === "p",
+    identity: row[4]
       ? {
-          indexName: row[3],
-          isReplIdent: row[4] === "t",
-          isPrimaryKey: row[5] === "t",
-          def: row[6],
+          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 };
 }

src/internal/getTablesInSchema.ts

@@ -0,0 +1,12 @@
+import escapeRegExp from "lodash/escapeRegExp";
+import { ident } from "./quote";
+
+/**
+ * Returns true if an SQL clause includes a properly quoted identifier (like
+ * index name).
+ */
+export function isNameInSql(name: string, sql: string): boolean {
+  return !!sql.match(
+    new RegExp("(\\W|^)" + escapeRegExp(ident(name).toString()) + "(\\W|$)"),
+  );
+}

src/internal/isNameInSql.ts

@@ -1,14 +1,19 @@
 import type { TableInfo } from "./getTablesInSchema";
-import { ident } from "./quote";
+import { isNameInSql } from "./isNameInSql";
 
-export function validateTables({
-  tables,
+/**
+ * Validates that the list of primary keys extracted from the dump earlier match
+ * the actual list of tables read from pg_catalog. This is mainly for sanity
+ * checking the dump parsing logic.
+ */
+export function validatePrimaryKeysFromDump({
   primaryKeys,
+  allTables,
 }: {
-  tables: TableInfo[];
   primaryKeys: string[];
+  allTables: TableInfo[];
 }): void {
-  const tablesWithoutIdentity = tables.filter(({ identity }) => !identity);
+  const tablesWithoutIdentity = allTables.filter((table) => !table.identity);
   if (tablesWithoutIdentity.length > 0) {
     throw (
       "To be moved, every table in the schema must have a PRIMARY KEY " +
@@ -18,10 +23,14 @@ export function validateTables({
     );
   }
 
+  const tables = allTables as Array<
+    TableInfo & { identity: NonNullable<TableInfo["identity"]> }
+  >;
+
   const tablesWithoutDumpEntries = tables.filter(
-    ({ identity }) =>
-      !primaryKeys.some((part) =>
-        part.includes(ident(identity!.indexName).toString()),
+    (table) =>
+      !primaryKeys.some((primaryKeySql) =>
+        isNameInSql(table.identity.indexName, primaryKeySql),
       ),
   );
   if (tablesWithoutDumpEntries.length > 0) {
@@ -30,15 +39,15 @@ export function validateTables({
       "are not found in the parsed post-data section of the dump. Is there a bug " +
       "in dump parsing logic? Missing indexes:\n" +
       tablesWithoutDumpEntries
-        .map(({ name, identity }) => `- ${identity?.indexName} (table ${name})`)
+        .map(({ name, identity }) => `- ${identity.indexName} (table ${name})`)
         .join("\n")
     );
   }
 
   const dumpEntriesWithoutTables = primaryKeys.filter(
-    (part) =>
-      !tables.some(({ identity }) =>
-        part.includes(ident(identity!.indexName).toString()),
+    (primaryKeySql) =>
+      !tables.some((table) =>
+        isNameInSql(table.identity.indexName, primaryKeySql),
       ),
   );
   if (dumpEntriesWithoutTables.length > 0) {