feat: enhance relationship handling for Pydantic to SQLModel conversion with support for dicts

Author: 50Bytes-dev

sqlmodel/_compat.py

@@ -343,7 +343,11 @@ def sqlmodel_validate(
         # Get and set any relationship objects
         if is_table_model_class(cls):
             for key in new_obj.__sqlmodel_relationships__:
-                value = getattr(use_obj, key, Undefined)
+                # Handle both dict and object access
+                if isinstance(use_obj, dict):
+                    value = use_obj.get(key, Undefined)
+                else:
+                    value = getattr(use_obj, key, Undefined)
                 if value is not Undefined:
                     setattr(new_obj, key, value)
         return new_obj

sqlmodel/main.py

@@ -924,7 +924,9 @@ def __setattr__(self, name: str, value: Any) -> None:
                 and name in self.__sqlmodel_relationships__
                 and value is not None
             ):
-                value = _convert_pydantic_to_table_model(value, name, self.__class__)
+                value = _convert_pydantic_to_table_model(
+                    value, name, self.__class__, self
+                )
 
             # Set in SQLAlchemy, before Pydantic to trigger events and updates
             if is_table_model_class(self.__class__) and is_instrumented(self, name):  # type: ignore[no-untyped-call]
@@ -1127,7 +1129,10 @@ def sqlmodel_update(
 
 
 def _convert_pydantic_to_table_model(
-    value: Any, relationship_name: str, owner_class: Type["SQLModel"]
+    value: Any,
+    relationship_name: str,
+    owner_class: Type["SQLModel"],
+    instance: Optional["SQLModel"] = None,
 ) -> Any:
     """
     Convert Pydantic objects to table models for relationship assignments.
@@ -1136,6 +1141,7 @@ def _convert_pydantic_to_table_model(
         value: The value being assigned to the relationship
         relationship_name: Name of the relationship attribute
         owner_class: The class that owns the relationship
+        instance: The SQLModel instance (for session context)
 
     Returns:
         Converted value(s) - table model instances instead of Pydantic objects
@@ -1185,7 +1191,7 @@ def _convert_pydantic_to_table_model(
             converted_items = []
             for item in value:
                 converted_item = _convert_single_pydantic_to_table_model(
-                    item, target_type
+                    item, target_type, instance
                 )
                 converted_items.append(converted_item)
             return converted_items
@@ -1198,21 +1204,24 @@ def _convert_pydantic_to_table_model(
                 resolved_type = default_registry._class_registry.get(target_type)
                 if resolved_type is not None:
                     target_type = resolved_type
-            except:
+            except Exception:
                 pass
 
-        return _convert_single_pydantic_to_table_model(value, target_type)
+        return _convert_single_pydantic_to_table_model(value, target_type, instance)
 
     return value
 
 
-def _convert_single_pydantic_to_table_model(item: Any, target_type: Any) -> Any:
+def _convert_single_pydantic_to_table_model(
+    item: Any, target_type: Any, instance: Optional["SQLModel"] = None
+) -> Any:
     """
     Convert a single Pydantic object to a table model.
 
     Args:
         item: The Pydantic object to convert
         target_type: The target table model type
+        instance: The SQLModel instance (for session context)
 
     Returns:
         Converted table model instance or original item if no conversion needed
@@ -1226,7 +1235,9 @@ def _convert_single_pydantic_to_table_model(item: Any, target_type: Any) -> Any:
         try:
             # Attempt to resolve forward reference from the default registry
             # This was part of the original logic and should be kept
-            resolved_type_from_registry = default_registry._class_registry.get(target_type)
+            resolved_type_from_registry = default_registry._class_registry.get(
+                target_type
+            )
             if resolved_type_from_registry is not None:
                 resolved_target_type = resolved_type_from_registry
         except Exception:
@@ -1235,9 +1246,14 @@ def _convert_single_pydantic_to_table_model(item: Any, target_type: Any) -> Any:
             # `_convert_pydantic_to_table_model` should provide a resolved type.
             # For safety, if it's still a string here, and item is a simple Pydantic model,
             # it's best to return item to avoid errors if no concrete type is found.
-            if isinstance(resolved_target_type, str) and isinstance(item, BaseModel) and hasattr(item, "__class__") and not is_table_model_class(item.__class__):
-                 return item # Fallback if no concrete type can be determined
-            pass # Continue if resolved_target_type is now a class or item is not a simple Pydantic model
+            if (
+                isinstance(resolved_target_type, str)
+                and isinstance(item, BaseModel)
+                and hasattr(item, "__class__")
+                and not is_table_model_class(item.__class__)
+            ):
+                return item  # Fallback if no concrete type can be determined
+            pass  # Continue if resolved_target_type is now a class or item is not a simple Pydantic model
 
     # If resolved_target_type is still a string and not a class, we cannot proceed with conversion.
     # This can happen if the forward reference cannot be resolved.
@@ -1253,7 +1269,8 @@ def _convert_single_pydantic_to_table_model(item: Any, target_type: Any) -> Any:
     if not (
         hasattr(resolved_target_type, "__mro__")
         and any(
-            hasattr(cls, "__sqlmodel_relationships__") for cls in resolved_target_type.__mro__
+            hasattr(cls, "__sqlmodel_relationships__")
+            for cls in resolved_target_type.__mro__
         )
     ):
         return item
@@ -1278,10 +1295,49 @@ def _convert_single_pydantic_to_table_model(item: Any, target_type: Any) -> Any:
                 # Pydantic v1
                 data = item.dict()
 
+            # If instance is available and item has an ID, try to find existing record
+            if instance is not None and "id" in data and data["id"] is not None:
+                from sqlalchemy.orm import object_session
+
+                session = object_session(instance)
+                if session is not None:
+                    # Try to find existing record by ID
+                    existing_record = session.get(resolved_target_type, data["id"])
+                    if existing_record is not None:
+                        # Update existing record with new data
+                        for key, value in data.items():
+                            if key != "id" and hasattr(existing_record, key):
+                                setattr(existing_record, key, value)
+                        return existing_record
+
             # Create new table model instance using resolved_target_type
             return resolved_target_type(**data)
         except Exception:
             # If conversion fails, return original item
             return item
 
+    # Check if item is a dictionary that should be converted to table model
+    elif isinstance(item, dict):
+        try:
+            # If instance is available and item has an ID, try to find existing record
+            if instance is not None and "id" in item and item["id"] is not None:
+                from sqlalchemy.orm import object_session
+
+                session = object_session(instance)
+                if session is not None:
+                    # Try to find existing record by ID
+                    existing_record = session.get(resolved_target_type, item["id"])
+                    if existing_record is not None:
+                        # Update existing record with new data
+                        for key, value in item.items():
+                            if key != "id" and hasattr(existing_record, key):
+                                setattr(existing_record, key, value)
+                        return existing_record
+
+            # Create new table model instance from dictionary
+            return resolved_target_type(**item)
+        except Exception:
+            # If conversion fails, return original item
+            return item
+
     return item

tests/test_relationships_set.py

@@ -1,7 +1,7 @@
 from sqlmodel import Field, Relationship, Session, SQLModel, create_engine
 
 
-def test_relationships_set():
+def test_relationships_set_pydantic():
     class Book(SQLModel, table=True):
         id: int = Field(default=None, primary_key=True)
         title: str
@@ -48,9 +48,50 @@ class IAuthorCreate(SQLModel):
         assert author.books[1].id is not None
         assert author.books[2].id is not None
 
-    author.books[0].title = "Updated Book One"
+
+def test_relationships_set_dict():
+    class Book(SQLModel, table=True):
+        id: int = Field(default=None, primary_key=True)
+        title: str
+        author_id: int = Field(foreign_key="author.id")
+        author: "Author" = Relationship(back_populates="books")
+
+    class IBookCreate(SQLModel):
+        title: str
+
+    class Author(SQLModel, table=True):
+        id: int = Field(default=None, primary_key=True)
+        name: str
+        books: list[Book] = Relationship(back_populates="author")
+
+    class IAuthorCreate(SQLModel):
+        name: str
+        books: list[IBookCreate] = []
+
+    book1 = IBookCreate(title="Book One")
+    book2 = IBookCreate(title="Book Two")
+    book3 = IBookCreate(title="Book Three")
+
+    author_data = IAuthorCreate(name="Author Name", books=[book1, book2, book3])
+
+    author = Author.model_validate(author_data.model_dump(exclude={"id"}))
+
+    engine = create_engine("sqlite://")
+
+    SQLModel.metadata.create_all(engine)
+
     with Session(engine) as session:
         session.add(author)
         session.commit()
         session.refresh(author)
-        assert author.books[0].title == "Updated Book One"
+        assert author.id is not None
+        assert len(author.books) == 3
+        assert author.books[0].title == "Book One"
+        assert author.books[1].title == "Book Two"
+        assert author.books[2].title == "Book Three"
+        assert author.books[0].author_id == author.id
+        assert author.books[1].author_id == author.id
+        assert author.books[2].author_id == author.id
+        assert author.books[0].id is not None
+        assert author.books[1].id is not None
+        assert author.books[2].id is not None

tests/test_relationships_update.py

@@ -13,7 +13,7 @@
 import pytest
 
 
-def test_relationships_update():
+def test_relationships_update_pydantic():
     """Test conversion of single Pydantic model to SQLModel with forward reference."""
 
     class IBookUpdate(BaseModel):
@@ -56,7 +56,7 @@ class Book(SQLModel, table=True):
 
         # Prepare the update data Pydantic model
         author_update_dto = IAuthorUpdate(
-            id=author_id, # This ID in DTO is informational
+            id=author_id,  # This ID in DTO is informational
             name="Updated Author",
             books=[IBookUpdate(id=book_id, title="Updated Book")],
         )
@@ -73,7 +73,7 @@ class Book(SQLModel, table=True):
                 book_to_update = session.get(Book, book_update_data.id)
 
                 if book_to_update:
-                    if book_update_data.title is not None: # Check if title is provided
+                    if book_update_data.title is not None:  # Check if title is provided
                         book_to_update.title = book_update_data.title
                     processed_books_list.append(book_to_update)
                 # else:
@@ -82,9 +82,76 @@ class Book(SQLModel, table=True):
             # Assign the list of (potentially updated) persistent Book SQLModel objects
             db_author.books = processed_books_list
 
-        session.add(db_author) # Add the updated instance to the session (marks it as dirty)
+        session.add(
+            db_author
+        )  # Add the updated instance to the session (marks it as dirty)
         session.commit()
-        session.refresh(db_author) # Refresh to get the latest state from DB
+        session.refresh(db_author)  # Refresh to get the latest state from DB
+
+        # Assertions on the original IDs and updated content
+        assert db_author.id == author_id
+        assert db_author.name == "Updated Author"
+        assert len(db_author.books) == 1
+        assert db_author.books[0].id == book_id
+        assert db_author.books[0].title == "Updated Book"
+
+
+def test_relationships_update_dict():
+    """Test conversion of single Pydantic model to SQLModel with forward reference."""
+
+    class IBookUpdate(BaseModel):
+        id: int
+        title: str | None = None
+
+    class IAuthorUpdate(BaseModel):
+        id: int
+        name: str | None = None
+        books: list[IBookUpdate] | None = None
+
+    class Author(SQLModel, table=True):
+        id: Optional[int] = Field(default=None, primary_key=True)
+        name: str
+        books: List["Book"] = Relationship(back_populates="author")
+
+    class Book(SQLModel, table=True):
+        id: Optional[int] = Field(default=None, primary_key=True)
+        title: str
+        author_id: Optional[int] = Field(default=None, foreign_key="author.id")
+        author: Optional["Author"] = Relationship(back_populates="books")
+
+    engine = create_engine("sqlite://", echo=False)
+    SQLModel.metadata.create_all(engine)
+
+    with Session(engine) as session:
+        book = Book(title="Test Book")
+        author = Author(name="Test Author", books=[book])
+        session.add(author)
+        session.commit()
+        session.refresh(author)
+
+        author_id = author.id
+        book_id = book.id
+
+    with Session(engine) as session:
+        # Fetch the existing author
+        db_author = session.get(Author, author_id)
+        assert db_author is not None, "Author to update was not found in the database."
+
+        # Prepare the update data Pydantic model
+        author_update_dto = IAuthorUpdate(
+            id=author_id,  # This ID in DTO is informational
+            name="Updated Author",
+            books=[IBookUpdate(id=book_id, title="Updated Book")],
+        )
+
+        update_data = author_update_dto.model_dump()
+
+        for field in update_data:
+            setattr(db_author, field, update_data[field])
+
+        session.add(db_author)
+        session.commit()
+        session.refresh(db_author)
 
         # Assertions on the original IDs and updated content
         assert db_author.id == author_id