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 activity_source_id UUID " "REFERENCES logic_activity_sources(id)" ) ) 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( "ALTER TABLE logic_workout_items " "ADD COLUMN IF NOT EXISTS measurement_type_snapshot VARCHAR(32) DEFAULT 'weight_reps'" ) ) connection.execute( text( "ALTER TABLE logic_workout_items " "ADD COLUMN IF NOT EXISTS category_snapshot VARCHAR(32) DEFAULT 'other'" ) ) connection.execute( text( "ALTER TABLE logic_workout_items " "ADD COLUMN IF NOT EXISTS equipment_snapshot VARCHAR(32) DEFAULT 'other'" ) ) 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") ) connection.execute( text( "UPDATE logic_workout_items " "SET measurement_type_snapshot = 'weight_reps' " "WHERE measurement_type_snapshot IS NULL" ) ) connection.execute( text( "UPDATE logic_workout_items " "SET category_snapshot = 'other' " "WHERE category_snapshot IS NULL" ) ) connection.execute( text( "UPDATE logic_workout_items " "SET equipment_snapshot = 'other' " "WHERE equipment_snapshot IS NULL" ) ) connection.execute( text("ALTER TABLE logic_workout_items ALTER COLUMN measurement_type_snapshot SET NOT NULL") ) connection.execute( text("ALTER TABLE logic_workout_items ALTER COLUMN category_snapshot SET NOT NULL") ) connection.execute( text("ALTER TABLE logic_workout_items ALTER COLUMN equipment_snapshot SET NOT NULL") ) drop_constraint_if_exists( connection, constraint_name="ck_workout_item_exactly_one_entity", table_name="logic_workout_items", ) add_check_constraint_if_missing( connection, constraint_name="ck_workout_item_exactly_one_entity", table_name="logic_workout_items", check_sql="(CASE WHEN activity_source_id IS NOT NULL THEN 1 ELSE 0 END + " "CASE WHEN exercise_id IS NOT NULL THEN 1 ELSE 0 END + " "CASE WHEN equipment_id IS NOT NULL THEN 1 ELSE 0 END) = 1", ) drop_constraint_if_exists( connection, constraint_name="ck_workout_item_source_kind", table_name="logic_workout_items", ) add_check_constraint_if_missing( connection, constraint_name="ck_workout_item_source_kind", table_name="logic_workout_items", check_sql="source_kind IN ('exercise', 'machine', 'equipment')", ) connection.execute( text("ALTER TABLE logic_workout_sets ADD COLUMN IF NOT EXISTS distance_km NUMERIC(8, 3)") ) 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 drop_constraint_if_exists( connection: Connection, *, constraint_name: str, table_name: str, ) -> None: exists = connection.execute( text("SELECT 1 FROM pg_constraint WHERE conname = :constraint_name"), {"constraint_name": constraint_name}, ).scalar() if not exists: return connection.execute(text(f"ALTER TABLE {table_name} DROP CONSTRAINT {constraint_name}")) def get_db() -> Generator[Session]: db: Any = SessionLocal() try: yield db finally: db.close()