WIP Prepare postgres for migration by replication

roles/postgresql_replica_logical.rb

@@ -0,0 +1,7 @@
+name "postgresql_replica_logical"
+
+run_list %w(
+  kosmos_postgresql::hostsfile
+  kosmos_postgresql::replica_logical
+  kosmos_postgresql::firewall
+)

site-cookbooks/kosmos_postgresql/attributes/default.rb

@@ -1,3 +1,8 @@
+node.default['kosmos_postgresql']['postgresql_version'] = "14"
+
 # This is set to false by default, and set to true in the server resource
 # for replicas.
 node.default['kosmos_postgresql']['ready_to_set_up_replica'] = false
+
+# Address space from which clients are allowed to connect
+node.default['kosmos_postgresql']['access_addr'] = "10.1.1.0/24"

site-cookbooks/kosmos_postgresql/files/create_publications.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+set -e
+
+echo "== Creating publication in each database =="
+
+for db in $(psql -Atqc "SELECT datname FROM pg_database WHERE datallowconn AND datname NOT IN ('template0','template1')"); do
+  echo "Processing DB: $db"
+
+  # Create publication (idempotent)
+  psql -d "$db" -v ON_ERROR_STOP=1 <<SQL
+DO \$\$
+BEGIN
+  IF NOT EXISTS (
+    SELECT 1 FROM pg_publication WHERE pubname = 'migrate_pub'
+  ) THEN
+    CREATE PUBLICATION migrate_pub FOR ALL TABLES;
+  END IF;
+END
+\$\$;
+SQL
+
+  # Create logical replication slot (idempotent-ish)
+  SLOT="migrate_slot_${db}"
+
+  if ! psql -d "$db" -Atqc "SELECT 1 FROM pg_replication_slots WHERE slot_name = '$SLOT'" | grep -q 1; then
+    echo "  Creating slot: $SLOT"
+    psql -d "$db" -c "SELECT pg_create_logical_replication_slot('$SLOT', 'pgoutput');"
+  else
+    echo "  Slot already exists: $SLOT"
+  fi
+
+done
+
+echo "== Done =="

site-cookbooks/kosmos_postgresql/files/drop_publications.sh

@@ -0,0 +1,33 @@
+set -e
+
+echo "== Dropping subscriptions slots and publications on PRIMARY =="
+
+for db in $(psql -Atqc "SELECT datname FROM pg_database WHERE datallowconn AND datname NOT IN ('template0','template1')"); do
+  echo "Processing DB: $db"
+
+  SLOT="migrate_slot_${db}"
+
+  # Drop slot if exists
+  if psql -d "$db" -Atqc "SELECT 1 FROM pg_replication_slots WHERE slot_name = '$SLOT'" | grep -q 1; then
+    echo "  Dropping slot: $SLOT"
+    psql -d "$db" -c "SELECT pg_drop_replication_slot('$SLOT');"
+  else
+    echo "  Slot not found: $SLOT"
+  fi
+
+  # Drop publication if exists
+  psql -d "$db" -v ON_ERROR_STOP=1 <<SQL
+DO \$\$
+BEGIN
+  IF EXISTS (
+    SELECT 1 FROM pg_publication WHERE pubname = 'migrate_pub'
+  ) THEN
+    DROP PUBLICATION migrate_pub;
+  END IF;
+END
+\$\$;
+SQL
+
+done
+
+echo "== Done =="

site-cookbooks/kosmos_postgresql/files/drop_subscriptions.sh

@@ -0,0 +1,28 @@
+set -e
+
+echo "== Dropping subscriptions on PG14 =="
+
+for db in $(psql -Atqc "SELECT datname FROM pg_database WHERE datallowconn AND datname NOT IN ('template0,'template1'')"); do
+  echo "Processing DB: $db"
+
+  SUB="migrate_sub_${db}"
+
+  # Disable first (important)
+  psql -d "$db" -c "ALTER SUBSCRIPTION $SUB DISABLE;" 2>/dev/null || true
+
+  # Drop subscription if exists
+  psql -d "$db" -v ON_ERROR_STOP=1 <<SQL
+DO \$\$
+BEGIN
+  IF EXISTS (
+    SELECT 1 FROM pg_subscription WHERE subname = '$SUB'
+  ) THEN
+    DROP SUBSCRIPTION $SUB;
+  END IF;
+END
+\$\$;
+SQL
+
+done
+
+echo "== Done =="

site-cookbooks/kosmos_postgresql/files/dump_all_databases.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+cd /tmp && \
+(pg_dumpall --globals-only > globals.sql) && \
+psql -Atqc "SELECT datname FROM pg_database WHERE datallowconn AND datname NOT IN (''template0'')" | \
+xargs -I{} -P4 sh -c "
+  pg_dump -Fd -j 4 -d \"{}\" -f dump_{} &&
+  tar -cf - dump_{} | zstd -19 -T0 > dump_{}.tar.zst &&
+  rm -rf dump_{}
+"

site-cookbooks/kosmos_postgresql/files/list_publications.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+for db in $(psql -Atqc "SELECT datname FROM pg_database WHERE datallowconn"); do
+  echo "DB: $db"
+  psql -d "$db" -Atqc "SELECT pubname FROM pg_publication;"
+done

site-cookbooks/kosmos_postgresql/files/list_replication_slots.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+psql -c "
+SELECT slot_name,
+pg_size_pretty(pg_wal_lsn_diff(pg_current_wal_lsn(), restart_lsn))
+FROM pg_replication_slots;"

site-cookbooks/kosmos_postgresql/files/list_subscriptions.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+for db in $(psql -Atqc "SELECT datname FROM pg_database WHERE datallowconn AND datname NOT IN ('template0','template1')"); do
+  echo "==== DB: $db ===="
+  psql -d "$db" -c "SELECT * FROM pg_stat_subscription;"
+done

site-cookbooks/kosmos_postgresql/files/restore_all_databases.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+cd /tmp
+for f in dump_*.tar.zst; do
+  db=$(echo $f | sed "s/dump_\(.*\)\.tar\.zst/\1/")
+  echo "Restoring $db"
+  zstd -d "$f" -c | tar -xf -
+  pg_restore -j 4 -d "$db" dump_$db
+done

site-cookbooks/kosmos_postgresql/libraries/helpers.rb

@@ -36,10 +36,16 @@ class Chef
       end
     end
 
-    def postgresql_service_name
+    def postgresql_version
-      postgresql_version = "12"
+      node['kosmos_postgresql']['postgresql_version']
+    end
 
+    def postgresql_service
       "postgresql@#{postgresql_version}-main"
     end
+
+    def postgresql_data_dir
+      "/var/lib/postgresql/#{postgresql_version}/main"
+    end
   end
 end

site-cookbooks/kosmos_postgresql/recipes/primary.rb

@@ -3,31 +3,41 @@
 # Recipe:: primary
 #
 
-postgresql_version = "12"
-postgresql_service = "postgresql@#{postgresql_version}-main"
-
-service postgresql_service do
-  supports restart: true, status: true, reload: true
-end
-
 postgresql_custom_server postgresql_version do
   role "primary"
 end
 
-postgresql_access "zerotier members" do
+cookbook_file "/usr/local/bin/pg_dump_all_databases" do
-  access_type "host"
+  source "dump_all_databases.sh"
-  access_db "all"
+  user "postgres"
-  access_user "all"
+  group "postgres"
-  access_addr "10.1.1.0/24"
+  mode "0744"
-  access_method "md5"
-  notifies :reload, "service[#{postgresql_service}]", :immediately
 end
 
-postgresql_access "zerotier members replication" do
+cookbook_file "/usr/local/bin/pg_create_replication_publications" do
-  access_type "host"
+  source "create_publications.sh"
-  access_db "replication"
+  user "postgres"
-  access_user "replication"
+  group "postgres"
-  access_addr "10.1.1.0/24"
+  mode "0744"
-  access_method "md5"
+end
-  notifies :reload, "service[#{postgresql_service}]", :immediately
+
+cookbook_file "/usr/local/bin/pg_drop_replication_publications" do
+  source "drop_publications.sh"
+  user "postgres"
+  group "postgres"
+  mode "0744"
+end
+
+cookbook_file "/usr/local/bin/pg_list_replication_publications" do
+  source "list_publications.sh"
+  user "postgres"
+  group "postgres"
+  mode "0744"
+end
+
+cookbook_file "/usr/local/bin/pg_list_replication_slots" do
+  source "list_replication_slots.sh"
+  user "postgres"
+  group "postgres"
+  mode "0744"
 end

site-cookbooks/kosmos_postgresql/recipes/replica.rb

@@ -3,54 +3,32 @@
 # Recipe:: replica
 #
 
-postgresql_version = "12"
+service postgresql_service do
-postgresql_service = "postgresql@#{postgresql_version}-main"
+  supports restart: true, status: true, reload: true
+end
 
 postgresql_custom_server postgresql_version do
   role "replica"
 end
 
-service postgresql_service do
-  supports restart: true, status: true, reload: true
-end
-
 postgresql_data_bag_item = data_bag_item('credentials', 'postgresql')
 
 primary = postgresql_primary
 
-unless primary.nil?
+if primary.nil?
-  # TODO
+  Chef::Log.warn("No PostgreSQL primary node found. Skipping replication setup.")
-  postgresql_data_dir = "/var/lib/postgresql/#{postgresql_version}/main"
+  return
+end
 
-  # FIXME get zerotier IP
+execute "set up replication" do
-  execute "set up replication" do
+  command <<-EOF
-    command <<-EOF
 systemctl stop #{postgresql_service}
 mv #{postgresql_data_dir} #{postgresql_data_dir}.old
 pg_basebackup -h pg.kosmos.local -U replication -D #{postgresql_data_dir} -R
 chown -R postgres:postgres #{postgresql_data_dir}
 systemctl start #{postgresql_service}
-    EOF
+  EOF
-    environment 'PGPASSWORD' => postgresql_data_bag_item['replication_password']
+  environment 'PGPASSWORD' => postgresql_data_bag_item['replication_password']
-    sensitive true
+  sensitive true
-    not_if { ::File.exist? "#{postgresql_data_dir}/standby.signal" }
+  not_if { ::File.exist? "#{postgresql_data_dir}/standby.signal" }
-  end
-
-  postgresql_access "zerotier members" do
-    access_type "host"
-    access_db "all"
-    access_user "all"
-    access_addr "10.1.1.0/24"
-    access_method "md5"
-    notifies :reload, "service[#{postgresql_service}]", :immediately
-  end
-
-  postgresql_access "zerotier members replication" do
-    access_type "host"
-    access_db "replication"
-    access_user "replication"
-    access_addr "10.1.1.0/24"
-    access_method "md5"
-    notifies :reload, "service[#{postgresql_service}]", :immediately
-  end
 end

site-cookbooks/kosmos_postgresql/recipes/replica_logical.rb

@@ -0,0 +1,50 @@
+#
+# Cookbook:: kosmos_postgresql
+# Recipe:: replica_logical
+#
+
+service postgresql_service do
+  supports restart: true, status: true, reload: true
+end
+
+postgresql_custom_server postgresql_version do
+  role "replica_logical"
+end
+
+postgresql_data_bag_item = data_bag_item('credentials', 'postgresql')
+
+primary = postgresql_primary
+
+if primary.nil?
+  Chef::Log.warn("No PostgreSQL primary node found. Skipping replication setup.")
+  return
+end
+
+template "/usr/local/bin/pg_create_replication_subscriptions" do
+  source "create_subscriptions.sh.erb"
+  user "postgres"
+  group "postgres"
+  mode "0740"
+  sensitive true
+end
+
+cookbook_file "/usr/local/bin/pg_drop_replication_subscriptions" do
+  source "drop_subscriptions.sh"
+  user "postgres"
+  group "postgres"
+  mode "0744"
+end
+
+cookbook_file "/usr/local/bin/pg_list_replication_subscriptions" do
+  source "list_subscriptions.sh"
+  user "postgres"
+  group "postgres"
+  mode "0744"
+end
+
+cookbook_file "/usr/local/bin/pg_restore_all_databases" do
+  source "restore_all_databases.sh"
+  user "postgres"
+  group "postgres"
+  mode "0744"
+end

site-cookbooks/kosmos_postgresql/resources/server.rb

@@ -56,7 +56,9 @@ action :create do
     timezone: "UTC", # default is GMT
     listen_addresses: "0.0.0.0",
     promote_trigger_file: "#{postgresql_data_dir}/failover.trigger",
-    wal_keep_segments: 256
+    wal_level: "logical",
+    wal_keep_size: 4096, # 256 segments, 16MB each
+    max_replication_slots: 16
   }
 
   postgresql_server_conf "main" do
@@ -70,6 +72,24 @@ action :create do
     replication true
     password postgresql_credentials['replication_password']
   end
+
+  postgresql_access "all members" do
+    access_type "host"
+    access_db "all"
+    access_user "all"
+    access_addr node['kosmos_postgresql']['access_addr']
+    access_method "md5"
+    notifies :reload, "service[#{postgresql_service}]", :immediately
+  end
+
+  postgresql_access "replication members" do
+    access_type "host"
+    access_db "replication"
+    access_user "replication"
+    access_addr node['kosmos_postgresql']['access_addr']
+    access_method "md5"
+    notifies :reload, "service[#{postgresql_service}]", :immediately
+  end
 end
 
 action_class do

site-cookbooks/kosmos_postgresql/templates/create_subscriptions.sh.erb

@@ -0,0 +1,33 @@
+set -e
+
+echo "== Creating subscriptions for all databases =="
+
+for db in $(psql -Atqc "SELECT datname FROM pg_database WHERE datallowconn AND datname NOT IN ('template0','template1')"); do
+  echo "Processing DB: $db"
+
+  SLOT="migrate_slot_${db}"
+  SUB="migrate_sub_${db}"
+
+  psql -d "$db" -v ON_ERROR_STOP=1 <<SQL
+DO \$\$
+BEGIN
+  IF NOT EXISTS (
+    SELECT 1 FROM pg_subscription WHERE subname = '$SUB'
+  ) THEN
+    CREATE SUBSCRIPTION $SUB
+    CONNECTION 'host=<%= @pg_host %> port=<%= @pg_port %> dbname=$db user=<%= @pg_user %> password=<%= @pg_pass %>'
+    PUBLICATION migrate_pub
+    WITH (
+      slot_name = '$SLOT',
+      create_slot = false,
+      copy_data = false,
+      enabled = true
+    );
+  END IF;
+END
+\$\$;
+SQL
+
+done
+
+echo "== Done =="