fix: correctly handle id_attribute with update (#502)

Correctly merge attributes onto existing instance when using `id_attribute` and `update`

Author: cofin

advanced_alchemy/repository/_async.py

@@ -1,3 +1,4 @@
+import contextlib
 import datetime
 import decimal
 import random
@@ -24,12 +25,14 @@
     Update,
     any_,
     delete,
+    inspect,
     over,
     select,
     text,
     update,
 )
 from sqlalchemy import func as sql_func
+from sqlalchemy.exc import MissingGreenlet, NoInspectionAvailable
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.ext.asyncio.scoping import async_scoped_session
 from sqlalchemy.orm import InstrumentedAttribute
@@ -45,6 +48,7 @@
     FilterableRepository,
     FilterableRepositoryProtocol,
     LoadSpec,
+    column_has_defaults,
     get_abstract_loader_options,
     get_instrumented_attr,
 )
@@ -1478,10 +1482,45 @@ async def update(
                 data,
                 id_attribute=id_attribute,
             )
-            # this will raise for not found, and will put the item in the session
-            await self.get(item_id, id_attribute=id_attribute, load=load, execution_options=execution_options)
-            # this will merge the inbound data to the instance we just put in the session
-            instance = await self._attach_to_session(data, strategy="merge")
+            existing_instance = await self.get(
+                item_id, id_attribute=id_attribute, load=load, execution_options=execution_options
+            )
+            mapper = None
+            with (
+                self.session.no_autoflush,
+                contextlib.suppress(MissingGreenlet, NoInspectionAvailable),
+            ):
+                mapper = inspect(data)
+                if mapper is not None:
+                    for column in mapper.mapper.columns:
+                        field_name = column.key
+                        new_field_value = getattr(data, field_name, MISSING)
+                        if new_field_value is not MISSING:
+                            # Skip setting columns with defaults/onupdate to None during updates
+                            # This prevents overwriting columns that should use their defaults
+                            if new_field_value is None and column_has_defaults(column):
+                                continue
+                            existing_field_value = getattr(existing_instance, field_name, MISSING)
+                            if existing_field_value is not MISSING and existing_field_value != new_field_value:
+                                setattr(existing_instance, field_name, new_field_value)
+
+                    # Handle relationships by merging objects into session first
+                    for relationship in mapper.mapper.relationships:
+                        if (new_value := getattr(data, relationship.key, MISSING)) is not MISSING:
+                            if isinstance(new_value, list):
+                                merged_values = [  # pyright: ignore
+                                    await self.session.merge(item, load=False)  # pyright: ignore
+                                    for item in new_value  # pyright: ignore
+                                ]
+                                setattr(existing_instance, relationship.key, merged_values)
+                            elif new_value is not None:
+                                merged_value = await self.session.merge(new_value, load=False)
+                                setattr(existing_instance, relationship.key, merged_value)
+                            else:
+                                setattr(existing_instance, relationship.key, new_value)
+
+            instance = await self._attach_to_session(existing_instance, strategy="merge")
+
             await self._flush_or_commit(auto_commit=auto_commit)
             await self._refresh(
                 instance,

advanced_alchemy/repository/_sync.py

@@ -1,5 +1,6 @@
 # Do not edit this file directly. It has been autogenerated from
 # advanced_alchemy/repository/_async.py
+import contextlib
 import datetime
 import decimal
 import random
@@ -26,12 +27,14 @@
     Update,
     any_,
     delete,
+    inspect,
     over,
     select,
     text,
     update,
 )
 from sqlalchemy import func as sql_func
+from sqlalchemy.exc import MissingGreenlet, NoInspectionAvailable
 from sqlalchemy.orm import InstrumentedAttribute, Session
 from sqlalchemy.orm.scoping import scoped_session
 from sqlalchemy.orm.strategy_options import _AbstractLoad  # pyright: ignore[reportPrivateUsage]
@@ -46,6 +49,7 @@
     FilterableRepository,
     FilterableRepositoryProtocol,
     LoadSpec,
+    column_has_defaults,
     get_abstract_loader_options,
     get_instrumented_attr,
 )
@@ -1479,10 +1483,45 @@ def update(
                 data,
                 id_attribute=id_attribute,
             )
-            # this will raise for not found, and will put the item in the session
-            self.get(item_id, id_attribute=id_attribute, load=load, execution_options=execution_options)
-            # this will merge the inbound data to the instance we just put in the session
-            instance = self._attach_to_session(data, strategy="merge")
+            existing_instance = self.get(
+                item_id, id_attribute=id_attribute, load=load, execution_options=execution_options
+            )
+            mapper = None
+            with (
+                self.session.no_autoflush,
+                contextlib.suppress(MissingGreenlet, NoInspectionAvailable),
+            ):
+                mapper = inspect(data)
+                if mapper is not None:
+                    for column in mapper.mapper.columns:
+                        field_name = column.key
+                        new_field_value = getattr(data, field_name, MISSING)
+                        if new_field_value is not MISSING:
+                            # Skip setting columns with defaults/onupdate to None during updates
+                            # This prevents overwriting columns that should use their defaults
+                            if new_field_value is None and column_has_defaults(column):
+                                continue
+                            existing_field_value = getattr(existing_instance, field_name, MISSING)
+                            if existing_field_value is not MISSING and existing_field_value != new_field_value:
+                                setattr(existing_instance, field_name, new_field_value)
+
+                    # Handle relationships by merging objects into session first
+                    for relationship in mapper.mapper.relationships:
+                        if (new_value := getattr(data, relationship.key, MISSING)) is not MISSING:
+                            if isinstance(new_value, list):
+                                merged_values = [  # pyright: ignore
+                                    self.session.merge(item, load=False)  # pyright: ignore
+                                    for item in new_value  # pyright: ignore
+                                ]
+                                setattr(existing_instance, relationship.key, merged_values)
+                            elif new_value is not None:
+                                merged_value = self.session.merge(new_value, load=False)
+                                setattr(existing_instance, relationship.key, merged_value)
+                            else:
+                                setattr(existing_instance, relationship.key, new_value)
+
+            instance = self._attach_to_session(existing_instance, strategy="merge")
+
             self._flush_or_commit(auto_commit=auto_commit)
             self._refresh(
                 instance,

advanced_alchemy/repository/_util.py

@@ -339,3 +339,26 @@ def _order_by_attribute(
         if not isinstance(statement, Select):
             return statement
         return statement.order_by(field.desc() if is_desc else field.asc())
+
+
+def column_has_defaults(column: Any) -> bool:
+    """Check if a column has any type of default value or update handler.
+
+    This includes:
+    - Python-side defaults (column.default)
+    - Server-side defaults (column.server_default)
+    - Python-side onupdate handlers (column.onupdate)
+    - Server-side onupdate handlers (column.server_onupdate)
+
+    Args:
+        column: SQLAlchemy column object to check
+
+    Returns:
+        bool: True if the column has any type of default or update handler
+    """
+    return (
+        column.default is not None
+        or column.server_default is not None
+        or column.onupdate is not None
+        or column.server_onupdate is not None
+    )

advanced_alchemy/service/_async.py

@@ -24,7 +24,7 @@
     SQLAlchemyAsyncQueryRepository,
 )
 from advanced_alchemy.repository._util import LoadSpec, model_from_dict
-from advanced_alchemy.repository.typing import ModelT, OrderingPair, SQLAlchemyAsyncRepositoryT
+from advanced_alchemy.repository.typing import MISSING, ModelT, OrderingPair, SQLAlchemyAsyncRepositoryT
 from advanced_alchemy.service._util import ResultConverter
 from advanced_alchemy.service.typing import (
     BulkModelDictT,
@@ -713,22 +713,33 @@ async def update(
         Returns:
             Updated representation.
         """
-        data = await self.to_model(data, "update")
-        if (
-            item_id is None
-            and self.repository.get_id_attribute_value(  # pyright: ignore[reportUnknownMemberType]
-                item=data,
-                id_attribute=id_attribute,
-            )
-            is None
-        ):
-            msg = (
-                "Could not identify ID attribute value.  One of the following is required: "
-                f"``item_id`` or ``data.{id_attribute or self.repository.id_attribute}``"
+        if is_dict(data) and item_id is not None:
+            existing_instance = await self.repository.get(
+                item_id, id_attribute=id_attribute, load=load, execution_options=execution_options
             )
-            raise RepositoryError(msg)
-        if item_id is not None:
-            data = self.repository.set_id_attribute_value(item_id=item_id, item=data, id_attribute=id_attribute)  # pyright: ignore[reportUnknownMemberType]
+            update_data = await self.to_model_on_update(data)
+            if is_dict(update_data):
+                for key, value in update_data.items():
+                    if getattr(existing_instance, key, MISSING) is not MISSING:
+                        setattr(existing_instance, key, value)
+            data = existing_instance
+        else:
+            data = await self.to_model(data, "update")
+            if (
+                item_id is None
+                and self.repository.get_id_attribute_value(  # pyright: ignore[reportUnknownMemberType]
+                    item=data,
+                    id_attribute=id_attribute,
+                )
+                is None
+            ):
+                msg = (
+                    "Could not identify ID attribute value.  One of the following is required: "
+                    f"``item_id`` or ``data.{id_attribute or self.repository.id_attribute}``"
+                )
+                raise RepositoryError(msg)
+            if item_id is not None:
+                data = self.repository.set_id_attribute_value(item_id=item_id, item=data, id_attribute=id_attribute)  # pyright: ignore[reportUnknownMemberType]
         return cast(
             "ModelT",
             await self.repository.update(

advanced_alchemy/service/_sync.py

@@ -23,7 +23,7 @@
 from advanced_alchemy.filters import StatementFilter
 from advanced_alchemy.repository import SQLAlchemySyncQueryRepository
 from advanced_alchemy.repository._util import LoadSpec, model_from_dict
-from advanced_alchemy.repository.typing import ModelT, OrderingPair, SQLAlchemySyncRepositoryT
+from advanced_alchemy.repository.typing import MISSING, ModelT, OrderingPair, SQLAlchemySyncRepositoryT
 from advanced_alchemy.service._util import ResultConverter
 from advanced_alchemy.service.typing import (
     BulkModelDictT,
@@ -712,22 +712,33 @@ def update(
         Returns:
             Updated representation.
         """
-        data = self.to_model(data, "update")
-        if (
-            item_id is None
-            and self.repository.get_id_attribute_value(  # pyright: ignore[reportUnknownMemberType]
-                item=data,
-                id_attribute=id_attribute,
-            )
-            is None
-        ):
-            msg = (
-                "Could not identify ID attribute value.  One of the following is required: "
-                f"``item_id`` or ``data.{id_attribute or self.repository.id_attribute}``"
+        if is_dict(data) and item_id is not None:
+            existing_instance = self.repository.get(
+                item_id, id_attribute=id_attribute, load=load, execution_options=execution_options
             )
-            raise RepositoryError(msg)
-        if item_id is not None:
-            data = self.repository.set_id_attribute_value(item_id=item_id, item=data, id_attribute=id_attribute)  # pyright: ignore[reportUnknownMemberType]
+            update_data = self.to_model_on_update(data)
+            if is_dict(update_data):
+                for key, value in update_data.items():
+                    if getattr(existing_instance, key, MISSING) is not MISSING:
+                        setattr(existing_instance, key, value)
+            data = existing_instance
+        else:
+            data = self.to_model(data, "update")
+            if (
+                item_id is None
+                and self.repository.get_id_attribute_value(  # pyright: ignore[reportUnknownMemberType]
+                    item=data,
+                    id_attribute=id_attribute,
+                )
+                is None
+            ):
+                msg = (
+                    "Could not identify ID attribute value.  One of the following is required: "
+                    f"``item_id`` or ``data.{id_attribute or self.repository.id_attribute}``"
+                )
+                raise RepositoryError(msg)
+            if item_id is not None:
+                data = self.repository.set_id_attribute_value(item_id=item_id, item=data, id_attribute=id_attribute)  # pyright: ignore[reportUnknownMemberType]
         return cast(
             "ModelT",
             self.repository.update(

tests/integration/test_repository.py

@@ -18,7 +18,7 @@
 from pytest_lazy_fixtures import lf
 from sqlalchemy import Engine, Table, and_, insert, select, text
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
-from sqlalchemy.orm import Session, sessionmaker
+from sqlalchemy.orm import Session, selectinload, sessionmaker
 from time_machine import travel
 
 from advanced_alchemy import base
@@ -2032,10 +2032,12 @@ async def test_lazy_load(
     assert len(tags_to_add) > 0  # pyright: ignore
     assert tags_to_add[0].id is not None  # pyright: ignore
     update_data["tags"] = tags_to_add  # type: ignore[assignment]
-    updated_obj = await maybe_async(item_repo.update(item_model(**update_data), auto_refresh=False))
+    await maybe_async(item_repo.update(item_model(**update_data), load=[selectinload(item_repo.model_type.tags)]))
+    # Refresh the object to ensure tags are loaded before assertions
+    refreshed_obj = await maybe_async(item_repo.get(first_item_id, load=[selectinload(item_repo.model_type.tags)]))
     await maybe_async(item_repo.session.commit())
-    assert len(updated_obj.tags) > 0
-    assert updated_obj.tags[0].name == "A new tag"
+    assert len(refreshed_obj.tags) > 0
+    assert refreshed_obj.tags[0].name == "A new tag"
 
 
 async def test_repo_health_check(author_repo: AnyAuthorRepository) -> None: