guideCity/backend/seed/seed_loader.py

"""CLI to load seed YAML content into the database, upserting by slug.

Usage:
    python -m seed.seed_loader --cities seed/cities.yaml
    python -m seed.seed_loader --file seed/moscow_red_square.yaml
"""

import argparse
from pathlib import Path

import yaml

from app.db.session import SessionLocal
from app.models.city import City
from app.models.place import Place
from app.models.place_content import PlaceContent


def load_cities(path: Path) -> None:
    data = yaml.safe_load(path.read_text(encoding="utf-8"))
    db = SessionLocal()
    try:
        for c in data["cities"]:
            city = db.query(City).filter(City.slug == c["slug"]).one_or_none()
            if city is None:
                city = City(slug=c["slug"], name_ru=c["name_ru"], name_en=c["name_en"])
                db.add(city)
            else:
                city.name_ru = c["name_ru"]
                city.name_en = c["name_en"]
            center = c.get("center")
            if center is not None:
                city.center_point = f"SRID=4326;POINT({center['lon']} {center['lat']})"
            db.flush()
        db.commit()
        print(f"Loaded {len(data['cities'])} cities from {path.name}")
    finally:
        db.close()


def _upsert_place(db, city_id: int, place_data: dict) -> Place:
    place = db.query(Place).filter(Place.slug == place_data["slug"]).one_or_none()
    if place is None:
        place = Place(city_id=city_id, slug=place_data["slug"], category=place_data["category"])
        db.add(place)
    else:
        place.category = place_data["category"]

    place.city_id = city_id
    place.address = place_data.get("address")
    place.district = place_data.get("district")
    place.built_year = place_data.get("built_year")
    place.architect_builder = place_data.get("architect_builder")
    place.architectural_style = place_data.get("architectural_style")
    place.notable_people = place_data.get("notable_people")

    loc = place_data["location"]
    place.location = f"SRID=4326;POINT({loc['lon']} {loc['lat']})"

    db.flush()
    return place


def _upsert_content(db, place_id: int, language: str, length: str, title: str, body: str) -> None:
    content = (
        db.query(PlaceContent)
        .filter(
            PlaceContent.place_id == place_id,
            PlaceContent.language == language,
            PlaceContent.length == length,
        )
        .one_or_none()
    )
    if content is None:
        db.add(PlaceContent(place_id=place_id, language=language, length=length, title=title, body=body))
    else:
        content.title = title
        content.body = body


def load_places(path: Path) -> None:
    data = yaml.safe_load(path.read_text(encoding="utf-8"))
    city_slug = data["city"]

    db = SessionLocal()
    try:
        city = db.query(City).filter(City.slug == city_slug).one_or_none()
        if city is None:
            raise SystemExit(
                f"City '{city_slug}' not found. Run `--cities seed/cities.yaml` first."
            )

        for place_data in data["places"]:
            place = _upsert_place(db, city.id, place_data)
            titles = place_data.get("title", {})
            for language, lengths in place_data["content"].items():
                title = titles.get(language, place_data["slug"])
                for length, body in lengths.items():
                    _upsert_content(db, place.id, language, length, title, body.strip())

        db.commit()
        print(f"Loaded {len(data['places'])} places from {path.name}")
    finally:
        db.close()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--cities", help="Path to a cities YAML file")
    group.add_argument("--file", help="Path to a places YAML file")
    args = parser.parse_args()

    if args.cities:
        load_cities(Path(args.cities))
    else:
        load_places(Path(args.file))


if __name__ == "__main__":
    main()