fix: skip redundant refresh/reload round trips in create() and update()

This commit is contained in:
2026-07-26 15:50:13 +02:00
committed by d3vyce
parent 19a823dc8b
commit adc0ff14c1
3 changed files with 110 additions and 15 deletions
+38 -15
View File
@@ -52,7 +52,6 @@ from .search import (
search_field_keys, search_field_keys,
) )
_ForUpdateMode: TypeAlias = bool | Literal["nowait", "skip_locked"] _ForUpdateMode: TypeAlias = bool | Literal["nowait", "skip_locked"]
@@ -174,17 +173,30 @@ class AsyncCrud(Generic[ModelType]):
return cls.default_load_options return cls.default_load_options
@classmethod @classmethod
async def _reload_with_options( def _capture_pk_values(cls: type[Self], instance: ModelType) -> dict[str, Any]:
cls: type[Self], session: AsyncSession, instance: ModelType """Capture PK values off instance — call before commit expires attributes."""
return {
cast(str, col.key): getattr(instance, cast(str, col.key))
for col in cls.model.__mapper__.primary_key
}
@classmethod
async def _reload_with_options_by_pk(
cls: type[Self], session: AsyncSession, pk_values: dict[str, Any]
) -> ModelType: ) -> ModelType:
"""Re-query instance by PK with default_load_options applied.""" """Re-query by previously captured PK values, with default_load_options applied."""
mapper = cls.model.__mapper__ # Only called when cls.default_load_options is set (see call sites).
pk_filters = [ pk_filters = [
getattr(cls.model, cast(str, col.key)) getattr(cls.model, key) == value for key, value in pk_values.items()
== getattr(instance, cast(str, col.key))
for col in mapper.primary_key
] ]
return await cls.get(session, filters=pk_filters) q = select(cls.model).where(and_(*pk_filters))
q = q.execution_options(populate_existing=True)
q = q.options(*cast(Sequence[ExecutableOption], cls.default_load_options))
result = await session.execute(q)
item = result.unique().scalar_one_or_none()
if item is None: # pragma: no cover — row was just flushed in this transaction
raise NotFoundError()
return cast(ModelType, item)
@classmethod @classmethod
async def _resolve_m2m( async def _resolve_m2m(
@@ -738,9 +750,14 @@ class AsyncCrud(Generic[ModelType]):
setattr(db_model, rel_attr, related_instances) setattr(db_model, rel_attr, related_instances)
session.add(db_model) session.add(db_model)
await session.refresh(db_model) pk_values: dict[str, Any] | None = None
if cls.default_load_options: if cls.default_load_options:
db_model = await cls._reload_with_options(session, db_model) await session.flush()
pk_values = cls._capture_pk_values(db_model)
if pk_values is not None:
db_model = await cls._reload_with_options_by_pk(session, pk_values)
else:
await session.refresh(db_model)
result = cast(ModelType, db_model) result = cast(ModelType, db_model)
if schema: if schema:
return Response(data=schema.model_validate(result)) return Response(data=schema.model_validate(result))
@@ -1105,9 +1122,15 @@ class AsyncCrud(Generic[ModelType]):
m2m_resolved = await cls._resolve_m2m(session, obj, only_set=True) m2m_resolved = await cls._resolve_m2m(session, obj, only_set=True)
for rel_attr, related_instances in m2m_resolved.items(): for rel_attr, related_instances in m2m_resolved.items():
setattr(db_model, rel_attr, related_instances) setattr(db_model, rel_attr, related_instances)
await session.refresh(db_model)
if cls.default_load_options: pk_values: dict[str, Any] | None = None
db_model = await cls._reload_with_options(session, db_model) if cls.default_load_options:
await session.flush()
pk_values = cls._capture_pk_values(db_model)
if pk_values is not None:
db_model = await cls._reload_with_options_by_pk(session, pk_values)
else:
await session.refresh(db_model)
if schema: if schema:
return Response(data=schema.model_validate(db_model)) return Response(data=schema.model_validate(db_model))
return db_model return db_model
+23
View File
@@ -476,3 +476,26 @@ async def db_session(engine):
# Drop tables after test # Drop tables after test
async with engine.begin() as conn: async with engine.begin() as conn:
await conn.run_sync(Base.metadata.drop_all) await conn.run_sync(Base.metadata.drop_all)
@pytest.fixture(scope="function")
async def db_session_expire_on_commit(engine):
"""Session with expire_on_commit=True (the SQLAlchemy default).
Attributes read off an instance after commit are expired and trigger an
implicit (sync) refresh under this setting — which fails under asyncio
with MissingGreenlet. The other ``db_session`` fixture uses
``expire_on_commit=False`` and would not catch that class of bug.
"""
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
session_factory = async_sessionmaker(engine, expire_on_commit=True)
session = session_factory()
try:
yield session
finally:
await session.close()
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.drop_all)
+49
View File
@@ -417,6 +417,55 @@ class TestDefaultLoadOptionsIntegration:
assert updated.role is not None assert updated.role is not None
assert updated.role.name == "admin" assert updated.role.name == "admin"
@pytest.mark.anyio
async def test_default_load_options_applied_to_create_expire_on_commit(
self, db_session_expire_on_commit: AsyncSession
):
"""create()'s reload uses captured PK values, not an expired instance attribute.
Regression test for MissingGreenlet: reading a PK off `db_model` after
commit under expire_on_commit=True (the SQLAlchemy default) would
trigger an implicit sync refresh, which fails under asyncio.
"""
UserWithDefaultLoad = CrudFactory(
User, default_load_options=[selectinload(User.role)]
)
role = await RoleCrud.create(
db_session_expire_on_commit, RoleCreate(name="admin")
)
user = await UserWithDefaultLoad.create(
db_session_expire_on_commit,
UserCreate(username="alice", email="alice@test.com", role_id=role.id),
)
assert user.role is not None
assert user.role.name == "admin"
@pytest.mark.anyio
async def test_default_load_options_applied_to_update_expire_on_commit(
self, db_session_expire_on_commit: AsyncSession
):
"""update()'s reload uses captured PK values, not an expired instance attribute.
Regression test for MissingGreenlet under expire_on_commit=True.
"""
UserWithDefaultLoad = CrudFactory(
User, default_load_options=[selectinload(User.role)]
)
role = await RoleCrud.create(
db_session_expire_on_commit, RoleCreate(name="admin")
)
user = await UserCrud.create(
db_session_expire_on_commit,
UserCreate(username="alice", email="alice@test.com"),
)
updated = await UserWithDefaultLoad.update(
db_session_expire_on_commit,
UserUpdate(role_id=role.id),
filters=[User.id == user.id],
)
assert updated.role is not None
assert updated.role.name == "admin"
@pytest.mark.anyio @pytest.mark.anyio
async def test_load_options_overrides_default_load_options( async def test_load_options_overrides_default_load_options(
self, db_session: AsyncSession self, db_session: AsyncSession