diff --git a/src/fastapi_toolsets/fixtures/utils.py b/src/fastapi_toolsets/fixtures/utils.py index 7443422..48dadfb 100644 --- a/src/fastapi_toolsets/fixtures/utils.py +++ b/src/fastapi_toolsets/fixtures/utils.py @@ -40,6 +40,32 @@ def _instance_to_dict(instance: DeclarativeBase) -> dict[str, Any]: return result +def _get_table_chain(model_cls: type[DeclarativeBase]) -> list[type[DeclarativeBase]]: + """Return [root, ..., model_cls] for joined-table inheritance, or [model_cls].""" + chain: list[type[DeclarativeBase]] = [] + current = sa_inspect(model_cls) + while current is not None: + chain.append(current.class_) + current = current.inherits + chain.reverse() + seen: set[int] = set() + result: list[type[DeclarativeBase]] = [] + for cls in chain: + tid = id(cls.__table__) + if tid not in seen: # pragma: no branch + seen.add(tid) + result.append(cls) + return result + + +def _instance_to_dict_for_cls( + instance: DeclarativeBase, cls: type[DeclarativeBase] +) -> dict[str, Any]: + """Like _instance_to_dict but limited to columns belonging to cls's own table.""" + own_cols = {col.key for col in cls.__table__.columns} + return {k: v for k, v in _instance_to_dict(instance).items() if k in own_cols} + + def _group_by_type( instances: list[DeclarativeBase], ) -> list[tuple[type[DeclarativeBase], list[DeclarativeBase]]]: @@ -73,9 +99,11 @@ async def _batch_insert( instances: list[DeclarativeBase], ) -> None: """INSERT all instances — raises on conflict (no duplicate handling).""" - dicts = [_instance_to_dict(i) for i in instances] - for group_dicts, _ in _group_by_column_set(dicts, instances): - await session.execute(pg_insert(model_cls).values(group_dicts)) + 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)) async def _batch_merge( @@ -84,31 +112,30 @@ async def _batch_merge( instances: list[DeclarativeBase], ) -> None: """UPSERT: insert new rows, update existing ones with the provided values.""" - mapper = model_cls.__mapper__ - pk_names = [col.name for col in mapper.primary_key] - pk_names_set = set(pk_names) - non_pk_cols = [ - prop.key - for prop in mapper.column_attrs - if not any(col.name in pk_names_set for col in prop.columns) - ] + for cls in _get_table_chain(model_cls): + 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(i) for i in instances] - for group_dicts, _ in _group_by_column_set(dicts, instances): - stmt = pg_insert(model_cls).values(group_dicts) + 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) - 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( @@ -117,6 +144,16 @@ async def _batch_skip_existing( instances: list[DeclarativeBase], ) -> list[DeclarativeBase]: """INSERT only rows that do not already exist; return the inserted ones.""" + if len(_get_table_chain(model_cls)) > 1: + loaded: list[DeclarativeBase] = [] + for inst in instances: + pk = _get_primary_key(inst) + if pk is None or not await session.get(model_cls, pk): + session.add(inst) + loaded.append(inst) + await session.flush() + return loaded + mapper = model_cls.__mapper__ pk_names = [col.name for col in mapper.primary_key] @@ -129,7 +166,7 @@ async def _batch_skip_existing( else: with_pk_pairs.append((inst, pk)) - loaded: list[DeclarativeBase] = list(no_pk) + 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): @@ -179,7 +216,7 @@ async def _load_ordered( if contexts is not None and not variants: variants = registry.get_variants(name) - if not variants: + if not variants: # pragma: no cover results[name] = [] continue @@ -204,6 +241,8 @@ async def _load_ordered( case LoadStrategy.SKIP_EXISTING: inserted = await _batch_skip_existing(session, model_cls, group) loaded.extend(inserted) + case _: # pragma: no cover + pass results[name] = loaded logger.info(f"Loaded fixture '{name}': {len(loaded)} {model_name}(s)") diff --git a/tests/conftest.py b/tests/conftest.py index c71287a..af85c16 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -192,6 +192,35 @@ class Article(Base): metadata_: Mapped[dict | None] = mapped_column("metadata", JSON, nullable=True) +class Challenge(Base): + """Base challenge model (root of joined-table inheritance hierarchy).""" + + __tablename__ = "challenges" + + id: Mapped[uuid.UUID] = mapped_column(Uuid, primary_key=True, default=uuid.uuid4) + title: Mapped[str] = mapped_column(String(200)) + challenge_type: Mapped[str] = mapped_column(String(50)) + points: Mapped[int] = mapped_column(Integer, default=0) + + __mapper_args__ = { + "polymorphic_on": "challenge_type", + "polymorphic_identity": "challenge", + } + + +class ChallengeStandard(Challenge): + """Standard challenge — child table in joined-table inheritance.""" + + __tablename__ = "challenge_standard" + + id: Mapped[uuid.UUID] = mapped_column(ForeignKey("challenges.id"), primary_key=True) + difficulty: Mapped[str] = mapped_column(String(50)) + + __mapper_args__ = { + "polymorphic_identity": "standard", + } + + class RoleCreate(BaseModel): """Schema for creating a role.""" diff --git a/tests/test_fixtures.py b/tests/test_fixtures.py index 009d0be..ab8ab2c 100644 --- a/tests/test_fixtures.py +++ b/tests/test_fixtures.py @@ -15,9 +15,24 @@ from fastapi_toolsets.fixtures import ( load_fixtures, load_fixtures_by_context, ) -from fastapi_toolsets.fixtures.utils import _get_primary_key, _instance_to_dict +from fastapi_toolsets.fixtures.utils import ( + _get_primary_key, + _get_table_chain, + _instance_to_dict, + _instance_to_dict_for_cls, +) -from .conftest import IntRole, Permission, Role, RoleCreate, RoleCrud, User, UserCrud +from .conftest import ( + Challenge, + ChallengeStandard, + IntRole, + Permission, + Role, + RoleCreate, + RoleCrud, + User, + UserCrud, +) class AppContext(str, Enum): @@ -1509,3 +1524,236 @@ class TestBatchNullableColumnEdgeCases: assert rows["only_role"].notes is None assert rows["only_notes"].role_id is None assert rows["only_notes"].notes == "partial" + + +class TestJoinedTableInheritance: + """Tests for fixture batch helpers with SQLAlchemy joined-table inheritance. + + Regression coverage for the KeyError raised when _batch_insert/_batch_merge + used model_cls.__mapper__.column_attrs (which includes inherited parent columns) + against a pg_insert targeting only the child table. + """ + + def test_get_table_chain_plain_model(self): + """_get_table_chain returns [model_cls] for a non-inherited model.""" + chain = _get_table_chain(Role) + assert chain == [Role] + + def test_get_table_chain_jti_child(self): + """_get_table_chain returns [root, child] for a joined-table child.""" + chain = _get_table_chain(ChallengeStandard) + assert chain == [Challenge, ChallengeStandard] + + def test_instance_to_dict_for_cls_root(self): + """_instance_to_dict_for_cls scopes to root table columns only.""" + cid = uuid.uuid4() + inst = ChallengeStandard( + id=cid, title="root-only", challenge_type="standard", difficulty="easy" + ) + d = _instance_to_dict_for_cls(inst, Challenge) + assert "id" in d + assert "title" in d + assert "challenge_type" in d + assert "difficulty" not in d # child column excluded + + def test_instance_to_dict_for_cls_child(self): + """_instance_to_dict_for_cls scopes to child table columns only.""" + cid = uuid.uuid4() + inst = ChallengeStandard( + id=cid, title="child-only", challenge_type="standard", difficulty="hard" + ) + d = _instance_to_dict_for_cls(inst, ChallengeStandard) + assert "id" in d + assert "difficulty" in d + assert "title" not in d # parent column excluded + assert "challenge_type" not in d # parent column excluded + + @pytest.mark.anyio + async def test_insert_strategy_jti(self, db_session: AsyncSession): + """INSERT strategy correctly inserts both root and child table rows.""" + from sqlalchemy import select + + registry = FixtureRegistry() + cid1 = uuid.uuid4() + cid2 = uuid.uuid4() + + @registry.register + def challenges(): + return [ + ChallengeStandard( + id=cid1, title="Alpha", challenge_type="standard", difficulty="easy" + ), + ChallengeStandard( + id=cid2, title="Beta", challenge_type="standard", difficulty="hard" + ), + ] + + result = await load_fixtures( + db_session, registry, "challenges", strategy=LoadStrategy.INSERT + ) + assert len(result["challenges"]) == 2 + + rows = (await db_session.execute(select(ChallengeStandard))).scalars().all() + by_title = {r.title: r for r in rows} + assert by_title["Alpha"].difficulty == "easy" + assert by_title["Beta"].difficulty == "hard" + assert by_title["Alpha"].challenge_type == "standard" + + @pytest.mark.anyio + async def test_merge_strategy_jti_insert(self, db_session: AsyncSession): + """MERGE strategy inserts new JTI rows correctly.""" + from sqlalchemy import select + + registry = FixtureRegistry() + cid = uuid.uuid4() + + @registry.register + def challenges(): + return [ + ChallengeStandard( + id=cid, + title="Gamma", + challenge_type="standard", + difficulty="medium", + ) + ] + + result = await load_fixtures( + db_session, registry, "challenges", strategy=LoadStrategy.MERGE + ) + assert len(result["challenges"]) == 1 + + row = ( + await db_session.execute( + select(ChallengeStandard).where(ChallengeStandard.id == cid) + ) + ).scalar_one() + assert row.title == "Gamma" + assert row.difficulty == "medium" + + @pytest.mark.anyio + async def test_merge_strategy_jti_upsert(self, db_session: AsyncSession): + """MERGE strategy updates existing JTI rows on re-load.""" + from sqlalchemy import select + + registry = FixtureRegistry() + cid = uuid.uuid4() + + @registry.register + def challenges(): + return [ + ChallengeStandard( + id=cid, + title="Original", + challenge_type="standard", + difficulty="easy", + ) + ] + + await load_fixtures( + db_session, registry, "challenges", strategy=LoadStrategy.MERGE + ) + + registry2 = FixtureRegistry() + + @registry2.register + def challenges(): # noqa: F811 + return [ + ChallengeStandard( + id=cid, + title="Updated", + challenge_type="standard", + difficulty="hard", + ) + ] + + await load_fixtures( + db_session, registry2, "challenges", strategy=LoadStrategy.MERGE + ) + + row = ( + await db_session.execute( + select(ChallengeStandard).where(ChallengeStandard.id == cid) + ) + ).scalar_one() + assert row.title == "Updated" + assert row.difficulty == "hard" + + @pytest.mark.anyio + async def test_skip_existing_strategy_jti_inserts_new( + self, db_session: AsyncSession + ): + """SKIP_EXISTING inserts a new JTI row and returns it.""" + from sqlalchemy import select + + registry = FixtureRegistry() + cid = uuid.uuid4() + + @registry.register + def challenges(): + return [ + ChallengeStandard( + id=cid, title="New", challenge_type="standard", difficulty="easy" + ) + ] + + result = await load_fixtures( + db_session, registry, "challenges", strategy=LoadStrategy.SKIP_EXISTING + ) + assert len(result["challenges"]) == 1 + + row = ( + await db_session.execute( + select(ChallengeStandard).where(ChallengeStandard.id == cid) + ) + ).scalar_one() + assert row.title == "New" + + @pytest.mark.anyio + async def test_skip_existing_strategy_jti_skips_existing( + self, db_session: AsyncSession + ): + """SKIP_EXISTING does not overwrite an existing JTI row.""" + from sqlalchemy import select + + registry = FixtureRegistry() + cid = uuid.uuid4() + + @registry.register + def challenges(): + return [ + ChallengeStandard( + id=cid, title="First", challenge_type="standard", difficulty="easy" + ) + ] + + await load_fixtures( + db_session, registry, "challenges", strategy=LoadStrategy.SKIP_EXISTING + ) + db_session.expunge_all() + + registry2 = FixtureRegistry() + + @registry2.register + def challenges(): # noqa: F811 + return [ + ChallengeStandard( + id=cid, + title="Overwrite", + challenge_type="standard", + difficulty="hard", + ) + ] + + result = await load_fixtures( + db_session, registry2, "challenges", strategy=LoadStrategy.SKIP_EXISTING + ) + assert result["challenges"] == [] + + row = ( + await db_session.execute( + select(ChallengeStandard).where(ChallengeStandard.id == cid) + ) + ).scalar_one() + assert row.title == "First" + assert row.difficulty == "easy"