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>

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
+        )