DateTime after_read_items timezone optimization

clickhouse_driver/columns/datetimecolumn.py

@@ -3,6 +3,7 @@
 from time import mktime
 
 from pytz import timezone as get_timezone, utc
+from tzlocal import get_localzone
 
 from .base import FormatColumn
 
@@ -18,27 +19,39 @@ def __init__(self, timezone=None, offset_naive=True, **kwargs):
         super(DateTimeColumn, self).__init__(**kwargs)
 
     def after_read_items(self, items, nulls_map=None):
-        timezone = self.timezone
+        tz = self.timezone
         fromtimestamp = datetime.fromtimestamp
 
+        # A bit ugly copy-paste. But it helps save time on items
+        # processing by avoiding lambda calls or if in loop.
         if self.offset_naive:
-            if nulls_map is None:
-                return tuple(
-                    fromtimestamp(item, timezone).replace(tzinfo=None)
-                    for item in items
-                )
+            if tz:
+                if nulls_map is None:
+                    return tuple(
+                        fromtimestamp(item, tz).replace(tzinfo=None)
+                        for item in items
+                    )
+                else:
+                    return tuple(
+                        (None if is_null else
+                         fromtimestamp(items[i], tz).replace(tzinfo=None))
+                        for i, is_null in enumerate(nulls_map)
+                    )
             else:
-                return tuple(
-                    (None if is_null else
-                     fromtimestamp(items[i], timezone).replace(tzinfo=None))
-                    for i, is_null in enumerate(nulls_map)
-                )
+                if nulls_map is None:
+                    return tuple(fromtimestamp(item) for item in items)
+                else:
+                    return tuple(
+                        (None if is_null else fromtimestamp(items[i]))
+                        for i, is_null in enumerate(nulls_map)
+                    )
+
         else:
             if nulls_map is None:
-                return tuple(fromtimestamp(item, timezone) for item in items)
+                return tuple(fromtimestamp(item, tz) for item in items)
             else:
                 return tuple(
-                    (None if is_null else fromtimestamp(items[i], timezone))
+                    (None if is_null else fromtimestamp(items[i], tz))
                     for i, is_null in enumerate(nulls_map)
                 )
 
@@ -86,7 +99,13 @@ def create_datetime_column(spec, column_options):
         offset_naive = False
     else:
         if not context.settings.get('use_client_time_zone', False):
-            tz_name = context.server_info.timezone
+            try:
+                local_timezone = get_localzone().zone
+            except Exception:
+                local_timezone = None
+
+            if local_timezone != context.server_info.timezone:
+                tz_name = context.server_info.timezone
 
     if tz_name:
         timezone = get_timezone(tz_name)

setup.py

@@ -117,6 +117,7 @@ def read_version():
         'pytz',
         'enum34; python_version<"3.4"',
         'ipaddress; python_version<"3.4"',
+        'tzlocal'
     ],
     ext_modules=extensions,
     extras_require={

tests/columns/test_datetime.py

@@ -4,13 +4,30 @@
 from time import tzset
 
 from mock import patch
-from pytz import timezone, utc
+from pytz import timezone, utc, UnknownTimeZoneError
+import tzlocal
 
 from tests.testcase import BaseTestCase
 from tests.util import require_server_version
 
 
-class DateTimeTestCase(BaseTestCase):
+class BaseDateTimeTestCase(BaseTestCase):
+    def setUp(self):
+        super(BaseDateTimeTestCase, self).setUp()
+
+        # Bust tzlocal cache.
+        try:
+            tzlocal.unix._cache_tz = None
+        except AttributeError:
+            pass
+
+        try:
+            tzlocal.win32._cache_tz = None
+        except AttributeError:
+            pass
+
+
+class DateTimeTestCase(BaseDateTimeTestCase):
     def test_simple(self):
         with self.create_table('a Date, b DateTime'):
             data = [(date(2012, 10, 25), datetime(2012, 10, 25, 14, 7, 19))]
@@ -64,8 +81,13 @@ def test_nullable_datetime(self):
             inserted = self.client.execute(query)
             self.assertEqual(inserted, data)
 
+    def test_handle_errors_from_tzlocal(self):
+        with patch('tzlocal.get_localzone') as mocked_get_localzone:
+            mocked_get_localzone.side_effect = UnknownTimeZoneError()
+            self.client.execute('SELECT now()')
+
 
-class DateTimeTimezonesTestCase(BaseTestCase):
+class DateTimeTimezonesTestCase(BaseDateTimeTestCase):
     @contextmanager
     def patch_env_tz(self, tz_name):
         # Although in many cases, changing the TZ environment variable may