diff --git a/src/infi/clickhouse_orm/query.py b/src/infi/clickhouse_orm/query.py index 1035dd4..5dfd660 100644 --- a/src/infi/clickhouse_orm/query.py +++ b/src/infi/clickhouse_orm/query.py @@ -3,6 +3,8 @@ import six import pytz from copy import copy from math import ceil + +from .engines import CollapsingMergeTree from .utils import comma_join @@ -243,6 +245,7 @@ class QuerySet(object): self._fields = model_cls.fields().keys() self._limits = None self._distinct = False + self._final = False def __iter__(self): """ @@ -290,9 +293,10 @@ class QuerySet(object): fields = comma_join('`%s`' % field for field in self._fields) ordering = '\nORDER BY ' + self.order_by_as_sql() if self._order_by else '' limit = '\nLIMIT %d, %d' % self._limits if self._limits else '' - params = (distinct, fields, self._model_cls.table_name(), + final = ' FINAL' if self._final else '' + params = (distinct, fields, self._model_cls.table_name(), final, self.conditions_as_sql(), ordering, limit) - return u'SELECT %s%s\nFROM `%s`\nWHERE %s%s%s' % params + return u'SELECT %s%s\nFROM `%s`%s\nWHERE %s%s%s' % params def order_by_as_sql(self): """ @@ -399,6 +403,18 @@ class QuerySet(object): qs._distinct = True return qs + def final(self): + """ + Adds a FINAL modifier to table, meaning data will be collapsed to final version. + Can be used with CollapsingMergeTree engine only + """ + if not issubclass(self._model_cls, CollapsingMergeTree): + raise TypeError('final() method can be used only with CollapsingMergeTree engine') + + qs = copy(self) + qs._final = True + return qs + def aggregate(self, *args, **kwargs): """ Returns an `AggregateQuerySet` over this query, with `args` serving as diff --git a/tests/base_test_with_data.py b/tests/base_test_with_data.py index 9f183cf..c3fc376 100644 --- a/tests/base_test_with_data.py +++ b/tests/base_test_with_data.py @@ -148,4 +148,4 @@ data = [ {"first_name": "Whitney", "last_name": "Scott", "birthday": "1971-07-04", "height": "1.70"}, {"first_name": "Wynter", "last_name": "Garcia", "birthday": "1975-01-10", "height": "1.69"}, {"first_name": "Yolanda", "last_name": "Duke", "birthday": "1997-02-25", "height": "1.74"} -]; +] diff --git a/tests/test_querysets.py b/tests/test_querysets.py index a4fef14..0fcc871 100644 --- a/tests/test_querysets.py +++ b/tests/test_querysets.py @@ -141,6 +141,20 @@ class QuerySetTestCase(TestCaseWithData): SampleModel(timestamp=now, num=4, color=Color.white), ]) + def _insert_sample_collapsing_model(self): + self.database.create_table(SampleCollapsingModel) + now = datetime.now() + self.database.insert([ + SampleCollapsingModel(timestamp=now, num=1, color=Color.red), + SampleCollapsingModel(timestamp=now, num=2, color=Color.red), + SampleCollapsingModel(timestamp=now, num=2, color=Color.red, sign=-1), + SampleCollapsingModel(timestamp=now, num=2, color=Color.green), + SampleCollapsingModel(timestamp=now, num=3, color=Color.white), + SampleCollapsingModel(timestamp=now, num=4, color=Color.white, sign=1), + SampleCollapsingModel(timestamp=now, num=4, color=Color.white, sign=-1), + SampleCollapsingModel(timestamp=now, num=4, color=Color.blue, sign=1), + ]) + def test_filter_enum_field(self): self._insert_sample_model() qs = SampleModel.objects_in(self.database) @@ -249,6 +263,17 @@ class QuerySetTestCase(TestCaseWithData): self._test_qs(qs[70:80], 10) self._test_qs(qs[80:], 20) + def test_final(self): + # Final can be used with CollapsingMergeTree engine only + with self.assertRaises(TypeError): + Person.objects_in(self.database).final() + + self._insert_sample_collapsing_model() + res = list(SampleCollapsingModel.objects_in(self.database).final().order_by('num')) + self.assertEqual(4, len(res)) + for item, exp_color in zip(res, (Color.red, Color.green, Color.white, Color.blue)): + self.assertEqual(exp_color, item.color) + class AggregateTestCase(TestCaseWithData): @@ -392,6 +417,13 @@ class SampleModel(Model): engine = MergeTree('materialized_date', ('materialized_date',)) +class SampleCollapsingModel(SampleModel): + + sign = UInt8Field(default=1) + + engine = CollapsingMergeTree('materialized_date', ('num',), 'sign') + + class Numbers(Model): number = UInt64Field()