From 7f38f8c6544a7bd1c242dd7872575633e6aebab7 Mon Sep 17 00:00:00 2001
From: M1ha <m1hacka@mail.ru>
Date: Sun, 2 Dec 2018 12:14:19 +0500
Subject: [PATCH] Added fina() QuerySet method Added tests for ite Not tested
 in case of https://github.com/Infinidat/infi.clickhouse_orm/issues/93

---
 src/infi/clickhouse_orm/query.py | 20 ++++++++++++++++++--
 tests/base_test_with_data.py     |  2 +-
 tests/test_querysets.py          | 32 ++++++++++++++++++++++++++++++++
 3 files changed, 51 insertions(+), 3 deletions(-)

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()