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

data/create_model.py
74 строки2 KB
pozzi88
initial
13 июл 2025, 13:42
13 июл 2025, 13:421981fb0
import sys
from pathlib import Path
import configparser

from sqlalchemy import create_engine, text
from sqlalchemy.exc import SQLAlchemyError

from data.models import Base

SCHEMA = "englishcard"
CFG_PATH = Path(__file__).resolve().parent.parent / "cfg" / "config.ini"


def read_cfg() -> dict[str, str]:
    parser = configparser.ConfigParser()
    if not parser.read(CFG_PATH):
        sys.exit(f"Config {CFG_PATH} not found.")
    try:
        return parser["database"]
    except KeyError:
        sys.exit("Section [database] missing in config.")


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


def ensure_schema(conn):
    conn.execute(text(f"CREATE SCHEMA IF NOT EXISTS {SCHEMA};"))


def create_all(engine_):
    with engine_.begin() as conn:
        ensure_schema(conn)
    Base.metadata.create_all(engine_)


def drop_schema(engine_):
    with engine_.begin() as conn:
        conn.execute(text(f"DROP SCHEMA IF EXISTS {SCHEMA} CASCADE;"))


def main():
    choice = input("1 — создать таблицы, 2 — удалить схему: ").strip()
    if choice not in {"1", "2"}:
        sys.exit("Нужно ввести 1 или 2.")

    eng = engine(read_cfg())

    try:
        if choice == "1":
            create_all(eng)
            print("✅ Таблицы созданы.")
        else:
            drop_schema(eng)
            print("🗑  Схема удалена.")

            follow = input("Пересоздать схему (1) или продолжить (2)? ").strip()
            if follow == "1":
                create_all(eng)
                print("✅ Схема пересоздана.")
            elif follow != "2":
                print("Неверный выбор, продолжаю без пересоздания.")
    except SQLAlchemyError as exc:
        sys.exit(f"Ошибка SQLAlchemy: {exc}")


if __name__ == "__main__":
    main()