diff --git a/cqlengine/columns.py b/cqlengine/columns.py index 10a5e1b9..dd7b3ba3 100644 --- a/cqlengine/columns.py +++ b/cqlengine/columns.py @@ -328,11 +328,9 @@ class DateTime(Column): else: raise ValidationError("'{}' is not a datetime object".format(value)) epoch = datetime(1970, 1, 1, tzinfo=value.tzinfo) - offset = 0 - if epoch.tzinfo: - offset_delta = epoch.tzinfo.utcoffset(epoch) - offset = offset_delta.days*24*3600 + offset_delta.seconds - return long(((value - epoch).total_seconds() - offset) * 1000) + offset = epoch.tzinfo.utcoffset(epoch).total_seconds() if epoch.tzinfo else 0 + + return long(((value - epoch).total_seconds() - offset) * 1000) class Date(Column): @@ -402,12 +400,7 @@ class TimeUUID(UUID): global _last_timestamp epoch = datetime(1970, 1, 1, tzinfo=dt.tzinfo) - - offset = 0 - if epoch.tzinfo: - offset_delta = epoch.tzinfo.utcoffset(epoch) - offset = offset_delta.days*24*3600 + offset_delta.seconds - + offset = epoch.tzinfo.utcoffset(epoch).total_seconds() if epoch.tzinfo else 0 timestamp = (dt - epoch).total_seconds() - offset node = None diff --git a/cqlengine/functions.py b/cqlengine/functions.py index eee1fcab..13661845 100644 --- a/cqlengine/functions.py +++ b/cqlengine/functions.py @@ -54,8 +54,10 @@ class MinTimeUUID(BaseQueryFunction): super(MinTimeUUID, self).__init__(value) def get_value(self): - epoch = datetime(1970, 1, 1) - return long((self.value - epoch).total_seconds() * 1000) + epoch = datetime(1970, 1, 1, tzinfo=self.value.tzinfo) + offset = epoch.tzinfo.utcoffset(epoch).total_seconds() if epoch.tzinfo else 0 + + return long(((self.value - epoch).total_seconds() - offset) * 1000) def get_dict(self, column): return {self.identifier: self.get_value()} @@ -79,8 +81,10 @@ class MaxTimeUUID(BaseQueryFunction): super(MaxTimeUUID, self).__init__(value) def get_value(self): - epoch = datetime(1970, 1, 1) - return long((self.value - epoch).total_seconds() * 1000) + epoch = datetime(1970, 1, 1, tzinfo=self.value.tzinfo) + offset = epoch.tzinfo.utcoffset(epoch).total_seconds() if epoch.tzinfo else 0 + + return long(((self.value - epoch).total_seconds() - offset) * 1000) def get_dict(self, column): return {self.identifier: self.get_value()} diff --git a/cqlengine/tests/query/test_queryset.py b/cqlengine/tests/query/test_queryset.py index 117c87c1..9004a28b 100644 --- a/cqlengine/tests/query/test_queryset.py +++ b/cqlengine/tests/query/test_queryset.py @@ -11,6 +11,25 @@ from cqlengine.management import delete_table from cqlengine.models import Model from cqlengine import columns from cqlengine import query +from datetime import timedelta +from datetime import tzinfo + + +class TzOffset(tzinfo): + """Minimal implementation of a timezone offset to help testing with timezone + aware datetimes. + """ + def __init__(self, offset): + self._offset = timedelta(hours=offset) + + def utcoffset(self, dt): + return self._offset + + def tzname(self, dt): + return 'TzOffset: {}'.format(self._offset.hours) + + def dst(self, dt): + return timedelta(0) class TestModel(Model): test_id = columns.Integer(primary_key=True) @@ -515,6 +534,49 @@ class TestMinMaxTimeUUIDFunctions(BaseCassEngTestCase): super(TestMinMaxTimeUUIDFunctions, cls).tearDownClass() delete_table(TimeUUIDQueryModel) + def test_tzaware_datetime_support(self): + """Test that using timezone aware datetime instances works with the + MinTimeUUID/MaxTimeUUID functions. + """ + pk = uuid4() + midpoint_utc = datetime.utcnow().replace(tzinfo=TzOffset(0)) + midpoint_helsinki = midpoint_utc.astimezone(TzOffset(3)) + + # Assert pre-condition that we have the same logical point in time + assert midpoint_utc.utctimetuple() == midpoint_helsinki.utctimetuple() + assert midpoint_utc.timetuple() != midpoint_helsinki.timetuple() + + TimeUUIDQueryModel.create( + partition=pk, + time=columns.TimeUUID.from_datetime(midpoint_utc - timedelta(minutes=1)), + data='1') + + TimeUUIDQueryModel.create( + partition=pk, + time=columns.TimeUUID.from_datetime(midpoint_utc), + data='2') + + TimeUUIDQueryModel.create( + partition=pk, + time=columns.TimeUUID.from_datetime(midpoint_utc + timedelta(minutes=1)), + data='3') + + assert ['1', '2'] == [o.data for o in TimeUUIDQueryModel.filter( + TimeUUIDQueryModel.partition == pk, + TimeUUIDQueryModel.time <= functions.MaxTimeUUID(midpoint_utc))] + + assert ['1', '2'] == [o.data for o in TimeUUIDQueryModel.filter( + TimeUUIDQueryModel.partition == pk, + TimeUUIDQueryModel.time <= functions.MaxTimeUUID(midpoint_helsinki))] + + assert ['2', '3'] == [o.data for o in TimeUUIDQueryModel.filter( + TimeUUIDQueryModel.partition == pk, + TimeUUIDQueryModel.time >= functions.MinTimeUUID(midpoint_utc))] + + assert ['2', '3'] == [o.data for o in TimeUUIDQueryModel.filter( + TimeUUIDQueryModel.partition == pk, + TimeUUIDQueryModel.time >= functions.MinTimeUUID(midpoint_helsinki))] + def test_success_case(self): """ Test that the min and max time uuid functions work as expected """ pk = uuid4()