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

1""" 

2Implementations of SQL functions for SQLite. 

3""" 

4 

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 

30 

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 

38 

39 

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) 

102 

103 

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 

121 

122 

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}") 

140 

141 

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}") 

160 

161 

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() 

167 

168 

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() 

174 

175 

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) 

192 

193 

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}") 

223 

224 

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) 

233 

234 

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 

242 

243 

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 

270 

271 

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 ) 

287 

288 

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) 

295 

296 

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)) 

303 

304 

305def _sqlite_acos(x): 

306 if x is None: 

307 return None 

308 return acos(x) 

309 

310 

311def _sqlite_asin(x): 

312 if x is None: 

313 return None 

314 return asin(x) 

315 

316 

317def _sqlite_atan(x): 

318 if x is None: 

319 return None 

320 return atan(x) 

321 

322 

323def _sqlite_atan2(y, x): 

324 if y is None or x is None: 

325 return None 

326 return atan2(y, x) 

327 

328 

329def _sqlite_bitxor(x, y): 

330 if x is None or y is None: 

331 return None 

332 return x ^ y 

333 

334 

335def _sqlite_ceiling(x): 

336 if x is None: 

337 return None 

338 return ceil(x) 

339 

340 

341def _sqlite_cos(x): 

342 if x is None: 

343 return None 

344 return cos(x) 

345 

346 

347def _sqlite_cot(x): 

348 if x is None: 

349 return None 

350 return 1 / tan(x) 

351 

352 

353def _sqlite_degrees(x): 

354 if x is None: 

355 return None 

356 return degrees(x) 

357 

358 

359def _sqlite_exp(x): 

360 if x is None: 

361 return None 

362 return exp(x) 

363 

364 

365def _sqlite_floor(x): 

366 if x is None: 

367 return None 

368 return floor(x) 

369 

370 

371def _sqlite_ln(x): 

372 if x is None: 

373 return None 

374 return log(x) 

375 

376 

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) 

382 

383 

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 

391 

392 

393def _sqlite_md5(text): 

394 if text is None: 

395 return None 

396 return md5(text.encode()).hexdigest() 

397 

398 

399def _sqlite_mod(x, y): 

400 if x is None or y is None: 

401 return None 

402 return fmod(x, y) 

403 

404 

405def _sqlite_pi(): 

406 return pi 

407 

408 

409def _sqlite_power(x, y): 

410 if x is None or y is None: 

411 return None 

412 return x**y 

413 

414 

415def _sqlite_radians(x): 

416 if x is None: 

417 return None 

418 return radians(x) 

419 

420 

421def _sqlite_repeat(text, count): 

422 if text is None or count is None: 

423 return None 

424 return text * count 

425 

426 

427def _sqlite_reverse(text): 

428 if text is None: 

429 return None 

430 return text[::-1] 

431 

432 

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] 

437 

438 

439def _sqlite_sha1(text): 

440 if text is None: 

441 return None 

442 return sha1(text.encode()).hexdigest() 

443 

444 

445def _sqlite_sha224(text): 

446 if text is None: 

447 return None 

448 return sha224(text.encode()).hexdigest() 

449 

450 

451def _sqlite_sha256(text): 

452 if text is None: 

453 return None 

454 return sha256(text.encode()).hexdigest() 

455 

456 

457def _sqlite_sha384(text): 

458 if text is None: 

459 return None 

460 return sha384(text.encode()).hexdigest() 

461 

462 

463def _sqlite_sha512(text): 

464 if text is None: 

465 return None 

466 return sha512(text.encode()).hexdigest() 

467 

468 

469def _sqlite_sign(x): 

470 if x is None: 

471 return None 

472 return (x > 0) - (x < 0) 

473 

474 

475def _sqlite_sin(x): 

476 if x is None: 

477 return None 

478 return sin(x) 

479 

480 

481def _sqlite_sqrt(x): 

482 if x is None: 

483 return None 

484 return sqrt(x) 

485 

486 

487def _sqlite_tan(x): 

488 if x is None: 

489 return None 

490 return tan(x) 

491 

492 

493class ListAggregate(list): 

494 step = list.append 

495 

496 

497class StdDevPop(ListAggregate): 

498 finalize = statistics.pstdev 

499 

500 

501class StdDevSamp(ListAggregate): 

502 finalize = statistics.stdev 

503 

504 

505class VarPop(ListAggregate): 

506 finalize = statistics.pvariance 

507 

508 

509class VarSamp(ListAggregate): 

510 finalize = statistics.variance