dns-sync/dns_sync/executor.py

from collections.abc import Iterable
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Protocol

from .adapters import (
    CERTIFICATES,
    DNS_RECORDS,
    PROXY_HOSTS,
    Capability,
    CapabilityOperation,
    Endpoint,
    EndpointPlan,
    NormalizedObject,
    PlanAction,
    PlanOperation,
)
from .persistence import SQLiteStateStore


_OBJECT_TYPE_ORDER = {
    CERTIFICATES: 0,
    PROXY_HOSTS: 1,
    DNS_RECORDS: 2,
}


class ExecutionAdapter(Protocol):
    endpoint: Endpoint
    capabilities: frozenset[Capability]

    def apply(self, action: PlanAction) -> str | None: ...

    def verify(self, object: NormalizedObject) -> bool: ...


@dataclass(frozen=True)
class ExecutedAction:
    endpoint_id: str
    operation: PlanOperation
    object_type: str
    logical_key: str


@dataclass(frozen=True)
class ExecutionResult:
    applied: tuple[ExecutedAction, ...]
    skipped: tuple[PlanAction, ...]


def _action_sort_key(
    action: PlanAction,
) -> tuple[int, str, str]:
    desired = action.desired
    if desired is None:
        return (999, action.endpoint_id, "")

    return (
        _OBJECT_TYPE_ORDER.get(desired.object_type, 100),
        action.endpoint_id,
        desired.logical_key,
    )


def _required_capability(
    adapter: ExecutionAdapter,
    object_type: str,
    operation: CapabilityOperation,
) -> None:
    capability = Capability(object_type, operation)
    if capability not in adapter.capabilities:
        raise RuntimeError(
            f"{adapter.endpoint.identity} lacks "
            f"{operation.value} capability for {object_type}"
        )


def _preflight(
    adapters: Iterable[ExecutionAdapter],
    plans: Iterable[EndpointPlan],
) -> tuple[
    dict[str, ExecutionAdapter],
    list[PlanAction],
    list[PlanAction],
]:
    adapters_by_endpoint: dict[str, ExecutionAdapter] = {}
    for adapter in adapters:
        endpoint_id = adapter.endpoint.identity
        if endpoint_id in adapters_by_endpoint:
            raise RuntimeError(
                f"Duplicate execution adapter: {endpoint_id}"
            )
        adapters_by_endpoint[endpoint_id] = adapter

    executable: list[PlanAction] = []
    skipped: list[PlanAction] = []

    for plan in plans:
        adapter = adapters_by_endpoint.get(plan.endpoint_id)
        if adapter is None:
            raise RuntimeError(
                f"No adapter for endpoint plan {plan.endpoint_id}"
            )

        for action in plan.actions:
            if action.endpoint_id != plan.endpoint_id:
                raise RuntimeError(
                    "Plan action endpoint does not match its endpoint plan"
                )

            if action.operation in {
                PlanOperation.NOOP,
                PlanOperation.SKIP,
            }:
                skipped.append(action)
                continue

            if action.operation not in {
                PlanOperation.CREATE,
                PlanOperation.UPDATE,
            }:
                raise RuntimeError(
                    "Unsupported executable plan operation: "
                    f"{action.operation.value}"
                )

            if action.desired is None:
                raise RuntimeError(
                    "Executable plan action requires desired state"
                )

            object_type = action.desired.object_type
            _required_capability(
                adapter,
                object_type,
                CapabilityOperation.APPLY,
            )
            _required_capability(
                adapter,
                object_type,
                CapabilityOperation.VERIFY,
            )
            executable.append(action)

    return adapters_by_endpoint, executable, skipped


def execute_endpoint_plans(
    adapters: Iterable[ExecutionAdapter],
    plans: Iterable[EndpointPlan],
    store: SQLiteStateStore,
    *,
    executed_at: datetime | None = None,
) -> ExecutionResult:
    timestamp = executed_at or datetime.now(timezone.utc)
    if timestamp.tzinfo is None:
        raise RuntimeError(
            "Execution timestamp must be timezone-aware"
        )
    timestamp = timestamp.astimezone(timezone.utc)

    adapters_by_endpoint, executable, skipped = _preflight(
        adapters,
        plans,
    )
    applied: list[ExecutedAction] = []

    for action in sorted(executable, key=_action_sort_key):
        adapter = adapters_by_endpoint[action.endpoint_id]
        desired = action.desired

        adapter.apply(action)
        store.record_applied(
            action.endpoint_id,
            desired,
            timestamp,
        )

        try:
            verified = adapter.verify(desired)
        except Exception as exc:
            error = str(exc)
            store.record_verification_failure(
                action.endpoint_id,
                desired,
                error,
                timestamp,
            )
            raise RuntimeError(
                f"{action.endpoint_id} "
                f"{desired.object_type} "
                f"{desired.logical_key}: verification failed"
            ) from exc

        if not verified:
            error = "adapter returned false"
            store.record_verification_failure(
                action.endpoint_id,
                desired,
                error,
                timestamp,
            )
            raise RuntimeError(
                f"{action.endpoint_id} "
                f"{desired.object_type} "
                f"{desired.logical_key}: verification failed"
            )

        store.record_verification_success(
            action.endpoint_id,
            desired,
            timestamp,
        )
        applied.append(
            ExecutedAction(
                endpoint_id=action.endpoint_id,
                operation=action.operation,
                object_type=desired.object_type,
                logical_key=desired.logical_key,
            )
        )

    return ExecutionResult(
        applied=tuple(applied),
        skipped=tuple(skipped),
    )