Skip to content

Commit

Permalink
fix: nested lookups (#455)
Browse files Browse the repository at this point in the history 
* fix: nested lookups

* fix: mongodb5

* version: 1.16.4
  • Loading branch information
@roman-right
roman-right authored Dec 20, 2022
1 parent 416fe6e commit f75e26a
Show file tree
Hide file tree
Showing 10 changed files with 187 additions and 54 deletions.
2 changes: 1 addition & 1 deletion beanie/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
from beanie.odm.views import View
from beanie.odm.union_doc import UnionDoc

__version__ = "1.16.3"
__version__ = "1.16.4"
__all__ = [
# ODM
"Document",
Expand Down
1 change: 1 addition & 0 deletions beanie/odm/queries/cursor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@ async def to_list(
if cursor is None:
raise RuntimeError("self.motor_cursor was not set")
motor_list: List[Dict[str, Any]] = self._get_cache()

if motor_list is None:
motor_list = await cursor.to_list(length)
self._set_cache(motor_list)
Expand Down
136 changes: 92 additions & 44 deletions beanie/odm/utils/find.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,9 @@
if TYPE_CHECKING:
from beanie import Document

# TODO: check if this is the most efficient way for
# appending subqueries to the queries var


def construct_lookup_queries(cls: Type["Document"]) -> List[Dict[str, Any]]:
if cls.get_model_type() == ModelType.UnionDoc:
Expand All @@ -27,64 +30,112 @@ def construct_query(
link_info: LinkInfo,
queries: List,
database_major_version: int,
parent_prefix: str = "",
):
field_path = ".".join(filter(None, (parent_prefix, link_info.field)))

if link_info.link_type in [
LinkTypes.DIRECT,
LinkTypes.OPTIONAL_DIRECT,
]:
queries += [
{
"$lookup": {
"from": link_info.model_class.get_motor_collection().name, # type: ignore
"localField": f"{field_path}.$id",
"foreignField": "_id",
"as": f"_link_{link_info.field}",
}
},
{
"$unwind": {
"path": f"$_link_{link_info.field}",
"preserveNullAndEmptyArrays": True,
}
},
{
"$set": {
field_path: {
"$cond": {
"if": {
"$ifNull": [
f"$_link_{link_info.field}",
False,
]
},
"then": f"$_link_{link_info.field}",
"else": f"${link_info.field}",
if database_major_version >= 5 or link_info.nested_links is None:
lookup_steps = [
{
"$lookup": {
"from": link_info.model_class.get_motor_collection().name,
# type: ignore
"localField": f"{link_info.field}.$id",
"foreignField": "_id",
"as": f"_link_{link_info.field}",
}
},
{
"$unwind": {
"path": f"$_link_{link_info.field}",
"preserveNullAndEmptyArrays": True,
}
},
{
"$set": {
link_info.field: {
"$cond": {
"if": {
"$ifNull": [
f"$_link_{link_info.field}",
False,
]
},
"then": f"$_link_{link_info.field}",
"else": f"${link_info.field}",
}
}
}
}
},
] # type: ignore
},
] # type: ignore
if link_info.nested_links is not None:
lookup_steps[0]["$lookup"]["pipeline"] = []
for nested_link in link_info.nested_links:
construct_query(
link_info=link_info.nested_links[nested_link],
queries=lookup_steps[0]["$lookup"]["pipeline"],
database_major_version=database_major_version,
)
queries += lookup_steps

if link_info.nested_links is not None:
else:
lookup_steps = [
{
"$lookup": {
"from": link_info.model_class.get_motor_collection().name,
"let": {"link_id": f"${link_info.field}.$id"},
"as": f"_link_{link_info.field}",
"pipeline": [
{
"$match": {
"$expr": {"$eq": ["$_id", "$$link_id"]}
}
},
],
}
},
{
"$unwind": {
"path": f"$_link_{link_info.field}",
"preserveNullAndEmptyArrays": True,
}
},
{
"$set": {
link_info.field: {
"$cond": {
"if": {
"$ifNull": [
f"$_link_{link_info.field}",
False,
]
},
"then": f"$_link_{link_info.field}",
"else": f"${link_info.field}",
}
}
}
},
]
for nested_link in link_info.nested_links:
construct_query(
link_info=link_info.nested_links[nested_link],
queries=queries,
queries=lookup_steps[0]["$lookup"]["pipeline"],
database_major_version=database_major_version,
parent_prefix=field_path,
)
queries += lookup_steps

else:
if database_major_version >= 5 or link_info.nested_links is None:
queries.append(
{
"$lookup": {
"from": link_info.model_class.get_motor_collection().name, # type: ignore
"localField": f"{field_path}.$id",
"from": link_info.model_class.get_motor_collection().name,
# type: ignore
"localField": f"{link_info.field}.$id",
"foreignField": "_id",
"as": field_path,
"as": link_info.field,
}
}
)
Expand All @@ -96,14 +147,13 @@ def construct_query(
link_info=link_info.nested_links[nested_link],
queries=queries[-1]["$lookup"]["pipeline"],
database_major_version=database_major_version,
parent_prefix="",
)
else:
lookup_step = {
"$lookup": {
"from": link_info.model_class.get_motor_collection().name,
"let": {"link_id": f"${field_path}.$id"},
"as": field_path,
"let": {"link_id": f"${link_info.field}.$id"},
"as": link_info.field,
"pipeline": [
{"$match": {"$expr": {"$in": ["$_id", "$$link_id"]}}},
],
Expand All @@ -115,8 +165,6 @@ def construct_query(
link_info=link_info.nested_links[nested_link],
queries=lookup_step["$lookup"]["pipeline"],
database_major_version=database_major_version,
parent_prefix="",
)
queries.append(lookup_step)

return queries
16 changes: 13 additions & 3 deletions beanie/odm/utils/init.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import importlib
import inspect
from copy import copy
from typing import Optional, List, Type, Union

from motor.motor_asyncio import AsyncIOMotorDatabase, AsyncIOMotorClient
Expand Down Expand Up @@ -168,8 +169,13 @@ def init_document_fields(cls) -> None:
Init class fields
:return: None
"""
cls.update_forward_refs()

def check_nested_links(link_info: LinkInfo):
def check_nested_links(
link_info: LinkInfo, prev_models: List[Type[BaseModel]]
):
if link_info.model_class in prev_models:
return
for k, v in link_info.model_class.__fields__.items():
nested_link_info = detect_link(v)
if nested_link_info is None:
Expand All @@ -178,7 +184,11 @@ def check_nested_links(link_info: LinkInfo):
if link_info.nested_links is None:
link_info.nested_links = {}
link_info.nested_links[v.name] = nested_link_info
check_nested_links(nested_link_info)
new_prev_models = copy(prev_models)
new_prev_models.append(link_info.model_class)
check_nested_links(
nested_link_info, prev_models=new_prev_models
)

if cls._link_fields is None:
cls._link_fields = {}
Expand All @@ -189,7 +199,7 @@ def check_nested_links(link_info: LinkInfo):
link_info = detect_link(v)
if link_info is not None:
cls._link_fields[v.name] = link_info
check_nested_links(link_info)
check_nested_links(link_info, prev_models=[])

cls._hidden_fields = cls.get_hidden_fields()

Expand Down
15 changes: 14 additions & 1 deletion docs/changelog.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,17 @@

Beanie project

## [1.16.4] - 2022-12-20

### Fix

- [[BUG] Initiating self-referencing documents with nested links breaks due to uncaught request loop](https://github.com/roman-right/beanie/issues/449)
- Nested lookups for direct links

### Implementation

- PR <https://github.com/roman-right/beanie/pull/455>

## [1.16.3] - 2022-12-19

### Fix
Expand Down Expand Up @@ -1167,4 +1178,6 @@ how specific type should be presented in the database

[1.16.2]: https://pypi.org/project/beanie/1.16.2

[1.16.3]: https://pypi.org/project/beanie/1.16.3
[1.16.3]: https://pypi.org/project/beanie/1.16.3

[1.16.4]: https://pypi.org/project/beanie/1.16.4
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[tool.poetry]
name = "beanie"
version = "1.16.3"
version = "1.16.4"
description = "Asynchronous Python ODM for MongoDB"
authors = ["Roman <[email protected]>"]
license = "Apache-2.0"
Expand Down
6 changes: 6 additions & 0 deletions tests/odm/conftest.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,9 @@
StateAndDecimalFieldModel,
Region,
UsersAddresses,
SelfLinked,
LoopedLinksA,
LoopedLinksB,
)
from tests.odm.views import TestView

Expand Down Expand Up @@ -208,6 +211,9 @@ async def init(loop, db):
StateAndDecimalFieldModel,
Region,
UsersAddresses,
SelfLinked,
LoopedLinksA,
LoopedLinksB,
]
await init_beanie(
database=db,
Expand Down
13 changes: 13 additions & 0 deletions tests/odm/models.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -640,3 +640,16 @@ class Settings:
"city": "$region_id.city",
"state": "$region_id.state",
}


class SelfLinked(Document):
item: Optional[Link["SelfLinked"]]
s: str


class LoopedLinksA(Document):
b: "LoopedLinksB"


class LoopedLinksB(Document):
a: Optional[LoopedLinksA]
Loading

0 comments on commit f75e26a

Please sign in to comment.