orchestrata_ultimate.semantic_graph.main_discovery

Reference › orchestrata_ultimate › semantic_graph › main_discovery

Client-side discovery of the current Memgraph MAIN data instance for the shared semantic graph (the memgraph database). No writer may pin a data-instance address, because Raft failover moves MAIN between the orchgraph-data-* instances and a restart once broke every hardcoded writer. resolve_main is the single place that answers which bolt URI is MAIN now, and the Memgraph adapter imports it. Resolution runs in a fixed order. First, if MEMGRAPH_HA_CONTROLLER_URL is set it asks the legacy controller's /main HTTP endpoint (_resolve_via_controller); the live cluster runs no such controller, so this is a no-op in production. Second, it probes each coordinator listed in MEMGRAPH_URIS with SHOW INSTANCES (_resolve_via_coordinators, with _as_bolt_uri normalising the scheme) until a row with role main answers, and returns that row's bolt address. Third, it falls back to MEMGRAPH_URI, which on the live cluster is a routing URI that the Bolt driver auto-routes to MAIN for writes, so the fallback is failover-safe but is not guaranteed to be a direct data-instance address. The result is cached per process, one discovery round trip rather than one per write; reset_cache clears it and is the test seam. See ha_controller for the legacy controller.

At a glance

Class diagram

No classes; see Functions.

Dependency diagram

API

Constants

Functions

Source

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

"""Client-side discovery of the current Memgraph MAIN data-instance bolt URI.

Reproducibility plan Phase R3: no writer may pin a specific data-instance
address (``bolt://10.53.1.123:7687`` etc) because Raft failover moves MAIN
between the ``orchgraph-data-*`` instances — a restart once failed MAIN over
from ``instance_2`` (``10.53.1.123``) to ``instance_1`` (``10.53.1.128``)
and every hardcoded writer broke. :func:`resolve_main` is the single place
that answers "which bolt URI is MAIN right now", so every writer can share it.

Resolution order:

1. ``MEMGRAPH_HA_CONTROLLER_URL`` -- if set, ask its ``/main`` HTTP endpoint
   (the shape documented in ``docs/memgraph-ha-cluster.md``). This is the
   legacy single-controller HA design; the live orchgraph cluster does not
   run one, so this branch is a no-op in production today but is kept for
   any future/alternate deployment that does.
2. ``MEMGRAPH_URIS`` -- a comma-separated list of coordinator bolt URIs.
   Each is probed with ``SHOW INSTANCES`` (a Raft admin query that only a
   *coordinator* answers -- see ``_as_bolt_uri``) until one returns a row
   with ``role == 'main'``; that row's ``bolt_server`` is the answer.
3. ``MEMGRAPH_URI`` -- last-resort fallback. On the live cluster this is
   itself a ``neo4j://`` routing URI to a coordinator, which the Bolt driver
   auto-routes to MAIN for writes, so even this fallback is failover-safe;
   it is not, however, guaranteed to be a direct ``bolt://<data-instance>``
   address the way (1) and (2) are.

The resolved URI is cached for the process by default (one discovery round
trip per process, not per write) -- call :func:`reset_cache` to clear it,
the injection seam used by tests.
"""

from __future__ import annotations

import json
import os
import threading
import urllib.request
from collections.abc import Mapping

from orchestrata_ultimate.semantic_graph.exceptions import ConnectionError

_cache_lock = threading.Lock()
_cached_main: str | None = None


def reset_cache() -> None:
    """Clear the process-level cached MAIN URI (test injection seam)."""
    global _cached_main
    with _cache_lock:
        _cached_main = None


def resolve_main(env: Mapping[str, str] | None = None, *, cache: bool = True) -> str:
    """Return the current Memgraph MAIN's bolt URI.

    Args:
        env: environment mapping to read ``MEMGRAPH_*`` keys from (defaults
            to ``os.environ``). Injectable so callers/tests never touch the
            real process environment.
        cache: when True (default), reuse a previously resolved URI for the
            life of the process instead of re-discovering on every call.

    Returns:
        A ``bolt://host:port`` (or, in the ``MEMGRAPH_URI`` fallback case,
        whatever scheme that env var carries) URI for the current MAIN.

    Raises:
        ConnectionError: none of the three resolution strategies produced a
            URI (no controller, no reachable coordinator, no fallback URI).
    """
    global _cached_main

    if cache:
        with _cache_lock:
            if _cached_main is not None:
                return _cached_main

    resolved = _resolve(env)

    if cache:
        with _cache_lock:
            _cached_main = resolved
    return resolved


def _resolve(env: Mapping[str, str] | None) -> str:
    active_env: Mapping[str, str] = os.environ if env is None else env

    controller_url = active_env.get("MEMGRAPH_HA_CONTROLLER_URL", "")
    if controller_url:
        uri = _resolve_via_controller(controller_url)
        if uri:
            return uri

    uris_raw = active_env.get("MEMGRAPH_URIS", "")
    coordinator_uris = [u.strip() for u in uris_raw.split(",") if u.strip()]
    if coordinator_uris:
        uri = _resolve_via_coordinators(coordinator_uris)
        if uri:
            return uri

    # No literal fallback (RUL-003): a bare 1-arg .get() (None when unset)
    # rather than .get("MEMGRAPH_URI", "") -- the `if fallback:` check below
    # already treats None the same as "", so behaviour is unchanged; this
    # only removes the literal default the scanner (and RUL-003) forbid.
    fallback = active_env.get("MEMGRAPH_URI")
    if fallback:
        return fallback

    raise ConnectionError(
        "Unable to resolve the Memgraph MAIN URI: none of "
        "MEMGRAPH_HA_CONTROLLER_URL, MEMGRAPH_URIS, or MEMGRAPH_URI are set."
    )


def _resolve_via_controller(controller_url: str) -> str | None:
    """Ask a legacy HA controller's ``/main`` endpoint for the MAIN URI."""
    try:
        with urllib.request.urlopen(
            f"{controller_url.rstrip('/')}/main", timeout=5
        ) as resp:
            data = json.loads(resp.read())
            uri = data.get("uri")
            return uri if isinstance(uri, str) and uri else None
    except OSError, ValueError:
        return None


def _resolve_via_coordinators(coordinator_uris: list[str]) -> str | None:
    """Probe each coordinator with ``SHOW INSTANCES`` for the MAIN row.

    The raw-driver coordinator probe lives in the allowlisted Memgraph HA seam
    (``ha_controller.probe_coordinator_for_main``); this module does only the
    env-driven orchestration, so it keeps no raw graph access of its own (D13).
    Imported lazily to avoid pulling the controller's FastAPI dependency into
    this lightweight, writer-path discovery module.
    """
    from orchestrata_ultimate.semantic_graph.ha_controller import (
        probe_coordinator_for_main,
    )

    for coordinator_uri in coordinator_uris:
        uri = probe_coordinator_for_main(_as_bolt_uri(coordinator_uri))
        if uri:
            return uri
    return None


def _as_bolt_uri(uri: str) -> str:
    """Normalize a coordinator URI to the ``bolt://`` scheme.

    ``SHOW INSTANCES`` is a Raft admin query that only a coordinator
    answers directly; Memgraph rejects it over a routed connection
    ("Only coordinator can run SHOW INSTANCES", verified live against
    ``neo4j://10.53.1.160:7691``). ``bolt://`` talks to the given address
    with no routing layer in between, so any configured scheme is
    normalized to it before connecting.
    """
    if "://" in uri:
        _, _, rest = uri.partition("://")
        return f"bolt://{rest}"
    return f"bolt://{uri}"