orchestrata_ultimate.knowledge_graph.graph_client

Reference › orchestrata_ultimate › knowledge_graph › graph_client

GraphClient is the low-level async access layer to the operational knowledge graph, the orchkg database on the orchgraph engine (a Memgraph fork). It is the class that most of the knowledge_graph package, the gateway RPC layer, the kernel and several volks write and read through (Spec 008, REQ-008-001). It is built from two mixins, GraphNodeCreationMixin and GraphFileQueryMixin, defined in graph_node_mixins; the client itself supplies connection handling and the generic node, relationship and query operations.

connect creates an async Bolt driver through core.graph_factory.create_async_graph using the pool size, lifetime and timeout from GraphClientConfig (see config) and verifies connectivity. It returns False and logs a warning on failure instead of raising. A driver passed to the constructor is treated as shared: connect only verifies it and close leaves it open. The neo4j name in the driver type is Bolt-client naming only; Neo4j itself is retired.

The generic operations are create_node, get_node, update_node, delete_node, create_relationship, delete_relationship, query, count_nodes and ensure_schema, which applies constraint and index statements. All of them funnel through _execute, which opens a session on the configured database and returns records as dictionaries. It raises GraphConnectionError when not connected. On transient concurrency errors, including the write-conflict abort that Memgraph's storage engine raises where Neo4j would block, it retries up to a fixed number of attempts with exponential backoff starting at 50 ms; any other failure becomes a GraphQueryError carrying the Cypher text. Exception types come from exceptions.

At a glance

Class diagram

Dependency diagram

API

Classes

GraphClient

Bases: GraphNodeCreationMixin, GraphFileQueryMixin · Decorators: traces(spec='008', story='US-001', reqs=['REQ-008-001']) · Line 44

Graph database client for the knowledge graph.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/knowledge_graph/graph_client.py

# mypy: ignore-errors
"""Graph client for the knowledge graph.

Handles connection, CRUD operations, schema management, and query execution
against the graph. Typed-node creation and file-query helpers are provided by
``GraphNodeCreationMixin`` and ``GraphFileQueryMixin`` (graph_node_mixins).
"""

from __future__ import annotations

import asyncio
import logging
from typing import TYPE_CHECKING, Any, LiteralString, cast

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.core.graph.retry import MAX_ATTEMPTS, is_mvcc_conflict_message
from orchestrata_ultimate.core.graph_factory import resolve_kg_dialect
from orchestrata_ultimate.knowledge_graph.config import GraphClientConfig
from orchestrata_ultimate.knowledge_graph.exceptions import (
    GraphConnectionError,
    GraphQueryError,
)
from orchestrata_ultimate.knowledge_graph.graph_node_mixins import (
    GraphFileQueryMixin,
    GraphNodeCreationMixin,
)
from orchestrata_ultimate.knowledge_graph.nodes import (
    cypher_record_to_node,
    node_to_cypher_props,
)
from orchestrata_ultimate.knowledge_graph.schema import (
    get_node_label,
    get_unified_ddl,
)
from orchestrata_ultimate.knowledge_graph.types import Node

if TYPE_CHECKING:
    import neo4j

logger = logging.getLogger(__name__)


@traces(spec="008", story="US-001", reqs=["REQ-008-001"])
class GraphClient(GraphNodeCreationMixin, GraphFileQueryMixin):
    """Graph database client for the knowledge graph.

    Provides async CRUD operations, schema management,
    and query execution.

    Typed-node creation (``create_file_node``, ``create_class_node``, etc.)
    and file-query helpers (``get_project_files``, ``delete_file_cascade``,
    etc.) are inherited from the mixin classes — all callers are unaffected.
    """

    def __init__(
        self,
        config: GraphClientConfig | None = None,
        driver: neo4j.AsyncDriver | None = None,
    ) -> None:
        self._config = config or GraphClientConfig()
        self._driver = driver
        self._connected = False
        self._owns_driver = driver is None

    @traces(spec="001", story="US-010", reqs=["REQ-001-US-010-013"])
    async def connect(self) -> bool:
        """Establish connection to the graph."""
        if self._driver and not self._owns_driver:
            try:
                await self._driver.verify_connectivity()
                self._connected = True
                return True
            except Exception as exc:
                logger.warning("Shared driver verify failed: %s", exc)
                self._connected = False
                return False
        try:
            from orchestrata_ultimate.core.graph_factory import create_async_graph

            self._driver = create_async_graph(
                self._config.uri,
                auth=(self._config.username, self._config.password),
                max_connection_pool_size=self._config.max_pool_size,
                max_connection_lifetime=self._config.max_connection_lifetime,
                connection_timeout=self._config.connection_timeout,
            )
            await self._driver.verify_connectivity()
            self._connected = True
            logger.info("Connected to the graph at %s", self._config.uri)
            return True
        except Exception as exc:
            logger.warning("Graph connection failed: %s", exc)
            self._connected = False
            return False

    async def close(self) -> None:
        """Close the graph connection."""
        if self._driver and self._owns_driver:
            await self._driver.close()
        self._driver = None
        self._connected = False

    @property
    def connected(self) -> bool:
        """Return True if currently connected to the graph."""
        return self._connected

    @traces(spec="008", story="US-001", reqs=["REQ-008-037"])
    async def ensure_schema(self, statements: list[str] | None = None) -> int:
        """Create constraints and indexes.

        Applies the canonical identity constraints (SPEC-106
        ``canonical_schema.get_all_ddl()`` — THE single identity source,
        generated from ``CANONICAL_MERGE_KEYS``) first, then the lookup
        layer still owned by ``schema.py`` (extra constraints, lookup
        indexes, fulltext indexes), deduplicated.

        Failure policy: only genuine already-exists conflicts (Neo4j
        ``*AlreadyExists`` / equivalent-schema-rule codes) are skipped;
        every other error raises ``GraphQueryError``. The previous blanket
        ``except Exception -> logger.debug`` swallowed real failures —
        against a dialect mismatch the KG silently booted with no indexes
        and no fulltext (plans/graph-consolidation-decision.md).

        Args:
            statements: Override the DDL set (tests / scoped migrations),
                mirroring ``canonical_schema.apply``.

        Returns:
            Count of DDL statements executed successfully.
        """
        if not self._connected or not self._driver:
            raise GraphConnectionError("Not connected to the graph")
        if statements is None:
            # Dialect at CALL time: the same client class now serves Neo4j
            # and orchgraph/Memgraph, whose schema grammar has no
            # ``IF NOT EXISTS`` and no ``CREATE TEXT INDEX ... FOR ... ON``.
            # Defaulting to the graph DDL made ensure_schema raise on the new
            # engine, and the failure policy below (correctly) does not
            # swallow it.
            statements = get_unified_ddl(dialect=resolve_kg_dialect())
        count = 0
        async with self._driver.session(database=self._config.database) as session:
            for ddl in statements:
                try:
                    await session.run(cast(LiteralString, ddl))
                    count += 1
                except Exception as exc:
                    # Neo4j reports this in the error CODE
                    # (Neo.ClientError.Schema.*AlreadyExists). orchgraph/
                    # Memgraph has one generic code and would say it in the
                    # MESSAGE — measured live 2026-07-26, re-creating an
                    # identical constraint or index on that engine is a silent
                    # no-op, so it does not normally reach here at all; the
                    # message check covers the forms that do (e.g. a plain
                    # index colliding with a constraint-backed one) and
                    # matches canonical_schema.apply's policy.
                    already = (
                        "AlreadyExists" in (getattr(exc, "code", "") or "")
                        or "already exists" in str(exc).lower()
                    )
                    if already:
                        # IF NOT EXISTS covers same-name duplicates; this is
                        # the equivalent-rule-under-a-different-name case.
                        logger.debug("Schema DDL already satisfied: %s", exc)
                        continue
                    logger.error("Schema DDL failed: %s -- %s", ddl, exc)
                    raise GraphQueryError(str(exc), query=ddl) from exc
        return count

    # -- Node CRUD ---------------------------------------------------------

    async def create_node(self, node: Node) -> str:
        """Create a generic node.

        Args:
            node: The Node to persist.

        Returns:
            The node ID.
        """
        if not self._connected or not self._driver:
            raise GraphConnectionError("Not connected")
        label = get_node_label(node.type.value)
        props = node_to_cypher_props(node)
        cypher = f"CREATE (n:{label} $props) RETURN n.id AS id"
        result = await self._execute("write", cypher, {"props": props})
        return result[0]["id"] if result else node.id

    async def get_node(self, node_id: str, node_type: str | None = None) -> Node | None:
        """Get a node by ID.

        Args:
            node_id: The ID of the node to retrieve.
            node_type: Optional node type (``types_enums.NodeType`` value,
                resolved via ``schema.get_node_label``) to plan a labeled
                index seek instead of an unlabeled ``MATCH (n {id: ...})`` --
                the latter cannot use the per-label unique constraints and
                plans an AllNodesScan on every call (measured live:
                ~7,900 dbHits vs. 15 for a labeled seek on the 768K-node KG).
                Pass it whenever the caller already knows the type; the bulk
                UNWIND writes in ``sync.py`` apply the same labeled-seek
                shape. Defaults to the prior unlabeled behaviour.

        Returns:
            The Node if found, otherwise None.
        """
        match_clause = "MATCH (n {id: $id})"
        if node_type is not None:
            label = get_node_label(node_type)
            match_clause = f"MATCH (n:{label} {{id: $id}})"
        cypher = f"{match_clause} RETURN properties(n) AS props"
        result = await self._execute("read", cypher, {"id": node_id})
        if result:
            return cypher_record_to_node(result[0].get("props", {}))
        return None

    async def update_node(
        self,
        node_id: str,
        properties: dict[str, Any],
        node_type: str | None = None,
    ) -> bool:
        """Update node properties.

        Args:
            node_id: The ID of the node to update.
            properties: Key/value pairs to merge onto the node.
            node_type: Optional node type -- see ``get_node``'s labeled-seek
                note. Defaults to the prior unlabeled behaviour.

        Returns:
            True if the node was found and updated.
        """
        match_clause = "MATCH (n {id: $id})"
        if node_type is not None:
            label = get_node_label(node_type)
            match_clause = f"MATCH (n:{label} {{id: $id}})"
        cypher = f"{match_clause} SET n += $props RETURN n.id AS id"
        result = await self._execute(
            "write", cypher, {"id": node_id, "props": properties}
        )
        return len(result) > 0

    @traces(spec="008", story="US-008-001", reqs=["REQ-008-089"])
    async def delete_node(self, node_id: str, node_type: str | None = None) -> bool:
        """Delete a node and all its relationships.

        Args:
            node_id: The ID of the node to delete.
            node_type: Optional node type -- see ``get_node``'s labeled-seek
                note. Defaults to the prior unlabeled behaviour.

        Returns:
            True if a node was deleted.
        """
        match_clause = "MATCH (n {id: $id})"
        if node_type is not None:
            label = get_node_label(node_type)
            match_clause = f"MATCH (n:{label} {{id: $id}})"
        cypher = f"{match_clause} DETACH DELETE n RETURN count(*) AS deleted"
        result = await self._execute("write", cypher, {"id": node_id})
        return result[0].get("deleted", 0) > 0 if result else False

    # -- Relationship CRUD -------------------------------------------------

    async def create_relationship(
        self,
        source_id: str,
        target_id: str,
        rel_type: str,
        properties: dict[str, Any] | None = None,
        source_type: str | None = None,
        target_type: str | None = None,
    ) -> bool:
        """Create a relationship between two nodes.

        Args:
            source_id: ID of the source node.
            target_id: ID of the target node.
            rel_type: Cypher relationship type label.
            properties: Optional properties to set on the relationship.
            source_type: Optional type of the source node -- see
                ``get_node``'s labeled-seek note. Defaults to the prior
                unlabeled behaviour.
            target_type: Optional type of the target node, independent of
                ``source_type`` (the two ends can be different node types).

        Returns:
            True if the relationship was created or already existed.
        """
        source_label = (
            f":{get_node_label(source_type)}" if source_type is not None else ""
        )
        target_label = (
            f":{get_node_label(target_type)}" if target_type is not None else ""
        )
        set_clause = ""
        params: dict[str, Any] = {"source_id": source_id, "target_id": target_id}
        if properties:
            # Neo4j rejects a parameter map embedded directly in a MERGE
            # pattern ("Parameter maps cannot be used in MERGE patterns");
            # properties must be applied afterwards via SET.
            set_clause = " SET r += $rel_props"
            params["rel_props"] = properties
        cypher = (
            f"MATCH (a{source_label} {{id: $source_id}}), (b{target_label} {{id: $target_id}}) "
            f"MERGE (a)-[r:{rel_type}]->(b)"
            f"{set_clause} "
            "RETURN type(r) AS rel_type"
        )
        result = await self._execute("write", cypher, params)
        return len(result) > 0

    async def delete_relationship(
        self,
        source_id: str,
        target_id: str,
        rel_type: str,
        source_type: str | None = None,
        target_type: str | None = None,
    ) -> bool:
        """Delete a relationship between two nodes.

        Args:
            source_id: ID of the source node.
            target_id: ID of the target node.
            rel_type: Cypher relationship type label.
            source_type: Optional type of the source node -- see
                ``get_node``'s labeled-seek note. Defaults to the prior
                unlabeled behaviour.
            target_type: Optional type of the target node.

        Returns:
            True if a relationship was deleted.
        """
        source_label = (
            f":{get_node_label(source_type)}" if source_type is not None else ""
        )
        target_label = (
            f":{get_node_label(target_type)}" if target_type is not None else ""
        )
        cypher = (
            f"MATCH (a{source_label} {{id: $source_id}})-[r:{rel_type}]->"
            f"(b{target_label} {{id: $target_id}}) "
            "DELETE r RETURN count(*) AS deleted"
        )
        params = {"source_id": source_id, "target_id": target_id}
        result = await self._execute("write", cypher, params)
        return result[0].get("deleted", 0) > 0 if result else False

    # -- Query -------------------------------------------------------------

    async def query(
        self, cypher: str, parameters: dict[str, Any] | None = None
    ) -> list[dict[str, Any]]:
        """Execute a Cypher query and return results.

        Args:
            cypher: The Cypher query string.
            parameters: Optional query parameters.

        Returns:
            List of result records as dicts.
        """
        return await self._execute("read", cypher, parameters or {})

    async def count_nodes(self, label: str | None = None) -> int:
        """Count nodes, optionally filtered by label.

        Args:
            label: Optional node label to filter by.

        Returns:
            The node count.
        """
        if label:
            cypher = f"MATCH (n:{label}) RETURN count(n) AS cnt"
        else:
            cypher = "MATCH (n) RETURN count(n) AS cnt"
        result = await self._execute("read", cypher, {})
        return result[0]["cnt"] if result else 0

    # -- Internal ----------------------------------------------------------

    async def _execute(
        self, mode: str, cypher: str, params: dict[str, Any]
    ) -> list[dict[str, Any]]:
        """Execute a Cypher statement.

        Retries on transient concurrency errors (e.g., nodes deleted by
        parallel tests) to keep the test suite stable under ``pytest -n auto``.

        Args:
            mode: ``'read'`` or ``'write'``.
            cypher: The Cypher query string.
            params: Query parameters.

        Returns:
            List of result records as dicts.

        Raises:
            GraphConnectionError: If not connected.
            GraphQueryError: If the query fails after retries.
        """
        if not self._connected or not self._driver:
            raise GraphConnectionError("Not connected to the graph")

        last_exc: Exception | None = None
        for attempt in range(MAX_ATTEMPTS):
            try:
                async with self._driver.session(
                    database=self._config.database
                ) as session:
                    result = await session.run(cast(LiteralString, cypher), params)
                    records = [dict(record) async for record in result]
                    return records
            except Exception as exc:
                last_exc = exc
                msg = str(exc)
                # Retry on transient concurrency errors caused by parallel
                # test cleanup deleting nodes mid-query, AND on a Memgraph
                # MVCC write-conflict abort (plans/neo4j-to-orchgraph-
                # migration-plan.md Phase 2 Step 2.1): Memgraph's storage
                # engine ABORTS one of two concurrently conflicting writers
                # instead of blocking on a lock the way Neo4j does.
                if self._is_transient_error(msg):
                    await asyncio.sleep(0.05 * (2**attempt))
                    continue
                raise GraphQueryError(msg, query=cypher)
        raise GraphQueryError(str(last_exc), query=cypher)

    @staticmethod
    def _is_transient_error(msg: str) -> bool:
        """Return True if *msg* is a transient concurrency error worth retrying.

        Matches Neo4j errors raised when parallel cleanup deletes or
        recreates nodes mid-query. Uses plain substring checks — the
        previous ``"Node.*already exists"`` treated ``.*`` as literal
        characters, so the retry branch never fired. Also matches a
        Memgraph MVCC write-conflict abort via
        ``core.graph.retry.is_mvcc_conflict_message`` (shared marker list
        with the ``core.graph`` repository adapters, not duplicated here).
        """
        return (
            "EntityNotFound" in msg
            or "already exists" in msg
            or is_mvcc_conflict_message(msg)
        )


# Alias for backward compatibility with generated tests
GraphClientAlias = GraphClient