Files
agentic-mobile-control/packages/cloud-platform/cloud/sql_repository.py
T

695 lines
24 KiB
Python

from __future__ import annotations
import json
import logging
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 cloud.observability import current_correlation_id
from core.models import utc_now
logger = logging.getLogger(__name__)
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],
*,
allow_device_takeover: bool = False,
) -> None:
with self._sessions.begin() as session:
session.execute(
delete(PooledDeviceRow).where(PooledDeviceRow.host_id == host_id)
)
if allow_device_takeover:
device_ids = [device.device_id for device in devices]
if device_ids:
session.execute(
delete(PooledDeviceRow).where(
PooledDeviceRow.host_id != host_id,
PooledDeviceRow.device_id.in_(device_ids),
)
)
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()
_log_task_lifecycle("assigned", task)
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()
_log_task_lifecycle("claimed", task)
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
_log_task_lifecycle("renewed", task)
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
_log_task_lifecycle("completed" if status == "done" else "failed", task)
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"
)
_log_task_lifecycle(
"retried" if task.status == "queued" else "failed",
task,
)
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 _log_task_lifecycle(event: str, task: ScheduledTaskRow) -> None:
logger.info(
"cloud task lifecycle",
extra={
"event": event,
"correlation_id": current_correlation_id(),
"task_id": task.id,
"host_id": task.assigned_host_id,
"device_id": task.assigned_device_id,
"attempt": task.attempt_count,
"lease_id": task.lease_id,
},
)
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),
)