Timezone support in DateTime type

.travis.yml

@@ -1,4 +1,9 @@
 env:
+  - VERSION=1.1.54380
+#  - VERSION=1.1.54378 client's image miss tzdata package: https://github.com/yandex/ClickHouse/commit/1bf49fe8446c7dea95beaef2b131e6c6708b0b62#diff-cc737435a5ba74620a889b7718f39a80
+  - VERSION=1.1.54343
+  - VERSION=1.1.54342
+#  - VERSION=1.1.54337 Broken network
   - VERSION=1.1.54327
   - VERSION=1.1.54310
   - VERSION=1.1.54304
@@ -17,11 +22,11 @@ cache: pip
 services:
   - docker
 before_install:
-  - docker run -d -p 127.0.0.1:9000:9000 --name test-clickhouse-server --ulimit nofile=262144:262144 yandex/clickhouse-server:$VERSION
+  - docker run -e "TZ=Europe/Moscow" -d -p 127.0.0.1:9000:9000 --name test-clickhouse-server --ulimit nofile=262144:262144 yandex/clickhouse-server:$VERSION
   - docker run -d --entrypoint "/bin/sh" --name test-clickhouse-client --link test-clickhouse-server:clickhouse-server yandex/clickhouse-client:$VERSION -c 'while :; do sleep 1; done'
   - docker ps -a
-  # Faking clickhouse-client real comminitation with container via docker exec.
-  - echo -e '#!/bin/bash\n\ndocker exec test-clickhouse-client clickhouse-client "$@"' | sudo tee /usr/local/bin/clickhouse-client > /dev/null
+  # Faking clickhouse-client real communication with container via docker exec.
+  - echo -e '#!/bin/bash\n\ndocker exec -e "`env | grep ^TZ=`" test-clickhouse-client clickhouse-client "$@"' | sudo tee /usr/local/bin/clickhouse-client > /dev/null
   - sudo chmod +x /usr/local/bin/clickhouse-client
   # Overriding setup.cfg. Set host=clickhouse-server
   - sed -i 's/^host=localhost$/host=clickhouse-server/' setup.cfg

setup.py

@@ -10,7 +10,7 @@
 
 PY34 = sys.version_info[0:2] >= (3, 4)
 
-install_requires = []
+install_requires = ['pytz']
 if not PY34:
     install_requires.append('enum34')
 

src/columns/datetimecolumn.py

@@ -1,7 +1,9 @@
 from calendar import timegm
 from datetime import datetime
+from time import mktime
+
+from pytz import timezone as get_timezone, utc
 
-from ..util.tzinfo import tzutc
 from .base import FormatColumn
 
 
@@ -10,10 +12,45 @@ class DateTimeColumn(FormatColumn):
     py_types = (datetime, )
     format = 'I'
 
-    utc = tzutc()
+    def __init__(self, timezone=None, **kwargs):
+        self.timezone = timezone
+        super(DateTimeColumn, self).__init__(**kwargs)
 
     def after_read_item(self, value):
-        return datetime.fromtimestamp(value, tz=self.utc).replace(tzinfo=None)
+        dt = datetime.fromtimestamp(value, self.timezone)
+        return dt.replace(tzinfo=None)
 
     def before_write_item(self, value):
-        return int(timegm(value.timetuple()))
+        if self.timezone:
+            # Set server's timezone for offset-naive datetime.
+            if value.tzinfo is None:
+                value = self.timezone.localize(value)
+
+            value = value.astimezone(utc)
+            return int(timegm(value.timetuple()))
+
+        else:
+            # If datetime is offset-aware use it's timezone.
+            if value.tzinfo is not None:
+                value = value.astimezone(utc)
+                return int(timegm(value.timetuple()))
+
+            return int(mktime(value.timetuple()))
+
+
+def create_datetime_column(spec, column_options):
+    context = column_options['context']
+
+    tz_name = timezone = None
+
+    # Use column's timezone if it's specified.
+    if spec[-1] == ')':
+        tz_name = spec[10:-2]
+    else:
+        if not context.settings.get('use_client_time_zone', False):
+            tz_name = context.server_info.timezone
+
+    if tz_name:
+        timezone = get_timezone(tz_name)
+
+    return DateTimeColumn(timezone=timezone, **column_options)

src/columns/service.py

@@ -1,7 +1,7 @@
 from .. import errors
 from .arraycolumn import create_array_column
 from .datecolumn import DateColumn
-from .datetimecolumn import DateTimeColumn
+from .datetimecolumn import create_datetime_column
 from . import exceptions as column_exceptions
 from .enumcolumn import create_enum_column
 from .floatcolumn import Float32, Float64
@@ -22,7 +22,7 @@
 
 
 column_by_type = {c.ch_type: c for c in [
-    DateColumn, DateTimeColumn, String, Float32, Float64,
+    DateColumn, String, Float32, Float64,
     Int8Column, Int16Column, Int32Column, Int64Column,
     UInt8Column, UInt16Column, UInt32Column, UInt64Column,
     NothingColumn, NullColumn, UUIDColumn,
@@ -45,6 +45,9 @@ def create_column_with_options(x):
     elif spec.startswith('Enum'):
         return create_enum_column(spec, column_options)
 
+    elif spec.startswith('DateTime'):
+        return create_datetime_column(spec, column_options)
+
     elif spec.startswith('Array'):
         return create_array_column(spec, create_column_with_options)
 

src/defines.py

@@ -23,4 +23,4 @@
 
 DBMS_NAME = 'ClickHouse'
 CLIENT_NAME = 'python-driver'
-CLIENT_VERSION = 54276
+CLIENT_VERSION = 54337

tests/columns/test_date.py

@@ -20,6 +20,6 @@ def test_do_not_use_timezone(self):
             inserted = self.emit_cli(query)
             self.assertEqual(inserted, '1970-01-02\n')
 
-            with patch.dict(os.environ, {'TZ': 'US/Hawaii'}, clear=True):
+            with patch.dict(os.environ, {'TZ': 'US/Hawaii'}):
                 inserted = self.client.execute(query)
                 self.assertEqual(inserted, data)

tests/columns/test_datetime.py

@@ -1,6 +1,13 @@
+from contextlib import contextmanager
 from datetime import date, datetime
+import os
+from time import tzset
+
+from mock import patch
+from pytz import timezone, utc
 
 from tests.testcase import BaseTestCase
+from tests.util import require_server_version
 
 
 class DateTimeTestCase(BaseTestCase):
@@ -56,3 +63,210 @@ def test_nullable_datetime(self):
 
             inserted = self.client.execute(query)
             self.assertEqual(inserted, data)
+
+
+class DateTimeTimezonesTestCase(BaseTestCase):
+    @contextmanager
+    def patch_env_tz(self, tz_name):
+        # Although in many cases, changing the TZ environment variable may
+        # affect the output of functions like localtime() without calling
+        # tzset(), this behavior should not be relied on.
+        # https://docs.python.org/3/library/time.html#time.tzset
+        with patch.dict(os.environ, {'TZ': tz_name}):
+            tzset()
+            yield
+
+        tzset()
+
+    # Asia/Kamchatka = UTC+12
+    # Asia/Novosibirsk = UTC+7
+    # Europe/Moscow = UTC+3
+
+    # 1500000000 second since epoch in Europe/Moscow.
+    # 1500010800 second since epoch in UTC.
+    dt = datetime(2017, 7, 14, 5, 40)
+    dt_tz = timezone('Asia/Kamchatka').localize(dt)
+
+    # INSERTs ans SELECTs must be the same as clickhouse-client's.
+
+    def test_use_server_timezone(self):
+        # Determine server timezone and calculate expected timestamp.
+        server_tz_name = self.client.execute('SELECT timezone()')[0][0]
+        offset = timezone(server_tz_name).utcoffset(self.dt).total_seconds()
+        timestamp = 1500010800 - int(offset)
+
+        with self.patch_env_tz('Asia/Novosibirsk'):
+            with self.create_table('a DateTime'):
+                self.client.execute(
+                    'INSERT INTO test (a) VALUES', [(self.dt, )]
+                )
+
+                self.emit_cli(
+                    "INSERT INTO test (a) VALUES ('2017-07-14 05:40:00')"
+                )
+
+                query = 'SELECT toInt32(a) FROM test'
+                inserted = self.emit_cli(query)
+                self.assertEqual(inserted, '{ts}\n{ts}\n'.format(ts=timestamp))
+
+                query = 'SELECT * FROM test'
+                inserted = self.emit_cli(query)
+                self.assertEqual(
+                    inserted,
+                    '2017-07-14 05:40:00\n2017-07-14 05:40:00\n'
+                )
+
+                inserted = self.client.execute(query)
+                self.assertEqual(inserted, [(self.dt, ), (self.dt, )])
+
+    def test_use_client_timezone(self):
+        settings = {'use_client_time_zone': True}
+
+        with self.patch_env_tz('Asia/Novosibirsk'):
+            with self.create_table('a DateTime'):
+                self.client.execute(
+                    'INSERT INTO test (a) VALUES', [(self.dt, )],
+                    settings=settings
+                )
+
+                self.emit_cli(
+                    "INSERT INTO test (a) VALUES ('2017-07-14 05:40:00')",
+                    use_client_time_zone=1
+                )
+
+                query = 'SELECT toInt32(a) FROM test'
+                inserted = self.emit_cli(query, use_client_time_zone=1)
+                # 1499985600 = 1500000000 - 4 * 3600
+                self.assertEqual(inserted, '1499985600\n1499985600\n')
+
+                query = 'SELECT * FROM test'
+                inserted = self.emit_cli(query, use_client_time_zone=1)
+                self.assertEqual(
+                    inserted,
+                    '2017-07-14 05:40:00\n2017-07-14 05:40:00\n'
+                )
+
+                inserted = self.client.execute(query, settings=settings)
+                self.assertEqual(inserted, [(self.dt, ), (self.dt, )])
+
+    @require_server_version(1, 1, 54337)
+    def test_datetime_with_timezone_use_server_timezone(self):
+        server_tz_name = self.client.execute('SELECT timezone()')[0][0]
+        offset = timezone(server_tz_name).utcoffset(self.dt)
+
+        with self.patch_env_tz('Asia/Novosibirsk'):
+            with self.create_table('a DateTime'):
+                self.client.execute(
+                    'INSERT INTO test (a) VALUES', [(self.dt_tz, )]
+                )
+
+                self.emit_cli(
+                    "INSERT INTO test (a) VALUES "
+                    "(toDateTime('2017-07-14 05:40:00', 'Asia/Kamchatka'))",
+                )
+
+                query = 'SELECT toInt32(a) FROM test'
+                inserted = self.emit_cli(query)
+                # 1499967600 = 1500000000 - 12 * 3600
+                self.assertEqual(inserted, '1499967600\n1499967600\n')
+
+                query = 'SELECT * FROM test'
+                inserted = self.emit_cli(query)
+
+                dt = (self.dt_tz.astimezone(utc) + offset).replace(tzinfo=None)
+                self.assertEqual(inserted, '{dt}\n{dt}\n'.format(dt=dt))
+
+                inserted = self.client.execute(query)
+                self.assertEqual(inserted, [(dt, ), (dt, )])
+
+    @require_server_version(1, 1, 54337)
+    def test_datetime_with_timezone_use_client_timezone(self):
+        settings = {'use_client_time_zone': True}
+
+        with self.patch_env_tz('Asia/Novosibirsk'):
+            with self.create_table('a DateTime'):
+                self.client.execute(
+                    'INSERT INTO test (a) VALUES', [(self.dt_tz, )],
+                    settings=settings
+                )
+
+                self.emit_cli(
+                    "INSERT INTO test (a) VALUES "
+                    "(toDateTime('2017-07-14 05:40:00', 'Asia/Kamchatka'))",
+                    use_client_time_zone=1
+                )
+
+                query = 'SELECT toInt32(a) FROM test'
+                inserted = self.emit_cli(query, use_client_time_zone=1)
+                # 1499967600 = 1500000000 - 12 * 3600
+                self.assertEqual(inserted, '1499967600\n1499967600\n')
+
+                query = 'SELECT * FROM test'
+                inserted = self.emit_cli(query, use_client_time_zone=1)
+                # 2017-07-14 00:40:00 = 2017-07-14 05:40:00 - 05:00:00
+                # (Kamchatka - Novosibirsk)
+                self.assertEqual(
+                    inserted,
+                    '2017-07-14 00:40:00\n2017-07-14 00:40:00\n'
+                )
+
+                inserted = self.client.execute(query, settings=settings)
+                dt = datetime(2017, 7, 14, 0, 40)
+                self.assertEqual(inserted, [(dt, ), (dt, )])
+
+    @require_server_version(1, 1, 54337)
+    def test_column_use_server_timezone(self):
+        with self.patch_env_tz('Europe/Moscow'):
+            with self.create_table("a DateTime('Asia/Novosibirsk')"):
+                self.client.execute(
+                    'INSERT INTO test (a) VALUES', [(self.dt, )]
+                )
+
+                self.emit_cli(
+                    "INSERT INTO test (a) VALUES ('2017-07-14 05:40:00')"
+                )
+
+                query = 'SELECT toInt32(a) FROM test'
+                inserted = self.emit_cli(query)
+                # 1499985600 = 1500000000 - 4 * 3600
+                self.assertEqual(inserted, '1499985600\n1499985600\n')
+
+                query = 'SELECT * FROM test'
+                inserted = self.emit_cli(query)
+                self.assertEqual(
+                    inserted,
+                    '2017-07-14 05:40:00\n2017-07-14 05:40:00\n'
+                )
+
+                inserted = self.client.execute(query)
+                self.assertEqual(inserted, [(self.dt, ), (self.dt, )])
+
+    @require_server_version(1, 1, 54337)
+    def test_column_use_client_timezone(self):
+        settings = {'use_client_time_zone': True}
+
+        with self.patch_env_tz('Europe/Moscow'):
+            with self.create_table("a DateTime('Asia/Novosibirsk')"):
+                self.client.execute(
+                    'INSERT INTO test (a) VALUES', [(self.dt, )],
+                    settings=settings
+                )
+                self.emit_cli(
+                    "INSERT INTO test (a) VALUES ('2017-07-14 05:40:00')",
+                    use_client_time_zone=1
+                )
+
+                query = 'SELECT toInt32(a) FROM test'
+                inserted = self.emit_cli(query, use_client_time_zone=1)
+                # 1499985600 = 1500000000 - 4 * 3600
+                self.assertEqual(inserted, '1499985600\n1499985600\n')
+
+                query = 'SELECT * FROM test'
+                inserted = self.emit_cli(query, use_client_time_zone=1)
+                self.assertEqual(
+                    inserted,
+                    '2017-07-14 05:40:00\n2017-07-14 05:40:00\n'
+                )
+
+                inserted = self.client.execute(query, settings=settings)
+                self.assertEqual(inserted, [(self.dt, ), (self.dt, )])