fix: EventSession.commit() wasted and N+1 reloads

src/fastapi_toolsets/models/watched.py

@@ -5,7 +5,7 @@ from collections.abc import Callable
 from enum import Enum
 from typing import Any
 
-from sqlalchemy import event
+from sqlalchemy import event, select, tuple_
 from sqlalchemy import inspect as sa_inspect
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm.attributes import set_committed_value as _sa_set_committed_value
@@ -29,14 +29,11 @@ _SESSION_DELETES = "_ft_deletes"
 _SESSION_UPDATES = "_ft_updates"
 _DEFERRED_STRATEGY_KEY = (("deferred", True), ("instrument", True))
 _EVENT_HANDLERS: dict[tuple[type, ModelEvent], list[Callable[..., Any]]] = {}
-_WATCHED_MODELS: set[type] = set()
-_WATCHED_CACHE: dict[type, bool] = {}
 _HANDLER_CACHE: dict[tuple[type, ModelEvent], list[Callable[..., Any]]] = {}
 
 
 def _invalidate_caches() -> None:
     """Clear lookup caches after handler registration."""
-    _WATCHED_CACHE.clear()
     _HANDLER_CACHE.clear()
 
 
@@ -56,24 +53,12 @@ def listens_for(
     def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
         for ev in evs:
             _EVENT_HANDLERS.setdefault((model_class, ev), []).append(fn)
-        _WATCHED_MODELS.add(model_class)
         _invalidate_caches()
         return fn
 
     return decorator
 
 
-def _is_watched(obj: Any) -> bool:
-    """Return True if *obj*'s type (or any ancestor) has registered handlers."""
-    cls = type(obj)
-    try:
-        return _WATCHED_CACHE[cls]
-    except KeyError:
-        result = any(klass in _WATCHED_MODELS for klass in cls.__mro__)
-        _WATCHED_CACHE[cls] = result
-        return result
-
-
 def _get_handlers(cls: type, ev: ModelEvent) -> list[Callable[..., Any]]:
     """Return registered handlers for *cls* and *ev*, walking the MRO."""
     key = (cls, ev)
@@ -144,18 +129,18 @@ def _upsert_changes(
 def _after_flush(session: Any, flush_context: Any) -> None:
     # New objects: capture reference. Attributes will be refreshed after commit.
     for obj in session.new:
-        if _is_watched(obj):
+        if _get_handlers(type(obj), ModelEvent.CREATE):
             session.info.setdefault(_SESSION_CREATES, []).append(obj)
 
     # Deleted objects: snapshot now while attributes are still loaded.
     for obj in session.deleted:
-        if _is_watched(obj):
+        if _get_handlers(type(obj), ModelEvent.DELETE):
             snapshot = _snapshot_column_attrs(obj)
             session.info.setdefault(_SESSION_DELETES, []).append((obj, snapshot))
 
     # Dirty objects: read old/new from SQLAlchemy attribute history.
     for obj in session.dirty:
-        if not _is_watched(obj):
+        if not _get_handlers(type(obj), ModelEvent.UPDATE):
             continue
 
         watched = _get_watched_fields(type(obj))
@@ -204,9 +189,18 @@ async def _invoke_callback(
         await result
 
 
-async def _reload_if_present(session: AsyncSession, obj: Any, state: Any) -> None:
-    """Re-populate *obj* from the DB if its row still exists."""
-    await session.get(type(obj), state.key[1], populate_existing=True)
+async def _batch_reload(
+    session: AsyncSession, model: type, pk_tuples: list[tuple[Any, ...]]
+) -> None:
+    """Re-populate all rows of *model* identified by *pk_tuples* in one round trip."""
+    pk_cols = sa_inspect(model).primary_key
+    where = (
+        pk_cols[0].in_([pk[0] for pk in pk_tuples])
+        if len(pk_cols) == 1
+        else tuple_(*pk_cols).in_(pk_tuples)
+    )
+    q = select(model).where(where).execution_options(populate_existing=True)
+    await session.execute(q)
 
 
 class EventSession(AsyncSession):
@@ -250,15 +244,36 @@ class EventSession(AsyncSession):
                 k: v for k, v in field_changes.items() if k not in create_ids
             }
 
-        # Dispatch CREATE callbacks.
+        # Resolve reloadable state up front and group PKs by model type so
+        # the post-commit reload is one query per type instead of one
+        # session.get() per object.
+        create_items: list[Any] = []
+        update_items: list[tuple[Any, dict[str, dict[str, Any]]]] = []
+        pk_by_type: dict[type, list[tuple[Any, ...]]] = {}
+
         for obj in creates:
-            try:
             state = sa_inspect(obj, raiseerr=False)
-                if (
-                    state is None or state.detached or state.transient
-                ):  # pragma: no cover
+            if state is None or state.detached or state.transient:  # pragma: no cover
                 continue
-                await _reload_if_present(self, obj, state)
+            create_items.append(obj)
+            pk_by_type.setdefault(type(obj), []).append(state.key[1])
+
+        for obj, changes in field_changes.values():
+            state = sa_inspect(obj, raiseerr=False)
+            if state is None or state.detached or state.transient:  # pragma: no cover
+                continue
+            update_items.append((obj, changes))
+            pk_by_type.setdefault(type(obj), []).append(state.key[1])
+
+        for model, pk_tuples in pk_by_type.items():
+            try:
+                await _batch_reload(self, model, pk_tuples)
+            except Exception as exc:
+                _logger.error(_CALLBACK_ERROR_MSG, exc_info=exc)
+
+        # Dispatch CREATE callbacks.
+        for obj in create_items:
+            try:
                 for handler in _get_handlers(type(obj), ModelEvent.CREATE):
                     await _invoke_callback(handler, obj, ModelEvent.CREATE, None)
             except Exception as exc:
@@ -275,14 +290,8 @@ class EventSession(AsyncSession):
                 _logger.error(_CALLBACK_ERROR_MSG, exc_info=exc)
 
         # Dispatch UPDATE callbacks.
-        for obj, changes in field_changes.values():
+        for obj, changes in update_items:
             try:
-                state = sa_inspect(obj, raiseerr=False)
-                if (
-                    state is None or state.detached or state.transient
-                ):  # pragma: no cover
-                    continue
-                await _reload_if_present(self, obj, state)
                 for handler in _get_handlers(type(obj), ModelEvent.UPDATE):
                     await _invoke_callback(handler, obj, ModelEvent.UPDATE, changes)
             except Exception as exc:

tests/test_models.py

@@ -25,13 +25,11 @@ from fastapi_toolsets.models.watched import (
     _SESSION_CREATES,
     _SESSION_DELETES,
     _SESSION_UPDATES,
-    _WATCHED_MODELS,
     EventSession,
     _after_flush,
     _after_rollback,
     _get_watched_fields,
     _invalidate_caches,
-    _is_watched,
     _snapshot_column_attrs,
     _upsert_changes,
 )
@@ -658,22 +656,6 @@ class TestWatchInheritance:
         assert "other" in _watch_inherit_events[0]["changes"]
 
 
-class TestIsWatched:
-    def test_watched_model_is_watched(self):
-        """_is_watched returns True for models with registered handlers."""
-        obj = WatchedModel(status="x", other="y")
-        assert _is_watched(obj) is True
-
-    def test_non_watched_model_is_not_watched(self):
-        """_is_watched returns False for models without registered handlers."""
-        assert _is_watched(object()) is False
-
-    def test_subclass_of_watched_model_is_watched(self):
-        """_is_watched returns True for subclasses of watched models (via MRO)."""
-        dog = PolyDog(name="Rex")
-        assert _is_watched(dog) is True
-
-
 class TestUpsertChanges:
     def test_inserts_new_entry(self):
         """New key is inserted with the full changes dict."""
@@ -715,7 +697,10 @@ class TestAfterFlush:
         """New watched objects are added to _SESSION_CREATES."""
         obj = object()
         session = SimpleNamespace(new=[obj], deleted=[], dirty=[], info={})
-        with patch("fastapi_toolsets.models.watched._is_watched", return_value=True):
+        with patch(
+            "fastapi_toolsets.models.watched._get_handlers",
+            return_value=[lambda *a: None],
+        ):
             _after_flush(session, None)
         assert session.info[_SESSION_CREATES] == [obj]
 
@@ -731,7 +716,10 @@ class TestAfterFlush:
         obj = object()
         session = SimpleNamespace(new=[], deleted=[obj], dirty=[], info={})
         with (
-            patch("fastapi_toolsets.models.watched._is_watched", return_value=True),
+            patch(
+                "fastapi_toolsets.models.watched._get_handlers",
+                return_value=[lambda *a: None],
+            ),
             patch(
                 "fastapi_toolsets.models.watched._snapshot_column_attrs",
                 return_value={"id": 1},
@@ -1023,28 +1011,19 @@ class TestEventCallbacks:
                 await other.commit()
             await engine.dispose()
 
-        real_get = mixin_session.get
-        real_refresh = mixin_session.refresh
+        real_batch_reload = _watched_module._batch_reload
 
-        def _matches_doomed(pk):
-            return pk == doomed_id or (isinstance(pk, tuple) and pk[0] == doomed_id)
-
-        async def racing_get(model, pk, *args, **kwargs):
-            if _matches_doomed(pk):
+        async def racing_batch_reload(session, model, pk_tuples):
+            if any(pk[0] == doomed_id for pk in pk_tuples):
                 await kill_doomed_row_once()
-            return await real_get(model, pk, *args, **kwargs)
+            return await real_batch_reload(session, model, pk_tuples)
 
-        async def racing_refresh(obj, *args, **kwargs):
-            if getattr(obj, "id", None) == doomed_id:
-                await kill_doomed_row_once()
-            return await real_refresh(obj, *args, **kwargs)
-
-        # Patch both possible reload mechanisms (session.get / session.refresh)
-        # so this test still exercises the race regardless of which one
-        # EventSession.commit() uses internally to pick up server defaults.
-        mixin_session.get = racing_get
-        mixin_session.refresh = racing_refresh
-        with patch.object(_watched_module._logger, "error") as mock_error:
+        # Patch the batched reload EventSession.commit() uses to pick up
+        # server defaults, so this test still exercises the race.
+        with (
+            patch.object(_watched_module, "_batch_reload", racing_batch_reload),
+            patch.object(_watched_module._logger, "error") as mock_error,
+        ):
             await mixin_session.commit()
             mock_error.assert_not_called()
 
@@ -1052,6 +1031,27 @@ class TestEventCallbacks:
         created_ids = {e["obj_id"] for e in _test_events if e["event"] == "create"}
         assert created_ids == {keep.id, doomed_id}
 
+    @pytest.mark.anyio
+    async def test_batch_reload_exception_is_logged_and_dispatch_continues(
+        self, mixin_session
+    ):
+        """A batched-reload failure is logged; CREATE handlers still fire."""
+        obj = WatchedModel(status="active", other="x")
+        mixin_session.add(obj)
+
+        async def failing_batch_reload(session, model, pk_tuples):
+            raise RuntimeError("reload failed")
+
+        with (
+            patch.object(_watched_module, "_batch_reload", failing_batch_reload),
+            patch.object(_watched_module._logger, "error") as mock_error,
+        ):
+            await mixin_session.commit()
+
+        mock_error.assert_called_once()
+        creates = [e for e in _test_events if e["event"] == "create"]
+        assert len(creates) == 1
+
 
 class TestTransientObject:
     """Create + delete within the same transaction should fire no events."""
@@ -1421,7 +1421,6 @@ class TestListensFor:
         for key in list(_EVENT_HANDLERS):
             if key[0] is ListenerModel:
                 del _EVENT_HANDLERS[key]
-        _WATCHED_MODELS.discard(ListenerModel)
         _invalidate_caches()
 
     @pytest.mark.anyio