celery 容灾升级

2026-07-22 14:48:29 +08:00
parent 3f1c4063b0
commit 69e7dec807
67 changed files with 6161 additions and 1958 deletions
@@ -0,0 +1,266 @@
+"""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")