mirror of
https://github.com/d3vyce/fastapi-toolsets.git
synced 2026-08-06 00:24:08 +00:00
Rework fixtures module (#333)
* feat: fixture refresh DB-generated values onto returned instances * refactor: replace utils.get_obj_by_attr/get_field_by_attr with registry.obj/field lookups * refactor: log fixture command output instead of print * chore: clean up fixture module
This commit is contained in:
d3vyce
@@ -1,17 +1,18 @@
|
||||
"""Fixture loading utilities for database seeding."""
|
||||
|
||||
from collections.abc import Callable, Sequence
|
||||
from collections.abc import Iterator
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
from typing import Any, cast
|
||||
|
||||
from sqlalchemy import Table, select
|
||||
from sqlalchemy import inspect as sa_inspect
|
||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.orm import DeclarativeBase
|
||||
from sqlalchemy.orm import DeclarativeBase, selectinload
|
||||
from sqlalchemy.orm.interfaces import ExecutableOption, ORMOption
|
||||
|
||||
from ..db import transaction
|
||||
from ..logger import get_logger
|
||||
from ..types import ModelType
|
||||
from .enum import LoadStrategy
|
||||
from .registry import FixtureRegistry, _normalize_contexts
|
||||
|
||||
@@ -93,17 +94,42 @@ def _group_by_column_set(
|
||||
return list(groups.values())
|
||||
|
||||
|
||||
def _grouped_table_dicts(
|
||||
model_cls: type[DeclarativeBase], instances: list[DeclarativeBase]
|
||||
) -> Iterator[
|
||||
tuple[type[DeclarativeBase], list[dict[str, Any]], list[DeclarativeBase]]
|
||||
]:
|
||||
"""Yield (cls, group_dicts, group_instances) per table in the inheritance
|
||||
chain and per column-set group, skipping empty groups.
|
||||
"""
|
||||
for cls in _get_table_chain(model_cls):
|
||||
dicts = [_instance_to_dict_for_cls(i, cls) for i in instances]
|
||||
for group_dicts, group_instances in _group_by_column_set(dicts, instances):
|
||||
if group_dicts and group_dicts[0]: # pragma: no branch
|
||||
yield cls, group_dicts, group_instances
|
||||
|
||||
|
||||
async def _batch_insert(
|
||||
session: AsyncSession,
|
||||
model_cls: type[DeclarativeBase],
|
||||
instances: list[DeclarativeBase],
|
||||
) -> None:
|
||||
"""INSERT all instances — raises on conflict (no duplicate handling)."""
|
||||
for cls in _get_table_chain(model_cls):
|
||||
dicts = [_instance_to_dict_for_cls(i, cls) for i in instances]
|
||||
for group_dicts, _ in _group_by_column_set(dicts, instances):
|
||||
if group_dicts and group_dicts[0]: # pragma: no branch
|
||||
await session.execute(pg_insert(cls).values(group_dicts))
|
||||
"""INSERT all instances, raises on conflict."""
|
||||
for cls, group_dicts, group_instances in _grouped_table_dicts(model_cls, instances):
|
||||
table = cast(Table, cls.__table__)
|
||||
missing_pk_cols = [
|
||||
col for col in table.primary_key.columns if col.key not in group_dicts[0]
|
||||
]
|
||||
if not missing_pk_cols:
|
||||
await session.execute(pg_insert(table), group_dicts)
|
||||
continue
|
||||
stmt = pg_insert(table).returning(
|
||||
*missing_pk_cols, sort_by_parameter_order=True
|
||||
)
|
||||
result = await session.execute(stmt, group_dicts)
|
||||
for inst, row in zip(group_instances, result):
|
||||
for col, val in zip(missing_pk_cols, row):
|
||||
setattr(inst, col.key, val)
|
||||
|
||||
|
||||
async def _batch_merge(
|
||||
@@ -112,30 +138,26 @@ async def _batch_merge(
|
||||
instances: list[DeclarativeBase],
|
||||
) -> None:
|
||||
"""UPSERT: insert new rows, update existing ones with the provided values."""
|
||||
for cls in _get_table_chain(model_cls):
|
||||
for cls, group_dicts, _ in _grouped_table_dicts(model_cls, instances):
|
||||
pk_names = [col.name for col in cls.__table__.primary_key]
|
||||
pk_names_set = set(pk_names)
|
||||
own_col_keys = {col.key for col in cls.__table__.columns}
|
||||
non_pk_cols = [k for k in own_col_keys if k not in pk_names_set]
|
||||
|
||||
dicts = [_instance_to_dict_for_cls(i, cls) for i in instances]
|
||||
for group_dicts, _ in _group_by_column_set(dicts, instances):
|
||||
if not group_dicts or not group_dicts[0]: # pragma: no cover
|
||||
continue
|
||||
stmt = pg_insert(cls).values(group_dicts)
|
||||
stmt = pg_insert(cls).values(group_dicts)
|
||||
|
||||
inserted_keys = set(group_dicts[0])
|
||||
update_cols = [col for col in non_pk_cols if col in inserted_keys]
|
||||
inserted_keys = set(group_dicts[0])
|
||||
update_cols = [col for col in non_pk_cols if col in inserted_keys]
|
||||
|
||||
if update_cols:
|
||||
stmt = stmt.on_conflict_do_update(
|
||||
index_elements=pk_names,
|
||||
set_={col: stmt.excluded[col] for col in update_cols},
|
||||
)
|
||||
else:
|
||||
stmt = stmt.on_conflict_do_nothing(index_elements=pk_names)
|
||||
if update_cols:
|
||||
stmt = stmt.on_conflict_do_update(
|
||||
index_elements=pk_names,
|
||||
set_={col: stmt.excluded[col] for col in update_cols},
|
||||
)
|
||||
else:
|
||||
stmt = stmt.on_conflict_do_nothing(index_elements=pk_names)
|
||||
|
||||
await session.execute(stmt)
|
||||
await session.execute(stmt)
|
||||
|
||||
|
||||
async def _batch_skip_existing(
|
||||
@@ -169,8 +191,14 @@ async def _batch_skip_existing(
|
||||
loaded = list(no_pk)
|
||||
if no_pk:
|
||||
no_pk_dicts = [_instance_to_dict(i) for i in no_pk]
|
||||
for group_dicts, _ in _group_by_column_set(no_pk_dicts, no_pk):
|
||||
await session.execute(pg_insert(model_cls).values(group_dicts))
|
||||
for group_dicts, group_instances in _group_by_column_set(no_pk_dicts, no_pk):
|
||||
stmt = pg_insert(cast(Table, model_cls.__table__)).returning(
|
||||
*mapper.primary_key, sort_by_parameter_order=True
|
||||
)
|
||||
result = await session.execute(stmt, group_dicts)
|
||||
for inst, row in zip(group_instances, result):
|
||||
for col, val in zip(mapper.primary_key, row):
|
||||
setattr(inst, col.key, val)
|
||||
|
||||
if with_pk_pairs:
|
||||
with_pk = [i for i, _ in with_pk_pairs]
|
||||
@@ -196,6 +224,64 @@ async def _batch_skip_existing(
|
||||
return loaded
|
||||
|
||||
|
||||
def _relationship_load_options(model: type[DeclarativeBase]) -> list[ExecutableOption]:
|
||||
"""Build selectinload options for all direct relationships on a model."""
|
||||
return [
|
||||
selectinload(getattr(model, rel.key)) for rel in model.__mapper__.relationships
|
||||
]
|
||||
|
||||
|
||||
async def _reload_with_relationships(
|
||||
session: AsyncSession,
|
||||
instances: list[DeclarativeBase],
|
||||
load_options: list[ExecutableOption],
|
||||
) -> list[DeclarativeBase]:
|
||||
"""Reload instances in a single bulk query with relationship eager-loading."""
|
||||
model = type(instances[0])
|
||||
mapper = model.__mapper__
|
||||
pk_cols = mapper.primary_key
|
||||
|
||||
if len(pk_cols) == 1:
|
||||
pk_attr = getattr(model, pk_cols[0].key)
|
||||
pks = [getattr(inst, pk_cols[0].key) for inst in instances]
|
||||
result = await session.execute(
|
||||
select(model).where(pk_attr.in_(pks)).options(*load_options)
|
||||
)
|
||||
by_pk = {getattr(row, pk_cols[0].key): row for row in result.unique().scalars()}
|
||||
return [by_pk[pk] for pk in pks]
|
||||
|
||||
# Composite PK: fall back to per-instance reload
|
||||
reloaded: list[DeclarativeBase] = []
|
||||
for instance in instances:
|
||||
pk = _get_primary_key(instance)
|
||||
refreshed = await session.get(
|
||||
model,
|
||||
pk,
|
||||
options=cast(list[ORMOption], load_options),
|
||||
populate_existing=True,
|
||||
)
|
||||
if refreshed is not None: # pragma: no branch
|
||||
reloaded.append(refreshed)
|
||||
return reloaded
|
||||
|
||||
|
||||
async def _refresh_loaded(
|
||||
session: AsyncSession, instances: list[DeclarativeBase]
|
||||
) -> list[DeclarativeBase]:
|
||||
"""Re-select freshly written rows, eager-loading relationships."""
|
||||
if not instances:
|
||||
return []
|
||||
refreshed: list[DeclarativeBase | None] = [None] * len(instances)
|
||||
for model_cls, group in _group_by_type(instances):
|
||||
positions = [i for i, inst in enumerate(instances) if type(inst) is model_cls]
|
||||
load_options = _relationship_load_options(model_cls)
|
||||
for pos, new in zip(
|
||||
positions, await _reload_with_relationships(session, group, load_options)
|
||||
):
|
||||
refreshed[pos] = new
|
||||
return cast(list[DeclarativeBase], refreshed)
|
||||
|
||||
|
||||
async def _load_ordered(
|
||||
session: AsyncSession,
|
||||
registry: FixtureRegistry,
|
||||
@@ -208,14 +294,11 @@ async def _load_ordered(
|
||||
|
||||
for name in ordered_names:
|
||||
variants = (
|
||||
registry.get_variants(name, *contexts)
|
||||
registry.get_load_variants(name, *contexts)
|
||||
if contexts is not None
|
||||
else registry.get_variants(name)
|
||||
)
|
||||
|
||||
if contexts is not None and not variants:
|
||||
variants = registry.get_variants(name)
|
||||
|
||||
if not variants: # pragma: no cover
|
||||
results[name] = []
|
||||
continue
|
||||
@@ -244,8 +327,10 @@ async def _load_ordered(
|
||||
case _: # pragma: no cover
|
||||
pass
|
||||
|
||||
loaded = await _refresh_loaded(session, loaded)
|
||||
|
||||
results[name] = loaded
|
||||
logger.info(f"Loaded fixture '{name}': {len(loaded)} {model_name}(s)")
|
||||
logger.info("Loaded fixture '%s': %d %s(s)", name, len(loaded), model_name)
|
||||
|
||||
return results
|
||||
|
||||
@@ -264,56 +349,6 @@ def _get_primary_key(instance: DeclarativeBase) -> Any | None:
|
||||
return None
|
||||
|
||||
|
||||
def get_obj_by_attr(
|
||||
fixtures: Callable[[], Sequence[ModelType]], attr_name: str, value: Any
|
||||
) -> ModelType:
|
||||
"""Get a SQLAlchemy model instance by matching an attribute value.
|
||||
|
||||
Args:
|
||||
fixtures: A fixture function registered via ``@registry.register``
|
||||
that returns a sequence of SQLAlchemy model instances.
|
||||
attr_name: Name of the attribute to match against.
|
||||
value: Value to match.
|
||||
|
||||
Returns:
|
||||
The first model instance where the attribute matches the given value.
|
||||
|
||||
Raises:
|
||||
StopIteration: If no matching object is found in the fixture group.
|
||||
"""
|
||||
try:
|
||||
return next(obj for obj in fixtures() if getattr(obj, attr_name) == value)
|
||||
except StopIteration:
|
||||
raise StopIteration(
|
||||
f"No object with {attr_name}={value} found in fixture '{getattr(fixtures, '__name__', repr(fixtures))}'"
|
||||
) from None
|
||||
|
||||
|
||||
def get_field_by_attr(
|
||||
fixtures: Callable[[], Sequence[ModelType]],
|
||||
attr_name: str,
|
||||
value: Any,
|
||||
*,
|
||||
field: str = "id",
|
||||
) -> Any:
|
||||
"""Get a single field value from a fixture object matched by an attribute.
|
||||
|
||||
Args:
|
||||
fixtures: A fixture function registered via ``@registry.register``
|
||||
that returns a sequence of SQLAlchemy model instances.
|
||||
attr_name: Name of the attribute to match against.
|
||||
value: Value to match.
|
||||
field: Attribute name to return from the matched object (default: ``"id"``).
|
||||
|
||||
Returns:
|
||||
The value of ``field`` on the first matching model instance.
|
||||
|
||||
Raises:
|
||||
StopIteration: If no matching object is found in the fixture group.
|
||||
"""
|
||||
return getattr(get_obj_by_attr(fixtures, attr_name, value), field)
|
||||
|
||||
|
||||
async def load_fixtures(
|
||||
session: AsyncSession,
|
||||
registry: FixtureRegistry,
|
||||
|
||||
Reference in New Issue
Block a user