diff --git a/sqlmodel/_compat.py b/sqlmodel/_compat.py index 230f8cc362..e204b1778a 100644 --- a/sqlmodel/_compat.py +++ b/sqlmodel/_compat.py @@ -179,6 +179,18 @@ def get_relationship_to( elif origin is list: use_annotation = get_args(annotation)[0] + # If a dict or Mapping, get the value type (second type argument) + elif origin is dict or origin is Mapping: + args = get_args(annotation) + if len(args) >= 2: + # For Dict[K, V] or Mapping[K, V], we want the value type (V) + use_annotation = args[1] + else: + raise ValueError( + f"Dict/Mapping relationship field '{name}' must have both key " + "and value type arguments (e.g., Dict[str, Model])" + ) + return get_relationship_to( name=name, rel_info=rel_info, annotation=use_annotation ) diff --git a/tests/test_dict_relationship_recursion.py b/tests/test_dict_relationship_recursion.py new file mode 100644 index 0000000000..9993b6843d --- /dev/null +++ b/tests/test_dict_relationship_recursion.py @@ -0,0 +1,60 @@ +"""Test for Dict relationship recursion bug fix.""" + +from typing import Dict + +from sqlalchemy.orm.collections import attribute_mapped_collection +from sqlmodel import Field, Relationship, SQLModel + + +def test_dict_relationship_pattern(): + """Test that Dict relationships with attribute_mapped_collection work.""" + + # Create a minimal reproduction of the pattern + # This should not raise a RecursionError + + class TestChild(SQLModel, table=True): + __tablename__ = "test_child" + id: int = Field(primary_key=True) + key: str = Field(nullable=False) + parent_id: int = Field(foreign_key="test_parent.id") + parent: "TestParent" = Relationship(back_populates="children") + + class TestParent(SQLModel, table=True): + __tablename__ = "test_parent" + id: int = Field(primary_key=True) + children: Dict[str, "TestChild"] = Relationship( + back_populates="parent", + sa_relationship_kwargs={ + "collection_class": attribute_mapped_collection("key") + }, + ) + + # If we got here without RecursionError, the bug is fixed + assert TestParent.__tablename__ == "test_parent" + assert TestChild.__tablename__ == "test_child" + + +def test_dict_relationship_with_optional(): + """Test that Optional[Dict[...]] relationships also work.""" + from typing import Optional + + class Child(SQLModel, table=True): + __tablename__ = "child" + id: int = Field(primary_key=True) + key: str = Field(nullable=False) + parent_id: int = Field(foreign_key="parent.id") + parent: Optional["Parent"] = Relationship(back_populates="children") + + class Parent(SQLModel, table=True): + __tablename__ = "parent" + id: int = Field(primary_key=True) + children: Optional[Dict[str, "Child"]] = Relationship( + back_populates="parent", + sa_relationship_kwargs={ + "collection_class": attribute_mapped_collection("key") + }, + ) + + # If we got here without RecursionError, the bug is fixed + assert Parent.__tablename__ == "parent" + assert Child.__tablename__ == "child"