From 00bf7eeb755f9df950a149908eade25637ad09e8 Mon Sep 17 00:00:00 2001
From: Itai Shirav <itais@infinidat.com>
Date: Sat, 16 May 2020 09:24:31 +0300
Subject: [PATCH] count() and paginate() - accept conditions as expressions and
 Q objects

---
 src/infi/clickhouse_orm/database.py | 10 ++++++++--
 tests/test_database.py              | 14 ++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/src/infi/clickhouse_orm/database.py b/src/infi/clickhouse_orm/database.py
index 54f153f..d42e224 100644
--- a/src/infi/clickhouse_orm/database.py
+++ b/src/infi/clickhouse_orm/database.py
@@ -249,9 +249,12 @@ class Database(object):
         - `model_class`: the model to count.
         - `conditions`: optional SQL conditions (contents of the WHERE clause).
         '''
+        from infi.clickhouse_orm.query import Q
         query = 'SELECT count() FROM $table'
         if conditions:
-            query += ' WHERE ' + conditions
+            if isinstance(conditions, Q):
+                conditions = conditions.to_sql(model_class)
+            query += ' WHERE ' + str(conditions)
         query = self._substitute(query, model_class)
         r = self._send(query)
         return int(r.text) if r.text else 0
@@ -303,6 +306,7 @@ class Database(object):
         The result is a namedtuple containing `objects` (list), `number_of_objects`,
         `pages_total`, `number` (of the current page), and `page_size`.
         '''
+        from infi.clickhouse_orm.query import Q
         count = self.count(model_class, conditions)
         pages_total = int(ceil(count / float(page_size)))
         if page_num == -1:
@@ -312,7 +316,9 @@ class Database(object):
         offset = (page_num - 1) * page_size
         query = 'SELECT * FROM $table'
         if conditions:
-            query += ' WHERE ' + conditions
+            if isinstance(conditions, Q):
+                conditions = conditions.to_sql(model_class)
+            query += ' WHERE ' + str(conditions)
         query += ' ORDER BY %s' % order_by
         query += ' LIMIT %d, %d' % (offset, page_size)
         query = self._substitute(query, model_class)
diff --git a/tests/test_database.py b/tests/test_database.py
index 30ab1c3..15d928f 100644
--- a/tests/test_database.py
+++ b/tests/test_database.py
@@ -8,6 +8,7 @@ from infi.clickhouse_orm.models import Model
 from infi.clickhouse_orm.engines import Memory
 from infi.clickhouse_orm.fields import *
 from infi.clickhouse_orm.funcs import F
+from infi.clickhouse_orm.query import Q
 from .base_test_with_data import *
 
 
@@ -47,9 +48,14 @@ class DatabaseTestCase(TestCaseWithData):
     def test_count(self):
         self.database.insert(self._sample_data())
         self.assertEqual(self.database.count(Person), 100)
+        # Conditions as string
         self.assertEqual(self.database.count(Person, "first_name = 'Courtney'"), 2)
         self.assertEqual(self.database.count(Person, "birthday > '2000-01-01'"), 22)
         self.assertEqual(self.database.count(Person, "birthday < '1970-03-01'"), 0)
+        # Conditions as expression
+        self.assertEqual(self.database.count(Person, Person.birthday > datetime.date(2000, 1, 1)), 22)
+        # Conditions as Q object
+        self.assertEqual(self.database.count(Person, Q(birthday__gt=datetime.date(2000, 1, 1))), 22)
 
     def test_select(self):
         self._insert_and_check(self._sample_data(), len(data))
@@ -146,8 +152,15 @@ class DatabaseTestCase(TestCaseWithData):
 
     def test_pagination_with_conditions(self):
         self._insert_and_check(self._sample_data(), len(data))
+        # Conditions as string
         page = self.database.paginate(Person, 'first_name, last_name', 1, 100, conditions="first_name < 'Ava'")
         self.assertEqual(page.number_of_objects, 10)
+        # Conditions as expression
+        page = self.database.paginate(Person, 'first_name, last_name', 1, 100, conditions=Person.first_name < 'Ava')
+        self.assertEqual(page.number_of_objects, 10)
+        # Conditions as Q object
+        page = self.database.paginate(Person, 'first_name, last_name', 1, 100, conditions=Q(first_name__lt='Ava'))
+        self.assertEqual(page.number_of_objects, 10)
 
     def test_special_chars(self):
         s = u'אבגד \\\'"`,.;éåäöšž\n\t\0\b\r'
@@ -263,3 +276,4 @@ class DatabaseTestCase(TestCaseWithData):
             self.assertEqual(model.table_name(), row.name)
             # Read a few records
             list(model.objects_in(self.database)[:10])
+