invix_test_task/api/app/bootstrap.py

79 lines
3.3 KiB
Python

import logging
from decimal import Decimal
from sqlalchemy import func, select, text
from app.db import SessionFactory, engine
from app.models import Base, Category, Product
from app.storage import ensure_bucket, object_exists, put_object
log = logging.getLogger("shop")
SEED = [
("Электроника", "Смартфон", "Карманный телефон для демо-каталога.", Decimal("50000.00"), "#1d4ed8", True),
("Электроника", "Наушники", "Беспроводные наушники.", Decimal("5000.00"), "#0f766e", True),
("Электроника", "Чехол", "Скрыт из каталога: is_active = false.", Decimal("900.00"), "#444444", False),
("Одежда", "Футболка", "Хлопковая футболка.", Decimal("1500.00"), "#b45309", True),
("Одежда", "Джинсы", "Прямые джинсы.", Decimal("4000.00"), "#1e3a8a", True),
]
def _svg(title: str, color: str) -> bytes:
return (
'<svg xmlns="http://www.w3.org/2000/svg" width="640" height="360">'
f'<rect width="100%" height="100%" fill="{color}"/>'
'<text x="50%" y="54%" fill="#ffffff" font-size="36" text-anchor="middle" '
f'font-family="sans-serif">{title}</text></svg>'
).encode()
async def init_db() -> None:
async with engine.connect() as lock_conn:
await lock_conn.execute(text("SELECT pg_advisory_lock(1, 42002)"))
await lock_conn.commit()
try:
async with engine.begin() as ddl:
await ddl.run_sync(Base.metadata.create_all)
await _seed()
finally:
await lock_conn.execute(text("SELECT pg_advisory_unlock(1, 42002)"))
await lock_conn.commit()
async def _seed() -> None:
await ensure_bucket()
async with SessionFactory() as session:
count = await session.scalar(select(func.count()).select_from(Product))
if count == 0:
categories: dict[str, Category] = {}
for name, title, description, price, color, active in SEED:
category = categories.get(name)
if category is None:
category = Category(name=name)
session.add(category)
await session.flush()
categories[name] = category
product = Product(
category_id=category.id,
title=title,
description=description,
price=price,
is_active=active,
)
session.add(product)
await session.flush()
product.image_key = f"products/{product.id}.svg"
await session.commit()
log.info("seeded catalog")
products = (await session.scalars(select(Product))).all()
colors = {title: color for _, title, _, _, color, _ in SEED}
for product in products:
if not product.image_key:
product.image_key = f"products/{product.id}.svg"
if await object_exists(product.image_key):
continue
color = colors.get(product.title, "#333333")
await put_object(product.image_key, _svg(product.title, color), "image/svg+xml")
await session.commit()