dns-sync/dns_sync/executor.py

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

from .adapters import (
    ApplyResult,
    Capability,
    CapabilityOperation,
    Endpoint,
    EndpointPlan,
    NormalizedObject,
    ObjectReference,
    PlanAction,
    PlanOperation,
    RetireRequest,
)
from .dependencies import order_plan_actions
from .persistence import SQLiteStateStore


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

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

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

    def retire(self, request: RetireRequest) -> None: ...


@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, ...]


@dataclass
class _LocalReferenceIndex:
    by_logical: dict[
        tuple[str, str],
        set[str],
    ]
    by_local: set[tuple[str, str]]


def _remember_local_object(
    index: _LocalReferenceIndex,
    object: NormalizedObject,
    *,
    authoritative: bool = False,
) -> None:
    local_key = object.local_key
    if local_key is None:
        return

    identity = (
        object.object_type,
        object.logical_key,
    )

    if authoritative:
        previous = index.by_logical.get(identity, set())
        for previous_local_key in previous:
            index.by_local.discard(
                (
                    object.object_type,
                    previous_local_key,
                )
            )
        index.by_logical[identity] = {local_key}
    else:
        index.by_logical.setdefault(
            identity,
            set(),
        ).add(local_key)

    index.by_local.add(
        (
            object.object_type,
            local_key,
        )
    )


def _local_reference_index(
    store: SQLiteStateStore,
    endpoint_id: str,
) -> _LocalReferenceIndex:
    index = _LocalReferenceIndex(
        by_logical={},
        by_local=set(),
    )
    observed_identities: set[
        tuple[str, str]
    ] = set()

    for observation in store.load_observations():
        if observation.endpoint_id != endpoint_id:
            continue

        object = observation.object
        if object.local_key is None:
            continue

        identity = (
            object.object_type,
            object.logical_key,
        )
        observed_identities.add(identity)
        _remember_local_object(index, object)

    for applied in store.load_applied_state(
        endpoint_id=endpoint_id,
    ):
        if not applied.is_verified:
            continue

        object = applied.object
        if object.local_key is None:
            continue

        identity = (
            object.object_type,
            object.logical_key,
        )

        # Endpoint observations are authoritative when present.
        # Verified applied state is only a fallback when this
        # logical identity is absent from the observation snapshot.
        if identity in observed_identities:
            continue

        _remember_local_object(index, object)

    return index


def _resolved_reference(
    endpoint_id: str,
    reference: ObjectReference,
    local_references: _LocalReferenceIndex,
) -> ObjectReference:
    if reference.local_key is not None:
        local_identity = (
            reference.object_type,
            reference.local_key,
        )

        if local_identity not in local_references.by_local:
            raise RuntimeError(
                "Unresolved endpoint-local reference for "
                f"{endpoint_id}: "
                f"{reference.relation} -> "
                f"{reference.object_type} "
                f"local key {reference.local_key!r}"
            )

        return reference

    identity = (
        reference.object_type,
        reference.logical_key,
    )
    candidates = local_references.by_logical.get(
        identity,
        set(),
    )

    if not candidates:
        raise RuntimeError(
            "Unresolved endpoint-local reference for "
            f"{endpoint_id}: "
            f"{reference.relation} -> "
            f"{reference.object_type} "
            f"{reference.logical_key!r}"
        )

    if len(candidates) > 1:
        raise RuntimeError(
            "Ambiguous endpoint-local reference for "
            f"{endpoint_id}: "
            f"{reference.relation} -> "
            f"{reference.object_type} "
            f"{reference.logical_key!r}: "
            f"{', '.join(sorted(candidates))}"
        )

    local_key = next(iter(candidates))

    return ObjectReference(
        relation=reference.relation,
        object_type=reference.object_type,
        local_key=local_key,
    )


def _resolve_action_references(
    action: PlanAction,
    local_references: _LocalReferenceIndex,
) -> PlanAction:
    desired = action.desired
    if desired is None:
        raise RuntimeError(
            "Executable plan action requires desired state"
        )

    if not desired.references:
        return action

    references = tuple(
        _resolved_reference(
            action.endpoint_id,
            reference,
            local_references,
        )
        for reference in desired.references
    )

    resolved_desired = replace(
        desired,
        references=references,
    )

    return replace(
        action,
        desired=resolved_desired,
    )


def _validated_apply_result(
    action: PlanAction,
    result: ApplyResult | None,
) -> ApplyResult:
    desired = action.desired
    if desired is None:
        raise RuntimeError(
            "Executable plan action requires desired state"
        )

    if not isinstance(result, ApplyResult):
        raise RuntimeError(
            f"{action.endpoint_id} "
            f"{desired.object_type} "
            f"{desired.logical_key}: "
            "adapter apply did not return ApplyResult"
        )

    applied = result.object

    if applied.object_type != desired.object_type:
        raise RuntimeError(
            "Applied object type does not match desired object type"
        )

    if applied.logical_key != desired.logical_key:
        raise RuntimeError(
            "Applied object logical key does not match desired "
            "logical key"
        )

    if applied.references != desired.references:
        raise RuntimeError(
            "Applied object references do not match resolved "
            "desired references"
        )

    if (
        action.operation != PlanOperation.REPLACE
        and result.retirements
    ):
        raise RuntimeError(
            "Only REPLACE actions may request deferred retirement"
        )

    for retirement in result.retirements:
        retired = retirement.object

        if retired.local_key is None:
            raise RuntimeError(
                "Deferred retirement requires endpoint-local identity"
            )

        if (
            retired.object_type == applied.object_type
            and retired.local_key == applied.local_key
        ):
            raise RuntimeError(
                "Deferred retirement cannot target applied object"
            )

    return result


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],
    list[tuple[str, RetireRequest]],
]:
    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] = []
    planned_retirements: list[
        tuple[str, RetireRequest]
    ] = []

    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 retirement in plan.retirements:
            planned_retirements.append(
                (
                    plan.endpoint_id,
                    retirement,
                )
            )

        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,
                PlanOperation.REPLACE,
            }:
                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,
        planned_retirements,
    )


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,
        planned_retirements,
    ) = _preflight(
        adapters,
        plans,
    )

    applied: list[ExecutedAction] = []
    pending_retirements: list[
        tuple[str, RetireRequest]
    ] = list(planned_retirements)

    local_references: dict[
        str,
        _LocalReferenceIndex,
    ] = {}

    for endpoint_id in adapters_by_endpoint:
        local_references[endpoint_id] = (
            _local_reference_index(
                store,
                endpoint_id,
            )
        )

    for action in skipped:
        if (
            action.operation == PlanOperation.NOOP
            and action.current is not None
            and action.current.local_key is not None
        ):
            _remember_local_object(
                local_references[action.endpoint_id],
                action.current,
                authoritative=True,
            )

    for action in order_plan_actions(executable):
        adapter = adapters_by_endpoint[action.endpoint_id]
        resolved_action = _resolve_action_references(
            action,
            local_references[action.endpoint_id],
        )
        desired = resolved_action.desired

        result = _validated_apply_result(
            resolved_action,
            adapter.apply(resolved_action),
        )
        applied_object = result.object

        store.record_applied(
            action.endpoint_id,
            applied_object,
            timestamp,
        )

        try:
            verified = adapter.verify(applied_object)
        except Exception as exc:
            error = str(exc)
            store.record_verification_failure(
                action.endpoint_id,
                applied_object,
                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,
                applied_object,
                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,
            applied_object,
            timestamp,
        )

        _remember_local_object(
            local_references[action.endpoint_id],
            applied_object,
            authoritative=True,
        )

        for retirement in result.retirements:
            pending_retirements.append(
                (
                    action.endpoint_id,
                    retirement,
                )
            )

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

    retired_identities: set[
        tuple[str, str, str]
    ] = set()

    for endpoint_id, retirement in pending_retirements:
        retired = retirement.object
        identity = (
            endpoint_id,
            retired.object_type,
            retired.local_key,
        )

        if identity in retired_identities:
            continue

        retired_identities.add(identity)

        adapter = adapters_by_endpoint[endpoint_id]

        retire = getattr(adapter, "retire", None)

        if retire is None:
            raise RuntimeError(
                f"{endpoint_id} adapter does not support "
                "deferred retirement"
            )

        retire(retirement)

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