Repository: kayak/pypika
Pull request #91 (2017-12-24 21:38:53)

Added ClickHouse WITH TOTALS and PREWHERE statements


Patch:
diff --git a/pypika/dialects.py b/pypika/dialects.py
index 75dc4bb..a4c02ad 100644
--- a/pypika/dialects.py
+++ b/pypika/dialects.py
@@ -63,3 +63,13 @@ class MSSQLQuery(Query):
     @classmethod
     def _builder(cls):
         return QueryBuilder(dialect=Dialects.MSSQL)
+
+
+class ClickHouseQuery(Query):
+    """
+    Defines a query class for use with Yandex ClickHouse.
+    """
+
+    @classmethod
+    def _builder(cls):
+        return QueryBuilder(dialect=Dialects.CLICKHOUSE)
diff --git a/pypika/enums.py b/pypika/enums.py
index 9679664..9a4472d 100644
--- a/pypika/enums.py
+++ b/pypika/enums.py
@@ -77,6 +77,7 @@ class SqlTypes(Enum):
 
 class Dialects(Enum):
     VERTICA = 'vertica'
+    CLICKHOUSE = 'clickhouse'
     ORACLE = 'oracle'
     MSSQL = 'mssql'
     MYSQL = 'mysql'
diff --git a/pypika/queries.py b/pypika/queries.py
index 12347d7..1a81008 100644
--- a/pypika/queries.py
+++ b/pypika/queries.py
@@ -191,7 +191,9 @@ class QueryBuilder(Selectable, Term):
         self._ignore = False
 
         self._wheres = None
+        self._prewheres = None
         self._groupbys = []
+        self._with_totals = False
         self._havings = None
         self._orderbys = []
         self._joins = []
@@ -305,6 +307,15 @@ class QueryBuilder(Selectable, Term):
     def ignore(self):
         self._ignore = True
 
+    @builder
+    def prewhere(self, criterion):
+        self._validate_term(criterion)
+
+        if self._prewheres:
+            self._prewheres &= criterion
+        else:
+            self._prewheres = criterion
+
     @builder
     def where(self, criterion):
         self._validate_term(criterion)
@@ -332,6 +343,10 @@ class QueryBuilder(Selectable, Term):
             self._validate_term(term)
             self._groupbys.append(term)
 
+    @builder
+    def with_totals(self):
+        self._with_totals = True
+
     @builder
     def rollup(self, *terms, **kwargs):
         for_mysql = 'mysql' == kwargs.get('vendor')
@@ -542,6 +557,9 @@ class QueryBuilder(Selectable, Term):
             querystring += " " + " ".join(join.get_sql(**kwargs)
                                           for join in self._joins)
 
+        if self._prewheres:
+            querystring += self._prewhere_sql(**kwargs)
+
         if self._wheres:
             querystring += self._where_sql(**kwargs)
 
@@ -625,6 +643,10 @@ class QueryBuilder(Selectable, Term):
             for clause in self._from
         ))
 
+    def _prewhere_sql(self, quote_char=None, **kwargs):
+        return ' PREWHERE {prewhere}'.format(
+            prewhere=self._prewheres.get_sql(quote_char=quote_char, subquery=True, **kwargs))
+
     def _where_sql(self, quote_char=None, **kwargs):
         return ' WHERE {where}'.format(where=self._wheres.get_sql(quote_char=quote_char, subquery=True, **kwargs))
 
@@ -647,7 +669,10 @@ class QueryBuilder(Selectable, Term):
             else:
                 clauses.append(field.get_sql(quote_char=quote_char, **kwargs))
 
-        return ' GROUP BY {groupby}'.format(groupby=','.join(clauses))
+        sql = ' GROUP BY {groupby}'.format(groupby=','.join(clauses))
+        if self._with_totals:
+            return sql + ' WITH TOTALS'
+        return sql
 
     def _orderby_sql(self, quote_char=None, **kwargs):
         """


Test patch:
diff --git a/pypika/tests/test_selects.py b/pypika/tests/test_selects.py
index 9dd80b2..73fe402 100644
--- a/pypika/tests/test_selects.py
+++ b/pypika/tests/test_selects.py
@@ -229,6 +229,22 @@ class WhereTests(unittest.TestCase):
         self.assertEqual("SELECT * FROM \"abc\" WHERE \"foo\" REGEX '^b'", str(q))
 
 
+class PreWhereTests(WhereTests):
+    t = Table('abc')
+
+    def test_prewhere_field_equals(self):
+        q1 = Query.from_(self.t).select('*').prewhere(self.t.foo == self.t.bar)
+        q2 = Query.from_(self.t).select('*').prewhere(self.t.foo.eq(self.t.bar))
+
+        self.assertEqual('SELECT * FROM "abc" PREWHERE "foo"="bar"', str(q1))
+        self.assertEqual('SELECT * FROM "abc" PREWHERE "foo"="bar"', str(q2))
+
+    def test_where_and_prewhere(self):
+        q = Query.from_(self.t).select('*').prewhere(self.t.foo == self.t.bar).where(self.t.foo == self.t.bar)
+
+        self.assertEqual('SELECT * FROM "abc" PREWHERE "foo"="bar" WHERE "foo"="bar"', str(q))
+
+
 class GroupByTests(unittest.TestCase):
     t = Table('abc')
     maxDiff = None
@@ -331,6 +347,16 @@ class GroupByTests(unittest.TestCase):
 
         self.assertEqual('SELECT "foo" FROM "abc" GROUP BY "foo"', str(q))
 
+    def test_group_by__single_with_totals(self):
+        q = Query.from_(self.t).groupby(self.t.foo).select(self.t.foo).with_totals()
+
+        self.assertEqual('SELECT "foo" FROM "abc" GROUP BY "foo" WITH TOTALS', str(q))
+
+    def test_groupby__multi_with_totals(self):
+        q = Query.from_(self.t).groupby(self.t.foo, self.t.bar).select(self.t.foo, self.t.bar).with_totals()
+
+        self.assertEqual('SELECT "foo","bar" FROM "abc" GROUP BY "foo","bar" WITH TOTALS', str(q))
+
 
 class HavingTests(unittest.TestCase):
     table_abc, table_efg = Tables('abc', 'efg')