1"""
2Managers for related objects.
3
4These managers provide the API for working with collections of related objects
5through foreign key and many-to-many relationships.
6"""
7
8from collections.abc import Sequence
9from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast
10
11if TYPE_CHECKING:
12 from collections.abc import Callable, Iterable
13
14 from plain.postgres.base import Model
15 from plain.postgres.fields.base import Field
16 from plain.postgres.fields.related import ForeignKeyField, ManyToManyField
17
18import builtins
19
20from plain.postgres import transaction
21from plain.postgres.db import get_connection
22from plain.postgres.dialect import quote_name
23from plain.postgres.expressions import RawSQL, Window
24from plain.postgres.functions import RowNumber
25from plain.postgres.lookups import GreaterThan, LessThanOrEqual
26from plain.postgres.query import QuerySet
27from plain.postgres.query_utils import Q
28from plain.postgres.utils import resolve_callables
29
30# TypeVar for generic manager support
31T = TypeVar("T", bound="Model")
32# TypeVar for custom QuerySet types (defaults to QuerySet[Any] when not specified)
33QS = TypeVar("QS", bound="QuerySet[Any]", default="QuerySet[Any]")
34
35
36def _filter_prefetch_queryset(
37 queryset: QuerySet, field_name: str, instances: Iterable[Model]
38) -> QuerySet:
39 filter_kwargs: dict[str, Any] = {f"{field_name}__in": instances}
40 predicate = Q(**filter_kwargs)
41 if queryset.sql_query.is_sliced:
42 # Use window functions for limited queryset prefetching
43 low_mark, high_mark = queryset.sql_query.low_mark, queryset.sql_query.high_mark
44 order_by = [
45 expr for expr, _ in queryset.sql_query.get_compiler().get_order_by()
46 ]
47 window = Window(RowNumber(), partition_by=field_name, order_by=order_by)
48 predicate &= GreaterThan(window, low_mark)
49 if high_mark is not None:
50 predicate &= LessThanOrEqual(window, high_mark)
51 queryset.sql_query.clear_limits()
52 return queryset.filter(predicate)
53
54
55class BaseRelatedManager(Generic[T, QS]):
56 """
57 Base class for all related object managers.
58
59 All related managers should have a 'query' property that returns a QuerySet.
60 """
61
62 @property
63 def query(self) -> QS:
64 """Access the QuerySet for this relationship."""
65 return self.get_queryset()
66
67 def get_queryset(self) -> QS:
68 """Return the QuerySet for this relationship."""
69 raise NotImplementedError("Subclasses must implement get_queryset()")
70
71
72class ReverseForeignKeyManager(BaseRelatedManager[T, QS]):
73 """
74 Manager for the reverse side of a foreign key relation.
75
76 This manager adds behaviors specific to foreign key relations.
77 """
78
79 # Type hints for attributes
80 model: type[T]
81 instance: Model
82 core_filters: dict[str, Model]
83 allow_null: bool
84
85 def __init__(
86 self, instance: Model, field: ForeignKeyField, related_model: type[Model]
87 ):
88 self.model = cast(type[T], related_model)
89 self.instance = instance
90 # Annotated here rather than as a class-body attribute: ForeignKeyField
91 # is itself a descriptor, so a class-body `field: ForeignKeyField`
92 # would route every `self.field` read through its __get__.
93 self.field: ForeignKeyField = field
94 self.core_filters = {self.field.name: instance}
95 self.allow_null = self.field.allow_null
96
97 def _check_fk_val(self) -> None:
98 field = self.field.target_field
99 if getattr(self.instance, field.name) is None:
100 raise ValueError(
101 f'"{self.instance!r}" needs to have a value for field '
102 f'"{field.name}" before this relationship can be used.'
103 )
104
105 def _apply_rel_filters(self, queryset: QuerySet) -> QuerySet:
106 """
107 Filter the queryset for the instance this manager is bound to.
108 """
109 queryset._defer_next_filter = True
110 queryset = queryset.filter(**self.core_filters)
111 target_field = self.field.target_field
112 rel_obj_id = getattr(self.instance, target_field.name)
113 if rel_obj_id is None:
114 return queryset.none()
115 queryset._known_related_objects = {self.field: {rel_obj_id: self.instance}}
116 return queryset
117
118 def _remove_prefetched_objects(self) -> None:
119 try:
120 self.instance._prefetched_objects_cache.pop(
121 self.field.remote_field.get_cache_name()
122 )
123 except AttributeError, KeyError:
124 pass # nothing to clear from cache
125
126 def get_queryset(self) -> QS:
127 # The instance must have a primary key value to be related against,
128 # though that isn't a guarantee it has actually been saved to the DB.
129 if self.instance.id is None:
130 raise ValueError(
131 f"{self.instance.__class__.__name__!r} instance needs to have a "
132 f"primary key value before this relationship can be used."
133 )
134 try:
135 return self.instance._prefetched_objects_cache[
136 self.field.remote_field.get_cache_name()
137 ]
138 except AttributeError, KeyError:
139 queryset = self.model.query
140 return cast(QS, self._apply_rel_filters(queryset))
141
142 def _get_prefetch_queryset(
143 self, instances: Iterable[Model], queryset: QuerySet | None = None
144 ) -> tuple[
145 QuerySet, Callable[[Model], Any], Callable[[Model], Any], bool, str, bool
146 ]:
147 if queryset is None:
148 queryset = self.model.query
149
150 rel_obj_attr = self.field.get_local_related_value
151 instance_attr = self.field.get_foreign_related_value
152 instances_dict = {instance_attr(inst): inst for inst in instances}
153 queryset = _filter_prefetch_queryset(queryset, self.field.name, instances)
154
155 # Since we just bypassed this class' get_queryset(), we must manage
156 # the reverse relation manually.
157 for rel_obj in queryset:
158 if not self.field.is_cached(rel_obj):
159 instance = instances_dict[rel_obj_attr(rel_obj)]
160 setattr(rel_obj, self.field.name, instance)
161 cache_name = self.field.remote_field.get_cache_name()
162 return queryset, rel_obj_attr, instance_attr, False, cache_name, False
163
164 def add(self, *objs: T, bulk: bool = True) -> None:
165 self._check_fk_val()
166 self._remove_prefetched_objects()
167
168 def check_and_update_obj(obj: Any) -> None:
169 if not isinstance(obj, self.model):
170 raise TypeError(
171 f"'{self.model.model_options.object_name}' instance expected, got {obj!r}"
172 )
173 setattr(obj, self.field.name, self.instance)
174
175 if bulk:
176 ids = []
177 for obj in objs:
178 check_and_update_obj(obj)
179 if obj._state.adding:
180 raise ValueError(
181 f"{obj!r} instance isn't saved. Use bulk=False or save "
182 "the object first."
183 )
184 ids.append(obj.id)
185 self.model._model_meta.base_queryset.filter(id__in=ids).update(
186 **{
187 self.field.name: self.instance,
188 }
189 )
190 else:
191 with transaction.atomic(savepoint=False):
192 for obj in objs:
193 check_and_update_obj(obj)
194 # bulk=False is the documented way to add unsaved objects,
195 # so an obj here may be new or already persisted.
196 if obj._state.adding:
197 obj.create()
198 else:
199 obj.update()
200
201 def create(self, **kwargs: Any) -> T:
202 self._check_fk_val()
203 kwargs[self.field.name] = self.instance
204 return self.model.query.create(**kwargs)
205
206 def get_or_create(
207 self, *, defaults: dict[str, Any] | None = None, **kwargs: Any
208 ) -> tuple[T, bool]:
209 self._check_fk_val()
210 kwargs[self.field.name] = self.instance
211 return self.model.query.get_or_create(defaults=defaults, **kwargs)
212
213 def upsert(
214 self,
215 *,
216 defaults: dict[str, Any] | None = None,
217 create_defaults: dict[str, Any] | None = None,
218 conflict_defaults: dict[str, Any] | None = None,
219 unique_fields: Sequence[Field[Any] | type[Model]],
220 **kwargs: Any,
221 ) -> tuple[T, bool]:
222 self._check_fk_val()
223 kwargs[self.field.name] = self.instance
224 return self.model.query.upsert(
225 defaults=defaults,
226 create_defaults=create_defaults,
227 conflict_defaults=conflict_defaults,
228 unique_fields=unique_fields,
229 **kwargs,
230 )
231
232 def remove(self, *objs: T, bulk: bool = True) -> None:
233 # remove() is only provided if the ForeignKeyField can have a value of null
234 if not self.allow_null:
235 raise AttributeError(
236 f"Cannot call remove() on a related manager for field "
237 f"{self.field.name} where null=False."
238 )
239 if not objs:
240 return
241 self._check_fk_val()
242 val = self.field.get_foreign_related_value(self.instance)
243 old_ids = set()
244 for obj in objs:
245 if not isinstance(obj, self.model):
246 raise TypeError(
247 f"'{self.model.model_options.object_name}' instance expected, got {obj!r}"
248 )
249 # Is obj actually part of this descriptor set?
250 if self.field.get_local_related_value(obj) == val:
251 old_ids.add(obj.id)
252 else:
253 raise self.field.remote_field.model.DoesNotExist(
254 f"{obj!r} is not related to {self.instance!r}."
255 )
256 self._clear(self.query.filter(id__in=old_ids), bulk)
257
258 def clear(self, *, bulk: bool = True) -> None:
259 # clear() is only provided if the ForeignKeyField can have a value of null
260 if not self.allow_null:
261 raise AttributeError(
262 f"Cannot call clear() on a related manager for field "
263 f"{self.field.name} where null=False."
264 )
265 self._check_fk_val()
266 self._clear(self.query, bulk)
267
268 def _clear(self, queryset: QuerySet, bulk: bool) -> None:
269 self._remove_prefetched_objects()
270 if bulk:
271 # `QuerySet.update()` is intrinsically atomic.
272 queryset.update(**{self.field.name: None})
273 else:
274 with transaction.atomic(savepoint=False):
275 for obj in queryset:
276 setattr(obj, self.field.name, None)
277 obj.update(fields=[self.field.name])
278
279 def set(self, objs: Any, *, bulk: bool = True, clear: bool = False) -> None:
280 self._check_fk_val()
281 # Force evaluation of `objs` in case it's a queryset whose value
282 # could be affected by `manager.clear()`. Refs #19816.
283 objs = tuple(objs)
284
285 if self.field.allow_null:
286 with transaction.atomic(savepoint=False):
287 if clear:
288 self.clear(bulk=bulk)
289 self.add(*objs, bulk=bulk)
290 else:
291 old_objs = set(self.query.all())
292 new_objs = []
293 for obj in objs:
294 if obj in old_objs:
295 old_objs.remove(obj)
296 else:
297 new_objs.append(obj)
298
299 self.remove(*old_objs, bulk=bulk)
300 self.add(*new_objs, bulk=bulk)
301 else:
302 self.add(*objs, bulk=bulk)
303
304
305class ManyToManyManager(BaseRelatedManager[T, QS]):
306 """
307 Manager for both forward and reverse sides of a many-to-many relation.
308
309 This manager handles both directions of many-to-many relations with
310 conditional logic for symmetrical relationships (which only apply to
311 forward relations).
312 """
313
314 # Type hints for attributes
315 model: type[T]
316 instance: Model
317 through: type[Model]
318 query_field_name: str
319 prefetch_cache_name: str
320 source_field_name: str
321 target_field_name: str
322 symmetrical: bool
323 core_filters: dict[str, Any]
324 related_val: Any
325
326 def __init__(
327 self,
328 instance: Model,
329 field: ManyToManyField,
330 through: type[Model],
331 related_model: type[Model],
332 is_reverse: bool,
333 symmetrical: bool = False,
334 ):
335 # Set direction-specific attributes
336 if is_reverse:
337 # Reverse: accessing from the target model back to the source
338 self.model = cast(type[T], related_model)
339 self.query_field_name = field.name
340 self.prefetch_cache_name = field.related_query_name()
341 self.source_field_name = field.m2m_reverse_field_name()
342 self.target_field_name = field.m2m_field_name()
343 self.symmetrical = False # Reverse relations are never symmetrical
344 else:
345 # Forward: accessing from the source model to the target
346 self.model = cast(type[T], related_model)
347 self.query_field_name = field.related_query_name()
348 self.prefetch_cache_name = field.name
349 self.source_field_name = field.m2m_field_name()
350 self.target_field_name = field.m2m_reverse_field_name()
351 self.symmetrical = symmetrical
352
353 # Initialize common M2M attributes
354 self.instance = instance
355 self.through = through
356
357 # M2M through model fields are always ForeignKey
358 self.source_field = cast(
359 "ForeignKeyField",
360 self.through._model_meta.get_forward_field(self.source_field_name),
361 )
362 self.target_field = cast(
363 "ForeignKeyField",
364 self.through._model_meta.get_forward_field(self.target_field_name),
365 )
366
367 source_fk = self.source_field
368 rh_name = source_fk.target_field.name
369 self.core_filters = {
370 f"{self.query_field_name}__{rh_name}": getattr(instance, rh_name)
371 }
372
373 self.related_val = source_fk.get_foreign_related_value(instance)
374 if self.related_val is None:
375 raise ValueError(
376 f'"{instance!r}" needs to have a value for field "{rh_name}" before '
377 "this many-to-many relationship can be used."
378 )
379 # The instance must have a primary key value to be related against.
380 if instance.id is None:
381 raise ValueError(
382 f"{instance.__class__.__name__!r} instance needs to have a primary key value before "
383 "a many-to-many relationship can be used."
384 )
385
386 def _apply_rel_filters(self, queryset: QuerySet) -> QuerySet:
387 """Filter the queryset for the instance this manager is bound to."""
388 queryset._defer_next_filter = True
389 return queryset._next_is_sticky().filter(**self.core_filters)
390
391 def _remove_prefetched_objects(self) -> None:
392 try:
393 self.instance._prefetched_objects_cache.pop(self.prefetch_cache_name)
394 except AttributeError, KeyError:
395 pass # nothing to clear from cache
396
397 def get_queryset(self) -> QS:
398 try:
399 return self.instance._prefetched_objects_cache[self.prefetch_cache_name]
400 except AttributeError, KeyError:
401 queryset = self.model.query
402 return cast(QS, self._apply_rel_filters(queryset))
403
404 def _get_prefetch_queryset(
405 self, instances: Iterable[Model], queryset: QuerySet | None = None
406 ) -> tuple[
407 QuerySet, Callable[[Model], Any], Callable[[Model], Any], bool, str, bool
408 ]:
409 if queryset is None:
410 queryset = self.model.query
411
412 queryset = _filter_prefetch_queryset(
413 queryset._next_is_sticky(), self.query_field_name, instances
414 )
415
416 # M2M: need to annotate the query in order to get the primary model
417 # that the secondary model was actually related to.
418 from typing import cast
419
420 from plain.postgres.fields.related import ForeignKeyField
421
422 fk = cast(
423 ForeignKeyField,
424 self.through._model_meta.get_forward_field(self.source_field_name),
425 ) # M2M through model fields are always ForeignKey
426 join_table = fk.model.model_options.db_table
427 qn = quote_name
428 # Expose the through-table FK column so the prefetch loader can match
429 # each secondary row back to the primary instance that owns it.
430 queryset = queryset.annotate(
431 **{
432 f"_prefetch_val_{fk.name}": RawSQL(
433 f"{qn(join_table)}.{qn(fk.column)}", []
434 )
435 }
436 )
437 conn = get_connection()
438 target_field = fk.target_field
439 return (
440 queryset,
441 lambda result: getattr(result, f"_prefetch_val_{fk.name}"),
442 lambda inst: target_field.get_db_prep_value(
443 target_field.value_from_object(inst), conn
444 ),
445 False,
446 self.prefetch_cache_name,
447 False,
448 )
449
450 def clear(self) -> None:
451 with transaction.atomic(savepoint=False):
452 self._remove_prefetched_objects()
453 filters = self._build_remove_filters(self.model.query)
454 self.through.query.filter(filters).delete()
455
456 def set(
457 self,
458 objs: Any,
459 *,
460 clear: bool = False,
461 through_defaults: dict[str, Any] | None = None,
462 ) -> None:
463 # Force evaluation of `objs` in case it's a queryset whose value
464 # could be affected by `manager.clear()`. Refs #19816.
465 objs = tuple(objs)
466
467 with transaction.atomic(savepoint=False):
468 if clear:
469 self.clear()
470 self.add(*objs, through_defaults=through_defaults)
471 else:
472 target_name = self.target_field.target_field.name
473 assert target_name is not None
474 old_ids = set(self.query.values_list(target_name, flat=True))
475
476 new_objs = []
477 for obj in objs:
478 fk_val = (
479 self.target_field.get_foreign_related_value(obj)
480 if isinstance(obj, self.model)
481 else self.target_field.get_prep_value(obj)
482 )
483 if fk_val in old_ids:
484 old_ids.remove(fk_val)
485 else:
486 new_objs.append(obj)
487
488 self.remove(*old_ids)
489 self.add(*new_objs, through_defaults=through_defaults)
490
491 def create(
492 self, *, through_defaults: dict[str, Any] | None = None, **kwargs: Any
493 ) -> T:
494 new_obj = self.model.query.create(**kwargs)
495 self.add(new_obj, through_defaults=through_defaults)
496 return new_obj
497
498 def get_or_create(
499 self,
500 *,
501 defaults: dict[str, Any] | None = None,
502 through_defaults: dict[str, Any] | None = None,
503 **kwargs: Any,
504 ) -> tuple[T, bool]:
505 obj, created = self.model.query.get_or_create(defaults=defaults, **kwargs)
506 # The get() isn't scoped to this relation, so an existing object may
507 # not be related to this instance yet. add() skips through-rows that
508 # are already there, so it's safe either way.
509 self.add(obj, through_defaults=through_defaults)
510 return obj, created
511
512 def upsert(
513 self,
514 *,
515 through_defaults: dict[str, Any] | None = None,
516 defaults: dict[str, Any] | None = None,
517 create_defaults: dict[str, Any] | None = None,
518 conflict_defaults: dict[str, Any] | None = None,
519 unique_fields: Sequence[Field[Any] | type[Model]],
520 **kwargs: Any,
521 ) -> tuple[T, bool]:
522 obj, created = self.model.query.upsert(
523 defaults=defaults,
524 create_defaults=create_defaults,
525 conflict_defaults=conflict_defaults,
526 unique_fields=unique_fields,
527 **kwargs,
528 )
529 # upsert() writes the target row, not the relationship, and the row may
530 # already exist without being related to this instance. add() skips
531 # through-rows that are already there, so it's safe either way.
532 self.add(obj, through_defaults=through_defaults)
533 return obj, created
534
535 def _get_target_ids(self, target_field_name: str, objs: Any) -> builtins.set[Any]:
536 """Return the set of ids of `objs` that the target field references."""
537 from typing import cast
538
539 from plain.postgres import Model
540 from plain.postgres.fields.related import ForeignKeyField
541
542 target_ids: set[Any] = set()
543 target_field = cast(
544 ForeignKeyField,
545 self.through._model_meta.get_forward_field(target_field_name),
546 ) # M2M through model fields are always ForeignKey
547 for obj in objs:
548 if isinstance(obj, self.model):
549 target_id = target_field.get_foreign_related_value(obj)
550 if target_id is None:
551 raise ValueError(
552 f'Cannot add "{obj!r}": the value for field "{target_field_name}" is None'
553 )
554 target_ids.add(target_id)
555 elif isinstance(obj, Model):
556 raise TypeError(
557 f"'{self.model.model_options.object_name}' instance expected, got {obj!r}"
558 )
559 else:
560 target_ids.add(target_field.get_prep_value(obj))
561 return target_ids
562
563 def _get_missing_target_ids(
564 self,
565 source_field_name: str,
566 target_field_name: str,
567 target_ids: builtins.set[Any],
568 ) -> builtins.set[Any]:
569 """Return the subset of ids of `objs` that aren't already assigned to this relationship."""
570 vals = self.through.query.values_list(target_field_name, flat=True).filter(
571 **{
572 source_field_name: self.related_val,
573 f"{target_field_name}__in": target_ids,
574 }
575 )
576 return target_ids.difference(vals)
577
578 def _add_items(
579 self,
580 source_field_name: str,
581 target_field_name: str,
582 *objs: Any,
583 through_defaults: dict[str, Any] | None = None,
584 ) -> None:
585 if not objs:
586 return
587
588 through_defaults = dict(resolve_callables(through_defaults or {}))
589 target_ids = self._get_target_ids(target_field_name, objs)
590
591 missing_target_ids = self._get_missing_target_ids(
592 source_field_name, target_field_name, target_ids
593 )
594 with transaction.atomic(savepoint=False):
595 # Add the ones that aren't there already.
596 self.through.query.bulk_create(
597 [
598 self.through(
599 **through_defaults,
600 **{
601 source_field_name: self.related_val,
602 target_field_name: target_id,
603 },
604 )
605 for target_id in missing_target_ids
606 ],
607 )
608
609 def _remove_items(
610 self, source_field_name: str, target_field_name: str, *objs: Any
611 ) -> None:
612 if not objs:
613 return
614
615 # Check that all the objects are of the right type
616 old_ids = set()
617 for obj in objs:
618 if isinstance(obj, self.model):
619 fk_val = self.target_field.get_foreign_related_value(obj)
620 old_ids.add(fk_val)
621 else:
622 old_ids.add(obj)
623
624 with transaction.atomic(savepoint=False):
625 target_model_qs = self.model.query
626 if target_model_qs._has_filters():
627 old_vals = target_model_qs.filter(
628 **{f"{self.target_field.target_field.name}__in": old_ids}
629 )
630 else:
631 old_vals = old_ids
632 filters = self._build_remove_filters(old_vals)
633 self.through.query.filter(filters).delete()
634
635 def _build_remove_filters(self, removed_vals: Any) -> Any:
636 filters = Q.create([(self.source_field_name, self.related_val)])
637 # No need to add a subquery condition if removed_vals is a QuerySet without
638 # filters.
639 removed_vals_filters = (
640 not isinstance(removed_vals, QuerySet) or removed_vals._has_filters()
641 )
642 if removed_vals_filters:
643 filters = filters & Q.create(
644 [(f"{self.target_field_name}__in", removed_vals)]
645 )
646 # Add symmetrical filters for forward symmetrical relations
647 if self.symmetrical:
648 symmetrical_filters = Q.create([(self.target_field_name, self.related_val)])
649 if removed_vals_filters:
650 symmetrical_filters = symmetrical_filters & Q.create(
651 [(f"{self.source_field_name}__in", removed_vals)]
652 )
653 filters = filters | symmetrical_filters
654 return filters
655
656 def add(self, *objs: T, through_defaults: dict[str, Any] | None = None) -> None:
657 self._remove_prefetched_objects()
658 with transaction.atomic(savepoint=False):
659 self._add_items(
660 self.source_field_name,
661 self.target_field_name,
662 *objs,
663 through_defaults=through_defaults,
664 )
665 # If this is a symmetrical m2m relation to self, add the mirror
666 # entry in the m2m table.
667 if self.symmetrical:
668 self._add_items(
669 self.target_field_name,
670 self.source_field_name,
671 *objs,
672 through_defaults=through_defaults,
673 )
674
675 def remove(self, *objs: T) -> None:
676 self._remove_prefetched_objects()
677 self._remove_items(self.source_field_name, self.target_field_name, *objs)