feat: add role-based permissions for secrets

projects/extension/ai/secrets.py

@@ -20,6 +20,32 @@ def get_guc_value(plpy, setting: str, default: str) -> str:
     return val
 
 
+def check_secret_permissions(plpy, secret_name: str) -> bool:
+    # check if the user has access to all secrets
+    plan = plpy.prepare(
+        """
+                        SELECT 1
+                        FROM ai.secret_permissions 
+                        WHERE name = '*'""",
+        [],
+    )
+    result = plan.execute([], 1)
+    if len(result) > 0:
+        return True
+
+    # check if the user has access to the specific secret
+    plan = plpy.prepare(
+        """
+                        SELECT 1
+                        FROM ai.secret_permissions 
+                        WHERE name = $1 
+                        """,
+        ["text"],
+    )
+    result = plan.execute([secret_name], 1)
+    return len(result) > 0
+
+
 def resolve_secret(plpy, secret_name: str) -> str:
     # first try the guc, then the secrets manager, then error
     secret_name_lower = secret_name.lower()
@@ -45,6 +71,10 @@ def reveal_secret(plpy, secret_name: str) -> str | None:
         plpy.error("secrets manager is not enabled")
         return None
 
+    if not check_secret_permissions(plpy, secret_name):
+        plpy.error(f"user does not have access to secret '{secret_name}'")
+        return None
+
     the_url = urljoin(
         get_guc_value(plpy, GUC_SECRETS_MANAGER_URL, ""),
         DEFAULT_SECRETS_MANAGER_PATH,

projects/extension/sql/ai--0.4.0.sql

@@ -114,6 +114,43 @@ end;
 $outer_migration_block$;
 
 
+-------------------------------------------------------------------------------
+-- 002-secret_permissions.sql
+do $outer_migration_block$ /*002-secret_permissions.sql*/
+declare
+    _sql text;
+    _migration record;
+    _migration_name text = $migration_name$002-secret_permissions.sql$migration_name$;
+    _migration_body text =
+$migration_body$
+create table ai._secret_permissions
+( 
+  name text not null check(name = '*' or name ~ '^[A-Za-z0-9_.]+$')
+, "role" text not null
+, primary key (name, "role")
+);
+perform pg_catalog.pg_extension_config_dump('ai._secret_permissions'::pg_catalog.regclass, '');
+--only admins will have access to this table
+revoke all on ai._secret_permissions from public;
+
+$migration_body$;
+begin
+    select * into _migration from ai.migration where "name" operator(pg_catalog.=) _migration_name;
+    if _migration is not null then
+        raise notice 'migration %s already applied. skipping.', _migration_name;
+        if _migration.body operator(pg_catalog.!=) _migration_body then
+            raise warning 'the contents of migration "%s" have changed', _migration_name;
+        end if;
+        return;
+    end if;
+    _sql = pg_catalog.format(E'do /*%s*/ $migration_body$\nbegin\n%s\nend;\n$migration_body$;', _migration_name, _migration_body);
+    execute _sql;
+    insert into ai.migration ("name", body, applied_at_version)
+    values (_migration_name, _migration_body, $version$0.4.0$version$);
+end;
+$outer_migration_block$;
+
+
 --------------------------------------------------------------------------------
 -- 001-openai.sql
 
@@ -3219,6 +3256,24 @@ language plpython3u volatile security invoker
 set search_path to pg_catalog, pg_temp;
 
 
+create or replace view ai.secret_permissions as
+SELECT * 
+FROM ai._secret_permissions
+WHERE to_regrole("role") is not null AND pg_has_role(current_user, "role", 'member');
+
+create or replace function ai.grant_secret(secret_name text, grant_to_role text) returns void
+as $func$
+    insert into ai._secret_permissions (name, "role") VALUES (secret_name, grant_to_role);
+$func$ language sql volatile security invoker
+set search_path to pg_catalog, pg_temp;
+
+create or replace function ai.revoke_secret(secret_name text, revoke_from_role text) returns void
+as $func$
+    delete from ai._secret_permissions where name = secret_name and "role" = revoke_from_role;
+$func$ language sql volatile security invoker
+set search_path to pg_catalog, pg_temp;
+
+
 
 --------------------------------------------------------------------------------
 -- 999-privileges.sql
@@ -3275,7 +3330,8 @@ begin
         and k.relkind in ('r', 'p', 'S', 'v') -- tables, sequences, and views
         and (admin, n.nspname, k.relname) not in
         (
-            (false, 'ai', 'migration') -- only admins get any access to this table
+            (false, 'ai', 'migration'), -- only admins get any access to this table
+            (false, 'ai', '_secret_permissions') -- only admins get any access to this table
         )
         order by n.nspname, k.relname
     )
@@ -3309,14 +3365,20 @@ begin
         and e.extname operator(pg_catalog.=) 'ai'
         and k.prokind in ('f', 'p')
         and case
-              when k.proname operator(pg_catalog.=) 'grant_ai_usage' then admin -- only admins get this function
+              when k.proname in ('grant_ai_usage', 'grant_secret', 'revoke_secret') then admin -- only admins get these function
               else true
             end
     )
     loop
         raise debug '%', _sql;
         execute _sql;
     end loop;
+
+    -- secret permissions
+    if admin then
+        -- grant access to all secrets to admin users
+        insert into ai.secret_permissions (name, "role") VALUES ('*', to_user);
+    end if;
 end
 $func$ language plpgsql volatile
 security invoker -- gotta have privs to give privs
@@ -3396,4 +3458,3 @@ $func$;
 
 
 
-

projects/extension/sql/idempotent/014-secrets.sql

@@ -8,3 +8,21 @@ as $python$
 $python$
 language plpython3u volatile security invoker
 set search_path to pg_catalog, pg_temp;
+
+
+create or replace view ai.secret_permissions as
+SELECT * 
+FROM ai._secret_permissions
+WHERE to_regrole("role") is not null AND pg_has_role(current_user, "role", 'member');
+
+create or replace function ai.grant_secret(secret_name text, grant_to_role text) returns void
+as $func$
+    insert into ai._secret_permissions (name, "role") VALUES (secret_name, grant_to_role);
+$func$ language sql volatile security invoker
+set search_path to pg_catalog, pg_temp;
+
+create or replace function ai.revoke_secret(secret_name text, revoke_from_role text) returns void
+as $func$
+    delete from ai._secret_permissions where name = secret_name and "role" = revoke_from_role;
+$func$ language sql volatile security invoker
+set search_path to pg_catalog, pg_temp;

projects/extension/sql/idempotent/999-privileges.sql

@@ -51,7 +51,8 @@ begin
         and k.relkind in ('r', 'p', 'S', 'v') -- tables, sequences, and views
         and (admin, n.nspname, k.relname) not in
         (
-            (false, 'ai', 'migration') -- only admins get any access to this table
+            (false, 'ai', 'migration'), -- only admins get any access to this table
+            (false, 'ai', '_secret_permissions') -- only admins get any access to this table
         )
         order by n.nspname, k.relname
     )
@@ -85,14 +86,20 @@ begin
         and e.extname operator(pg_catalog.=) 'ai'
         and k.prokind in ('f', 'p')
         and case
-              when k.proname operator(pg_catalog.=) 'grant_ai_usage' then admin -- only admins get this function
+              when k.proname in ('grant_ai_usage', 'grant_secret', 'revoke_secret') then admin -- only admins get these function
               else true
             end
     )
     loop
         raise debug '%', _sql;
         execute _sql;
     end loop;
+
+    -- secret permissions
+    if admin then
+        -- grant access to all secrets to admin users
+        insert into ai.secret_permissions (name, "role") VALUES ('*', to_user);
+    end if;
 end
 $func$ language plpgsql volatile
 security invoker -- gotta have privs to give privs
@@ -169,4 +176,3 @@ begin
     end loop;
 end
 $func$;
-

projects/extension/sql/incremental/002-secret_permissions.sql

@@ -0,0 +1,9 @@
+create table ai._secret_permissions
+( 
+  name text not null check(name = '*' or name ~ '^[A-Za-z0-9_.]+$')
+, "role" text not null
+, primary key (name, "role")
+);
+perform pg_catalog.pg_extension_config_dump('ai._secret_permissions'::pg_catalog.regclass, '');
+--only admins will have access to this table
+revoke all on ai._secret_permissions from public;

projects/extension/tests/contents/output.expected

@@ -24,6 +24,7 @@ CREATE EXTENSION
  function ai.execute_vectorizer(integer)
  function ai.formatting_python_template(text)
  function ai.grant_ai_usage(name,boolean)
+ function ai.grant_secret(text,text)
  function ai.grant_to()
  function ai.grant_to(name[])
  function ai.indexing_default()
@@ -48,6 +49,7 @@ CREATE EXTENSION
  function ai._resolve_indexing_default()
  function ai._resolve_scheduling_default()
  function ai.reveal_secret(text)
+ function ai.revoke_secret(text,text)
  function ai.scheduling_default()
  function ai.scheduling_none()
  function ai.scheduling_timescaledb(interval,timestamp with time zone,boolean,text)
@@ -75,10 +77,30 @@ CREATE EXTENSION
  function ai._vectorizer_vector_index_exists(name,name,jsonb)
  sequence ai.vectorizer_id_seq
  table ai.migration
+ table ai._secret_permissions
  table ai.vectorizer
  table ai.vectorizer_errors
+ view ai.secret_permissions
  view ai.vectorizer_status
-(74 rows)
+(78 rows)
+
+                                    Table "ai._secret_permissions"
+ Column | Type | Collation | Nullable | Default | Storage  | Compression | Stats target | Description 
+--------+------+-----------+----------+---------+----------+-------------+--------------+-------------
+ name   | text |           | not null |         | extended |             |              | 
+ role   | text |           | not null |         | extended |             |              | 
+Indexes:
+    "_secret_permissions_pkey" PRIMARY KEY, btree (name, role)
+Check constraints:
+    "_secret_permissions_name_check" CHECK (name = '*'::text OR name ~ '^[A-Za-z0-9_.]+$'::text)
+Access method: heap
+
+             Index "ai._secret_permissions_pkey"
+ Column | Type | Key? | Definition | Storage  | Stats target 
+--------+------+------+------------+----------+--------------
+ name   | text | yes  | name       | extended | 
+ role   | text | yes  | role       | extended | 
+primary key, btree, for table "ai._secret_permissions"
 
                                                               Table "ai.migration"
        Column       |           Type           | Collation | Nullable |      Default      | Storage  | Compression | Stats target | Description 
@@ -97,6 +119,17 @@ Access method: heap
  name   | text | yes  | name       | extended | 
 primary key, btree, for table "ai.migration"
 
+                      View "ai.secret_permissions"
+ Column | Type | Collation | Nullable | Default | Storage  | Description 
+--------+------+-----------+----------+---------+----------+-------------
+ name   | text |           |          |         | extended | 
+ role   | text |           |          |         | extended | 
+View definition:
+ SELECT name,
+    role
+   FROM ai._secret_permissions
+  WHERE to_regrole(role) IS NOT NULL AND pg_has_role(CURRENT_USER, role::name, 'member'::text);
+
                                                           Table "ai.vectorizer"
     Column     |  Type   | Collation | Nullable |             Default              | Storage  | Compression | Stats target | Description 
 ---------------+---------+-----------+----------+----------------------------------+----------+-------------+--------------+-------------

projects/extension/tests/privileges/function.expected

@@ -176,6 +176,10 @@
  f       | bob   | execute   | no      | ai     | grant_ai_usage(to_user name, admin boolean)
  f       | fred  | execute   | no      | ai     | grant_ai_usage(to_user name, admin boolean)
  f       | jill  | execute   | no      | ai     | grant_ai_usage(to_user name, admin boolean)
+ f       | alice | execute   | YES     | ai     | grant_secret(secret_name text, grant_to_role text)
+ f       | bob   | execute   | no      | ai     | grant_secret(secret_name text, grant_to_role text)
+ f       | fred  | execute   | no      | ai     | grant_secret(secret_name text, grant_to_role text)
+ f       | jill  | execute   | no      | ai     | grant_secret(secret_name text, grant_to_role text)
  f       | alice | execute   | YES     | ai     | grant_to()
  f       | bob   | execute   | no      | ai     | grant_to()
  f       | fred  | execute   | no      | ai     | grant_to()
@@ -264,6 +268,10 @@
  f       | bob   | execute   | no      | ai     | reveal_secret(secret_name text)
  f       | fred  | execute   | no      | ai     | reveal_secret(secret_name text)
  f       | jill  | execute   | YES     | ai     | reveal_secret(secret_name text)
+ f       | alice | execute   | YES     | ai     | revoke_secret(secret_name text, revoke_from_role text)
+ f       | bob   | execute   | no      | ai     | revoke_secret(secret_name text, revoke_from_role text)
+ f       | fred  | execute   | no      | ai     | revoke_secret(secret_name text, revoke_from_role text)
+ f       | jill  | execute   | no      | ai     | revoke_secret(secret_name text, revoke_from_role text)
  f       | alice | execute   | YES     | ai     | scheduling_default()
  f       | bob   | execute   | no      | ai     | scheduling_default()
  f       | fred  | execute   | no      | ai     | scheduling_default()
@@ -280,5 +288,5 @@
  f       | bob   | execute   | no      | ai     | vectorizer_queue_pending(vectorizer_id integer)
  f       | fred  | execute   | no      | ai     | vectorizer_queue_pending(vectorizer_id integer)
  f       | jill  | execute   | YES     | ai     | vectorizer_queue_pending(vectorizer_id integer)
-(280 rows)
+(288 rows)
 

projects/extension/tests/privileges/table.expected

@@ -1,5 +1,21 @@
  schema |        table         | user  | privilege | granted 
 --------+----------------------+-------+-----------+---------
+ ai     | _secret_permissions  | alice | delete    | YES
+ ai     | _secret_permissions  | alice | insert    | YES
+ ai     | _secret_permissions  | alice | select    | YES
+ ai     | _secret_permissions  | alice | update    | YES
+ ai     | _secret_permissions  | bob   | delete    | no
+ ai     | _secret_permissions  | bob   | insert    | no
+ ai     | _secret_permissions  | bob   | select    | no
+ ai     | _secret_permissions  | bob   | update    | no
+ ai     | _secret_permissions  | fred  | delete    | no
+ ai     | _secret_permissions  | fred  | insert    | no
+ ai     | _secret_permissions  | fred  | select    | no
+ ai     | _secret_permissions  | fred  | update    | no
+ ai     | _secret_permissions  | jill  | delete    | no
+ ai     | _secret_permissions  | jill  | insert    | no
+ ai     | _secret_permissions  | jill  | select    | no
+ ai     | _secret_permissions  | jill  | update    | no
  ai     | _vectorizer_q_1      | alice | delete    | YES
  ai     | _vectorizer_q_1      | alice | insert    | YES
  ai     | _vectorizer_q_1      | alice | select    | YES
@@ -96,5 +112,5 @@
  wiki   | post_embedding_store | jill  | insert    | YES
  wiki   | post_embedding_store | jill  | select    | YES
  wiki   | post_embedding_store | jill  | update    | YES
-(96 rows)
+(112 rows)