diff --git a/src/sqlservice/model.py b/src/sqlservice/model.py index 85e99c1..af11a52 100644 --- a/src/sqlservice/model.py +++ b/src/sqlservice/model.py @@ -63,7 +63,11 @@ 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. @@ -71,11 +75,15 @@ def to_dict( 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: diff --git a/tests/fixtures.py b/tests/fixtures.py index b1de901..4cde68e 100644 --- a/tests/fixtures.py +++ b/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): diff --git a/tests/test_model.py b/tests/test_model.py index 3ccd15b..8cd26d6 100644 --- a/tests/test_model.py +++ b/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(