add model_dump options to sqlmodel_update

Author: A00474880

sqlmodel/main.py

@@ -984,25 +984,17 @@ def sqlmodel_update(
         obj: builtins.dict[str, Any] | BaseModel,
         *,
         update: builtins.dict[str, Any] | None = None,
+        **model_dump_kwargs,
     ) -> _TSQLModel:
-        use_update = (update or {}).copy()
-        if isinstance(obj, dict):
-            for key, value in {**obj, **use_update}.items():
-                if key in get_model_fields(self):
-                    setattr(self, key, value)
-        elif isinstance(obj, BaseModel):
-            for key in get_model_fields(obj):
-                if key in use_update:
-                    value = use_update.pop(key)
-                else:
-                    value = getattr(obj, key)
-                setattr(self, key, value)
-            for remaining_key, value in use_update.items():
-                if remaining_key in get_model_fields(self):
-                    setattr(self, remaining_key, value)
-        else:
+        if not (isinstance(obj, dict) or isinstance(obj, BaseModel)):
             raise ValueError(
                 "Can't use sqlmodel_update() with something that "
                 f"is not a dict or SQLModel or Pydantic model: {obj}"
             )
+        if isinstance(obj, BaseModel):
+            obj = obj.model_dump(**model_dump_kwargs)
+        use_update = (update or {}).copy()
+        for key, value in {**obj, **use_update}.items():
+            if key in get_model_fields(self):
+                setattr(self, key, value)
         return self

tests/test_update.py

@@ -1,14 +1,17 @@
+from pytest import raises
 from sqlmodel import Field, SQLModel
 
 
 def test_sqlmodel_update():
     class Organization(SQLModel, table=True):
         id: int = Field(default=None, primary_key=True)
         name: str
+        city: str
         headquarters: str
 
     class OrganizationUpdate(SQLModel):
         name: str
+        city: str | None = None
 
     org = Organization(name="Example Org", city="New York", headquarters="NYC HQ")
     org_in = OrganizationUpdate(name="Updated org")
@@ -17,4 +20,13 @@ class OrganizationUpdate(SQLModel):
         update={
             "headquarters": "-",  # This field is in Organization, but not in OrganizationUpdate
         },
+        exclude_unset=True
     )
+    # fields that should stay the same
+    assert org.city == "New York"
+    #fields that should be updated
+    assert org.name == "Updated org"
+    assert org.headquarters == "-"
+    # test raise value error when passing in updates other than dict or BaseModel
+    with raises(ValueError):
+        org.sqlmodel_update(["Boston"])