dns-sync/dns_sync/orchestrator.py

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

from .adapters import (
    DNS_RECORDS,
    PROXY_HOSTS,
    Capability,
    CapabilityOperation,
    DiscoveredObject,
    Endpoint,
    EndpointPlan,
    NormalizedObject,
    Observation,
)
from .engine import hostname_matches_domain
from .object_normalizer import normalize_observations
from .persistence import SQLiteStateStore
from .projections import project_proxy_hosts_to_dns
from .reconciliation import (
    build_endpoint_plans,
    reconcile_observations,
)


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

    def connect(self) -> None: ...

    def discover(
        self,
        object_type: str,
    ) -> list[DiscoveredObject]: ...


@dataclass(frozen=True)
class DiscoveryFailure:
    endpoint_id: str
    object_type: str
    error: str


@dataclass(frozen=True)
class ReconciliationRun:
    observations: tuple[Observation, ...]
    desired: tuple[NormalizedObject, ...]
    plans: tuple[EndpointPlan, ...]
    failures: tuple[DiscoveryFailure, ...]


def _discovery_types(
    adapter: ReconciliationAdapter,
) -> tuple[str, ...]:
    return tuple(
        sorted(
            capability.object_type
            for capability in adapter.capabilities
            if capability.operation == CapabilityOperation.DISCOVER
        )
    )


def _configured_observations(
    store: SQLiteStateStore,
    endpoint_ids: set[str],
) -> list[Observation]:
    return [
        observation
        for observation in store.load_observations()
        if observation.endpoint_id in endpoint_ids
    ]


def _filter_projected_dns(
    observations: Iterable[Observation],
    *,
    domain_matching: bool,
    domain: str | None,
) -> list[Observation]:
    rows = list(observations)
    if not domain_matching:
        return rows
    if not domain:
        raise RuntimeError(
            "Domain filtering enabled without a configured DOMAIN"
        )

    return [
        observation
        for observation in rows
        if hostname_matches_domain(
            observation.object.logical_key,
            domain,
        )
    ]


def discover_reconcile_plan(
    adapters: Iterable[ReconciliationAdapter],
    store: SQLiteStateStore,
    *,
    target_ip: str,
    include_disabled: bool = False,
    domain_matching: bool = False,
    domain: str | None = None,
    run_at: datetime | None = None,
) -> ReconciliationRun:
    timestamp = run_at or datetime.now(timezone.utc)
    if timestamp.tzinfo is None:
        raise RuntimeError(
            "Reconciliation run timestamp must be timezone-aware"
        )
    timestamp = timestamp.astimezone(timezone.utc)

    adapter_list = list(adapters)
    adapters_by_endpoint: dict[str, ReconciliationAdapter] = {}
    for adapter in adapter_list:
        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

    failures: list[DiscoveryFailure] = []

    for endpoint_id, adapter in sorted(
        adapters_by_endpoint.items()
    ):
        object_types = _discovery_types(adapter)
        if not object_types:
            continue

        try:
            adapter.connect()
        except Exception as exc:
            for object_type in object_types:
                error = str(exc)
                store.record_discovery_failure(
                    endpoint_id,
                    object_type,
                    error,
                    timestamp,
                )
                failures.append(
                    DiscoveryFailure(
                        endpoint_id,
                        object_type,
                        error,
                    )
                )
            continue

        for object_type in object_types:
            try:
                discovered = adapter.discover(object_type)
                observations = normalize_observations(
                    endpoint_id,
                    discovered,
                    timestamp,
                )
                store.replace_observations(
                    endpoint_id,
                    object_type,
                    observations,
                    timestamp,
                )
            except Exception as exc:
                error = str(exc)
                store.record_discovery_failure(
                    endpoint_id,
                    object_type,
                    error,
                    timestamp,
                )
                failures.append(
                    DiscoveryFailure(
                        endpoint_id,
                        object_type,
                        error,
                    )
                )

    observations = _configured_observations(
        store,
        set(adapters_by_endpoint),
    )
    proxy_observations = [
        observation
        for observation in observations
        if observation.object.object_type == PROXY_HOSTS
    ]
    projected_dns = project_proxy_hosts_to_dns(
        proxy_observations,
        target_ip,
        include_disabled=include_disabled,
    )
    projected_dns = _filter_projected_dns(
        projected_dns,
        domain_matching=domain_matching,
        domain=domain,
    )

    desired = reconcile_observations(projected_dns)
    store.replace_desired_state(
        DNS_RECORDS,
        desired,
        timestamp,
    )
    plans = build_endpoint_plans(
        adapter_list,
        desired,
        observations,
    )

    return ReconciliationRun(
        observations=tuple(observations),
        desired=tuple(desired),
        plans=plans,
        failures=tuple(failures),
    )