* timezone support moved to base DateTimeField class, timezone accepts string value or pytz.timezone

* test timezones and DateTime64 precision

src/infi/clickhouse_orm/fields.py

@@ -1,14 +1,13 @@
 from __future__ import unicode_literals
 import datetime
-from typing import List
-
+from typing import List, Union
 import iso8601
 import pytz
 from calendar import timegm
 from decimal import Decimal, localcontext
 from uuid import UUID
 from logging import getLogger
-from pytz import UnknownTimeZoneError
+from pytz import BaseTzInfo
 from .utils import escape, parse_array, comma_join, string_or_func, get_subclass_names
 from .funcs import F, FunctionOperatorsMixin
 from ipaddress import IPv4Address, IPv6Address
@@ -197,9 +196,23 @@ class DateTimeField(Field):
     class_default = datetime.datetime.fromtimestamp(0, pytz.utc)
     db_type = 'DateTime'
 
+    def __init__(self, default=None, alias=None, materialized=None, readonly=None, codec=None,
+                 timezone: Union[BaseTzInfo, str] = None):
+        super().__init__(default, alias, materialized, readonly, codec)
+        # assert not timezone, 'Temporarily field timezone is not supported'
+        if timezone:
+            timezone = timezone if isinstance(timezone, BaseTzInfo) else pytz.timezone(timezone)
+        self.timezone: BaseTzInfo = timezone
+
+    def get_db_type_args(self) -> List[str]:
+        args = []
+        if self.timezone:
+            args.append(escape(self.timezone.zone))
+        return args
+
     def to_python(self, value, timezone_in_use):
         if isinstance(value, datetime.datetime):
-            return value.astimezone(pytz.utc) if value.tzinfo else value.replace(tzinfo=pytz.utc)
+            return value if value.tzinfo else value.replace(tzinfo=pytz.utc)
         if isinstance(value, datetime.date):
             return datetime.datetime(value.year, value.month, value.day, tzinfo=pytz.utc)
         if isinstance(value, int):
@@ -222,7 +235,7 @@ class DateTimeField(Field):
             # convert naive to aware
             if dt.tzinfo is None or dt.tzinfo.utcoffset(dt) is None:
                 dt = timezone_in_use.localize(dt)
-            return dt.astimezone(pytz.utc)
+            return dt
         raise ValueError('Invalid value for %s - %r' % (self.__class__.__name__, value))
 
     def to_db_string(self, value, quote=True):
@@ -233,19 +246,15 @@ class DateTime64Field(DateTimeField):
     db_type = 'DateTime64'
 
     def __init__(self, default=None, alias=None, materialized=None, readonly=None, codec=None,
-                 precision: int = 6, timezone: str = None):
-        super().__init__(default, alias, materialized, readonly, codec)
+                 timezone: Union[BaseTzInfo, str] = None, precision: int = 6):
+        super().__init__(default, alias, materialized, readonly, codec, timezone)
         assert precision is None or isinstance(precision, int), 'Precision must be int type'
-        assert timezone is None or isinstance(timezone, str), 'Timezone must be string type'
-        if timezone:
-            pytz.timezone(timezone)
         self.precision = precision
-        self.timezone = timezone
 
     def get_db_type_args(self) -> List[str]:
         args = [str(self.precision)]
         if self.timezone:
-            args.append(escape(self.timezone))
+            args.append(escape(self.timezone.zone))
         return args
 
     def to_db_string(self, value: datetime.datetime, quote=True):
@@ -254,8 +263,30 @@ class DateTime64Field(DateTimeField):
 
         Returns string in 0000000000.000000 format, where remainder digits count is equal to precision
         """
-        width = 11 + self.precision
-        return escape(f'{value.timestamp():0{width}.{self.precision}f}', quote)
+        return escape(
+            '{timestamp:0{width}.{precision}f}'.format(
+                timestamp=value.timestamp(),
+                width=11 + self.precision,
+                precision=6),
+            quote
+        )
+
+    def to_python(self, value, timezone_in_use):
+        try:
+            return super().to_python(value, timezone_in_use)
+        except ValueError:
+            if isinstance(value, (int, float)):
+                return datetime.datetime.utcfromtimestamp(value).replace(tzinfo=pytz.utc)
+            if isinstance(value, str):
+                if value.split('.')[0] == '0000-00-00 00:00:00':
+                    return self.class_default
+                if len(value.split('.')[0]) == 10:
+                    try:
+                        value = float(value)
+                        return datetime.datetime.utcfromtimestamp(value).replace(tzinfo=pytz.utc)
+                    except ValueError:
+                        pass
+            raise
 
 
 class BaseIntField(Field):

src/infi/clickhouse_orm/models.py

@@ -206,8 +206,10 @@ class ModelBase(type):
             return orm_fields.BaseEnumField.create_ad_hoc_field(db_type)
         # DateTime with timezone
         if db_type.startswith('DateTime('):
-            # Some functions return DateTimeField with timezone in brackets
-            return orm_fields.DateTimeField()
+            timezone = db_type[9:-1]
+            return orm_fields.DateTimeField(
+                timezone=timezone[1:-1] if timezone else None
+            )
         # DateTime64
         if db_type.startswith('DateTime64('):
             precision, *timezone = [s.strip() for s in db_type[11:-1].split(',')]
@@ -382,14 +384,15 @@ class Model(metaclass=ModelBase):
 
         - `line`: the TSV-formatted data.
         - `field_names`: names of the model fields in the data.
-        - `timezone_in_use`: the timezone to use when parsing dates and datetimes.
+        - `timezone_in_use`: the timezone to use when parsing dates and datetimes. Some fields use their own timezones.
         - `database`: if given, sets the database that this instance belongs to.
         '''
         values = iter(parse_tsv(line))
         kwargs = {}
         for name in field_names:
             field = getattr(cls, name)
-            kwargs[name] = field.to_python(next(values), timezone_in_use)
+            field_timezone = getattr(field, 'timezone', None) or timezone_in_use
+            kwargs[name] = field.to_python(next(values), field_timezone)
 
         obj = cls(**kwargs)
         if database is not None:

tests/test_datetime_fields.py

@@ -23,11 +23,14 @@ class DateFieldsTest(unittest.TestCase):
             ModelWithDate(
                 date_field='2016-08-30',
                 datetime_field='2016-08-30 03:50:00',
-                datetime64_field='2016-08-30 03:50:00.001'),
+                datetime64_field='2016-08-30 03:50:00.123456',
+                datetime64_3_field='2016-08-30 03:50:00.123456'
+            ),
             ModelWithDate(
                 date_field='2016-08-31',
                 datetime_field='2016-08-31 01:30:00',
-                datetime64_field='2016-08-31 01:30:00.002')
+                datetime64_field='2016-08-31 01:30:00.123456',
+                datetime64_3_field='2016-08-31 01:30:00.123456')
         ])
 
         # toStartOfHour returns DateTime('Asia/Yekaterinburg') in my case, so I test it here to
@@ -36,17 +39,70 @@ class DateFieldsTest(unittest.TestCase):
         self.assertEqual(len(results), 2)
         self.assertEqual(results[0].date_field, datetime.date(2016, 8, 30))
         self.assertEqual(results[0].datetime_field, datetime.datetime(2016, 8, 30, 3, 50, 0, tzinfo=pytz.UTC))
-        self.assertEqual(results[0].datetime64_field, datetime.datetime(2016, 8, 30, 3, 50, 0, 1000, tzinfo=pytz.UTC))
         self.assertEqual(results[0].hour_start, datetime.datetime(2016, 8, 30, 3, 0, 0, tzinfo=pytz.UTC))
         self.assertEqual(results[1].date_field, datetime.date(2016, 8, 31))
         self.assertEqual(results[1].datetime_field, datetime.datetime(2016, 8, 31, 1, 30, 0, tzinfo=pytz.UTC))
-        self.assertEqual(results[1].datetime64_field, datetime.datetime(2016, 8, 31, 1, 30, 0, 2000, tzinfo=pytz.UTC))
         self.assertEqual(results[1].hour_start, datetime.datetime(2016, 8, 31, 1, 0, 0, tzinfo=pytz.UTC))
 
+        self.assertEqual(results[0].datetime64_field, datetime.datetime(2016, 8, 30, 3, 50, 0, 123456, tzinfo=pytz.UTC))
+        self.assertEqual(results[0].datetime64_3_field, datetime.datetime(2016, 8, 30, 3, 50, 0, 123000,
+                                                                          tzinfo=pytz.UTC))
+        self.assertEqual(results[1].datetime64_field, datetime.datetime(2016, 8, 31, 1, 30, 0, 123456, tzinfo=pytz.UTC))
+        self.assertEqual(results[1].datetime64_3_field, datetime.datetime(2016, 8, 31, 1, 30, 0, 123000,
+                                                                          tzinfo=pytz.UTC))
+
 
 class ModelWithDate(Model):
     date_field = DateField()
     datetime_field = DateTimeField()
     datetime64_field = DateTime64Field()
+    datetime64_3_field = DateTime64Field(precision=3)
 
     engine = MergeTree('date_field', ('date_field',))
+
+
+class ModelWithTz(Model):
+    datetime_no_tz_field = DateTimeField()  # server tz
+    datetime_tz_field = DateTimeField(timezone='Europe/Madrid')
+    datetime_utc_field = DateTimeField(timezone=pytz.UTC)
+
+    engine = MergeTree('datetime_no_tz_field', ('datetime_no_tz_field',))
+
+
+class DateTimeFieldWithTzTest(unittest.TestCase):
+    def setUp(self):
+        self.database = Database('test-db', log_statements=True)
+        self.database.create_table(ModelWithTz)
+
+    def tearDown(self):
+        self.database.drop_database()
+
+    def test_ad_hoc_model(self):
+        self.database.insert([
+            ModelWithTz(
+                datetime_no_tz_field='2020-06-11 04:00:00',
+                datetime_tz_field='2020-06-11 04:00:00',
+                datetime_utc_field='2020-06-11 04:00:00',
+            ),
+            ModelWithTz(
+                datetime_no_tz_field='2020-06-11 07:00:00+0300',
+                datetime_tz_field='2020-06-11 07:00:00+0300',
+                datetime_utc_field='2020-06-11 07:00:00+0300',
+            ),
+        ])
+        query = 'SELECT * from $db.modelwithtz ORDER BY datetime_no_tz_field'
+        results = list(self.database.select(query))
+
+        self.assertEqual(results[0].datetime_no_tz_field, datetime.datetime(2020, 6, 11, 4, 0, 0, tzinfo=pytz.UTC))
+        self.assertEqual(results[0].datetime_tz_field, datetime.datetime(2020, 6, 11, 4, 0, 0, tzinfo=pytz.UTC))
+        self.assertEqual(results[0].datetime_utc_field, datetime.datetime(2020, 6, 11, 4, 0, 0, tzinfo=pytz.UTC))
+        self.assertEqual(results[1].datetime_no_tz_field, datetime.datetime(2020, 6, 11, 4, 0, 0, tzinfo=pytz.UTC))
+        self.assertEqual(results[1].datetime_tz_field, datetime.datetime(2020, 6, 11, 4, 0, 0, tzinfo=pytz.UTC))
+        self.assertEqual(results[1].datetime_utc_field, datetime.datetime(2020, 6, 11, 4, 0, 0, tzinfo=pytz.UTC))
+
+        self.assertEqual(results[0].datetime_no_tz_field.tzinfo.zone, self.database.server_timezone.zone)
+        self.assertEqual(results[0].datetime_tz_field.tzinfo.zone, pytz.timezone('Europe/Madrid').zone)
+        self.assertEqual(results[0].datetime_utc_field.tzinfo.zone, pytz.timezone('UTC').zone)
+        self.assertEqual(results[1].datetime_no_tz_field.tzinfo.zone, self.database.server_timezone.zone)
+        self.assertEqual(results[1].datetime_tz_field.tzinfo.zone, pytz.timezone('Europe/Madrid').zone)
+        self.assertEqual(results[1].datetime_utc_field.tzinfo.zone, pytz.timezone('UTC').zone)

tests/test_nullable_fields.py

@@ -38,7 +38,7 @@ class NullableFieldsTest(unittest.TestCase):
             if value == '\\N':
                 self.assertIsNone(dt)
             else:
-                self.assertEqual(dt.tzinfo, pytz.utc)
+                self.assertTrue(dt.tzinfo)
             # Verify that conversion to and from db string does not change value
             dt2 = f.to_python(f.to_db_string(dt, quote=False), pytz.utc)
             self.assertEqual(dt, dt2)

tests/test_simple_fields.py

@@ -6,18 +6,21 @@ import pytz
 
 
 class SimpleFieldsTest(unittest.TestCase):
+    epoch = datetime(1970, 1, 1, tzinfo=pytz.utc)
+    # Valid values
+    dates = [
+        date(1970, 1, 1), datetime(1970, 1, 1), epoch,
+        epoch.astimezone(pytz.timezone('US/Eastern')), epoch.astimezone(pytz.timezone('Asia/Jerusalem')),
+        '1970-01-01 00:00:00', '1970-01-17 00:00:17', '0000-00-00 00:00:00', 0,
+        '2017-07-26T08:31:05', '2017-07-26T08:31:05Z', '2017-07-26 08:31',
+        '2017-07-26T13:31:05+05', '2017-07-26 13:31:05+0500'
+    ]
 
     def test_datetime_field(self):
         f = DateTimeField()
-        epoch = datetime(1970, 1, 1, tzinfo=pytz.utc)
-        # Valid values
-        for value in (date(1970, 1, 1), datetime(1970, 1, 1), epoch,
-                      epoch.astimezone(pytz.timezone('US/Eastern')), epoch.astimezone(pytz.timezone('Asia/Jerusalem')),
-                      '1970-01-01 00:00:00', '1970-01-17 00:00:17', '0000-00-00 00:00:00', 0,
-                      '2017-07-26T08:31:05', '2017-07-26T08:31:05Z', '2017-07-26 08:31',
-                      '2017-07-26T13:31:05+05', '2017-07-26 13:31:05+0500'):
+        for value in self.dates:
             dt = f.to_python(value, pytz.utc)
-            self.assertEqual(dt.tzinfo, pytz.utc)
+            self.assertTrue(dt.tzinfo)
             # Verify that conversion to and from db string does not change value
             dt2 = f.to_python(f.to_db_string(dt, quote=False), pytz.utc)
             self.assertEqual(dt, dt2)
@@ -27,6 +30,28 @@ class SimpleFieldsTest(unittest.TestCase):
             with self.assertRaises(ValueError):
                 f.to_python(value, pytz.utc)
 
+    def test_datetime64_field(self):
+        f = DateTime64Field()
+        epoch = datetime(1970, 1, 1, tzinfo=pytz.utc)
+        # Valid values
+        for value in self.dates + [
+            datetime(1970, 1, 1, microsecond=100000),
+            datetime(1970, 1, 1, microsecond=100000).astimezone(pytz.timezone('US/Eastern')),
+            '1970-01-01 00:00:00.1', '1970-01-17 00:00:17.1', '0000-00-00 00:00:00.1', 0.1,
+            '2017-07-26T08:31:05.1', '2017-07-26T08:31:05.1Z', '2017-07-26 08:31.1',
+            '2017-07-26T13:31:05.1+05', '2017-07-26 13:31:05.1+0500'
+        ]:
+            dt = f.to_python(value, pytz.utc)
+            self.assertTrue(dt.tzinfo)
+            # Verify that conversion to and from db string does not change value
+            dt2 = f.to_python(f.to_db_string(dt, quote=False), pytz.utc)
+            self.assertEqual(dt, dt2)
+        # Invalid values
+        for value in ('nope', '21/7/1999',
+                      '2017-01 15:06:00', '2017-01-01X15:06:00', '2017-13-01T15:06:00'):
+            with self.assertRaises(ValueError):
+                f.to_python(value, pytz.utc)
+
     def test_date_field(self):
         f = DateField()
         epoch = date(1970, 1, 1)