database/crud.py - master - rnj/SuaiTestBot

database/crud.py
492 строки15 KB
Dmitriy Pozdeev
init commit
03 июн 2026, 23:58
03 июн 2026, 23:58110bad2
"""
CRUD operations for the schedule bot (15+ queries).
"""
from dataclasses import dataclass
from typing import List, Optional

from sqlalchemy import delete, select
from sqlalchemy.orm import joinedload

from database.database import Session
from model.discipline import Discipline
from model.group import Group
from model.lesson import Lesson
from model.lesson_group import LessonGroup
from model.teacher import Teacher
from model.user import User


@dataclass
class LessonInfo:
    lesson: Lesson
    teacher_name: str
    discipline_name: str


def _lesson_infos(rows) -> List[LessonInfo]:
    return [
        LessonInfo(
            lesson=lesson,
            teacher_name=teacher.full_name,
            discipline_name=discipline.name,
        )
        for lesson, teacher, discipline in rows
    ]


def _base_lesson_query():
    return (
        select(Lesson, Teacher, Discipline)
        .join(Teacher, Lesson.teacher_id == Teacher.id)
        .join(Discipline, Lesson.discipline_id == Discipline.id)
    )


def _order_lessons(query):
    return query.order_by(
        Lesson.week_type,
        Lesson.day_of_week,
        Lesson.slot_number,
    )


# ---------- Users ----------
def add_user(telegram_id: int, username: Optional[str] = None, is_admin: bool = False) -> User:
    with Session() as session:
        user = User(telegram_id=telegram_id, username=username, is_admin=is_admin)
        session.add(user)
        session.commit()
        session.refresh(user)
        return user


def get_user_by_telegram_id(telegram_id: int) -> Optional[User]:
    with Session() as session:
        return session.execute(
            select(User).where(User.telegram_id == telegram_id)
        ).scalar_one_or_none()


def user_exists(telegram_id: int) -> bool:
    with Session() as session:
        return (
            session.execute(select(User).where(User.telegram_id == telegram_id)).first()
            is not None
        )


def is_admin(telegram_id: int) -> bool:
    with Session() as session:
        user = session.execute(
            select(User).where(User.telegram_id == telegram_id)
        ).scalar_one_or_none()
        return user is not None and user.is_admin


def set_subscription(telegram_id: int, subscribed: bool) -> bool:
    with Session() as session:
        user = session.execute(
            select(User).where(User.telegram_id == telegram_id)
        ).scalar_one_or_none()
        if not user:
            return False
        user.is_subscribed = subscribed
        session.commit()
        return True


def get_subscribed_users() -> List[User]:
    with Session() as session:
        return session.execute(
            select(User).where(User.is_subscribed == True)  # noqa: E712
        ).scalars().all()


def set_user_group(telegram_id: int, group_id: int) -> bool:
    with Session() as session:
        user = session.execute(
            select(User).where(User.telegram_id == telegram_id)
        ).scalar_one_or_none()
        if not user:
            return False
        user.group_id = group_id
        session.commit()
        return True


def list_admins() -> List[User]:
    with Session() as session:
        return session.execute(select(User).where(User.is_admin == True)).scalars().all()  # noqa: E712


def promote_admin(telegram_id: int) -> Optional[User]:
    with Session() as session:
        user = session.execute(
            select(User).where(User.telegram_id == telegram_id)
        ).scalar_one_or_none()
        if not user:
            user = User(telegram_id=telegram_id, is_admin=True)
            session.add(user)
        else:
            user.is_admin = True
        session.commit()
        session.refresh(user)
        return user


# ---------- Groups ----------
def get_all_groups() -> List[Group]:
    with Session() as session:
        return session.execute(select(Group).order_by(Group.course, Group.number)).scalars().all()


def get_group_by_id(group_id: int) -> Optional[Group]:
    with Session() as session:
        return session.get(Group, group_id)


# ---------- Teachers ----------
def create_teacher(full_name: str, phone: str, email: str) -> Teacher:
    with Session() as session:
        teacher = Teacher(full_name=full_name, phone=phone, email=email)
        session.add(teacher)
        session.commit()
        session.refresh(teacher)
        return teacher


def get_all_teachers() -> List[Teacher]:
    with Session() as session:
        return session.execute(select(Teacher).order_by(Teacher.full_name)).scalars().all()


def get_teacher_by_id(teacher_id: int) -> Optional[Teacher]:
    with Session() as session:
        return session.get(Teacher, teacher_id)


def update_teacher(teacher_id: int, full_name: str, phone: str, email: str) -> Optional[Teacher]:
    with Session() as session:
        teacher = session.get(Teacher, teacher_id)
        if not teacher:
            return None
        teacher.full_name = full_name
        teacher.phone = phone
        teacher.email = email
        session.commit()
        session.refresh(teacher)
        return teacher


def delete_teacher(teacher_id: int) -> bool:
    with Session() as session:
        teacher = session.get(Teacher, teacher_id)
        if not teacher:
            return False
        session.delete(teacher)
        session.commit()
        return True


# ---------- Disciplines ----------
def add_discipline(name: str, short_name: str, hours: int) -> Discipline:
    with Session() as session:
        discipline = Discipline(name=name, short_name=short_name, hours=hours)
        session.add(discipline)
        session.commit()
        session.refresh(discipline)
        return discipline


def get_discipline_by_id(discipline_id: int) -> Optional[Discipline]:
    with Session() as session:
        return session.get(Discipline, discipline_id)


def get_all_disciplines() -> List[Discipline]:
    with Session() as session:
        return session.execute(select(Discipline).order_by(Discipline.name)).scalars().all()


def update_discipline(
    discipline_id: int, name: str, short_name: str, hours: int
) -> Optional[Discipline]:
    with Session() as session:
        discipline = session.get(Discipline, discipline_id)
        if not discipline:
            return None
        discipline.name = name
        discipline.short_name = short_name
        discipline.hours = hours
        session.commit()
        session.refresh(discipline)
        return discipline


def delete_discipline(discipline_id: int) -> bool:
    with Session() as session:
        discipline = session.get(Discipline, discipline_id)
        if not discipline:
            return False
        session.delete(discipline)
        session.commit()
        return True


# ---------- Lessons (недельное расписание) ----------
def _group_has_slot_conflict(
    session, group_ids: List[int], week_type: str, day_of_week: int, slot_number: int,
    exclude_lesson_id: Optional[int] = None,
) -> bool:
    q = (
        select(Lesson.id)
        .join(LessonGroup, Lesson.id == LessonGroup.lesson_id)
        .where(
            LessonGroup.group_id.in_(group_ids),
            Lesson.week_type == week_type,
            Lesson.day_of_week == day_of_week,
            Lesson.slot_number == slot_number,
        )
    )
    if exclude_lesson_id:
        q = q.where(Lesson.id != exclude_lesson_id)
    return session.execute(q).first() is not None


def create_lesson(
    week_type: str,
    day_of_week: int,
    slot_number: int,
    room: str,
    lesson_type: str,
    discipline_id: int,
    teacher_id: int,
    group_ids: List[int],
) -> Lesson:
    with Session() as session:
        if _group_has_slot_conflict(session, group_ids, week_type, day_of_week, slot_number):
            raise ValueError("slot_conflict")
        lesson = Lesson(
            week_type=week_type,
            day_of_week=day_of_week,
            slot_number=slot_number,
            room=room,
            lesson_type=lesson_type,
            discipline_id=discipline_id,
            teacher_id=teacher_id,
        )
        session.add(lesson)
        session.flush()
        for gid in group_ids:
            session.add(LessonGroup(lesson_id=lesson.id, group_id=gid))
        session.commit()
        session.refresh(lesson)
        return lesson


def get_lesson_by_id(lesson_id: int) -> Optional[Lesson]:
    with Session() as session:
        return session.execute(
            select(Lesson)
            .options(joinedload(Lesson.group_links))
            .where(Lesson.id == lesson_id)
        ).unique().scalar_one_or_none()


def get_all_lessons() -> List[Lesson]:
    with Session() as session:
        return session.execute(
            _order_lessons(select(Lesson))
        ).scalars().all()


def get_lessons_for_slot(
    group_id: int, week_type: str, day_of_week: int, slot_number: Optional[int] = None
) -> List[LessonInfo]:
    with Session() as session:
        q = (
            _base_lesson_query()
            .join(LessonGroup, Lesson.id == LessonGroup.lesson_id)
            .where(
                LessonGroup.group_id == group_id,
                Lesson.week_type == week_type,
                Lesson.day_of_week == day_of_week,
            )
        )
        if slot_number is not None:
            q = q.where(Lesson.slot_number == slot_number)
        rows = session.execute(_order_lessons(q)).all()
        return _lesson_infos(rows)


def get_lessons_for_week_type(group_id: int, week_type: str) -> List[LessonInfo]:
    with Session() as session:
        rows = session.execute(
            _order_lessons(
                _base_lesson_query()
                .join(LessonGroup, Lesson.id == LessonGroup.lesson_id)
                .where(LessonGroup.group_id == group_id, Lesson.week_type == week_type)
            )
        ).all()
        return _lesson_infos(rows)


def get_lessons_for_group(group_id: int) -> List[LessonInfo]:
    with Session() as session:
        rows = session.execute(
            _order_lessons(
                _base_lesson_query()
                .join(LessonGroup, Lesson.id == LessonGroup.lesson_id)
                .where(LessonGroup.group_id == group_id)
            )
        ).all()
        return _lesson_infos(rows)


def get_teachers_for_group(group_id: int) -> List[Teacher]:
    with Session() as session:
        return session.execute(
            select(Teacher)
            .join(Lesson, Lesson.teacher_id == Teacher.id)
            .join(LessonGroup, Lesson.id == LessonGroup.lesson_id)
            .where(LessonGroup.group_id == group_id)
            .distinct()
            .order_by(Teacher.full_name)
        ).scalars().all()


def get_disciplines_for_group(group_id: int) -> List[Discipline]:
    with Session() as session:
        return session.execute(
            select(Discipline)
            .join(Lesson, Lesson.discipline_id == Discipline.id)
            .join(LessonGroup, Lesson.id == LessonGroup.lesson_id)
            .where(LessonGroup.group_id == group_id)
            .distinct()
            .order_by(Discipline.name)
        ).scalars().all()


def get_lessons_by_teacher(name_part: str) -> List[LessonInfo]:
    with Session() as session:
        pattern = f"%{name_part}%"
        rows = session.execute(
            _order_lessons(_base_lesson_query().where(Teacher.full_name.ilike(pattern)))
        ).all()
        return _lesson_infos(rows)


def get_lessons_by_teacher_for_group(teacher_id: int, group_id: int) -> List[LessonInfo]:
    with Session() as session:
        rows = session.execute(
            _order_lessons(
                _base_lesson_query()
                .join(LessonGroup, Lesson.id == LessonGroup.lesson_id)
                .where(Lesson.teacher_id == teacher_id, LessonGroup.group_id == group_id)
            )
        ).all()
        return _lesson_infos(rows)


def get_lessons_by_teacher_name_for_group(name_part: str, group_id: int) -> List[LessonInfo]:
    with Session() as session:
        pattern = f"%{name_part}%"
        rows = session.execute(
            _order_lessons(
                _base_lesson_query()
                .join(LessonGroup, Lesson.id == LessonGroup.lesson_id)
                .where(Teacher.full_name.ilike(pattern), LessonGroup.group_id == group_id)
            )
        ).all()
        return _lesson_infos(rows)


def get_lessons_by_group(group_id: int) -> List[LessonInfo]:
    return get_lessons_for_group(group_id)


def get_lessons_by_discipline(name_part: str) -> List[LessonInfo]:
    with Session() as session:
        pattern = f"%{name_part}%"
        rows = session.execute(
            _order_lessons(_base_lesson_query().where(Discipline.name.ilike(pattern)))
        ).all()
        return _lesson_infos(rows)


def get_lessons_by_discipline_for_group(discipline_id: int, group_id: int) -> List[LessonInfo]:
    with Session() as session:
        rows = session.execute(
            _order_lessons(
                _base_lesson_query()
                .join(LessonGroup, Lesson.id == LessonGroup.lesson_id)
                .where(Lesson.discipline_id == discipline_id, LessonGroup.group_id == group_id)
            )
        ).all()
        return _lesson_infos(rows)


def get_lessons_by_discipline_name_for_group(name_part: str, group_id: int) -> List[LessonInfo]:
    with Session() as session:
        pattern = f"%{name_part}%"
        rows = session.execute(
            _order_lessons(
                _base_lesson_query()
                .join(LessonGroup, Lesson.id == LessonGroup.lesson_id)
                .where(Discipline.name.ilike(pattern), LessonGroup.group_id == group_id)
            )
        ).all()
        return _lesson_infos(rows)


def update_lesson(
    lesson_id: int,
    week_type: str,
    day_of_week: int,
    slot_number: int,
    room: str,
    lesson_type: str,
    discipline_id: int,
    teacher_id: int,
    group_ids: Optional[List[int]] = None,
) -> Optional[Lesson]:
    with Session() as session:
        lesson = session.get(Lesson, lesson_id)
        if not lesson:
            return None
        check_groups = group_ids
        if check_groups is None:
            check_groups = [
                lg.group_id
                for lg in session.execute(
                    select(LessonGroup).where(LessonGroup.lesson_id == lesson_id)
                ).scalars().all()
            ]
        if _group_has_slot_conflict(
            session, check_groups, week_type, day_of_week, slot_number, exclude_lesson_id=lesson_id
        ):
            raise ValueError("slot_conflict")
        lesson.week_type = week_type
        lesson.day_of_week = day_of_week
        lesson.slot_number = slot_number
        lesson.room = room
        lesson.lesson_type = lesson_type
        lesson.discipline_id = discipline_id
        lesson.teacher_id = teacher_id
        if group_ids is not None:
            session.execute(delete(LessonGroup).where(LessonGroup.lesson_id == lesson_id))
            for gid in group_ids:
                session.add(LessonGroup(lesson_id=lesson_id, group_id=gid))
        session.commit()
        session.refresh(lesson)
        return lesson


def delete_lesson(lesson_id: int) -> bool:
    with Session() as session:
        lesson = session.get(Lesson, lesson_id)
        if not lesson:
            return False
        session.delete(lesson)
        session.commit()
        return True