test: assert sqlmodel_update fluent identity and Session.exec non-scalar queries

tests/test_query.py

@@ -1,6 +1,6 @@
 import pytest
 from sqlalchemy import Engine
-from sqlmodel import Field, Session, SQLModel
+from sqlmodel import Field, Session, SQLModel, select
 
 
 def test_query(database_engine: Engine):
@@ -25,3 +25,57 @@ class Hero(SQLModel, table=True):
             query_hero = session.query(Hero).first()
         assert query_hero
         assert query_hero.name == hero_1.name
+
+
+def test_exec_multi_column_select(database_engine: Engine):
+    class Hero(SQLModel, table=True):
+        id: int | None = Field(default=None, primary_key=True)
+        name: str
+        secret_name: str
+
+    hero_1 = Hero(name="Deadpond", secret_name="Dive Wilson")
+    engine = database_engine
+    SQLModel.metadata.create_all(engine)
+
+    with Session(engine) as session:
+        session.add(hero_1)
+        session.commit()
+        session.refresh(hero_1)
+
+    with Session(engine) as session:
+        statement = select(Hero.id, Hero.name)
+        result = session.exec(statement)
+        assert result is not None
+        rows = result.all()
+        assert len(rows) == 1
+        assert rows[0] == (hero_1.id, "Deadpond")
+
+
+def test_exec_update_statement(database_engine: Engine):
+    from sqlalchemy import update
+
+    class Hero(SQLModel, table=True):
+        id: int | None = Field(default=None, primary_key=True)
+        name: str
+        secret_name: str
+
+    hero_1 = Hero(name="Deadpond", secret_name="Dive Wilson")
+    engine = database_engine
+    SQLModel.metadata.create_all(engine)
+
+    with Session(engine) as session:
+        session.add(hero_1)
+        session.commit()
+        session.refresh(hero_1)
+
+    with Session(engine) as session:
+        statement = update(Hero).where(Hero.name == "Deadpond").values(name="Alivepond")
+        cursor_result = session.exec(statement)
+        assert cursor_result is not None
+        assert cursor_result.rowcount == 1
+        session.commit()
+
+    with Session(engine) as session:
+        updated = session.exec(select(Hero).where(Hero.id == hero_1.id)).first()
+        assert updated is not None
+        assert updated.name == "Alivepond"

tests/test_update.py

@@ -1,3 +1,4 @@
+import pytest
 from sqlmodel import Field, SQLModel
 
 
@@ -12,9 +13,65 @@ class OrganizationUpdate(SQLModel):
 
     org = Organization(name="Example Org", city="New York", headquarters="NYC HQ")
     org_in = OrganizationUpdate(name="Updated org")
-    org.sqlmodel_update(
+    result = org.sqlmodel_update(
         org_in,
         update={
             "headquarters": "-",  # This field is in Organization, but not in OrganizationUpdate
         },
     )
+    assert result is org
+    assert org.name == "Updated org"
+    assert org.headquarters == "-"
+    assert org.id is None
+
+
+def test_sqlmodel_update_dict():
+    class Organization(SQLModel, table=True):
+        id: int = Field(default=None, primary_key=True)
+        name: str
+        headquarters: str
+
+    org = Organization(name="Example Org", headquarters="NYC HQ")
+    result = org.sqlmodel_update(
+        {"name": "New Name", "extra_field": "val"},
+        update={"headquarters": "Berlin", "other_extra": "val2"},
+    )
+    assert result is org
+    assert org.name == "New Name"
+    assert org.headquarters == "Berlin"
+    assert not hasattr(org, "extra_field")
+    assert not hasattr(org, "other_extra")
+
+
+def test_sqlmodel_update_base_model_extra_update_keys():
+    class Organization(SQLModel, table=True):
+        id: int = Field(default=None, primary_key=True)
+        name: str
+        headquarters: str
+
+    class OrganizationUpdate(SQLModel):
+        name: str
+
+    org = Organization(name="Example Org", headquarters="NYC HQ")
+    org_in = OrganizationUpdate(name="Updated org")
+    result = org.sqlmodel_update(
+        org_in,
+        update={
+            "headquarters": "Berlin",
+            "non_model_field": "ignored",
+        },
+    )
+    assert result is org
+    assert org.name == "Updated org"
+    assert org.headquarters == "Berlin"
+    assert not hasattr(org, "non_model_field")
+
+
+def test_sqlmodel_update_invalid_type_raises_value_error():
+    class Organization(SQLModel, table=True):
+        id: int = Field(default=None, primary_key=True)
+        name: str
+
+    org = Organization(name="Example Org")
+    with pytest.raises(ValueError, match="Can't use sqlmodel_update"):
+        org.sqlmodel_update(42)  # type: ignore