Add max depth parameter to be able load nested relationships on a model

src/sqlservice/model.py

@@ -63,19 +63,27 @@ def pk(self) -> t.Tuple[t.Any, ...]:
         return self.__mapper__.primary_key_from_instance(self)
 
     def to_dict(
-        self, *, exclude_relationships: bool = False, lazyload: bool = False
+        self,
+        *,
+        lazyload: bool = False,
+        exclude_relationships: bool = False,
+        include_nested_relationships: bool = False,
     ) -> t.Dict[str, t.Any]:
         """
         Serialize ORM loaded data to dictionary.
 
         Only the loaded data, i.e. data previously fetched from the database, will be serialized.
         Lazy-loaded columns and relationships will be excluded to avoid extra database queries.
 
-        By default, only table columns will be included. To include relationship fields, set
-        ``include_relationships=True``. This will nest ``to_dict()`` calls to the relationship
-        models.
+        By default, table columns and relationships will be included while nested relationships
+        will be excluded. To exclude relationships, set ``exclude_relationships=True``. To
+        include nested relationships, set ``include_nested_relationships=True``.
         """
-        serializer = ModelSerializer(exclude_relationships=exclude_relationships, lazyload=lazyload)
+        serializer = ModelSerializer(
+            exclude_relationships=exclude_relationships,
+            lazyload=lazyload,
+            include_nested_relationships=include_nested_relationships,
+        )
         return serializer.to_dict(self)
 
     def __iter__(self):
@@ -110,12 +118,19 @@ def delete(cls) -> Delete:
 
 
 class ModelSerializer:
-    def __init__(self, *, exclude_relationships: bool = False, lazyload: bool = False):
-        self.exclude_relationships = exclude_relationships
+    def __init__(
+        self,
+        *,
+        lazyload: bool = False,
+        exclude_relationships: bool = False,
+        include_nested_relationships: bool = False,
+    ):
         self.lazyload = lazyload
+        self.exclude_relationships = exclude_relationships
+        self.include_nested_relationships = include_nested_relationships
 
     def to_dict(self, model: ModelBase) -> t.Dict[str, t.Any]:
-        ctx: t.Dict[str, t.Any] = {"seen": set()}
+        ctx: t.Dict[str, t.Any] = {"seen": set(), "cache": {}, "depth": 0}
         return self.from_value(ctx, model)
 
     def from_value(self, ctx: dict, value: t.Any) -> t.Any:
@@ -129,16 +144,33 @@ def from_value(self, ctx: dict, value: t.Any) -> t.Any:
 
     def from_model(self, ctx: dict, value: t.Any) -> t.Dict[str, t.Any]:
         ctx.setdefault("seen", set())
+        ctx.setdefault("cache", {})
+        ctx.setdefault("depth", 0)
+
+        # Return the cached data if the model has already been seen
+        if value in ctx["seen"]:
+            # Return the cached data to break the cycle
+            return ctx["cache"][id(value)].copy()
+
+        # Add the model to the seen and path
         ctx["seen"].add(value)
 
+        data: t.Dict[str, t.Any] = {}
+        ctx["cache"][id(value)] = data
+
         state: orm.state.InstanceState = sa.inspect(value)
         mapper: orm.Mapper = sa.inspect(type(value))
 
+        current_depth = ctx["depth"]
+        include_relationships = not self.exclude_relationships
+
+        if current_depth > 0:
+            include_relationships = self.include_nested_relationships
+
         fields = mapper.columns.keys()
-        if not self.exclude_relationships:
+        if include_relationships:
             fields += mapper.relationships.keys()
 
-        data = {}
         for key in fields:
             loaded_value = state.attrs[key].loaded_value
 
@@ -152,8 +184,15 @@ def from_model(self, ctx: dict, value: t.Any) -> t.Dict[str, t.Any]:
             ):
                 continue
 
+            is_relationship = key in mapper.relationships.keys()
+            if is_relationship:
+                ctx["depth"] += 1
+
             data[key] = self.from_value(ctx, loaded_value)
 
+            if is_relationship:
+                ctx["depth"] -= 1
+
         return data
 
     def from_dict(self, ctx: dict, value: dict) -> dict:

tests/fixtures.py

@@ -36,6 +36,21 @@ class Model(ModelBase):
     pass
 
 
+class Book(Model):
+    __tablename__ = "books"
+
+    id: Mapped[int] = mapped_column(sa.Integer(), primary_key=True)
+    title: Mapped[str] = mapped_column(sa.String())
+    created_by: Mapped[int] = mapped_column(sa.Integer(), sa.ForeignKey("users.id"), nullable=False)
+    created_by_user: Mapped["User"] = relationship(
+        "User", foreign_keys=[created_by], back_populates="created_books"
+    )
+    updated_by: Mapped[int] = mapped_column(sa.Integer(), sa.ForeignKey("users.id"), nullable=False)
+    updated_by_user: Mapped["User"] = relationship(
+        "User", foreign_keys=[updated_by], back_populates="updated_books"
+    )
+
+
 class User(Model):
     __tablename__ = "users"
 
@@ -48,6 +63,12 @@ class User(Model):
         "GroupMembership", back_populates="user"
     )
     items: Mapped[t.List["GroupMembership"]] = relationship("Item")
+    created_books: Mapped[t.List["Book"]] = relationship(
+        "Book", foreign_keys=[Book.created_by], back_populates="created_by_user"
+    )
+    updated_books: Mapped[t.List["Book"]] = relationship(
+        "Book", foreign_keys=[Book.updated_by], back_populates="updated_by_user"
+    )
 
 
 class Address(Model):

tests/test_model.py

@@ -9,7 +9,7 @@
 
 from sqlservice import Database, ModelBase, ModelMeta, as_declarative, declarative_base
 
-from .fixtures import Address, Group, GroupMembership, Item, Note, User
+from .fixtures import Address, Book, Group, GroupMembership, Item, Note, User
 
 
 parametrize = pytest.mark.parametrize
@@ -217,6 +217,8 @@ def test_model_to_dict__with_non_persisted_model(model: ModelBase, args: dict, e
                 "addresses": [{"id": 1, "user_id": 1, "addr": "a", "zip_code": "1"}],
                 "group_memberships": [],
                 "items": [],
+                "created_books": [],
+                "updated_books": [],
             },
             id="1:M_all_loaded_include_relationships",
         ),
@@ -245,7 +247,7 @@ def test_model_to_dict__with_non_persisted_model(model: ModelBase, args: dict, e
                     GroupMembership(group=Group(name="g2")),
                 ],
             ),
-            {},
+            {"include_nested_relationships": True},
             [orm.joinedload("*")],
             {
                 "id": 1,
@@ -260,6 +262,8 @@ def test_model_to_dict__with_non_persisted_model(model: ModelBase, args: dict, e
                     {"group_id": 2, "user_id": 1, "group": {"id": 2, "name": "g2"}},
                 ],
                 "items": [],
+                "created_books": [],
+                "updated_books": [],
             },
             id="nested_relationships_all_loaded",
         ),
@@ -282,7 +286,7 @@ def test_model_to_dict__with_non_persisted_model(model: ModelBase, args: dict, e
                     )
                 ],
             ),
-            {"lazyload": True},
+            {"lazyload": True, "include_nested_relationships": True},
             [orm.lazyload("*")],
             {
                 "id": 1,
@@ -323,9 +327,103 @@ def test_model_to_dict__with_non_persisted_model(model: ModelBase, args: dict, e
                         },
                     }
                 ],
+                "created_books": [],
+                "updated_books": [],
             },
             id="nested_relationships_lazy_loaded",
         ),
+        param(
+            Book(
+                id=1,
+                title="t",
+                created_by_user=User(id=1, name="n"),
+                updated_by_user=User(id=2, name="n"),
+            ),
+            {"lazyload": True},
+            [],
+            {
+                "id": 1,
+                "title": "t",
+                "created_by": 1,
+                "updated_by": 2,
+                "created_by_user": {"id": 1, "name": "n", "active": True},
+                "updated_by_user": {"id": 2, "name": "n", "active": True},
+            },
+            id="exclude_nested_relationships_lazy_loaded_with_circular_reference",
+        ),
+        param(
+            Book(
+                id=1,
+                title="t",
+                created_by_user=User(id=1, name="n"),
+                updated_by_user=User(id=2, name="n"),
+            ),
+            {},
+            [],
+            {
+                "id": 1,
+                "title": "t",
+                "created_by": 1,
+                "updated_by": 2,
+            },
+            id="exclude_nested_relationships_with_circular_reference",
+        ),
+        param(
+            Book(
+                id=1,
+                title="t",
+                created_by_user=User(id=1, name="n"),
+                updated_by_user=User(id=2, name="n"),
+            ),
+            {"lazyload": True, "include_nested_relationships": True},
+            [],
+            {
+                "id": 1,
+                "title": "t",
+                "created_by": 1,
+                "updated_by": 2,
+                "created_by_user": {
+                    "id": 1,
+                    "name": "n",
+                    "active": True,
+                    "addresses": [],
+                    "group_memberships": [],
+                    "items": [],
+                    "created_books": [{"id": 1, "title": "t", "created_by": 1, "updated_by": 2}],
+                    "updated_books": [],
+                },
+                "updated_by_user": {
+                    "id": 2,
+                    "name": "n",
+                    "active": True,
+                    "addresses": [],
+                    "group_memberships": [],
+                    "items": [],
+                    "created_books": [],
+                    "updated_books": [
+                        {
+                            "id": 1,
+                            "title": "t",
+                            "created_by": 1,
+                            "updated_by": 2,
+                            "created_by_user": {
+                                "id": 1,
+                                "name": "n",
+                                "active": True,
+                                "addresses": [],
+                                "group_memberships": [],
+                                "items": [],
+                                "created_books": [
+                                    {"id": 1, "title": "t", "created_by": 1, "updated_by": 2}
+                                ],
+                                "updated_books": [],
+                            },
+                        }
+                    ],
+                },
+            },
+            id="include_nested_relationships_lazy_loaded_with_circular_reference",
+        ),
     ],
 )
 def test_model_to_dict__loaded_from_database(