sqlakit 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
sqlakit/_query.py ADDED
@@ -0,0 +1,1156 @@
1
+ from __future__ import annotations
2
+
3
+ import base64
4
+ import binascii
5
+ import hashlib
6
+ import json
7
+ from collections.abc import Mapping
8
+ from dataclasses import dataclass
9
+ from typing import (
10
+ TYPE_CHECKING,
11
+ Any,
12
+ Generic,
13
+ NamedTuple,
14
+ Protocol,
15
+ Self,
16
+ TypeVar,
17
+ cast,
18
+ )
19
+
20
+ import sqlalchemy as sa
21
+ from sqlalchemy import exc as sa_exc
22
+ from sqlalchemy.orm import (
23
+ InstrumentedAttribute,
24
+ contains_eager,
25
+ joinedload,
26
+ selectinload,
27
+ subqueryload,
28
+ )
29
+ from sqlalchemy.orm.exc import MultipleResultsFound, NoResultFound
30
+ from sqlalchemy.sql import operators
31
+
32
+ from ._model import resolve_alias, soft_delete_column
33
+ from .exceptions import (
34
+ BulkQueryError,
35
+ InstanceNotFoundError,
36
+ InvalidCursorError,
37
+ InvalidOrderFieldError,
38
+ KeyLookupError,
39
+ MultipleInstancesFoundError,
40
+ NullCursorValueError,
41
+ PageItemsMismatchError,
42
+ RawStatementError,
43
+ UncomparableOrderingError,
44
+ UnknownOrderFieldError,
45
+ UnorderedPageError,
46
+ )
47
+
48
+ if TYPE_CHECKING:
49
+ from collections.abc import Callable, Iterable, Iterator, Sequence
50
+
51
+ from sqlalchemy.sql._typing import _ColumnExpressionArgument
52
+ from sqlalchemy.sql.roles import ReturnsRowsRole
53
+ from sqlalchemy.sql.selectable import ForUpdateParameter
54
+
55
+ __all__ = [
56
+ "BaseQuery",
57
+ "CursorPage",
58
+ "OrderBy",
59
+ "Page",
60
+ "one_row",
61
+ "one_row_or_none",
62
+ "orderable",
63
+ "ordered",
64
+ ]
65
+
66
+ HIDDEN = "hidden"
67
+ """The rows a soft delete marked are left out, as a read does by default."""
68
+
69
+ INCLUDED = "included"
70
+ """They are read alongside the rest."""
71
+
72
+ ONLY = "only"
73
+ """They are the only rows read."""
74
+
75
+ ModelT = TypeVar("ModelT")
76
+ OtherT = TypeVar("OtherT")
77
+ RowT = TypeVar("RowT")
78
+ RowT_co = TypeVar("RowT_co", covariant=True)
79
+
80
+
81
+ @dataclass(frozen=True, slots=True)
82
+ class Page(Generic[ModelT]):
83
+ """One page of rows, and how many there are in total."""
84
+
85
+ items: Sequence[ModelT]
86
+ total: int | None
87
+ """How many rows match, or None for a page read without counting them."""
88
+
89
+ limit: int
90
+ offset: int
91
+
92
+ has_next: bool = False
93
+ """Whether a page follows this one."""
94
+
95
+ def __post_init__(self) -> None:
96
+ if self.total is not None:
97
+ object.__setattr__(
98
+ self, "has_next", self.offset + len(self.items) < self.total
99
+ )
100
+
101
+ def map(self, transform: Callable[[ModelT], OtherT]) -> Page[OtherT]:
102
+ """Return the page with every row put through ``transform``.
103
+
104
+ ```python
105
+ page = User.query.page(limit=20).map(UserResponse.from_user)
106
+ ```
107
+ """
108
+ return self.with_items([transform(item) for item in self.items])
109
+
110
+ def map_all(
111
+ self,
112
+ transform: Callable[[Sequence[ModelT]], Sequence[OtherT]],
113
+ ) -> Page[OtherT]:
114
+ """Return the page with the rows put through ``transform`` together.
115
+
116
+ For work that reads better in one go than row by row: one query for
117
+ what the rows refer to, one call to a serializer that takes a list.
118
+ """
119
+ return self.with_items(transform(self.items))
120
+
121
+ def with_items(self, items: Sequence[OtherT]) -> Page[OtherT]:
122
+ """Return the page carrying these rows instead, counts unchanged.
123
+
124
+ What an asynchronous transform needs: `page.with_items(await serialize(...))`.
125
+
126
+ Raises:
127
+ PageItemsMismatchError: if there is not one item per row.
128
+
129
+ """
130
+ if len(items) != len(self.items):
131
+ raise PageItemsMismatchError(len(self.items), len(items))
132
+ return Page(
133
+ items=items,
134
+ total=self.total,
135
+ limit=self.limit,
136
+ offset=self.offset,
137
+ has_next=self.has_next,
138
+ )
139
+
140
+
141
+ @dataclass(frozen=True, slots=True)
142
+ class CursorPage(Generic[ModelT]):
143
+ """One page of rows, and the cursors that read the ones either side."""
144
+
145
+ items: Sequence[ModelT]
146
+ next_cursor: str | None = None
147
+ """Hand it back as ``cursor`` to read on."""
148
+
149
+ previous_cursor: str | None = None
150
+ """Hand it back as ``cursor`` to read the page in front of this one."""
151
+
152
+ @property
153
+ def has_next(self) -> bool:
154
+ """Whether a page follows this one."""
155
+ return self.next_cursor is not None
156
+
157
+ @property
158
+ def has_previous(self) -> bool:
159
+ """Whether a page comes before this one."""
160
+ return self.previous_cursor is not None
161
+
162
+ def map(self, transform: Callable[[ModelT], OtherT]) -> CursorPage[OtherT]:
163
+ """Return the page with every row put through ``transform``."""
164
+ return self.with_items([transform(item) for item in self.items])
165
+
166
+ def map_all(
167
+ self,
168
+ transform: Callable[[Sequence[ModelT]], Sequence[OtherT]],
169
+ ) -> CursorPage[OtherT]:
170
+ """Return the page with the rows put through ``transform`` together."""
171
+ return self.with_items(transform(self.items))
172
+
173
+ def with_items(self, items: Sequence[OtherT]) -> CursorPage[OtherT]:
174
+ """Return the page carrying these rows instead, cursors unchanged.
175
+
176
+ Raises:
177
+ PageItemsMismatchError: if there is not one item per row.
178
+
179
+ """
180
+ if len(items) != len(self.items):
181
+ raise PageItemsMismatchError(len(self.items), len(items))
182
+ return CursorPage(
183
+ items=items,
184
+ next_cursor=self.next_cursor,
185
+ previous_cursor=self.previous_cursor,
186
+ )
187
+
188
+
189
+ class OrderBy(NamedTuple):
190
+ """A field to order by that lives in another table.
191
+
192
+ Name the table, or the relationship that reaches it, and a query ordered by
193
+ that field joins it once, however many fields name it:
194
+
195
+ ```python
196
+ {"team": OrderBy(Team.name, join=cls.team)}
197
+ ```
198
+
199
+ Join what holds one row. A collection multiplies the rows, and a page of
200
+ multiplied rows counts wrong: join a subquery that aggregates them instead,
201
+ with the ``on`` it needs.
202
+ """
203
+
204
+ expression: Any
205
+ join: Any = None
206
+ on: Any = None
207
+
208
+
209
+ class SupportsClause(Protocol):
210
+ """Anything SQLAlchemy will take a statement from, ours included."""
211
+
212
+ def __clause_element__(self) -> Any: ... # noqa: ANN401
213
+
214
+
215
+ class _Ordering(NamedTuple):
216
+ """One column of a cursor's ordering: how to sort it, how to read it back."""
217
+
218
+ column: sa.ColumnElement[Any]
219
+ descending: bool
220
+ attribute: str
221
+
222
+
223
+ class BaseQuery(Generic[ModelT]):
224
+ """The statements a query builds, without running any of them.
225
+
226
+ Every builder method returns a new query, so one can be kept around and
227
+ branched from, the way a SQLAlchemy `Select` can.
228
+ """
229
+
230
+ def __init__(
231
+ self,
232
+ model: type[ModelT],
233
+ db: Any, # noqa: ANN401
234
+ select: sa.Select[Any] | None = None,
235
+ ) -> None:
236
+ self.model = model
237
+ self.db = db
238
+ self._select = sa.select(model) if select is None else select
239
+ self._statement: Any = None
240
+ self.filtered = True
241
+ self.deleted = HIDDEN
242
+
243
+ def using(self, target: str | Any) -> Self: # noqa: ANN401
244
+ """Run this query on another database, named or handed over.
245
+
246
+ ```python
247
+ User.query.using("replica").order_by("name").page(limit=20)
248
+ ```
249
+
250
+ It pins the query: reads and writes both go where you said, whatever the
251
+ routers would have answered.
252
+ """
253
+ db = resolve_alias(self.model, target) if isinstance(target, str) else target
254
+ query = self._copy()
255
+ query.db = db
256
+ return query
257
+
258
+ @property
259
+ def select(self) -> sa.Select[Any]:
260
+ """The SQLAlchemy `Select` this query has built so far.
261
+
262
+ Read it for what the builders do not cover, and hand the result to
263
+ `with_select`. Building a new query is how a query changes: this one
264
+ cannot be assigned to.
265
+ """
266
+ return self._select
267
+
268
+ def with_select(self, select: sa.Select[Any]) -> Self:
269
+ """Return a query like this one, over the given select.
270
+
271
+ This is what every builder is made of. Use it in a method of your own
272
+ when the statement needs SQLAlchemy the builders do not cover.
273
+
274
+ Raises:
275
+ RawStatementError: if the query is a statement of its own, which
276
+ has no select to replace.
277
+
278
+ """
279
+ self._reject_statement(
280
+ "with_select",
281
+ "there is no select to replace, so build the statement you want and "
282
+ "hand it to `from_statement`",
283
+ )
284
+ query = self._copy()
285
+ query._select = select # noqa: SLF001 - a copy of this class
286
+ return query
287
+
288
+ def _copy(self) -> Self:
289
+ """Return this query again, statement and all, for a caller to adjust."""
290
+ query = object.__new__(type(self))
291
+ query.__dict__.update(self.__dict__)
292
+ return query
293
+
294
+ # --- Building ---
295
+
296
+ @property
297
+ def is_ordered(self) -> bool:
298
+ """Whether this query carries an ordering.
299
+
300
+ What a method of your own asks before it adds a default one, since
301
+ `page` and `cursor_page` refuse a query with no ordering.
302
+ """
303
+ return bool(self._select._order_by_clauses) # noqa: SLF001
304
+
305
+ def unfiltered(self) -> Self:
306
+ """Drop the model's own `__query_filter__` for this query.
307
+
308
+ The rows a [soft delete](models.md#soft-deletes) marked stay hidden.
309
+ They are a separate switch, `with_deleted()`.
310
+ """
311
+ query = self._copy()
312
+ query.filtered = False
313
+ return query
314
+
315
+ def with_deleted(self) -> Self:
316
+ """Include the rows a [soft delete](models.md#soft-deletes) marked.
317
+
318
+ Every other filter the model carries stays on.
319
+ """
320
+ query = self._copy()
321
+ query.deleted = INCLUDED
322
+ return query
323
+
324
+ def only_deleted(self) -> Self:
325
+ """Read only the rows a soft delete marked.
326
+
327
+ ```python
328
+ Note.query.only_deleted().delete(force=True) # empty the bin
329
+ ```
330
+ """
331
+ query = self._copy()
332
+ query.deleted = ONLY
333
+ return query
334
+
335
+ def from_statement(self, statement: ReturnsRowsRole | SupportsClause) -> Self:
336
+ """Take the rows of a statement of your own, mapped onto the model.
337
+
338
+ ```python
339
+ User.query.from_statement(sa.text("SELECT * FROM users WHERE ...")).all()
340
+ ```
341
+
342
+ Nothing can be added afterwards, and ``__query_filter__`` is not applied:
343
+ what the statement selects is what comes back.
344
+ """
345
+ query = self._copy()
346
+ query._statement = self._select.from_statement( # noqa: SLF001
347
+ cast("ReturnsRowsRole", statement)
348
+ )
349
+ return query
350
+
351
+ def from_sql(self, template: str, **context: Any) -> Self: # noqa: ANN401
352
+ """Take the rows of a SQL template, mapped onto the model.
353
+
354
+ ```python
355
+ User.query.from_sql("users/active.sql", team="red").all()
356
+ ```
357
+
358
+ Read from the database this query runs on, and rendered for its dialect. As
359
+ with `from_statement`, nothing can be added afterwards and
360
+ ``__query_filter__`` is not applied.
361
+ """
362
+ return self.from_statement(self.db.sql.from_file(template, **context).statement)
363
+
364
+ def where(self, *criteria: _ColumnExpressionArgument[bool]) -> Self:
365
+ """Narrow the rows, as `Select.where` does."""
366
+ self._reject_statement("where")
367
+ return self.with_select(self._select.where(*criteria))
368
+
369
+ def filter_by(self, **values: Any) -> Self: # noqa: ANN401
370
+ """Narrow the rows by equality, as `Select.filter_by` does."""
371
+ self._reject_statement("filter_by")
372
+ return self.with_select(self._select.filter_by(**values))
373
+
374
+ def join(
375
+ self,
376
+ target: Any, # noqa: ANN401
377
+ onclause: Any = None, # noqa: ANN401
378
+ *,
379
+ isouter: bool = False,
380
+ full: bool = False,
381
+ ) -> Self:
382
+ """Join another table in."""
383
+ self._reject_statement("join")
384
+ return self.with_select(
385
+ self._select.join(target, onclause, isouter=isouter, full=full)
386
+ )
387
+
388
+ def outerjoin(self, target: Any, onclause: Any = None) -> Self: # noqa: ANN401
389
+ """Join another table in, keeping the rows without a match."""
390
+ return self.join(target, onclause, isouter=True)
391
+
392
+ def select_from(self, *froms: Any) -> Self: # noqa: ANN401
393
+ """Name what the query selects from, when the joins do not say it."""
394
+ self._reject_statement("select_from")
395
+ return self.with_select(self._select.select_from(*froms))
396
+
397
+ def group_by(self, *criteria: Any) -> Self: # noqa: ANN401
398
+ """Group the rows."""
399
+ self._reject_statement("group_by")
400
+ return self.with_select(self._select.group_by(*criteria))
401
+
402
+ def having(self, *criteria: _ColumnExpressionArgument[bool]) -> Self:
403
+ """Narrow the groups."""
404
+ self._reject_statement("having")
405
+ return self.with_select(self._select.having(*criteria))
406
+
407
+ def distinct(self) -> Self:
408
+ """Drop the duplicate rows."""
409
+ self._reject_statement("distinct")
410
+ return self.with_select(self._select.distinct())
411
+
412
+ def execution_options(self, **options: Any) -> Self: # noqa: ANN401
413
+ """Set execution options, `yield_per` and `populate_existing` among them."""
414
+ return self.with_select(self._select.execution_options(**options))
415
+
416
+ def order_by(
417
+ self,
418
+ *criteria: Any, # noqa: ANN401
419
+ ci_fields: Sequence[str] = (),
420
+ ) -> Self:
421
+ """Order the rows, by columns or by the sort strings a request carries.
422
+
423
+ A string is `name`, `name.desc`, or `name.desc.nulls_last`. Names are looked
424
+ up in what the model offers, so a field nobody meant to sort by is refused
425
+ rather than turned into SQL:
426
+
427
+ ```python
428
+ User.query.order_by(request.sort).page(limit=20)
429
+ User.query.order_by("team", "created_at.desc").cursor_page(limit=20)
430
+ User.query.order_by(User.name.desc()).all()
431
+ ```
432
+
433
+ A `None` is skipped and a list is taken apart, so a request that names no
434
+ sort, or several, passes straight through.
435
+
436
+ ``ci_fields`` names the fields to compare without regard to case. It sorts by
437
+ `lower(...)`, which a cursor cannot page: use it with `page`, or fold the case
438
+ in ``__orderable__`` and index it.
439
+
440
+ A model sorts by its own mapped columns. `orderable` says how to offer
441
+ others, including fields that are not columns at all.
442
+
443
+ Raises:
444
+ UnknownOrderFieldError: if a string names a field the model does not
445
+ offer.
446
+
447
+ """
448
+ self._reject_statement("order_by")
449
+ return self.with_select(
450
+ ordered(self._select, self._orderable(), criteria, ci_fields)
451
+ )
452
+
453
+ def _directed(self, column: Any, *, descending: bool) -> Any: # noqa: ANN401
454
+ """Return the criterion that orders by a column one way or the other.
455
+
456
+ Takes what `order_by` takes: a column, or the name of a field the model
457
+ offers.
458
+ """
459
+ if isinstance(column, str):
460
+ name, _, _ = _parse_sort_field(column)
461
+ return f"{name}.desc" if descending else f"{name}.asc"
462
+ return sa.desc(column) if descending else sa.asc(column)
463
+
464
+ def _orderable(self) -> Mapping[str, Any]:
465
+ """Return what this model orders by, which a subclass may narrow."""
466
+ return orderable(self.model)
467
+
468
+ def limit(self, limit: int) -> Self:
469
+ """Take at most this many rows."""
470
+ self._reject_statement("limit")
471
+ return self.with_select(self._select.limit(limit))
472
+
473
+ def offset(self, offset: int) -> Self:
474
+ """Skip this many rows."""
475
+ self._reject_statement("offset")
476
+ return self.with_select(self._select.offset(offset))
477
+
478
+ def options(self, *options: Any) -> Self: # noqa: ANN401
479
+ """Apply loader options, as `Select.options` does."""
480
+ return self.with_select(self._select.options(*options))
481
+
482
+ def joinedload(self, *keys: Any) -> Self: # noqa: ANN401
483
+ """Load a relationship, and the ones below it, with a JOIN."""
484
+ return self.options(_chain(joinedload, keys))
485
+
486
+ def selectinload(self, *keys: Any) -> Self: # noqa: ANN401
487
+ """Load a relationship, and the ones below it, with a second SELECT."""
488
+ return self.options(_chain(selectinload, keys))
489
+
490
+ def subqueryload(self, *keys: Any) -> Self: # noqa: ANN401
491
+ """Load a relationship, and the ones below it, with a subquery."""
492
+ return self.options(_chain(subqueryload, keys))
493
+
494
+ def contains_eager(self, *keys: Any) -> Self: # noqa: ANN401
495
+ """Read a relationship from a join this query already makes."""
496
+ return self.options(_chain(contains_eager, keys))
497
+
498
+ def with_for_update(
499
+ self,
500
+ *,
501
+ nowait: bool = False,
502
+ read: bool = False,
503
+ of: Any = None, # noqa: ANN401
504
+ skip_locked: bool = False,
505
+ key_share: bool = False,
506
+ ) -> Self:
507
+ """Lock the matched rows.
508
+
509
+ Args:
510
+ nowait: Fail rather than wait for a row someone else holds.
511
+ read: Take a shared lock instead of an exclusive one.
512
+ of: Lock the rows of these tables only.
513
+ skip_locked: Pass over the rows someone else holds.
514
+ key_share: Take the weakest lock that still blocks key changes.
515
+
516
+ """
517
+ parameters: ForUpdateParameter = {
518
+ "nowait": nowait,
519
+ "read": read,
520
+ "of": of,
521
+ "skip_locked": skip_locked,
522
+ "key_share": key_share,
523
+ }
524
+ self._reject_statement("with_for_update")
525
+ return self.with_select(self._select.with_for_update(**parameters))
526
+
527
+ # --- Statements ---
528
+
529
+ def _lookup_options(self) -> Sequence[Any]:
530
+ """Return the loader options a lookup by key carries over."""
531
+ self._reject_statement("get")
532
+ select = self._select
533
+ carried = {
534
+ "where": select.whereclause is not None,
535
+ "limit": select._limit_clause is not None, # noqa: SLF001
536
+ "offset": select._offset_clause is not None, # noqa: SLF001
537
+ "order_by": bool(select._order_by_clauses), # noqa: SLF001
538
+ "join": bool(select._setup_joins), # noqa: SLF001
539
+ }
540
+ present = tuple(name for name, carries in carried.items() if carries)
541
+ if present:
542
+ raise KeyLookupError(present)
543
+ return select._with_options # noqa: SLF001
544
+
545
+ def _lock(self) -> Any: # noqa: ANN401
546
+ """Return the lock this query asks for, if it asks for one."""
547
+ return self._select._for_update_arg # noqa: SLF001
548
+
549
+ def _lookup_statement(self, ident: Any) -> sa.Select[Any] | None: # noqa: ANN401
550
+ """Return the select a filtered model needs, or None to look up by key.
551
+
552
+ A model that hides rows, through ``__query_filter__`` or a soft delete,
553
+ hides them from a lookup too, and the session cannot answer that from the
554
+ identity map: it knows the row, not whether the filter still admits it.
555
+ """
556
+ unfiltered = (
557
+ getattr(self.model, "__query_filter__", None) is None or not self.filtered
558
+ )
559
+ column = soft_delete_column(self.model)
560
+ if unfiltered and (column is None or self.deleted == INCLUDED):
561
+ return None
562
+ mapper = sa.inspect(self.model, raiseerr=True)
563
+ if isinstance(ident, Mapping):
564
+ criteria = [mapper.columns[name] == value for name, value in ident.items()]
565
+ else:
566
+ values = ident if isinstance(ident, tuple) else (ident,)
567
+ criteria = [
568
+ column == value
569
+ for column, value in zip(mapper.primary_key, values, strict=True)
570
+ ]
571
+ return self._filtered().where(*criteria)
572
+
573
+ def _executable(self) -> Any: # noqa: ANN401
574
+ return self._filtered() if self._statement is None else self._statement
575
+
576
+ def _filtered(self) -> sa.Select[Any]:
577
+ """Return the select, with the two filters a model can put on a read.
578
+
579
+ ``__query_filter__``, which hides rows for good, and the one hiding rows a
580
+ soft delete marked. They lift separately, with `unfiltered()` and
581
+ `with_deleted()`.
582
+ """
583
+ select = self._select
584
+ criterion = getattr(self.model, "__query_filter__", None)
585
+ if criterion is not None and self.filtered:
586
+ select = select.where(criterion())
587
+ column = soft_delete_column(self.model)
588
+ if column is None or self.deleted == INCLUDED:
589
+ return select
590
+ marked = getattr(self.model, column)
591
+ return select.where(
592
+ marked.is_not(None) if self.deleted == ONLY else marked.is_(None)
593
+ )
594
+
595
+ def _reject_statement(self, method: str, advice: str | None = None) -> None:
596
+ if self._statement is not None:
597
+ raise RawStatementError(method, advice)
598
+
599
+ def _count_statement(self) -> sa.Select[tuple[int]]:
600
+ self._reject_statement("count")
601
+ return sa.select(sa.func.count()).select_from(
602
+ self._filtered().order_by(None).subquery()
603
+ )
604
+
605
+ def _exists_statement(self) -> sa.Select[tuple[bool]]:
606
+ self._reject_statement("exists")
607
+ return sa.select(sa.exists(self._filtered().order_by(None).limit(1).subquery()))
608
+
609
+ def _page_statement(self, *, limit: int, offset: int) -> sa.Select[Any]:
610
+ self._reject_statement("page")
611
+ self._require_ordering()
612
+ return self._tiebroken(self._filtered()).limit(limit).offset(offset)
613
+
614
+ def _tiebroken(self, select: sa.Select[Any]) -> sa.Select[Any]:
615
+ """Append the key, so rows with equal sort values keep their places.
616
+
617
+ Without it the database is free to put a tied row on either page, so one row
618
+ comes twice and another never comes.
619
+ """
620
+ clauses = select._order_by_clauses # noqa: SLF001
621
+ ordered = {_column_identity(clause) for clause in clauses}
622
+ descending = _direction(clauses[-1])[1] if clauses else False
623
+ missing = [
624
+ column
625
+ for column in map(_as_column, self._cursor_key())
626
+ if _column_identity(column) not in ordered
627
+ ]
628
+ if not missing:
629
+ return select
630
+ return select.order_by(
631
+ *(sa.desc(column) if descending else sa.asc(column) for column in missing)
632
+ )
633
+
634
+ def _cursor_statement(self, *, limit: int, cursor: str | None) -> sa.Select[Any]:
635
+ """One row more than asked for, so the caller knows whether more follow.
636
+
637
+ Backwards is the same walk with every direction flipped: the rows nearest the
638
+ cursor come first, the limit cuts the far end, and `_cursor_page` turns them
639
+ around again.
640
+ """
641
+ self._reject_statement("cursor_page")
642
+ self._require_ordering()
643
+ ordering = self._keyset_ordering()
644
+ ordering_key = _ordering_key(ordering)
645
+ backwards = cursor is not None and _is_backwards(cursor)
646
+ if backwards:
647
+ ordering = [
648
+ item._replace(descending=not item.descending) for item in ordering
649
+ ]
650
+ select = (
651
+ self._filtered()
652
+ .order_by(None)
653
+ .order_by(
654
+ *(
655
+ item.column.desc() if item.descending else item.column.asc()
656
+ for item in ordering
657
+ )
658
+ )
659
+ )
660
+ if cursor is not None:
661
+ select = select.where(
662
+ _seek(ordering, _decode(cursor, ordering, ordering_key))
663
+ )
664
+ return select.limit(limit + 1)
665
+
666
+ def _require_ordering(self) -> None:
667
+ """Refuse a page the database is free to order as it likes."""
668
+ if not self.is_ordered:
669
+ raise UnorderedPageError
670
+
671
+ def _keyset_ordering(self) -> list[_Ordering]:
672
+ """Return the ordering a cursor walks: what was asked, then the key.
673
+
674
+ A cursor compares rows by the columns they are ordered on, so the ordering
675
+ has to end in something unique, or rows sharing the last value fall on both
676
+ sides of a page boundary. The key is the primary key unless the model names
677
+ another with ``__cursor_key__``; columns already in the ordering stay where
678
+ they are.
679
+ """
680
+ mapper = sa.inspect(self.model, raiseerr=True)
681
+ ordering: list[_Ordering] = []
682
+ seen: set[str] = set()
683
+ for clause in self._select._order_by_clauses: # noqa: SLF001
684
+ column, descending = _direction(clause)
685
+ try:
686
+ attribute = mapper.get_property_by_column(column).key
687
+ except (AttributeError, sa_exc.InvalidRequestError):
688
+ raise UncomparableOrderingError(clause) from None
689
+ if attribute not in seen:
690
+ seen.add(attribute)
691
+ ordering.append(_Ordering(column, descending, attribute))
692
+ for entry in self._cursor_key():
693
+ column = _as_column(entry)
694
+ attribute = mapper.get_property_by_column(column).key
695
+ if attribute not in seen:
696
+ seen.add(attribute)
697
+ descending = ordering[-1].descending if ordering else False
698
+ ordering.append(_Ordering(column, descending, attribute))
699
+ return ordering
700
+
701
+ def _update_statement(self, values: Mapping[str, Any]) -> sa.Update:
702
+ return (
703
+ sa.update(self.model).where(*self._bulk_criteria("update")).values(**values)
704
+ )
705
+
706
+ def _delete_statement(self, *, force: bool = False) -> sa.Delete | sa.Update:
707
+ """Return the statement that deletes, or the one that marks.
708
+
709
+ A model that soft-deletes is marked rather than removed, so that the
710
+ bulk path and `Model.delete()` agree. ``force`` removes the rows.
711
+ """
712
+ criteria = self._bulk_criteria("delete")
713
+ column = soft_delete_column(self.model)
714
+ if column is None or force:
715
+ return sa.delete(self.model).where(*criteria)
716
+ return sa.update(self.model).where(*criteria).values({column: sa.func.now()})
717
+
718
+ def _bulk_criteria(self, method: str) -> list[Any]:
719
+ """Return what a bulk statement carries over: the narrowing, nothing else."""
720
+ self._reject_statement(method)
721
+ select = self._filtered()
722
+ carried = {
723
+ "limit": select._limit_clause is not None, # noqa: SLF001
724
+ "offset": select._offset_clause is not None, # noqa: SLF001
725
+ "order_by": bool(select._order_by_clauses), # noqa: SLF001
726
+ "join": bool(select._setup_joins), # noqa: SLF001
727
+ }
728
+ dropped = tuple(name for name, present in carried.items() if present)
729
+ if dropped:
730
+ raise BulkQueryError(method, dropped)
731
+ return [] if select.whereclause is None else [select.whereclause]
732
+
733
+ def _columns_select(self, columns: Sequence[Any]) -> sa.Select[Any]:
734
+ """Return the query, reading the given columns instead of whole rows."""
735
+ self._reject_statement("only_columns")
736
+ return self._filtered().with_only_columns(*columns)
737
+
738
+ def _cursor_key(self) -> Sequence[Any]:
739
+ """Return the columns that make a cursor unique."""
740
+ key = getattr(self.model, "__cursor_key__", None)
741
+ if key is not None:
742
+ return key()
743
+ return sa.inspect(self.model, raiseerr=True).primary_key
744
+
745
+ def _cursor_page(
746
+ self,
747
+ rows: list[ModelT],
748
+ *,
749
+ limit: int,
750
+ cursor: str | None = None,
751
+ ) -> CursorPage[ModelT]:
752
+ """Trim the extra row, and read the cursors off the page.
753
+
754
+ Every cursor points at a row of this page, so one only comes back when there
755
+ is a row to point at.
756
+ """
757
+ backwards = cursor is not None and _is_backwards(cursor)
758
+ more = len(rows) > limit
759
+ items = rows[:limit]
760
+ if backwards:
761
+ items.reverse()
762
+ if not items:
763
+ return CursorPage(items=items)
764
+ if backwards:
765
+ # The page ahead is the one this request came from.
766
+ return CursorPage(
767
+ items=items,
768
+ next_cursor=self._cursor_at(items[-1]),
769
+ previous_cursor=self._cursor_at(items[0], backwards=True)
770
+ if more
771
+ else None,
772
+ )
773
+ return CursorPage(
774
+ items=items,
775
+ next_cursor=self._cursor_at(items[-1]) if more else None,
776
+ previous_cursor=self._cursor_at(items[0], backwards=True)
777
+ if cursor is not None
778
+ else None,
779
+ )
780
+
781
+ def _cursor_at(self, row: ModelT, *, backwards: bool = False) -> str:
782
+ """Return the cursor that reads on from this row, one way or the other.
783
+
784
+ Raises:
785
+ NullCursorValueError: if the row is NULL in a column of the ordering,
786
+ which nothing compares against.
787
+ UncomparableOrderingError: if the ordering names a column the rows do
788
+ not carry, such as one belonging to a joined table.
789
+
790
+ """
791
+ ordering = self._keyset_ordering()
792
+ values = []
793
+ for item in ordering:
794
+ value = getattr(row, item.attribute)
795
+ if value is None:
796
+ raise NullCursorValueError(item.attribute)
797
+ values.append(value)
798
+ return _encode(values, backwards=backwards, ordering=_ordering_key(ordering))
799
+
800
+
801
+ def orderable(model: type[Any]) -> Mapping[str, Any]:
802
+ """Return the fields a model can be ordered by name.
803
+
804
+ Every mapped column, unless the model says otherwise with ``__orderable__``:
805
+ a tuple of the names it allows, or a classmethod returning a mapping of name
806
+ to what sorts by it.
807
+ """
808
+ fields = getattr(model, "__orderable__", None)
809
+ if fields is None:
810
+ mapper = sa.inspect(model, raiseerr=True)
811
+ return {attr.key: getattr(model, attr.key) for attr in mapper.column_attrs}
812
+ if callable(fields):
813
+ return fields()
814
+ return {name: _mapped_column(model, name) for name in fields}
815
+
816
+
817
+ def _mapped_column(model: type[Any], name: str) -> Any: # noqa: ANN401
818
+ """Return the model's column of that name, or say the declaration is wrong."""
819
+ column = getattr(model, name, None)
820
+ if not isinstance(column, InstrumentedAttribute):
821
+ raise InvalidOrderFieldError(model.__name__, name)
822
+ return column
823
+
824
+
825
+ def ordered(
826
+ select: sa.Select[Any],
827
+ fields: Mapping[str, Any],
828
+ criteria: Iterable[Any],
829
+ ci_fields: Sequence[str] = (),
830
+ ) -> sa.Select[Any]:
831
+ """Return the statement ordered by these criteria, joining what they need.
832
+
833
+ Raises:
834
+ UnknownOrderFieldError: if a name is not one of the fields.
835
+
836
+ """
837
+ named = list(_flatten(criteria))
838
+ if not named:
839
+ return select
840
+ ci = set(ci_fields)
841
+ clauses = []
842
+ for criterion in named:
843
+ clause, join = _ordering_for(criterion, fields, ci)
844
+ clauses.append(clause)
845
+ if join is not None:
846
+ target, onclause = join
847
+ if not _is_joined(select, target):
848
+ select = select.join(target, onclause)
849
+ return select.order_by(*clauses)
850
+
851
+
852
+ def _ordering_for(
853
+ criterion: Any, # noqa: ANN401
854
+ fields: Mapping[str, Any],
855
+ ci: set[str],
856
+ ) -> tuple[Any, Any]:
857
+ """Return the clause a criterion stands for, and the table it needs."""
858
+ if isinstance(criterion, OrderBy):
859
+ return criterion.expression, (criterion.join, criterion.on)
860
+ if not isinstance(criterion, str):
861
+ return criterion, None
862
+ name, descending, nulls = _parse_sort_field(criterion)
863
+ if name not in fields:
864
+ raise UnknownOrderFieldError(name, list(fields))
865
+ field = fields[name]
866
+ if isinstance(field, OrderBy):
867
+ column, join = field.expression, (field.join, field.on)
868
+ else:
869
+ column, join = field, None
870
+ if name in ci:
871
+ column = _case_insensitive(column)
872
+ return _sort_clause(column, descending=descending, nulls=nulls), join
873
+
874
+
875
+ def _chain(loader: Any, keys: Sequence[Any]) -> Any: # noqa: ANN401
876
+ option = loader(keys[0])
877
+ for key in keys[1:]:
878
+ option = getattr(option, loader.__name__)(key)
879
+ return option
880
+
881
+
882
+ def _direction(clause: Any) -> tuple[sa.ColumnElement[Any], bool]: # noqa: ANN401
883
+ """Return the column an ORDER BY clause names, and whether it runs backwards.
884
+
885
+ `desc()`, `asc()` and the `nulls_*()` variants each wrap the column in one
886
+ more expression, so unwrap until the column itself is left.
887
+ """
888
+ descending = False
889
+ element = clause
890
+ while isinstance(element, sa.UnaryExpression):
891
+ if element.modifier is operators.desc_op:
892
+ descending = True
893
+ elif element.modifier is operators.asc_op:
894
+ descending = False
895
+ element = element.element
896
+ return element, descending
897
+
898
+
899
+ def _column_identity(entry: Any) -> tuple[str | None, str | None]: # noqa: ANN401
900
+ """Return the table and name an ordering clause sorts by.
901
+
902
+ `teams.id` and `players.id` are one name and two columns, and taking one for
903
+ the other drops the tiebreaker a page needs.
904
+ """
905
+ column, _ = _direction(entry)
906
+ table = getattr(column, "table", None)
907
+ return getattr(table, "fullname", None), getattr(column, "key", None)
908
+
909
+
910
+ def _parse_sort_field(field: str) -> tuple[str, bool, str | None]:
911
+ """Split `name[.direction[.nulls]]` into what an ORDER BY needs."""
912
+ name, _, rest = field.partition(".")
913
+ direction, _, nulls = rest.partition(".")
914
+ return name, direction.lower() == "desc", nulls.lower() or None
915
+
916
+
917
+ def _sort_clause(column: Any, *, descending: bool, nulls: str | None) -> Any: # noqa: ANN401
918
+ """Return the ORDER BY clause for one sort field.
919
+
920
+ The direction goes underneath a modifier the model set, such as
921
+ ``sa.nulls_last()``, or the SQL comes out as `NULLS LAST DESC`.
922
+ """
923
+ column, wrapped = _split_nulls(column)
924
+ clause = sa.desc(column) if descending else sa.asc(column)
925
+ nulls = nulls or wrapped
926
+ if nulls == "nulls_first":
927
+ return clause.nulls_first()
928
+ if nulls == "nulls_last":
929
+ return clause.nulls_last()
930
+ return clause
931
+
932
+
933
+ def _flatten(criteria: Iterable[Any]) -> Iterator[Any]:
934
+ """Yield the ordering criteria, taking lists apart and dropping the Nones."""
935
+ for criterion in criteria:
936
+ if criterion is None:
937
+ continue
938
+ if isinstance(criterion, OrderBy):
939
+ yield criterion
940
+ elif isinstance(criterion, (list, tuple, set, frozenset)):
941
+ yield from _flatten(criterion)
942
+ else:
943
+ yield criterion
944
+
945
+
946
+ def _is_joined(select: sa.Select[Any], target: Any) -> bool: # noqa: ANN401
947
+ """Whether this statement already reaches what a field needs."""
948
+ wanted = _join_identity(target)
949
+ if wanted is None:
950
+ return False
951
+ joined = {_join_identity(join[0]) for join in select._setup_joins} # noqa: SLF001
952
+ joined |= {_join_identity(entity) for entity in select.columns_clause_froms}
953
+ return wanted in joined
954
+
955
+
956
+ def _join_identity(target: Any) -> str | None: # noqa: ANN401
957
+ """Return what a join target stands for, as a name that tells two apart.
958
+
959
+ Two aliases of one table are two things to join, and the table's name says
960
+ they are one, so the alias's name counts when there is one.
961
+ """
962
+ selectable = _selectable_of(target)
963
+ if selectable is None:
964
+ return None
965
+ name = getattr(selectable, "fullname", None) or getattr(selectable, "name", None)
966
+ return None if name is None else str(name)
967
+
968
+
969
+ def _selectable_of(target: Any) -> Any: # noqa: ANN401
970
+ """Return the table, alias or subquery a join target stands for."""
971
+ if isinstance(target, sa.Table | sa.Alias | sa.Subquery):
972
+ return target
973
+ inspected = sa.inspect(target, raiseerr=False)
974
+ if inspected is None:
975
+ return None
976
+ relationship = getattr(inspected, "property", None)
977
+ if relationship is not None:
978
+ return getattr(getattr(relationship, "mapper", None), "selectable", None)
979
+ return getattr(inspected, "selectable", None)
980
+
981
+
982
+ def _case_insensitive(column: Any) -> Any: # noqa: ANN401
983
+ """Return the column folded to one case, if it holds text at all."""
984
+ inner, nulls = _split_nulls(column)
985
+ if not isinstance(getattr(inner, "type", None), sa.String):
986
+ return column
987
+ folded = sa.func.lower(inner)
988
+ if nulls == "nulls_last":
989
+ return sa.nulls_last(folded)
990
+ if nulls == "nulls_first":
991
+ return sa.nulls_first(folded)
992
+ return folded
993
+
994
+
995
+ def _split_nulls(column: Any) -> tuple[Any, str | None]: # noqa: ANN401
996
+ """Separate a column from the nulls modifier wrapped around it."""
997
+ if isinstance(column, sa.UnaryExpression):
998
+ if column.modifier is operators.nulls_last_op:
999
+ return column.element, "nulls_last"
1000
+ if column.modifier is operators.nulls_first_op:
1001
+ return column.element, "nulls_first"
1002
+ return column, None
1003
+
1004
+
1005
+ def _as_column(entry: Any) -> sa.ColumnElement[Any]: # noqa: ANN401
1006
+ """Return the column an attribute stands for, and pass a column through."""
1007
+ element = getattr(entry, "__clause_element__", None)
1008
+ return entry if element is None else element()
1009
+
1010
+
1011
+ def _seek(
1012
+ ordering: Sequence[_Ordering],
1013
+ values: Sequence[Any],
1014
+ ) -> sa.ColumnElement[bool]:
1015
+ """Everything after the row the cursor points at, in this ordering.
1016
+
1017
+ One direction is a row comparison, which an index matches. Mixed directions
1018
+ become the equivalent chain of comparisons.
1019
+ """
1020
+ descending = {item.descending for item in ordering}
1021
+ if len(descending) == 1:
1022
+ columns = sa.tuple_(*(item.column for item in ordering))
1023
+ row = sa.tuple_(*values)
1024
+ return columns < row if descending.pop() else columns > row
1025
+
1026
+ terms = []
1027
+ for index, item in enumerate(ordering):
1028
+ column = item.column
1029
+ after = column < values[index] if item.descending else column > values[index]
1030
+ equal = [
1031
+ earlier.column == values[position]
1032
+ for position, earlier in enumerate(ordering[:index])
1033
+ ]
1034
+ terms.append(sa.and_(*equal, after))
1035
+ return sa.or_(*terms)
1036
+
1037
+
1038
+ def _ordering_key(ordering: Sequence[_Ordering]) -> str:
1039
+ """Return a short fingerprint of the ordering, for the cursor to carry.
1040
+
1041
+ A cursor only makes sense under the ordering that produced it, and the
1042
+ values alone cannot tell a different one apart when the types line up.
1043
+ """
1044
+ spelled = "|".join(
1045
+ f"{item.attribute}.{'desc' if item.descending else 'asc'}" for item in ordering
1046
+ )
1047
+ return hashlib.blake2s(spelled.encode(), digest_size=4).hexdigest()
1048
+
1049
+
1050
+ def _encode(values: Sequence[Any], *, backwards: bool = False, ordering: str) -> str:
1051
+ """Return the row's values, and the way to read from them, as one token.
1052
+
1053
+ The direction rides along so that a caller has one thing to hand back
1054
+ whichever page they asked for, and cannot ask for both at once.
1055
+ """
1056
+ payload = json.dumps(
1057
+ {"v": [_as_json(value) for value in values], "b": backwards, "o": ordering},
1058
+ separators=(",", ":"),
1059
+ )
1060
+ return base64.urlsafe_b64encode(payload.encode()).decode().rstrip("=")
1061
+
1062
+
1063
+ def _payload(cursor: str) -> dict[str, Any]:
1064
+ """Return what a cursor holds, or raise if it was not made here."""
1065
+ try:
1066
+ padded = cursor + "=" * (-len(cursor) % 4)
1067
+ payload = json.loads(base64.urlsafe_b64decode(padded.encode()).decode())
1068
+ except (binascii.Error, UnicodeDecodeError, json.JSONDecodeError, ValueError):
1069
+ raise InvalidCursorError from None
1070
+ if not isinstance(payload, dict) or not isinstance(payload.get("v"), list):
1071
+ raise InvalidCursorError
1072
+ return payload
1073
+
1074
+
1075
+ def _is_backwards(cursor: str) -> bool:
1076
+ """Whether this cursor reads the page in front of the one it came from."""
1077
+ return bool(_payload(cursor).get("b"))
1078
+
1079
+
1080
+ def _decode(
1081
+ cursor: str,
1082
+ ordering: Sequence[_Ordering],
1083
+ ordering_key: str,
1084
+ ) -> list[Any]:
1085
+ """Return the values a cursor carries, in the types its columns hold.
1086
+
1087
+ Raises:
1088
+ InvalidCursorError: if the cursor came from another ordering, or from
1089
+ somewhere else entirely.
1090
+
1091
+ """
1092
+ payload = _payload(cursor)
1093
+ values = payload["v"]
1094
+ if payload.get("o") != ordering_key or len(values) != len(ordering):
1095
+ raise InvalidCursorError
1096
+ try:
1097
+ return [
1098
+ _from_json(value, item.column.type)
1099
+ for value, item in zip(values, ordering, strict=True)
1100
+ ]
1101
+ except ValueError:
1102
+ raise InvalidCursorError from None
1103
+
1104
+
1105
+ def _as_json(value: Any) -> Any: # noqa: ANN401
1106
+ if isinstance(value, str | int | float | bool | None):
1107
+ return value
1108
+ return str(value)
1109
+
1110
+
1111
+ def _from_json(value: Any, type_: sa.types.TypeEngine[Any]) -> Any: # noqa: ANN401
1112
+ if value is None:
1113
+ raise InvalidCursorError
1114
+ python_type = type_.python_type
1115
+ if isinstance(value, python_type):
1116
+ return value
1117
+ if hasattr(python_type, "fromisoformat"):
1118
+ return python_type.fromisoformat(value)
1119
+ return python_type(value)
1120
+
1121
+
1122
+ class _Rows(Protocol[RowT_co]):
1123
+ """What the helpers below need of a result: one row, or none."""
1124
+
1125
+ def one(self) -> RowT_co: ...
1126
+
1127
+ def one_or_none(self) -> RowT_co | None: ...
1128
+
1129
+
1130
+ def one_row(rows: _Rows[RowT], name: str) -> RowT:
1131
+ """Return the single row, saying what had none or too many.
1132
+
1133
+ Raises:
1134
+ InstanceNotFoundError: if there is none.
1135
+ MultipleInstancesFoundError: if there is more than one.
1136
+
1137
+ """
1138
+ try:
1139
+ return rows.one()
1140
+ except NoResultFound:
1141
+ raise InstanceNotFoundError(name) from None
1142
+ except MultipleResultsFound:
1143
+ raise MultipleInstancesFoundError(name) from None
1144
+
1145
+
1146
+ def one_row_or_none(rows: _Rows[RowT], name: str) -> RowT | None:
1147
+ """Return the single row or None, saying what had too many.
1148
+
1149
+ Raises:
1150
+ MultipleInstancesFoundError: if there is more than one.
1151
+
1152
+ """
1153
+ try:
1154
+ return rows.one_or_none()
1155
+ except MultipleResultsFound:
1156
+ raise MultipleInstancesFoundError(name) from None