Skip to content

Commit e351a30

Browse files
jayvynljayvynl
committed
feat: support prewhere, basically completed, need unit tests
1 parent 6a33a37 commit e351a30

File tree

3 files changed

+237
-14
lines changed

3 files changed

+237
-14
lines changed

‎CHANGELOG.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
### unreleased
1+
### 1.2.0
22

33
- feat: #72 support window functions.
4-
- feat: #80 support [prewhere](https://clickhouse.com/docs/en/sql-reference/statements/select/prewhere).
4+
- feat: #80 support [prewhere clause](https://clickhouse.com/docs/en/sql-reference/statements/select/prewhere).
55

66
### 1.1.7
77

‎clickhouse_backend/models/sql/compiler.py‎

Lines changed: 229 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
import itertools
22

3+
from django.core.exceptions import EmptyResultSet
4+
from django.db import NotSupportedError
35
from django.db.models.fields import AutoFieldMixin
46
from django.db.models.sql import compiler
57

@@ -9,6 +11,13 @@
911

1012
if compat.dj_ge42:
1113
from django.core.exceptions import FullResultSet
14+
else:
15+
16+
class FullResultSet(Exception):
17+
"""A database query predicate is matches everything."""
18+
19+
pass
20+
1221

1322
# Max rows you can insert using expression as value.
1423
MAX_ROWS_INSERT_USE_EXPRESSION = 1000
@@ -39,26 +48,234 @@ def _add_settings_sql(self, sql, params):
3948
return sql, params
4049

4150
def _compile_where(self, table):
42-
if compat.dj_ge42:
43-
try:
44-
where, params = self.compile(self.query.where)
45-
except FullResultSet:
46-
where, params = "", ()
47-
else:
51+
try:
4852
where, params = self.compile(self.query.where)
53+
except FullResultSet:
54+
where, params = "", ()
4955
if where:
5056
where = where.replace(table + ".", "")
5157
else:
52-
where = "1"
58+
where = "TRUE"
5359
return where, params
5460

5561

5662
class SQLCompiler(ClickhouseMixin, compiler.SQLCompiler):
57-
def as_sql(self, *args, **kwargs):
58-
sql, params = super().as_sql(*args, **kwargs)
59-
sql, params = self._add_settings_sql(sql, params)
60-
sql, params = self._add_explain_sql(sql, params)
61-
return sql, params
63+
def pre_sql_setup(self, with_col_aliases=False):
64+
"""
65+
Do any necessary class setup immediately prior to producing SQL. This
66+
is for things that can't necessarily be done in __init__ because we
67+
might not have all the pieces in place at that time.
68+
"""
69+
if compat.dj_ge42:
70+
self.setup_query(with_col_aliases=with_col_aliases)
71+
(
72+
self.where,
73+
self.having,
74+
self.qualify,
75+
) = self.query.where.split_having_qualify(
76+
must_group_by=self.query.group_by is not None
77+
)
78+
(
79+
self.prewhere,
80+
prehaving,
81+
prequalify,
82+
) = self.query.prewhere.split_having_qualify(
83+
must_group_by=self.query.group_by is not None
84+
)
85+
# Check before ClickHouse complain.
86+
# DB::Exception: Window function is found in PREWHERE in query. (ILLEGAL_AGGREGATION)
87+
if prequalify:
88+
raise NotSupportedError(
89+
"Window function is disallowed in the prewhere clause."
90+
)
91+
else:
92+
self.setup_query()
93+
self.where, self.having = self.query.where.split_having()
94+
self.prewhere, prehaving = self.query.where.split_having()
95+
# Check before ClickHouse complain.
96+
# DB::Exception: Aggregate function is found in PREWHERE in query. (ILLEGAL_AGGREGATION)
97+
if prehaving:
98+
raise NotSupportedError(
99+
"Aggregate function is disallowed in the prewhere clause."
100+
)
101+
order_by = self.get_order_by()
102+
extra_select = self.get_extra_select(order_by, self.select)
103+
self.has_extra_select = bool(extra_select)
104+
group_by = self.get_group_by(self.select + extra_select, order_by)
105+
return extra_select, order_by, group_by
106+
107+
def as_sql(self, with_limits=True, with_col_aliases=False):
108+
"""
109+
Create the SQL for this query. Return the SQL string and list of
110+
parameters.
111+
112+
If 'with_limits' is False, any limit/offset information is not included
113+
in the query.
114+
"""
115+
refcounts_before = self.query.alias_refcount.copy()
116+
try:
117+
combinator = self.query.combinator
118+
if compat.dj_ge42:
119+
extra_select, order_by, group_by = self.pre_sql_setup(
120+
with_col_aliases=with_col_aliases or bool(combinator),
121+
)
122+
else:
123+
extra_select, order_by, group_by = self.pre_sql_setup()
124+
# Is a LIMIT/OFFSET clause needed?
125+
with_limit_offset = with_limits and self.query.is_sliced
126+
if combinator:
127+
result, params = self.get_combinator_sql(
128+
combinator, self.query.combinator_all
129+
)
130+
# Django >= 4.2 have this branch
131+
elif compat.dj_ge42 and self.qualify:
132+
result, params = self.get_qualify_sql()
133+
order_by = None
134+
else:
135+
distinct_fields, distinct_params = self.get_distinct()
136+
# This must come after 'select', 'ordering', and 'distinct'
137+
# (see docstring of get_from_clause() for details).
138+
from_, f_params = self.get_from_clause()
139+
try:
140+
where, w_params = (
141+
self.compile(self.where) if self.where is not None else ("", [])
142+
)
143+
except EmptyResultSet:
144+
if compat.dj_ge42 and self.elide_empty:
145+
raise
146+
# Use a predicate that's always False.
147+
where, w_params = "FALSE", []
148+
except FullResultSet:
149+
where, w_params = "", []
150+
try:
151+
having, h_params = (
152+
self.compile(self.having)
153+
if self.having is not None
154+
else ("", [])
155+
)
156+
except FullResultSet:
157+
having, h_params = "", []
158+
# v1.2.0 new feature, support prewhere clause.
159+
# refer https://clickhouse.com/docs/en/sql-reference/statements/select/prewhere
160+
try:
161+
prewhere, p_params = (
162+
self.compile(self.prewhere)
163+
if self.prewhere is not None
164+
else ("", [])
165+
)
166+
except EmptyResultSet:
167+
if compat.dj_ge42 and self.elide_empty:
168+
raise
169+
# Use a predicate that's always False.
170+
prewhere, p_params = "FALSE", []
171+
except FullResultSet:
172+
prewhere, p_params = "", []
173+
result = ["SELECT"]
174+
params = []
175+
176+
if self.query.distinct:
177+
distinct_result, distinct_params = self.connection.ops.distinct_sql(
178+
distinct_fields,
179+
distinct_params,
180+
)
181+
result += distinct_result
182+
params += distinct_params
183+
184+
out_cols = []
185+
for _, (s_sql, s_params), alias in self.select + extra_select:
186+
if alias:
187+
s_sql = "%s AS %s" % (
188+
s_sql,
189+
self.connection.ops.quote_name(alias),
190+
)
191+
params.extend(s_params)
192+
out_cols.append(s_sql)
193+
194+
result += [", ".join(out_cols)]
195+
if from_:
196+
result += ["FROM", *from_]
197+
params.extend(f_params)
198+
199+
if prewhere:
200+
result.append("PREWHERE %s" % prewhere)
201+
params.extend(p_params)
202+
203+
if where:
204+
result.append("WHERE %s" % where)
205+
params.extend(w_params)
206+
207+
grouping = []
208+
for g_sql, g_params in group_by:
209+
grouping.append(g_sql)
210+
params.extend(g_params)
211+
if grouping:
212+
if distinct_fields:
213+
raise NotImplementedError(
214+
"annotate() + distinct(fields) is not implemented."
215+
)
216+
result.append("GROUP BY %s" % ", ".join(grouping))
217+
if self._meta_ordering:
218+
order_by = None
219+
if having:
220+
result.append("HAVING %s" % having)
221+
params.extend(h_params)
222+
223+
if order_by:
224+
ordering = []
225+
for _, (o_sql, o_params, _) in order_by:
226+
ordering.append(o_sql)
227+
params.extend(o_params)
228+
order_by_sql = "ORDER BY %s" % ", ".join(ordering)
229+
result.append(order_by_sql)
230+
231+
if with_limit_offset:
232+
result.append(
233+
self.connection.ops.limit_offset_sql(
234+
self.query.low_mark, self.query.high_mark
235+
)
236+
)
237+
238+
if self.query.subquery and extra_select:
239+
# If the query is used as a subquery, the extra selects would
240+
# result in more columns than the left-hand side expression is
241+
# expecting. This can happen when a subquery uses a combination
242+
# of order_by() and distinct(), forcing the ordering expressions
243+
# to be selected as well. Wrap the query in another subquery
244+
# to exclude extraneous selects.
245+
sub_selects = []
246+
sub_params = []
247+
for index, (select, _, alias) in enumerate(self.select, start=1):
248+
if alias:
249+
sub_selects.append(
250+
"%s.%s"
251+
% (
252+
self.connection.ops.quote_name("subquery"),
253+
self.connection.ops.quote_name(alias),
254+
)
255+
)
256+
else:
257+
select_clone = select.relabeled_clone(
258+
{select.alias: "subquery"}
259+
)
260+
subselect, subparams = select_clone.as_sql(
261+
self, self.connection
262+
)
263+
sub_selects.append(subselect)
264+
sub_params.extend(subparams)
265+
sql = "SELECT %s FROM (%s) subquery" % (
266+
", ".join(sub_selects),
267+
" ".join(result),
268+
)
269+
params = tuple(sub_params + params)
270+
else:
271+
sql = " ".join(result)
272+
params = tuple(params)
273+
sql, params = self._add_settings_sql(sql, params)
274+
sql, params = self._add_explain_sql(sql, params)
275+
return sql, params
276+
finally:
277+
# Finally do cleanup - get rid of the joins we created above.
278+
self.query.reset_refcounts(refcounts_before)
62279

63280

64281
class SQLInsertCompiler(compiler.SQLInsertCompiler):

‎clickhouse_backend/models/sql/query.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,12 @@ def add_prewhere(self, q_object):
5353
self.prewhere.add(clause, AND)
5454
self.demote_joins(existing_inner)
5555

56+
if not compat.dj_ge42:
57+
58+
@property
59+
def is_sliced(self):
60+
return self.low_mark != 0 or self.high_mark is not None
61+
5662

5763
def clone_decorator(cls):
5864
old_clone = cls.clone

0 commit comments

Comments
 (0)