workout_watcher/services/logic/app/db.py

from collections.abc import Generator
from time import sleep
from typing import Any

from sqlalchemy import Connection, MetaData, create_engine, text
from sqlalchemy.exc import OperationalError
from sqlalchemy.orm import Session, sessionmaker

from app.core import settings

engine = create_engine(
    settings.database_url,
    pool_pre_ping=True,
    connect_args={"connect_timeout": 3},
)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


def create_schema(metadata: MetaData, attempts: int = 30, delay_seconds: int = 2) -> None:
    last_error: OperationalError | None = None
    for attempt in range(1, attempts + 1):
        try:
            print(f"Connecting to database, attempt {attempt}/{attempts}", flush=True)
            with engine.begin() as connection:
                connection.execute(text("SELECT 1"))
                metadata.create_all(bind=connection)
                upgrade_existing_schema(connection)
            print("Database schema is ready", flush=True)
            return
        except OperationalError as exc:
            last_error = exc
            print(f"Database is not ready: {exc}", flush=True)
            sleep(delay_seconds)
    if last_error:
        raise last_error


def upgrade_existing_schema(connection: Connection) -> None:
    """Apply safe additive upgrades for DBs created by pre-Alembic create_all.

    Local Compose databases may already have the initial tables without the
    tidy-wolf columns. SQLAlchemy create_all intentionally does not ALTER
    existing tables, so keep this idempotent compatibility upgrade until
    migrations are wired into service startup.
    """
    connection.execute(
        text("ALTER TABLE logic_workouts ADD COLUMN IF NOT EXISTS status VARCHAR(20)")
    )
    connection.execute(
        text("ALTER TABLE logic_workouts ADD COLUMN IF NOT EXISTS total_sets INTEGER")
    )
    connection.execute(
        text("ALTER TABLE logic_workouts ADD COLUMN IF NOT EXISTS total_volume NUMERIC(12, 2)")
    )
    connection.execute(
        text(
            """
            UPDATE logic_workouts
            SET status = CASE WHEN finished_at IS NULL THEN 'active' ELSE 'finished' END
            WHERE status IS NULL
            """
        )
    )
    connection.execute(text("UPDATE logic_workouts SET total_sets = 0 WHERE total_sets IS NULL"))
    connection.execute(
        text("UPDATE logic_workouts SET total_volume = 0 WHERE total_volume IS NULL")
    )
    connection.execute(text("ALTER TABLE logic_workouts ALTER COLUMN status SET DEFAULT 'active'"))
    connection.execute(text("ALTER TABLE logic_workouts ALTER COLUMN status SET NOT NULL"))
    connection.execute(text("ALTER TABLE logic_workouts ALTER COLUMN total_sets SET DEFAULT 0"))
    connection.execute(text("ALTER TABLE logic_workouts ALTER COLUMN total_sets SET NOT NULL"))
    connection.execute(text("ALTER TABLE logic_workouts ALTER COLUMN total_volume SET DEFAULT 0"))
    connection.execute(text("ALTER TABLE logic_workouts ALTER COLUMN total_volume SET NOT NULL"))
    connection.execute(
        text("CREATE INDEX IF NOT EXISTS ix_logic_workouts_status ON logic_workouts (status)")
    )
    add_check_constraint_if_missing(
        connection,
        constraint_name="ck_workout_status",
        table_name="logic_workouts",
        check_sql="status IN ('active', 'finished', 'discarded')",
    )

    connection.execute(
        text("ALTER TABLE logic_workout_items ADD COLUMN IF NOT EXISTS source_kind VARCHAR(20)")
    )
    connection.execute(
        text("ALTER TABLE logic_workout_items ADD COLUMN IF NOT EXISTS title_snapshot VARCHAR(160)")
    )
    connection.execute(
        text("ALTER TABLE logic_workout_items ADD COLUMN IF NOT EXISTS image_s3_url_snapshot TEXT")
    )
    connection.execute(
        text(
            """
            UPDATE logic_workout_items AS item
            SET
                source_kind = 'exercise',
                title_snapshot = exercise.name,
                image_s3_url_snapshot = exercise.image_s3_url
            FROM logic_exercises AS exercise
            WHERE item.exercise_id = exercise.id
              AND (item.source_kind IS NULL OR item.title_snapshot IS NULL)
            """
        )
    )
    connection.execute(
        text(
            """
            UPDATE logic_workout_items AS item
            SET
                source_kind = 'equipment',
                title_snapshot = equipment.name,
                image_s3_url_snapshot = equipment.image_s3_url
            FROM logic_equipment AS equipment
            WHERE item.equipment_id = equipment.id
              AND (item.source_kind IS NULL OR item.title_snapshot IS NULL)
            """
        )
    )
    connection.execute(
        text(
            """
            UPDATE logic_workout_items
            SET source_kind = CASE WHEN exercise_id IS NOT NULL THEN 'exercise' ELSE 'equipment' END
            WHERE source_kind IS NULL
            """
        )
    )
    connection.execute(
        text(
            "UPDATE logic_workout_items "
            "SET title_snapshot = 'Без названия' "
            "WHERE title_snapshot IS NULL"
        )
    )
    connection.execute(
        text("ALTER TABLE logic_workout_items ALTER COLUMN source_kind SET NOT NULL")
    )
    connection.execute(
        text("ALTER TABLE logic_workout_items ALTER COLUMN title_snapshot SET NOT NULL")
    )
    add_check_constraint_if_missing(
        connection,
        constraint_name="ck_workout_item_source_kind",
        table_name="logic_workout_items",
        check_sql="source_kind IN ('exercise', 'equipment')",
    )

    connection.execute(
        text(
            """
            UPDATE logic_workouts AS workout
            SET
                total_sets = totals.total_sets,
                total_volume = totals.total_volume,
                estimated_calories = totals.estimated_calories
            FROM (
                SELECT
                    workout_source.id AS workout_id,
                    COUNT(set_row.id) AS total_sets,
                    COALESCE(SUM(set_row.weight * set_row.reps), 0) AS total_volume,
                    COALESCE(SUM(set_row.calories), 0) AS estimated_calories
                FROM logic_workouts AS workout_source
                LEFT JOIN logic_workout_items AS item ON item.workout_id = workout_source.id
                LEFT JOIN logic_workout_sets AS set_row ON set_row.workout_item_id = item.id
                GROUP BY workout_source.id
            ) AS totals
            WHERE workout.id = totals.workout_id
            """
        )
    )


def add_check_constraint_if_missing(
    connection: Connection,
    *,
    constraint_name: str,
    table_name: str,
    check_sql: str,
) -> None:
    exists = connection.execute(
        text("SELECT 1 FROM pg_constraint WHERE conname = :constraint_name"),
        {"constraint_name": constraint_name},
    ).scalar()
    if exists:
        return
    connection.execute(
        text(
            f"ALTER TABLE {table_name} "
            f"ADD CONSTRAINT {constraint_name} "
            f"CHECK ({check_sql})"
        )
    )


def get_db() -> Generator[Session]:
    db: Any = SessionLocal()
    try:
        yield db
    finally:
        db.close()