From bae232b2a7d0643d4e91928f3cd0f72d6312bc42 Mon Sep 17 00:00:00 2001
From: Bobbie Soedirgo <bobbie@soedirgo.dev>
Date: Wed, 19 Feb 2025 21:16:34 +0800
Subject: [PATCH] fix: update migration to support vault 0.2.8 and above

---
 ...221207154255_create_pgsodium_and_vault.sql | 60 +++++++++++--------
 ...0529180330_alter_api_roles_for_inherit.sql |  7 ++-
 .../20250218031949_pgsodium_mask_role.sql     | 44 ++++++++------
 3 files changed, 66 insertions(+), 45 deletions(-)

diff --git a/migrations/db/migrations/20221207154255_create_pgsodium_and_vault.sql b/migrations/db/migrations/20221207154255_create_pgsodium_and_vault.sql
index 3d3867ab5..f93cc211a 100644
--- a/migrations/db/migrations/20221207154255_create_pgsodium_and_vault.sql
+++ b/migrations/db/migrations/20221207154255_create_pgsodium_and_vault.sql
@@ -5,34 +5,44 @@ DECLARE
   pgsodium_exists boolean;
   vault_exists boolean;
 BEGIN
-  pgsodium_exists = (
-    select count(*) = 1 
-    from pg_available_extensions 
-    where name = 'pgsodium'
-    and default_version in ('3.1.6', '3.1.7', '3.1.8', '3.1.9')
-  );
-  
-  vault_exists = (
+  IF EXISTS (SELECT FROM pg_available_extensions WHERE name = 'supabase_vault' AND default_version != '0.2.8') THEN
+    CREATE EXTENSION IF NOT EXISTS supabase_vault;
+
+    -- for some reason extension custom scripts aren't run during AMI build, so
+    -- we manually run it here
+    GRANT USAGE ON SCHEMA vault TO postgres WITH GRANT OPTION;
+    GRANT SELECT, DELETE ON vault.secrets, vault.decrypted_secrets TO postgres WITH GRANT OPTION;
+    GRANT EXECUTE ON FUNCTION vault.create_secret, vault.update_secret, vault._crypto_aead_det_decrypt TO postgres WITH GRANT OPTION;
+  ELSE
+    pgsodium_exists = (
       select count(*) = 1 
       from pg_available_extensions 
-      where name = 'supabase_vault'
-  );
-
-  IF pgsodium_exists 
-  THEN
-    create extension if not exists pgsodium;
-
-    grant pgsodium_keyiduser to postgres with admin option;
-    grant pgsodium_keyholder to postgres with admin option;
-    grant pgsodium_keymaker  to postgres with admin option;
-
-    grant execute on function pgsodium.crypto_aead_det_decrypt(bytea, bytea, uuid, bytea) to service_role;
-    grant execute on function pgsodium.crypto_aead_det_encrypt(bytea, bytea, uuid, bytea) to service_role;
-    grant execute on function pgsodium.crypto_aead_det_keygen to service_role;
-
-    IF vault_exists
+      where name = 'pgsodium'
+      and default_version in ('3.1.6', '3.1.7', '3.1.8', '3.1.9')
+    );
+    
+    vault_exists = (
+        select count(*) = 1 
+        from pg_available_extensions 
+        where name = 'supabase_vault'
+    );
+  
+    IF pgsodium_exists 
     THEN
-      create extension if not exists supabase_vault;
+      create extension if not exists pgsodium;
+  
+      grant pgsodium_keyiduser to postgres with admin option;
+      grant pgsodium_keyholder to postgres with admin option;
+      grant pgsodium_keymaker  to postgres with admin option;
+  
+      grant execute on function pgsodium.crypto_aead_det_decrypt(bytea, bytea, uuid, bytea) to service_role;
+      grant execute on function pgsodium.crypto_aead_det_encrypt(bytea, bytea, uuid, bytea) to service_role;
+      grant execute on function pgsodium.crypto_aead_det_keygen to service_role;
+  
+      IF vault_exists
+      THEN
+        create extension if not exists supabase_vault;
+      END IF;
     END IF;
   END IF;
 END $$;
diff --git a/migrations/db/migrations/20230529180330_alter_api_roles_for_inherit.sql b/migrations/db/migrations/20230529180330_alter_api_roles_for_inherit.sql
index 4df82e3f4..d8291bbaa 100644
--- a/migrations/db/migrations/20230529180330_alter_api_roles_for_inherit.sql
+++ b/migrations/db/migrations/20230529180330_alter_api_roles_for_inherit.sql
@@ -4,7 +4,12 @@ ALTER ROLE authenticated inherit;
 ALTER ROLE anon inherit;
 ALTER ROLE service_role inherit;
 
-GRANT pgsodium_keyholder to service_role;
+DO $$
+BEGIN
+  IF EXISTS (SELECT FROM pg_roles WHERE rolname = 'pgsodium_keyholder') THEN
+    GRANT pgsodium_keyholder to service_role;
+  END IF;
+END $$;
 
 -- migrate:down
 
diff --git a/migrations/db/migrations/20250218031949_pgsodium_mask_role.sql b/migrations/db/migrations/20250218031949_pgsodium_mask_role.sql
index c4a3eadf4..f44fa9866 100644
--- a/migrations/db/migrations/20250218031949_pgsodium_mask_role.sql
+++ b/migrations/db/migrations/20250218031949_pgsodium_mask_role.sql
@@ -1,25 +1,31 @@
 -- migrate:up
-CREATE OR REPLACE FUNCTION pgsodium.mask_role(masked_role regrole, source_name text, view_name text)
-RETURNS void
-LANGUAGE plpgsql
-SECURITY DEFINER
-SET search_path TO ''
-AS $function$
+
+DO $$
 BEGIN
-  EXECUTE format(
-    'GRANT SELECT ON pgsodium.key TO %s',
-    masked_role);
+  IF EXISTS (SELECT FROM pg_extension WHERE extname = 'pgsodium') THEN
+    CREATE OR REPLACE FUNCTION pgsodium.mask_role(masked_role regrole, source_name text, view_name text)
+    RETURNS void
+    LANGUAGE plpgsql
+    SECURITY DEFINER
+    SET search_path TO ''
+    AS $function$
+    BEGIN
+      EXECUTE format(
+        'GRANT SELECT ON pgsodium.key TO %s',
+        masked_role);
 
-  EXECUTE format(
-    'GRANT pgsodium_keyiduser, pgsodium_keyholder TO %s',
-    masked_role);
+      EXECUTE format(
+        'GRANT pgsodium_keyiduser, pgsodium_keyholder TO %s',
+        masked_role);
 
-  EXECUTE format(
-    'GRANT ALL ON %I TO %s',
-    view_name,
-    masked_role);
-  RETURN;
-END
-$function$;
+      EXECUTE format(
+        'GRANT ALL ON %I TO %s',
+        view_name,
+        masked_role);
+      RETURN;
+    END
+    $function$;
+  END IF;
+END $$;
 
 -- migrate:down