db/database.py - main - o.polulyakh/RunTracker_Bot

db/database.py
215 строк8 KB
Oksana Polulyakh
verse 0.1.2
02 янв 2026, 18:49
02 янв 2026, 18:494c52575
from datetime import date

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from db.models import Base
from core.config import DSN
from contextlib import contextmanager
from db.models import User, Workout

engine = create_engine(DSN)
SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False)


def create_table():
    """
    Создает все таблицы в базе данных.
    
    Эта функция использует метаданные SQLAlchemy для создания
    всех таблиц, определенных в моделях, если они еще не существуют.
    """
    Base.metadata.create_all(engine)


def drop_table():
    """
    Удаляет все таблицы из базы данных.
    
    Осторожно: эта функция удаляет все данные из таблиц.
    Используйте только для тестирования или сброса базы данных.
    """
    Base.metadata.drop_all(engine)


def get_session():
    """
    Создает и возвращает новую сессию SQLAlchemy.
    
    Returns:
        Session: Новая сессия для работы с базой данных
    """
    return SessionLocal()


def get_user_data(telegram_id: int) -> dict | None:
    """
    Получает данные пользователя по его Telegram ID.
    
    Args:
        telegram_id (int): Уникальный идентификатор пользователя в Telegram
    
    Returns:
        dict | None: Словарь с данными пользователя или None, если пользователь не найден
    """
    with session_scope() as session:
        user = session.query(User).filter(User.telegram_id == telegram_id).first()
        if user:
            return user.to_dict()
        return None


def get_user_stats(user_id: int) -> dict:
    """
    Рассчитывает и возвращает статистику тренировок пользователя.
    
    Args:
        user_id (int): Идентификатор пользователя в базе данных
    
    Returns:
        dict: Словарь со статистикой, содержащий:
            - count (int): количество тренировок
            - total_km (float): общая дистанция в км
            - total_min (int): общее время в минутах
            - avg_pace (float): средний темп в мин/км
            - best_pace (float): лучший темп в мин/км
            - max_distance (float): максимальная дистанция в км
    """
    with session_scope() as session:
        workouts = session.query(Workout).filter(Workout.user_id == user_id).all()
        if not workouts:
            return {"count": 0}

        total_km = sum(w.distance for w in workouts)
        total_min = sum(w.duration_min for w in workouts)
        paces = [w.pace for w in workouts if w.pace > 0]

        return {
            "count": len(workouts),
            "total_km": total_km,
            "total_min": total_min,
            "avg_pace": total_min / total_km if total_km > 0 else 0,
            "best_pace": min(paces) if paces else 0,
            "max_distance": max(w.distance for w in workouts),
        }


def get_user_workouts(user_id: int) -> list[dict]:
    """
    Получает список всех тренировок пользователя, отсортированных по дате.
    
    Args:
        user_id (int): Идентификатор пользователя в базе данных
    
    Returns:
        list[dict]: Список словарей с данными тренировок, 
                   отсортированный по дате (сначала новые)
    """
    with session_scope() as session:
        workouts = (
            session.query(Workout)
            .filter(Workout.user_id == user_id)
            .order_by(Workout.date_time.desc())
            .all()
        )
        return [w.to_dict() for w in workouts]


def create_user(telegram_id: int, first_name: str, last_name: str | None) -> dict:
    """
    Создает нового пользователя в базе данных.
    
    Args:
        telegram_id (int): Уникальный идентификатор пользователя в Telegram
        first_name (str): Имя пользователя
        last_name (str | None): Фамилия пользователя (может быть None)
    
    Returns:
        dict: Словарь с данными созданного пользователя
    
    Note:
        По умолчанию всем пользователям присваивается роль "runner"
    """
    with session_scope() as session:
        user = User(
            telegram_id=telegram_id,
            first_name=first_name,
            last_name=last_name,
            role="runner"
        )
        session.add(user)
        session.flush()
        return user.to_dict()


@contextmanager
def session_scope():
    """
    Контекстный менеджер для работы с сессией SQLAlchemy.
    
    Обеспечивает автоматическое подтверждение транзакции при успешном выполнении
    и откат при возникновении исключения. В любом случае сессия закрывается.
    
    Yields:
        Session: Сессия SQLAlchemy для выполнения операций с базой данных
        # Сессия автоматически коммитится при выходе из блока
    """
    session = SessionLocal()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


def create_workout(user_id: int, distance: float, duration_min: int, date_time: date) -> dict:
    """
    Создает новую запись о тренировке в базе данных.
    
    Args:
        user_id (int): Идентификатор пользователя
        distance (float): Дистанция в километрах
        duration_min (int): Время тренировки в минутах
        date_time (date): Дата тренировки
    
    Returns:
        dict: Словарь с данными созданной тренировки, включая рассчитанный темп
    """
    with session_scope() as session:
        workout = Workout(
            user_id=user_id,
            distance=distance,
            duration_min=duration_min,
            date_time=date_time
        )
        session.add(workout)
        session.flush()
        return workout.to_dict()


def delete_workout(workout_id: int, user_id: int) -> bool:
    """
    Удаляет тренировку из базы данных.
    
    Args:
        workout_id (int): Идентификатор тренировки для удаления
        user_id (int): Идентификатор владельца тренировки (для безопасности)
    
    Returns:
        bool: True, если тренировка была успешно удалена, False если тренировка не найдена
    
    Note:
        Функция проверяет, что тренировка принадлежит указанному пользователю,
        чтобы предотвратить удаление чужих тренировок.
    """
    with session_scope() as session:
        workout = session.query(Workout).filter(
            Workout.id == workout_id,
            Workout.user_id == user_id
        ).first()
        if workout:
            session.delete(workout)
            return True
        return False