dns-sync/dns_sync/reconciliation.py

from collections import defaultdict
from collections.abc import Iterable, Mapping, Sequence
from typing import Protocol

from .adapters import (
    Capability,
    CapabilityOperation,
    DNS_RECORDS,
    Endpoint,
    EndpointPlan,
    NormalizedObject,
    Observation,
    PlanAction,
    PlanOperation,
)


ObservationGroupKey = tuple[str, str]


class ReconciliationConflict(RuntimeError):
    """Raised when observations cannot be reconciled safely."""


class PlanningConflict(RuntimeError):
    """Raised when an endpoint cannot be planned safely."""


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


class ReconciliationPolicy(Protocol):
    object_type: str

    def reconcile(
        self,
        logical_key: str,
        observations: Sequence[Observation],
    ) -> NormalizedObject: ...


def _observation_sort_key(
    observation: Observation,
) -> tuple[str, str, str]:
    return (
        observation.endpoint_id,
        observation.object.local_key or "",
        observation.observed_at.isoformat(),
    )


def group_observations(
    observations: Iterable[Observation],
) -> dict[ObservationGroupKey, tuple[Observation, ...]]:
    grouped: dict[
        ObservationGroupKey,
        list[Observation],
    ] = defaultdict(list)

    for observation in observations:
        normalized = observation.object
        if not observation.endpoint_id:
            raise RuntimeError("Observation endpoint_id cannot be empty")
        if not normalized.object_type:
            raise RuntimeError("Observation object_type cannot be empty")
        if not normalized.logical_key:
            raise RuntimeError("Observation logical_key cannot be empty")

        key = (
            normalized.object_type,
            normalized.logical_key,
        )
        grouped[key].append(observation)

    return {
        key: tuple(sorted(rows, key=_observation_sort_key))
        for key, rows in sorted(grouped.items())
    }


class DnsUnionPolicy:
    object_type = DNS_RECORDS

    def reconcile(
        self,
        logical_key: str,
        observations: Sequence[Observation],
    ) -> NormalizedObject:
        if not observations:
            raise RuntimeError(
                "DNS union policy requires at least one observation"
            )

        first = observations[0].object
        conflicting_endpoints = {
            observation.endpoint_id
            for observation in observations
            if observation.object.payload != first.payload
        }
        if conflicting_endpoints:
            conflicting_endpoints.add(observations[0].endpoint_id)
            endpoint_list = ", ".join(
                sorted(conflicting_endpoints)
            )
            raise ReconciliationConflict(
                f"Conflicting {DNS_RECORDS} observations for "
                f"{logical_key!r} from {endpoint_list}"
            )

        return NormalizedObject(
            object_type=DNS_RECORDS,
            logical_key=logical_key,
            payload=dict(first.payload),
        )


DEFAULT_POLICIES: Mapping[str, ReconciliationPolicy] = {
    DNS_RECORDS: DnsUnionPolicy(),
}


def reconcile_observations(
    observations: Iterable[Observation],
    policies: Mapping[str, ReconciliationPolicy] | None = None,
) -> list[NormalizedObject]:
    selected_policies = (
        DEFAULT_POLICIES
        if policies is None
        else policies
    )
    desired: list[NormalizedObject] = []

    for (object_type, logical_key), rows in (
        group_observations(observations).items()
    ):
        policy = selected_policies.get(object_type)
        if policy is None:
            raise RuntimeError(
                f"No reconciliation policy for {object_type!r}"
            )
        if policy.object_type != object_type:
            raise RuntimeError(
                f"Policy for {object_type!r} advertises "
                f"{policy.object_type!r}"
            )

        desired.append(policy.reconcile(logical_key, rows))

    return desired


def build_endpoint_plans(
    adapters: Iterable[PlanningAdapter],
    desired_objects: Iterable[NormalizedObject],
    observations: Iterable[Observation],
) -> tuple[EndpointPlan, ...]:
    desired_by_identity: dict[
        tuple[str, str],
        NormalizedObject,
    ] = {}

    for desired in desired_objects:
        if desired.local_key is not None:
            raise RuntimeError(
                "Desired objects in plans cannot contain "
                "an endpoint-local key"
            )

        identity = (
            desired.object_type,
            desired.logical_key,
        )
        if identity in desired_by_identity:
            raise RuntimeError(
                "Duplicate desired object identity: "
                f"{desired.object_type} {desired.logical_key}"
            )
        desired_by_identity[identity] = desired

    current_by_identity: dict[
        tuple[str, str, str],
        list[Observation],
    ] = defaultdict(list)
    for observation in observations:
        current_by_identity[
            (
                observation.endpoint_id,
                observation.object.object_type,
                observation.object.logical_key,
            )
        ].append(observation)

    adapters_by_endpoint: dict[str, PlanningAdapter] = {}
    for adapter in adapters:
        endpoint_id = adapter.endpoint.identity
        if endpoint_id in adapters_by_endpoint:
            raise RuntimeError(
                f"Duplicate adapter endpoint: {endpoint_id}"
            )
        adapters_by_endpoint[endpoint_id] = adapter

    plans: list[EndpointPlan] = []
    for endpoint_id, adapter in sorted(
        adapters_by_endpoint.items()
    ):
        apply_types = {
            capability.object_type
            for capability in adapter.capabilities
            if capability.operation == CapabilityOperation.APPLY
        }
        if not apply_types:
            continue

        actions: list[PlanAction] = []
        for identity, desired in sorted(
            desired_by_identity.items()
        ):
            object_type, logical_key = identity
            if object_type not in apply_types:
                continue

            current_rows = current_by_identity.get(
                (endpoint_id, object_type, logical_key),
                [],
            )
            if len(current_rows) > 1:
                raise PlanningConflict(
                    f"Multiple endpoint objects for {endpoint_id} "
                    f"{object_type} {logical_key!r}"
                )

            current = (
                current_rows[0].object
                if current_rows
                else None
            )
            if current is None:
                operation = PlanOperation.CREATE
            elif current.payload == desired.payload:
                operation = PlanOperation.NOOP
            else:
                operation = PlanOperation.UPDATE

            actions.append(
                PlanAction(
                    endpoint_id=endpoint_id,
                    operation=operation,
                    desired=desired,
                    current=current,
                )
            )

        plans.append(
            EndpointPlan(
                endpoint_id=endpoint_id,
                actions=tuple(actions),
            )
        )

    return tuple(plans)