seed vehicle trims catalog

app/db/seed.py

@@ -1,4 +1,5 @@
 import asyncio
+import argparse
 from datetime import date
 from decimal import Decimal
 
@@ -7,25 +8,25 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import selectinload
 
 from app.db.session import async_session_factory
-from app.models.car import Car, CarMake, CarModel
+from app.models.car import Car, CarMake, CarModel, CarTrim
 from app.models.expense import FuelEntry, ServiceEntry, ServiceType
 from app.models.user import User
-from app.services.catalog_data import CAR_CATALOG
+from app.services.catalog_data import CAR_CATALOG, CAR_TRIMS, COMMON_TRIMS, MAKE_COUNTRIES
 
 
 MOCK_PLATE_PREFIX = "MOCK"
 
 MOCK_CARS = [
-    ("KIA Sportage", "KIA", "Sportage", 2021, "gasoline", 36200, Decimal("2450000")),
-    ("Toyota Camry", "Toyota", "Camry", 2020, "gasoline", 58400, Decimal("2850000")),
-    ("Hyundai Tucson", "Hyundai", "Tucson", 2022, "gasoline", 27100, Decimal("2750000")),
-    ("Volkswagen Tiguan", "Volkswagen", "Tiguan", 2019, "gasoline", 73400, Decimal("2300000")),
-    ("BMW X3", "BMW", "X3", 2021, "diesel", 48900, Decimal("4350000")),
-    ("Mercedes GLC", "Mercedes", "GLC", 2020, "gasoline", 52200, Decimal("4500000")),
-    ("Nissan X-Trail", "Nissan", "X-Trail", 2018, "gasoline", 91400, Decimal("1850000")),
-    ("Skoda Octavia", "Skoda", "Octavia", 2021, "gasoline", 46800, Decimal("2050000")),
-    ("Tesla Model 3", "Tesla", "Model 3", 2022, "electric", 33800, Decimal("3900000")),
-    ("Haval Jolion", "Haval", "Jolion", 2023, "gasoline", 19600, Decimal("2150000")),
+    ("KIA Sportage", "KIA", "Sportage", "GT-Line 1.6T DCT AWD", 2021, "gasoline", 36200, Decimal("2450000")),
+    ("Toyota Camry", "Toyota", "Camry", "Comfort 2.5 AT", 2020, "gasoline", 58400, Decimal("2850000")),
+    ("Hyundai Tucson", "Hyundai", "Tucson", "Prestige 2.0 AT AWD", 2022, "gasoline", 27100, Decimal("2750000")),
+    ("Volkswagen Tiguan", "Volkswagen", "Tiguan", "Status 2.0 TSI DSG 4Motion", 2019, "gasoline", 73400, Decimal("2300000")),
+    ("BMW X3", "BMW", "X3", "20d xDrive", 2021, "diesel", 48900, Decimal("4350000")),
+    ("Mercedes GLC", "Mercedes", "GLC", "GLC 300 4MATIC AMG Line", 2020, "gasoline", 52200, Decimal("4500000")),
+    ("Nissan X-Trail", "Nissan", "X-Trail", "Premium", 2018, "gasoline", 91400, Decimal("1850000")),
+    ("Skoda Octavia", "Skoda", "Octavia", "Comfort", 2021, "gasoline", 46800, Decimal("2050000")),
+    ("Tesla Model 3", "Tesla", "Model 3", "Long Range AWD", 2022, "electric", 33800, Decimal("3900000")),
+    ("Haval Jolion", "Haval", "Jolion", "Premium", 2023, "gasoline", 19600, Decimal("2150000")),
 ]
 
 
@@ -43,15 +44,40 @@ async def seed_catalog(session: AsyncSession) -> None:
     for make_name, model_names in CAR_CATALOG.items():
         make = existing.get(make_name)
         if make is None:
-            make = CarMake(name=make_name)
+            make = CarMake(name=make_name, country=MAKE_COUNTRIES.get(make_name))
             session.add(make)
             await session.flush()
             existing_models = set()
         else:
+            if not make.country and MAKE_COUNTRIES.get(make_name):
+                make.country = MAKE_COUNTRIES[make_name]
             existing_models = {model.name for model in make.models}
         for model_name in model_names:
             if model_name not in existing_models:
                 session.add(CarModel(make_id=make.id, name=model_name))
+    await session.flush()
+    await seed_trims(session)
+
+
+async def seed_trims(session: AsyncSession) -> None:
+    result = await session.execute(
+        select(CarMake).options(selectinload(CarMake.models).selectinload(CarModel.trims))
+    )
+    makes = {make.name: make for make in result.scalars()}
+    for make in makes.values():
+        for model in make.models:
+            trim_rows = CAR_TRIMS.get((make.name, model.name)) or [
+                {**item, "body_type": infer_body_type(model.name)} for item in COMMON_TRIMS
+            ]
+            existing = {trim.name for trim in model.trims}
+            for row in trim_rows:
+                if row["name"] not in existing:
+                    session.add(CarTrim(model_id=model.id, **row))
+
+
+def infer_body_type(model_name: str) -> str:
+    suv_markers = ("Q", "X", "CX", "CR-V", "RAV", "Tiguan", "Tucson", "Sportage", "Jolion")
+    return "SUV" if any(marker in model_name for marker in suv_markers) else "sedan"
 
 
 async def pick_owner(session: AsyncSession) -> User:
@@ -82,12 +108,16 @@ async def seed_mock_usage(session: AsyncSession, owner: User) -> None:
     await clear_previous_mock(session)
     today = date.today()
 
-    for index, (name, make, model, year, fuel_type, start_odo, price) in enumerate(MOCK_CARS, start=1):
+    for index, (name, make, model, trim, year, fuel_type, start_odo, price) in enumerate(
+        MOCK_CARS,
+        start=1,
+    ):
         car = Car(
             owner_id=owner.id,
             name=name,
             make=make,
             model=model,
+            trim=trim,
             year=year,
             plate_number=f"{MOCK_PLATE_PREFIX}-{index:02d}",
             fuel_type=fuel_type,
@@ -201,9 +231,13 @@ async def seed_mock_usage(session: AsyncSession, owner: User) -> None:
         car.current_odometer = odometer
 
 
-async def main() -> None:
+async def main(catalog_only: bool = False) -> None:
     async with async_session_factory() as session:
         await seed_catalog(session)
+        if catalog_only:
+            await session.commit()
+            print("Seeded vehicle catalog")
+            return
         owner = await pick_owner(session)
         await seed_mock_usage(session, owner)
         await session.commit()
@@ -211,4 +245,7 @@ async def main() -> None:
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--catalog-only", action="store_true")
+    args = parser.parse_args()
+    asyncio.run(main(catalog_only=args.catalog_only))