plain.postgres 0.84.0__py3-none-any.whl

plain/postgres/query_utils.py

@@ -0,0 +1,456 @@
+"""
+Various data structures used in query construction.
+
+Factored out from plain.postgres.query to avoid making the main module very
+large and/or so that they can be used by other modules without getting into
+circular import difficulties.
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+import logging
+from collections.abc import Callable, Generator
+from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple, Self
+
+from plain.postgres.constants import LOOKUP_SEP
+from plain.postgres.db import DatabaseError
+from plain.postgres.exceptions import FieldError
+from plain.utils import tree
+
+if TYPE_CHECKING:
+    from plain.postgres.base import Model
+    from plain.postgres.connection import DatabaseConnection
+    from plain.postgres.fields import Field
+    from plain.postgres.fields.related import ForeignKeyField
+    from plain.postgres.fields.reverse_related import ForeignObjectRel
+    from plain.postgres.lookups import Lookup, Transform
+    from plain.postgres.meta import Meta
+    from plain.postgres.sql.compiler import SQLCompiler
+    from plain.postgres.sql.where import WhereNode
+
+logger = logging.getLogger("plain.postgres")
+
+
+class PathInfo(NamedTuple):
+    """Information about a relation path when converting lookups (fk__somecol).
+
+    Describes the relation in Model terms (Meta and Fields for both
+    sides of the relation). The join_field is the field backing the relation.
+    """
+
+    from_meta: Meta
+    to_meta: Meta
+    target_fields: tuple[Field, ...]
+    join_field: ForeignKeyField | ForeignObjectRel
+    m2m: bool
+    direct: bool
+    filtered_relation: FilteredRelation | None
+
+
+def subclasses(cls: type) -> Generator[type]:
+    yield cls
+    for subclass in cls.__subclasses__():
+        yield from subclasses(subclass)
+
+
+class Q(tree.Node):
+    """
+    Encapsulate filters as objects that can then be combined logically (using
+    `&` and `|`).
+    """
+
+    # Connection types
+    AND = "AND"
+    OR = "OR"
+    XOR = "XOR"
+    default = AND
+    conditional = True
+
+    def __init__(
+        self,
+        *args: Any,
+        _connector: str | None = None,
+        _negated: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(
+            children=[*args, *sorted(kwargs.items())],
+            connector=_connector,
+            negated=_negated,
+        )
+
+    def _combine(self, other: Any, conn: str) -> Q:
+        if getattr(other, "conditional", False) is False:
+            raise TypeError(other)
+        if not self:
+            return other.copy()
+        if not other and isinstance(other, Q):
+            return self.copy()
+
+        obj = self.create(connector=conn)
+        obj.add(self, conn)
+        obj.add(other, conn)
+        return obj
+
+    def __or__(self, other: Any) -> Q:
+        return self._combine(other, self.OR)
+
+    def __and__(self, other: Any) -> Q:
+        return self._combine(other, self.AND)
+
+    def __xor__(self, other: Any) -> Q:
+        return self._combine(other, self.XOR)
+
+    def __invert__(self) -> Q:
+        obj = self.copy()
+        obj.negate()
+        return obj
+
+    def resolve_expression(
+        self,
+        query: Any = None,
+        allow_joins: bool = True,
+        reuse: Any = None,
+        summarize: bool = False,
+        for_save: bool = False,
+    ) -> WhereNode:
+        # We must promote any new joins to left outer joins so that when Q is
+        # used as an expression, rows aren't filtered due to joins.
+        clause, joins = query._add_q(
+            self,
+            reuse,
+            allow_joins=allow_joins,
+            split_subq=False,
+            check_filterable=False,
+            summarize=summarize,
+        )
+        query.promote_joins(joins)
+        return clause
+
+    def flatten(self) -> Generator[Any]:
+        """
+        Recursively yield this Q object and all subexpressions, in depth-first
+        order.
+        """
+        yield self
+        for child in self.children:
+            if isinstance(child, tuple):
+                # Use the lookup.
+                child = child[1]
+            if hasattr(child, "flatten"):
+                yield from child.flatten()
+            else:
+                yield child
+
+    def check(self, against: dict[str, Any]) -> bool:
+        """
+        Do a database query to check if the expressions of the Q instance
+        matches against the expressions.
+        """
+        # Avoid circular imports.
+        from plain.postgres.expressions import ResolvableExpression, Value
+        from plain.postgres.fields import BooleanField
+        from plain.postgres.functions import Coalesce
+        from plain.postgres.sql import SINGLE, Query
+
+        query = Query(None)
+        for name, value in against.items():
+            if not isinstance(value, ResolvableExpression):
+                value = Value(value)
+            query.add_annotation(value, name, select=False)
+        query.add_annotation(Value(1), "_check")
+        # This will raise a FieldError if a field is missing in "against".
+        query.add_q(Q(Coalesce(self, True, output_field=BooleanField())))
+        compiler = query.get_compiler()
+        try:
+            return compiler.execute_sql(SINGLE) is not None
+        except DatabaseError as e:
+            logger.warning("Got a database error calling check() on %r: %s", self, e)
+            return True
+
+    def deconstruct(self) -> tuple[str, tuple[Any, ...], dict[str, Any]]:
+        path = f"{self.__class__.__module__}.{self.__class__.__name__}"
+        if path.startswith("plain.postgres.query_utils"):
+            path = path.replace("plain.postgres.query_utils", "plain.postgres")
+        args = tuple(self.children)
+        kwargs: dict[str, Any] = {}
+        if self.connector != self.default:
+            kwargs["_connector"] = self.connector
+        if self.negated:
+            kwargs["_negated"] = True
+        return path, args, kwargs
+
+
+class class_or_instance_method:
+    """
+    Hook used in RegisterLookupMixin to return partial functions depending on
+    the caller type (instance or class of models.Field).
+    """
+
+    def __init__(self, class_method: Any, instance_method: Any) -> None:
+        self.class_method = class_method
+        self.instance_method = instance_method
+
+    def __get__(self, instance: Any, owner: type) -> Any:
+        if instance is None:
+            return functools.partial(self.class_method, owner)
+        return functools.partial(self.instance_method, instance)
+
+
+class RegisterLookupMixin:
+    class_lookups: ClassVar[dict[str, type[Lookup | Transform]]]
+
+    def _get_lookup(self, lookup_name: str) -> type[Lookup | Transform] | None:
+        return self.get_lookups().get(lookup_name, None)
+
+    @functools.cache
+    def get_class_lookups(cls: type[Self]) -> dict[str, type[Lookup | Transform]]:
+        class_lookups = [
+            parent.__dict__.get("class_lookups", {}) for parent in inspect.getmro(cls)
+        ]
+        return cls.merge_dicts(class_lookups)
+
+    def get_instance_lookups(self) -> dict[str, type[Lookup | Transform]]:
+        class_lookups = self.get_class_lookups()
+        if instance_lookups := getattr(self, "instance_lookups", None):
+            return {**class_lookups, **instance_lookups}
+        return class_lookups
+
+    get_lookups = class_or_instance_method(get_class_lookups, get_instance_lookups)
+    get_class_lookups: ClassVar[classmethod[Any, ..., Any]] = classmethod(
+        get_class_lookups
+    )
+
+    def get_lookup(self, lookup_name: str) -> type[Lookup] | None:
+        from plain.postgres.lookups import Lookup
+
+        found = self._get_lookup(lookup_name)
+        if found is None:
+            # output_field is a Field which inherits from RegisterLookupMixin
+            if output_field := getattr(self, "output_field", None):
+                return output_field.get_lookup(lookup_name)
+        if found is not None and not issubclass(found, Lookup):
+            return None
+        return found
+
+    def get_transform(
+        self, lookup_name: str
+    ) -> type[Transform] | Callable[..., Any] | None:
+        from plain.postgres.lookups import Transform
+
+        found = self._get_lookup(lookup_name)
+        if found is None:
+            # output_field is a Field which inherits from RegisterLookupMixin
+            if output_field := getattr(self, "output_field", None):
+                return output_field.get_transform(lookup_name)
+        if found is not None and not issubclass(found, Transform):
+            return None
+        return found
+
+    @staticmethod
+    def merge_dicts(
+        dicts: list[dict[str, type[Lookup | Transform]]],
+    ) -> dict[str, type[Lookup | Transform]]:
+        """
+        Merge dicts in reverse to preference the order of the original list. e.g.,
+        merge_dicts([a, b]) will preference the keys in 'a' over those in 'b'.
+        """
+        merged: dict[str, type[Lookup | Transform]] = {}
+        for d in reversed(dicts):
+            merged.update(d)
+        return merged
+
+    @classmethod
+    def _clear_cached_class_lookups(cls: type[Self]) -> None:
+        for subclass in subclasses(cls):
+            if cached := getattr(subclass, "get_class_lookups", None):
+                cached.cache_clear()
+
+    def register_class_lookup(
+        cls: type[Self],
+        lookup: type[Lookup | Transform],
+        lookup_name: str | None = None,
+    ) -> type[Lookup | Transform]:
+        if lookup_name is None:
+            lookup_name = lookup.lookup_name
+        assert lookup_name is not None, "lookup_name must be set on the lookup class"
+        if "class_lookups" not in cls.__dict__:
+            cls.class_lookups = {}
+        cls.class_lookups[lookup_name] = lookup
+        cls._clear_cached_class_lookups()
+        return lookup
+
+    def register_instance_lookup(
+        self, lookup: type[Lookup | Transform], lookup_name: str | None = None
+    ) -> type[Lookup | Transform]:
+        if lookup_name is None:
+            lookup_name = lookup.lookup_name
+        if "instance_lookups" not in self.__dict__:
+            self.instance_lookups = {}
+        self.instance_lookups[lookup_name] = lookup
+        return lookup
+
+    register_lookup = class_or_instance_method(
+        register_class_lookup, register_instance_lookup
+    )
+    register_class_lookup: ClassVar[classmethod[Any, ..., Any]] = classmethod(
+        register_class_lookup
+    )
+
+    def _unregister_class_lookup(
+        cls: type[Self],
+        lookup: type[Lookup | Transform],
+        lookup_name: str | None = None,
+    ) -> None:
+        """
+        Remove given lookup from cls lookups. For use in tests only as it's
+        not thread-safe.
+        """
+        if lookup_name is None:
+            lookup_name = lookup.lookup_name
+        assert lookup_name is not None, "lookup_name must be set on the lookup class"
+        del cls.class_lookups[lookup_name]
+        cls._clear_cached_class_lookups()
+
+    def _unregister_instance_lookup(
+        self, lookup: type[Lookup | Transform], lookup_name: str | None = None
+    ) -> None:
+        """
+        Remove given lookup from instance lookups. For use in tests only as
+        it's not thread-safe.
+        """
+        if lookup_name is None:
+            lookup_name = lookup.lookup_name
+        del self.instance_lookups[lookup_name]
+
+    _unregister_lookup = class_or_instance_method(
+        _unregister_class_lookup, _unregister_instance_lookup
+    )
+    _unregister_class_lookup: ClassVar[classmethod[Any, ..., Any]] = classmethod(
+        _unregister_class_lookup
+    )
+
+
+def select_related_descend(
+    field: Any,
+    restricted: bool | None,
+    requested: dict[str, Any] | None,
+    select_mask: Any,
+    reverse: bool = False,
+) -> bool:
+    """
+    Return True if this field should be used to descend deeper for
+    select_related() purposes. Used by both the query construction code
+    (compiler.get_related_selections()) and the model instance creation code
+    (compiler.klass_info).
+
+    Arguments:
+     * field - the field to be checked
+     * restricted - a boolean field, indicating if the field list has been
+       manually restricted using a requested clause)
+     * requested - The select_related() dictionary.
+     * select_mask - the dictionary of selected fields.
+     * reverse - boolean, True if we are checking a reverse select related
+    """
+    from plain.postgres.fields.related import RelatedField
+
+    if not isinstance(field, RelatedField):
+        return False
+    if restricted:
+        assert requested is not None, "requested must be provided when restricted=True"
+        if reverse and field.related_query_name() not in requested:
+            return False
+        if not reverse and field.name not in requested:
+            return False
+    if not restricted and field.allow_null:
+        return False
+    if (
+        restricted
+        and select_mask
+        and field.name in requested  # type: ignore[operator]
+        and field not in select_mask
+    ):
+        raise FieldError(
+            f"Field {field.model.model_options.object_name}.{field.name} cannot be both "
+            "deferred and traversed using select_related at the same time."
+        )
+    return True
+
+
+def refs_expression(
+    lookup_parts: list[str], annotations: dict[str, Any]
+) -> tuple[str | None, tuple[str, ...]]:
+    """
+    Check if the lookup_parts contains references to the given annotations set.
+    Because the LOOKUP_SEP is contained in the default annotation names, check
+    each prefix of the lookup_parts for a match.
+    """
+    for n in range(1, len(lookup_parts) + 1):
+        level_n_lookup = LOOKUP_SEP.join(lookup_parts[0:n])
+        if annotations.get(level_n_lookup):
+            return level_n_lookup, tuple(lookup_parts[n:])
+    return None, ()
+
+
+def check_rel_lookup_compatibility(
+    model: type[Model], target_meta: Meta, field: Field | ForeignObjectRel
+) -> bool:
+    """
+    Check that model is compatible with target_meta. Compatibility
+    is OK if:
+      1) model and meta.model match (where proxy inheritance is removed)
+      2) model is parent of meta's model or the other way around
+    """
+
+    def check(meta: Meta) -> bool:
+        return model == meta.model
+
+    # If the field is a primary key, then doing a query against the field's
+    # model is ok, too. Consider the case:
+    # class Restaurant(models.Model):
+    #     place = OneToOneField(Place, primary_key=True):
+    # Restaurant.query.filter(id__in=Restaurant.query.all()).
+    # If we didn't have the primary key check, then id__in (== place__in) would
+    # give Place's meta as the target meta, but Restaurant isn't compatible
+    # with that. This logic applies only to primary keys, as when doing __in=qs,
+    # we are going to turn this into __in=qs.values('id') later on.
+    return check(target_meta) or (
+        getattr(field, "primary_key", False) and check(field.model._model_meta)
+    )
+
+
+class FilteredRelation:
+    """Specify custom filtering in the ON clause of SQL joins."""
+
+    def __init__(self, relation_name: str, *, condition: Q = Q()) -> None:
+        if not relation_name:
+            raise ValueError("relation_name cannot be empty.")
+        self.relation_name = relation_name
+        self.alias: str | None = None
+        if not isinstance(condition, Q):
+            raise ValueError("condition argument must be a Q() instance.")
+        self.condition = condition
+        self.path: list[str] = []
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, self.__class__):
+            return NotImplemented
+        return (
+            self.relation_name == other.relation_name
+            and self.alias == other.alias
+            and self.condition == other.condition
+        )
+
+    def clone(self) -> FilteredRelation:
+        clone = FilteredRelation(self.relation_name, condition=self.condition)
+        clone.alias = self.alias
+        clone.path = self.path[:]
+        return clone
+
+    def as_sql(self, compiler: SQLCompiler, connection: DatabaseConnection) -> Any:
+        # Resolve the condition in Join.filtered_relation.
+        query = compiler.query
+        where = query.build_filtered_relation_q(self.condition, reuse=set(self.path))
+        return compiler.compile(where)

plain/postgres/registry.py

@@ -0,0 +1,217 @@
+from __future__ import annotations
+
+import functools
+import warnings
+from collections import defaultdict
+from collections.abc import Callable
+from functools import partial
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from plain.postgres.base import Model
+
+
+class ModelsRegistryNotReady(Exception):
+    """The plain.postgres registry is not populated yet"""
+
+    pass
+
+
+class ModelsRegistry:
+    def __init__(self) -> None:
+        # Mapping of app labels => model names => model classes. Every time a
+        # model is imported, ModelBase.__new__ calls packages.register_model which
+        # creates an entry in all_models. All imported models are registered,
+        # regardless of whether they're defined in an installed application
+        # and whether the registry has been populated. Since it isn't possible
+        # to reimport a module safely (it could reexecute initialization code)
+        # all_models is never overridden or reset.
+        self.all_models: defaultdict[str, dict[str, type[Model]]] = defaultdict(dict)
+
+        # Maps ("package_label", "modelname") tuples to lists of functions to be
+        # called when the corresponding model is ready. Used by this class's
+        # `lazy_model_operation()` and `do_pending_operations()` methods.
+        self._pending_operations: defaultdict[
+            tuple[str, str], list[Callable[[type[Model]], None]]
+        ] = defaultdict(list)
+
+        self.ready: bool = False
+
+    def check_ready(self) -> None:
+        """Raise an exception if all models haven't been imported yet."""
+        if not self.ready:
+            raise ModelsRegistryNotReady("Models aren't loaded yet.")
+
+    # This method is performance-critical at least for Plain's test suite.
+    @functools.cache
+    def get_models(self, *, package_label: str = "") -> list[type[Model]]:
+        """
+        Return a list of all installed models.
+
+        By default, the following models aren't included:
+
+        - auto-created models for many-to-many relations without
+          an explicit intermediate table,
+
+        Set the corresponding keyword argument to True to include such models.
+        """
+
+        self.check_ready()
+
+        models = []
+
+        # Get models for a single package
+        if package_label:
+            package_models = self.all_models[package_label]
+            for model in package_models.values():
+                models.append(model)
+            return models
+
+        # Get models for all packages
+        for package_models in self.all_models.values():
+            for model in package_models.values():
+                models.append(model)
+
+        return models
+
+    def get_model(
+        self,
+        package_label: str,
+        model_name: str | None = None,
+        require_ready: bool = True,
+    ) -> type[Model]:
+        """
+        Return the model matching the given package_label and model_name.
+
+        As a shortcut, package_label may be in the form <package_label>.<model_name>.
+
+        model_name is case-insensitive.
+
+        Raise LookupError if no application exists with this label, or no
+        model exists with this name in the application. Raise ValueError if
+        called with a single argument that doesn't contain exactly one dot.
+        """
+
+        if require_ready:
+            self.check_ready()
+
+        if model_name is None:
+            package_label, model_name = package_label.split(".")
+
+        package_models = self.all_models[package_label]
+        return package_models[model_name.lower()]
+
+    def register_model(self, package_label: str, model: type[Model]) -> None:
+        # Since this method is called when models are imported, it cannot
+        # perform imports because of the risk of import loops. It mustn't
+        # call get_package_config().
+        model_name = model.model_options.model_name
+        app_models = self.all_models[package_label]
+        if model_name in app_models:
+            if (
+                model.__name__ == app_models[model_name].__name__
+                and model.__module__ == app_models[model_name].__module__
+            ):
+                warnings.warn(
+                    f"Model '{package_label}.{model_name}' was already registered. Reloading models is not "
+                    "advised as it can lead to inconsistencies, most notably with "
+                    "related models.",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+            else:
+                raise RuntimeError(
+                    f"Conflicting '{model_name}' models in application '{package_label}': {app_models[model_name]} and {model}."
+                )
+        app_models[model_name] = model
+        self.do_pending_operations(model)
+        self.clear_cache()
+
+    def _get_registered_model(self, package_label: str, model_name: str) -> type[Model]:
+        """
+        Similar to get_model(), but doesn't require that an app exists with
+        the given package_label.
+
+        It's safe to call this method at import time, even while the registry
+        is being populated.
+        """
+        model = self.all_models[package_label].get(model_name.lower())
+        if model is None:
+            raise LookupError(f"Model '{package_label}.{model_name}' not registered.")
+        return model
+
+    def clear_cache(self) -> None:
+        """
+        Clear all internal caches, for methods that alter the app registry.
+
+        This is mostly used in tests.
+        """
+        # Call expire cache on each model. This will purge
+        # the relation tree and the fields cache.
+        self.get_models.cache_clear()
+        if self.ready:
+            # Circumvent self.get_models() to prevent that the cache is refilled.
+            # This particularly prevents that an empty value is cached while cloning.
+            for package_models in self.all_models.values():
+                for model in package_models.values():
+                    model._model_meta._expire_cache()
+
+    def lazy_model_operation(
+        self, function: Callable[..., None], *model_keys: tuple[str, str]
+    ) -> None:
+        """
+        Take a function and a number of ("package_label", "modelname") tuples, and
+        when all the corresponding models have been imported and registered,
+        call the function with the model classes as its arguments.
+
+        The function passed to this method must accept exactly n models as
+        arguments, where n=len(model_keys).
+        """
+        # Base case: no arguments, just execute the function.
+        if not model_keys:
+            function()
+        # Recursive case: take the head of model_keys, wait for the
+        # corresponding model class to be imported and registered, then apply
+        # that argument to the supplied function. Pass the resulting partial
+        # to lazy_model_operation() along with the remaining model args and
+        # repeat until all models are loaded and all arguments are applied.
+        else:
+            next_model, *more_models = model_keys
+
+            # This will be executed after the class corresponding to next_model
+            # has been imported and registered.
+            def apply_next_model(model: type[Model]) -> None:
+                next_function = partial(function, model)
+                self.lazy_model_operation(next_function, *more_models)
+
+            # If the model has already been imported and registered, partially
+            # apply it to the function now. If not, add it to the list of
+            # pending operations for the model, where it will be executed with
+            # the model class as its sole argument once the model is ready.
+            try:
+                model_class = self._get_registered_model(*next_model)
+            except LookupError:
+                self._pending_operations[next_model].append(apply_next_model)
+            else:
+                apply_next_model(model_class)
+
+    def do_pending_operations(self, model: type[Model]) -> None:
+        """
+        Take a newly-prepared model and pass it to each function waiting for
+        it. This is called at the very end of Models.register_model().
+        """
+        key = model.model_options.package_label, model.model_options.model_name
+        for function in self._pending_operations.pop(key, []):
+            function(model)
+
+
+models_registry = ModelsRegistry()
+
+
+# Decorator to register a model (using the internal registry for the correct state).
+def register_model[M: "Model"](model_class: type[M]) -> type[M]:
+    model_class._model_meta.models_registry.register_model(
+        model_class.model_options.package_label,
+        model_class,
+    )
+    return model_class