csv_loader.py - master - BigYahu/dododo

csv_loader.py
55 строк3 KB
BigYahu
create: config.py, db_models.py, migrate.py, csv_loader.py, pytest.py
14 май 2026, 10:18
14 май 2026, 10:1857212ab
# csv_loader.py
import pandas as pd
from sqlalchemy import select
from sqlalchemy.orm import Session
from db_models import get_engine, get_session_factory, Product

def load_csv_to_db(csv_path: str, db_uri: str = None, chunk_size: int = 1000):
    engine = get_engine(db_uri)
    SessionLocal = get_session_factory(engine)

    # Ожидаем, что CSV имеет колонки: sku,name,description,price,currency,stock,category
    required_cols = {"sku", "name", "price"}
    df = pd.read_csv(csv_path)

    if not required_cols.issubset(set(df.columns)):
        missing = required_cols - set(df.columns)
        raise ValueError(f"CSV пропускает обязательные столбцы: {missing}")

    with SessionLocal() as session:
        # транзакция на всю загрузку
        with session.begin():
            for idx, row in df.iterrows():
                sku = str(row["sku"]).strip()
                name = str(row["name"]).strip()
                price = float(row["price"])
                description = str(row.get("description", "")) if pd.notna(row.get("description")) else None
                currency = str(row.get("currency", "USD")) if pd.notna(row.get("currency", "USD")) else "USD"
                stock = int(row.get("stock", 0)) if pd.notna(row.get("stock", 0)) else 0
                category = str(row.get("category")) if "category" in df.columns and pd.notna(row.get("category")) else None

                # попытка найти существующий товар по sku
                stmt = select(Product).where(Product.sku == sku)
                existing = session.execute(stmt).scalar_one_or_none()

                if existing:
                    # апдейт существующего
                    existing.name = name
                    existing.description = description
                    existing.price = price
                    existing.currency = currency
                    existing.stock = stock
                    existing.category = category
                else:
                    p = Product(
                        sku=sku,
                        name=name,
                        description=description,
                        price=price,
                        currency=currency,
                        stock=stock,
                        category=category
                    )
                    session.add(p)
        # коммит осуществляется контекстом with session.begin()
    print("CSV успешно загружен в БД.")