feat(cloud-store): add Alembic baseline migration

2026-07-12 16:55:46 +08:00
parent 64aa9b39bc
commit 24d9dbf38b
10 changed files with 328 additions and 2 deletions
@@ -0,0 +1,112 @@
+"""Create the cloud repository schema and lease fields."""
+
+from __future__ import annotations
+
+from alembic import op
+import sqlalchemy as sa
+
+
+revision = "0001_cloud_repository"
+down_revision = None
+branch_labels = None
+depends_on = None
+
+
+LEASE_COLUMNS = (
+    sa.Column("attempt_count", sa.Integer(), nullable=False, server_default="0"),
+    sa.Column("lease_id", sa.String(), nullable=True),
+    sa.Column("lease_expires_at", sa.String(), nullable=True),
+    sa.Column("failure_reason", sa.Text(), nullable=True),
+    sa.Column("result_json", sa.Text(), nullable=True),
+    sa.Column("updated_at", sa.String(), nullable=True),
+)
+
+
+def upgrade() -> None:
+    inspector = sa.inspect(op.get_bind())
+    tables = set(inspector.get_table_names())
+    if "host_registrations" not in tables:
+        _create_base_tables()
+    else:
+        existing_columns = {
+            column["name"] for column in inspector.get_columns("scheduled_tasks")
+        }
+        with op.batch_alter_table("scheduled_tasks") as batch:
+            for column in LEASE_COLUMNS:
+                if column.name not in existing_columns:
+                    batch.add_column(column.copy())
+
+    inspector = sa.inspect(op.get_bind())
+    if "task_attempts" not in inspector.get_table_names():
+        _create_attempts_table()
+
+
+def downgrade() -> None:
+    inspector = sa.inspect(op.get_bind())
+    tables = set(inspector.get_table_names())
+    if "task_attempts" in tables:
+        op.drop_table("task_attempts")
+    if "scheduled_tasks" not in tables:
+        return
+    existing_columns = {
+        column["name"] for column in inspector.get_columns("scheduled_tasks")
+    }
+    with op.batch_alter_table("scheduled_tasks") as batch:
+        for column in reversed(LEASE_COLUMNS):
+            if column.name in existing_columns:
+                batch.drop_column(column.name)
+
+
+def _create_base_tables() -> None:
+    op.create_table(
+        "host_registrations",
+        sa.Column("host_id", sa.String(), primary_key=True),
+        sa.Column("address", sa.String(), nullable=True),
+        sa.Column("last_seen_at", sa.String(), nullable=False),
+    )
+    op.create_table(
+        "pooled_devices",
+        sa.Column("host_id", sa.String(), primary_key=True),
+        sa.Column("device_id", sa.String(), primary_key=True),
+        sa.Column("driver_type", sa.String(), nullable=False),
+        sa.Column("status", sa.String(), nullable=False),
+        sa.Column("capability_tags_json", sa.Text(), nullable=False),
+        sa.Column("synced_at", sa.String(), nullable=True),
+    )
+    op.create_table(
+        "scheduled_tasks",
+        sa.Column("id", sa.String(), primary_key=True),
+        sa.Column("goal", sa.Text(), nullable=True),
+        sa.Column("workflow_definition_id", sa.String(), nullable=True),
+        sa.Column("constraints_json", sa.Text(), nullable=False),
+        sa.Column("status", sa.String(), nullable=False),
+        sa.Column("assigned_device_id", sa.String(), nullable=True),
+        sa.Column("assigned_host_id", sa.String(), nullable=True),
+        *[column.copy() for column in LEASE_COLUMNS],
+        sa.Column("created_at", sa.String(), nullable=False),
+    )
+    op.create_table(
+        "plugins",
+        sa.Column("name", sa.String(), primary_key=True),
+        sa.Column("version", sa.String(), nullable=False),
+        sa.Column("entry_point_kind", sa.String(), nullable=False),
+        sa.Column("target", sa.String(), nullable=False),
+        sa.Column("wired", sa.Integer(), nullable=False),
+    )
+
+
+def _create_attempts_table() -> None:
+    op.create_table(
+        "task_attempts",
+        sa.Column("task_id", sa.String(), primary_key=True),
+        sa.Column("attempt", sa.Integer(), primary_key=True),
+        sa.Column("lease_id", sa.String(), nullable=False),
+        sa.Column("host_id", sa.String(), nullable=False),
+        sa.Column("device_id", sa.String(), nullable=False),
+        sa.Column("status", sa.String(), nullable=False),
+        sa.Column("lease_expires_at", sa.String(), nullable=False),
+        sa.Column("created_at", sa.String(), nullable=False),
+        sa.Column("completed_at", sa.String(), nullable=True),
+        sa.Column("failure_reason", sa.Text(), nullable=True),
+        sa.Column("result_json", sa.Text(), nullable=True),
+    )
@@ -0,0 +1 @@
+"""Cloud database schema revisions."""