1from typing import TYPE_CHECKING, Any
  2
  3from plain.postgres.lookups import (
  4    Exact,
  5    GreaterThan,
  6    GreaterThanOrEqual,
  7    In,
  8    IsNull,
  9    LessThan,
 10    LessThanOrEqual,
 11    Lookup,
 12)
 13
 14if TYPE_CHECKING:
 15    from plain.postgres.connection import DatabaseConnection
 16    from plain.postgres.sql.compiler import SQLCompiler
 17
 18
 19def _related_target_field(lhs: Any) -> Any:
 20    """The remote field a relational lhs filters against (e.g. the target's id),
 21    or None when the lhs is not a relation.
 22
 23    Narrowing on the relation types keeps this explicit and type-checked rather
 24    than probing for a ``path_infos`` attribute with hasattr().
 25    """
 26    from plain.postgres.fields.related import RelatedField
 27    from plain.postgres.fields.reverse_related import ForeignObjectRel
 28
 29    output_field = lhs.output_field
 30    if isinstance(output_field, RelatedField | ForeignObjectRel):
 31        return output_field.path_infos[-1].target_field
 32    return None
 33
 34
 35def get_normalized_value(value: Any, lhs: Any) -> Any:
 36    from plain.postgres import Model
 37
 38    if isinstance(value, Model):
 39        if value.id is None:
 40            raise ValueError("Model instances passed to related filters must be saved.")
 41        # FK relations are always single-column, targeting the remote `id`.
 42        source = lhs.output_field.path_infos[-1].target_field
 43        return getattr(value, source.name)
 44    return value
 45
 46
 47class RelatedIn(In):
 48    def get_prep_lookup(self) -> list[Any]:
 49        if self.rhs_is_direct_value():
 50            self.rhs = [get_normalized_value(val, self.lhs) for val in self.rhs]
 51            # We need to run the related field's get_prep_value(). Consider
 52            # case ForeignKeyField to IntegerField given value 'abc'. The
 53            # ForeignKeyField itself doesn't have validation for non-integers,
 54            # so we must run validation using the target field.
 55            if target_field := _related_target_field(self.lhs):
 56                self.rhs = [target_field.get_prep_value(v) for v in self.rhs]
 57        return super().get_prep_lookup()
 58
 59
 60class RelatedLookupMixin(Lookup):
 61    # Type hints for attributes/methods expected from Lookup base class
 62    lhs: Any
 63    rhs: Any
 64    prepare_rhs: bool
 65    lookup_name: str | None
 66
 67    def get_prep_lookup(self) -> Any:
 68        if not hasattr(self.rhs, "resolve_expression"):
 69            self.rhs = get_normalized_value(self.rhs, self.lhs)
 70            # We need to run the related field's get_prep_value(). Consider case
 71            # ForeignKeyField to IntegerField given value 'abc'. The ForeignKeyField itself
 72            # doesn't have validation for non-integers, so we must run validation
 73            # using the target field.
 74            if self.prepare_rhs and (target_field := _related_target_field(self.lhs)):
 75                self.rhs = target_field.get_prep_value(self.rhs)
 76
 77        return super().get_prep_lookup()
 78
 79    def as_sql(
 80        self, compiler: SQLCompiler, connection: DatabaseConnection
 81    ) -> tuple[str, list[Any]]:
 82        sql, params = super().as_sql(compiler, connection)
 83        return sql, list(params)
 84
 85
 86class RelatedExact(RelatedLookupMixin, Exact):
 87    pass
 88
 89
 90class RelatedLessThan(RelatedLookupMixin, LessThan):
 91    pass
 92
 93
 94class RelatedGreaterThan(RelatedLookupMixin, GreaterThan):
 95    pass
 96
 97
 98class RelatedGreaterThanOrEqual(RelatedLookupMixin, GreaterThanOrEqual):
 99    pass
100
101
102class RelatedLessThanOrEqual(RelatedLookupMixin, LessThanOrEqual):
103    pass
104
105
106class RelatedIsNull(RelatedLookupMixin, IsNull):
107    pass