diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 97ad3bd53..e97d1071b 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,6 +8,13 @@ Changelog 1.1 === +Unreleased +---------- + +Added +^^^^^ +- Support forward relation ``in`` and ``not_in`` filters with model instances. (#2323) + 1.1.9 ----- diff --git a/docs/query.rst b/docs/query.rst index 2179544c4..57c9c7e43 100644 --- a/docs/query.rst +++ b/docs/query.rst @@ -122,6 +122,15 @@ query around related entities: # Gets all tournaments where teams with "junior" in their name are participating await Tournament.filter(event__participants__name__icontains='junior').distinct() +Forward foreign key and one-to-one relations support ``in`` and ``not_in`` with +related model instances or referenced values. Model instances use the relation's +``to_field`` when it points to a unique field other than the primary key. + +.. code-block:: python3 + + await Event.filter(tournament__in=[first_tournament, second_tournament]) + await Event.filter(tournament__not_in=[first_tournament]) + Usually you not only want to filter by related data, but also get that related data as well. You could do it using ``.prefetch_related()``: diff --git a/examples/relations.py b/examples/relations.py index 7a7557519..636c1443e 100644 --- a/examples/relations.py +++ b/examples/relations.py @@ -6,8 +6,9 @@ to get this related objects """ -from tortoise import Tortoise, fields, run_async +from tortoise import fields from tortoise.exceptions import NoValuesFetched +from tortoise.expressions import Subquery from tortoise.models import Model @@ -17,7 +18,7 @@ class Tournament(Model): events: fields.ReverseRelation["Event"] - def __str__(self): + def __str__(self) -> str: return self.name @@ -32,7 +33,10 @@ class Event(Model): "models.Team", related_name="events", through="event_team" ) - def __str__(self): + tournament_id: int + address: fields.BackwardOneToOneRelation["Address"] + + def __str__(self) -> str: return self.name @@ -43,8 +47,9 @@ class Address(Model): event: fields.OneToOneRelation[Event] = fields.OneToOneField( Event, on_delete=fields.OnDelete.CASCADE, related_name="address", primary_key=True ) + event_id: int - def __str__(self): + def __str__(self) -> str: return f"Address({self.city}, {self.street})" @@ -54,14 +59,11 @@ class Team(Model): events: fields.ManyToManyRelation[Event] - def __str__(self): + def __str__(self) -> str: return self.name -async def run(): - await Tortoise.init(db_url="sqlite://:memory:", modules={"models": ["__main__"]}) - await Tortoise.generate_schemas() - +async def run() -> None: tournament = Tournament(name="New Tournament") await tournament.save() await Event(name="Without participants", tournament_id=tournament.id).save() @@ -95,14 +97,24 @@ async def run(): "participants", "tournament" ) ) - print(await participants[0].fetch_related("events")) - print(await Team.fetch_for_list(participants, "events")) + await participants[0].fetch_related("events") + print(list(participants[0].events)) + + await Team.fetch_for_list(participants, "events") + for team in participants: + print(team.name, list(team.events)) print(await Team.filter(events__tournament__id=tournament.id)) print(await Event.filter(tournament=tournament)) + print(await Event.filter(tournament__in=[tournament])) + + print(await Event.filter(address__city__in=Subquery(Address.all().values("city")))) + + print(await Tournament.filter(events__in=[event])) + print( await Tournament.filter(events__name__in=["Test", "Prod"]) .order_by("-events__participants__name") @@ -115,12 +127,16 @@ async def run(): print(await Address.filter(event=event).first()) - event_reload1 = await Event.filter(id=event.id).first() + event_reload1 = await Event.filter(id=event.id).get() print(await event_reload1.address) - event_reload2 = await Event.filter(id=event.id).prefetch_related("address").first() + event_reload2 = await Event.filter(id=event.id).prefetch_related("address").get() print(event_reload2.address) + print(dict(await event_reload2.address.event)) + if __name__ == "__main__": - run_async(run()) + from tortoise.contrib.test import run_in_memory_sqlite + + run_in_memory_sqlite(run) diff --git a/tests/fields/test_relation_membership.py b/tests/fields/test_relation_membership.py new file mode 100644 index 000000000..301f8ea25 --- /dev/null +++ b/tests/fields/test_relation_membership.py @@ -0,0 +1,141 @@ +from collections.abc import Iterable + +import pytest + +from tests import testmodels +from tortoise.exceptions import FieldError +from tortoise.expressions import F, Subquery + + +@pytest.mark.asyncio +@pytest.mark.parametrize("lookup", ["in", "not_in"]) +@pytest.mark.parametrize("input_kind", ["models", "iterator", "mixed"]) +async def test_foreign_key_membership(db, lookup: str, input_kind: str) -> None: + tournaments = [await testmodels.Tournament.create(name=str(i)) for i in range(3)] + relations = [await testmodels.MinRelation.create(tournament=t) for t in tournaments] + values: Iterable[testmodels.Tournament | int] = tournaments[:2] + if input_kind == "iterator": + values = iter(values) + elif input_kind == "mixed": + values = [tournaments[0], tournaments[1].pk] + + result = await testmodels.MinRelation.filter(**{f"tournament__{lookup}": values}).order_by("id") + assert result == (relations[:2] if lookup == "in" else relations[2:]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("lookup", ["in", "not_in"]) +@pytest.mark.parametrize("relation_model", [testmodels.Student, testmodels.Principal]) +async def test_membership_uses_referenced_field( + db, lookup: str, relation_model: type[testmodels.Student] | type[testmodels.Principal] +) -> None: + first = await testmodels.School.create(name="first", id=101) + second = await testmodels.School.create(name="second", id=202) + first_relation = await relation_model.create(name="first", school=first) + second_relation = await relation_model.create(name="second", school=second) + + result = await relation_model.filter(**{f"school__{lookup}": [first]}) + assert result == [first_relation if lookup == "in" else second_relation] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("lookup", ["in", "not_in"]) +async def test_membership_with_string_primary_keys(db, lookup: str) -> None: + first = await testmodels.CharPkModel.create(id="first") + second = await testmodels.CharPkModel.create(id="second") + first_relation = await testmodels.CharFkRelatedModel.create(model=first) + second_relation = await testmodels.CharFkRelatedModel.create(model=second) + + result = await testmodels.CharFkRelatedModel.filter(**{f"model__{lookup}": [first]}) + assert result == [first_relation if lookup == "in" else second_relation] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("lookup", ["in", "not_in"]) +async def test_membership_with_renamed_uuid_columns(db, lookup: str) -> None: + first = await testmodels.UUIDPkSourceModel.create() + second = await testmodels.UUIDPkSourceModel.create() + first_relation = await testmodels.UUIDFkRelatedSourceModel.create(model=first) + second_relation = await testmodels.UUIDFkRelatedSourceModel.create(model=second) + + result = await testmodels.UUIDFkRelatedSourceModel.filter(**{f"model__{lookup}": [first]}) + assert result == [first_relation if lookup == "in" else second_relation] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("lookup", ["in", "not_in"]) +async def test_membership_with_empty_values(db, lookup: str) -> None: + tournament = await testmodels.Tournament.create(name="first") + relation = await testmodels.MinRelation.create(tournament=tournament) + + result = await testmodels.MinRelation.filter(**{f"tournament__{lookup}": []}) + assert result == ([] if lookup == "in" else [relation]) + + +@pytest.mark.asyncio +async def test_membership_with_subquery(db) -> None: + first = await testmodels.Tournament.create(name="first") + second = await testmodels.Tournament.create(name="second") + relation = await testmodels.MinRelation.create(tournament=first) + await testmodels.MinRelation.create(tournament=second) + selected = Subquery(testmodels.Tournament.filter(name="first").values("id")) + + assert await testmodels.MinRelation.filter(tournament__in=selected) == [relation] + + +@pytest.mark.asyncio +async def test_existing_foreign_key_filters(db) -> None: + first = await testmodels.Tournament.create(name="first") + second = await testmodels.Tournament.create(name="second") + relation = await testmodels.MinRelation.create(tournament=first) + await testmodels.MinRelation.create(tournament=second) + + assert await testmodels.MinRelation.filter(tournament=first) == [relation] + assert await testmodels.MinRelation.filter(tournament_id__in=[first.pk]) == [relation] + assert await testmodels.MinRelation.filter(tournament__name__in=["first"]) == [relation] + + +@pytest.mark.asyncio +async def test_existing_annotation_filter_name(db) -> None: + tournament = await testmodels.Tournament.create(name="first") + relation = await testmodels.MinRelation.create(tournament=tournament) + await testmodels.MinRelation.create(tournament=tournament) + + result = await ( + testmodels.MinRelation.annotate(tournament__in=F("id")) + .filter(tournament__in=relation.pk) + .values("id") + ) + assert result == [{"id": relation.pk}] + + +@pytest.mark.asyncio +async def test_trailing_lookup_separator_still_raises(db) -> None: + tournament = await testmodels.Tournament.create(name="first") + await testmodels.MinRelation.create(tournament=tournament) + + with pytest.raises(FieldError, match="Unknown filter param"): + await testmodels.MinRelation.filter(**{"tournament__": tournament}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("lookup", ["in", "not_in"]) +async def test_nullable_relation_membership(db, lookup: str) -> None: + parent = await testmodels.UUIDPkSourceModel.create() + linked = await testmodels.UUIDFkRelatedNullSourceModel.create(model=parent) + unlinked = await testmodels.UUIDFkRelatedNullSourceModel.create(model=None) + + result = await testmodels.UUIDFkRelatedNullSourceModel.filter(**{f"model__{lookup}": [parent]}) + assert result == ([linked] if lookup == "in" else [unlinked]) + + +@pytest.mark.asyncio +async def test_nested_forward_relation_membership(db) -> None: + first = await testmodels.School.create(name="first", id=101) + second = await testmodels.School.create(name="second", id=202) + selected = await testmodels.Student.create(name="first", school=first) + await testmodels.Student.create(name="second", school=second) + await testmodels.Principal.create(name="first", school=first) + await testmodels.Principal.create(name="second", school=second) + + assert await testmodels.Student.filter(school__principal__school__in=[first]) == [selected] diff --git a/tortoise/__init__.py b/tortoise/__init__.py index 848842271..13d3d4130 100644 --- a/tortoise/__init__.py +++ b/tortoise/__init__.py @@ -8,8 +8,6 @@ from types import ModuleType from typing import TYPE_CHECKING, Any -from anyio import from_thread - from tortoise.apps import Apps from tortoise.backends.base.client import BaseDBAsyncClient from tortoise.backends.base.config_generator import expand_db_url @@ -536,11 +534,17 @@ def _init_timezone(cls, use_tz: bool, timezone: str) -> None: _reset_timezone_cache() -def run_async(coro: Coroutine) -> None: +def run_async(coro: Coroutine, *, verbose: bool = False) -> None: """ Simple async runner that cleans up DB connections on exit. This is meant for simple scripts. + :param coro: The coroutine to run. + :param verbose: If ``True``, set the ``tortoise`` logger to ``DEBUG`` + so that generated SQL and runtime details are printed. If no + logging handler is configured, a basic one is added to the root + logger. Other loggers are left untouched. + Usage:: from tortoise import Tortoise, run_async @@ -555,8 +559,19 @@ async def do_stuff(): run_async(do_stuff()) """ + # Imports `anyio` inside function to improve the tortoise module startup performance + # (it's only used by this convenience helper). + from anyio import from_thread + from tortoise.context import get_current_context + if verbose: + logging.basicConfig( + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + logger.setLevel(logging.DEBUG) + async def main() -> None: try: await coro diff --git a/tortoise/contrib/test/__init__.py b/tortoise/contrib/test/__init__.py index b55864a45..1dce563a8 100644 --- a/tortoise/contrib/test/__init__.py +++ b/tortoise/contrib/test/__init__.py @@ -25,13 +25,13 @@ async def test_sqlite_only(db): from __future__ import annotations import inspect -import typing +import sys from collections.abc import Callable, Coroutine from functools import partial, wraps -from typing import TYPE_CHECKING, ParamSpec, TypeVar, cast +from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, cast, overload from unittest import SkipTest, expectedFailure, skip, skipIf, skipUnless -from tortoise import Tortoise +from tortoise import Tortoise, run_async from tortoise.connection import get_connection from tortoise.context import TortoiseContext, tortoise_test_context @@ -157,12 +157,10 @@ def _topological_sort_models(models: list[type[Model]]) -> list[type[Model]]: return sorted_models -_FT = TypeVar("_FT", bound=Callable[..., typing.Any]) +_FT = TypeVar("_FT", bound=Callable[..., Any]) -def requireCapability( - connection_name: str = "models", **conditions: typing.Any -) -> Callable[[_FT], _FT]: +def requireCapability(connection_name: str = "models", **conditions: Any) -> Callable[[_FT], _FT]: """ Skip a test if the required capabilities are not matched. @@ -204,14 +202,14 @@ def check_capabilities() -> None: if inspect.iscoroutinefunction(test_item): @wraps(test_item) - async def skip_wrapper(*args: typing.Any, **kwargs: typing.Any) -> typing.Any: + async def skip_wrapper(*args: Any, **kwargs: Any) -> Any: check_capabilities() return await test_item(*args, **kwargs) else: @wraps(test_item) - def skip_wrapper(*args: typing.Any, **kwargs: typing.Any) -> typing.Any: + def skip_wrapper(*args: Any, **kwargs: Any) -> Any: check_capabilities() return test_item(*args, **kwargs) @@ -235,11 +233,11 @@ def skip_wrapper(*args: typing.Any, **kwargs: typing.Any) -> typing.Any: return decorator -@typing.overload +@overload def init_memory_sqlite(models: ModulesConfigType | None = None) -> AsyncFuncDeco: ... -@typing.overload +@overload def init_memory_sqlite(models: AsyncFunc) -> AsyncFunc: ... @@ -300,3 +298,33 @@ async def runner(*args, **kwargs) -> T: else: models = cast(list, models) return partial(wrapper, ms=models) + + +def run_in_memory_sqlite( + func: AsyncFunc, + *, + models: ModulesConfigType | None = None, + verbose: bool | None = None, +) -> None: + """ + Run an async function using an in-memory SQLite database. + + Equivalent to decorating ``func`` with :func:`init_memory_sqlite` and + passing the result to :func:`~tortoise.run_async`:: + + async def run(): + ... + + if __name__ == "__main__": + run_in_memory_sqlite(run) + + :param func: The async function to run. + :param models: Modules to load models from. Defaults to ``["__main__"]``. + :param verbose: If ``None`` (default), enable debug logging when + ``--verbose`` is present in ``sys.argv``. Pass ``True`` or ``False`` + to force it. See :func:`run_async` for details. + """ + if verbose is None: + verbose = "--verbose" in sys.argv + wrapped = init_memory_sqlite(models)(func) + run_async(wrapped(), verbose=verbose) diff --git a/tortoise/expressions.py b/tortoise/expressions.py index d47ce1f4b..d81cb6085 100644 --- a/tortoise/expressions.py +++ b/tortoise/expressions.py @@ -435,23 +435,51 @@ def _resolve_regular_kwarg( return modifier def _get_actual_filter_params( - self, resolve_context: ResolveContext, key: str, value: Table | FilterInfoDict + self, resolve_context: ResolveContext, key: str, value: Any ) -> tuple[str, Any]: filter_key = key + filter_value: Any model_meta = resolve_context.model._meta - if key in model_meta.fk_fields or key in model_meta.o2o_fields: - field_object = cast(RelationalField, model_meta.fields_map[key]) + relation_name, separator, lookup = key.partition("__") + if ( + (relation_name in model_meta.fk_fields or relation_name in model_meta.o2o_fields) + and (not separator or lookup in ("in", "not_in")) + and ( + # Keep bare relation conversion (e.g. `filter(tournament=obj)`), but + # let exact field or annotation names win over relation lookups. e.g. + # + # .annotate(tournament__in=F("id")).filter(tournament__in=relation.pk) + # + # Here `tournament__in` refers to the annotation, not the FK lookup. + not lookup + or (key not in model_meta.filters and key not in resolve_context.custom_filters) + ) + ): + field_object = cast(RelationalField, model_meta.fields_map[relation_name]) filter_key = cast(str, field_object.source_field) to_field_name = field_object.to_field_instance.model_field_name - filter_value = ( - getattr(value, to_field_name) - if isinstance(value, field_object.related_model) - else value - ) + if lookup: + filter_key = f"{filter_key}__{lookup}" + filter_value = ( + value + if isinstance(value, (Expression, Term)) + else [ + getattr(item, to_field_name) + if isinstance(item, field_object.related_model) + else item + for item in value + ] + ) + else: + filter_value = ( + getattr(value, to_field_name) + if isinstance(value, field_object.related_model) + else value + ) elif key in model_meta.m2m_fields: filter_value = getattr(value, "pk", value) elif ( - key.split("__")[0] in model_meta.fetch_fields + relation_name in model_meta.fetch_fields or key in resolve_context.custom_filters or key in model_meta.filters ): diff --git a/tortoise/models.py b/tortoise/models.py index 19721a232..06fea29f3 100644 --- a/tortoise/models.py +++ b/tortoise/models.py @@ -4,7 +4,7 @@ import inspect import re import sys -from collections.abc import Awaitable, Callable, Generator, Iterable +from collections.abc import Awaitable, Callable, Generator, Iterable, Iterator from copy import copy, deepcopy from functools import partial from typing import TYPE_CHECKING, Any, TypedDict, TypeVar, cast @@ -876,7 +876,7 @@ def __hash__(self) -> int: raise TypeError("Model instances without id are unhashable") return hash(self.pk) - def __iter__(self) -> Iterable[tuple[str, Any]]: + def __iter__(self) -> Iterator[tuple[str, Any]]: for field in self._meta.db_fields: yield field, getattr(self, field)