"""add celery runtime fencing fields Revision ID: 7cf645f7c418 Revises: d8ebe79ab575 Create Date: 2026-07-22 13:45:57.949417 """ from __future__ import annotations from collections.abc import Sequence from alembic import op import sqlalchemy as sa from sqlalchemy.engine.reflection import Inspector # revision identifiers, used by Alembic. revision: str = "7cf645f7c418" down_revision: str | None = "d8ebe79ab575" branch_labels: str | Sequence[str] | None = None depends_on: str | Sequence[str] | None = None _SHOT_SEGMENT_TABLE = "shot_replicate_segments" _SHOT_TASK_SET_TABLE = "shot_replicate_task_sets" _ANALYSIS_ATTEMPT_COLUMN = "analysis_attempt_no" # SQLAlchemy Inspector caches reflected metadata. Always create a fresh # inspector after DDL so partially applied migrations are detected correctly. def _inspector() -> Inspector: return sa.inspect(op.get_bind()) def _require_table(table_name: str) -> None: if not _inspector().has_table(table_name): raise RuntimeError( f"Required table {table_name!r} does not exist; " "refusing to mark migration 7cf645f7c418 as applied incompletely." ) def _column_names(table_name: str) -> set[str]: return { str(column["name"]) for column in _inspector().get_columns(table_name) } def _add_column_if_missing(table_name: str, column: sa.Column[object]) -> bool: """Add one column only when its name is absent. Returns True when DDL was executed and False when the column already exists. """ if column.name in _column_names(table_name): return False op.add_column(table_name, column) return True def _drop_column_if_exists(table_name: str, column_name: str) -> bool: """Drop one column only when both the table and column still exist.""" inspector = _inspector() if not inspector.has_table(table_name): return False if column_name not in { str(column["name"]) for column in inspector.get_columns(table_name) }: return False op.drop_column(table_name, column_name) return True def _index_definitions(table_name: str) -> dict[str, dict[str, object]]: return { str(index["name"]): index for index in _inspector().get_indexes(table_name) if index.get("name") } def _ensure_index( index_name: str, table_name: str, columns: Sequence[str], *, unique: bool = False, ) -> None: """Create an index if absent and reject a conflicting same-name index.""" required_columns = tuple(columns) existing_columns = _column_names(table_name) missing_columns = [name for name in required_columns if name not in existing_columns] if missing_columns: raise RuntimeError( f"Cannot create index {index_name!r}: table {table_name!r} " f"is missing columns {missing_columns!r}." ) existing = _index_definitions(table_name).get(index_name) if existing is not None: reflected_columns = tuple( str(name) for name in (existing.get("column_names") or []) ) reflected_unique = bool(existing.get("unique", False)) if reflected_columns != required_columns or reflected_unique != unique: raise RuntimeError( f"Index {index_name!r} already exists with an unexpected definition: " f"columns={reflected_columns!r}, unique={reflected_unique!r}; " f"expected columns={required_columns!r}, unique={unique!r}." ) return op.create_index( index_name, table_name, list(required_columns), unique=unique, ) def _drop_index_if_exists(index_name: str, table_name: str) -> bool: """Drop one index only when the table and named index still exist.""" inspector = _inspector() if not inspector.has_table(table_name): return False existing_names = { str(index["name"]) for index in inspector.get_indexes(table_name) if index.get("name") } if index_name not in existing_names: return False op.drop_index(index_name, table_name=table_name) return True def _ensure_analysis_attempt_column(table_name: str) -> None: """Create/backfill the non-null attempt counter safely for existing rows.""" _add_column_if_missing( table_name, sa.Column( _ANALYSIS_ATTEMPT_COLUMN, sa.Integer(), nullable=False, server_default=sa.text("1"), ), ) # Also repairs a partially applied/manual migration where the column exists # but contains NULL values or still carries the temporary database default. quoted_table = op.get_bind().dialect.identifier_preparer.quote(table_name) quoted_column = op.get_bind().dialect.identifier_preparer.quote( _ANALYSIS_ATTEMPT_COLUMN ) op.execute( sa.text( f"UPDATE {quoted_table} " f"SET {quoted_column} = 1 " f"WHERE {quoted_column} IS NULL" ) ) op.alter_column( table_name, _ANALYSIS_ATTEMPT_COLUMN, existing_type=sa.Integer(), nullable=False, server_default=None, ) def upgrade() -> None: """Add only the Celery fencing fields required by shot replication.""" _require_table(_SHOT_SEGMENT_TABLE) _require_table(_SHOT_TASK_SET_TABLE) # Segment split fencing. _add_column_if_missing( _SHOT_SEGMENT_TABLE, sa.Column("split_claim_token", sa.String(length=64), nullable=True), ) # Segment analysis fencing. _ensure_analysis_attempt_column(_SHOT_SEGMENT_TABLE) _add_column_if_missing( _SHOT_SEGMENT_TABLE, sa.Column("analysis_claim_token", sa.String(length=64), nullable=True), ) _add_column_if_missing( _SHOT_SEGMENT_TABLE, sa.Column("analysis_started_at", sa.DateTime(timezone=True), nullable=True), ) _add_column_if_missing( _SHOT_SEGMENT_TABLE, sa.Column("analysis_lease_until", sa.DateTime(timezone=True), nullable=True), ) _ensure_index( "idx_shot_segments_analysis_lease", _SHOT_SEGMENT_TABLE, ("analysis_status", "analysis_lease_until"), ) _ensure_index( "idx_shot_segments_split_lease", _SHOT_SEGMENT_TABLE, ("split_status", "split_lease_until"), ) # Task-set analysis fencing. _ensure_analysis_attempt_column(_SHOT_TASK_SET_TABLE) _add_column_if_missing( _SHOT_TASK_SET_TABLE, sa.Column("analysis_claim_token", sa.String(length=64), nullable=True), ) _add_column_if_missing( _SHOT_TASK_SET_TABLE, sa.Column("analysis_started_at", sa.DateTime(timezone=True), nullable=True), ) _add_column_if_missing( _SHOT_TASK_SET_TABLE, sa.Column("analysis_lease_until", sa.DateTime(timezone=True), nullable=True), ) _ensure_index( "idx_shot_task_sets_analysis_lease", _SHOT_TASK_SET_TABLE, ("analysis_status", "analysis_lease_until"), ) def downgrade() -> None: """Remove only fields and indexes introduced by this revision. Every operation is guarded so a partially reverted database does not fail merely because an index or column is already absent. """ _drop_index_if_exists( "idx_shot_task_sets_analysis_lease", _SHOT_TASK_SET_TABLE, ) _drop_column_if_exists(_SHOT_TASK_SET_TABLE, "analysis_lease_until") _drop_column_if_exists(_SHOT_TASK_SET_TABLE, "analysis_started_at") _drop_column_if_exists(_SHOT_TASK_SET_TABLE, "analysis_claim_token") _drop_column_if_exists(_SHOT_TASK_SET_TABLE, _ANALYSIS_ATTEMPT_COLUMN) _drop_index_if_exists( "idx_shot_segments_split_lease", _SHOT_SEGMENT_TABLE, ) _drop_index_if_exists( "idx_shot_segments_analysis_lease", _SHOT_SEGMENT_TABLE, ) _drop_column_if_exists(_SHOT_SEGMENT_TABLE, "analysis_lease_until") _drop_column_if_exists(_SHOT_SEGMENT_TABLE, "analysis_started_at") _drop_column_if_exists(_SHOT_SEGMENT_TABLE, "analysis_claim_token") _drop_column_if_exists(_SHOT_SEGMENT_TABLE, _ANALYSIS_ATTEMPT_COLUMN) _drop_column_if_exists(_SHOT_SEGMENT_TABLE, "split_claim_token")