solfuse-privacy-preserving-federated-orc/src/solfuse/registry.py

from __future__ import annotations

import json
import sqlite3
from pathlib import Path
from typing import Any

from jsonschema import ValidationError, validate
from packaging.version import InvalidVersion, Version

from .models import AdapterContract


class GraphOfContractsRegistry:
    def __init__(self, path: str | Path = ":memory:") -> None:
        self.path = str(path)
        self._conn = sqlite3.connect(self.path)
        self._conn.row_factory = sqlite3.Row
        self._ensure_schema()

    def _ensure_schema(self) -> None:
        self._conn.execute(
            """
            CREATE TABLE IF NOT EXISTS contracts (
                name TEXT NOT NULL,
                version TEXT NOT NULL,
                adapter_type TEXT NOT NULL,
                input_schema TEXT NOT NULL,
                output_schema TEXT NOT NULL,
                capabilities TEXT NOT NULL,
                metadata TEXT NOT NULL,
                fingerprint TEXT NOT NULL,
                PRIMARY KEY (name, version)
            )
            """
        )
        self._conn.execute(
            """
            CREATE TABLE IF NOT EXISTS edges (
                source_name TEXT NOT NULL,
                source_version TEXT NOT NULL,
                target_name TEXT NOT NULL,
                target_version TEXT NOT NULL,
                relation TEXT NOT NULL,
                PRIMARY KEY (source_name, source_version, target_name, target_version, relation)
            )
            """
        )
        self._conn.execute(
            """
            CREATE TABLE IF NOT EXISTS conformance_checks (
                contract_name TEXT NOT NULL,
                contract_version TEXT NOT NULL,
                payload_kind TEXT NOT NULL,
                passed INTEGER NOT NULL,
                error TEXT,
                payload TEXT NOT NULL,
                created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
            )
            """
        )
        self._conn.commit()

    def register(self, contract: AdapterContract) -> AdapterContract:
        self._conn.execute(
            """
            INSERT OR REPLACE INTO contracts
            (name, version, adapter_type, input_schema, output_schema, capabilities, metadata, fingerprint)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                contract.name,
                contract.version,
                contract.adapter_type,
                json.dumps(contract.input_schema, sort_keys=True),
                json.dumps(contract.output_schema, sort_keys=True),
                json.dumps(contract.capabilities, sort_keys=True),
                json.dumps(contract.metadata, sort_keys=True),
                contract.fingerprint,
            ),
        )
        self._conn.commit()
        return contract

    def link(self, source: tuple[str, str], target: tuple[str, str], relation: str = "compatible_with") -> None:
        self._conn.execute(
            "INSERT OR REPLACE INTO edges VALUES (?, ?, ?, ?, ?)",
            (source[0], source[1], target[0], target[1], relation),
        )
        self._conn.commit()

    def latest(self, name: str) -> AdapterContract:
        versions = self.versions(name)
        if not versions:
            raise KeyError(name)
        latest_version = versions[-1]
        return self.get(name, latest_version)

    def get(self, name: str, version: str) -> AdapterContract:
        row = self._conn.execute(
            "SELECT * FROM contracts WHERE name = ? AND version = ?",
            (name, version),
        ).fetchone()
        if row is None:
            raise KeyError(f"{name}@{version}")
        return self._row_to_contract(row)

    def versions(self, name: str) -> list[str]:
        rows = self._conn.execute("SELECT version FROM contracts WHERE name = ? ORDER BY version", (name,)).fetchall()
        versions = [row["version"] for row in rows]
        try:
            return sorted(versions, key=Version)
        except InvalidVersion:
            return sorted(versions)

    def validate_payload(self, name: str, version: str, payload: dict[str, Any], *, kind: str = "input") -> None:
        contract = self.get(name, version)
        schema = contract.input_schema if kind == "input" else contract.output_schema
        try:
            validate(instance=payload, schema=schema)
            passed = True
            error = None
        except ValidationError as exc:
            passed = False
            error = str(exc)
            self._record_conformance(contract, kind, payload, passed, error)
            raise
        self._record_conformance(contract, kind, payload, passed, error)

    def _record_conformance(self, contract: AdapterContract, kind: str, payload: dict[str, Any], passed: bool, error: str | None) -> None:
        self._conn.execute(
            """
            INSERT INTO conformance_checks
            (contract_name, contract_version, payload_kind, passed, error, payload)
            VALUES (?, ?, ?, ?, ?, ?)
            """,
            (contract.name, contract.version, kind, int(passed), error, json.dumps(payload, sort_keys=True)),
        )
        self._conn.commit()

    def conformance_report(self, name: str) -> dict[str, Any]:
        rows = self._conn.execute(
            "SELECT contract_version, payload_kind, passed, COUNT(*) AS count FROM conformance_checks WHERE contract_name = ? GROUP BY contract_version, payload_kind, passed",
            (name,),
        ).fetchall()
        return {
            "contract": name,
            "checks": [dict(row) for row in rows],
        }

    def _row_to_contract(self, row: sqlite3.Row) -> AdapterContract:
        return AdapterContract(
            name=row["name"],
            version=row["version"],
            adapter_type=row["adapter_type"],
            input_schema=json.loads(row["input_schema"]),
            output_schema=json.loads(row["output_schema"]),
            capabilities=json.loads(row["capabilities"]),
            metadata=json.loads(row["metadata"]),
        )