From 7f5f740e2a3c19b10af26c8a896a922f675a245b Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Sun, 12 Jul 2026 16:59:08 +0800 Subject: [PATCH] test(cloud-store): verify schema migrations --- .../cloud-control-plane-integration/tasks.md | 2 +- .../cloud/migrations/alembic.ini | 1 + .../versions/0001_cloud_repository.py | 37 +++-- packages/cloud-platform/cloud/schema.py | 58 ++++++++ tests/test_cloud_migrations.py | 140 ++++++++++++++++++ 5 files changed, 224 insertions(+), 14 deletions(-) create mode 100644 packages/cloud-platform/cloud/schema.py create mode 100644 tests/test_cloud_migrations.py diff --git a/openspec/changes/cloud-control-plane-integration/tasks.md b/openspec/changes/cloud-control-plane-integration/tasks.md index c8ec09e..0f3052b 100644 --- a/openspec/changes/cloud-control-plane-integration/tasks.md +++ b/openspec/changes/cloud-control-plane-integration/tasks.md @@ -12,7 +12,7 @@ - [x] 2.2 Implement SQLAlchemy models and a repository adapter that preserves existing `CloudStore` observable behavior. - [x] 2.3 Add PostgreSQL and SQLite database URL support with engine/session lifecycle owned by the cloud application. - [x] 2.4 Add Alembic configuration and a baseline migration that preserves existing host, device, task, and plugin data while adding lease/attempt/result fields. -- [ ] 2.5 Add forward/downgrade migration tests and schema-version readiness checks. +- [x] 2.5 Add forward/downgrade migration tests and schema-version readiness checks. - [ ] 2.6 Run repository contract tests against SQLite and PostgreSQL, including rollback and process-restart cases. ## 3. Lease-Backed Scheduling diff --git a/packages/cloud-platform/cloud/migrations/alembic.ini b/packages/cloud-platform/cloud/migrations/alembic.ini index e2b5fbb..38e1425 100644 --- a/packages/cloud-platform/cloud/migrations/alembic.ini +++ b/packages/cloud-platform/cloud/migrations/alembic.ini @@ -1,6 +1,7 @@ [alembic] script_location = %(here)s prepend_sys_path = . +path_separator = os sqlalchemy.url = sqlite:///cloud/cloud.sqlite3 [loggers] diff --git a/packages/cloud-platform/cloud/migrations/versions/0001_cloud_repository.py b/packages/cloud-platform/cloud/migrations/versions/0001_cloud_repository.py index 6b574f4..7b45575 100644 --- a/packages/cloud-platform/cloud/migrations/versions/0001_cloud_repository.py +++ b/packages/cloud-platform/cloud/migrations/versions/0001_cloud_repository.py @@ -12,13 +12,13 @@ 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), +LEASE_COLUMN_NAMES = ( + "attempt_count", + "lease_id", + "lease_expires_at", + "failure_reason", + "result_json", + "updated_at", ) @@ -32,9 +32,9 @@ def upgrade() -> None: column["name"] for column in inspector.get_columns("scheduled_tasks") } with op.batch_alter_table("scheduled_tasks") as batch: - for column in LEASE_COLUMNS: + for column in _lease_columns(): if column.name not in existing_columns: - batch.add_column(column.copy()) + batch.add_column(column) inspector = sa.inspect(op.get_bind()) if "task_attempts" not in inspector.get_table_names(): @@ -52,9 +52,9 @@ def downgrade() -> None: 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) + for column_name in reversed(LEASE_COLUMN_NAMES): + if column_name in existing_columns: + batch.drop_column(column_name) def _create_base_tables() -> None: @@ -82,7 +82,7 @@ def _create_base_tables() -> None: 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], + *_lease_columns(), sa.Column("created_at", sa.String(), nullable=False), ) op.create_table( @@ -110,3 +110,14 @@ def _create_attempts_table() -> None: sa.Column("failure_reason", sa.Text(), nullable=True), sa.Column("result_json", sa.Text(), nullable=True), ) + + +def _lease_columns() -> tuple[sa.Column, ...]: + return ( + 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), + ) diff --git a/packages/cloud-platform/cloud/schema.py b/packages/cloud-platform/cloud/schema.py new file mode 100644 index 0000000..4149b93 --- /dev/null +++ b/packages/cloud-platform/cloud/schema.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +from pathlib import Path + +from alembic import command +from alembic.config import Config +from alembic.runtime.migration import MigrationContext + +from cloud.database import create_database_engine, normalize_database_url + + +HEAD_REVISION = "0001_cloud_repository" + + +class SchemaVersionError(RuntimeError): + """Raised when the database schema is not at the required revision.""" + + +def upgrade_database(database_url: str, revision: str = "head") -> None: + command.upgrade(_alembic_config(database_url), revision) + + +def downgrade_database(database_url: str, revision: str = "base") -> None: + command.downgrade(_alembic_config(database_url), revision) + + +def current_revision(database_url: str) -> str | None: + engine = create_database_engine(database_url) + try: + with engine.connect() as connection: + context = MigrationContext.configure(connection) + return context.get_current_revision() + finally: + engine.dispose() + + +def is_schema_current(database_url: str) -> bool: + return current_revision(database_url) == HEAD_REVISION + + +def require_current_schema(database_url: str) -> None: + revision = current_revision(database_url) + if revision != HEAD_REVISION: + raise SchemaVersionError( + f"cloud database schema is {revision or 'unversioned'}; " + f"required revision is {HEAD_REVISION}" + ) + + +def _alembic_config(database_url: str) -> Config: + migrations_dir = Path(__file__).resolve().parent / "migrations" + config = Config(str(migrations_dir / "alembic.ini")) + config.set_main_option("script_location", str(migrations_dir)) + config.set_main_option( + "sqlalchemy.url", + normalize_database_url(database_url).replace("%", "%%"), + ) + return config diff --git a/tests/test_cloud_migrations.py b/tests/test_cloud_migrations.py new file mode 100644 index 0000000..f9a7459 --- /dev/null +++ b/tests/test_cloud_migrations.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +import json + +import pytest +from sqlalchemy import create_engine, inspect, text + +from cloud.schema import ( + HEAD_REVISION, + SchemaVersionError, + current_revision, + downgrade_database, + require_current_schema, + upgrade_database, +) + + +def _database_url(tmp_path) -> str: + return f"sqlite:///{(tmp_path / 'cloud.sqlite3').as_posix()}" + + +def test_forward_and_downgrade_migrations_on_empty_database(tmp_path) -> None: + database_url = _database_url(tmp_path) + + upgrade_database(database_url) + + engine = create_engine(database_url) + try: + table_names = set(inspect(engine).get_table_names()) + assert { + "host_registrations", + "pooled_devices", + "scheduled_tasks", + "plugins", + "task_attempts", + } <= table_names + assert current_revision(database_url) == HEAD_REVISION + finally: + engine.dispose() + + downgrade_database(database_url) + + engine = create_engine(database_url) + try: + inspector = inspect(engine) + assert "task_attempts" not in inspector.get_table_names() + task_columns = { + column["name"] for column in inspector.get_columns("scheduled_tasks") + } + assert "lease_id" not in task_columns + assert current_revision(database_url) is None + finally: + engine.dispose() + + +def test_legacy_data_survives_upgrade_and_downgrade(tmp_path) -> None: + database_url = _database_url(tmp_path) + engine = create_engine(database_url) + try: + with engine.begin() as connection: + _create_legacy_schema(connection) + connection.execute( + text( + "insert into host_registrations " + "(host_id, address, last_seen_at) values " + "('host-a', 'local', '2026-01-01T00:00:00+00:00')" + ) + ) + connection.execute( + text( + "insert into scheduled_tasks " + "(id, goal, workflow_definition_id, constraints_json, status, " + "assigned_device_id, assigned_host_id, created_at) values " + "('task-a', 'goal', null, :constraints, 'queued', null, null, " + "'2026-01-01T00:00:00+00:00')" + ), + {"constraints": json.dumps({"driver_type": None, "capability_tags": []})}, + ) + finally: + engine.dispose() + + upgrade_database(database_url) + engine = create_engine(database_url) + try: + with engine.connect() as connection: + assert connection.scalar(text("select count(*) from host_registrations")) == 1 + assert connection.scalar(text("select count(*) from scheduled_tasks")) == 1 + task_columns = { + column["name"] for column in inspect(engine).get_columns("scheduled_tasks") + } + assert {"attempt_count", "lease_id", "result_json"} <= task_columns + finally: + engine.dispose() + + downgrade_database(database_url) + engine = create_engine(database_url) + try: + with engine.connect() as connection: + assert connection.scalar(text("select count(*) from host_registrations")) == 1 + assert connection.scalar(text("select count(*) from scheduled_tasks")) == 1 + task_columns = { + column["name"] for column in inspect(engine).get_columns("scheduled_tasks") + } + assert "attempt_count" not in task_columns + finally: + engine.dispose() + + +def test_schema_readiness_requires_head_revision(tmp_path) -> None: + database_url = _database_url(tmp_path) + + with pytest.raises(SchemaVersionError, match="unversioned"): + require_current_schema(database_url) + + upgrade_database(database_url) + require_current_schema(database_url) + + +def _create_legacy_schema(connection) -> None: + connection.exec_driver_sql( + "create table host_registrations (" + "host_id text primary key, address text, last_seen_at text not null)" + ) + connection.exec_driver_sql( + "create table pooled_devices (" + "device_id text not null, host_id text not null, driver_type text not null, " + "status text not null, capability_tags_json text not null, synced_at text, " + "primary key (host_id, device_id))" + ) + connection.exec_driver_sql( + "create table scheduled_tasks (" + "id text primary key, goal text, workflow_definition_id text, " + "constraints_json text not null, status text not null, " + "assigned_device_id text, assigned_host_id text, created_at text not null)" + ) + connection.exec_driver_sql( + "create table plugins (" + "name text primary key, version text not null, entry_point_kind text not null, " + "target text not null, wired integer not null)" + )