Refs #33308 -- Added psycopg_any.IsolationLevel.

django · Dec 12, 2022 · 1d90c9b · 1d90c9b
1 parent 2f38f7b
commit 1d90c9b
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 9 deletions.
diff --git a/django/db/backends/postgresql/base.py b/django/db/backends/postgresql/base.py
@@ -48,6 +48,7 @@ def psycopg2_version():
 from .features import DatabaseFeatures  # NOQA
 from .introspection import DatabaseIntrospection  # NOQA
 from .operations import DatabaseOperations  # NOQA
+from .psycopg_any import IsolationLevel  # NOQA
 from .schema import DatabaseSchemaEditor  # NOQA
 
 psycopg2.extensions.register_adapter(SafeString, psycopg2.extensions.QuotedString)
@@ -212,22 +213,30 @@ def get_connection_params(self):
 
     @async_unsafe
     def get_new_connection(self, conn_params):
-        connection = Database.connect(**conn_params)
-
         # self.isolation_level must be set:
         # - after connecting to the database in order to obtain the database's
         #   default when no value is explicitly specified in options.
         # - before calling _set_autocommit() because if autocommit is on, that
         #   will set connection.isolation_level to ISOLATION_LEVEL_AUTOCOMMIT.
         options = self.settings_dict["OPTIONS"]
+        set_isolation_level = False
         try:
-            self.isolation_level = options["isolation_level"]
+            isolation_level_value = options["isolation_level"]
         except KeyError:
-            self.isolation_level = connection.isolation_level
+            self.isolation_level = IsolationLevel.READ_COMMITTED
         else:
             # Set the isolation level to the value from OPTIONS.
-            if self.isolation_level != connection.isolation_level:
-                connection.set_session(isolation_level=self.isolation_level)
+            try:
+                self.isolation_level = IsolationLevel(isolation_level_value)
+                set_isolation_level = True
+            except ValueError:
+                raise ImproperlyConfigured(
+                    f"Invalid transaction isolation level {isolation_level_value} "
+                    f"specified. Use one of the IsolationLevel values."
+                )
+        connection = Database.connect(**conn_params)
+        if set_isolation_level:
+            connection.isolation_level = self.isolation_level
         # Register dummy loads() to avoid a round trip from psycopg2's decode
         # to json.dumps() to json.loads(), when using a custom decoder in
         # JSONField.

diff --git a/django/db/backends/postgresql/psycopg_any.py b/django/db/backends/postgresql/psycopg_any.py
@@ -1,3 +1,5 @@
+from enum import IntEnum
+
 from psycopg2 import errors, extensions, sql  # NOQA
 from psycopg2.extras import DateRange, DateTimeRange, DateTimeTZRange, Inet  # NOQA
 from psycopg2.extras import Json as Jsonb  # NOQA
@@ -6,6 +8,13 @@
 RANGE_TYPES = (DateRange, DateTimeRange, DateTimeTZRange, NumericRange)
 
 
+class IsolationLevel(IntEnum):
+    READ_UNCOMMITTED = extensions.ISOLATION_LEVEL_READ_UNCOMMITTED
+    READ_COMMITTED = extensions.ISOLATION_LEVEL_READ_COMMITTED
+    REPEATABLE_READ = extensions.ISOLATION_LEVEL_REPEATABLE_READ
+    SERIALIZABLE = extensions.ISOLATION_LEVEL_SERIALIZABLE
+
+
 def _quote(value, connection=None):
     adapted = extensions.adapt(value)
     if hasattr(adapted, "encoding"):

diff --git a/tests/backends/postgresql/tests.py b/tests/backends/postgresql/tests.py
@@ -223,7 +223,7 @@ def test_connect_isolation_level(self):
         The transaction level can be configured with
         DATABASES ['OPTIONS']['isolation_level'].
         """
-        from psycopg2.extensions import ISOLATION_LEVEL_SERIALIZABLE as serializable
+        from django.db.backends.postgresql.psycopg_any import IsolationLevel
 
         # Since this is a django.test.TestCase, a transaction is in progress
         # and the isolation level isn't reported as 0. This test assumes that
@@ -232,15 +232,31 @@ def test_connect_isolation_level(self):
         self.assertIsNone(connection.connection.isolation_level)
 
         new_connection = connection.copy()
-        new_connection.settings_dict["OPTIONS"]["isolation_level"] = serializable
+        new_connection.settings_dict["OPTIONS"][
+            "isolation_level"
+        ] = IsolationLevel.SERIALIZABLE
         try:
             # Start a transaction so the isolation level isn't reported as 0.
             new_connection.set_autocommit(False)
             # Check the level on the psycopg2 connection, not the Django wrapper.
-            self.assertEqual(new_connection.connection.isolation_level, serializable)
+            self.assertEqual(
+                new_connection.connection.isolation_level,
+                IsolationLevel.SERIALIZABLE,
+            )
         finally:
             new_connection.close()
 
+    def test_connect_invalid_isolation_level(self):
+        self.assertIsNone(connection.connection.isolation_level)
+        new_connection = connection.copy()
+        new_connection.settings_dict["OPTIONS"]["isolation_level"] = -1
+        msg = (
+            "Invalid transaction isolation level -1 specified. Use one of the "
+            "IsolationLevel values."
+        )
+        with self.assertRaisesMessage(ImproperlyConfigured, msg):
+            new_connection.ensure_connection()
+
     def test_connect_no_is_usable_checks(self):
         new_connection = connection.copy()
         try: