fix(pgmq): resolve drop_queue overloading bug

ansible/files/postgresql_extension_custom_scripts/pgmq/after-create.sql

@@ -18,25 +18,29 @@ begin
     this update is backwards compatible with version 1.4.4 but should be removed once we're on
     physical backups everywhere
 */
--- Detach and delete the official function
-if extversion = '1.4.4' then
-  alter extension pgmq drop function pgmq.drop_queue;
-  drop function pgmq.drop_queue;
-else -- 1.5.1+
-  alter extension pgmq drop function pgmq.drop_queue(TEXT);
-  drop function pgmq.drop_queue(TEXT);
-end if;
+  -- detach and drop both historical drop_queue signatures (1.4.4 only ever
+  -- has (text, boolean); 1.5.0+ has both (text) and (text, boolean))
+  begin
+    alter extension pgmq drop function pgmq.drop_queue(text);
+  exception when others then null;
+  end;
+  begin
+    alter extension pgmq drop function pgmq.drop_queue(text, boolean);
+  exception when others then null;
+  end;
+
+  drop function if exists pgmq.drop_queue(text);
+  drop function if exists pgmq.drop_queue(text, boolean);
 
 -- Create and reattach the patched function
-CREATE FUNCTION pgmq.drop_queue(queue_name TEXT)
+CREATE FUNCTION pgmq.drop_queue(queue_name TEXT, partitioned BOOLEAN DEFAULT FALSE)
 RETURNS BOOLEAN AS $func$
 DECLARE
     qtable TEXT := pgmq.format_table_name(queue_name, 'q');
     qtable_seq TEXT := qtable || '_msg_id_seq';
     fq_qtable TEXT := 'pgmq.' || qtable;
     atable TEXT := pgmq.format_table_name(queue_name, 'a');
     fq_atable TEXT := 'pgmq.' || atable;
-    partitioned BOOLEAN;
 BEGIN
     EXECUTE FORMAT(
         $QUERY$
@@ -140,11 +144,7 @@ BEGIN
 END;
 $func$ LANGUAGE plpgsql;
 
-if extversion = '1.4.4' then
-  alter extension pgmq add function pgmq.drop_queue;
-else -- 1.5.1+
-  alter extension pgmq add function pgmq.drop_queue(TEXT);
-end if;
+  alter extension pgmq add function pgmq.drop_queue(text, boolean);
 
 
   update pg_extension set extowner = 'postgres'::regrole where extname = 'pgmq';

nix/ext/tests/pgmq-drop-queue-overload.nix

@@ -0,0 +1,69 @@
+{ self, pkgs }:
+let
+  pname = "pgmq";
+  inherit (pkgs) lib;
+  system = pkgs.pkgsLinux.stdenv.hostPlatform.system;
+  testLib = import ./lib.nix { inherit self pkgs; };
+
+  installedExtension = self.legacyPackages.${system}."psql_15".exts."${pname}";
+  # boundary versions only - the fix doesn't branch on extversion, so
+  # oldest/newest is enough to catch a regression
+  versions = lib.unique [
+    (lib.head installedExtension.versions)
+    (lib.last installedExtension.versions)
+  ];
+in
+pkgs.testers.runNixOSTest {
+  name = "pgmq-drop-queue-overload";
+  nodes.server =
+    { ... }:
+    {
+      imports = [
+        (testLib.makeSupabaseTestConfig {
+          majorVersion = "15";
+        })
+      ];
+    };
+  testScript =
+    { ... }:
+    let
+      versionList = lib.concatStringsSep ", " (map (v: ''"${v}"'') versions);
+    in
+    # python
+    ''
+      versions = [${versionList}]
+
+      def sql(query):
+          return server.succeed(
+              "psql -U supabase_admin -d postgres -t -A -c \"" + query.replace('"', '\\"') + "\""
+          ).strip()
+
+      def assert_single_merged_overload(version):
+          ok = sql(
+              "select count(*) = 1 "
+              "  and bool_and(d.objid is not null) "
+              "  and bool_and(pg_get_function_identity_arguments(p.oid) = 'queue_name text, partitioned boolean') "
+              "from pg_proc p "
+              "join pg_depend d on d.objid = p.oid and d.deptype = 'e' "
+              "  and d.refobjid = (select oid from pg_extension where extname = 'pgmq') "
+              "where p.pronamespace = 'pgmq'::regnamespace and p.proname = 'drop_queue';"
+          )
+          assert ok == "t", f"[{version}] expected exactly one merged, extension-owned drop_queue(text, boolean)"
+
+      start_all()
+      server.wait_for_unit("supabase-db-init.service")
+
+      for version in versions:
+          with subtest(f"install pgmq {version}"):
+              server.succeed("psql -U supabase_admin -d postgres -c 'DROP EXTENSION IF EXISTS pgmq;'")
+              server.succeed(
+                  f"psql -U supabase_admin -d postgres -c \"CREATE EXTENSION pgmq WITH VERSION '{version}' CASCADE;\""
+              )
+
+              assert_single_merged_overload(version)
+
+              qname = f"q_{version.replace('.', '_')}"
+              sql(f"select pgmq.create('{qname}_a'); select pgmq.drop_queue('{qname}_a');")
+              sql(f"select pgmq.create('{qname}_b'); select pgmq.drop_queue('{qname}_b', false);")
+    '';
+}

nix/tests/expected/pgmq.out

@@ -175,7 +175,6 @@ order by
  pgmq        | delete                        | postgres
  pgmq        | detach_archive                | postgres
  pgmq        | drop_queue                    | postgres
- pgmq        | drop_queue                    | postgres
  pgmq        | format_table_name             | postgres
  pgmq        | list_queues                   | postgres
  pgmq        | metrics                       | postgres
@@ -198,7 +197,7 @@ order by
  pgmq        | send_batch                    | postgres
  pgmq        | set_vt                        | postgres
  pgmq        | validate_queue_name           | postgres
-(40 rows)
+(39 rows)
 
 -- assert search_path is preserved after after-create script is run
 show search_path;