orchestrata_ultimate.session.kg_sync._base

Reference › orchestrata_ultimate › session › kg_sync › _base

_SyncerBase owns the graph connection state and the Cypher execution primitives that every stage mixin uses. The constructor takes the project root and optional graph URI, user, password, database and project id. connect returns whether the connection succeeded, connected reports the state and close ends it. _run(cypher, **params) executes one statement and returns the rows. _run_batch runs a statement over a list of parameter dictionaries in chunks, with _run_chunk handling a chunk (it takes a recursion depth, so a failing chunk can be retried in smaller pieces), and _wire_rel_chunked wires relationships in chunks and records errors in a SyncStats. The stage modules declare the same method names as stubs so they type-check, and the combined KnowledgeGraphSyncer supplies these implementations.

At a glance

Class diagram

Dependency diagram

API

Classes

_SyncerBase

Line 23

Owns graph connection state and Cypher execution primitives.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/session/kg_sync/_base.py

# mypy: ignore-errors
"""Connection base for the Knowledge Graph Syncer.

Owns all connection state (``self._*``) and the low-level Cypher
execution primitives shared by every stage mixin. Composed LAST in the
``KnowledgeGraphSyncer`` MRO so the stage mixins can call ``self._run``,
``self._run_batch``, etc.
"""

import logging
from pathlib import Path
from typing import LiteralString, cast

from orchestrata_ultimate.core.graph import BoltSyncRepository, GraphRepositoryConfig
from orchestrata_ultimate.core.graph.retry import is_mvcc_conflict
from orchestrata_ultimate.core.graph_factory import resolve_kg_database, resolve_uri
from orchestrata_ultimate.core.tracing.traceability import traces
from orchestrata_ultimate.session.kg_sync._common import BATCH_SIZE, SyncStats

logger = logging.getLogger(__name__)


class _SyncerBase:
    """Owns graph connection state and Cypher execution primitives.

    Escape hatch (backlog #163 phase 5): this class needs explicit
    ``begin_transaction()``/``commit()``/``rollback()`` control for its
    bespoke deadlock-retry (``_run_chunk``'s ``DeadlockDetected`` branch)
    and transaction-memory chunk-splitting (``MemoryPool`` / "Transaction
    closed" recursion) — neither is exposed by the
    ``SyncGraphRepository``/``UnitOfWork`` protocol, which only offers
    managed (single-statement) reads/writes or a commit-on-clean-exit
    unit of work. ``connect()`` therefore builds and owns a
    ``BoltSyncRepository`` directly, and every method that needs
    transaction-level control reaches through its documented ``.driver``
    escape hatch (``core.graph.adapters.bolt_sync.BoltSyncRepository.driver``)
    for the raw ``neo4j.Driver`` this class's Cypher primitives are built on.
    """

    def __init__(
        self,
        project_root: str,
        graph_uri: str | None = None,
        graph_user: str | None = None,
        graph_password: str | None = None,
        graph_database: str | None = None,
        project_id: str | None = None,
        connection_timeout: int = 30,
    ) -> None:
        self._root = Path(project_root)
        self._uri = graph_uri or resolve_uri()
        self._user = graph_user or ""
        self._password = graph_password or ""
        self._database = resolve_kg_database(graph_database)
        self._connection_timeout = connection_timeout
        self._repo: BoltSyncRepository | None = None
        self._project_id = project_id or self._root.name
        self._collected_traces: list[dict] = []
        self._collected_implements: list[dict] = []

    # ── Connection ─────────────────────────────────────────────

    def connect(self) -> bool:
        """Connect to graph. Returns True on success."""
        try:
            repo = BoltSyncRepository(
                GraphRepositoryConfig(
                    uri=self._uri,
                    auth=(self._user, self._password),
                    database=self._database,
                    driver_kwargs={"connection_timeout": self._connection_timeout},
                )
            )
            repo.verify_connectivity()
            self._repo = repo
            return True
        except Exception as e:
            logger.warning("Graph connection failed: %s", e)
            self._repo = None
            return False

    def close(self) -> None:
        """Close graph connection."""
        if self._repo:
            self._repo.close()
            self._repo = None

    @property
    def connected(self) -> bool:
        return self._repo is not None

    def _run(self, cypher: str, **params) -> list[dict]:
        """Execute a Cypher query and return results.

        Uses explicit transactions so errors propagate immediately
        instead of being swallowed by the managed transaction retry loop.
        Write queries use write transactions, reads use read transactions.
        """
        if not self._repo:
            raise RuntimeError("Not connected to graph")

        upper = cypher.upper()
        any(kw in upper for kw in ("MERGE", "CREATE", "SET ", "DELETE", "REMOVE"))

        with self._repo.driver.session(database=self._database) as session:
            tx = session.begin_transaction()
            try:
                # cypher is always a source-literal query string; cast narrows
                # the widened `str` param type back to LiteralString for tx.run().
                result = tx.run(cast(LiteralString, cypher), **params)
                records = list(result)
                tx.commit()
                return [dict(r) for r in records]
            except Exception:
                try:
                    tx.rollback()
                except Exception:
                    pass
                raise

    @traces(spec="008", story="US-001", reqs=["REQ-008-023", "REQ-008-039"])
    def _run_batch(self, cypher: str, batch: list[dict], **extra_params) -> list[dict]:
        """Execute a batched Cypher query using UNWIND $batch.

        Splits into chunks of BATCH_SIZE and runs each in an
        explicit write transaction for retry safety. If a chunk
        hits the graph transaction memory limit, it is split in
        half and retried recursively.
        """
        if not self._repo:
            raise RuntimeError("Not connected to graph")
        if not batch:
            return []
        results = []
        for i in range(0, len(batch), BATCH_SIZE):
            chunk = batch[i : i + BATCH_SIZE]
            results.extend(self._run_chunk(cypher, chunk, extra_params))
        return results

    def _run_chunk(
        self, cypher: str, chunk: list[dict], extra_params: dict, depth: int = 0
    ) -> list[dict]:
        """Run a single chunk with explicit transaction, splitting on memory errors."""
        from neo4j.exceptions import Neo4jError

        if not self._repo:
            raise RuntimeError("Not connected to graph")

        with self._repo.driver.session(database=self._database) as session:
            tx = session.begin_transaction()
            try:
                result = tx.run(
                    cast(LiteralString, cypher), batch=chunk, **extra_params
                )
                records = [dict(r) for r in result]
                tx.commit()
                return records
            except Neo4jError as exc:
                try:
                    tx.rollback()
                except Exception:
                    pass
                exc_str = str(exc)
                # Graph serializes conflicting writers and reports
                # DeadlockDetected when it cannot; orchgraph/Memgraph's MVCC
                # storage ABORTS one of them with a transient error instead.
                # Both are the same situation — a losing writer that must
                # re-run — so both take this retry. This class owns raw
                # transactions (see the class docstring's escape-hatch note)
                # and therefore does NOT inherit the repository adapter's
                # retry_on_mvcc_conflict; folding is_mvcc_conflict in here is
                # what keeps the two paths equivalent.
                if (
                    "DeadlockDetected" in exc_str or is_mvcc_conflict(exc)
                ) and depth < 4:
                    import time

                    logger.info(
                        "Write conflict (attempt %d), retrying chunk of %d",
                        depth + 1,
                        len(chunk),
                    )
                    time.sleep(0.5 * (depth + 1))
                    return self._run_chunk(cypher, chunk, extra_params, depth + 1)
                if (
                    ("MemoryPool" in exc_str or "Transaction closed" in exc_str)
                    and depth < 4
                    and len(chunk) > 1
                ):
                    mid = len(chunk) // 2
                    logger.info(
                        "Transaction memory limit hit (chunk=%d), splitting to %d",
                        len(chunk),
                        mid,
                    )
                    left = self._run_chunk(cypher, chunk[:mid], extra_params, depth + 1)
                    right = self._run_chunk(
                        cypher, chunk[mid:], extra_params, depth + 1
                    )
                    return left + right
                raise
            except Exception:
                tx.rollback()
                raise

    def _wire_rel_chunked(self, cypher: str, stats: SyncStats, **params) -> None:
        """Run a relationship-wiring query in chunks to stay under 64 MiB tx limit."""
        chunk_size = BATCH_SIZE
        while True:
            result = self._run(cypher, limit=chunk_size, **params)
            cnt = result[0].get("cnt", 0) if result else 0
            stats.relationships_created += cnt
            if cnt < chunk_size:
                break