From 0b78ac3fc7bd9f0c57518d0c1a153582318edd59 Mon Sep 17 00:00:00 2001
From: Mike Crute <mike@crute.us>
Date: Mon, 5 Dec 2022 20:26:37 -0800
Subject: [PATCH] Fixed #34200 -- Made the session role configurable on
 PostgreSQL.

---
 django/db/backends/postgresql/base.py | 25 ++++++++++++++++++++-----
 docs/ref/databases.txt                | 21 +++++++++++++++++++++
 docs/releases/4.2.txt                 |  6 ++++++
 tests/backends/postgresql/tests.py    | 17 ++++++++++++++++-
 4 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/django/db/backends/postgresql/base.py b/django/db/backends/postgresql/base.py
index 99403f53228..17a3c7a3777 100644
--- a/django/db/backends/postgresql/base.py
+++ b/django/db/backends/postgresql/base.py
@@ -221,6 +221,7 @@ class DatabaseWrapper(BaseDatabaseWrapper):
         else:
             conn_params = {**settings_dict["OPTIONS"]}
 
+        conn_params.pop("assume_role", None)
         conn_params.pop("isolation_level", None)
         if settings_dict["USER"]:
             conn_params["user"] = settings_dict["USER"]
@@ -288,14 +289,28 @@ class DatabaseWrapper(BaseDatabaseWrapper):
             return True
         return False
 
+    def ensure_role(self):
+        if self.connection is None:
+            return False
+        if new_role := self.settings_dict.get("OPTIONS", {}).get("assume_role"):
+            with self.connection.cursor() as cursor:
+                sql = self.ops.compose_sql("SET ROLE %s", [new_role])
+                cursor.execute(sql)
+            return True
+        return False
+
     def init_connection_state(self):
         super().init_connection_state()
 
-        timezone_changed = self.ensure_timezone()
-        if timezone_changed:
-            # Commit after setting the time zone (see #17062)
-            if not self.get_autocommit():
-                self.connection.commit()
+        # Commit after setting the time zone.
+        commit_tz = self.ensure_timezone()
+        # Set the role on the connection. This is useful if the credential used
+        # to login is not the same as the role that owns database resources. As
+        # can be the case when using temporary or ephemeral credentials.
+        commit_role = self.ensure_role()
+
+        if (commit_role or commit_tz) and not self.get_autocommit():
+            self.connection.commit()
 
     @async_unsafe
     def create_cursor(self, name=None):
diff --git a/docs/ref/databases.txt b/docs/ref/databases.txt
index d62adbe8322..ee573c5f658 100644
--- a/docs/ref/databases.txt
+++ b/docs/ref/databases.txt
@@ -230,6 +230,27 @@ configuration in :setting:`DATABASES`::
 
     ``IsolationLevel`` was added.
 
+.. _database-role:
+
+Role
+----
+
+.. versionadded:: 4.2
+
+If you need to use a different role for database connections than the role use
+to establish the connection, set it in the :setting:`OPTIONS` part of your
+database configuration in :setting:`DATABASES`::
+
+    DATABASES = {
+        "default": {
+            "ENGINE": "django.db.backends.postgresql",
+            # ...
+            "OPTIONS": {
+                "assume_role": "my_application_role",
+            },
+        },
+    }
+
 Indexes for ``varchar`` and ``text`` columns
 --------------------------------------------
 
diff --git a/docs/releases/4.2.txt b/docs/releases/4.2.txt
index 40ca5ba1a70..c5228baaa37 100644
--- a/docs/releases/4.2.txt
+++ b/docs/releases/4.2.txt
@@ -224,6 +224,12 @@ CSRF
 
 * ...
 
+Database backends
+~~~~~~~~~~~~~~~~~
+
+* The new ``"assume_role"`` option is now supported in :setting:`OPTIONS` on
+  PostgreSQL to allow specifying the :ref:`session role <database-role>`.
+
 Decorators
 ~~~~~~~~~~
 
diff --git a/tests/backends/postgresql/tests.py b/tests/backends/postgresql/tests.py
index 7e1a2d000d5..7b0e0bfb522 100644
--- a/tests/backends/postgresql/tests.py
+++ b/tests/backends/postgresql/tests.py
@@ -15,7 +15,7 @@ from django.db.backends.base.base import BaseDatabaseWrapper
 from django.test import TestCase, override_settings
 
 try:
-    from django.db.backends.postgresql.psycopg_any import is_psycopg3
+    from django.db.backends.postgresql.psycopg_any import errors, is_psycopg3
 except ImportError:
     is_psycopg3 = False
 
@@ -262,6 +262,21 @@ class Tests(TestCase):
         with self.assertRaisesMessage(ImproperlyConfigured, msg):
             new_connection.ensure_connection()
 
+    def test_connect_role(self):
+        """
+        The session role can be configured with DATABASES
+        ["OPTIONS"]["assume_role"].
+        """
+        try:
+            custom_role = "django_nonexistent_role"
+            new_connection = connection.copy()
+            new_connection.settings_dict["OPTIONS"]["assume_role"] = custom_role
+            msg = f'role "{custom_role}" does not exist'
+            with self.assertRaisesMessage(errors.InvalidParameterValue, msg):
+                new_connection.connect()
+        finally:
+            new_connection.close()
+
     def test_connect_no_is_usable_checks(self):
         new_connection = connection.copy()
         try: