Add database_privileges.replace_for_roles RPC function.

db/roles/operations/update.py

@@ -0,0 +1,8 @@
+import json
+from db.connection import exec_msar_func
+
+
+def replace_database_privileges_for_roles(conn, privileges):
+    return exec_msar_func(
+        conn, 'replace_database_privileges_for_roles', json.dumps(privileges)
+    ).fetchone()[0]

db/sql/00_msar.sql

@@ -497,6 +497,28 @@ END;
 $$ LANGUAGE plpgsql RETURNS NULL ON NULL INPUT;
 
 
+CREATE OR REPLACE FUNCTION msar.get_role_name(rol_oid oid) RETURNS TEXT AS $$/*
+Return the UNQUOTED name of a given role.
+
+If the role does not exist, an exception will be raised.
+
+Args:
+  rol_oid: The OID of the role.
+*/
+DECLARE rol_name text;
+BEGIN
+  SELECT rolname INTO rol_name FROM pg_roles WHERE oid=rol_oid;
+
+  IF rol_name IS NULL THEN
+    RAISE EXCEPTION 'Role with OID % does not exist', rol_oid
+    USING ERRCODE = '42704'; -- undefined_object
+  END IF;
+
+  RETURN rol_name;
+END;
+$$ LANGUAGE plpgsql RETURNS NULL ON NULL INPUT;
+
+
 CREATE OR REPLACE FUNCTION msar.get_constraint_type_api_code(contype char) RETURNS TEXT AS $$/*
 This function returns a string that represents the constraint type code used to describe
 constraints when listing them within the Mathesar API.
@@ -897,6 +919,7 @@ FROM (
 $$ LANGUAGE SQL;
 
 
+DROP FUNCTION IF EXISTS msar.role_info_table();
 CREATE OR REPLACE FUNCTION msar.list_schema_privileges(sch_id regnamespace) RETURNS jsonb AS $$/*
 Given a schema, returns a json array of objects with direct, non-default schema privileges
 
@@ -925,7 +948,7 @@ $$ LANGUAGE SQL STABLE RETURNS NULL ON NULL INPUT;
 
 CREATE OR REPLACE FUNCTION msar.role_info_table() RETURNS TABLE
 (
-  oid oid, -- The OID of the role.
+  oid bigint, -- The OID of the role.
   name name, -- Name of the role.
   super boolean, -- Whether the role has SUPERUSER status.
   inherits boolean, -- Whether the role has INHERIT attribute.
@@ -1026,7 +1049,7 @@ Each returned JSON object in the array has the form:
 WITH priv_cte AS (
   SELECT
     jsonb_build_object(
-      'role_oid', pgr.oid,
+      'role_oid', pgr.oid::bigint,
       'direct',  jsonb_agg(acl.privilege_type)
     ) AS p
   FROM
@@ -1051,7 +1074,7 @@ The returned JSON object has the form:
   }
 */
 SELECT jsonb_build_object(
-  'owner_oid', pgd.datdba,
+  'owner_oid', pgd.datdba::bigint,
   'current_role_db_priv', array_remove(
     ARRAY[
       CASE WHEN has_database_privilege(pgd.oid, 'CREATE') THEN 'CREATE' END,
@@ -1136,6 +1159,50 @@ END;
 $$ LANGUAGE plpgsql;
 
 
+CREATE OR REPLACE FUNCTION
+msar.build_database_privilege_replace_expr(rol_id regrole, privileges_ jsonb) RETURNS TEXT AS $$
+SELECT string_agg(
+  format(
+    concat(
+      CASE WHEN privileges_ ? val THEN 'GRANT' ELSE 'REVOKE' END,
+      ' %1$s ON DATABASE %2$I ',
+      CASE WHEN privileges_ ? val THEN 'TO' ELSE 'FROM' END,
+      ' %3$I'
+    ),
+    val,
+    current_database(),
+    msar.get_role_name(rol_id)
+  ),
+  E';\n'
+) || E';\n'
+FROM unnest(ARRAY['CONNECT', 'CREATE', 'TEMPORARY']) as x(val);
+$$ LANGUAGE SQL STABLE RETURNS NULL ON NULL INPUT;
+
+
+CREATE OR REPLACE FUNCTION
+msar.replace_database_privileges_for_roles(priv_spec jsonb) RETURNS jsonb AS $$/*
+Grant/Revoke privileges for a set of roles on the current database.
+
+Args:
+  priv_spec: An array defining the privileges to grant or revoke for each role.
+
+Each object in the priv_spec should have the form:
+{role_oid: <int>, privileges: SET<"CONNECT"|"CREATE"|"TEMPORARY">}
+
+Any privilege that exists in the privileges subarray will be granted. Any which is missing will be
+revoked.
+*/
+BEGIN
+EXECUTE string_agg(
+  msar.build_database_privilege_replace_expr(role_oid, direct),
+  E';\n'
+) || ';'
+FROM jsonb_to_recordset(priv_spec) AS x(role_oid regrole, direct jsonb);
+RETURN msar.list_db_priv(current_database());
+END;
+$$ LANGUAGE plpgsql RETURNS NULL ON NULL INPUT;
+
+
 ----------------------------------------------------------------------------------------------------
 ----------------------------------------------------------------------------------------------------
 -- ALTER SCHEMA FUNCTIONS

db/sql/test_00_msar.sql

@@ -4422,3 +4422,126 @@ BEGIN
   );
 END;
 $$ LANGUAGE plpgsql;
+
+
+-- msar.replace_database_privileges_for_roles ------------------------------------------------------
+
+CREATE OR REPLACE FUNCTION
+test_replace_database_privileges_for_roles_basic() RETURNS SETOF TEXT AS $$/*
+Happy path, smoke test.
+*/
+DECLARE
+  alice_id oid;
+  bob_id oid;
+BEGIN
+  CREATE ROLE "Alice";
+  CREATE ROLE bob;
+  alice_id := '"Alice"'::regrole::oid;
+  bob_id := 'bob'::regrole::oid;
+
+  RETURN NEXT set_eq(
+    format(
+      $t1$SELECT jsonb_array_elements_text(direct) FROM jsonb_to_recordset(
+        msar.replace_database_privileges_for_roles(jsonb_build_array(jsonb_build_object(
+          'role_oid', %1$s, 'direct', jsonb_build_array('CONNECT', 'CREATE')))))
+        AS x(direct jsonb, role_oid regrole)
+      WHERE role_oid=%1$s $t1$,
+      alice_id
+    ),
+    ARRAY['CONNECT', 'CREATE'],
+    'Response should contain updated role info in correct form'
+  );
+  RETURN NEXT set_eq(
+    concat(
+      'SELECT privilege_type FROM pg_database, LATERAL aclexplode(pg_database.datacl) acl',
+      format(' WHERE acl.grantee=%s;', alice_id)
+    ),
+    ARRAY['CONNECT', 'CREATE'],
+    'Privileges should be updated for actual database properly'
+  );
+  RETURN NEXT set_eq(
+    format(
+      $t2$SELECT jsonb_array_elements_text(direct) FROM jsonb_to_recordset(
+        msar.replace_database_privileges_for_roles(jsonb_build_array(jsonb_build_object(
+              'role_oid', %1$s, 'direct', jsonb_build_array('CONNECT')))))
+        AS x(direct jsonb, role_oid regrole)
+      WHERE role_oid=%1$s $t2$,
+      bob_id
+    ),
+    ARRAY['CONNECT'],
+    'Response should contain updated role info in correct form'
+  );
+  RETURN NEXT set_eq(
+    concat(
+      'SELECT privilege_type FROM pg_database, LATERAL aclexplode(pg_database.datacl) acl',
+      format(' WHERE acl.grantee=%s;', alice_id)
+    ),
+    ARRAY['CONNECT', 'CREATE'],
+    'Alice''s privileges should be left alone properly'
+  );
+  RETURN NEXT set_eq(
+    concat(
+      'SELECT privilege_type FROM pg_database, LATERAL aclexplode(pg_database.datacl) acl',
+      format(' WHERE acl.grantee=%s;', bob_id)
+    ),
+    ARRAY['CONNECT'],
+    'Privileges should be updated for actual database properly'
+  );
+END;
+$$ LANGUAGE plpgsql;
+
+
+CREATE OR REPLACE FUNCTION __setup_alice_and_bob_preloaded() RETURNS SETOF TEXT AS $$
+BEGIN
+  CREATE ROLE "Alice";
+  GRANT CONNECT, CREATE ON DATABASE mathesar_testing TO "Alice";
+  CREATE ROLE bob;
+  GRANT CONNECT ON DATABASE mathesar_testing TO bob;
+END;
+$$ LANGUAGE plpgsql;
+
+
+CREATE OR REPLACE FUNCTION
+test_replace_database_privileges_for_roles_multi_ops() RETURNS SETOF TEXT AS $$/*
+Test that we can add/revoke multiple privileges to/from multiple roles simultaneously.
+*/
+DECLARE
+  alice_id oid;
+  bob_id oid;
+BEGIN
+  PERFORM __setup_alice_and_bob_preloaded();
+  alice_id := '"Alice"'::regrole::oid;
+  bob_id := 'bob'::regrole::oid;
+  RETURN NEXT set_eq(
+    -- Revoke CREATE from Alice, Grant CREATE to Bob.
+    format(
+      $t1$SELECT jsonb_array_elements_text(direct) FROM jsonb_to_recordset(
+        msar.replace_database_privileges_for_roles(jsonb_build_array(
+          jsonb_build_object('role_oid', %1$s, 'direct', jsonb_build_array('CONNECT')),
+          jsonb_build_object('role_oid', %2$s, 'direct', jsonb_build_array('CONNECT', 'CREATE')))))
+        AS x(direct jsonb, role_oid regrole)
+      WHERE role_oid=%1$s $t1$,
+      alice_id,
+      bob_id
+    ),
+    ARRAY['CONNECT'], -- This only checks form of Alice's info in response.
+    'Response should contain updated role info in correct form'
+  );
+  RETURN NEXT set_eq(
+    concat(
+      'SELECT privilege_type FROM pg_database, LATERAL aclexplode(pg_database.datacl) acl',
+      format(' WHERE acl.grantee=%s;', alice_id)
+    ),
+    ARRAY['CONNECT'],
+    'Alice''s privileges should be updated for actual database properly'
+  );
+  RETURN NEXT set_eq(
+    concat(
+      'SELECT privilege_type FROM pg_database, LATERAL aclexplode(pg_database.datacl) acl',
+      format(' WHERE acl.grantee=%s;', bob_id)
+    ),
+    ARRAY['CONNECT', 'CREATE'],
+    'Bob''s privileges should be updated for actual database properly'
+  );
+END;
+$$ LANGUAGE plpgsql;

db/tests/roles/operations/test_update.py

@@ -0,0 +1,14 @@
+import json
+from unittest.mock import patch
+from db.roles.operations import update as rupdate
+
+
+def test_replace_database_privileges_for_roles():
+    priv_spec = [{"role_oid": 1234, "privileges": ["CONNECT", "CREATE"]}]
+    with patch.object(rupdate, 'exec_msar_func') as mock_exec:
+        mock_exec.return_value.fetchone = lambda: ('a', 'b')
+        result = rupdate.replace_database_privileges_for_roles('conn', priv_spec)
+    mock_exec.assert_called_once_with(
+        'conn', 'replace_database_privileges_for_roles', json.dumps(priv_spec)
+    )
+    assert result == 'a'

docs/docs/api/rpc.md

@@ -84,6 +84,7 @@ To use an RPC function:
       members:
       - list_direct
       - get_owner_oid_and_curr_role_db_priv
+      - replace_for_roles
       - DBPrivileges
       - CurrentDBPrivileges
 

mathesar/rpc/database_privileges.py

@@ -1,9 +1,10 @@
-from typing import TypedDict
+from typing import Literal, TypedDict
 
 from modernrpc.core import rpc_method, REQUEST_KEY
 from modernrpc.auth.basic import http_basic_auth_login_required
 
 from db.roles.operations.select import list_db_priv, get_curr_role_db_priv
+from db.roles.operations.update import replace_database_privileges_for_roles
 from mathesar.rpc.utils import connect
 from mathesar.models.base import Database
 from mathesar.rpc.exceptions.handlers import handle_rpc_exceptions
@@ -18,7 +19,7 @@ class DBPrivileges(TypedDict):
         direct: A list of database privileges for the afforementioned role_oid.
     """
     role_oid: int
-    direct: list[str]
+    direct: list[Literal['CONNECT', 'CREATE', 'TEMPORARY']]
 
     @classmethod
     def from_dict(cls, d):
@@ -86,3 +87,36 @@ def get_owner_oid_and_curr_role_db_priv(*, database_id: int, **kwargs) -> Curren
         db_name = Database.objects.get(id=database_id).name
         curr_role_db_priv = get_curr_role_db_priv(db_name, conn)
     return CurrentDBPrivileges.from_dict(curr_role_db_priv)
+
+
+@rpc_method(name="database_privileges.replace_for_roles")
+@http_basic_auth_login_required
+@handle_rpc_exceptions
+def replace_for_roles(
+        *, privileges: list[DBPrivileges], database_id: int, **kwargs
+) -> list[DBPrivileges]:
+    """
+    Replace direct database privileges for roles.
+
+    Possible privileges are `CONNECT`, `CREATE`, and `TEMPORARY`.
+
+    Only roles which are included in a passed `DBPrivileges` object are
+    affected.
+
+    WARNING: Any privilege included in the `direct` list for a role
+    is GRANTed, and any privilege not included is REVOKEd.
+
+    Args:
+        privileges: The new privilege sets for roles.
+        database_id: The Django id of the database.
+
+    Returns:
+        A list of all non-default privileges on the database after the
+        operation.
+    """
+    user = kwargs.get(REQUEST_KEY).user
+    with connect(database_id, user) as conn:
+        raw_db_priv = replace_database_privileges_for_roles(
+            conn, [DBPrivileges.from_dict(i) for i in privileges]
+        )
+    return [DBPrivileges.from_dict(i) for i in raw_db_priv]

mathesar/tests/rpc/test_database_privileges.py

@@ -0,0 +1,53 @@
+"""
+This file tests the database_privileges RPC functions.
+
+Fixtures:
+    rf(pytest-django): Provides mocked `Request` objects.
+    monkeypatch(pytest): Lets you monkeypatch an object for testing.
+"""
+from contextlib import contextmanager
+
+from mathesar.rpc import database_privileges
+from mathesar.models.users import User
+
+
+def test_database_privileges_set_for_roles(rf, monkeypatch):
+    _username = 'alice'
+    _password = 'pass1234'
+    _database_id = 2
+    _privileges = [{"role_oid": 12345, "direct": ["CONNECT"]}]
+    request = rf.post('/api/rpc/v0/', data={})
+    request.user = User(username=_username, password=_password)
+
+    @contextmanager
+    def mock_connect(database_id, user):
+        if database_id == _database_id and user.username == _username:
+            try:
+                yield True
+            finally:
+                pass
+        else:
+            raise AssertionError('incorrect parameters passed')
+
+    def mock_replace_privileges(
+            conn,
+            privileges,
+    ):
+        if privileges != _privileges:
+            raise AssertionError('incorrect parameters passed')
+        return _privileges + [{"role_oid": 67890, "direct": ["CONNECT", "TEMPORARY"]}]
+
+    monkeypatch.setattr(database_privileges, 'connect', mock_connect)
+    monkeypatch.setattr(
+        database_privileges,
+        'replace_database_privileges_for_roles',
+        mock_replace_privileges
+    )
+    expect_response = [
+        {"role_oid": 12345, "direct": ["CONNECT"]},
+        {"role_oid": 67890, "direct": ["CONNECT", "TEMPORARY"]}
+    ]
+    actual_response = database_privileges.replace_for_roles(
+        privileges=_privileges, database_id=_database_id, request=request
+    )
+    assert actual_response == expect_response

mathesar/tests/rpc/test_endpoints.py

@@ -158,6 +158,11 @@
         "database_privileges.get_owner_oid_and_curr_role_db_priv",
         [user_is_authenticated]
     ),
+    (
+        database_privileges.replace_for_roles,
+        "database_privileges.replace_for_roles",
+        [user_is_authenticated]
+    ),
     (
         database_setup.create_new,
         "database_setup.create_new",