first commit
This commit is contained in:
214
app/db/seed.py
Normal file
214
app/db/seed.py
Normal file
@@ -0,0 +1,214 @@
|
||||
import asyncio
|
||||
from datetime import date
|
||||
from decimal import Decimal
|
||||
|
||||
from sqlalchemy import delete, select
|
||||
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.expense import FuelEntry, ServiceEntry, ServiceType
|
||||
from app.models.user import User
|
||||
from app.services.catalog_data import CAR_CATALOG
|
||||
|
||||
|
||||
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")),
|
||||
]
|
||||
|
||||
|
||||
def month_shift(base: date, months_back: int) -> date:
|
||||
month_index = base.year * 12 + base.month - 1 - months_back
|
||||
year = month_index // 12
|
||||
month = month_index % 12 + 1
|
||||
return date(year, month, min(base.day, 24))
|
||||
|
||||
|
||||
async def seed_catalog(session: AsyncSession) -> None:
|
||||
result = await session.execute(select(CarMake).options(selectinload(CarMake.models)))
|
||||
existing = {make.name: make for make in result.scalars()}
|
||||
|
||||
for make_name, model_names in CAR_CATALOG.items():
|
||||
make = existing.get(make_name)
|
||||
if make is None:
|
||||
make = CarMake(name=make_name)
|
||||
session.add(make)
|
||||
await session.flush()
|
||||
existing_models = set()
|
||||
else:
|
||||
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))
|
||||
|
||||
|
||||
async def pick_owner(session: AsyncSession) -> User:
|
||||
result = await session.execute(select(User).where(User.telegram_id == 1))
|
||||
user = result.scalar_one_or_none()
|
||||
if user:
|
||||
return user
|
||||
|
||||
user = User(telegram_id=1, username="demo", first_name="Demo")
|
||||
session.add(user)
|
||||
await session.flush()
|
||||
return user
|
||||
|
||||
|
||||
async def clear_previous_mock(session: AsyncSession) -> None:
|
||||
result = await session.execute(
|
||||
select(Car.id).where(Car.plate_number.like(f"{MOCK_PLATE_PREFIX}-%"))
|
||||
)
|
||||
car_ids = list(result.scalars())
|
||||
if not car_ids:
|
||||
return
|
||||
await session.execute(delete(ServiceEntry).where(ServiceEntry.car_id.in_(car_ids)))
|
||||
await session.execute(delete(FuelEntry).where(FuelEntry.car_id.in_(car_ids)))
|
||||
await session.execute(delete(Car).where(Car.id.in_(car_ids)))
|
||||
|
||||
|
||||
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):
|
||||
car = Car(
|
||||
owner_id=owner.id,
|
||||
name=name,
|
||||
make=make,
|
||||
model=model,
|
||||
year=year,
|
||||
plate_number=f"{MOCK_PLATE_PREFIX}-{index:02d}",
|
||||
fuel_type=fuel_type,
|
||||
purchase_date=date(year, min(index, 12), 10),
|
||||
purchase_price=price,
|
||||
current_odometer=start_odo,
|
||||
)
|
||||
session.add(car)
|
||||
await session.flush()
|
||||
|
||||
odometer = start_odo
|
||||
monthly_km = 820 + index * 95
|
||||
consumption = Decimal("7.2") + Decimal(index % 5) * Decimal("0.45")
|
||||
if fuel_type == "electric":
|
||||
consumption = Decimal("0")
|
||||
|
||||
for months_back in range(11, -1, -1):
|
||||
base_day = month_shift(today.replace(day=15), months_back)
|
||||
odometer += monthly_km
|
||||
|
||||
if fuel_type == "electric":
|
||||
energy_cost = Decimal(110 + index * 8 + months_back % 3 * 15)
|
||||
session.add(
|
||||
ServiceEntry(
|
||||
car_id=car.id,
|
||||
entry_date=base_day,
|
||||
odometer=odometer,
|
||||
service_type=ServiceType.other,
|
||||
title="Зарядка и парковка",
|
||||
category="charging",
|
||||
vendor="EV network",
|
||||
total_cost=energy_cost,
|
||||
)
|
||||
)
|
||||
else:
|
||||
liters_per_month = Decimal(monthly_km) * consumption / Decimal(100)
|
||||
price_per_liter = Decimal("58.50") + Decimal(index % 4) * Decimal("2.10")
|
||||
for fill in range(2):
|
||||
fill_date = date(base_day.year, base_day.month, 8 if fill == 0 else 22)
|
||||
liters = (liters_per_month / 2).quantize(Decimal("0.001"))
|
||||
session.add(
|
||||
FuelEntry(
|
||||
car_id=car.id,
|
||||
entry_date=fill_date,
|
||||
odometer=odometer - (monthly_km // 2 if fill == 0 else 0),
|
||||
liters=liters,
|
||||
price_per_liter=price_per_liter,
|
||||
total_cost=(liters * price_per_liter).quantize(Decimal("0.01")),
|
||||
station=["Shell", "Lukoil", "Gazprom", "Rosneft", "Neste"][index % 5],
|
||||
is_full_tank=True,
|
||||
)
|
||||
)
|
||||
|
||||
if months_back in {11, 5}:
|
||||
session.add(
|
||||
ServiceEntry(
|
||||
car_id=car.id,
|
||||
entry_date=date(base_day.year, base_day.month, 12),
|
||||
odometer=odometer,
|
||||
service_type=ServiceType.maintenance,
|
||||
title="Плановое ТО",
|
||||
category="regular",
|
||||
vendor="Service center",
|
||||
total_cost=Decimal(12500 + index * 950),
|
||||
next_due_odometer=odometer + 10000,
|
||||
)
|
||||
)
|
||||
|
||||
if months_back in {8, 2}:
|
||||
session.add(
|
||||
ServiceEntry(
|
||||
car_id=car.id,
|
||||
entry_date=date(base_day.year, base_day.month, 18),
|
||||
odometer=odometer,
|
||||
service_type=ServiceType.tire,
|
||||
title="Сезонная замена шин",
|
||||
category="tires",
|
||||
vendor="Tire shop",
|
||||
total_cost=Decimal(4200 + index * 180),
|
||||
)
|
||||
)
|
||||
|
||||
if months_back == 6:
|
||||
session.add(
|
||||
ServiceEntry(
|
||||
car_id=car.id,
|
||||
entry_date=date(base_day.year, base_day.month, 20),
|
||||
odometer=odometer,
|
||||
service_type=ServiceType.insurance,
|
||||
title="ОСАГО / страховка",
|
||||
category="insurance",
|
||||
vendor="Insurance",
|
||||
total_cost=Decimal(18200 + index * 1150),
|
||||
)
|
||||
)
|
||||
|
||||
if months_back == index % 10:
|
||||
session.add(
|
||||
ServiceEntry(
|
||||
car_id=car.id,
|
||||
entry_date=date(base_day.year, base_day.month, 25),
|
||||
odometer=odometer,
|
||||
service_type=ServiceType.repair,
|
||||
title=["Тормозные колодки", "Диагностика подвески", "Замена АКБ", "Ремонт кондиционера"][index % 4],
|
||||
category="repair",
|
||||
vendor="Garage",
|
||||
total_cost=Decimal(7200 + index * 1350),
|
||||
)
|
||||
)
|
||||
|
||||
car.current_odometer = odometer
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
async with async_session_factory() as session:
|
||||
await seed_catalog(session)
|
||||
owner = await pick_owner(session)
|
||||
await seed_mock_usage(session, owner)
|
||||
await session.commit()
|
||||
print(f"Seeded catalog and mock usage for owner_id={owner.id}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user