krz/hutch-stats

 1"""recent backfill scope and actor fields"""
 2
 3from __future__ import annotations
 4
 5from alembic import op
 6import sqlalchemy as sa
 7from sqlalchemy import inspect
 8
 9
10revision = "20260411_0004"
11down_revision = "20260411_0003"
12branch_labels = None
13depends_on = None
14
15
16def _table_names() -> set[str]:
17    return set(inspect(op.get_bind()).get_table_names())
18
19
20def _column_names(table_name: str) -> set[str]:
21    return {column["name"] for column in inspect(op.get_bind()).get_columns(table_name)}
22
23
24def _service_backfill_needs_upgrade() -> bool:
25    if "service_backfill_states" not in _table_names():
26        return False
27    columns = _column_names("service_backfill_states")
28    if "scope" not in columns:
29        return True
30    unique_constraints = {
31        constraint["name"]
32        for constraint in inspect(op.get_bind()).get_unique_constraints("service_backfill_states")
33    }
34    return "uq_service_backfill_state_actor_service_scope" not in unique_constraints
35
36
37def _upgrade_service_backfill_states() -> None:
38    op.execute(
39        sa.text(
40            """
41            CREATE TABLE service_backfill_states__alembic_new (
42                id INTEGER NOT NULL PRIMARY KEY,
43                actor VARCHAR(255) NOT NULL,
44                service VARCHAR(32) NOT NULL,
45                scope VARCHAR(16) NOT NULL,
46                cursor_json JSON,
47                status VARCHAR(32) NOT NULL,
48                started_at DATETIME,
49                completed_at DATETIME,
50                last_error TEXT,
51                updated_at DATETIME NOT NULL,
52                CONSTRAINT uq_service_backfill_state_actor_service_scope UNIQUE (actor, service, scope)
53            )
54            """
55        )
56    )
57    op.execute(
58        sa.text(
59            """
60            INSERT INTO service_backfill_states__alembic_new
61                (id, actor, service, scope, cursor_json, status, started_at, completed_at, last_error, updated_at)
62            SELECT id, actor, service, 'full', cursor_json, status, started_at, completed_at, last_error, updated_at
63            FROM service_backfill_states
64            """
65        )
66    )
67    op.execute(sa.text("DROP TABLE service_backfill_states"))
68    op.execute(sa.text("ALTER TABLE service_backfill_states__alembic_new RENAME TO service_backfill_states"))
69
70
71def upgrade() -> None:
72    if "tracked_actors" in _table_names():
73        columns = _column_names("tracked_actors")
74        if "recent_backfill_status" not in columns:
75            op.add_column("tracked_actors", sa.Column("recent_backfill_status", sa.String(length=32), nullable=True))
76            op.execute(sa.text("UPDATE tracked_actors SET recent_backfill_status = 'pending' WHERE recent_backfill_status IS NULL"))
77        if "recent_backfill_started_at" not in columns:
78            op.add_column("tracked_actors", sa.Column("recent_backfill_started_at", sa.DateTime(timezone=True), nullable=True))
79        if "recent_backfill_completed_at" not in columns:
80            op.add_column("tracked_actors", sa.Column("recent_backfill_completed_at", sa.DateTime(timezone=True), nullable=True))
81        if "last_recent_backfill_error" not in columns:
82            op.add_column("tracked_actors", sa.Column("last_recent_backfill_error", sa.Text(), nullable=True))
83
84    if _service_backfill_needs_upgrade():
85        _upgrade_service_backfill_states()
86
87
88def downgrade() -> None:
89    pass