From 651326d54b17a4e96894a9bafb70089662bfc70d Mon Sep 17 00:00:00 2001 From: d3vyce Date: Sun, 26 Jul 2026 10:08:33 -0400 Subject: [PATCH] fix: EventSession.commit() wasted and N+1 reloads --- src/fastapi_toolsets/models/watched.py | 81 ++++++++++++++------------ tests/test_models.py | 79 +++++++++++++------------ 2 files changed, 84 insertions(+), 76 deletions(-) diff --git a/src/fastapi_toolsets/models/watched.py b/src/fastapi_toolsets/models/watched.py index 083a17d..566f008 100644 --- a/src/fastapi_toolsets/models/watched.py +++ b/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: + state = sa_inspect(obj, raiseerr=False) + if state is None or state.detached or state.transient: # pragma: no cover + continue + 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: - 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.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: diff --git a/tests/test_models.py b/tests/test_models.py index 6c60f04..42a5062 100644 --- a/tests/test_models.py +++ b/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