karldns.de/karldns/store/records.py

from __future__ import annotations

import sqlite3
from dataclasses import dataclass
from datetime import UTC, datetime

from ..models import UpdateResult
from .base import StoreConnectionMixin
from .types import RecordRow, UpdateEventRow, record_row, update_event_row


@dataclass(frozen=True, slots=True)
class UpdateEventInput:
    hostname: str
    ipv4: str | None
    ipv6: str | None
    status: str
    detail: str
    source_ip: str


@dataclass(frozen=True, slots=True)
class RecordUpdateInput:
    hostname: str
    ipv4: str | None
    ipv6: str | None
    update_ipv4: bool = True
    update_ipv6: bool = True


def _update_event_input(event: UpdateEventInput | str, legacy_values: tuple[object, ...]) -> UpdateEventInput:
    if isinstance(event, UpdateEventInput):
        return event
    if len(legacy_values) != 5:
        raise TypeError("legacy update event values must be ipv4, ipv6, status, detail, source_ip")
    ipv4, ipv6, status, detail, source_ip = legacy_values
    return UpdateEventInput(
        hostname=event,
        ipv4=str(ipv4) if ipv4 is not None else None,
        ipv6=str(ipv6) if ipv6 is not None else None,
        status=str(status),
        detail=str(detail),
        source_ip=str(source_ip),
    )


def _planned_record_update(
    update: RecordUpdateInput,
    *,
    current_ipv4: str | None,
    current_ipv6: str | None,
    existing: bool,
) -> UpdateResult:
    effective_ipv4 = update.ipv4 if update.update_ipv4 else current_ipv4
    effective_ipv6 = update.ipv6 if update.update_ipv6 else current_ipv6
    return UpdateResult(
        hostname=update.hostname,
        ipv4=effective_ipv4,
        ipv6=effective_ipv6,
        changed=not existing or current_ipv4 != effective_ipv4 or current_ipv6 != effective_ipv6,
        update_ipv4=update.update_ipv4,
        update_ipv6=update.update_ipv6,
    )


class RecordsEventsMixin(StoreConnectionMixin):
    def upsert(
        self,
        hostname: str,
        ipv4: str | None,
        ipv6: str | None,
        update_ipv4: bool = True,
        update_ipv6: bool = True,
    ) -> UpdateResult:
        now = datetime.now(UTC).isoformat()
        update = RecordUpdateInput(hostname, ipv4, ipv6, update_ipv4, update_ipv6)
        with self.connect() as conn:
            existing = conn.execute("SELECT ipv4, ipv6 FROM records WHERE hostname = ?", (hostname,)).fetchone()
            current_ipv4 = existing["ipv4"] if existing else None
            current_ipv6 = existing["ipv6"] if existing else None
            result = _planned_record_update(
                update,
                current_ipv4=current_ipv4,
                current_ipv6=current_ipv6,
                existing=existing is not None,
            )
            conn.execute(
                """
                INSERT INTO records(hostname, ipv4, ipv6, updated_at)
                VALUES (?, ?, ?, ?)
                ON CONFLICT(hostname) DO UPDATE SET
                  ipv4=excluded.ipv4,
                  ipv6=excluded.ipv6,
                  updated_at=excluded.updated_at
                """,
                (result.hostname, result.ipv4, result.ipv6, now),
            )
        return result

    def commit_update_result(self, result: UpdateResult) -> UpdateResult:
        now = datetime.now(UTC).isoformat()
        with self.connect() as conn:
            conn.execute(
                """
                INSERT INTO records(hostname, ipv4, ipv6, updated_at)
                VALUES (?, ?, ?, ?)
                ON CONFLICT(hostname) DO UPDATE SET
                  ipv4=excluded.ipv4,
                  ipv6=excluded.ipv6,
                  updated_at=excluded.updated_at
                """,
                (result.hostname, result.ipv4, result.ipv6, now),
            )
        return result

    def preview_update(
        self,
        hostname: str,
        ipv4: str | None,
        ipv6: str | None,
        update_ipv4: bool = True,
        update_ipv6: bool = True,
    ) -> UpdateResult:
        update = RecordUpdateInput(hostname, ipv4, ipv6, update_ipv4, update_ipv6)
        with self.connect() as conn:
            existing = conn.execute("SELECT ipv4, ipv6 FROM records WHERE hostname = ?", (hostname,)).fetchone()
        current_ipv4 = existing["ipv4"] if existing else None
        current_ipv6 = existing["ipv6"] if existing else None
        return _planned_record_update(
            update,
            current_ipv4=current_ipv4,
            current_ipv6=current_ipv6,
            existing=existing is not None,
        )

    def list_records(self) -> list[RecordRow]:
        return self.list_records_page(limit=10_000, offset=0)

    def list_records_page(self, limit: int = 100, offset: int = 0) -> list[RecordRow]:
        with self.connect() as conn:
            rows = conn.execute(
                """
                SELECT hostname, ipv4, ipv6, updated_at
                FROM records
                ORDER BY hostname
                LIMIT ? OFFSET ?
                """,
                (limit, offset),
            ).fetchall()
        return [record_row(row) for row in rows]

    def count_records(self) -> int:
        with self.connect() as conn:
            row = conn.execute("SELECT COUNT(*) FROM records").fetchone()
        return int(row[0])

    def log_update_event(self, event: UpdateEventInput | str, *legacy_values: object) -> None:
        event = _update_event_input(event, legacy_values)
        with self.connect() as conn:
            conn.execute(
                """
                INSERT INTO update_events(hostname, ipv4, ipv6, status, detail, source_ip, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    event.hostname,
                    event.ipv4,
                    event.ipv6,
                    event.status,
                    event.detail[:500],
                    event.source_ip,
                    datetime.now(UTC).isoformat(),
                ),
            )

    def log_update_rollup(self, event: UpdateEventInput | str, *legacy_values: object) -> None:
        event = _update_event_input(event, legacy_values)
        now = datetime.now(UTC).isoformat()
        with self.connect() as conn:
            conn.execute(
                """
                INSERT INTO update_event_rollups(
                  hostname, ipv4, ipv6, status, detail, source_ip, first_seen_at, last_seen_at, count
                )
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, 1)
                ON CONFLICT(hostname, ipv4, ipv6, status, detail, source_ip) DO UPDATE SET
                  last_seen_at=excluded.last_seen_at,
                  count=count + 1
                """,
                (
                    event.hostname,
                    event.ipv4 or "",
                    event.ipv6 or "",
                    event.status,
                    event.detail[:500],
                    event.source_ip,
                    now,
                    now,
                ),
            )

    def list_update_events(self, limit: int = 100, offset: int = 0) -> list[UpdateEventRow]:
        return self.list_update_events_page(limit=limit, offset=offset)

    def list_update_events_page(self, limit: int = 100, offset: int = 0) -> list[UpdateEventRow]:
        with self.connect() as conn:
            rows = conn.execute(
                """
                SELECT hostname, ipv4, ipv6, status, detail, source_ip, created_at, count
                FROM (
                    SELECT hostname, ipv4, ipv6, status, detail, source_ip, created_at, 1 AS count
                    FROM update_events
                    UNION ALL
                    SELECT hostname, ipv4, ipv6, status, detail, source_ip, last_seen_at AS created_at, count
                    FROM update_event_rollups
                )
                ORDER BY created_at DESC
                LIMIT ? OFFSET ?
                """,
                (limit, offset),
            ).fetchall()
        return [update_event_row(row) for row in rows]

    def list_update_events_before(self, limit: int = 100, before: str | None = None) -> list[UpdateEventRow]:
        where_clause = "WHERE created_at < ?" if before else ""
        params: tuple[object, ...] = (before, limit) if before else (limit,)
        with self.connect() as conn:
            rows = conn.execute(
                f"""
                SELECT hostname, ipv4, ipv6, status, detail, source_ip, created_at, count
                FROM (
                    SELECT hostname, ipv4, ipv6, status, detail, source_ip, created_at, 1 AS count
                    FROM update_events
                    UNION ALL
                    SELECT hostname, ipv4, ipv6, status, detail, source_ip, last_seen_at AS created_at, count
                    FROM update_event_rollups
                )
                {where_clause}
                ORDER BY created_at DESC
                LIMIT ?
                """,
                params,
            ).fetchall()
        return [update_event_row(row) for row in rows]

    def count_update_events(self) -> int:
        with self.connect() as conn:
            event_count = conn.execute("SELECT COUNT(*) FROM update_events").fetchone()[0]
            rollup_count = conn.execute("SELECT COUNT(*) FROM update_event_rollups").fetchone()[0]
        return int(event_count) + int(rollup_count)

    def list_update_events_for_hostname(self, hostname: str, limit: int = 25) -> list[UpdateEventRow]:
        with self.connect() as conn:
            event_rows = conn.execute(
                """
                SELECT hostname, ipv4, ipv6, status, detail, source_ip, created_at, 1 AS count
                FROM update_events
                WHERE hostname = ?
                ORDER BY created_at DESC
                LIMIT ?
                """,
                (hostname, limit),
            ).fetchall()
            rollup_rows = conn.execute(
                """
                SELECT hostname, ipv4, ipv6, status, detail, source_ip, last_seen_at AS created_at, count
                FROM update_event_rollups
                WHERE hostname = ?
                ORDER BY last_seen_at DESC
                LIMIT ?
                """,
                (hostname, limit),
            ).fetchall()
        return _merge_event_rows(event_rows, rollup_rows, limit, 0)


def _merge_event_rows(
    event_rows: list[sqlite3.Row],
    rollup_rows: list[sqlite3.Row],
    limit: int,
    offset: int,
) -> list[UpdateEventRow]:
    rows = [update_event_row(row) for row in event_rows]
    rows.extend(update_event_row(row) for row in rollup_rows)
    rows.sort(key=lambda row: str(row["created_at"]), reverse=True)
    return rows[offset : offset + limit]