- Add nullable cancel_requested_at column (migration 0012) - Widen ScheduledTaskStatus/TerminalTaskStatus to include cancelled - Add CancellationRequestStatus + request_task_cancellation() to CloudRepository protocol and SQLAlchemy implementation - renew_lease() now returns LeaseRenewalResult, surfacing whether cancellation is pending, instead of a bare status string - reap_expired_leases() resolves pending-cancellation tasks to cancelled instead of requeuing/failing them - record_task_result() accepts cancelled and clears cancel_requested_at on any terminal write Note: internal_api/api.py's renew_assignment route still compares renew_lease()'s return value against a bare string; it will be updated in the next task (Internal Host<->Cloud protocol) to consume LeaseRenewalResult and populate the new cancel_requested wire field.
59 lines
1.8 KiB
Python
59 lines
1.8 KiB
Python
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 = "0012_task_cancellation"
|
|
|
|
|
|
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
|