656 lines
23 KiB
Python
656 lines
23 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import asdict
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
from sqlalchemy import Engine, delete, func, select
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
from cloud.db_models import (
|
|
Base,
|
|
HostRow,
|
|
PluginRow,
|
|
PooledDeviceRow,
|
|
ScheduledTaskRow,
|
|
TaskAttemptRow,
|
|
)
|
|
from core.models import utc_now
|
|
|
|
|
|
class SQLAlchemyCloudRepository:
|
|
"""SQLAlchemy adapter preserving the existing CloudStore CRUD surface."""
|
|
|
|
def __init__(self, engine: Engine, *, create_schema: bool = True) -> None:
|
|
self.engine = engine
|
|
self._sessions = sessionmaker(bind=engine, expire_on_commit=False)
|
|
if create_schema:
|
|
Base.metadata.create_all(engine)
|
|
|
|
def upsert_host(
|
|
self,
|
|
host_id: str,
|
|
*,
|
|
address: str | None,
|
|
last_seen_at: datetime,
|
|
) -> None:
|
|
with self._sessions.begin() as session:
|
|
row = session.get(HostRow, host_id)
|
|
if row is None:
|
|
session.add(
|
|
HostRow(
|
|
host_id=host_id,
|
|
address=address,
|
|
last_seen_at=_iso(last_seen_at),
|
|
)
|
|
)
|
|
return
|
|
if address is not None:
|
|
row.address = address
|
|
row.last_seen_at = _iso(last_seen_at)
|
|
|
|
def replace_host_devices(
|
|
self,
|
|
host_id: str,
|
|
devices: list[Any],
|
|
) -> None:
|
|
with self._sessions.begin() as session:
|
|
session.execute(
|
|
delete(PooledDeviceRow).where(PooledDeviceRow.host_id == host_id)
|
|
)
|
|
session.add_all(
|
|
[
|
|
PooledDeviceRow(
|
|
device_id=device.device_id,
|
|
host_id=device.host_id,
|
|
driver_type=device.driver_type,
|
|
status=device.status,
|
|
capability_tags_json=json.dumps(
|
|
list(device.capability_tags),
|
|
ensure_ascii=False,
|
|
),
|
|
synced_at=_iso(device.synced_at) if device.synced_at else None,
|
|
)
|
|
for device in devices
|
|
]
|
|
)
|
|
|
|
def list_hosts(self) -> list[Any]:
|
|
with self._sessions() as session:
|
|
rows = session.scalars(select(HostRow).order_by(HostRow.host_id)).all()
|
|
return [_host_from_row(row) for row in rows]
|
|
|
|
def get_host(self, host_id: str) -> Any | None:
|
|
with self._sessions() as session:
|
|
row = session.get(HostRow, host_id)
|
|
return _host_from_row(row) if row else None
|
|
|
|
def list_devices(self) -> list[Any]:
|
|
with self._sessions() as session:
|
|
rows = session.scalars(
|
|
select(PooledDeviceRow).order_by(
|
|
PooledDeviceRow.host_id,
|
|
PooledDeviceRow.device_id,
|
|
)
|
|
).all()
|
|
return [_device_from_row(row) for row in rows]
|
|
|
|
def get_device(self, device_id: str) -> Any | None:
|
|
with self._sessions() as session:
|
|
row = session.scalars(
|
|
select(PooledDeviceRow)
|
|
.where(PooledDeviceRow.device_id == device_id)
|
|
.order_by(PooledDeviceRow.host_id)
|
|
.limit(1)
|
|
).first()
|
|
return _device_from_row(row) if row else None
|
|
|
|
def enqueue_task(self, task: Any) -> None:
|
|
with self._sessions.begin() as session:
|
|
session.add(
|
|
ScheduledTaskRow(
|
|
id=task.id,
|
|
goal=task.goal,
|
|
workflow_definition_id=task.workflow_definition_id,
|
|
constraints_json=json.dumps(
|
|
asdict(task.constraints),
|
|
ensure_ascii=False,
|
|
),
|
|
status=task.status,
|
|
assigned_device_id=task.assigned_device_id,
|
|
assigned_host_id=task.assigned_host_id,
|
|
attempt_count=task.attempt_count,
|
|
lease_id=task.lease_id,
|
|
lease_expires_at=(
|
|
_iso(task.lease_expires_at) if task.lease_expires_at else None
|
|
),
|
|
failure_reason=task.failure_reason,
|
|
result_json=(
|
|
json.dumps(task.terminal_result, ensure_ascii=False)
|
|
if task.terminal_result is not None
|
|
else None
|
|
),
|
|
updated_at=_iso(task.updated_at) if task.updated_at else None,
|
|
created_at=_iso(task.created_at),
|
|
)
|
|
)
|
|
|
|
def list_queued_tasks(self) -> list[Any]:
|
|
with self._sessions() as session:
|
|
rows = session.scalars(
|
|
select(ScheduledTaskRow)
|
|
.where(ScheduledTaskRow.status == "queued")
|
|
.order_by(ScheduledTaskRow.created_at, ScheduledTaskRow.id)
|
|
).all()
|
|
return [_task_from_row(row) for row in rows]
|
|
|
|
def get_task(self, task_id: str) -> Any | None:
|
|
with self._sessions() as session:
|
|
row = session.get(ScheduledTaskRow, task_id)
|
|
return _task_from_row(row) if row else None
|
|
|
|
def update_task(
|
|
self,
|
|
task_id: str,
|
|
*,
|
|
status: str | None = None,
|
|
assigned_device_id: str | None = None,
|
|
assigned_host_id: str | None = None,
|
|
) -> None:
|
|
with self._sessions.begin() as session:
|
|
row = session.get(ScheduledTaskRow, task_id)
|
|
if row is None:
|
|
return
|
|
if status is not None:
|
|
row.status = status
|
|
if assigned_device_id is not None:
|
|
row.assigned_device_id = assigned_device_id
|
|
if assigned_host_id is not None:
|
|
row.assigned_host_id = assigned_host_id
|
|
|
|
def count_queued_tasks(self) -> int:
|
|
with self._sessions() as session:
|
|
count = session.scalar(
|
|
select(func.count())
|
|
.select_from(ScheduledTaskRow)
|
|
.where(ScheduledTaskRow.status == "queued")
|
|
)
|
|
return int(count or 0)
|
|
|
|
def save_plugin(self, manifest: Any, *, wired: bool) -> None:
|
|
with self._sessions.begin() as session:
|
|
row = session.get(PluginRow, manifest.name)
|
|
if row is None:
|
|
session.add(
|
|
PluginRow(
|
|
name=manifest.name,
|
|
version=manifest.version,
|
|
entry_point_kind=manifest.entry_point_kind,
|
|
target=manifest.target,
|
|
wired=1 if wired else 0,
|
|
)
|
|
)
|
|
return
|
|
row.version = manifest.version
|
|
row.entry_point_kind = manifest.entry_point_kind
|
|
row.target = manifest.target
|
|
row.wired = 1 if wired else 0
|
|
|
|
def list_plugins(self) -> list[tuple[Any, bool]]:
|
|
with self._sessions() as session:
|
|
rows = session.scalars(select(PluginRow).order_by(PluginRow.name)).all()
|
|
return [_plugin_from_row(row) for row in rows]
|
|
|
|
def get_plugin(self, name: str) -> tuple[Any, bool] | None:
|
|
with self._sessions() as session:
|
|
row = session.get(PluginRow, name)
|
|
return _plugin_from_row(row) if row else None
|
|
|
|
def list_reserved_device_ids(self, *, now: datetime) -> set[str]:
|
|
with self._sessions() as session:
|
|
device_ids = session.scalars(
|
|
select(ScheduledTaskRow.assigned_device_id).where(
|
|
ScheduledTaskRow.status.in_(("assigned", "dispatched")),
|
|
ScheduledTaskRow.assigned_device_id.is_not(None),
|
|
ScheduledTaskRow.lease_expires_at.is_not(None),
|
|
ScheduledTaskRow.lease_expires_at > _iso(now),
|
|
)
|
|
).all()
|
|
return {device_id for device_id in device_ids if device_id is not None}
|
|
|
|
def assign_task(
|
|
self,
|
|
*,
|
|
task_id: str,
|
|
host_id: str,
|
|
device_id: str,
|
|
lease_id: str,
|
|
lease_expires_at: datetime,
|
|
now: datetime,
|
|
) -> Any | None:
|
|
with self._sessions.begin() as session:
|
|
task_statement = select(ScheduledTaskRow).where(
|
|
ScheduledTaskRow.id == task_id
|
|
)
|
|
device_statement = select(PooledDeviceRow).where(
|
|
PooledDeviceRow.host_id == host_id,
|
|
PooledDeviceRow.device_id == device_id,
|
|
)
|
|
if self.engine.dialect.name == "postgresql":
|
|
task_statement = task_statement.with_for_update(skip_locked=True)
|
|
device_statement = device_statement.with_for_update()
|
|
|
|
task = session.scalars(task_statement).first()
|
|
if task is None or task.status != "queued":
|
|
return None
|
|
device = session.scalars(device_statement).first()
|
|
if device is None or device.status != "idle":
|
|
return None
|
|
|
|
active_reservation = session.scalar(
|
|
select(ScheduledTaskRow.id)
|
|
.where(
|
|
ScheduledTaskRow.id != task_id,
|
|
ScheduledTaskRow.assigned_device_id == device_id,
|
|
ScheduledTaskRow.status.in_(("assigned", "dispatched")),
|
|
ScheduledTaskRow.lease_expires_at.is_not(None),
|
|
ScheduledTaskRow.lease_expires_at > _iso(now),
|
|
)
|
|
.limit(1)
|
|
)
|
|
if active_reservation is not None:
|
|
return None
|
|
|
|
attempt = task.attempt_count + 1
|
|
task.status = "assigned"
|
|
task.assigned_host_id = host_id
|
|
task.assigned_device_id = device_id
|
|
task.attempt_count = attempt
|
|
task.lease_id = lease_id
|
|
task.lease_expires_at = _iso(lease_expires_at)
|
|
task.failure_reason = None
|
|
task.result_json = None
|
|
task.updated_at = _iso(now)
|
|
session.add(
|
|
TaskAttemptRow(
|
|
task_id=task.id,
|
|
attempt=attempt,
|
|
lease_id=lease_id,
|
|
host_id=host_id,
|
|
device_id=device_id,
|
|
status="assigned",
|
|
lease_expires_at=_iso(lease_expires_at),
|
|
created_at=_iso(now),
|
|
completed_at=None,
|
|
failure_reason=None,
|
|
result_json=None,
|
|
)
|
|
)
|
|
session.flush()
|
|
return _leased_assignment_from_row(task)
|
|
|
|
def claim_assignment(
|
|
self,
|
|
*,
|
|
host_id: str,
|
|
now: datetime,
|
|
) -> Any | None:
|
|
with self._sessions.begin() as session:
|
|
statement = (
|
|
select(ScheduledTaskRow)
|
|
.where(
|
|
ScheduledTaskRow.status == "assigned",
|
|
ScheduledTaskRow.assigned_host_id == host_id,
|
|
ScheduledTaskRow.lease_id.is_not(None),
|
|
ScheduledTaskRow.lease_expires_at.is_not(None),
|
|
ScheduledTaskRow.lease_expires_at > _iso(now),
|
|
)
|
|
.order_by(ScheduledTaskRow.created_at, ScheduledTaskRow.id)
|
|
.limit(1)
|
|
)
|
|
if self.engine.dialect.name == "postgresql":
|
|
statement = statement.with_for_update(skip_locked=True)
|
|
|
|
task = session.scalars(statement).first()
|
|
if task is None:
|
|
return None
|
|
attempt = session.get(
|
|
TaskAttemptRow,
|
|
(task.id, task.attempt_count),
|
|
with_for_update=self.engine.dialect.name == "postgresql",
|
|
)
|
|
if (
|
|
attempt is None
|
|
or attempt.status != "assigned"
|
|
or attempt.lease_id != task.lease_id
|
|
):
|
|
return None
|
|
|
|
task.status = "dispatched"
|
|
task.updated_at = _iso(now)
|
|
attempt.status = "dispatched"
|
|
session.flush()
|
|
return _leased_assignment_from_row(task)
|
|
|
|
def renew_lease(
|
|
self,
|
|
*,
|
|
task_id: str,
|
|
attempt: int,
|
|
lease_id: str,
|
|
host_id: str,
|
|
lease_expires_at: datetime,
|
|
now: datetime,
|
|
) -> str:
|
|
with self._sessions.begin() as session:
|
|
task = session.get(
|
|
ScheduledTaskRow,
|
|
task_id,
|
|
with_for_update=self.engine.dialect.name == "postgresql",
|
|
)
|
|
if task is None:
|
|
return "not_found"
|
|
if (
|
|
task.status not in {"assigned", "dispatched"}
|
|
or task.attempt_count != attempt
|
|
or task.lease_id != lease_id
|
|
or task.assigned_host_id != host_id
|
|
):
|
|
return "conflict"
|
|
current_expiry = _parse_dt(task.lease_expires_at)
|
|
if current_expiry is None or current_expiry <= now:
|
|
return "expired"
|
|
if lease_expires_at <= now:
|
|
return "conflict"
|
|
|
|
attempt_row = session.get(
|
|
TaskAttemptRow,
|
|
(task_id, attempt),
|
|
with_for_update=self.engine.dialect.name == "postgresql",
|
|
)
|
|
if (
|
|
attempt_row is None
|
|
or attempt_row.status not in {"assigned", "dispatched"}
|
|
or attempt_row.lease_id != lease_id
|
|
or attempt_row.host_id != host_id
|
|
):
|
|
return "conflict"
|
|
|
|
renewed_until = _iso(lease_expires_at)
|
|
task.lease_expires_at = renewed_until
|
|
task.updated_at = _iso(now)
|
|
attempt_row.lease_expires_at = renewed_until
|
|
return "renewed"
|
|
|
|
def record_task_result(
|
|
self,
|
|
*,
|
|
task_id: str,
|
|
attempt: int,
|
|
lease_id: str,
|
|
host_id: str,
|
|
status: str,
|
|
failure_reason: str | None,
|
|
terminal_result: dict[str, Any] | None,
|
|
completed_at: datetime,
|
|
) -> str:
|
|
with self._sessions.begin() as session:
|
|
task = session.get(
|
|
ScheduledTaskRow,
|
|
task_id,
|
|
with_for_update=self.engine.dialect.name == "postgresql",
|
|
)
|
|
if task is None:
|
|
return "conflict"
|
|
attempt_row = session.get(
|
|
TaskAttemptRow,
|
|
(task_id, attempt),
|
|
with_for_update=self.engine.dialect.name == "postgresql",
|
|
)
|
|
if (
|
|
attempt_row is None
|
|
or task.attempt_count != attempt
|
|
or task.lease_id != lease_id
|
|
or task.assigned_host_id != host_id
|
|
or attempt_row.lease_id != lease_id
|
|
or attempt_row.host_id != host_id
|
|
):
|
|
return "conflict"
|
|
|
|
if task.status in {"done", "failed"}:
|
|
if (
|
|
task.status == status
|
|
and task.failure_reason == failure_reason
|
|
and _parse_json_object(task.result_json) == terminal_result
|
|
and attempt_row.status == status
|
|
):
|
|
return "already_recorded"
|
|
return "conflict"
|
|
if task.status not in {"assigned", "dispatched"}:
|
|
return "conflict"
|
|
current_expiry = _parse_dt(task.lease_expires_at)
|
|
if current_expiry is None or current_expiry <= completed_at:
|
|
return "conflict"
|
|
if status not in {"done", "failed"}:
|
|
return "conflict"
|
|
|
|
result_json = (
|
|
json.dumps(terminal_result, ensure_ascii=False)
|
|
if terminal_result is not None
|
|
else None
|
|
)
|
|
completed_at_iso = _iso(completed_at)
|
|
task.status = status
|
|
task.failure_reason = failure_reason
|
|
task.result_json = result_json
|
|
task.updated_at = completed_at_iso
|
|
attempt_row.status = status
|
|
attempt_row.completed_at = completed_at_iso
|
|
attempt_row.failure_reason = failure_reason
|
|
attempt_row.result_json = result_json
|
|
return "recorded"
|
|
|
|
def reap_expired_leases(
|
|
self,
|
|
*,
|
|
now: datetime,
|
|
max_attempts: int,
|
|
) -> list[str]:
|
|
if max_attempts < 1:
|
|
raise ValueError("max_attempts must be at least 1")
|
|
|
|
with self._sessions.begin() as session:
|
|
statement = (
|
|
select(ScheduledTaskRow)
|
|
.where(
|
|
ScheduledTaskRow.status.in_(("assigned", "dispatched")),
|
|
ScheduledTaskRow.lease_expires_at.is_not(None),
|
|
ScheduledTaskRow.lease_expires_at <= _iso(now),
|
|
)
|
|
.order_by(ScheduledTaskRow.lease_expires_at, ScheduledTaskRow.id)
|
|
)
|
|
if self.engine.dialect.name == "postgresql":
|
|
statement = statement.with_for_update(skip_locked=True)
|
|
|
|
tasks = session.scalars(statement).all()
|
|
reaped_task_ids: list[str] = []
|
|
for task in tasks:
|
|
attempt_row = session.get(
|
|
TaskAttemptRow,
|
|
(task.id, task.attempt_count),
|
|
with_for_update=self.engine.dialect.name == "postgresql",
|
|
)
|
|
if attempt_row is None:
|
|
continue
|
|
|
|
attempt_row.status = "expired"
|
|
attempt_row.completed_at = _iso(now)
|
|
attempt_row.failure_reason = "lease expired"
|
|
task.updated_at = _iso(now)
|
|
task.lease_id = None
|
|
task.lease_expires_at = None
|
|
task.result_json = None
|
|
if task.attempt_count < max_attempts:
|
|
task.status = "queued"
|
|
task.assigned_host_id = None
|
|
task.assigned_device_id = None
|
|
task.failure_reason = None
|
|
else:
|
|
task.status = "failed"
|
|
task.failure_reason = (
|
|
f"lease expired after {task.attempt_count} attempts"
|
|
)
|
|
reaped_task_ids.append(task.id)
|
|
return reaped_task_ids
|
|
|
|
def list_task_attempts(self, task_id: str) -> list[Any]:
|
|
with self._sessions() as session:
|
|
rows = session.scalars(
|
|
select(TaskAttemptRow)
|
|
.where(TaskAttemptRow.task_id == task_id)
|
|
.order_by(TaskAttemptRow.attempt)
|
|
).all()
|
|
return [_task_attempt_from_row(row) for row in rows]
|
|
|
|
def health_check(self) -> None:
|
|
with self._sessions() as session:
|
|
session.execute(select(1))
|
|
|
|
def close(self) -> None:
|
|
self.engine.dispose()
|
|
|
|
|
|
def _iso(value: datetime) -> str:
|
|
return value.isoformat()
|
|
|
|
|
|
def _parse_dt(value: str | None) -> datetime | None:
|
|
if not value:
|
|
return None
|
|
try:
|
|
return datetime.fromisoformat(value)
|
|
except ValueError:
|
|
return None
|
|
|
|
|
|
def _host_from_row(row: HostRow) -> Any:
|
|
from cloud.pool import HostRegistration
|
|
|
|
return HostRegistration(
|
|
host_id=row.host_id,
|
|
address=row.address,
|
|
last_seen_at=_parse_dt(row.last_seen_at) or utc_now(),
|
|
)
|
|
|
|
|
|
def _device_from_row(row: PooledDeviceRow) -> Any:
|
|
from cloud.pool import PooledDevice
|
|
|
|
try:
|
|
tags = list(json.loads(row.capability_tags_json))
|
|
except TypeError, ValueError:
|
|
tags = []
|
|
return PooledDevice(
|
|
device_id=row.device_id,
|
|
host_id=row.host_id,
|
|
driver_type=row.driver_type,
|
|
status=row.status,
|
|
capability_tags=tags,
|
|
synced_at=_parse_dt(row.synced_at),
|
|
)
|
|
|
|
|
|
def _task_from_row(row: ScheduledTaskRow) -> Any:
|
|
from cloud.scheduler import ScheduledTask, TaskConstraints
|
|
|
|
try:
|
|
constraints_data = json.loads(row.constraints_json)
|
|
except TypeError, ValueError:
|
|
constraints_data = {}
|
|
terminal_result = _parse_json_object(row.result_json)
|
|
return ScheduledTask(
|
|
id=row.id,
|
|
goal=row.goal,
|
|
workflow_definition_id=row.workflow_definition_id,
|
|
constraints=TaskConstraints(
|
|
driver_type=constraints_data.get("driver_type"),
|
|
capability_tags=list(constraints_data.get("capability_tags") or []),
|
|
),
|
|
status=row.status,
|
|
assigned_device_id=row.assigned_device_id,
|
|
assigned_host_id=row.assigned_host_id,
|
|
attempt_count=row.attempt_count,
|
|
lease_id=row.lease_id,
|
|
lease_expires_at=_parse_dt(row.lease_expires_at),
|
|
terminal_result=terminal_result,
|
|
failure_reason=row.failure_reason,
|
|
updated_at=_parse_dt(row.updated_at),
|
|
created_at=_parse_dt(row.created_at) or utc_now(),
|
|
)
|
|
|
|
|
|
def _task_attempt_from_row(row: TaskAttemptRow) -> Any:
|
|
from cloud.repository import TaskAttemptRecord
|
|
|
|
return TaskAttemptRecord(
|
|
task_id=row.task_id,
|
|
attempt=row.attempt,
|
|
lease_id=row.lease_id,
|
|
host_id=row.host_id,
|
|
device_id=row.device_id,
|
|
status=row.status,
|
|
lease_expires_at=_parse_dt(row.lease_expires_at) or utc_now(),
|
|
created_at=_parse_dt(row.created_at) or utc_now(),
|
|
completed_at=_parse_dt(row.completed_at),
|
|
failure_reason=row.failure_reason,
|
|
terminal_result=_parse_json_object(row.result_json),
|
|
)
|
|
|
|
|
|
def _leased_assignment_from_row(row: ScheduledTaskRow) -> Any:
|
|
from cloud.repository import LeasedAssignment
|
|
|
|
lease_expires_at = _parse_dt(row.lease_expires_at)
|
|
if (
|
|
row.lease_id is None
|
|
or lease_expires_at is None
|
|
or row.assigned_host_id is None
|
|
or row.assigned_device_id is None
|
|
):
|
|
raise ValueError(f"task {row.id!r} does not contain a complete lease")
|
|
return LeasedAssignment(
|
|
task_id=row.id,
|
|
attempt=row.attempt_count,
|
|
lease_id=row.lease_id,
|
|
lease_expires_at=lease_expires_at,
|
|
host_id=row.assigned_host_id,
|
|
device_id=row.assigned_device_id,
|
|
goal=row.goal,
|
|
workflow_definition_id=row.workflow_definition_id,
|
|
)
|
|
|
|
|
|
def _parse_json_object(value: str | None) -> dict[str, Any] | None:
|
|
if value is None:
|
|
return None
|
|
try:
|
|
parsed = json.loads(value)
|
|
except TypeError, ValueError:
|
|
return None
|
|
return parsed if isinstance(parsed, dict) else None
|
|
|
|
|
|
def _plugin_from_row(row: PluginRow) -> tuple[Any, bool]:
|
|
from cloud.plugins import PluginManifest
|
|
|
|
return (
|
|
PluginManifest(
|
|
name=row.name,
|
|
version=row.version,
|
|
entry_point_kind=row.entry_point_kind,
|
|
target=row.target,
|
|
),
|
|
bool(row.wired),
|
|
)
|