From bbc0f59801ac8845e684148d0956b5c4c86f39e2 Mon Sep 17 00:00:00 2001
From: Varik Matevosyan <varikmatevosyan@gmail.com>
Date: Thu, 21 Nov 2024 14:20:41 +0400
Subject: [PATCH] clean up db before running pg_upgrade

---
 model/lantern/lantern_server.rb               | 34 ++++++++++++++++++-
 prog/lantern/lantern_server_nexus.rb          | 17 ++++++----
 rhizome/lantern/bin/exec_all                  |  2 +-
 rhizome/lantern/bin/setup_ssl                 |  2 +-
 rhizome/lantern/lib/common.rb                 | 11 +++---
 spec/model/lantern/lantern_server_spec.rb     |  9 ++++-
 .../prog/lantern/lantern_server_nexus_spec.rb | 10 ++++--
 7 files changed, 67 insertions(+), 18 deletions(-)

diff --git a/model/lantern/lantern_server.rb b/model/lantern/lantern_server.rb
index c08136002..121d3f57a 100644
--- a/model/lantern/lantern_server.rb
+++ b/model/lantern/lantern_server.rb
@@ -242,7 +242,7 @@ def prewarm_indexes_query
   end
 
   def list_all_databases
-    vm.sshable.cmd("sudo docker compose -f #{Config.compose_file} exec postgresql psql -U postgres -P \"footer=off\" -c 'SELECT datname from pg_database' | tail -n +3 | grep -v 'template0' | grep -v 'template1'")
+    vm.sshable.cmd("sudo docker compose -f #{Config.compose_file} exec postgresql psql -U postgres -t -c 'SELECT datname FROM pg_database WHERE datistemplate=FALSE'")
       .chomp
       .strip
       .split("\n")
@@ -302,6 +302,38 @@ def start_container
     vm.sshable.cmd("sudo docker compose -f #{Config.compose_file} up -d")
   end
 
+  def prepare_database_for_upgrade
+    # drop lantern_extras extension because it has breaking binary changes
+    # drop all indexes which does not have `dim` specified in params or are old external indexes
+    # which use _experimental_index_path param
+    commands = <<SQL
+DROP EXTENSION IF EXISTS lantern_extras;
+DO $$
+DECLARE
+    drop_query TEXT;
+BEGIN
+    FOR drop_query IN
+        SELECT 'DROP INDEX "' || i.relname || '";'
+        FROM pg_class t
+        JOIN pg_index ix ON t.oid = ix.indrelid
+        JOIN pg_class i ON i.oid = ix.indexrelid
+        JOIN pg_am a ON i.relam = a.oid
+        JOIN pg_namespace n ON n.oid = i.relnamespace
+        WHERE a.amname = 'lantern_hnsw'
+        AND (
+          lower(pg_get_indexdef(i.oid)) NOT SIMILAR TO '%"?dim"?=%'
+          OR
+          lower(pg_get_indexdef(i.oid)) SIMILAR TO '%"?_experimental_index_path"?=%'
+        )
+    LOOP
+        EXECUTE drop_query;
+    END LOOP;
+END
+$$;
+SQL
+    run_query_all(commands)
+  end
+
   # def failover_target
   #   nil
   # end
diff --git a/prog/lantern/lantern_server_nexus.rb b/prog/lantern/lantern_server_nexus.rb
index 9d8df48a7..302f352db 100644
--- a/prog/lantern/lantern_server_nexus.rb
+++ b/prog/lantern/lantern_server_nexus.rb
@@ -227,10 +227,14 @@ def before_run
         incr_run_pg_upgrade
       end
 
-      if lantern_server.resource.logical_replication && !lantern_server.resource.parent.representative_server.domain.nil?
-        # prepare for fast switchover
-        lantern_server.add_domain_to_stack(lantern_server.resource.parent.representative_server.domain, strand)
-        incr_setup_ssl
+      if lantern_server.resource.logical_replication
+        lantern_server.resource.drop_ddl_log_trigger
+
+        if !lantern_server.resource.parent.representative_server.domain.nil?
+          # prepare for fast switchover
+          lantern_server.add_domain_to_stack(lantern_server.resource.parent.representative_server.domain, strand)
+          incr_setup_ssl
+        end
       end
 
       hop_wait_timeline_available
@@ -241,8 +245,8 @@ def before_run
 
   label def run_pg_upgrade
     decr_run_pg_upgrade
-    lantern_server.resource.drop_ddl_log_trigger
     pg_upgrade_info = strand.stack.first["pg_upgrade"]
+    lantern_server.prepare_database_for_upgrade
     vm.sshable.cmd(
       "common/bin/daemonizer 'sudo lantern/bin/run_pg_upgrade' pg_upgrade",
       stdin: JSON.generate({
@@ -402,7 +406,8 @@ def before_run
         dns_token: Config.cf_token,
         dns_zone_id: Config.cf_zone_id,
         dns_email: Config.lantern_dns_email,
-        domain: frame["domain"] || lantern_server.domain
+        domain: frame["domain"] || lantern_server.domain,
+        pg_version: lantern_server.resource.pg_version
       }))
     when "Failed"
       logs = JSON.parse(vm.sshable.cmd("common/bin/daemonizer --logs setup_ssl"))
diff --git a/rhizome/lantern/bin/exec_all b/rhizome/lantern/bin/exec_all
index 514a9c4a5..807e18971 100755
--- a/rhizome/lantern/bin/exec_all
+++ b/rhizome/lantern/bin/exec_all
@@ -6,7 +6,7 @@ require "yaml"
 require_relative "../../common/lib/util"
 require_relative "../lib/common"
 
-all_dbs = (r "docker compose -f #{$compose_file} exec postgresql psql -U postgres -P \"footer=off\" -c 'SELECT datname from pg_database' | tail -n +3 | grep -v 'template0' | grep -v 'template1'").strip!.split("\n")
+all_dbs = (r "docker compose -f #{$compose_file} exec postgresql psql -U postgres -t -c 'SELECT datname FROM pg_database WHERE datistemplate=FALSE'").strip!.split("\n")
 all_dbs.each do |db|
   $stdout.puts r "docker compose -f #{$compose_file} exec -T postgresql psql -U postgres -t --csv #{db}", stdin: $stdin
 end
diff --git a/rhizome/lantern/bin/setup_ssl b/rhizome/lantern/bin/setup_ssl
index 46c247a9b..54a07bfaa 100755
--- a/rhizome/lantern/bin/setup_ssl
+++ b/rhizome/lantern/bin/setup_ssl
@@ -8,4 +8,4 @@ require_relative "../lib/common"
 
 $configure_hash = JSON.parse($stdin.read)
 
-configure_tls($configure_hash["domain"], $configure_hash["dns_email"], $configure_hash["dns_token"], $configure_hash["dns_zone_id"], "dns_cf")
+configure_tls($configure_hash["domain"], $configure_hash["dns_email"], $configure_hash["dns_token"], $configure_hash["dns_zone_id"], "dns_cf", $configure_hash["pg_version"])
diff --git a/rhizome/lantern/lib/common.rb b/rhizome/lantern/lib/common.rb
index 4fe29f2e2..144e6d2f1 100755
--- a/rhizome/lantern/lib/common.rb
+++ b/rhizome/lantern/lib/common.rb
@@ -16,7 +16,7 @@ def configure_gcr(container_image)
 end
 
 def update_extensions_in_sql
-  all_dbs = (r "docker compose -f #{$compose_file} exec postgresql psql -U postgres -P \"footer=off\" -c 'SELECT datname from pg_database' | tail -n +3 | grep -v 'template0' | grep -v 'template1'").strip!.split("\n")
+  all_dbs = (r "docker compose -f #{$compose_file} exec postgresql psql -U postgres -t -c 'SELECT datname FROM pg_database WHERE datistemplate=FALSE'").strip!.split("\n")
   tmp_cmd = ""
   all_dbs.each do |db|
     tmp_cmd += "docker compose -f #{$compose_file} exec postgresql psql -U postgres -f /lantern-init.sql #{db} \n"
@@ -93,7 +93,7 @@ def tls_already_configured?(domain)
   !r("(test -f #{$datadir}/server.key && test -f #{$datadir}/server.crt && echo 1) || echo ''").chomp.empty?
 end
 
-def configure_tls(domain, email, dns_token, dns_zone_id, provider)
+def configure_tls(domain, email, dns_token, dns_zone_id, provider, pg_version)
   puts "Configuring TLS for domain #{domain}"
 
   if !tls_already_configured?(domain)
@@ -107,11 +107,12 @@ def configure_tls(domain, email, dns_token, dns_zone_id, provider)
     r "#{env} /root/.acme.sh/acme.sh --server letsencrypt --issue --dns #{provider} -d #{domain}"
     reload_cmd = "sudo docker compose -f #{$compose_file} exec postgresql psql -U postgres -c 'SELECT pg_reload_conf()' && sudo docker compose -f #{$compose_file} exec postgresql psql -p6432 -U postgres pgbouncer -c RELOAD"
     r "/root/.acme.sh/acme.sh --install-cert -d #{domain} --key-file #{$datadir}/server.key  --fullchain-file #{$datadir}/server.crt --reloadcmd \"#{reload_cmd}\""
-    chown_with_daemon_user "#{$datadir}/server.key", pg_version
-    chown_with_daemon_user "#{$datadir}/server.crt", pg_version
-    r "sudo chmod 600 #{$datadir}/server.key"
   end
 
+  chown_with_daemon_user "#{$datadir}/server.key", pg_version
+  chown_with_daemon_user "#{$datadir}/server.crt", pg_version
+  r "sudo chmod 600 #{$datadir}/server.key"
+
   append_env([
     ["POSTGRESQL_ENABLE_TLS", "yes"],
     ["POSTGRESQL_TLS_CERT_FILE", "/bitnami/postgresql/server.crt"],
diff --git a/spec/model/lantern/lantern_server_spec.rb b/spec/model/lantern/lantern_server_spec.rb
index 9ad3d1593..715ddd5a8 100644
--- a/spec/model/lantern/lantern_server_spec.rb
+++ b/spec/model/lantern/lantern_server_spec.rb
@@ -691,7 +691,7 @@
 
   describe "#list_all_databases" do
     it "returns list of all databases" do
-      expect(lantern_server.vm.sshable).to receive(:cmd).with("sudo docker compose -f /var/lib/lantern/docker-compose.yaml exec postgresql psql -U postgres -P \"footer=off\" -c 'SELECT datname from pg_database' | tail -n +3 | grep -v 'template0' | grep -v 'template1'").and_return("postgres\ndb2\n")
+      expect(lantern_server.vm.sshable).to receive(:cmd).with("sudo docker compose -f /var/lib/lantern/docker-compose.yaml exec postgresql psql -U postgres -t -c 'SELECT datname FROM pg_database WHERE datistemplate=FALSE'").and_return("postgres\ndb2\n")
       expect(lantern_server.list_all_databases).to eq(["postgres", "db2"])
     end
   end
@@ -906,4 +906,11 @@
       expect { lantern_server.remove_domain_from_stack(strand) }.not_to raise_error
     end
   end
+
+  describe "#prepare_database_for_upgrade" do
+    it "cleans database before running pg_upgrade" do
+      expect(lantern_server).to receive(:run_query_all)
+      expect { lantern_server.prepare_database_for_upgrade }.not_to raise_error
+    end
+  end
 end
diff --git a/spec/prog/lantern/lantern_server_nexus_spec.rb b/spec/prog/lantern/lantern_server_nexus_spec.rb
index fdcc3b818..5cac960ed 100644
--- a/spec/prog/lantern/lantern_server_nexus_spec.rb
+++ b/spec/prog/lantern/lantern_server_nexus_spec.rb
@@ -423,6 +423,7 @@
       expect(lantern_server.resource).to receive(:parent).and_return(parent_reosurce).at_least(:once)
       expect(lantern_server.resource).to receive(:logical_replication).and_return(true)
       expect(lantern_server.resource).to receive(:allow_timeline_access_to_bucket)
+      expect(lantern_server.resource).to receive(:drop_ddl_log_trigger)
       expect(lantern_server).to receive(:run_query).and_return("f")
       expect(lantern_server).to receive(:timeline_id=)
       expect(lantern_server).to receive(:timeline_access=).with("push")
@@ -443,6 +444,7 @@
       expect(lantern_server.resource).to receive(:parent).and_return(parent_reosurce).at_least(:once)
       expect(lantern_server.resource).to receive(:logical_replication).and_return(true)
       expect(lantern_server.resource).to receive(:allow_timeline_access_to_bucket)
+      expect(lantern_server.resource).to receive(:drop_ddl_log_trigger)
       expect(lantern_server).to receive(:add_domain_to_stack).with(parent_reosurce.representative_server.domain, nx.strand)
       expect(nx).to receive(:incr_setup_ssl)
       expect(lantern_server).to receive(:run_query).and_return("f")
@@ -647,7 +649,8 @@
         dns_token: Config.cf_token,
         dns_zone_id: Config.cf_zone_id,
         dns_email: Config.lantern_dns_email,
-        domain: "db.lantern.dev"
+        domain: "db.lantern.dev",
+        pg_version: lantern_server.resource.pg_version
       }))
       expect { nx.setup_ssl }.to nap(10)
     end
@@ -662,7 +665,8 @@
         dns_token: Config.cf_token,
         dns_zone_id: Config.cf_zone_id,
         dns_email: Config.lantern_dns_email,
-        domain: lantern_server.domain
+        domain: lantern_server.domain,
+        pg_version: lantern_server.resource.pg_version
       }))
       expect { nx.setup_ssl }.to nap(10)
     end
@@ -1067,7 +1071,7 @@
       frame = {"pg_upgrade" => {"lantern_version" => "0.5.0", "extras_version" => "0.5.0", "minor_version" => "1", "pg_version" => 17}}
       expect(nx.strand).to receive(:stack).and_return([frame]).at_least(:once)
       expect(lantern_server).to receive(:container_image).and_return(image).at_least(:once)
-      expect(lantern_server.resource).to receive(:drop_ddl_log_trigger)
+      expect(lantern_server).to receive(:prepare_database_for_upgrade)
       expect(lantern_server.vm.sshable).to receive(:cmd).with("common/bin/daemonizer 'sudo lantern/bin/run_pg_upgrade' pg_upgrade", stdin: JSON.generate(
         container_image: lantern_server.container_image,
         pg_version: 17,