pg partman support

README.md

@@ -50,6 +50,12 @@ heroku pg:psql -a your-app-name -c "SELECT extname, extversion FROM pg_extension
 
 For each extension listed, enable it on your PlanetScale database before starting the migration. See the [PlanetScale Postgres extensions documentation](https://planetscale.com/docs/postgres/extensions) for supported extensions and how to enable them. If you need help, [contact us](https://planetscale.com/contact).
 
+If your Heroku database uses `pg_partman`, the migrator copies the partitioned
+table schema and uses pg_partman's metadata dump function to recreate partition
+maintenance configuration on PlanetScale. Bucardo then replicates the
+application partition tables, but it intentionally does not replicate
+pg_partman's internal schemas such as `partman` or `pg_partman`.
+
 ### 4. Check for blocking vacuum processes
 
 Bucardo creates triggers on your Heroku tables to track changes. In rare cases, a long-running autovacuum process can block trigger creation, which can also block your application's queries. Before starting the migration, check for wraparound vacuum processes:

scripts/mk-bucardo-repl.sh

@@ -33,12 +33,59 @@ if [ -z "$PRIMARY" -o -z "$REPLICA" ]
 then usage 1
 fi
 
+internal_schema_filter_sql() {
+  cat <<'SQL'
+    n.nspname <> 'information_schema'
+    AND n.nspname <> 'bucardo'
+    AND n.nspname <> 'heroku_ext'
+    AND n.nspname <> 'partman'
+    AND n.nspname <> 'pg_partman'
+    AND left(n.nspname, 3) <> 'pg_'
+SQL
+}
+
+sql_identifier() {
+  printf '"%s"' "$(printf "%s" "$1" | sed 's/"/""/g')"
+}
+
 # Copy the schema from the primary to the (soon to be) replica.
 if [ "$SKIP_SCHEMA" -eq 0 ]; then
   echo "Copying schema from primary to replica..."
   pg_dump --no-owner --no-privileges --no-publications --no-subscriptions --schema-only "$PRIMARY" |
   grep -v -E "^COMMENT ON EXTENSION " |
   psql "$REPLICA" -a --set ON_ERROR_STOP=1
+
+  PG_PARTMAN_SCHEMA=$(psql "$PRIMARY" -A -t -c "
+    SELECT n.nspname
+    FROM pg_namespace n
+    WHERE n.nspname IN ('partman', 'pg_partman')
+      AND EXISTS (
+        SELECT 1
+        FROM pg_class c
+        WHERE c.relnamespace = n.oid
+          AND c.relname = 'part_config'
+      )
+      AND EXISTS (
+        SELECT 1
+        FROM pg_proc p
+        WHERE p.pronamespace = n.oid
+          AND p.proname = 'dump_partitioned_table_definition'
+      )
+    ORDER BY n.nspname
+    LIMIT 1;")
+
+  if [ -n "$PG_PARTMAN_SCHEMA" ]; then
+    echo "Detected pg_partman metadata; recreating partition maintenance configuration on replica..."
+    pg_partman_ident=$(sql_identifier "$PG_PARTMAN_SCHEMA")
+    PG_PARTMAN_RECREATE_SQL=$(psql "$PRIMARY" -A -t -c "
+      SELECT ${pg_partman_ident}.dump_partitioned_table_definition(parent_table)
+      FROM ${pg_partman_ident}.part_config
+      ORDER BY parent_table;")
+    if [ -n "$PG_PARTMAN_RECREATE_SQL" ]; then
+      printf "%s\n" "$PG_PARTMAN_RECREATE_SQL" |
+      psql "$REPLICA" -a --set ON_ERROR_STOP=1
+    fi
+  fi
 else
   echo "Skipping schema copy (--skip-schema flag set)"
 fi
@@ -60,9 +107,51 @@ bucardo add database "planetscale" \
   password="$(echo "$REPLICA" | cut -d ":" -f 3 | cut -d "@" -f 1)" \
   dbname="$(echo "$REPLICA" | cut -d "/" -f 4 | cut -d "?" -f 1)"
 
-# Add all the sequences and tables to Bucardo.
-bucardo add all sequences --relgroup "planetscale_import"
-bucardo add all tables --relgroup "planetscale_import"
+# Add application sequences and tables to Bucardo. Extension/internal schemas
+# such as pg_partman are schema-copied above but should not be replicated.
+REPLICATION_SCHEMAS=$(psql "$PRIMARY" -A -t -c "
+  SELECT n.nspname
+  FROM pg_namespace n
+  WHERE $(internal_schema_filter_sql)
+    AND EXISTS (
+      SELECT 1
+      FROM pg_class c
+      WHERE c.relnamespace = n.oid
+        AND c.relkind IN ('r', 'p', 'S')
+    )
+  ORDER BY n.nspname;")
+
+if [ -z "$REPLICATION_SCHEMAS" ]; then
+  echo "No application schemas with tables or sequences were found for Bucardo replication" >&2
+  exit 1
+fi
+
+printf "%s\n" "$REPLICATION_SCHEMAS" | while IFS= read -r schema
+do
+  [ -z "$schema" ] && continue
+  echo "Adding Bucardo relations from schema: $schema"
+  bucardo add all sequences db=heroku -n "$schema" relgroup=planetscale_import
+done
+
+REPLICATION_TABLES=$(psql "$PRIMARY" -A -t -c "
+  SELECT format('%I.%I', n.nspname, c.relname)
+  FROM pg_class c
+  JOIN pg_namespace n ON n.oid = c.relnamespace
+  WHERE $(internal_schema_filter_sql)
+    AND c.relkind = 'r'
+  ORDER BY n.nspname, c.relname;")
+
+if [ -z "$REPLICATION_TABLES" ]; then
+  echo "No application tables were found for Bucardo replication" >&2
+  exit 1
+fi
+
+printf "%s\n" "$REPLICATION_TABLES" | while IFS= read -r table
+do
+  [ -z "$table" ] && continue
+  echo "Adding Bucardo table: $table"
+  bucardo add table "$table" db=heroku relgroup=planetscale_import
+done
 
 # Bucardo 5.6 does not filter out PostgreSQL generated columns when issuing
 # COPY against the target, which fails with "column ... is a generated column /
@@ -76,9 +165,7 @@ GENERATED_TABLES=$(psql "$PRIMARY" -A -t -F"|" -c "
   FROM pg_attribute a
   JOIN pg_class c ON c.oid = a.attrelid
   JOIN pg_namespace n ON n.oid = c.relnamespace
-  WHERE n.nspname <> 'information_schema'
-    AND n.nspname <> 'bucardo'
-    AND left(n.nspname, 3) <> 'pg_'
+  WHERE $(internal_schema_filter_sql)
     AND c.relkind = 'r'
     AND a.attnum > 0 AND NOT a.attisdropped
     AND a.attgenerated <> ''
@@ -106,10 +193,10 @@ fi
 # Add the sync configuration to Bucardo.
 if [ "$NO_INITIAL_COPY" -eq 0 ]; then
   echo "Configuring sync with initial data copy..."
-  bucardo add sync "planetscale_import" dbs="heroku,planetscale" onetimecopy=1 relgroup="planetscale_import"
+  bucardo add sync "planetscale_import" dbs="heroku,planetscale" onetimecopy=1 checktime=5 relgroup="planetscale_import"
 else
   echo "Configuring sync without initial copy (--no-initial-copy flag set)..."
-  bucardo add sync "planetscale_import" dbs="heroku,planetscale" onetimecopy=0 relgroup="planetscale_import"
+  bucardo add sync "planetscale_import" dbs="heroku,planetscale" onetimecopy=0 checktime=5 relgroup="planetscale_import"
 fi
 
 # Give Bucardo enough time to validate all tables across both databases.

status-server/server.rb

@@ -335,6 +335,9 @@ def check_tables_with_generated_columns
           "JOIN pg_namespace n ON n.oid = c.relnamespace " \
           "WHERE n.nspname <> 'information_schema' " \
           "  AND n.nspname <> 'bucardo' " \
+          "  AND n.nspname <> 'heroku_ext' " \
+          "  AND n.nspname <> 'partman' " \
+          "  AND n.nspname <> 'pg_partman' " \
           "  AND left(n.nspname, 3) <> 'pg_' " \
           "  AND c.relkind = 'r' " \
           "  AND a.attnum > 0 AND NOT a.attisdropped " \

test/test_pg_partman_setup.sh

@@ -0,0 +1,150 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+PROJECT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
+TMP_DIR="$(mktemp -d)"
+FAKE_BIN_DIR="$TMP_DIR/bin"
+LOG_FILE="$TMP_DIR/commands.log"
+
+cleanup() {
+  rm -rf "$TMP_DIR"
+}
+trap cleanup EXIT
+
+mkdir -p "$FAKE_BIN_DIR"
+
+cat > "$FAKE_BIN_DIR/pg_dump" <<'SH'
+#!/bin/sh
+printf "pg_dump %s\n" "$*" >> "$COMMAND_LOG"
+cat <<'SQL'
+CREATE SCHEMA public;
+CREATE SCHEMA app_private;
+CREATE SCHEMA partman;
+CREATE EXTENSION pg_partman WITH SCHEMA partman;
+CREATE TABLE public.partman_events (id bigint, created_at timestamptz NOT NULL, payload jsonb NOT NULL, PRIMARY KEY (id, created_at)) PARTITION BY RANGE (created_at);
+CREATE TABLE app_private.docs (id bigint PRIMARY KEY, title text NOT NULL);
+CREATE TABLE partman.template_public_partman_events (LIKE public.partman_events);
+SQL
+SH
+chmod +x "$FAKE_BIN_DIR/pg_dump"
+
+cat > "$FAKE_BIN_DIR/psql" <<'SH'
+#!/bin/sh
+query=""
+printf "psql %s\n" "$*" >> "$COMMAND_LOG"
+while [ "$#" -gt 0 ]; do
+  if [ "$1" = "-c" ]; then
+    shift
+    query="$1"
+  fi
+  shift
+done
+
+case "$query" in
+  *"dump_partitioned_table_definition(parent_table)"*"part_config"*)
+    cat <<'SQL'
+SELECT partman.create_parent(
+  p_parent_table := 'public.partman_events',
+  p_control := 'created_at',
+  p_interval := '1 month',
+  p_type := 'range'
+);
+SQL
+    ;;
+  *"part_config"*"pg_proc"*"dump_partitioned_table_definition"*)
+    echo "partman"
+    ;;
+  *"SELECT n.nspname"*"c.relkind IN ('r', 'p', 'S')"*)
+    echo "app_private"
+    echo "public"
+    ;;
+  *"SELECT format('%I.%I', n.nspname, c.relname)"*"c.relkind = 'r'"*)
+    echo "app_private.docs"
+    echo "public.partman_events_p20260101"
+    echo "public.partman_events_p20260201"
+    echo "public.posts"
+    ;;
+  *"SELECT DISTINCT n.nspname, c.relname"*)
+    echo "public|posts"
+    echo "app_private|docs"
+    ;;
+  *"string_agg(quote_ident(attname)"*)
+    echo "id, title"
+    ;;
+esac
+
+if [ -z "$query" ]; then
+  printf "psql stdin start\n" >> "$COMMAND_LOG"
+  cat >> "$COMMAND_LOG"
+  printf "psql stdin end\n" >> "$COMMAND_LOG"
+fi
+exit 0
+SH
+chmod +x "$FAKE_BIN_DIR/psql"
+
+cat > "$FAKE_BIN_DIR/bucardo" <<'SH'
+#!/bin/sh
+printf "bucardo %s\n" "$*" >> "$COMMAND_LOG"
+if [ "$1" = "status" ]; then
+  cat <<'OUT'
+Status                   : Active
+Onetimecopy              : No
+Current state            : Good
+OUT
+fi
+SH
+chmod +x "$FAKE_BIN_DIR/bucardo"
+
+export COMMAND_LOG="$LOG_FILE"
+export PATH="$FAKE_BIN_DIR:$PATH"
+
+sh "$PROJECT_DIR/scripts/mk-bucardo-repl.sh" \
+  --primary "postgres://user:pass@primary:5432/source" \
+  --replica "postgres://user:pass@replica:5432/target" \
+  >/dev/null
+
+PASS=0
+FAIL=0
+
+pass() { PASS=$((PASS + 1)); printf "  PASS %s\n" "$*"; }
+fail() { FAIL=$((FAIL + 1)); printf "  FAIL %s\n" "$*" >&2; }
+
+assert_log_contains() {
+  local pattern="$1" message="$2"
+  if grep -Fq -- "$pattern" "$LOG_FILE"; then
+    pass "$message"
+  else
+    fail "$message"
+  fi
+}
+
+assert_log_absent() {
+  local pattern="$1" message="$2"
+  if grep -Fq -- "$pattern" "$LOG_FILE"; then
+    fail "$message"
+  else
+    pass "$message"
+  fi
+}
+
+assert_log_contains "CREATE EXTENSION pg_partman WITH SCHEMA partman;" "schema copy preserves pg_partman extension"
+assert_log_contains "SELECT partman.create_parent(" "pg_partman recreation SQL is applied to replica"
+assert_log_contains "bucardo add table public.posts db=heroku relgroup=planetscale_import" "Bucardo includes public application table"
+assert_log_contains "bucardo add table app_private.docs db=heroku relgroup=planetscale_import" "Bucardo includes non-public application table"
+assert_log_contains "bucardo add table public.partman_events_p20260101 db=heroku relgroup=planetscale_import" "Bucardo includes partition leaf tables"
+assert_log_absent "bucardo add table public.partman_events db=heroku relgroup=planetscale_import" "Bucardo excludes partitioned parent table"
+assert_log_absent "bucardo add table partman.template_public_partman_events" "Bucardo excludes pg_partman template table"
+assert_log_absent "bucardo add all sequences db=heroku -n partman" "Bucardo excludes pg_partman internal sequences"
+assert_log_contains "bucardo add sync planetscale_import dbs=heroku,planetscale onetimecopy=1 checktime=5 relgroup=planetscale_import" "Bucardo sync has periodic checktime fallback"
+assert_log_contains "bucardo add customcols public.posts SELECT id, title db=planetscale" "Generated-column customcols still run for public tables"
+assert_log_contains "bucardo add customcols app_private.docs SELECT id, title db=planetscale" "Generated-column customcols still run for non-public tables"
+
+if [ "$FAIL" -gt 0 ]; then
+  printf "\nCommand log:\n" >&2
+  sed 's/^/  /' "$LOG_FILE" >&2
+  printf "\n%d passed, %d failed\n" "$PASS" "$FAIL" >&2
+  exit 1
+fi
+
+printf "\n%d passed, %d failed\n" "$PASS" "$FAIL"