test(cloud-store): verify schema migrations

This commit is contained in:
2026-07-12 16:59:08 +08:00
parent 24d9dbf38b
commit 7f5f740e2a
5 changed files with 224 additions and 14 deletions
@@ -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
@@ -1,6 +1,7 @@
[alembic]
script_location = %(here)s
prepend_sys_path = .
path_separator = os
sqlalchemy.url = sqlite:///cloud/cloud.sqlite3
[loggers]
@@ -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),
)
+58
View File
@@ -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
+140
View File
@@ -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)"
)