267 lines
8.1 KiB
Python
267 lines
8.1 KiB
Python
"""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")
|