From 6d7b6250c55c1850a623cf47bd0d97738cf2e1a1 Mon Sep 17 00:00:00 2001
From: utapyngo <ut@pyngo.tom.ru>
Date: Mon, 20 Jan 2020 19:53:38 +0700
Subject: [PATCH] Support for using LIMIT N BY feature

See https://clickhouse.yandex/docs/en/query_language/select/#limit-by-clause
---
 src/infi/clickhouse_orm/query.py | 24 ++++++++++++++++++++++++
 tests/test_querysets.py          |  6 ++++++
 2 files changed, 30 insertions(+)

diff --git a/src/infi/clickhouse_orm/query.py b/src/infi/clickhouse_orm/query.py
index 16de5ba..9298deb 100644
--- a/src/infi/clickhouse_orm/query.py
+++ b/src/infi/clickhouse_orm/query.py
@@ -293,6 +293,8 @@ class QuerySet(object):
         self._grouping_with_totals = False
         self._fields = model_cls.fields().keys()
         self._limits = None
+        self._limit_by = None
+        self._limit_by_fields = None
         self._distinct = False
         self._final = False
 
@@ -332,6 +334,24 @@ class QuerySet(object):
             qs._limits = (start, stop - start)
             return qs
 
+    def limit_by(self, offset_limit, *fields):
+        if isinstance(offset_limit, six.integer_types):
+            # Single limit
+            assert offset_limit >= 0, 'negative limits are not supported'
+            qs = copy(self)
+            qs._limit_by = (0, offset_limit)
+            qs._limit_by_fields = fields
+            return qs
+        else:
+            # Offset, limit
+            offset = offset_limit[0]
+            limit = offset_limit[1]
+            assert offset >= 0 and limit >= 0, 'negative limits are not supported'
+            qs = copy(self)
+            qs._limit_by = (offset, limit)
+            qs._limit_by_fields = fields
+            return qs
+
     def select_fields_as_sql(self):
         """
         Returns the selected fields or expressions as a SQL string.
@@ -369,6 +389,10 @@ class QuerySet(object):
         if self._limits:
             sql += '\nLIMIT %d, %d' % self._limits
 
+        if self._limit_by:
+            sql += '\nLIMIT %d, %d' % self._limit_by
+            sql += ' BY %s' % comma_join('`%s`' % field for field in self._limit_by_fields)
+
         return sql
 
     def order_by_as_sql(self):
diff --git a/tests/test_querysets.py b/tests/test_querysets.py
index b17933b..e1a7f08 100644
--- a/tests/test_querysets.py
+++ b/tests/test_querysets.py
@@ -432,6 +432,12 @@ class AggregateTestCase(TestCaseWithData):
         qs = Mdl.objects_in(self.database).filter(the__next__number__gt=1)
         self.assertEqual(qs.conditions_as_sql(), 'the__next__number > 1')
 
+    def test_limit_by(self):
+        qs = Person.objects_in(self.database).aggregate('first_name', 'last_name', 'height', n='count()').\
+            order_by('first_name', '-height').limit_by(1, 'first_name')
+        self.assertEqual(qs.count(), 94)
+        self.assertEqual(list(qs)[89].last_name, 'Bowen')
+
 
 Color = Enum('Color', u'red blue green yellow brown white black')