1 files changed, 44 insertions, 8 deletions
diff --git a/postgresqleu/util/db.py b/postgresqleu/util/db.py
index c9d8f5c3..e0a07381 100644
--- a/postgresqleu/util/db.py
+++ b/postgresqleu/util/db.py
@@ -1,33 +1,43 @@
 from django.db import connection
+from django.conf import settings
 import collections
 
+from psycopg2.tz import LocalTimezone
 
-def exec_no_result(query, params=None):
+
+def get_native_cursor():
+    # Remove djangos dumb prevention of us to change the timezone
     curs = connection.cursor()
+    curs.cursor.tzinfo_factory = LocalTimezone
+    return curs
+
+
+def exec_no_result(query, params=None):
+    curs = get_native_cursor()
     curs.execute(query, params)
 
 
 def exec_to_list(query, params=None):
-    curs = connection.cursor()
+    curs = get_native_cursor()
     curs.execute(query, params)
     return curs.fetchall()
 
 
 def exec_to_single_list(query, params=None):
-    curs = connection.cursor()
+    curs = get_native_cursor()
     curs.execute(query, params)
     return [r[0] for r in curs.fetchall()]
 
 
 def exec_to_dict(query, params=None):
-    curs = connection.cursor()
+    curs = get_native_cursor()
     curs.execute(query, params)
     columns = [col[0] for col in curs.description]
     return [dict(list(zip(columns, row)))for row in curs.fetchall()]
 
 
 def exec_to_scalar(query, params=None):
-    curs = connection.cursor()
+    curs = get_native_cursor()
     curs.execute(query, params)
     r = curs.fetchone()
     if r:
@@ -44,20 +54,20 @@ def conditional_exec_to_scalar(condition, query, params=None):
 
 
 def exec_to_keyed_dict(query, params=None):
-    curs = connection.cursor()
+    curs = get_native_cursor()
     curs.execute(query, params)
     columns = [col[0] for col in curs.description]
     return {r[columns[0]]: r for r in (dict(list(zip(columns, row)))for row in curs.fetchall())}
 
 
 def exec_to_keyed_scalar(query, params=None):
-    curs = connection.cursor()
+    curs = get_native_cursor()
     curs.execute(query, params)
     return dict(curs.fetchall())
 
 
 def exec_to_grouped_dict(query, params=None):
-    curs = connection.cursor()
+    curs = get_native_cursor()
     curs.execute(query, params)
     columns = [col[0] for col in curs.description[1:]]
     full = collections.OrderedDict()
@@ -73,3 +83,29 @@ def exec_to_grouped_dict(query, params=None):
     if last:
         full[last] = curr
     return full
+
+
+class ensure_conference_timezone():
+    """
+    This context handler will set the timezone *in PostgreSQL* to the one from the
+    specified conference, and reset it back to UTC when it's done. During this time,
+    calls from django's own ORM will fail, it's intended only to wrap native SQL
+    querying that needs to use the appropriate timezone.
+    If None is specified as the conference, set the timezone to the global one for
+    the installation (otherwise the default is UTC).
+    """
+    def __init__(self, conference):
+        if conference is None:
+            self.tzname = settings.TIME_ZONE
+        else:
+            self.tzname = conference.tzname
+
+    def __enter__(self):
+        c = connection.cursor()
+        c.execute("SET TIMEZONE=%(tz)s", {
+            'tz': self.tzname,
+        })
+        return c
+
+    def __exit__(self, *args):
+        connection.cursor().execute("SET TIMEZONE='UTC'")