From d36738b343983202dd2c333cff60c0b9624c05aa Mon Sep 17 00:00:00 2001
From: Paul Cioanca <paul.cioanca@supabase.io>
Date: Sun, 15 Dec 2024 17:49:35 +0200
Subject: [PATCH] fix: pgmq perms+data

---
 .../pg_upgrade_scripts/complete.sh            | 74 ++++++++++++++++++-
 .../20241215003910_backfill_pgmq_metadata.sql | 29 ++++++++
 2 files changed, 102 insertions(+), 1 deletion(-)
 create mode 100644 migrations/db/migrations/20241215003910_backfill_pgmq_metadata.sql

diff --git a/ansible/files/admin_api_scripts/pg_upgrade_scripts/complete.sh b/ansible/files/admin_api_scripts/pg_upgrade_scripts/complete.sh
index c8542c685..515c490f6 100755
--- a/ansible/files/admin_api_scripts/pg_upgrade_scripts/complete.sh
+++ b/ansible/files/admin_api_scripts/pg_upgrade_scripts/complete.sh
@@ -79,7 +79,79 @@ EOF
         run_sql -c "$RECREATE_PG_CRON_QUERY"
     fi
 
-    # #incident-2024-09-12-project-upgrades-are-temporarily-disabled
+    # Patching pgmq ownership as it resets during upgrade
+    HAS_PGMQ=$(run_sql -A -t -c "select count(*) > 0 from pg_extension where extname = 'pgmq';")
+    if [ "$HAS_PGMQ" = "t" ]; then
+        PATCH_PGMQ_QUERY=$(cat <<EOF
+        do \$\$
+        declare
+            tbl record;
+            seq_name text;
+            new_seq_name text;
+            archive_table_name text;
+        begin
+            -- Loop through each table in the pgmq schema starting with 'q_'
+            -- Rebuild the pkey column's default to avoid pg_dumpall segfaults
+            for tbl in
+                select c.relname as table_name
+                from pg_catalog.pg_attribute a
+                join pg_catalog.pg_class c on c.oid = a.attrelid
+                join pg_catalog.pg_namespace n on n.oid = c.relnamespace
+                where n.nspname = 'pgmq'
+                    and c.relname like 'q_%'
+                    and a.attname = 'msg_id'
+                    and a.attidentity in ('a', 'd') -- 'a' for ALWAYS, 'd' for BY DEFAULT
+            loop
+                -- Check if msg_id is an IDENTITY column for idempotency
+                -- Define sequence names
+                seq_name := 'pgmq.' || format ('"%s_msg_id_seq"', tbl.table_name);
+                new_seq_name := 'pgmq.' || format ('"%s_msg_id_seq2"', tbl.table_name);
+                archive_table_name := regexp_replace(tbl.table_name, '^q_', 'a_');
+                -- Execute dynamic SQL to perform the required operations
+                execute format('
+                    create sequence %s;
+                    select setval(''%s'', nextval(''%s''));
+                    alter table %s."%s" alter column msg_id drop identity;
+                    alter table %s."%s" alter column msg_id set default nextval(''%s'');
+                    alter sequence %s rename to "%s";',
+                    -- Parameters for format placeholders
+                    new_seq_name,
+                    new_seq_name, seq_name,
+                    'pgmq', tbl.table_name,
+                    'pgmq', tbl.table_name,
+                    new_seq_name,
+                    -- alter seq
+                    new_seq_name, 
+                    tbl.table_name || '_msg_id_seq'
+                );
+            end loop;
+            -- No tables should be owned by the extension.
+            -- We want them to be included in logical backups
+            for tbl in
+                select c.relname as table_name
+                from pg_class c
+                join pg_depend d
+                    on c.oid = d.objid
+                join pg_extension e
+                    on d.refobjid = e.oid
+                where 
+                c.relkind in ('r', 'p', 'u')
+                and e.extname = 'pgmq'
+                and (c.relname like 'q_%' or c.relname like 'a_%')
+            loop
+            execute format('
+                alter extension pgmq drop table pgmq."%s";',
+                tbl.table_name
+            );
+            end loop;
+        end \$\$;
+EOF
+        )
+
+        run_sql -c "$PATCH_PGMQ_QUERY"
+        run_sql -c "update pg_extension set extowner = 'postgres'::regrole where extname = 'pgmq';"
+    fi
+
     run_sql -c "grant pg_read_all_data, pg_signal_backend to postgres"
 }
 
diff --git a/migrations/db/migrations/20241215003910_backfill_pgmq_metadata.sql b/migrations/db/migrations/20241215003910_backfill_pgmq_metadata.sql
new file mode 100644
index 000000000..652be2a1d
--- /dev/null
+++ b/migrations/db/migrations/20241215003910_backfill_pgmq_metadata.sql
@@ -0,0 +1,29 @@
+-- migrate:up
+do $$
+begin
+    -- Check if the pgmq.meta table exists
+    if exists (
+        select 1
+        from pg_catalog.pg_class c
+        join pg_catalog.pg_namespace n on c.relnamespace = n.oid
+        where n.nspname = 'pgmq' and c.relname = 'meta'
+    ) then
+        -- Insert data into pgmq.meta for all tables matching the naming pattern 'pgmq.q_<queue_name>'
+        insert into pgmq.meta (queue_name, is_partitioned, is_unlogged, created_at)
+        select
+            substring(c.relname from 3) as queue_name,
+            false as is_partitioned,
+            case when c.relpersistence = 'u' then true else false end as is_unlogged,
+            now() as created_at
+        from
+			pg_catalog.pg_class c
+        	join pg_catalog.pg_namespace n
+				on c.relnamespace = n.oid
+        where
+			n.nspname = 'pgmq'
+			and c.relname like 'q_%'
+	 		and c.relkind in ('r', 'p', 'u');
+    end if;
+end $$;
+
+-- migrate:down