Support model instances in forward relation membership filters

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
 -----
 

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()``:

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)

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]

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