From 7380ac57340653854bc2cfe0ed80298cdac6061d Mon Sep 17 00:00:00 2001 From: Sarah Boyce <42296566+sarahboyce@users.noreply.github.com> Date: Mon, 19 Aug 2024 16:51:31 +0200 Subject: [PATCH] Fixed #35688 -- Restored timezone and role setters to be PostgreSQL DatabaseWrapper methods. Following the addition of PostgreSQL connection pool support in Refs #33497, the methods for configuring the database role and timezone were moved to module-level functions. This change prevented subclasses of DatabaseWrapper from overriding these methods as needed, for example, when creating wrappers for other PostgreSQL-based backends. Thank you Christian Hardenberg for the report and to Florian Apolloner and Natalia Bidart for the review. Regression in fad334e1a9b54ea1acb8cce02a25934c5acfe99f. Co-authored-by: Natalia <124304+nessita@users.noreply.github.com> --- django/db/backends/postgresql/base.py | 46 ++++++++++++--------------- docs/releases/5.1.1.txt | 4 +++ tests/backends/postgresql/tests.py | 46 +++++++++++++++++++++++++++ 3 files changed, 71 insertions(+), 25 deletions(-) diff --git a/django/db/backends/postgresql/base.py b/django/db/backends/postgresql/base.py index e97ab6aa89..c864cab57a 100644 --- a/django/db/backends/postgresql/base.py +++ b/django/db/backends/postgresql/base.py @@ -86,24 +86,6 @@ def _get_varchar_column(data): return "varchar(%(max_length)s)" % data -def ensure_timezone(connection, ops, timezone_name): - conn_timezone_name = connection.info.parameter_status("TimeZone") - if timezone_name and conn_timezone_name != timezone_name: - with connection.cursor() as cursor: - cursor.execute(ops.set_time_zone_sql(), [timezone_name]) - return True - return False - - -def ensure_role(connection, ops, role_name): - if role_name: - with connection.cursor() as cursor: - sql = ops.compose_sql("SET ROLE %s", [role_name]) - cursor.execute(sql) - return True - return False - - class DatabaseWrapper(BaseDatabaseWrapper): vendor = "postgresql" display_name = "PostgreSQL" @@ -364,21 +346,35 @@ class DatabaseWrapper(BaseDatabaseWrapper): self.close_pool() if self.connection is None: return False - return ensure_timezone(self.connection, self.ops, self.timezone_name) + return self._configure_timezone(self.connection) + + def _configure_timezone(self, connection): + conn_timezone_name = connection.info.parameter_status("TimeZone") + timezone_name = self.timezone_name + if timezone_name and conn_timezone_name != timezone_name: + with connection.cursor() as cursor: + cursor.execute(self.ops.set_time_zone_sql(), [timezone_name]) + return True + return False + + def _configure_role(self, connection): + if new_role := self.settings_dict["OPTIONS"].get("assume_role"): + with connection.cursor() as cursor: + sql = self.ops.compose_sql("SET ROLE %s", [new_role]) + cursor.execute(sql) + return True + return False def _configure_connection(self, connection): # This function is called from init_connection_state and from the - # psycopg pool itself after a connection is opened. Make sure that - # whatever is done here does not access anything on self aside from - # variables. + # psycopg pool itself after a connection is opened. # Commit after setting the time zone. - commit_tz = ensure_timezone(connection, self.ops, self.timezone_name) + commit_tz = self._configure_timezone(connection) # 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. - role_name = self.settings_dict["OPTIONS"].get("assume_role") - commit_role = ensure_role(connection, self.ops, role_name) + commit_role = self._configure_role(connection) return commit_role or commit_tz diff --git a/docs/releases/5.1.1.txt b/docs/releases/5.1.1.txt index b6f9d23c35..b1c9ec4176 100644 --- a/docs/releases/5.1.1.txt +++ b/docs/releases/5.1.1.txt @@ -31,3 +31,7 @@ Bugfixes * Adjusted the deprecation warning ``stacklevel`` in ``FieldCacheMixin.get_cache_name()`` to correctly point to the offending call site (:ticket:`35405`). + +* Restored, following a regression in Django 5.1, the ability to override the + timezone and role setting behavior used within the ``init_connection_state`` + method of the PostgreSQL backend (:ticket:`35688`). diff --git a/tests/backends/postgresql/tests.py b/tests/backends/postgresql/tests.py index 0b4f580612..37c5ee562b 100644 --- a/tests/backends/postgresql/tests.py +++ b/tests/backends/postgresql/tests.py @@ -567,3 +567,49 @@ class Tests(TestCase): ) finally: new_connection.close() + + def test_bypass_timezone_configuration(self): + from django.db.backends.postgresql.base import DatabaseWrapper + + class CustomDatabaseWrapper(DatabaseWrapper): + def _configure_timezone(self, connection): + return False + + for Wrapper, commit in [ + (DatabaseWrapper, True), + (CustomDatabaseWrapper, False), + ]: + with self.subTest(wrapper=Wrapper, commit=commit): + new_connection = no_pool_connection() + self.addCleanup(new_connection.close) + + # Set the database default time zone to be different from + # the time zone in new_connection.settings_dict. + with new_connection.cursor() as cursor: + cursor.execute("RESET TIMEZONE") + cursor.execute("SHOW TIMEZONE") + db_default_tz = cursor.fetchone()[0] + new_tz = "Europe/Paris" if db_default_tz == "UTC" else "UTC" + new_connection.timezone_name = new_tz + + settings = new_connection.settings_dict.copy() + conn = new_connection.connection + self.assertIs(Wrapper(settings)._configure_connection(conn), commit) + + def test_bypass_role_configuration(self): + from django.db.backends.postgresql.base import DatabaseWrapper + + class CustomDatabaseWrapper(DatabaseWrapper): + def _configure_role(self, connection): + return False + + new_connection = no_pool_connection() + self.addCleanup(new_connection.close) + new_connection.connect() + + settings = new_connection.settings_dict.copy() + settings["OPTIONS"]["assume_role"] = "django_nonexistent_role" + conn = new_connection.connection + self.assertIs( + CustomDatabaseWrapper(settings)._configure_connection(conn), False + )