database/db_session.py - master - REDOSSSSS/tg-bot-dental

database/db_session.py
100 строк3 KB
iwan
done project
04 июн 2026, 00:09
04 июн 2026, 00:0964f2830
from __future__ import annotations

import os
from datetime import time

from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker

from database.models import AdminSettings, Base, Contact, Service, User, UserRole

_engine = None
_SessionLocal = None


def get_database_url() -> str:
    return os.getenv("DATABASE_URL", "sqlite:///dental_bot.db")


def get_engine(database_url: str | None = None):
    global _engine, _SessionLocal
    url = database_url or get_database_url()
    if _engine is None or str(_engine.url) != url:
        connect_args = {}
        if url.startswith("sqlite"):
            connect_args["check_same_thread"] = False
        _engine = create_engine(url, connect_args=connect_args)
        _SessionLocal = sessionmaker(
            bind=_engine, autoflush=False, autocommit=False, expire_on_commit=False
        )
    return _engine


def get_session_factory(database_url: str | None = None):
    get_engine(database_url)
    return _SessionLocal


def init_db(database_url: str | None = None) -> None:
    engine = get_engine(database_url)
    Base.metadata.create_all(bind=engine)


def seed_db(session: Session, default_admin_id: int | None = None) -> None:
    if session.query(Contact).count() == 0:
        session.add(
            Contact(
                phone="+7 (000) 000-00-00",
                address="г. Москва, ул. Стоматологическая, 10",
                work_schedule_text="Пн–Пт: 9:00–20:00, Сб: 10:00–16:00",
            )
        )

    if session.query(AdminSettings).count() == 0:
        session.add(
            AdminSettings(
                work_start=time(10, 0),
                work_end=time(15, 0),
                slot_interval_minutes=30,
            )
        )
    else:
        settings = session.query(AdminSettings).first()
        if settings and settings.slot_interval_minutes != 30:
            settings.slot_interval_minutes = 30

    if session.query(Service).count() == 0:
        session.add_all(
            [
                Service(
                    name="Консультация стоматолога",
                    description="Осмотр, диагностика, план лечения",
                    price_rub=1500,
                ),
                Service(
                    name="Профессиональная гигиена",
                    description="Удаление налёта, полировка",
                    price_rub=4500,
                ),
                Service(
                    name="Лечение кариеса",
                    description="Пломбирование одного зуба",
                    price_rub=5500,
                ),
            ]
        )

    if default_admin_id:
        user = session.query(User).filter(User.telegram_id == default_admin_id).first()
        if user is None:
            session.add(
                User(
                    telegram_id=default_admin_id,
                    role=UserRole.admin,
                    is_subscribed=False,
                )
            )
        elif user.role != UserRole.admin:
            user.role = UserRole.admin

    session.commit()