Add --prune flag to seed command to remove stale data
Without --prune, seeds continue to only upsert (add/update). With --prune, routes, encounters, and bosses not present in the seed JSON files are deleted from the database. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -124,11 +124,14 @@ async def upsert_routes(
|
||||
session: AsyncSession,
|
||||
version_group_id: int,
|
||||
routes: list[dict],
|
||||
*,
|
||||
prune: bool = False,
|
||||
) -> dict[str, int]:
|
||||
"""Upsert route records for a version group, return {name: id} mapping.
|
||||
|
||||
Handles hierarchical routes: routes with 'children' are parent routes,
|
||||
and their children get parent_route_id set accordingly.
|
||||
When prune is True, deletes routes not present in the seed data.
|
||||
"""
|
||||
# First pass: upsert all parent routes (without parent_route_id)
|
||||
for route in routes:
|
||||
@@ -185,6 +188,27 @@ async def upsert_routes(
|
||||
|
||||
await session.flush()
|
||||
|
||||
if prune:
|
||||
seed_names: set[str] = set()
|
||||
for route in routes:
|
||||
seed_names.add(route["name"])
|
||||
for child in route.get("children", []):
|
||||
seed_names.add(child["name"])
|
||||
|
||||
pruned = await session.execute(
|
||||
delete(Route)
|
||||
.where(
|
||||
Route.version_group_id == version_group_id,
|
||||
Route.name.not_in(seed_names),
|
||||
)
|
||||
.returning(Route.id)
|
||||
)
|
||||
pruned_count = len(pruned.all())
|
||||
if pruned_count:
|
||||
print(f" Pruned {pruned_count} stale route(s)")
|
||||
|
||||
await session.flush()
|
||||
|
||||
# Return full mapping including children
|
||||
result = await session.execute(
|
||||
select(Route.name, Route.id).where(Route.version_group_id == version_group_id)
|
||||
@@ -233,8 +257,15 @@ async def upsert_route_encounters(
|
||||
encounters: list[dict],
|
||||
dex_to_id: dict[int, int],
|
||||
game_id: int,
|
||||
*,
|
||||
prune: bool = False,
|
||||
) -> int:
|
||||
"""Upsert encounters for a route and game, return count of upserted rows."""
|
||||
"""Upsert encounters for a route and game, return count of upserted rows.
|
||||
|
||||
When prune is True, deletes encounters not present in the seed data.
|
||||
"""
|
||||
seed_keys: set[tuple[int, str, str]] = set()
|
||||
|
||||
count = 0
|
||||
for enc in encounters:
|
||||
pokemon_id = dex_to_id.get(enc["pokeapi_id"])
|
||||
@@ -245,6 +276,7 @@ async def upsert_route_encounters(
|
||||
conditions = enc.get("conditions")
|
||||
if conditions:
|
||||
for condition_name, rate in conditions.items():
|
||||
seed_keys.add((pokemon_id, enc["method"], condition_name))
|
||||
await _upsert_single_encounter(
|
||||
session,
|
||||
route_id,
|
||||
@@ -258,6 +290,7 @@ async def upsert_route_encounters(
|
||||
)
|
||||
count += 1
|
||||
else:
|
||||
seed_keys.add((pokemon_id, enc["method"], ""))
|
||||
await _upsert_single_encounter(
|
||||
session,
|
||||
route_id,
|
||||
@@ -270,6 +303,23 @@ async def upsert_route_encounters(
|
||||
)
|
||||
count += 1
|
||||
|
||||
if prune:
|
||||
existing = await session.execute(
|
||||
select(RouteEncounter).where(
|
||||
RouteEncounter.route_id == route_id,
|
||||
RouteEncounter.game_id == game_id,
|
||||
)
|
||||
)
|
||||
stale_ids = [
|
||||
row.id
|
||||
for row in existing.scalars()
|
||||
if (row.pokemon_id, row.encounter_method, row.condition) not in seed_keys
|
||||
]
|
||||
if stale_ids:
|
||||
await session.execute(
|
||||
delete(RouteEncounter).where(RouteEncounter.id.in_(stale_ids))
|
||||
)
|
||||
|
||||
return count
|
||||
|
||||
|
||||
@@ -280,8 +330,13 @@ async def upsert_bosses(
|
||||
dex_to_id: dict[int, int],
|
||||
route_name_to_id: dict[str, int] | None = None,
|
||||
slug_to_game_id: dict[str, int] | None = None,
|
||||
*,
|
||||
prune: bool = False,
|
||||
) -> int:
|
||||
"""Upsert boss battles for a version group, return count of bosses upserted."""
|
||||
"""Upsert boss battles for a version group, return count of bosses upserted.
|
||||
|
||||
When prune is True, deletes boss battles not present in the seed data.
|
||||
"""
|
||||
count = 0
|
||||
for boss in bosses:
|
||||
# Resolve after_route_name to an ID
|
||||
@@ -364,6 +419,20 @@ async def upsert_bosses(
|
||||
|
||||
count += 1
|
||||
|
||||
if prune:
|
||||
seed_orders = {boss["order"] for boss in bosses}
|
||||
pruned = await session.execute(
|
||||
delete(BossBattle)
|
||||
.where(
|
||||
BossBattle.version_group_id == version_group_id,
|
||||
BossBattle.order.not_in(seed_orders),
|
||||
)
|
||||
.returning(BossBattle.id)
|
||||
)
|
||||
pruned_count = len(pruned.all())
|
||||
if pruned_count:
|
||||
print(f" Pruned {pruned_count} stale boss battle(s)")
|
||||
|
||||
await session.flush()
|
||||
return count
|
||||
|
||||
|
||||
Reference in New Issue
Block a user