Skip to content

Commit b5c9d86

Browse files
xzkostyanxzkostyan
authored
Merge pull request #383 from qyber1/add-with-fill
Add with fill
2 parents 4f1a43f + 4da140d commit b5c9d86

File tree

5 files changed

+670
-1
lines changed

5 files changed

+670
-1
lines changed

‎clickhouse_sqlalchemy/drivers/base.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
from .. import Table, types
1919
from .. import engines
20+
from ..ext.clauses import WithFill
2021
from ..util import compat
2122

2223
# Column specifications
@@ -174,6 +175,23 @@ def visit_lambda(self, lambda_, **kw):
174175

175176
return text
176177

178+
def visit_with_fill(self, with_fill: WithFill, **kw):
179+
column = with_fill.column
180+
text = f'{column} WITH FILL'
181+
182+
if with_fill.from_ is not None:
183+
184+
text += ' FROM ' + self.process(with_fill.from_, **kw)
185+
186+
if with_fill.to is not None:
187+
188+
text += ' TO ' + self.process(with_fill.to, **kw)
189+
190+
if with_fill.step is not None:
191+
text += ' STEP ' + self.process(with_fill.step, **kw)
192+
193+
return text
194+
177195
def visit_extract(self, extract, **kw):
178196
field = self.extract_map.get(extract.field, extract.field)
179197
column = self.process(extract.expr, **kw)

‎clickhouse_sqlalchemy/ext/clauses.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,5 +57,17 @@ def __init__(self, func):
5757
self.func = func
5858

5959

60+
class WithFill(ColumnElement):
61+
"""Represent a ``WITH FILL`` expression."""
62+
__visit_name__ = 'with_fill'
63+
64+
def __init__(self, column, from_=None, to=None, step=None):
65+
66+
self.column = column
67+
self.from_ = from_
68+
self.to = to
69+
self.step = step
70+
71+
6072
class ArrayJoin(ClauseList):
6173
__visit_name__ = 'array_join'

‎clickhouse_sqlalchemy/types/common.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,10 +107,22 @@ class Float64(Float):
107107
class Date(types.Date, ClickHouseTypeEngine):
108108
__visit_name__ = 'date'
109109

110+
def literal_processor(self, dialect):
111+
def process(value):
112+
return "'%s'" % value
113+
114+
return process
115+
110116

111117
class DateTime(types.Date, ClickHouseTypeEngine):
112118
__visit_name__ = 'datetime'
113119

120+
def literal_processor(self, dialect):
121+
def process(value):
122+
return "'%s'" % value
123+
124+
return process
125+
114126

115127
class DateTime64(DateTime, ClickHouseTypeEngine):
116128
__visit_name__ = 'datetime64'
@@ -120,6 +132,12 @@ def __init__(self, precision=3, timezone=None):
120132
self.timezone = timezone
121133
super(DateTime64, self).__init__()
122134

135+
def literal_processor(self, dialect):
136+
def process(value):
137+
return "'%s'" % value
138+
139+
return process
140+
123141

124142
class Enum(types.Enum, ClickHouseTypeEngine):
125143
__visit_name__ = 'enum'

‎tests/orm/test_select.py‎

Lines changed: 296 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
from datetime import date, datetime as dt
12
from sqlalchemy import Column, exc, func, literal, select, text, tuple_
23

34
from clickhouse_sqlalchemy import types, Table, engines
@@ -146,6 +147,301 @@ def test_lambda_functions(self):
146147
") AS test"
147148
)
148149

150+
def test_with_fill(self):
151+
columns = [Column('y', types.Int32), Column('z', types.Int32)]
152+
table = self._make_table(*columns)
153+
154+
query = self.session.query(table.c.x) \
155+
.order_by(WithFill(table.c.x))
156+
157+
self.assertEqual(
158+
self.compile(query),
159+
'SELECT t1.x AS t1_x FROM t1 ORDER BY t1.x '
160+
'WITH FILL'
161+
)
162+
163+
def test_with_fill_from(self):
164+
columns = [Column('y', types.Int32), Column('z', types.Int32)]
165+
table = self._make_table(*columns)
166+
167+
query = self.session.query(table.c.x) \
168+
.order_by(WithFill(table.c.x, from_=literal(2)))
169+
170+
self.assertEqual(
171+
self.compile(query),
172+
'SELECT t1.x AS t1_x FROM t1 ORDER BY t1.x '
173+
'WITH FILL FROM %(param_1)s'
174+
)
175+
176+
self.assertEqual(
177+
self.compile(query, literal_binds=True),
178+
'SELECT t1.x AS t1_x FROM t1 ORDER BY t1.x '
179+
'WITH FILL FROM 2'
180+
)
181+
182+
def test_with_fill_from_to(self):
183+
columns = [Column('y', types.Int32), Column('z', types.Int32)]
184+
table = self._make_table(*columns)
185+
186+
query = self.session.query(table.c.x) \
187+
.order_by(WithFill(table.c.x, from_=literal(2), to=literal(5)))
188+
189+
self.assertEqual(
190+
self.compile(query),
191+
'SELECT t1.x AS t1_x FROM t1 ORDER BY t1.x '
192+
'WITH FILL FROM %(param_1)s TO %(param_2)s'
193+
)
194+
195+
self.assertEqual(
196+
self.compile(query, literal_binds=True),
197+
'SELECT t1.x AS t1_x FROM t1 ORDER BY t1.x '
198+
'WITH FILL FROM 2 TO 5'
199+
)
200+
201+
def test_with_fill_from_to_step(self):
202+
columns = [Column('y', types.Int32), Column('z', types.Int32)]
203+
table = self._make_table(*columns)
204+
205+
query = self.session.query(table.c.x) \
206+
.order_by(
207+
WithFill(
208+
table.c.x, from_=literal(1), to=literal(5), step=literal(1)
209+
)
210+
)
211+
212+
self.assertEqual(
213+
self.compile(query),
214+
'SELECT t1.x AS t1_x FROM t1 ORDER BY t1.x '
215+
'WITH FILL FROM %(param_1)s TO %(param_2)s STEP %(param_3)s'
216+
)
217+
218+
self.assertEqual(
219+
self.compile(query, literal_binds=True),
220+
'SELECT t1.x AS t1_x FROM t1 ORDER BY t1.x '
221+
'WITH FILL FROM 1 TO 5 STEP 1'
222+
)
223+
224+
def test_with_fill_from_to_step_multiply(self):
225+
columns = [Column('y', types.Int32), Column('z', types.Int32)]
226+
table = self._make_table(*columns)
227+
228+
query = self.session.query(table.c.x) \
229+
.order_by(
230+
WithFill(
231+
table.c.x, from_=literal(1), to=literal(5), step=literal(1)
232+
),
233+
WithFill(
234+
table.c.z, from_=literal(2), to=literal(6), step=literal(1)
235+
)
236+
)
237+
238+
self.assertEqual(
239+
self.compile(query),
240+
'SELECT t1.x AS t1_x FROM t1 ORDER BY t1.x '
241+
'WITH FILL '
242+
'FROM %(param_1)s TO %(param_2)s STEP %(param_3)s, t1.z '
243+
'WITH FILL '
244+
'FROM %(param_4)s TO %(param_5)s STEP %(param_6)s'
245+
)
246+
self.assertEqual(
247+
self.compile(query, literal_binds=True),
248+
'SELECT t1.x AS t1_x FROM t1 ORDER BY t1.x '
249+
'WITH FILL '
250+
'FROM 1 TO 5 STEP 1, '
251+
't1.z WITH FILL FROM 2 TO 6 STEP 1'
252+
)
253+
254+
def test_with_fill_from_date(self):
255+
columns = [Column('y', types.Int32), Column('date', types.DateTime)]
256+
table = self._make_table(*columns)
257+
258+
date_from = date(year=2025, month=1, day=1)
259+
query = self.session.query(table.c.x) \
260+
.order_by(
261+
WithFill(
262+
table.c.date,
263+
from_=func.toDateTime(
264+
literal(date_from, type_=types.DateTime)
265+
)
266+
)
267+
)
268+
269+
self.assertEqual(
270+
self.compile(query),
271+
"SELECT t1.x AS t1_x FROM t1 ORDER BY t1.date "
272+
"WITH FILL "
273+
"FROM toDateTime(%(param_1)s)"
274+
)
275+
276+
self.assertEqual(
277+
self.compile(query, literal_binds=True),
278+
"SELECT t1.x AS t1_x FROM t1 ORDER BY t1.date "
279+
"WITH FILL "
280+
"FROM toDateTime('2025-01-01')"
281+
)
282+
dt_from = dt(year=2025, month=1, day=1, hour=1, minute=1, second=1)
283+
query = self.session.query(table.c.x) \
284+
.order_by(
285+
WithFill(
286+
table.c.date, from_=func.toDateTime(
287+
literal(dt_from, type_=types.Date))
288+
)
289+
)
290+
291+
self.assertEqual(
292+
self.compile(query),
293+
"SELECT t1.x AS t1_x FROM t1 ORDER BY t1.date "
294+
"WITH FILL "
295+
"FROM toDateTime(%(param_1)s)"
296+
)
297+
298+
self.assertEqual(
299+
self.compile(query, literal_binds=True),
300+
"SELECT t1.x AS t1_x FROM t1 ORDER BY t1.date "
301+
"WITH FILL "
302+
"FROM toDateTime('2025-01-01 01:01:01')"
303+
)
304+
305+
def test_with_fill_from_to_step_date(self):
306+
columns = [Column('y', types.Int32), Column('date', types.Date)]
307+
table = self._make_table(*columns)
308+
309+
date_from = date(year=2025, month=1, day=1)
310+
date_to = date(year=2025, month=2, day=1)
311+
query = self.session.query(table.c.x) \
312+
.order_by(
313+
WithFill(
314+
table.c.date,
315+
from_=func.toDate(literal(date_from, type_=types.Date)),
316+
to=func.toDate(literal(date_to, type_=types.Date)),
317+
step=literal(1)
318+
)
319+
)
320+
321+
self.assertEqual(
322+
self.compile(query),
323+
"SELECT t1.x AS t1_x FROM t1 ORDER BY t1.date "
324+
"WITH FILL "
325+
"FROM toDate(%(param_1)s) "
326+
"TO toDate(%(param_2)s) "
327+
"STEP %(param_3)s"
328+
)
329+
330+
self.assertEqual(
331+
self.compile(query, literal_binds=True),
332+
"SELECT t1.x AS t1_x FROM t1 ORDER BY t1.date "
333+
"WITH FILL "
334+
"FROM toDate('2025-01-01') "
335+
"TO toDate('2025-02-01') "
336+
"STEP 1"
337+
)
338+
339+
def test_with_fill_from_to_step_with_datetime(self):
340+
columns = [
341+
Column('y', types.Int32),
342+
Column('dt', types.DateTime)
343+
]
344+
table = self._make_table(*columns)
345+
346+
date_from = date(year=2025, month=1, day=1)
347+
date_to = date(year=2025, month=2, day=1)
348+
query = self.session.query(table.c.x) \
349+
.order_by(
350+
WithFill(
351+
table.c.dt,
352+
from_=func.toDateTime(literal(date_from, type_=types.Date)),
353+
to=func.toDateTime(literal(date_to, type_=types.Date)),
354+
step=literal(1)
355+
)
356+
)
357+
358+
self.assertEqual(
359+
self.compile(query),
360+
"SELECT t1.x AS t1_x FROM t1 ORDER BY t1.dt "
361+
"WITH FILL "
362+
"FROM toDateTime(%(param_1)s) "
363+
"TO toDateTime(%(param_2)s) "
364+
"STEP %(param_3)s"
365+
)
366+
367+
dt_from = dt(year=2025, month=1, day=1, hour=1, minute=1, second=1)
368+
dt_to = dt(year=2025, month=2, day=1)
369+
query = self.session.query(table.c.x) \
370+
.order_by(
371+
WithFill(
372+
table.c.dt,
373+
from_=func.toDateTime(
374+
literal(dt_from, type_=types.Date)
375+
),
376+
to=func.toDateTime(literal(dt_to, type_=types.Date)),
377+
step=literal(1)
378+
)
379+
)
380+
381+
self.assertEqual(
382+
self.compile(query),
383+
"SELECT t1.x AS t1_x FROM t1 ORDER BY t1.dt "
384+
"WITH FILL "
385+
"FROM toDateTime(%(param_1)s) "
386+
"TO toDateTime(%(param_2)s) "
387+
"STEP %(param_3)s"
388+
)
389+
self.assertEqual(
390+
self.compile(query, literal_binds=True),
391+
"SELECT t1.x AS t1_x FROM t1 ORDER BY t1.dt "
392+
"WITH FILL "
393+
"FROM toDateTime('2025-01-01 01:01:01') "
394+
"TO toDateTime('2025-02-01 00:00:00') "
395+
"STEP 1"
396+
)
397+
398+
def test_with_fill_from_to_step_with_datetime64(self):
399+
columns = [
400+
Column('y', types.Int32),
401+
Column('dt', types.DateTime64)
402+
]
403+
table = self._make_table(*columns)
404+
405+
date_from = date(year=2025, month=1, day=1)
406+
date_to = date(year=2025, month=2, day=1)
407+
query = self.session.query(table.c.x) \
408+
.order_by(
409+
WithFill(
410+
table.c.dt,
411+
from_=func.toDateTime64(
412+
literal(
413+
date_from,
414+
type_=types.DateTime64
415+
), table.c.dt.type.precision
416+
),
417+
to=func.toDateTime64(
418+
literal(
419+
date_to,
420+
type_=types.DateTime64
421+
), table.c.dt.type.precision),
422+
step=literal(1)
423+
)
424+
)
425+
426+
self.assertEqual(
427+
self.compile(query),
428+
"SELECT t1.x AS t1_x FROM t1 ORDER BY t1.dt "
429+
"WITH FILL FROM "
430+
"toDateTime64(%(param_1)s, %(toDateTime64_1)s)"
431+
" TO "
432+
"toDateTime64(%(param_2)s, %(toDateTime64_2)s)"
433+
" STEP %(param_3)s"
434+
)
435+
436+
self.assertEqual(
437+
self.compile(query, literal_binds=True),
438+
"SELECT t1.x AS t1_x FROM t1 ORDER BY t1.dt "
439+
"WITH FILL "
440+
"FROM toDateTime64('2025-01-01', 3) "
441+
"TO toDateTime64('2025-02-01', 3) "
442+
"STEP 1"
443+
)
444+
149445

150446
class JoinTestCase(CompilationTestCase):
151447
def test_joins(self):

0 commit comments

Comments
 (0)