Coverage for /Users/davegaeddert/Developer/dropseed/plain/plain-models/plain/models/backends/sqlite3/_functions.py: 29%
346 statements
« prev ^ index » next coverage.py v7.6.9, created at 2024-12-23 11:16 -0600
« prev ^ index » next coverage.py v7.6.9, created at 2024-12-23 11:16 -0600
1"""
2Implementations of SQL functions for SQLite.
3"""
5import functools
6import random
7import statistics
8import zoneinfo
9from datetime import timedelta
10from hashlib import md5, sha1, sha224, sha256, sha384, sha512
11from math import (
12 acos,
13 asin,
14 atan,
15 atan2,
16 ceil,
17 cos,
18 degrees,
19 exp,
20 floor,
21 fmod,
22 log,
23 pi,
24 radians,
25 sin,
26 sqrt,
27 tan,
28)
29from re import search as re_search
31from plain.models.backends.utils import (
32 split_tzname_delta,
33 typecast_time,
34 typecast_timestamp,
35)
36from plain.utils import timezone
37from plain.utils.duration import duration_microseconds
40def register(connection):
41 create_deterministic_function = functools.partial(
42 connection.create_function,
43 deterministic=True,
44 )
45 create_deterministic_function("plain_date_extract", 2, _sqlite_datetime_extract)
46 create_deterministic_function("plain_date_trunc", 4, _sqlite_date_trunc)
47 create_deterministic_function(
48 "plain_datetime_cast_date", 3, _sqlite_datetime_cast_date
49 )
50 create_deterministic_function(
51 "plain_datetime_cast_time", 3, _sqlite_datetime_cast_time
52 )
53 create_deterministic_function("plain_datetime_extract", 4, _sqlite_datetime_extract)
54 create_deterministic_function("plain_datetime_trunc", 4, _sqlite_datetime_trunc)
55 create_deterministic_function("plain_time_extract", 2, _sqlite_time_extract)
56 create_deterministic_function("plain_time_trunc", 4, _sqlite_time_trunc)
57 create_deterministic_function("plain_time_diff", 2, _sqlite_time_diff)
58 create_deterministic_function("plain_timestamp_diff", 2, _sqlite_timestamp_diff)
59 create_deterministic_function("plain_format_dtdelta", 3, _sqlite_format_dtdelta)
60 create_deterministic_function("regexp", 2, _sqlite_regexp)
61 create_deterministic_function("BITXOR", 2, _sqlite_bitxor)
62 create_deterministic_function("COT", 1, _sqlite_cot)
63 create_deterministic_function("LPAD", 3, _sqlite_lpad)
64 create_deterministic_function("MD5", 1, _sqlite_md5)
65 create_deterministic_function("REPEAT", 2, _sqlite_repeat)
66 create_deterministic_function("REVERSE", 1, _sqlite_reverse)
67 create_deterministic_function("RPAD", 3, _sqlite_rpad)
68 create_deterministic_function("SHA1", 1, _sqlite_sha1)
69 create_deterministic_function("SHA224", 1, _sqlite_sha224)
70 create_deterministic_function("SHA256", 1, _sqlite_sha256)
71 create_deterministic_function("SHA384", 1, _sqlite_sha384)
72 create_deterministic_function("SHA512", 1, _sqlite_sha512)
73 create_deterministic_function("SIGN", 1, _sqlite_sign)
74 # Don't use the built-in RANDOM() function because it returns a value
75 # in the range [-1 * 2^63, 2^63 - 1] instead of [0, 1).
76 connection.create_function("RAND", 0, random.random)
77 connection.create_aggregate("STDDEV_POP", 1, StdDevPop)
78 connection.create_aggregate("STDDEV_SAMP", 1, StdDevSamp)
79 connection.create_aggregate("VAR_POP", 1, VarPop)
80 connection.create_aggregate("VAR_SAMP", 1, VarSamp)
81 # Some math functions are enabled by default in SQLite 3.35+.
82 sql = "select sqlite_compileoption_used('ENABLE_MATH_FUNCTIONS')"
83 if not connection.execute(sql).fetchone()[0]:
84 create_deterministic_function("ACOS", 1, _sqlite_acos)
85 create_deterministic_function("ASIN", 1, _sqlite_asin)
86 create_deterministic_function("ATAN", 1, _sqlite_atan)
87 create_deterministic_function("ATAN2", 2, _sqlite_atan2)
88 create_deterministic_function("CEILING", 1, _sqlite_ceiling)
89 create_deterministic_function("COS", 1, _sqlite_cos)
90 create_deterministic_function("DEGREES", 1, _sqlite_degrees)
91 create_deterministic_function("EXP", 1, _sqlite_exp)
92 create_deterministic_function("FLOOR", 1, _sqlite_floor)
93 create_deterministic_function("LN", 1, _sqlite_ln)
94 create_deterministic_function("LOG", 2, _sqlite_log)
95 create_deterministic_function("MOD", 2, _sqlite_mod)
96 create_deterministic_function("PI", 0, _sqlite_pi)
97 create_deterministic_function("POWER", 2, _sqlite_power)
98 create_deterministic_function("RADIANS", 1, _sqlite_radians)
99 create_deterministic_function("SIN", 1, _sqlite_sin)
100 create_deterministic_function("SQRT", 1, _sqlite_sqrt)
101 create_deterministic_function("TAN", 1, _sqlite_tan)
104def _sqlite_datetime_parse(dt, tzname=None, conn_tzname=None):
105 if dt is None:
106 return None
107 try:
108 dt = typecast_timestamp(dt)
109 except (TypeError, ValueError):
110 return None
111 if conn_tzname:
112 dt = dt.replace(tzinfo=zoneinfo.ZoneInfo(conn_tzname))
113 if tzname is not None and tzname != conn_tzname:
114 tzname, sign, offset = split_tzname_delta(tzname)
115 if offset:
116 hours, minutes = offset.split(":")
117 offset_delta = timedelta(hours=int(hours), minutes=int(minutes))
118 dt += offset_delta if sign == "+" else -offset_delta
119 dt = timezone.localtime(dt, zoneinfo.ZoneInfo(tzname))
120 return dt
123def _sqlite_date_trunc(lookup_type, dt, tzname, conn_tzname):
124 dt = _sqlite_datetime_parse(dt, tzname, conn_tzname)
125 if dt is None:
126 return None
127 if lookup_type == "year":
128 return f"{dt.year:04d}-01-01"
129 elif lookup_type == "quarter":
130 month_in_quarter = dt.month - (dt.month - 1) % 3
131 return f"{dt.year:04d}-{month_in_quarter:02d}-01"
132 elif lookup_type == "month":
133 return f"{dt.year:04d}-{dt.month:02d}-01"
134 elif lookup_type == "week":
135 dt -= timedelta(days=dt.weekday())
136 return f"{dt.year:04d}-{dt.month:02d}-{dt.day:02d}"
137 elif lookup_type == "day":
138 return f"{dt.year:04d}-{dt.month:02d}-{dt.day:02d}"
139 raise ValueError(f"Unsupported lookup type: {lookup_type!r}")
142def _sqlite_time_trunc(lookup_type, dt, tzname, conn_tzname):
143 if dt is None:
144 return None
145 dt_parsed = _sqlite_datetime_parse(dt, tzname, conn_tzname)
146 if dt_parsed is None:
147 try:
148 dt = typecast_time(dt)
149 except (ValueError, TypeError):
150 return None
151 else:
152 dt = dt_parsed
153 if lookup_type == "hour":
154 return f"{dt.hour:02d}:00:00"
155 elif lookup_type == "minute":
156 return f"{dt.hour:02d}:{dt.minute:02d}:00"
157 elif lookup_type == "second":
158 return f"{dt.hour:02d}:{dt.minute:02d}:{dt.second:02d}"
159 raise ValueError(f"Unsupported lookup type: {lookup_type!r}")
162def _sqlite_datetime_cast_date(dt, tzname, conn_tzname):
163 dt = _sqlite_datetime_parse(dt, tzname, conn_tzname)
164 if dt is None:
165 return None
166 return dt.date().isoformat()
169def _sqlite_datetime_cast_time(dt, tzname, conn_tzname):
170 dt = _sqlite_datetime_parse(dt, tzname, conn_tzname)
171 if dt is None:
172 return None
173 return dt.time().isoformat()
176def _sqlite_datetime_extract(lookup_type, dt, tzname=None, conn_tzname=None):
177 dt = _sqlite_datetime_parse(dt, tzname, conn_tzname)
178 if dt is None:
179 return None
180 if lookup_type == "week_day":
181 return (dt.isoweekday() % 7) + 1
182 elif lookup_type == "iso_week_day":
183 return dt.isoweekday()
184 elif lookup_type == "week":
185 return dt.isocalendar().week
186 elif lookup_type == "quarter":
187 return ceil(dt.month / 3)
188 elif lookup_type == "iso_year":
189 return dt.isocalendar().year
190 else:
191 return getattr(dt, lookup_type)
194def _sqlite_datetime_trunc(lookup_type, dt, tzname, conn_tzname):
195 dt = _sqlite_datetime_parse(dt, tzname, conn_tzname)
196 if dt is None:
197 return None
198 if lookup_type == "year":
199 return f"{dt.year:04d}-01-01 00:00:00"
200 elif lookup_type == "quarter":
201 month_in_quarter = dt.month - (dt.month - 1) % 3
202 return f"{dt.year:04d}-{month_in_quarter:02d}-01 00:00:00"
203 elif lookup_type == "month":
204 return f"{dt.year:04d}-{dt.month:02d}-01 00:00:00"
205 elif lookup_type == "week":
206 dt -= timedelta(days=dt.weekday())
207 return f"{dt.year:04d}-{dt.month:02d}-{dt.day:02d} 00:00:00"
208 elif lookup_type == "day":
209 return f"{dt.year:04d}-{dt.month:02d}-{dt.day:02d} 00:00:00"
210 elif lookup_type == "hour":
211 return f"{dt.year:04d}-{dt.month:02d}-{dt.day:02d} {dt.hour:02d}:00:00"
212 elif lookup_type == "minute":
213 return (
214 f"{dt.year:04d}-{dt.month:02d}-{dt.day:02d} "
215 f"{dt.hour:02d}:{dt.minute:02d}:00"
216 )
217 elif lookup_type == "second":
218 return (
219 f"{dt.year:04d}-{dt.month:02d}-{dt.day:02d} "
220 f"{dt.hour:02d}:{dt.minute:02d}:{dt.second:02d}"
221 )
222 raise ValueError(f"Unsupported lookup type: {lookup_type!r}")
225def _sqlite_time_extract(lookup_type, dt):
226 if dt is None:
227 return None
228 try:
229 dt = typecast_time(dt)
230 except (ValueError, TypeError):
231 return None
232 return getattr(dt, lookup_type)
235def _sqlite_prepare_dtdelta_param(conn, param):
236 if conn in ["+", "-"]:
237 if isinstance(param, int):
238 return timedelta(0, 0, param)
239 else:
240 return typecast_timestamp(param)
241 return param
244def _sqlite_format_dtdelta(connector, lhs, rhs):
245 """
246 LHS and RHS can be either:
247 - An integer number of microseconds
248 - A string representing a datetime
249 - A scalar value, e.g. float
250 """
251 if connector is None or lhs is None or rhs is None:
252 return None
253 connector = connector.strip()
254 try:
255 real_lhs = _sqlite_prepare_dtdelta_param(connector, lhs)
256 real_rhs = _sqlite_prepare_dtdelta_param(connector, rhs)
257 except (ValueError, TypeError):
258 return None
259 if connector == "+":
260 # typecast_timestamp() returns a date or a datetime without timezone.
261 # It will be formatted as "%Y-%m-%d" or "%Y-%m-%d %H:%M:%S[.%f]"
262 out = str(real_lhs + real_rhs)
263 elif connector == "-":
264 out = str(real_lhs - real_rhs)
265 elif connector == "*":
266 out = real_lhs * real_rhs
267 else:
268 out = real_lhs / real_rhs
269 return out
272def _sqlite_time_diff(lhs, rhs):
273 if lhs is None or rhs is None:
274 return None
275 left = typecast_time(lhs)
276 right = typecast_time(rhs)
277 return (
278 (left.hour * 60 * 60 * 1000000)
279 + (left.minute * 60 * 1000000)
280 + (left.second * 1000000)
281 + (left.microsecond)
282 - (right.hour * 60 * 60 * 1000000)
283 - (right.minute * 60 * 1000000)
284 - (right.second * 1000000)
285 - (right.microsecond)
286 )
289def _sqlite_timestamp_diff(lhs, rhs):
290 if lhs is None or rhs is None:
291 return None
292 left = typecast_timestamp(lhs)
293 right = typecast_timestamp(rhs)
294 return duration_microseconds(left - right)
297def _sqlite_regexp(pattern, string):
298 if pattern is None or string is None:
299 return None
300 if not isinstance(string, str):
301 string = str(string)
302 return bool(re_search(pattern, string))
305def _sqlite_acos(x):
306 if x is None:
307 return None
308 return acos(x)
311def _sqlite_asin(x):
312 if x is None:
313 return None
314 return asin(x)
317def _sqlite_atan(x):
318 if x is None:
319 return None
320 return atan(x)
323def _sqlite_atan2(y, x):
324 if y is None or x is None:
325 return None
326 return atan2(y, x)
329def _sqlite_bitxor(x, y):
330 if x is None or y is None:
331 return None
332 return x ^ y
335def _sqlite_ceiling(x):
336 if x is None:
337 return None
338 return ceil(x)
341def _sqlite_cos(x):
342 if x is None:
343 return None
344 return cos(x)
347def _sqlite_cot(x):
348 if x is None:
349 return None
350 return 1 / tan(x)
353def _sqlite_degrees(x):
354 if x is None:
355 return None
356 return degrees(x)
359def _sqlite_exp(x):
360 if x is None:
361 return None
362 return exp(x)
365def _sqlite_floor(x):
366 if x is None:
367 return None
368 return floor(x)
371def _sqlite_ln(x):
372 if x is None:
373 return None
374 return log(x)
377def _sqlite_log(base, x):
378 if base is None or x is None:
379 return None
380 # Arguments reversed to match SQL standard.
381 return log(x, base)
384def _sqlite_lpad(text, length, fill_text):
385 if text is None or length is None or fill_text is None:
386 return None
387 delta = length - len(text)
388 if delta <= 0:
389 return text[:length]
390 return (fill_text * length)[:delta] + text
393def _sqlite_md5(text):
394 if text is None:
395 return None
396 return md5(text.encode()).hexdigest()
399def _sqlite_mod(x, y):
400 if x is None or y is None:
401 return None
402 return fmod(x, y)
405def _sqlite_pi():
406 return pi
409def _sqlite_power(x, y):
410 if x is None or y is None:
411 return None
412 return x**y
415def _sqlite_radians(x):
416 if x is None:
417 return None
418 return radians(x)
421def _sqlite_repeat(text, count):
422 if text is None or count is None:
423 return None
424 return text * count
427def _sqlite_reverse(text):
428 if text is None:
429 return None
430 return text[::-1]
433def _sqlite_rpad(text, length, fill_text):
434 if text is None or length is None or fill_text is None:
435 return None
436 return (text + fill_text * length)[:length]
439def _sqlite_sha1(text):
440 if text is None:
441 return None
442 return sha1(text.encode()).hexdigest()
445def _sqlite_sha224(text):
446 if text is None:
447 return None
448 return sha224(text.encode()).hexdigest()
451def _sqlite_sha256(text):
452 if text is None:
453 return None
454 return sha256(text.encode()).hexdigest()
457def _sqlite_sha384(text):
458 if text is None:
459 return None
460 return sha384(text.encode()).hexdigest()
463def _sqlite_sha512(text):
464 if text is None:
465 return None
466 return sha512(text.encode()).hexdigest()
469def _sqlite_sign(x):
470 if x is None:
471 return None
472 return (x > 0) - (x < 0)
475def _sqlite_sin(x):
476 if x is None:
477 return None
478 return sin(x)
481def _sqlite_sqrt(x):
482 if x is None:
483 return None
484 return sqrt(x)
487def _sqlite_tan(x):
488 if x is None:
489 return None
490 return tan(x)
493class ListAggregate(list):
494 step = list.append
497class StdDevPop(ListAggregate):
498 finalize = statistics.pstdev
501class StdDevSamp(ListAggregate):
502 finalize = statistics.stdev
505class VarPop(ListAggregate):
506 finalize = statistics.pvariance
509class VarSamp(ListAggregate):
510 finalize = statistics.variance