Ticket #17062: 17062.diff

django/db/backends/postgresql_psycopg2/base.py

@@ -176 +176 @@
                 conn_params['port'] = settings_dict['PORT']
             self.connection = Database.connect(**conn_params)
             self.connection.set_client_encoding('UTF8')
-            self.connection.set_isolation_level(self.isolation_level)
+            wanted_tz = 'UTC' if settings.USE_TZ else settings_dict.get('TIME_ZONE')
+            db_default_tz = self.connection.get_parameter_status('TimeZone')
+            if wanted_tz and db_default_tz <> wanted_tz:
+                # We must set this in autocommit mode. Otherwise a rollback
+                # will roll back also the SET TIME ZONE.
+                self.connection.set_isolation_level(
+                    psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+                self.connection.cursor().execute(
+                    "SET TIME ZONE %s", (wanted_tz,))
+
+            self.connection.set_isolation_level(
+                 psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)
+            self._get_pg_version()
             connection_created.send(sender=self.__class__, connection=self)
         cursor = self.connection.cursor()
         cursor.tzinfo_factory = utc_tzinfo_factory if settings.USE_TZ else None
-        if new_connection:
-            tz = 'UTC' if settings.USE_TZ else settings_dict.get('TIME_ZONE')
-            if tz:
-                cursor.execute("SET TIME ZONE %s", [tz])
-            self._get_pg_version()
         return CursorWrapper(cursor)
 
     def _enter_transaction_management(self, managed):

tests/regressiontests/backends/tests.py

@@ -10 +10 @@
     IntegrityError, transaction)
 from django.db.backends.signals import connection_created
 from django.db.backends.postgresql_psycopg2 import version as pg_version
+from django.db.utils import ConnectionHandler, DatabaseError
 from django.test import TestCase, skipUnlessDBFeature, TransactionTestCase
 from django.utils import unittest
 
@@ -229 +230 @@
         conn = OlderConnectionMock()
         self.assertEqual(pg_version.get_version(conn), 80300)
 
+class PostgresNewConnectionTest(TestCase):
+    """
+    PostgreSQL will roll back SET TIME ZONE if the transaction containing
+    the SET TIME ZONE command is rolled back. Test that this does not happen.
+    """
+    @unittest.skipUnless(connection.vendor == 'postgresql',
+                         "Test valid only for PostgreSQL")
+    @unittest.skipUnless(connection.isolation_level > 0,
+                         "Test valid only if not using autocommit")
+    def test_connect_and_rollback(self):
+        new_connections = ConnectionHandler(settings.DATABASES)
+        new_connection = new_connections[DEFAULT_DB_ALIAS]
+        try:
+            # Ensure the database default time zone is different than
+            # the time zone in new_connection.settings_dict. We can
+            # get the default time zone by reset & show.
+            cursor = new_connection.cursor()
+            cursor.execute("reset timezone")
+            cursor.execute("show timezone")
+            db_default_tz = cursor.fetchone()[0]
+            new_connection.close()
+            new_tz = db_default_tz == 'UTC' and 'GMT+1' or 'UTC'
+
+            # Fetch a new connection with the new_tz as default
+            # time zone, run a query and rollback.
+            new_connection.settings_dict['TIME_ZONE'] = new_tz
+            new_connection.enter_transaction_management()
+            cursor = new_connection.cursor()
+            cursor.execute("select 1")
+            new_connection.rollback()
+
+            # Now lets see if the rollback rolled back the SET TIME ZONE.
+            encoding = cursor.fetchone()[0]
+            cursor.execute("show timezone")
+            tz = cursor.fetchone()[0]
+            self.assertEqual(new_tz, tz)
+        finally:
+            try:
+                new_connection.close()
+            except DatabaseError:
+                pass
+
+
 # Unfortunately with sqlite3 the in-memory test database cannot be
 # closed, and so it cannot be re-opened during testing, and so we
 # sadly disable this test for now.