From 74b836abf541bb83047c5f07fbb6f49783d46475 Mon Sep 17 00:00:00 2001
From: Aymeric Augustin <aymeric.augustin@m4x.org>
Date: Sun, 20 Nov 2011 16:12:33 +0000
Subject: [PATCH] Fixed #17062 -- Ensured that the effect of SET TIME ZONE
 isn't lost when the first transation is rolled back under PostgreSQL. Thanks
 Anssi for the patch.

git-svn-id: http://code.djangoproject.com/svn/django/trunk@17128 bcc190cf-cafb-0310-a4f2-bffc1f526a37
---
 .../db/backends/postgresql_psycopg2/base.py   | 14 +++----
 tests/regressiontests/backends/tests.py       | 42 +++++++++++++++++++
 2 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/django/db/backends/postgresql_psycopg2/base.py b/django/db/backends/postgresql_psycopg2/base.py
index dbd49c55e1..c81623720e 100644
--- a/django/db/backends/postgresql_psycopg2/base.py
+++ b/django/db/backends/postgresql_psycopg2/base.py
@@ -153,10 +153,8 @@ class DatabaseWrapper(BaseDatabaseWrapper):
     pg_version = property(_get_pg_version)
 
     def _cursor(self):
-        new_connection = False
         settings_dict = self.settings_dict
         if self.connection is None:
-            new_connection = True
             if settings_dict['NAME'] == '':
                 from django.core.exceptions import ImproperlyConfigured
                 raise ImproperlyConfigured("You need to specify NAME in your Django settings file.")
@@ -176,15 +174,17 @@ class DatabaseWrapper(BaseDatabaseWrapper):
                 conn_params['port'] = settings_dict['PORT']
             self.connection = Database.connect(**conn_params)
             self.connection.set_client_encoding('UTF8')
+            # Set the time zone in autocommit mode (see #17062)
+            tz = 'UTC' if settings.USE_TZ else settings_dict.get('TIME_ZONE')
+            if tz:
+                self.connection.set_isolation_level(
+                        psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+                self.connection.cursor().execute("SET TIME ZONE %s", [tz])
             self.connection.set_isolation_level(self.isolation_level)
+            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):
diff --git a/tests/regressiontests/backends/tests.py b/tests/regressiontests/backends/tests.py
index f2bd71df3a..a37d566048 100644
--- a/tests/regressiontests/backends/tests.py
+++ b/tests/regressiontests/backends/tests.py
@@ -10,6 +10,7 @@ from django.db import (backend, connection, connections, DEFAULT_DB_ALIAS,
     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,6 +230,47 @@ class PostgresVersionTest(TestCase):
         conn = OlderConnectionMock()
         self.assertEqual(pg_version.get_version(conn), 80300)
 
+class PostgresNewConnectionTest(TestCase):
+    """
+    #17062: PostgreSQL shouldn't roll back SET TIME ZONE, even if the first
+    transaction is rolled back.
+    """
+    @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_tz = 'Europe/Paris' if db_default_tz == 'UTC' else 'UTC'
+            new_connection.close()
+
+            # 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()
+            new_connection.rollback()
+
+            # Now let's see if the rollback rolled back the SET TIME ZONE.
+            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.