Eric-ontology/app/repository.py

from __future__ import annotations

from pathlib import Path
from threading import RLock

from rdflib import Graph
from rdflib.query import ResultRow
from rdflib.term import Node


class RDFRepository:
    """Loads the ontology and data files into one local, queryable RDF graph."""

    def __init__(self, root: Path):
        self.root = root
        self.graph = Graph()
        self._lock = RLock()
        self.loaded_files: list[Path] = []
        self.reload()

    def reload(self) -> None:
        paths = [
            self.root / "ontology" / "global-shipment.ttl",
            *sorted((self.root / "ontology").glob("*-extension.ttl")),
            *sorted((self.root / "data").glob("*.ttl")),
        ]

        graph = Graph()
        loaded_files: list[Path] = []
        for path in paths:
            if path.exists() and path not in loaded_files:
                graph.parse(path, format="turtle")
                loaded_files.append(path)

        with self._lock:
            self.graph = graph
            self.loaded_files = loaded_files
            self.base_triple_count = len(graph)

    def add_triples(self, triples: list[tuple[Node, Node, Node]]) -> int:
        with self._lock:
            before = len(self.graph)
            for triple in triples:
                self.graph.add(triple)
            return len(self.graph) - before

    def query_file(self, path: Path) -> dict[str, list[dict[str, str]]]:
        if not path.exists():
            raise FileNotFoundError(path)
        return self.query(path.read_text(encoding="utf-8"))

    def query(self, sparql: str) -> dict[str, list[dict[str, str]]]:
        with self._lock:
            result = self.graph.query(sparql)

        variables = [str(variable) for variable in result.vars or []]
        rows = [self._row_to_dict(row, variables) for row in result]
        return {"rows": rows}

    @staticmethod
    def _row_to_dict(row: ResultRow, variables: list[str]) -> dict[str, str]:
        values: dict[str, str] = {}
        for variable in variables:
            value = row.get(variable)
            values[variable] = "" if value is None else str(value)
        return values

    @property
    def triple_count(self) -> int:
        return len(self.graph)

    @property
    def in_memory_triple_count(self) -> int:
        return self.triple_count - self.base_triple_count

    @property
    def loaded_file_names(self) -> list[str]:
        return [str(path.relative_to(self.root)) for path in self.loaded_files]