Merge pull request #1 from time-loop/feature/ENG_OPER-10194_schema_ownership

feat(db-ops): Allows schemas to have ownership set at creation [ENG_OPER-10194]

Author: bricklenclickup

sql/functions/microsharding_ensure_exist.sql

@@ -1,6 +1,9 @@
+DROP FUNCTION IF EXISTS microsharding_ensure_exist(integer, integer);
+
 CREATE OR REPLACE FUNCTION microsharding_ensure_exist(
   shard_from integer,
-  shard_to integer = NULL
+  shard_to integer = NULL,
+  schema_owner varchar(63) DEFAULT NULL
 ) RETURNS SETOF regnamespace
 LANGUAGE plpgsql
 SET search_path FROM CURRENT
@@ -13,6 +16,16 @@ BEGIN
   IF shard_from < 0 OR shard_to > 9999 THEN
     RAISE EXCEPTION 'Invalid shard_from or shard_to';
   END IF;
+
+  IF schema_owner IS NOT NULL
+    AND NOT EXISTS (
+      SELECT 1 FROM pg_catalog.pg_roles WHERE rolname = schema_owner
+      UNION
+      SELECT 1 FROM pg_catalog.pg_user WHERE usename = schema_owner
+    ) THEN
+      RAISE EXCEPTION 'Invalid schema_owner: %', schema_owner;
+  END IF;
+
   FOR shard IN
     WITH shards AS (
       SELECT microsharding_schema_name_(n) AS shard
@@ -22,11 +35,17 @@ BEGIN
     WHERE NOT EXISTS (SELECT 1 FROM pg_catalog.pg_namespace WHERE nspname = shards.shard)
     ORDER BY shards.shard
   LOOP
-    EXECUTE format('CREATE SCHEMA %I', shard);
+
+    IF schema_owner IS NOT NULL THEN
+      EXECUTE format('CREATE SCHEMA IF NOT EXISTS %I AUTHORIZATION %I ', shard, schema_owner);
+    ELSE
+      EXECUTE format('CREATE SCHEMA IF NOT EXISTS %I', shard);
+    END IF;
+
     RETURN NEXT shard::regnamespace;
   END LOOP;
 END;
 $$;
 
-COMMENT ON FUNCTION microsharding_ensure_exist(integer, integer)
+COMMENT ON FUNCTION microsharding_ensure_exist(integer, integer, varchar)
   IS 'Creates shards (schemas) in the range shard_from..shard_to (inclusive).';

sql/pg-microsharding-down.sql

@@ -1,19 +1,20 @@
-DROP FUNCTION microsharding_active_shards_();
-DROP FUNCTION microsharding_advisory_lock_();
-DROP FUNCTION microsharding_debug_fdw_create(text, text[]);
-DROP FUNCTION microsharding_debug_fdw_drop(text);
-DROP FUNCTION microsharding_debug_fdw_schemas_(text);
-DROP FUNCTION microsharding_debug_views_create(text, text);
-DROP FUNCTION microsharding_debug_views_drop(text);
-DROP FUNCTION microsharding_do_on_each(text);
-DROP FUNCTION microsharding_ensure_absent(integer, integer);
-DROP FUNCTION microsharding_ensure_active_shards_(text[]);
-DROP FUNCTION microsharding_ensure_active(integer, integer);
-DROP FUNCTION microsharding_ensure_exist(integer, integer);
-DROP FUNCTION microsharding_ensure_inactive(integer, integer);
-DROP FUNCTION microsharding_fmt_(text, integer);
-DROP FUNCTION microsharding_list_active_shards();
-DROP FUNCTION microsharding_migration_after(text, text, text);
-DROP FUNCTION microsharding_migration_before(text);
-DROP FUNCTION microsharding_notice_locks_(text, timestamptz);
-DROP FUNCTION microsharding_schema_name_(integer);
+DROP FUNCTION IF EXISTS microsharding_active_shards_();
+DROP FUNCTION IF EXISTS microsharding_advisory_lock_();
+DROP FUNCTION IF EXISTS microsharding_debug_fdw_create(text, text[]);
+DROP FUNCTION IF EXISTS microsharding_debug_fdw_drop(text);
+DROP FUNCTION IF EXISTS microsharding_debug_fdw_schemas_(text);
+DROP FUNCTION IF EXISTS microsharding_debug_views_create(text, text);
+DROP FUNCTION IF EXISTS microsharding_debug_views_drop(text);
+DROP FUNCTION IF EXISTS microsharding_do_on_each(text);
+DROP FUNCTION IF EXISTS microsharding_ensure_absent(integer, integer);
+DROP FUNCTION IF EXISTS microsharding_ensure_active_shards_(text[]);
+DROP FUNCTION IF EXISTS microsharding_ensure_active(integer, integer);
+DROP FUNCTION IF EXISTS microsharding_ensure_exist(integer, integer);
+DROP FUNCTION IF EXISTS microsharding_ensure_exist(integer, integer, varchar);
+DROP FUNCTION IF EXISTS microsharding_ensure_inactive(integer, integer);
+DROP FUNCTION IF EXISTS microsharding_fmt_(text, integer);
+DROP FUNCTION IF EXISTS microsharding_list_active_shards();
+DROP FUNCTION IF EXISTS microsharding_migration_after(text, text, text);
+DROP FUNCTION IF EXISTS microsharding_migration_before(text);
+DROP FUNCTION IF EXISTS microsharding_notice_locks_(text, timestamptz);
+DROP FUNCTION IF EXISTS microsharding_schema_name_(integer);