Fixed #34200 -- Made the session role configurable on PostgreSQL.

This commit is contained in:
Mike Crute 2022-12-05 20:26:37 -08:00 committed by Mariusz Felisiak
parent 2a14b8df39
commit 0b78ac3fc7
4 changed files with 63 additions and 6 deletions

View File

@ -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):

View File

@ -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
--------------------------------------------

View File

@ -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
~~~~~~~~~~

View File

@ -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: