data/load_fixtures.py - master - pozzi88/sqlpy-diplom

data/load_fixtures.py
155 строк5 KB
pozzi88
initial
13 июл 2025, 13:42
13 июл 2025, 13:421981fb0

import json
import sys
from pathlib import Path
import configparser

from sqlalchemy import create_engine, select, func
from sqlalchemy.orm import sessionmaker

from data.models import (
    Base,
    User,
    EnglishWord,
    RussianWord,
    RelatedWord,
)

SCHEMA_NAME = "englishcard"


def read_config() -> dict[str, str]:
    cfg_path = Path(__file__).resolve().parent.parent / "cfg" / "config.ini"
    parser = configparser.ConfigParser()
    if not parser.read(cfg_path):
        sys.exit(f"Config file {cfg_path} not found.")
    try:
        return parser["database"]
    except KeyError:
        sys.exit("Section [database] not found in config.")


def make_engine(dbcfg: dict[str, str]):
    url = (
        "postgresql+psycopg2://"
        f"{dbcfg['user']}:{dbcfg['password']}"
        f"@{dbcfg['host']}:{dbcfg['port']}/{dbcfg['name']}"
    )
    return create_engine(url, future=True)


def get_or_create_user(session, name: str, score: int = 0) -> User:
    user = session.scalar(select(User).where(func.lower(User.name) == name.lower()))
    if user:
        return user
    user = User(name=name, score=score)
    session.add(user)
    session.flush() 
    return user


def get_or_create_eng(session, user_id: int, word: str) -> EnglishWord:
    obj = session.scalar(
        select(EnglishWord).where(func.lower(EnglishWord.word) == word.lower())
    )
    if obj:
        return obj
    obj = EnglishWord(word=word, created_by=user_id)
    session.add(obj)
    session.flush()
    return obj


def get_or_create_ru(session, user_id: int, word: str) -> RussianWord:
    obj = session.scalar(
        select(RussianWord).where(func.lower(RussianWord.word) == word.lower())
    )
    if obj:
        return obj
    obj = RussianWord(word=word, created_by=user_id)
    session.add(obj)
    session.flush()
    return obj


def load_fixtures(session, fixtures_path: Path) -> None:
    with fixtures_path.open(encoding="utf-8") as fp:
        data = json.load(fp)

    user_data = data["user"]
    default_user = get_or_create_user(
        session, name=user_data["name"], score=user_data.get("score", 0)
    )


    for pair in data["pairs"]:
        en_word = get_or_create_eng(session, default_user.id, pair["en"])
        ru_word = get_or_create_ru(session, default_user.id, pair["ru"])


        exists = session.scalar(
            select(RelatedWord).where(
                RelatedWord.englishword_id == en_word.id,
                RelatedWord.russianword_id == ru_word.id,
            )
        )
        if not exists:
            session.add(
                RelatedWord(
                    englishword_id=en_word.id,
                    russianword_id=ru_word.id,
                )
            )


    for word in data["single_en"]:
        get_or_create_eng(session, default_user.id, word)


    for word in data["single_ru"]:
        get_or_create_ru(session, default_user.id, word)


def print_stats(session):
    total_users = session.scalar(select(func.count()).select_from(User))
    total_en = session.scalar(select(func.count()).select_from(EnglishWord))
    total_ru = session.scalar(select(func.count()).select_from(RussianWord))
    total_pairs = session.scalar(select(func.count()).select_from(RelatedWord))

    print(f"\nПользователей:       {total_users}")
    print(f"Английских слов:     {total_en}")
    print(f"Русских слов:        {total_ru}")
    print(f"Связей EN↔RU:        {total_pairs}")


def main():
    fixtures_path = (
        Path(__file__).resolve().parent.parent / "fixtures" / "word_base.json"
    )
    if not fixtures_path.exists():
        sys.exit(f"{fixtures_path} not found.")

    engine = make_engine(read_config())
    Session = sessionmaker(engine, autoflush=False, expire_on_commit=False)

    with Session() as session:
        load_fixtures(session, fixtures_path)
        session.commit()
        print("\n✅ Данные из fixtures загружены в БД.")

        print_stats(session)
        answer = input("\nЖелаете ли проверить выборкой? (Y/N): ").strip().lower()
        if answer == "y":
            print("\nПример: первые 5 пар EN↔RU")
            rows = (
                session.query(EnglishWord.word, RussianWord.word)
                .join(RelatedWord, RelatedWord.englishword_id == EnglishWord.id)
                .join(RussianWord, RussianWord.id == RelatedWord.russianword_id)
                .limit(5)
                .all()
            )
            for en, ru in rows:
                print(f"{en} → {ru}")


if __name__ == "__main__":
    main()