Julian Tabel
053dece33e
Add seeder support for boss battles so new database instances come pre-populated. Adds --export-bosses CLI flag to dump boss data from the database to JSON seed files, and loads those files during normal seeding. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
296 lines
9.6 KiB
Python
296 lines
9.6 KiB
Python
"""Database upsert helpers for seed data."""
|
|
|
|
from sqlalchemy import delete, select
|
|
from sqlalchemy.dialects.postgresql import insert
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.models.boss_battle import BossBattle
|
|
from app.models.boss_pokemon import BossPokemon
|
|
from app.models.evolution import Evolution
|
|
from app.models.game import Game
|
|
from app.models.pokemon import Pokemon
|
|
from app.models.route import Route
|
|
from app.models.route_encounter import RouteEncounter
|
|
from app.models.version_group import VersionGroup
|
|
|
|
|
|
async def upsert_version_groups(
|
|
session: AsyncSession,
|
|
vg_data: dict[str, dict],
|
|
) -> dict[str, int]:
|
|
"""Upsert version group records, return {slug: id} mapping."""
|
|
for vg_slug, vg_info in vg_data.items():
|
|
vg_name = " / ".join(
|
|
g["name"].replace("Pokemon ", "")
|
|
for g in vg_info["games"].values()
|
|
)
|
|
stmt = insert(VersionGroup).values(
|
|
name=vg_name,
|
|
slug=vg_slug,
|
|
).on_conflict_do_update(
|
|
index_elements=["slug"],
|
|
set_={"name": vg_name},
|
|
)
|
|
await session.execute(stmt)
|
|
|
|
await session.flush()
|
|
|
|
result = await session.execute(select(VersionGroup.slug, VersionGroup.id))
|
|
return {row.slug: row.id for row in result}
|
|
|
|
|
|
async def upsert_games(
|
|
session: AsyncSession,
|
|
games: list[dict],
|
|
slug_to_vg_id: dict[str, int] | None = None,
|
|
) -> dict[str, int]:
|
|
"""Upsert game records, return {slug: id} mapping."""
|
|
for game in games:
|
|
values = {
|
|
"name": game["name"],
|
|
"slug": game["slug"],
|
|
"generation": game["generation"],
|
|
"region": game["region"],
|
|
"release_year": game.get("release_year"),
|
|
"color": game.get("color"),
|
|
}
|
|
update_set = {
|
|
"name": game["name"],
|
|
"generation": game["generation"],
|
|
"region": game["region"],
|
|
"release_year": game.get("release_year"),
|
|
"color": game.get("color"),
|
|
}
|
|
if slug_to_vg_id is not None:
|
|
vg_id = slug_to_vg_id.get(game["slug"])
|
|
if vg_id is not None:
|
|
values["version_group_id"] = vg_id
|
|
update_set["version_group_id"] = vg_id
|
|
|
|
stmt = insert(Game).values(**values).on_conflict_do_update(
|
|
index_elements=["slug"],
|
|
set_=update_set,
|
|
)
|
|
await session.execute(stmt)
|
|
|
|
await session.flush()
|
|
|
|
result = await session.execute(select(Game.slug, Game.id))
|
|
return {row.slug: row.id for row in result}
|
|
|
|
|
|
async def upsert_pokemon(session: AsyncSession, pokemon_list: list[dict]) -> dict[int, int]:
|
|
"""Upsert pokemon records, return {pokeapi_id: id} mapping."""
|
|
for poke in pokemon_list:
|
|
stmt = insert(Pokemon).values(
|
|
pokeapi_id=poke["pokeapi_id"],
|
|
national_dex=poke["national_dex"],
|
|
name=poke["name"],
|
|
types=poke["types"],
|
|
sprite_url=poke.get("sprite_url"),
|
|
).on_conflict_do_update(
|
|
index_elements=["pokeapi_id"],
|
|
set_={
|
|
"national_dex": poke["national_dex"],
|
|
"name": poke["name"],
|
|
"types": poke["types"],
|
|
"sprite_url": poke.get("sprite_url"),
|
|
},
|
|
)
|
|
await session.execute(stmt)
|
|
|
|
await session.flush()
|
|
|
|
result = await session.execute(select(Pokemon.pokeapi_id, Pokemon.id))
|
|
return {row.pokeapi_id: row.id for row in result}
|
|
|
|
|
|
async def upsert_routes(
|
|
session: AsyncSession,
|
|
version_group_id: int,
|
|
routes: list[dict],
|
|
) -> 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.
|
|
"""
|
|
# First pass: upsert all parent routes (without parent_route_id)
|
|
for route in routes:
|
|
stmt = insert(Route).values(
|
|
name=route["name"],
|
|
version_group_id=version_group_id,
|
|
order=route["order"],
|
|
parent_route_id=None, # Parent routes have no parent
|
|
).on_conflict_do_update(
|
|
constraint="uq_routes_version_group_name",
|
|
set_={"order": route["order"], "parent_route_id": None},
|
|
)
|
|
await session.execute(stmt)
|
|
|
|
await session.flush()
|
|
|
|
# Get mapping of parent routes
|
|
result = await session.execute(
|
|
select(Route.name, Route.id).where(Route.version_group_id == version_group_id)
|
|
)
|
|
name_to_id = {row.name: row.id for row in result}
|
|
|
|
# Second pass: upsert child routes with parent_route_id
|
|
for route in routes:
|
|
children = route.get("children", [])
|
|
if not children:
|
|
continue
|
|
|
|
parent_id = name_to_id[route["name"]]
|
|
for child in children:
|
|
stmt = insert(Route).values(
|
|
name=child["name"],
|
|
version_group_id=version_group_id,
|
|
order=child["order"],
|
|
parent_route_id=parent_id,
|
|
pinwheel_zone=child.get("pinwheel_zone"),
|
|
).on_conflict_do_update(
|
|
constraint="uq_routes_version_group_name",
|
|
set_={
|
|
"order": child["order"],
|
|
"parent_route_id": parent_id,
|
|
"pinwheel_zone": child.get("pinwheel_zone"),
|
|
},
|
|
)
|
|
await session.execute(stmt)
|
|
|
|
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)
|
|
)
|
|
return {row.name: row.id for row in result}
|
|
|
|
|
|
async def upsert_route_encounters(
|
|
session: AsyncSession,
|
|
route_id: int,
|
|
encounters: list[dict],
|
|
dex_to_id: dict[int, int],
|
|
game_id: int,
|
|
) -> int:
|
|
"""Upsert encounters for a route and game, return count of upserted rows."""
|
|
count = 0
|
|
for enc in encounters:
|
|
pokemon_id = dex_to_id.get(enc["pokeapi_id"])
|
|
if pokemon_id is None:
|
|
print(f" Warning: no pokemon_id for pokeapi_id {enc['pokeapi_id']}")
|
|
continue
|
|
|
|
stmt = insert(RouteEncounter).values(
|
|
route_id=route_id,
|
|
pokemon_id=pokemon_id,
|
|
game_id=game_id,
|
|
encounter_method=enc["method"],
|
|
encounter_rate=enc["encounter_rate"],
|
|
min_level=enc["min_level"],
|
|
max_level=enc["max_level"],
|
|
).on_conflict_do_update(
|
|
constraint="uq_route_pokemon_method_game",
|
|
set_={
|
|
"encounter_rate": enc["encounter_rate"],
|
|
"min_level": enc["min_level"],
|
|
"max_level": enc["max_level"],
|
|
},
|
|
)
|
|
await session.execute(stmt)
|
|
count += 1
|
|
|
|
return count
|
|
|
|
|
|
async def upsert_bosses(
|
|
session: AsyncSession,
|
|
version_group_id: int,
|
|
bosses: list[dict],
|
|
dex_to_id: dict[int, int],
|
|
) -> int:
|
|
"""Upsert boss battles for a version group, return count of bosses upserted."""
|
|
count = 0
|
|
for boss in bosses:
|
|
# Upsert the boss battle on (version_group_id, order) conflict
|
|
stmt = insert(BossBattle).values(
|
|
version_group_id=version_group_id,
|
|
name=boss["name"],
|
|
boss_type=boss["boss_type"],
|
|
badge_name=boss.get("badge_name"),
|
|
badge_image_url=boss.get("badge_image_url"),
|
|
level_cap=boss["level_cap"],
|
|
order=boss["order"],
|
|
location=boss["location"],
|
|
sprite_url=boss.get("sprite_url"),
|
|
).on_conflict_do_update(
|
|
constraint="uq_boss_battles_version_group_order",
|
|
set_={
|
|
"name": boss["name"],
|
|
"boss_type": boss["boss_type"],
|
|
"badge_name": boss.get("badge_name"),
|
|
"badge_image_url": boss.get("badge_image_url"),
|
|
"level_cap": boss["level_cap"],
|
|
"location": boss["location"],
|
|
"sprite_url": boss.get("sprite_url"),
|
|
},
|
|
).returning(BossBattle.id)
|
|
result = await session.execute(stmt)
|
|
boss_id = result.scalar_one()
|
|
|
|
# Delete existing boss_pokemon for this boss, then re-insert
|
|
await session.execute(
|
|
delete(BossPokemon).where(BossPokemon.boss_battle_id == boss_id)
|
|
)
|
|
for bp in boss.get("pokemon", []):
|
|
pokemon_id = dex_to_id.get(bp["pokeapi_id"])
|
|
if pokemon_id is None:
|
|
print(f" Warning: no pokemon_id for pokeapi_id {bp['pokeapi_id']}")
|
|
continue
|
|
session.add(BossPokemon(
|
|
boss_battle_id=boss_id,
|
|
pokemon_id=pokemon_id,
|
|
level=bp["level"],
|
|
order=bp["order"],
|
|
))
|
|
|
|
count += 1
|
|
|
|
await session.flush()
|
|
return count
|
|
|
|
|
|
async def upsert_evolutions(
|
|
session: AsyncSession,
|
|
evolutions: list[dict],
|
|
dex_to_id: dict[int, int],
|
|
) -> int:
|
|
"""Upsert evolution pairs, return count of upserted rows."""
|
|
await session.execute(delete(Evolution))
|
|
|
|
count = 0
|
|
for evo in evolutions:
|
|
from_id = dex_to_id.get(evo["from_pokeapi_id"])
|
|
to_id = dex_to_id.get(evo["to_pokeapi_id"])
|
|
if from_id is None or to_id is None:
|
|
continue
|
|
|
|
evolution = Evolution(
|
|
from_pokemon_id=from_id,
|
|
to_pokemon_id=to_id,
|
|
trigger=evo["trigger"],
|
|
min_level=evo.get("min_level"),
|
|
item=evo.get("item"),
|
|
held_item=evo.get("held_item"),
|
|
condition=evo.get("condition"),
|
|
region=evo.get("region"),
|
|
)
|
|
session.add(evolution)
|
|
count += 1
|
|
|
|
await session.flush()
|
|
return count
|