From 4273f2bf28751ec91ced0a748fbb1a11ef541b42 Mon Sep 17 00:00:00 2001 From: matt Date: Wed, 1 Jul 2026 15:46:18 +0100 Subject: [PATCH 01/10] initial bold+routing --- docs/source/module.rst | 27 ++++ pyproject.toml | 5 +- python/mgclient/__init__.py | 37 ++++++ python/mgclient/routing.py | 240 ++++++++++++++++++++++++++++++++++++ setup.py | 7 +- src/mgclientmodule.c | 4 +- test/conftest.py | 143 +++++++++++++++++++++ test/test_connection.py | 92 +++----------- test/test_routing.py | 147 ++++++++++++++++++++++ 9 files changed, 621 insertions(+), 81 deletions(-) create mode 100644 python/mgclient/__init__.py create mode 100644 python/mgclient/routing.py create mode 100644 test/conftest.py create mode 100644 test/test_routing.py diff --git a/docs/source/module.rst b/docs/source/module.rst index 2a7e329e..75fb3099 100644 --- a/docs/source/module.rst +++ b/docs/source/module.rst @@ -11,6 +11,33 @@ The module interface respects the DB-API 2.0 standard defined in :pep:`249`. See :ref:`lazy-connections` section to learn about advantages and limitations of using the ``lazy`` parameter. +##################### +Client-side routing +##################### + +When connecting to a Memgraph high-availability cluster, passing +``routing=True`` to :func:`connect` makes the connection routing-aware: the +given ``host``/``port`` must point at a cluster coordinator, from which the +current cluster topology is fetched (via a Bolt ``ROUTE`` message) and a data +instance matching the requested ``access_mode`` is selected and connected to. + +.. data:: mgclient.ACCESS_MODE_WRITE + + ``access_mode`` value selecting a server that accepts writes (the cluster + main). This is the default. + +.. data:: mgclient.ACCESS_MODE_READ + + ``access_mode`` value selecting a server that serves reads (a replica). + +If the addresses a cluster advertises are not directly reachable by the client +(for example when the cluster is reached through a proxy or a port-forward), +pass a ``resolver`` callable that maps an advertised ``"host:port"`` address to +an iterable of ``"host:port"`` targets to try. + +For lower-level access to the routing table itself, see +:meth:`Connection.get_routing_table`. + ################ Module constants ################ diff --git a/pyproject.toml b/pyproject.toml index d2660189..421dbd68 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,9 +46,8 @@ Documentation = "https://memgraph.github.io/pymgclient" [tool.setuptools] py-modules = ["tools.pymgclient_build_meta"] - -[tool.setuptools.packages.find] -include = ["tools*"] +packages = ["tools", "mgclient"] +package-dir = { mgclient = "python/mgclient" } [tool.setuptools.dynamic] version = { attr = "tools.pymgclient_build_meta.__version__" } diff --git a/python/mgclient/__init__.py b/python/mgclient/__init__.py new file mode 100644 index 00000000..e4b235bd --- /dev/null +++ b/python/mgclient/__init__.py @@ -0,0 +1,37 @@ +# Copyright (c) 2016-2020 Memgraph Ltd. [https://memgraph.com] +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Memgraph database adapter for Python, compliant with the DB-API 2.0 +specification described by :pep:`249`. + +The compiled functionality lives in the :mod:`mgclient._mgclient` C extension +and is re-exported here unchanged, so ``import mgclient`` behaves exactly as it +always has. On top of that this package adds optional client-side routing for +Memgraph high-availability clusters (see :func:`connect` and the routing +constants below). +""" + +# Re-export everything the C extension provides (connect, Connection, Cursor, +# the exception hierarchy, the type objects and all module constants) so that +# the public ``mgclient`` API is unchanged. +from mgclient._mgclient import * # noqa: F401,F403 + +# The routing-aware ``connect`` wrapper replaces the C ``connect`` imported +# above; with routing disabled (the default) it delegates straight to it, so +# existing code is unaffected. +from mgclient.routing import ( # noqa: F401,E402 + ACCESS_MODE_READ, + ACCESS_MODE_WRITE, + connect, +) diff --git a/python/mgclient/routing.py b/python/mgclient/routing.py new file mode 100644 index 00000000..2bbb1c9c --- /dev/null +++ b/python/mgclient/routing.py @@ -0,0 +1,240 @@ +# Copyright (c) 2016-2020 Memgraph Ltd. [https://memgraph.com] +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Client-side routing for Memgraph high-availability clusters. + +This module builds a small routing layer on top of the primitives provided by +the :mod:`mgclient._mgclient` C extension (:func:`connect` and +:meth:`Connection.get_routing_table`). It is pure Python and opens no +connections beyond the ones the C extension already makes. +""" + +from mgclient._mgclient import connect as _base_connect +from mgclient._mgclient import OperationalError + +#: Route the connection to a server that accepts writes (the cluster main). +ACCESS_MODE_WRITE = "WRITE" +#: Route the connection to a server that serves reads (a replica). +ACCESS_MODE_READ = "READ" + +_ACCESS_MODES = (ACCESS_MODE_WRITE, ACCESS_MODE_READ) + + +def _default_resolver(address): + """Identity resolver: the advertised address is used as-is.""" + return [address] + + +def _split_host_port(address): + """Split a routing-table ``"host:port"`` address into ``(host, port)``.""" + host, separator, port = address.rpartition(":") + if not separator or not host: + raise ValueError(f"invalid server address in routing table: {address!r}") + try: + port = int(port) + except ValueError: + raise ValueError( + f"invalid port in routing table address: {address!r}" + ) from None + return host, port + + +def _routed_connect(access_mode, resolver, routing_context, kwargs): + if isinstance(access_mode, str): + access_mode = access_mode.upper() + if access_mode not in _ACCESS_MODES: + raise ValueError( + f"access_mode must be one of {_ACCESS_MODES!r}, got {access_mode!r}" + ) + if resolver is None: + resolver = _default_resolver + + # Bootstrap: connect to the coordinator and fetch the routing table. ROUTE + # requires an idle, non-lazy session, so drop "lazy" for this connection + # only (the final connection still honours the caller's choice). + bootstrap_kwargs = dict(kwargs) + bootstrap_kwargs.pop("lazy", None) + + coordinator = _base_connect(**bootstrap_kwargs) + try: + context = routing_context if routing_context is not None else {} + table = coordinator.get_routing_table(routing_context=context) + finally: + coordinator.close() + + # Collect the advertised addresses of every server that serves the + # requested access mode, preserving the coordinator's ordering. + advertised = [] + for server in table.get("servers", []): + if server.get("role") == access_mode: + advertised.extend(server.get("addresses", [])) + + if not advertised: + raise OperationalError( + f"no {access_mode} server available in the routing table" + ) + + # Map advertised addresses to reachable connect targets, de-duplicating + # while keeping order so failover is deterministic. + candidates = [] + for address in advertised: + for target in resolver(address): + if target not in candidates: + candidates.append(target) + + # The final connection reuses the caller's parameters (auth, SSL, lazy, ...) + # but is directed at the resolved data instance. "address" (a numeric IP + # for the coordinator) never applies to the routed target. + final_kwargs = dict(kwargs) + final_kwargs.pop("address", None) + + errors = [] + for target in candidates: + host, port = _split_host_port(target) + attempt = dict(final_kwargs) + attempt["host"] = host + attempt["port"] = port + try: + return _base_connect(**attempt) + except OperationalError as exc: + errors.append(f"{target}: {exc}") + + raise OperationalError( + f"could not connect to any {access_mode} server; tried " + + "; ".join(errors) + ) + + +def connect( + *, + routing=False, + access_mode=ACCESS_MODE_WRITE, + resolver=None, + routing_context=None, + **kwargs, +): + """Make a new connection to a Memgraph server and return a + :class:`Connection` object. + + Currently recognized connection parameters are: + + * :obj:`host` + + DNS resolvable name of host to connect to. Exactly one of host and + address parameters must be specified. + + * :obj:`address` + + Numeric IP address of host to connect to. This should be in the + standard IPv4 address format. You can also use IPv6 if your machine + supports it. Exactly one of host and address parameters must be + specified. + + * :obj:`port` + + Port number to connect to at the server host. + + * :obj:`username` + + Username to connect as. + + * :obj:`password` + + Password to be used if the server demands password authentication. + + * :obj:`client_name` + + Alternate name and version of the client to send to server. Default + is set by the underlying mgclient library. + + * :obj:`sslmode` + + Whether a secure connection will be negotiated with the server; + either :const:`mgclient.MG_SSLMODE_DISABLE` (default) or + :const:`mgclient.MG_SSLMODE_REQUIRE`. + + * :obj:`sslcert` + + File name of the client SSL certificate. Ignored if an SSL + connection is not made. + + * :obj:`sslkey` + + Location of the secret key used for the client certificate. Ignored + if an SSL connection is not made. + + * :obj:`trust_callback` + + A callable taking four arguments (hostname, IP address, public key + type and fingerprint), called after the SSL handshake. If it returns + ``False`` the connection is terminated. Useful for implementing a + TOFU (trust on first use) mechanism. + + * :obj:`lazy` + + If ``True``, a lazy connection is made. Default is ``False``. + + When ``routing`` is ``True`` the connection is made in a routing-aware way, + suitable for a Memgraph high-availability cluster: + + * The ``host``/``port`` (or ``address``) parameters must point at a + cluster coordinator. A Bolt ``ROUTE`` message is sent to it to obtain + the current cluster topology. + + * Depending on ``access_mode`` a server is selected from the routing + table -- :data:`ACCESS_MODE_WRITE` (the default) selects the main and + :data:`ACCESS_MODE_READ` selects a replica. + + * A normal connection is then opened to that server and returned. If a + selected server cannot be reached, the remaining candidates are tried + in turn; an :exc:`OperationalError` is raised only if none succeed. + + All other connection parameters (authentication, SSL, ``lazy``, ...) are + applied to the returned data-instance connection. + + Extra routing parameters: + + * :obj:`routing` + + If ``True``, enable client-side routing as described above. + Defaults to ``False``. + + * :obj:`access_mode` + + Either :data:`ACCESS_MODE_WRITE` or :data:`ACCESS_MODE_READ`. + Defaults to :data:`ACCESS_MODE_WRITE`. + + * :obj:`resolver` + + Optional callable mapping an advertised ``"host:port"`` address + from the routing table to an iterable of ``"host:port"`` targets to + actually try (in order). Defaults to using the advertised address + unchanged. This is the hook to use when the addresses the cluster + advertises are not directly reachable by the client (for example + behind a proxy or port-forward). + + * :obj:`routing_context` + + Optional :class:`dict` passed through to + :meth:`Connection.get_routing_table` as the routing context. + """ + if not routing: + return _base_connect(**kwargs) + + return _routed_connect( + access_mode=access_mode, + resolver=resolver, + routing_context=routing_context, + kwargs=kwargs, + ) diff --git a/setup.py b/setup.py index 492b569b..c7f7598e 100644 --- a/setup.py +++ b/setup.py @@ -36,7 +36,12 @@ # 1. The mgclient library which is the official Memgraph client library. # 2. The mgclient python extension module which is a wrapper around the # client library. -EXTENSION_NAME = "mgclient" +# +# The compiled extension is installed as ``mgclient._mgclient`` and the +# pure-Python ``mgclient`` package (see python/mgclient) re-exports it so that +# ``import mgclient`` keeps working exactly as before, while adding the +# client-side routing helpers on top. +EXTENSION_NAME = "mgclient._mgclient" sources = [str(path) for path in Path("src").glob("*.c")] diff --git a/src/mgclientmodule.c b/src/mgclientmodule.c index 4c97106d..78241d51 100644 --- a/src/mgclientmodule.c +++ b/src/mgclientmodule.c @@ -281,13 +281,13 @@ static PyMethodDef mgclient_methods[] = { {NULL, NULL, 0, NULL}}; static struct PyModuleDef mgclient_module = {.m_base = PyModuleDef_HEAD_INIT, - .m_name = "mgclient", + .m_name = "mgclient._mgclient", .m_doc = NULL, .m_size = -1, .m_methods = mgclient_methods, .m_slots = NULL}; -PyMODINIT_FUNC PyInit_mgclient(void) { +PyMODINIT_FUNC PyInit__mgclient(void) { PyObject *m; if (!(m = PyModule_Create(&mgclient_module))) { return NULL; diff --git a/test/conftest.py b/test/conftest.py new file mode 100644 index 00000000..bdb08764 --- /dev/null +++ b/test/conftest.py @@ -0,0 +1,143 @@ +# Copyright (c) 2016-2020 Memgraph Ltd. [https://memgraph.com] +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared fixtures for the test suite. + +The high-availability fixtures here are used by both ``test_connection.py`` +(the ``get_routing_table`` primitive) and ``test_routing.py`` (the client-side +routing built on top of it). +""" + +import os +import time + +import mgclient +import pytest + +from common import ( + MEMGRAPH_HA_COORDINATOR_HOST, + MEMGRAPH_HA_COORDINATOR_PORT, +) + +# Topology created by the "Run Memgraph HA Cluster" CI step. The fixture below +# must agree with the container names and ports used there. +HA_COORDINATORS = ["mg-coord1", "mg-coord2", "mg-coord3"] +HA_DATA_INSTANCES = [("instance_1", "mg-data1"), ("instance_2", "mg-data2")] +HA_MAIN = "instance_1" +HA_BOLT_PORT = 7687 +HA_COORDINATOR_PORT = 12121 +HA_MANAGEMENT_PORT = 13011 +HA_REPLICATION_PORT = 10000 + + +def _ha_admin(conn, query): + """Run a coordinator admin query, tolerating idempotent re-runs.""" + cursor = conn.cursor() + try: + cursor.execute(query) + try: + cursor.fetchall() + except mgclient.Error: + pass + except mgclient.DatabaseError as exc: + # Re-running setup against an already-configured cluster is fine. + print(f"HA setup query ignored error: {exc}") + + +@pytest.fixture(scope="module") +def ha_cluster(): + host = MEMGRAPH_HA_COORDINATOR_HOST + port = MEMGRAPH_HA_COORDINATOR_PORT + + conn = mgclient.connect(host=host, port=port) + conn.autocommit = True + + # mg-coord1 is the bootstrap coordinator; add the remaining ones. + for cid, name in enumerate(HA_COORDINATORS, start=1): + if cid == 1: + continue + _ha_admin( + conn, + f'ADD COORDINATOR {cid} WITH CONFIG ' + f'{{"bolt_server": "{name}:{HA_BOLT_PORT}", ' + f'"coordinator_server": "{name}:{HA_COORDINATOR_PORT}", ' + f'"management_server": "{name}:{HA_MANAGEMENT_PORT}"}}', + ) + + for name, data_host in HA_DATA_INSTANCES: + _ha_admin( + conn, + f'REGISTER INSTANCE {name} WITH CONFIG ' + f'{{"bolt_server": "{data_host}:{HA_BOLT_PORT}", ' + f'"management_server": "{data_host}:{HA_MANAGEMENT_PORT}", ' + f'"replication_server": "{data_host}:{HA_REPLICATION_PORT}"}}', + ) + + _ha_admin(conn, f"SET INSTANCE {HA_MAIN} TO MAIN") + + # Wait for the cluster to converge and advertise all roles: the main + # (WRITE), the replica (READ) and the coordinators (ROUTE). + timeout = 60 + for _ in range(timeout): + table = conn.get_routing_table() + roles = {server["role"] for server in table["servers"]} + if {"READ", "WRITE", "ROUTE"} <= roles: + break + time.sleep(1) + else: + conn.close() + raise RuntimeError(f"HA cluster did not converge: {table}") + + yield host, port + + conn.close() + + +def _parse_address_map(raw): + """Parse ``adv1=target1,adv2=target2`` into a ``dict``.""" + mapping = {} + for pair in raw.split(","): + pair = pair.strip() + if not pair: + continue + advertised, _, target = pair.partition("=") + mapping[advertised.strip()] = target.strip() + return mapping + + +def resolve_ha_address(address): + """Map a routing-table address to a locally reachable one. + + In CI the cluster runs on the same Docker network as the test runner, so + the advertised addresses (e.g. ``mg-data1:7687``) resolve directly and this + is the identity function. When running against a cluster whose advertised + addresses are not directly reachable (for example a Kubernetes cluster + reached through ``kubectl port-forward``), set ``MEMGRAPH_HA_ADDRESS_MAP`` + to a comma-separated list of ``advertised=target`` pairs, e.g.:: + + MEMGRAPH_HA_ADDRESS_MAP="memgraph-data-0.default.svc.cluster.local:7687=127.0.0.1:17690,..." + """ + raw = os.getenv("MEMGRAPH_HA_ADDRESS_MAP") + mapping = _parse_address_map(raw) if raw else {} + return mapping.get(address, address) + + +@pytest.fixture(scope="module") +def ha_resolver(): + """A ``resolver`` callable suitable for the current environment.""" + + def resolver(address): + return [resolve_ha_address(address)] + + return resolver diff --git a/test/test_connection.py b/test/test_connection.py index 282c1f2a..d4d6964b 100644 --- a/test/test_connection.py +++ b/test/test_connection.py @@ -15,7 +15,6 @@ import mgclient import pytest import tempfile -import time from common import ( start_memgraph, @@ -23,8 +22,6 @@ requires_ssl_enabled, requires_ssl_disabled, requires_ha_cluster, - MEMGRAPH_HA_COORDINATOR_HOST, - MEMGRAPH_HA_COORDINATOR_PORT, ) from OpenSSL import crypto @@ -125,78 +122,8 @@ def test_get_routing_table_closed_connection(memgraph_server): conn.get_routing_table() -# Topology created by the "Run Memgraph HA Cluster" CI step. The fixture below -# must agree with the container names and ports used there. -HA_COORDINATORS = ["mg-coord1", "mg-coord2", "mg-coord3"] -HA_DATA_INSTANCES = [("instance_1", "mg-data1"), ("instance_2", "mg-data2")] -HA_MAIN = "instance_1" -HA_BOLT_PORT = 7687 -HA_COORDINATOR_PORT = 12121 -HA_MANAGEMENT_PORT = 13011 -HA_REPLICATION_PORT = 10000 - - -def _ha_admin(conn, query): - """Run a coordinator admin query, tolerating idempotent re-runs.""" - cursor = conn.cursor() - try: - cursor.execute(query) - try: - cursor.fetchall() - except mgclient.Error: - pass - except mgclient.DatabaseError as exc: - # Re-running setup against an already-configured cluster is fine. - print(f"HA setup query ignored error: {exc}") - - -@pytest.fixture(scope="module") -def ha_cluster(): - host = MEMGRAPH_HA_COORDINATOR_HOST - port = MEMGRAPH_HA_COORDINATOR_PORT - - conn = mgclient.connect(host=host, port=port) - conn.autocommit = True - - # mg-coord1 is the bootstrap coordinator; add the remaining ones. - for cid, name in enumerate(HA_COORDINATORS, start=1): - if cid == 1: - continue - _ha_admin( - conn, - f'ADD COORDINATOR {cid} WITH CONFIG ' - f'{{"bolt_server": "{name}:{HA_BOLT_PORT}", ' - f'"coordinator_server": "{name}:{HA_COORDINATOR_PORT}", ' - f'"management_server": "{name}:{HA_MANAGEMENT_PORT}"}}', - ) - - for name, data_host in HA_DATA_INSTANCES: - _ha_admin( - conn, - f'REGISTER INSTANCE {name} WITH CONFIG ' - f'{{"bolt_server": "{data_host}:{HA_BOLT_PORT}", ' - f'"management_server": "{data_host}:{HA_MANAGEMENT_PORT}", ' - f'"replication_server": "{data_host}:{HA_REPLICATION_PORT}"}}', - ) - - _ha_admin(conn, f"SET INSTANCE {HA_MAIN} TO MAIN") - - # Wait for the cluster to converge and advertise all roles: the main - # (WRITE), the replica (READ) and the coordinators (ROUTE). - timeout = 60 - for _ in range(timeout): - table = conn.get_routing_table() - roles = {server["role"] for server in table["servers"]} - if {"READ", "WRITE", "ROUTE"} <= roles: - break - time.sleep(1) - else: - conn.close() - raise RuntimeError(f"HA cluster did not converge: {table}") - - yield host, port - - conn.close() +# The ``ha_cluster`` fixture lives in conftest.py so it can be shared with the +# client-side routing tests in test_routing.py. @requires_ha_cluster @@ -228,6 +155,21 @@ def test_connect_insecure_success(memgraph_server): assert conn.status == mgclient.CONN_STATUS_READY +@requires_ssl_disabled +def test_connect_routing_false_is_plain_connection(memgraph_server): + # routing=False must be indistinguishable from a plain connection: it goes + # straight to the given host/port with no ROUTE round-trip. + host, port, sslmode, _ = memgraph_server + conn = mgclient.connect(host=host, port=port, sslmode=sslmode, routing=False) + + assert conn.status == mgclient.CONN_STATUS_READY + + cursor = conn.cursor() + cursor.execute("RETURN 1") + assert cursor.fetchall() == [(1,)] + conn.close() + + @requires_ssl_disabled def test_connection_secure_fail(memgraph_server): # server doesn't use SSL diff --git a/test/test_routing.py b/test/test_routing.py new file mode 100644 index 00000000..d171fe85 --- /dev/null +++ b/test/test_routing.py @@ -0,0 +1,147 @@ +# Copyright (c) 2016-2020 Memgraph Ltd. [https://memgraph.com] +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for client-side routing (``connect(routing=True, ...)``). + +These run against a Memgraph high-availability cluster; they are skipped unless +``MEMGRAPH_HA_COORDINATOR_HOST`` is set. See ``conftest.resolve_ha_address`` +for how the advertised cluster addresses are mapped to reachable ones. +""" + +import mgclient +import pytest + +from common import requires_ha_cluster +from conftest import resolve_ha_address + + +def _replication_role(conn): + """Return 'main' or 'replica' for the instance a connection is bound to.""" + conn.autocommit = True + cursor = conn.cursor() + cursor.execute("SHOW REPLICATION ROLE") + return cursor.fetchall()[0][0] + + +def test_access_mode_constants(): + # The access-mode constants match the routing-table role names so they read + # naturally and can be used interchangeably with the plain strings. + assert mgclient.ACCESS_MODE_WRITE == "WRITE" + assert mgclient.ACCESS_MODE_READ == "READ" + + +def test_connect_routing_invalid_access_mode(): + # Validation happens before any connection is attempted, so this needs no + # server. + with pytest.raises(ValueError): + mgclient.connect( + host="127.0.0.1", port=7687, routing=True, access_mode="SIDEWAYS" + ) + + +def test_connect_rejects_positional_arguments(): + # connect() takes only keyword arguments, exactly like the underlying C + # connect, so a positional host is a TypeError rather than being silently + # treated as the routing flag. + with pytest.raises(TypeError): + mgclient.connect("127.0.0.1", 7687) + + +@requires_ha_cluster +def test_connect_routing_write_reaches_main(ha_cluster, ha_resolver): + host, port = ha_cluster + conn = mgclient.connect( + host=host, port=port, routing=True, access_mode="WRITE", resolver=ha_resolver + ) + assert conn.status == mgclient.CONN_STATUS_READY + assert _replication_role(conn) == "main" + conn.close() + + +@requires_ha_cluster +def test_connect_routing_read_reaches_replica(ha_cluster, ha_resolver): + host, port = ha_cluster + conn = mgclient.connect( + host=host, port=port, routing=True, access_mode="READ", resolver=ha_resolver + ) + assert conn.status == mgclient.CONN_STATUS_READY + assert _replication_role(conn) == "replica" + conn.close() + + +@requires_ha_cluster +def test_connect_routing_default_access_mode_is_write(ha_cluster, ha_resolver): + host, port = ha_cluster + conn = mgclient.connect(host=host, port=port, routing=True, resolver=ha_resolver) + assert _replication_role(conn) == "main" + conn.close() + + +@requires_ha_cluster +def test_connect_routing_access_mode_is_case_insensitive(ha_cluster, ha_resolver): + host, port = ha_cluster + conn = mgclient.connect( + host=host, port=port, routing=True, access_mode="read", resolver=ha_resolver + ) + assert _replication_role(conn) == "replica" + conn.close() + + +@requires_ha_cluster +def test_connect_routing_resolver_receives_advertised_addresses(ha_cluster): + host, port = ha_cluster + + seen = [] + + def resolver(address): + seen.append(address) + return [resolve_ha_address(address)] + + conn = mgclient.connect( + host=host, port=port, routing=True, access_mode="WRITE", resolver=resolver + ) + # The resolver is consulted with the advertised "host:port" addresses from + # the routing table. + assert seen + assert all(":" in address for address in seen) + assert _replication_role(conn) == "main" + conn.close() + + +@requires_ha_cluster +def test_connect_routing_fails_over_across_candidates(ha_cluster): + host, port = ha_cluster + + # The first candidate is unreachable; routing must fall back to the next. + def resolver(address): + return ["127.0.0.1:1", resolve_ha_address(address)] + + conn = mgclient.connect( + host=host, port=port, routing=True, access_mode="WRITE", resolver=resolver + ) + assert _replication_role(conn) == "main" + conn.close() + + +@requires_ha_cluster +def test_connect_routing_all_candidates_unreachable(ha_cluster): + host, port = ha_cluster + + def resolver(address): + return ["127.0.0.1:1"] + + with pytest.raises(mgclient.OperationalError): + mgclient.connect( + host=host, port=port, routing=True, access_mode="WRITE", resolver=resolver + ) From a6396a32a0e8abf2220adec21cb4d5333083106e Mon Sep 17 00:00:00 2001 From: matt Date: Wed, 1 Jul 2026 22:18:14 +0100 Subject: [PATCH 02/10] added router --- docs/source/module.rst | 8 + python/mgclient/__init__.py | 1 + python/mgclient/routing.py | 333 +++++++++++++++++++++++++++++------- test/test_routing.py | 164 ++++++++++++++++++ 4 files changed, 442 insertions(+), 64 deletions(-) diff --git a/docs/source/module.rst b/docs/source/module.rst index 75fb3099..bf747702 100644 --- a/docs/source/module.rst +++ b/docs/source/module.rst @@ -35,6 +35,14 @@ If the addresses a cluster advertises are not directly reachable by the client pass a ``resolver`` callable that maps an advertised ``"host:port"`` address to an iterable of ``"host:port"`` targets to try. +``connect(routing=True, ...)`` performs a fresh routing lookup on every call. +For a long-lived router that caches the routing table (honouring its TTL), +balances reads across replicas and fails over across coordinators, use the +:class:`Router` class: + +.. autoclass:: mgclient.Router + :members: connect, refresh, routing_table + For lower-level access to the routing table itself, see :meth:`Connection.get_routing_table`. diff --git a/python/mgclient/__init__.py b/python/mgclient/__init__.py index e4b235bd..f9705763 100644 --- a/python/mgclient/__init__.py +++ b/python/mgclient/__init__.py @@ -33,5 +33,6 @@ from mgclient.routing import ( # noqa: F401,E402 ACCESS_MODE_READ, ACCESS_MODE_WRITE, + Router, connect, ) diff --git a/python/mgclient/routing.py b/python/mgclient/routing.py index 2bbb1c9c..5735fc25 100644 --- a/python/mgclient/routing.py +++ b/python/mgclient/routing.py @@ -18,8 +18,16 @@ the :mod:`mgclient._mgclient` C extension (:func:`connect` and :meth:`Connection.get_routing_table`). It is pure Python and opens no connections beyond the ones the C extension already makes. + +The :class:`Router` class is the reusable engine: it caches the routing table +(honouring its TTL), balances reads across replicas, and fails over across +coordinators when refreshing. :func:`connect` with ``routing=True`` is a +one-shot convenience built on top of it. """ +import threading +import time + from mgclient._mgclient import connect as _base_connect from mgclient._mgclient import OperationalError @@ -36,6 +44,16 @@ def _default_resolver(address): return [address] +def _normalize_access_mode(access_mode): + if isinstance(access_mode, str): + access_mode = access_mode.upper() + if access_mode not in _ACCESS_MODES: + raise ValueError( + f"access_mode must be one of {_ACCESS_MODES!r}, got {access_mode!r}" + ) + return access_mode + + def _split_host_port(address): """Split a routing-table ``"host:port"`` address into ``(host, port)``.""" host, separator, port = address.rpartition(":") @@ -50,70 +68,257 @@ def _split_host_port(address): return host, port -def _routed_connect(access_mode, resolver, routing_context, kwargs): - if isinstance(access_mode, str): - access_mode = access_mode.upper() - if access_mode not in _ACCESS_MODES: - raise ValueError( - f"access_mode must be one of {_ACCESS_MODES!r}, got {access_mode!r}" +class RoutingTable: + """A parsed routing table: advertised addresses grouped by role. + + Addresses are the raw ``"host:port"`` strings advertised by the + coordinator, before any :class:`Router` ``resolver`` is applied. + """ + + def __init__(self, write, read, route, ttl): + self.write = list(write) + self.read = list(read) + self.route = list(route) + self.ttl = ttl + + @classmethod + def parse(cls, raw): + """Build a :class:`RoutingTable` from a ``get_routing_table`` result.""" + buckets = { + ACCESS_MODE_WRITE: [], + ACCESS_MODE_READ: [], + "ROUTE": [], + } + for server in raw.get("servers", []): + role = server.get("role") + if role in buckets: + buckets[role].extend(server.get("addresses", [])) + return cls( + buckets[ACCESS_MODE_WRITE], + buckets[ACCESS_MODE_READ], + buckets["ROUTE"], + int(raw.get("ttl", 0)), ) - if resolver is None: - resolver = _default_resolver - # Bootstrap: connect to the coordinator and fetch the routing table. ROUTE - # requires an idle, non-lazy session, so drop "lazy" for this connection - # only (the final connection still honours the caller's choice). - bootstrap_kwargs = dict(kwargs) - bootstrap_kwargs.pop("lazy", None) + def servers(self, access_mode): + """Return the advertised addresses serving ``access_mode``.""" + if access_mode == ACCESS_MODE_WRITE: + return self.write + return self.read + + +class Router: + """Routing engine for a Memgraph high-availability cluster. + + A :class:`Router` is created against one or more coordinators and is meant + to be long-lived and reused. It fetches the cluster routing table via a + Bolt ``ROUTE`` message and caches it until its TTL expires; :meth:`connect` + then hands out ordinary :class:`Connection` objects bound to the + appropriate data instance. + + Reads are balanced across the available replicas in round-robin order, and + both the ``ROUTE`` refresh (across coordinators) and the data-instance + connect (across a role's servers) fail over across candidates. If every + candidate for a requested access mode is unreachable, the cached table is + refreshed and the attempt retried once before giving up. + + All connection parameters other than ``host``/``port``/``address`` (for + example ``username``, ``password`` and the SSL options) are stored and + reused for both the coordinator and the data-instance connections. + + Parameters: + + * :obj:`host` / :obj:`address` / :obj:`port` + + The seed coordinator to contact for the first routing-table fetch. + + * :obj:`resolver` + + Optional callable mapping an advertised ``"host:port"`` address to + an iterable of ``"host:port"`` targets to try, in order. Defaults + to using the advertised address unchanged. + + * :obj:`routing_context` + + Optional :class:`dict` passed through to + :meth:`Connection.get_routing_table`. + """ + + def __init__( + self, + *, + host=None, + address=None, + port=None, + resolver=None, + routing_context=None, + **connect_kwargs, + ): + # The seed coordinator address. "address" (a numeric IP) and "host" are + # interchangeable connect targets; store whichever was given. + self._seed_host = address or host + self._seed_port = port + + self._resolver = resolver if resolver is not None else _default_resolver + self._routing_context = ( + routing_context if routing_context is not None else {} + ) + + # Parameters reused for the data-instance connection (honours "lazy") + # and for the coordinator connection (ROUTE needs an idle, non-lazy + # session, so "lazy" is dropped there). + self._data_kwargs = dict(connect_kwargs) + self._coordinator_kwargs = dict(connect_kwargs) + self._coordinator_kwargs.pop("lazy", None) + + self._table = None + self._expires_at = 0.0 + self._read_index = 0 + self._refresh_count = 0 + self._lock = threading.Lock() + + # -- routing table management ------------------------------------------- + + def _router_targets_locked(self): + """Coordinator ``(host, port)`` targets to try when refreshing. + + The seed comes first, followed by the ROUTE-role servers from the last + known table (resolved through ``resolver``), so a refresh survives the + seed coordinator going away. + """ + targets = [] + if self._seed_host is not None: + targets.append((self._seed_host, self._seed_port)) + if self._table is not None: + for advertised in self._table.route: + for target in self._resolver(advertised): + host_port = _split_host_port(target) + if host_port not in targets: + targets.append(host_port) + return targets + + def _refresh_locked(self): + """Fetch a fresh routing table; assumes ``self._lock`` is held.""" + errors = [] + for host, port in self._router_targets_locked(): + try: + coordinator = _base_connect( + host=host, port=port, **self._coordinator_kwargs + ) + except OperationalError as exc: + errors.append(f"{host}:{port}: {exc}") + continue + try: + raw = coordinator.get_routing_table( + routing_context=self._routing_context + ) + finally: + coordinator.close() + self._table = RoutingTable.parse(raw) + self._expires_at = time.monotonic() + max(self._table.ttl, 0) + self._refresh_count += 1 + return - coordinator = _base_connect(**bootstrap_kwargs) - try: - context = routing_context if routing_context is not None else {} - table = coordinator.get_routing_table(routing_context=context) - finally: - coordinator.close() - - # Collect the advertised addresses of every server that serves the - # requested access mode, preserving the coordinator's ordering. - advertised = [] - for server in table.get("servers", []): - if server.get("role") == access_mode: - advertised.extend(server.get("addresses", [])) - - if not advertised: raise OperationalError( - f"no {access_mode} server available in the routing table" + "could not fetch routing table from any coordinator: " + + "; ".join(errors) ) - # Map advertised addresses to reachable connect targets, de-duplicating - # while keeping order so failover is deterministic. - candidates = [] - for address in advertised: - for target in resolver(address): - if target not in candidates: - candidates.append(target) - - # The final connection reuses the caller's parameters (auth, SSL, lazy, ...) - # but is directed at the resolved data instance. "address" (a numeric IP - # for the coordinator) never applies to the routed target. - final_kwargs = dict(kwargs) - final_kwargs.pop("address", None) - - errors = [] - for target in candidates: - host, port = _split_host_port(target) - attempt = dict(final_kwargs) - attempt["host"] = host - attempt["port"] = port - try: - return _base_connect(**attempt) - except OperationalError as exc: - errors.append(f"{target}: {exc}") - - raise OperationalError( - f"could not connect to any {access_mode} server; tried " - + "; ".join(errors) - ) + def _ensure_fresh(self): + with self._lock: + if self._table is None or time.monotonic() >= self._expires_at: + self._refresh_locked() + + def refresh(self): + """Force an immediate refresh of the cached routing table.""" + with self._lock: + self._refresh_locked() + + @property + def routing_table(self): + """A snapshot of the (refreshed if stale) routing table as a dict.""" + self._ensure_fresh() + with self._lock: + table = self._table + return { + "ttl": table.ttl, + "write": list(table.write), + "read": list(table.read), + "route": list(table.route), + } + + # -- server selection ---------------------------------------------------- + + def _ordered_targets(self, access_mode): + """Resolved ``"host:port"`` candidates for ``access_mode``, in order. + + Reads round-robin across replicas so successive READ connections start + at different servers; the remaining servers stay as failover + candidates. + """ + with self._lock: + if self._table is None: + servers = [] + else: + servers = self._table.servers(access_mode) + if access_mode == ACCESS_MODE_READ and servers: + start = self._read_index % len(servers) + self._read_index += 1 + servers = servers[start:] + servers[:start] + else: + servers = list(servers) + + candidates = [] + for advertised in servers: + for target in self._resolver(advertised): + if target not in candidates: + candidates.append(target) + return candidates + + def _try_connect(self, candidates): + errors = [] + for target in candidates: + host, port = _split_host_port(target) + attempt = dict(self._data_kwargs) + attempt["host"] = host + attempt["port"] = port + try: + return _base_connect(**attempt), None + except OperationalError as exc: + errors.append(f"{target}: {exc}") + return None, "; ".join(errors) + + def connect(self, access_mode=ACCESS_MODE_WRITE): + """Open a connection to a data instance serving ``access_mode``. + + Returns an ordinary :class:`Connection`. Raises + :exc:`OperationalError` if no server for the requested access mode can + be reached, even after refreshing the routing table. + """ + access_mode = _normalize_access_mode(access_mode) + + last_error = None + # Two attempts: the second runs against a freshly refreshed table in + # case the topology changed (e.g. a failover) since it was cached. + for attempt in range(2): + self._ensure_fresh() + candidates = self._ordered_targets(access_mode) + if candidates: + connection, error = self._try_connect(candidates) + if connection is not None: + return connection + last_error = error + else: + last_error = f"no {access_mode} server in the routing table" + + if attempt == 0: + # Selected servers were unreachable (or none were listed); + # discard the cached table and retry with a fresh one. + self.refresh() + + raise OperationalError( + f"could not connect to any {access_mode} server: {last_error}" + ) def connect( @@ -200,6 +405,10 @@ def connect( selected server cannot be reached, the remaining candidates are tried in turn; an :exc:`OperationalError` is raised only if none succeed. + This is a one-shot convenience: it performs a fresh ``ROUTE`` lookup on + every call. For a long-lived, caching, load-balancing router, use + :class:`Router` directly. + All other connection parameters (authentication, SSL, ``lazy``, ...) are applied to the returned data-instance connection. @@ -232,9 +441,5 @@ def connect( if not routing: return _base_connect(**kwargs) - return _routed_connect( - access_mode=access_mode, - resolver=resolver, - routing_context=routing_context, - kwargs=kwargs, - ) + router = Router(resolver=resolver, routing_context=routing_context, **kwargs) + return router.connect(access_mode=access_mode) diff --git a/test/test_routing.py b/test/test_routing.py index d171fe85..764a7b21 100644 --- a/test/test_routing.py +++ b/test/test_routing.py @@ -24,6 +24,7 @@ from common import requires_ha_cluster from conftest import resolve_ha_address +from mgclient.routing import Router, RoutingTable def _replication_role(conn): @@ -145,3 +146,166 @@ def resolver(address): mgclient.connect( host=host, port=port, routing=True, access_mode="WRITE", resolver=resolver ) + + +# --------------------------------------------------------------------------- +# Router: cached routing table, load balancing and failover (no cluster needed +# for the pure-logic tests below). +# --------------------------------------------------------------------------- + + +def test_routing_table_parse_groups_addresses_by_role(): + raw = { + "ttl": 120, + "servers": [ + {"role": "WRITE", "addresses": ["m:7687"]}, + {"role": "READ", "addresses": ["r1:7687", "r2:7687"]}, + {"role": "ROUTE", "addresses": ["c1:7687", "c2:7687"]}, + {"role": "SOMETHING_ELSE", "addresses": ["x:7687"]}, + ], + } + table = RoutingTable.parse(raw) + + assert table.write == ["m:7687"] + assert table.read == ["r1:7687", "r2:7687"] + assert table.route == ["c1:7687", "c2:7687"] + assert table.ttl == 120 + + +def _router_with_cached_table(table): + # Router.__init__ opens no connection; we inject a table and a far-future + # expiry to exercise selection logic without a cluster. + router = Router(host="unused", port=7687) + router._table = table + router._expires_at = float("inf") + return router + + +def test_read_targets_round_robin_across_replicas(): + table = RoutingTable(["m:7687"], ["r1:7687", "r2:7687", "r3:7687"], ["c:7687"], 300) + router = _router_with_cached_table(table) + + firsts = [router._ordered_targets("READ")[0] for _ in range(6)] + + # Each READ selection starts at the next replica, cycling through them all. + assert firsts == [ + "r1:7687", + "r2:7687", + "r3:7687", + "r1:7687", + "r2:7687", + "r3:7687", + ] + # Every replica remains a failover candidate regardless of starting point. + assert set(router._ordered_targets("READ")) == {"r1:7687", "r2:7687", "r3:7687"} + + +def test_write_targets_do_not_rotate(): + table = RoutingTable(["m:7687"], ["r1:7687", "r2:7687"], ["c:7687"], 300) + router = _router_with_cached_table(table) + + assert router._ordered_targets("WRITE") == ["m:7687"] + assert router._ordered_targets("WRITE") == ["m:7687"] + + +@requires_ha_cluster +def test_router_connect_reaches_correct_instances(ha_cluster, ha_resolver): + host, port = ha_cluster + router = Router(host=host, port=port, resolver=ha_resolver) + + writer = router.connect(access_mode="WRITE") + assert _replication_role(writer) == "main" + writer.close() + + reader = router.connect(access_mode="READ") + assert _replication_role(reader) == "replica" + reader.close() + + +@requires_ha_cluster +def test_router_caches_routing_table(ha_cluster, ha_resolver): + host, port = ha_cluster + router = Router(host=host, port=port, resolver=ha_resolver) + + router.connect(access_mode="WRITE").close() + assert router._refresh_count == 1 + + # Subsequent connects within the TTL are served from the cached table, with + # no further ROUTE round-trips to a coordinator. + router.connect(access_mode="READ").close() + router.connect(access_mode="WRITE").close() + assert router._refresh_count == 1 + + +@requires_ha_cluster +def test_router_refresh_forces_new_lookup(ha_cluster, ha_resolver): + host, port = ha_cluster + router = Router(host=host, port=port, resolver=ha_resolver) + + router.connect(access_mode="WRITE").close() + before = router._refresh_count + + router.refresh() + assert router._refresh_count == before + 1 + + +@requires_ha_cluster +def test_router_expired_table_is_refreshed(ha_cluster, ha_resolver): + host, port = ha_cluster + router = Router(host=host, port=port, resolver=ha_resolver) + + router.connect(access_mode="WRITE").close() + before = router._refresh_count + + # Simulate the TTL having elapsed. + router._expires_at = 0.0 + router.connect(access_mode="WRITE").close() + assert router._refresh_count == before + 1 + + +@requires_ha_cluster +def test_router_survives_seed_coordinator_failure(ha_cluster, ha_resolver): + # After the first refresh the table lists every coordinator. If the seed + # coordinator later becomes unreachable, a refresh must fall back to a + # ROUTE-role server from the cached table. + host, port = ha_cluster + router = Router(host=host, port=port, resolver=ha_resolver) + + router.connect(access_mode="WRITE").close() + + # Break the seed; refresh should still succeed via a cached coordinator. + router._seed_host, router._seed_port = "127.0.0.1", 1 + router.refresh() + + writer = router.connect(access_mode="WRITE") + assert _replication_role(writer) == "main" + writer.close() + + +@requires_ha_cluster +def test_router_refreshes_when_targets_unreachable(ha_cluster): + host, port = ha_cluster + + # The coordinator (seed) is reachable, but the data instances are not, so + # routing should refresh and retry before giving up. + def resolver(address): + return ["127.0.0.1:1"] + + router = Router(host=host, port=port, resolver=resolver) + + with pytest.raises(mgclient.OperationalError): + router.connect(access_mode="WRITE") + + assert router._refresh_count >= 2 + + +@requires_ha_cluster +def test_router_routing_table_property(ha_cluster, ha_resolver): + host, port = ha_cluster + router = Router(host=host, port=port, resolver=ha_resolver) + + table = router.routing_table + assert table["write"] + assert table["read"] + assert table["route"] + assert isinstance(table["ttl"], int) From 66b3b35bc2918a2e5659719079a8636b8c52d1cf Mon Sep 17 00:00:00 2001 From: matt Date: Mon, 6 Jul 2026 14:46:43 +0100 Subject: [PATCH 03/10] managed retry logic --- docs/source/module.rst | 17 +++- python/mgclient/__init__.py | 2 + python/mgclient/routing.py | 158 +++++++++++++++++++++++++++++- test/test_routing.py | 189 +++++++++++++++++++++++++++++++++++- 4 files changed, 363 insertions(+), 3 deletions(-) diff --git a/docs/source/module.rst b/docs/source/module.rst index bf747702..1eb9131e 100644 --- a/docs/source/module.rst +++ b/docs/source/module.rst @@ -41,7 +41,22 @@ balances reads across replicas and fails over across coordinators, use the :class:`Router` class: .. autoclass:: mgclient.Router - :members: connect, refresh, routing_table + :members: connect, execute_read, execute_write, refresh, routing_table + +:meth:`Router.execute_read` and :meth:`Router.execute_write` are *managed +transactions*: they run your unit of work against the right instance and +automatically retry transient failover conditions (a new main being elected, a +replica catching up, an instance dropping mid-request) with a routing refresh +and capped exponential backoff. A write that committed on the main but could +not reach a synchronous replica is treated as success rather than retried +(retrying would duplicate it). + +The classification used for retries is also exposed for building your own retry +loops: + +.. autofunction:: mgclient.is_transient_error + +.. autofunction:: mgclient.is_committed_on_main_error For lower-level access to the routing table itself, see :meth:`Connection.get_routing_table`. diff --git a/python/mgclient/__init__.py b/python/mgclient/__init__.py index f9705763..5d93e742 100644 --- a/python/mgclient/__init__.py +++ b/python/mgclient/__init__.py @@ -35,4 +35,6 @@ ACCESS_MODE_WRITE, Router, connect, + is_committed_on_main_error, + is_transient_error, ) diff --git a/python/mgclient/routing.py b/python/mgclient/routing.py index 5735fc25..30fb4161 100644 --- a/python/mgclient/routing.py +++ b/python/mgclient/routing.py @@ -29,7 +29,7 @@ import time from mgclient._mgclient import connect as _base_connect -from mgclient._mgclient import OperationalError +from mgclient._mgclient import Error, OperationalError #: Route the connection to a server that accepts writes (the cluster main). ACCESS_MODE_WRITE = "WRITE" @@ -38,6 +38,63 @@ _ACCESS_MODES = (ACCESS_MODE_WRITE, ACCESS_MODE_READ) +# Default retry budget for managed transactions (see Router.execute_read / +# execute_write). Retries are generous because a cluster stays available as a +# whole while individual instances flap during a failover, but the backoff is +# capped so total time is bounded. +_DEFAULT_MAX_RETRIES = 8 +_DEFAULT_RETRY_BACKOFF = 1.0 +_DEFAULT_RETRY_BACKOFF_CAP = 15.0 + +# Substrings that mark a transient cluster condition worth retrying: a failover +# in progress, a replica catching up, or an instance dropped/killed mid-request. +# They clear once the cluster settles. Matched case-insensitively against the +# exception message. Adapted from hard-won experience running the neo4j driver +# against a Memgraph HA cluster under chaos. +_TRANSIENT_MESSAGE_FRAGMENTS = ( + "replication exception", # SYNC replica lagging/unreachable + "forbidden on the", # write forbidden on a replica or the main mid-election + "setting up a new main", # coordinator electing a new main; no writer yet + "please retry the query", # Memgraph's explicit "retry later" during failover + "database shutdown", # instance being restarted mid-transaction + "asked to abort", # transaction aborted by a shutting-down instance + "defunct connection", # connection dropped underneath the driver + "failed to read from", # partial/dropped Bolt read + "failed to receive chunk", # connection reset mid-Bolt-stream + "connection reset", # peer went away + "broken pipe", + "couldn't connect", # candidate instance not accepting connections yet + "connection refused", + "routing information", # unable to retrieve routing info during a failover + # Raised by Router itself while the cluster has no server for a role yet. + "no write server", + "no read server", + "could not connect to any", + "could not fetch routing table", +) + + +def is_transient_error(exc): + """True for a transient HA condition worth retrying after a short backoff. + + These arise during a failover, while a replica catches up, or when an + instance is dropped mid-request; they clear once the cluster reconverges. + """ + message = str(getattr(exc, "message", "") or exc).lower() + return any(fragment in message for fragment in _TRANSIENT_MESSAGE_FRAGMENTS) + + +def is_committed_on_main_error(exc): + """True only when a write committed on the main but a SYNC replica was down. + + Narrower than :func:`is_transient_error`: this is the one case where the + write is already durable (Memgraph says so in the message), so it is safe to + treat as success. Retrying it instead would duplicate the write, so managed + writes must check this *before* the generic transient check. + """ + message = str(getattr(exc, "message", "") or exc).lower() + return "replication exception" in message and "committed on the main" in message + def _default_resolver(address): """Identity resolver: the advertised address is used as-is.""" @@ -152,6 +209,9 @@ def __init__( port=None, resolver=None, routing_context=None, + max_retries=_DEFAULT_MAX_RETRIES, + retry_backoff=_DEFAULT_RETRY_BACKOFF, + retry_backoff_cap=_DEFAULT_RETRY_BACKOFF_CAP, **connect_kwargs, ): # The seed coordinator address. "address" (a numeric IP) and "host" are @@ -164,6 +224,11 @@ def __init__( routing_context if routing_context is not None else {} ) + # Managed-transaction retry budget (see execute_read / execute_write). + self._max_retries = max_retries + self._retry_backoff = retry_backoff + self._retry_backoff_cap = retry_backoff_cap + # Parameters reused for the data-instance connection (honours "lazy") # and for the coordinator connection (ROUTE needs an idle, non-lazy # session, so "lazy" is dropped there). @@ -320,6 +385,97 @@ def connect(self, access_mode=ACCESS_MODE_WRITE): f"could not connect to any {access_mode} server: {last_error}" ) + # -- managed transactions ------------------------------------------------ + + def _retry_delay(self, attempt): + """Capped exponential backoff: base, 2*base, 4*base, ... up to the cap.""" + return min( + self._retry_backoff * (2 ** (attempt - 1)), + self._retry_backoff_cap, + ) + + def _refresh_quietly(self): + """Refresh the routing table, ignoring failures (a coordinator may be + briefly unreachable mid-failover; the next attempt tries again).""" + try: + self.refresh() + except Error: + pass + + def _run_managed(self, access_mode, work): + writing = access_mode == ACCESS_MODE_WRITE + last_exc = None + + for attempt in range(1, self._max_retries + 1): + connection = None + result = None + try: + connection = self.connect(access_mode=access_mode) + cursor = connection.cursor() + if writing: + # Explicit transaction so the whole unit of work is atomic + # and any SYNC-replication error surfaces at commit (after + # ``work`` has produced its result). + connection.autocommit = False + result = work(cursor) + connection.commit() + else: + connection.autocommit = True + result = work(cursor) + return result + except Error as exc: + last_exc = exc + # A write that committed on the main but couldn't reach a SYNC + # replica is durable -- treat it as success rather than retrying + # (a retry would duplicate the write). + if writing and is_committed_on_main_error(exc): + return result + if attempt == self._max_retries or not is_transient_error(exc): + raise + finally: + if connection is not None: + connection.close() + + # Transient failure: refresh routing (so a retry re-routes to the + # new main) and back off before trying again. + self._refresh_quietly() + time.sleep(self._retry_delay(attempt)) + + # Unreachable in practice (the loop either returns or raises), but keep + # a definite fallback. + raise last_exc + + def execute_read(self, work): + """Run ``work(cursor)`` as a managed read against a replica. + + ``work`` receives a :class:`Cursor` from a freshly routed READ + connection and returns whatever the caller wants; its return value is + returned from :meth:`execute_read`. On a transient cluster condition + (see :func:`is_transient_error`) the routing table is refreshed and the + work is retried with capped exponential backoff, up to ``max_retries``. + + ``work`` may be called more than once, so it should be free of side + effects other than the database operations themselves. + """ + return self._run_managed(ACCESS_MODE_READ, work) + + def execute_write(self, work): + """Run ``work(cursor)`` as a managed write against the main. + + Like :meth:`execute_read`, but the work runs inside an explicit + transaction that is committed for you, and it is routed to the main. + + Transient failover conditions are retried (with a routing refresh and + backoff between attempts). The one exception is a write that committed + on the main but could not reach a SYNC replica: that write is durable, + so it is treated as success and *not* retried (retrying would duplicate + it). + + ``work`` may be called more than once, so it should be free of side + effects other than the database operations themselves. + """ + return self._run_managed(ACCESS_MODE_WRITE, work) + def connect( *, diff --git a/test/test_routing.py b/test/test_routing.py index 764a7b21..cfcd3831 100644 --- a/test/test_routing.py +++ b/test/test_routing.py @@ -24,7 +24,13 @@ from common import requires_ha_cluster from conftest import resolve_ha_address -from mgclient.routing import Router, RoutingTable +import mgclient.routing as routing +from mgclient.routing import ( + Router, + RoutingTable, + is_committed_on_main_error, + is_transient_error, +) def _replication_role(conn): @@ -309,3 +315,184 @@ def test_router_routing_table_property(ha_cluster, ha_resolver): assert table["read"] assert table["route"] assert isinstance(table["ttl"], int) + + +# --------------------------------------------------------------------------- +# Managed retry / transaction functions (no cluster needed for these). +# --------------------------------------------------------------------------- + +# Real failover messages observed against a chaos cluster. +_FAILOVER_MESSAGES = [ + "Write queries currently forbidden on the main instance. The cluster is in " + "the process of setting up a new main instance, please retry the query " + "later on.", + "could not connect to any WRITE server: no WRITE server in the routing table", + "memgraph-data-0:7687: couldn't connect to host: Connection refused", + "failed to receive chunk size", +] + +_COMMITTED_ON_MAIN_MESSAGE = ( + "Replication Exception: Failed to replicate to SYNC replica 'instance_1': " + "replica is not reachable or not in sync with the main. Replica will be " + "recovered automatically. Transaction is still committed on the main " + "instance and other alive replicas." +) + + +@pytest.mark.parametrize("message", _FAILOVER_MESSAGES) +def test_is_transient_error_matches_failover_conditions(message): + assert is_transient_error(mgclient.DatabaseError(message)) + + +def test_is_transient_error_false_for_query_errors(): + assert not is_transient_error(mgclient.DatabaseError("Syntax error near 'FOO'")) + + +def test_committed_on_main_error_needs_both_markers(): + assert is_committed_on_main_error(mgclient.DatabaseError(_COMMITTED_ON_MAIN_MESSAGE)) + # A replication error without the durability note is not committed-on-main. + assert not is_committed_on_main_error( + mgclient.DatabaseError("Replication Exception: something unrelated") + ) + assert not is_committed_on_main_error(mgclient.DatabaseError("Syntax error")) + + +class _FakeCursor: + def execute(self, *args, **kwargs): + pass + + def fetchall(self): + return [] + + +class _FakeConnection: + """A stand-in connection so the retry loop can be tested without a server.""" + + commit_error = None + + def __init__(self): + self.autocommit = False + + def cursor(self): + return _FakeCursor() + + def commit(self): + if self.commit_error is not None: + raise self.commit_error + + def close(self): + pass + + +def _router_with_stubbed_connect(monkeypatch, connection_factory): + router = Router(host="unused", port=7687) + monkeypatch.setattr( + router, "connect", lambda access_mode=routing.ACCESS_MODE_WRITE: connection_factory() + ) + monkeypatch.setattr(router, "refresh", lambda: None) + monkeypatch.setattr(routing.time, "sleep", lambda _seconds: None) + return router + + +def test_execute_write_retries_transient_then_succeeds(monkeypatch): + router = _router_with_stubbed_connect(monkeypatch, _FakeConnection) + + calls = {"n": 0} + + def work(cursor): + calls["n"] += 1 + if calls["n"] < 3: + raise mgclient.DatabaseError( + "setting up a new main, please retry the query later on" + ) + return "done" + + assert router.execute_write(work) == "done" + assert calls["n"] == 3 + + +def test_execute_write_treats_committed_on_main_as_success(monkeypatch): + class CommitFails(_FakeConnection): + commit_error = mgclient.DatabaseError(_COMMITTED_ON_MAIN_MESSAGE) + + router = _router_with_stubbed_connect(monkeypatch, CommitFails) + + calls = {"n": 0} + + def work(cursor): + calls["n"] += 1 + return "created" + + # The write is durable on the main, so it counts as success and must NOT be + # retried (retrying would duplicate the write). + assert router.execute_write(work) == "created" + assert calls["n"] == 1 + + +def test_execute_write_raises_non_transient_immediately(monkeypatch): + router = _router_with_stubbed_connect(monkeypatch, _FakeConnection) + + calls = {"n": 0} + + def work(cursor): + calls["n"] += 1 + raise mgclient.DatabaseError("Syntax error near 'FOO'") + + with pytest.raises(mgclient.DatabaseError): + router.execute_write(work) + assert calls["n"] == 1 + + +def test_execute_write_gives_up_after_max_retries(monkeypatch): + router = _router_with_stubbed_connect(monkeypatch, _FakeConnection) + + calls = {"n": 0} + + def work(cursor): + calls["n"] += 1 + raise mgclient.DatabaseError("please retry the query later on") + + with pytest.raises(mgclient.DatabaseError): + router.execute_write(work) + assert calls["n"] == router._max_retries + + +def test_execute_read_retries_transient_then_succeeds(monkeypatch): + router = _router_with_stubbed_connect(monkeypatch, _FakeConnection) + + calls = {"n": 0} + + def work(cursor): + calls["n"] += 1 + if calls["n"] < 2: + raise mgclient.OperationalError( + "could not connect to any READ server: no READ server in the routing table" + ) + return 7 + + assert router.execute_read(work) == 7 + assert calls["n"] == 2 + + +@requires_ha_cluster +def test_execute_write_and_read_roundtrip(ha_cluster, ha_resolver): + host, port = ha_cluster + router = Router(host=host, port=port, resolver=ha_resolver) + + def create(cursor): + cursor.execute( + "CREATE (n:MgExecTest {tag: $tag}) RETURN n.tag", {"tag": "exec"} + ) + return cursor.fetchall()[0][0] + + assert router.execute_write(create) == "exec" + + def count(cursor): + cursor.execute("MATCH (n:MgExecTest {tag: 'exec'}) RETURN count(n)") + return cursor.fetchall()[0][0] + + assert router.execute_read(count) >= 1 + + router.execute_write( + lambda cursor: cursor.execute("MATCH (n:MgExecTest) DETACH DELETE n") + ) From 0cc0f96e1101c7ad21f878d06fba13ac9a5cad0b Mon Sep 17 00:00:00 2001 From: matt Date: Mon, 6 Jul 2026 15:56:38 +0100 Subject: [PATCH 04/10] update usage --- docs/source/usage.rst | 60 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/docs/source/usage.rst b/docs/source/usage.rst index 1e4fd596..4641f5a1 100644 --- a/docs/source/usage.rst +++ b/docs/source/usage.rst @@ -124,3 +124,63 @@ will be immediately committed and no rollback is possible. A few commands (``CREATE INDEX``, ``CREATE USER`` and similar) require to be run outside any transaction. To set the connection in *autocommit* mode, set :attr:`.autocommit` property of the connection to ``True``. + +############################################## +Connecting to a high-availability cluster +############################################## + +When connecting to a Memgraph high-availability (HA) cluster you connect to a +*coordinator* and let :mod:`mgclient` route each query to the right data +instance: writes to the main, reads to a replica. Pass ``routing=True`` to +:func:`.connect` and choose an ``access_mode``:: + + >>> import mgclient + + # Routed to the main (writes) + >>> conn = mgclient.connect(host="coordinator-1", port=7687, + ... username="user", password="pass", + ... routing=True, + ... access_mode=mgclient.ACCESS_MODE_WRITE) + >>> cursor = conn.cursor() + >>> cursor.execute("CREATE (:Person {name: $name})", {"name": "Ada"}) + >>> conn.commit() + + # Routed to a replica (reads) + >>> reader = mgclient.connect(host="coordinator-1", port=7687, + ... username="user", password="pass", + ... routing=True, + ... access_mode=mgclient.ACCESS_MODE_READ) + +Each :func:`.connect` call above performs a fresh routing lookup. For a +long-lived client, use a :class:`Router` instead: it caches the routing table, +balances reads across replicas, and fails over across coordinators:: + + >>> router = mgclient.Router(host="coordinator-1", port=7687, + ... username="user", password="pass") + >>> writer = router.connect(access_mode=mgclient.ACCESS_MODE_WRITE) + >>> reader = router.connect(access_mode=mgclient.ACCESS_MODE_READ) + +To keep working through failovers (a new main being elected, a replica catching +up, an instance dropping), run your queries through the *managed* transaction +helpers, which retry transient conditions automatically:: + + >>> def create_person(cursor): + ... cursor.execute("CREATE (p:Person {name: $n}) RETURN id(p)", {"n": "Ada"}) + ... return cursor.fetchall()[0][0] + >>> new_id = router.execute_write(create_person) # committed for you + + >>> def count_people(cursor): + ... cursor.execute("MATCH (:Person) RETURN count(*)") + ... return cursor.fetchall()[0][0] + >>> total = router.execute_read(count_people) + +If the addresses the cluster advertises are not directly reachable from the +client (for example in-cluster Kubernetes names reached from outside), pass a +``resolver`` that maps an advertised ``"host:port"`` to reachable targets:: + + >>> def resolver(advertised): + ... mapping = {"memgraph-data-0.default.svc.cluster.local:7687": "10.0.0.5:7687"} + ... return [mapping.get(advertised, advertised)] + >>> router = mgclient.Router(host="coordinator-1", port=7687, resolver=resolver) + +See the :doc:`module` reference for the full routing API. From 6e08a974a63c4306da417aced763ccc20a70b815 Mon Sep 17 00:00:00 2001 From: matt Date: Mon, 6 Jul 2026 16:54:13 +0100 Subject: [PATCH 05/10] added a warning for write --- python/mgclient/routing.py | 17 +++++++++++++++-- test/test_routing.py | 17 +++++++++++++++++ 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/python/mgclient/routing.py b/python/mgclient/routing.py index 30fb4161..84741e52 100644 --- a/python/mgclient/routing.py +++ b/python/mgclient/routing.py @@ -25,12 +25,15 @@ one-shot convenience built on top of it. """ +import logging import threading import time from mgclient._mgclient import connect as _base_connect from mgclient._mgclient import Error, OperationalError +_logger = logging.getLogger(__name__) + #: Route the connection to a server that accepts writes (the cluster main). ACCESS_MODE_WRITE = "WRITE" #: Route the connection to a server that serves reads (a replica). @@ -427,8 +430,17 @@ def _run_managed(self, access_mode, work): last_exc = exc # A write that committed on the main but couldn't reach a SYNC # replica is durable -- treat it as success rather than retrying - # (a retry would duplicate the write). + # (a retry would duplicate the write). The synchronous + # replication guarantee was not met, though, so surface that as + # a warning rather than hiding it. if writing and is_committed_on_main_error(exc): + _logger.warning( + "write committed on the main but a SYNC replica was " + "unreachable; treating as success. The write is durable " + "on the main, but the synchronous-replication guarantee " + "was not met: %s", + exc, + ) return result if attempt == self._max_retries or not is_transient_error(exc): raise @@ -469,7 +481,8 @@ def execute_write(self, work): backoff between attempts). The one exception is a write that committed on the main but could not reach a SYNC replica: that write is durable, so it is treated as success and *not* retried (retrying would duplicate - it). + it). Because the synchronous-replication guarantee was not met in that + case, a warning is logged on the ``mgclient.routing`` logger. ``work`` may be called more than once, so it should be free of side effects other than the database operations themselves. diff --git a/test/test_routing.py b/test/test_routing.py index cfcd3831..503a60a8 100644 --- a/test/test_routing.py +++ b/test/test_routing.py @@ -429,6 +429,23 @@ def work(cursor): assert calls["n"] == 1 +def test_execute_write_warns_when_committed_on_main(monkeypatch, caplog): + class CommitFails(_FakeConnection): + commit_error = mgclient.DatabaseError(_COMMITTED_ON_MAIN_MESSAGE) + + router = _router_with_stubbed_connect(monkeypatch, CommitFails) + + # Behaviour is unchanged (still returns the result), but the weakened + # durability guarantee must be surfaced as a warning rather than hidden. + with caplog.at_level("WARNING", logger="mgclient.routing"): + assert router.execute_write(lambda cursor: "created") == "created" + + assert any( + "committed on the main" in record.getMessage().lower() + for record in caplog.records + ) + + def test_execute_write_raises_non_transient_immediately(monkeypatch): router = _router_with_stubbed_connect(monkeypatch, _FakeConnection) From b915e9b061f2a8c43b32499cd5dbba5fad119beb Mon Sep 17 00:00:00 2001 From: matt Date: Mon, 6 Jul 2026 17:52:47 +0100 Subject: [PATCH 06/10] some tidying up --- python/mgclient/__init__.py | 15 +++------------ python/mgclient/routing.py | 2 +- setup.py | 7 +------ src/mgclientmodule.c | 2 +- test/conftest.py | 2 +- test/test_connection.py | 6 +----- test/test_routing.py | 7 ++----- 7 files changed, 10 insertions(+), 31 deletions(-) diff --git a/python/mgclient/__init__.py b/python/mgclient/__init__.py index 5d93e742..61a4ab1a 100644 --- a/python/mgclient/__init__.py +++ b/python/mgclient/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2016-2020 Memgraph Ltd. [https://memgraph.com] +# Copyright (c) 2016-2026 Memgraph Ltd. [https://memgraph.com] # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,22 +14,13 @@ """Memgraph database adapter for Python, compliant with the DB-API 2.0 specification described by :pep:`249`. - -The compiled functionality lives in the :mod:`mgclient._mgclient` C extension -and is re-exported here unchanged, so ``import mgclient`` behaves exactly as it -always has. On top of that this package adds optional client-side routing for -Memgraph high-availability clusters (see :func:`connect` and the routing -constants below). """ -# Re-export everything the C extension provides (connect, Connection, Cursor, -# the exception hierarchy, the type objects and all module constants) so that -# the public ``mgclient`` API is unchanged. from mgclient._mgclient import * # noqa: F401,F403 # The routing-aware ``connect`` wrapper replaces the C ``connect`` imported -# above; with routing disabled (the default) it delegates straight to it, so -# existing code is unaffected. +# above; with routing disabled (the default) it delegates straight to the C +# ``connect`` function. from mgclient.routing import ( # noqa: F401,E402 ACCESS_MODE_READ, ACCESS_MODE_WRITE, diff --git a/python/mgclient/routing.py b/python/mgclient/routing.py index 84741e52..7741c27d 100644 --- a/python/mgclient/routing.py +++ b/python/mgclient/routing.py @@ -1,4 +1,4 @@ -# Copyright (c) 2016-2020 Memgraph Ltd. [https://memgraph.com] +# Copyright (c) 2016-2026 Memgraph Ltd. [https://memgraph.com] # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/setup.py b/setup.py index c7f7598e..f9862092 100644 --- a/setup.py +++ b/setup.py @@ -1,4 +1,4 @@ -# Copyright (c) 2016-2020 Memgraph Ltd. [https://memgraph.com] +# Copyright (c) 2016-2026 Memgraph Ltd. [https://memgraph.com] # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -36,11 +36,6 @@ # 1. The mgclient library which is the official Memgraph client library. # 2. The mgclient python extension module which is a wrapper around the # client library. -# -# The compiled extension is installed as ``mgclient._mgclient`` and the -# pure-Python ``mgclient`` package (see python/mgclient) re-exports it so that -# ``import mgclient`` keeps working exactly as before, while adding the -# client-side routing helpers on top. EXTENSION_NAME = "mgclient._mgclient" sources = [str(path) for path in Path("src").glob("*.c")] diff --git a/src/mgclientmodule.c b/src/mgclientmodule.c index 78241d51..4c299756 100644 --- a/src/mgclientmodule.c +++ b/src/mgclientmodule.c @@ -1,4 +1,4 @@ -// Copyright (c) 2016-2020 Memgraph Ltd. [https://memgraph.com] +// Copyright (c) 2016-2026 Memgraph Ltd. [https://memgraph.com] // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. diff --git a/test/conftest.py b/test/conftest.py index bdb08764..bc12fcd4 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -1,4 +1,4 @@ -# Copyright (c) 2016-2020 Memgraph Ltd. [https://memgraph.com] +# Copyright (c) 2016-2026 Memgraph Ltd. [https://memgraph.com] # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/test/test_connection.py b/test/test_connection.py index d4d6964b..7574669b 100644 --- a/test/test_connection.py +++ b/test/test_connection.py @@ -1,4 +1,4 @@ -# Copyright (c) 2016-2020 Memgraph Ltd. [https://memgraph.com] +# Copyright (c) 2016-2026 Memgraph Ltd. [https://memgraph.com] # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -122,10 +122,6 @@ def test_get_routing_table_closed_connection(memgraph_server): conn.get_routing_table() -# The ``ha_cluster`` fixture lives in conftest.py so it can be shared with the -# client-side routing tests in test_routing.py. - - @requires_ha_cluster def test_get_routing_table_ha(ha_cluster): host, port = ha_cluster diff --git a/test/test_routing.py b/test/test_routing.py index 503a60a8..a813780d 100644 --- a/test/test_routing.py +++ b/test/test_routing.py @@ -1,4 +1,4 @@ -# Copyright (c) 2016-2020 Memgraph Ltd. [https://memgraph.com] +# Copyright (c) 2016-2026 Memgraph Ltd. [https://memgraph.com] # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -423,8 +423,6 @@ def work(cursor): calls["n"] += 1 return "created" - # The write is durable on the main, so it counts as success and must NOT be - # retried (retrying would duplicate the write). assert router.execute_write(work) == "created" assert calls["n"] == 1 @@ -435,8 +433,7 @@ class CommitFails(_FakeConnection): router = _router_with_stubbed_connect(monkeypatch, CommitFails) - # Behaviour is unchanged (still returns the result), but the weakened - # durability guarantee must be surfaced as a warning rather than hidden. + # Warn if committed on main, but sync error (still return result successfully) with caplog.at_level("WARNING", logger="mgclient.routing"): assert router.execute_write(lambda cursor: "created") == "created" From bb7377f3de5b027cdd91422ea2b05c28e3cb3316 Mon Sep 17 00:00:00 2001 From: matt Date: Mon, 6 Jul 2026 17:59:35 +0100 Subject: [PATCH 07/10] tidy up messages --- python/mgclient/routing.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/python/mgclient/routing.py b/python/mgclient/routing.py index 7741c27d..1066399c 100644 --- a/python/mgclient/routing.py +++ b/python/mgclient/routing.py @@ -61,14 +61,11 @@ "please retry the query", # Memgraph's explicit "retry later" during failover "database shutdown", # instance being restarted mid-transaction "asked to abort", # transaction aborted by a shutting-down instance - "defunct connection", # connection dropped underneath the driver - "failed to read from", # partial/dropped Bolt read "failed to receive chunk", # connection reset mid-Bolt-stream "connection reset", # peer went away "broken pipe", "couldn't connect", # candidate instance not accepting connections yet "connection refused", - "routing information", # unable to retrieve routing info during a failover # Raised by Router itself while the cluster has no server for a role yet. "no write server", "no read server", From 6867f86ebaef97d58d91527e596147127c7babe4 Mon Sep 17 00:00:00 2001 From: matt Date: Tue, 7 Jul 2026 17:48:09 +0100 Subject: [PATCH 08/10] Update TransientError logic --- python/mgclient/routing.py | 52 +++++++++++++++++++------------------- src/connection-int.c | 4 ++- src/connection.c | 6 ++--- src/exceptions.h | 1 + src/mgclientmodule.c | 8 ++++++ test/test_connection.py | 8 ++++++ test/test_routing.py | 23 ++++++++++++----- 7 files changed, 65 insertions(+), 37 deletions(-) diff --git a/python/mgclient/routing.py b/python/mgclient/routing.py index 1066399c..fdc038bb 100644 --- a/python/mgclient/routing.py +++ b/python/mgclient/routing.py @@ -30,7 +30,7 @@ import time from mgclient._mgclient import connect as _base_connect -from mgclient._mgclient import Error, OperationalError +from mgclient._mgclient import Error, OperationalError, TransientError _logger = logging.getLogger(__name__) @@ -49,37 +49,36 @@ _DEFAULT_RETRY_BACKOFF = 1.0 _DEFAULT_RETRY_BACKOFF_CAP = 15.0 -# Substrings that mark a transient cluster condition worth retrying: a failover -# in progress, a replica catching up, or an instance dropped/killed mid-request. -# They clear once the cluster settles. Matched case-insensitively against the -# exception message. Adapted from hard-won experience running the neo4j driver -# against a Memgraph HA cluster under chaos. +# Best-effort message fallback for transient conditions the TransientError +# category does not (yet) capture. +# TODO(matt): remove top three fragments if `WriteQueryOnMainException` becomes a `TransientError` in Memgraph, as it currently provides `ClientError`. _TRANSIENT_MESSAGE_FRAGMENTS = ( - "replication exception", # SYNC replica lagging/unreachable - "forbidden on the", # write forbidden on a replica or the main mid-election - "setting up a new main", # coordinator electing a new main; no writer yet - "please retry the query", # Memgraph's explicit "retry later" during failover - "database shutdown", # instance being restarted mid-transaction - "asked to abort", # transaction aborted by a shutting-down instance - "failed to receive chunk", # connection reset mid-Bolt-stream - "connection reset", # peer went away - "broken pipe", - "couldn't connect", # candidate instance not accepting connections yet + "forbidden on the", # write forbidden on the main mid-failover + "setting up a new main", # older Memgraph wording of the same condition + "please retry the query", # Memgraph's explicit "retry later" during failover + # Low-level transport drops: these are not Bolt FAILUREs, so they never + # carry an error code/category at all. + "failed to receive chunk", # connection reset mid-Bolt-stream + "couldn't connect", # instance not accepting connections yet "connection refused", - # Raised by Router itself while the cluster has no server for a role yet. - "no write server", - "no read server", - "could not connect to any", - "could not fetch routing table", ) + def is_transient_error(exc): """True for a transient HA condition worth retrying after a short backoff. These arise during a failover, while a replica catches up, or when an instance is dropped mid-request; they clear once the cluster reconverges. + + The server's error category is authoritative: an :exc:`mgclient.TransientError` + (Memgraph's ``TransientError`` Bolt code) is always transient. The message + fragments below are a best-effort fallback for conditions the category does + not (yet) capture -- notably Memgraph errors that are misclassified as + ``ClientError`` and low-level transport drops that never carry a Bolt code. """ + if isinstance(exc, TransientError): + return True message = str(getattr(exc, "message", "") or exc).lower() return any(fragment in message for fragment in _TRANSIENT_MESSAGE_FRAGMENTS) @@ -284,7 +283,7 @@ def _refresh_locked(self): self._refresh_count += 1 return - raise OperationalError( + raise TransientError( "could not fetch routing table from any coordinator: " + "; ".join(errors) ) @@ -357,8 +356,9 @@ def connect(self, access_mode=ACCESS_MODE_WRITE): """Open a connection to a data instance serving ``access_mode``. Returns an ordinary :class:`Connection`. Raises - :exc:`OperationalError` if no server for the requested access mode can - be reached, even after refreshing the routing table. + :exc:`TransientError` if no server for the requested access mode can + be reached, even after refreshing the routing table (this is a + transient cluster condition, e.g. a failover in progress). """ access_mode = _normalize_access_mode(access_mode) @@ -381,7 +381,7 @@ def connect(self, access_mode=ACCESS_MODE_WRITE): # discard the cached table and retry with a fresh one. self.refresh() - raise OperationalError( + raise TransientError( f"could not connect to any {access_mode} server: {last_error}" ) @@ -569,7 +569,7 @@ def connect( * A normal connection is then opened to that server and returned. If a selected server cannot be reached, the remaining candidates are tried - in turn; an :exc:`OperationalError` is raised only if none succeed. + in turn; a :exc:`TransientError` is raised only if none succeed. This is a one-shot convenience: it performs a fresh ``ROUTE`` lookup on every call. For a long-lived, caching, load-balancing router, use diff --git a/src/connection-int.c b/src/connection-int.c index 6ff9b720..aa68c70b 100644 --- a/src/connection-int.c +++ b/src/connection-int.c @@ -36,7 +36,9 @@ void connection_handle_error(ConnectionObject *conn, int error) { error == MG_ERROR_CLIENT_ERROR) { conn->status = CONN_STATUS_READY; } - PyErr_SetString(DatabaseError, mg_session_error(conn->session)); + PyObject *exc = + (error == MG_ERROR_TRANSIENT_ERROR) ? TransientError : DatabaseError; + PyErr_SetString(exc, mg_session_error(conn->session)); } int connection_run_without_results(ConnectionObject *conn, const char *query) { diff --git a/src/connection.c b/src/connection.c index 7a961987..b3026d97 100644 --- a/src/connection.c +++ b/src/connection.c @@ -113,9 +113,9 @@ static int connection_init(ConnectionObject *conn, PyObject *args, int status = mg_connect(params, &session); mg_session_params_destroy(params); if (status != 0) { - // TODO(mtomic): maybe convert MG_ERROR_* codes to different kinds of - // Python exceptions - PyErr_SetString(OperationalError, mg_session_error(session)); + PyObject *exc = + (status == MG_ERROR_TRANSIENT_ERROR) ? TransientError : OperationalError; + PyErr_SetString(exc, mg_session_error(session)); mg_session_destroy(session); return -1; } diff --git a/src/exceptions.h b/src/exceptions.h index 1c8a5f81..22fa71b2 100644 --- a/src/exceptions.h +++ b/src/exceptions.h @@ -21,6 +21,7 @@ extern PyObject *InterfaceError; extern PyObject *DatabaseError; extern PyObject *DataError; extern PyObject *OperationalError; +extern PyObject *TransientError; extern PyObject *IntegrityError; extern PyObject *InternalError; extern PyObject *ProgrammingError; diff --git a/src/mgclientmodule.c b/src/mgclientmodule.c index 4c299756..cf8d0dce 100644 --- a/src/mgclientmodule.c +++ b/src/mgclientmodule.c @@ -34,6 +34,7 @@ PyObject *InterfaceError; PyObject *DatabaseError; PyObject *DataError; PyObject *OperationalError; +PyObject *TransientError; PyObject *IntegrityError; PyObject *InternalError; PyObject *ProgrammingError; @@ -56,6 +57,11 @@ PyDoc_STRVAR( "Exception raised for errors related to the database's operation, not " "necessarily under the control of the programmer (e.g. unexpected " "disconnect, failed allocation)."); +PyDoc_STRVAR( + TransientError_doc, + "Exception raised for transient errors that may succeed if the operation " + "is retried (e.g. during a high-availability failover). A subclass of " + ":exc:`OperationalError`."); PyDoc_STRVAR( IntegrityError_doc, "Exception raised when the relational integrity of the database is " @@ -86,6 +92,8 @@ int add_module_exceptions(PyObject *module) { {"mgclient.DataError", &DataError, &DatabaseError, DataError_doc}, {"mgclient.OperationalError", &OperationalError, &DatabaseError, OperationalError_doc}, + {"mgclient.TransientError", &TransientError, &OperationalError, + TransientError_doc}, {"mgclient.IntegrityError", &IntegrityError, &DatabaseError, IntegrityError_doc}, {"mgclient.InternalError", &InternalError, &DatabaseError, diff --git a/test/test_connection.py b/test/test_connection.py index 7574669b..e67fa780 100644 --- a/test/test_connection.py +++ b/test/test_connection.py @@ -88,6 +88,14 @@ def test_connect_args_validation(): ) +def test_transient_error_hierarchy(): + # TransientError is a distinct type, but a subclass of OperationalError so `except DatabaseError` should still catch it. + assert issubclass(mgclient.TransientError, mgclient.OperationalError) + assert issubclass(mgclient.TransientError, mgclient.DatabaseError) + assert issubclass(mgclient.TransientError, mgclient.Error) + assert mgclient.TransientError is not mgclient.OperationalError + + @requires_ssl_disabled def test_get_routing_table_args_validation(memgraph_server): host, port, sslmode, _ = memgraph_server diff --git a/test/test_routing.py b/test/test_routing.py index a813780d..259d6d2c 100644 --- a/test/test_routing.py +++ b/test/test_routing.py @@ -148,7 +148,9 @@ def test_connect_routing_all_candidates_unreachable(ha_cluster): def resolver(address): return ["127.0.0.1:1"] - with pytest.raises(mgclient.OperationalError): + # Being unable to reach any server is a transient cluster condition, so the + # Router raises TransientError (a subclass of OperationalError). + with pytest.raises(mgclient.TransientError): mgclient.connect( host=host, port=port, routing=True, access_mode="WRITE", resolver=resolver ) @@ -299,7 +301,7 @@ def resolver(address): router = Router(host=host, port=port, resolver=resolver) - with pytest.raises(mgclient.OperationalError): + with pytest.raises(mgclient.TransientError): router.connect(access_mode="WRITE") assert router._refresh_count >= 2 @@ -321,12 +323,13 @@ def test_router_routing_table_property(ha_cluster, ha_resolver): # Managed retry / transaction functions (no cluster needed for these). # --------------------------------------------------------------------------- -# Real failover messages observed against a chaos cluster. +# Real failover messages that the TransientError category does NOT cover, so +# they must still be recognised by message: Memgraph errors misclassified as +# ClientError, and low-level transport drops (no Bolt code at all). _FAILOVER_MESSAGES = [ "Write queries currently forbidden on the main instance. The cluster is in " "the process of setting up a new main instance, please retry the query " "later on.", - "could not connect to any WRITE server: no WRITE server in the routing table", "memgraph-data-0:7687: couldn't connect to host: Connection refused", "failed to receive chunk size", ] @@ -348,6 +351,12 @@ def test_is_transient_error_false_for_query_errors(): assert not is_transient_error(mgclient.DatabaseError("Syntax error near 'FOO'")) +def test_is_transient_error_recognizes_transient_error_type(): + assert is_transient_error( + mgclient.TransientError("a server transient error with no matching fragment") + ) + + def test_committed_on_main_error_needs_both_markers(): assert is_committed_on_main_error(mgclient.DatabaseError(_COMMITTED_ON_MAIN_MESSAGE)) # A replication error without the durability note is not committed-on-main. @@ -479,9 +488,9 @@ def test_execute_read_retries_transient_then_succeeds(monkeypatch): def work(cursor): calls["n"] += 1 if calls["n"] < 2: - raise mgclient.OperationalError( - "could not connect to any READ server: no READ server in the routing table" - ) + # A transient failover condition is now surfaced as TransientError + # (recognised by type, not message). + raise mgclient.TransientError("no READ server in the routing table") return 7 assert router.execute_read(work) == 7 From 44e404f15c3d8b5090040fbafaba9605daf2cb2b Mon Sep 17 00:00:00 2001 From: matt Date: Tue, 7 Jul 2026 18:00:31 +0100 Subject: [PATCH 09/10] make transport errors defined as TransientErrors --- python/mgclient/routing.py | 25 ++++++++----------------- src/connection-int.c | 15 ++++++++++++++- src/connection.c | 2 +- src/connection.h | 2 ++ test/test_connection.py | 9 +++++++++ test/test_routing.py | 8 +++----- 6 files changed, 37 insertions(+), 24 deletions(-) diff --git a/python/mgclient/routing.py b/python/mgclient/routing.py index fdc038bb..9cf21a2b 100644 --- a/python/mgclient/routing.py +++ b/python/mgclient/routing.py @@ -41,26 +41,17 @@ _ACCESS_MODES = (ACCESS_MODE_WRITE, ACCESS_MODE_READ) -# Default retry budget for managed transactions (see Router.execute_read / -# execute_write). Retries are generous because a cluster stays available as a -# whole while individual instances flap during a failover, but the backoff is -# capped so total time is bounded. _DEFAULT_MAX_RETRIES = 8 _DEFAULT_RETRY_BACKOFF = 1.0 _DEFAULT_RETRY_BACKOFF_CAP = 15.0 -# Best-effort message fallback for transient conditions the TransientError -# category does not (yet) capture. -# TODO(matt): remove top three fragments if `WriteQueryOnMainException` becomes a `TransientError` in Memgraph, as it currently provides `ClientError`. + # Best-effort message fallback for transient conditions the TransientError + # category does not (yet) capture. + # TODO(matt): remove these if `WriteQueryOnMainException` becomes a `TransientError` in Memgraph, as it currently provides `ClientError`. _TRANSIENT_MESSAGE_FRAGMENTS = ( "forbidden on the", # write forbidden on the main mid-failover "setting up a new main", # older Memgraph wording of the same condition "please retry the query", # Memgraph's explicit "retry later" during failover - # Low-level transport drops: these are not Bolt FAILUREs, so they never - # carry an error code/category at all. - "failed to receive chunk", # connection reset mid-Bolt-stream - "couldn't connect", # instance not accepting connections yet - "connection refused", ) @@ -71,11 +62,11 @@ def is_transient_error(exc): These arise during a failover, while a replica catches up, or when an instance is dropped mid-request; they clear once the cluster reconverges. - The server's error category is authoritative: an :exc:`mgclient.TransientError` - (Memgraph's ``TransientError`` Bolt code) is always transient. The message - fragments below are a best-effort fallback for conditions the category does - not (yet) capture -- notably Memgraph errors that are misclassified as - ``ClientError`` and low-level transport drops that never carry a Bolt code. + The exception type is authoritative: an :exc:`mgclient.TransientError` + (Memgraph's ``TransientError`` Bolt code, or a low-level transport/connection + failure) is always transient. The message fragments are only a fallback for + :class:`WriteQueryOnMainException`, which Memgraph currently misclassifies as + ``ClientError`` -- see :data:`_TRANSIENT_MESSAGE_FRAGMENTS`. """ if isinstance(exc, TransientError): return True diff --git a/src/connection-int.c b/src/connection-int.c index aa68c70b..63b04cd1 100644 --- a/src/connection-int.c +++ b/src/connection-int.c @@ -28,6 +28,19 @@ int connection_raise_if_bad_status(const ConnectionObject *conn) { return 0; } +int connection_error_is_transient(int error) { + switch (error) { + case MG_ERROR_TRANSIENT_ERROR: + case MG_ERROR_SEND_FAILED: + case MG_ERROR_RECV_FAILED: + case MG_ERROR_NETWORK_FAILURE: + case MG_ERROR_SOCKET: + return 1; + default: + return 0; + } +} + void connection_handle_error(ConnectionObject *conn, int error) { if (mg_session_status(conn->session) == MG_SESSION_BAD) { conn->status = CONN_STATUS_BAD; @@ -37,7 +50,7 @@ void connection_handle_error(ConnectionObject *conn, int error) { conn->status = CONN_STATUS_READY; } PyObject *exc = - (error == MG_ERROR_TRANSIENT_ERROR) ? TransientError : DatabaseError; + connection_error_is_transient(error) ? TransientError : DatabaseError; PyErr_SetString(exc, mg_session_error(conn->session)); } diff --git a/src/connection.c b/src/connection.c index b3026d97..daa941c5 100644 --- a/src/connection.c +++ b/src/connection.c @@ -114,7 +114,7 @@ static int connection_init(ConnectionObject *conn, PyObject *args, mg_session_params_destroy(params); if (status != 0) { PyObject *exc = - (status == MG_ERROR_TRANSIENT_ERROR) ? TransientError : OperationalError; + connection_error_is_transient(status) ? TransientError : OperationalError; PyErr_SetString(exc, mg_session_error(session)); mg_session_destroy(session); return -1; diff --git a/src/connection.h b/src/connection.h index 8cb78b67..30c7f21e 100644 --- a/src/connection.h +++ b/src/connection.h @@ -42,6 +42,8 @@ extern PyTypeObject ConnectionType; int connection_raise_if_bad_status(const ConnectionObject *conn); +int connection_error_is_transient(int error); + void connection_handle_error(ConnectionObject *conn, int error); int connection_run_without_results(ConnectionObject *conn, const char *query); diff --git a/test/test_connection.py b/test/test_connection.py index e67fa780..b73f9125 100644 --- a/test/test_connection.py +++ b/test/test_connection.py @@ -88,6 +88,15 @@ def test_connect_args_validation(): ) +def test_connection_refused_is_transient(): + # A connection that can't be established is a low-level transport failure + # (no Bolt code); it is surfaced as TransientError (a subclass of + # OperationalError), since it is retryable in an HA cluster. Port 1 has + # nothing listening, so the connect is refused. + with pytest.raises(mgclient.TransientError): + mgclient.connect(host="127.0.0.1", port=1) + + def test_transient_error_hierarchy(): # TransientError is a distinct type, but a subclass of OperationalError so `except DatabaseError` should still catch it. assert issubclass(mgclient.TransientError, mgclient.OperationalError) diff --git a/test/test_routing.py b/test/test_routing.py index 259d6d2c..13aa9c08 100644 --- a/test/test_routing.py +++ b/test/test_routing.py @@ -323,15 +323,13 @@ def test_router_routing_table_property(ha_cluster, ha_resolver): # Managed retry / transaction functions (no cluster needed for these). # --------------------------------------------------------------------------- -# Real failover messages that the TransientError category does NOT cover, so -# they must still be recognised by message: Memgraph errors misclassified as -# ClientError, and low-level transport drops (no Bolt code at all). +# The only transient condition still recognised by message rather than type: +# WriteQueryOnMainException, which Memgraph misclassifies as ClientError. +# (Transport drops are now surfaced as TransientError and covered by type.) _FAILOVER_MESSAGES = [ "Write queries currently forbidden on the main instance. The cluster is in " "the process of setting up a new main instance, please retry the query " "later on.", - "memgraph-data-0:7687: couldn't connect to host: Connection refused", - "failed to receive chunk size", ] _COMMITTED_ON_MAIN_MESSAGE = ( From 1193b7a03742b550e945106aebf2025ca481bbe2 Mon Sep 17 00:00:00 2001 From: matt Date: Wed, 8 Jul 2026 12:13:42 +0100 Subject: [PATCH 10/10] remove message parsing --- python/mgclient/routing.py | 26 +++++++------------------- test/test_routing.py | 35 ++++++++++------------------------- 2 files changed, 17 insertions(+), 44 deletions(-) diff --git a/python/mgclient/routing.py b/python/mgclient/routing.py index 9cf21a2b..05655cac 100644 --- a/python/mgclient/routing.py +++ b/python/mgclient/routing.py @@ -45,16 +45,6 @@ _DEFAULT_RETRY_BACKOFF = 1.0 _DEFAULT_RETRY_BACKOFF_CAP = 15.0 - # Best-effort message fallback for transient conditions the TransientError - # category does not (yet) capture. - # TODO(matt): remove these if `WriteQueryOnMainException` becomes a `TransientError` in Memgraph, as it currently provides `ClientError`. -_TRANSIENT_MESSAGE_FRAGMENTS = ( - "forbidden on the", # write forbidden on the main mid-failover - "setting up a new main", # older Memgraph wording of the same condition - "please retry the query", # Memgraph's explicit "retry later" during failover -) - - def is_transient_error(exc): """True for a transient HA condition worth retrying after a short backoff. @@ -62,16 +52,14 @@ def is_transient_error(exc): These arise during a failover, while a replica catches up, or when an instance is dropped mid-request; they clear once the cluster reconverges. - The exception type is authoritative: an :exc:`mgclient.TransientError` - (Memgraph's ``TransientError`` Bolt code, or a low-level transport/connection - failure) is always transient. The message fragments are only a fallback for - :class:`WriteQueryOnMainException`, which Memgraph currently misclassifies as - ``ClientError`` -- see :data:`_TRANSIENT_MESSAGE_FRAGMENTS`. + Classification is purely by type: the driver surfaces every transient + condition as :exc:`mgclient.TransientError` (Memgraph's ``TransientError`` + Bolt code, or a low-level transport/connection failure). Errors Memgraph + reports as ``ClientError`` -- including a write briefly forbidden while a new + main is being elected -- are treated as non-transient, like any other + client error. """ - if isinstance(exc, TransientError): - return True - message = str(getattr(exc, "message", "") or exc).lower() - return any(fragment in message for fragment in _TRANSIENT_MESSAGE_FRAGMENTS) + return isinstance(exc, TransientError) def is_committed_on_main_error(exc): diff --git a/test/test_routing.py b/test/test_routing.py index 13aa9c08..fcd27107 100644 --- a/test/test_routing.py +++ b/test/test_routing.py @@ -323,15 +323,6 @@ def test_router_routing_table_property(ha_cluster, ha_resolver): # Managed retry / transaction functions (no cluster needed for these). # --------------------------------------------------------------------------- -# The only transient condition still recognised by message rather than type: -# WriteQueryOnMainException, which Memgraph misclassifies as ClientError. -# (Transport drops are now surfaced as TransientError and covered by type.) -_FAILOVER_MESSAGES = [ - "Write queries currently forbidden on the main instance. The cluster is in " - "the process of setting up a new main instance, please retry the query " - "later on.", -] - _COMMITTED_ON_MAIN_MESSAGE = ( "Replication Exception: Failed to replicate to SYNC replica 'instance_1': " "replica is not reachable or not in sync with the main. Replica will be " @@ -339,22 +330,18 @@ def test_router_routing_table_property(ha_cluster, ha_resolver): "instance and other alive replicas." ) - -@pytest.mark.parametrize("message", _FAILOVER_MESSAGES) -def test_is_transient_error_matches_failover_conditions(message): - assert is_transient_error(mgclient.DatabaseError(message)) - - -def test_is_transient_error_false_for_query_errors(): - assert not is_transient_error(mgclient.DatabaseError("Syntax error near 'FOO'")) - - def test_is_transient_error_recognizes_transient_error_type(): assert is_transient_error( - mgclient.TransientError("a server transient error with no matching fragment") + mgclient.TransientError("a server transient error, any message") ) +def test_is_transient_error_false_for_non_transient_types(): + # Matches TransientError only -- not its parent types + assert not is_transient_error(mgclient.DatabaseError("Syntax error near 'FOO'")) + assert not is_transient_error(mgclient.OperationalError("unexpected disconnect")) + + def test_committed_on_main_error_needs_both_markers(): assert is_committed_on_main_error(mgclient.DatabaseError(_COMMITTED_ON_MAIN_MESSAGE)) # A replication error without the durability note is not committed-on-main. @@ -409,9 +396,7 @@ def test_execute_write_retries_transient_then_succeeds(monkeypatch): def work(cursor): calls["n"] += 1 if calls["n"] < 3: - raise mgclient.DatabaseError( - "setting up a new main, please retry the query later on" - ) + raise mgclient.TransientError("transient failover condition") return "done" assert router.execute_write(work) == "done" @@ -471,9 +456,9 @@ def test_execute_write_gives_up_after_max_retries(monkeypatch): def work(cursor): calls["n"] += 1 - raise mgclient.DatabaseError("please retry the query later on") + raise mgclient.TransientError("persistent transient condition") - with pytest.raises(mgclient.DatabaseError): + with pytest.raises(mgclient.TransientError): router.execute_write(work) assert calls["n"] == router._max_retries