diff --git a/.github/workflows/reusable_buildtest.yml b/.github/workflows/reusable_buildtest.yml index 9a00e53f..9ba85d01 100644 --- a/.github/workflows/reusable_buildtest.yml +++ b/.github/workflows/reusable_buildtest.yml @@ -70,49 +70,20 @@ jobs: run: | # High availability is a Memgraph Enterprise feature. The license # secrets are unavailable on PRs from forks, so skip the cluster (and - # the HA test) with a warning rather than failing the whole build. + # the HA tests) with a warning rather than failing the whole build. if [ -z "$MEMGRAPH_ENTERPRISE_LICENSE" ] || [ -z "$MEMGRAPH_ORGANIZATION_NAME" ]; then - echo "::warning::Memgraph enterprise license secrets are not set; skipping the HA cluster and the get_routing_table HA test." + echo "::warning::Memgraph enterprise license secrets are not set; skipping the HA cluster and the routing tests." exit 0 fi - # Two data instances... - for name in mg-data1 mg-data2; do - docker run -d --rm --name "$name" --network pymgclient-network \ - -e MEMGRAPH_ENTERPRISE_LICENSE="$MEMGRAPH_ENTERPRISE_LICENSE" \ - -e MEMGRAPH_ORGANIZATION_NAME="$MEMGRAPH_ORGANIZATION_NAME" \ - memgraph/memgraph:latest \ - --bolt-port=7687 \ - --management-port=13011 \ - --telemetry-enabled=false - done - - # ...and three coordinators for a Raft quorum. mg-coord1 is the - # bootstrap leader; the test harness adds the other two coordinators, - # registers the instances and elects a main once the containers are up. - for id in 1 2 3; do - docker run -d --rm --name "mg-coord${id}" --network pymgclient-network \ - -e MEMGRAPH_ENTERPRISE_LICENSE="$MEMGRAPH_ENTERPRISE_LICENSE" \ - -e MEMGRAPH_ORGANIZATION_NAME="$MEMGRAPH_ORGANIZATION_NAME" \ - memgraph/memgraph:latest \ - --bolt-port=7687 \ - --coordinator-id="${id}" \ - --coordinator-port=12121 \ - --coordinator-hostname="mg-coord${id}" \ - --management-port=13011 \ - --telemetry-enabled=false - done - - # Fail with a clear message if the cluster didn't come up, rather than - # as an opaque DNS error when the test connects to the coordinator. - sleep 10 - for name in mg-data1 mg-data2 mg-coord1 mg-coord2 mg-coord3; do - if [ "$(docker inspect -f '{{.State.Running}}' "$name" 2>/dev/null)" != "true" ]; then - echo "::error::Memgraph HA container $name is not running" - docker logs "$name" || true - exit 1 - fi - done + # Start a 3-data + 3-coordinator cluster on the test runner's Docker + # network, using the shared script from the mgclient submodule. In + # bridge mode it addresses instances by container name (mg-*:7687), + # which the test container resolves; the script also registers the + # topology and waits for the cluster to converge. + MEMGRAPH_HA_NETWORK=pymgclient-network \ + MEMGRAPH_HA_IMAGE=memgraph/memgraph:latest \ + bash mgclient/tool/ha_cluster.sh start - name: Build Docker Image run: | @@ -163,7 +134,7 @@ jobs: # when the enterprise license secrets are unavailable, e.g. forks). if [ "$(docker inspect -f '{{.State.Running}}' mg-coord1 2>/dev/null)" = "true" ]; then docker exec -i testcontainer \ - bash -c "export MEMGRAPH_HA_COORDINATOR_HOST=mg-coord1 && cd /home/memgraph/pymgclient && python$python_version -m pytest -v test/test_connection.py -k get_routing_table_ha" + bash -c "export MEMGRAPH_HA_COORDINATOR_HOST=mg-coord1 && cd /home/memgraph/pymgclient && python$python_version -m pytest -v test/test_routing.py" fi done @@ -192,9 +163,7 @@ jobs: docker wait memgraph || echo "Memgraph container does not exist" docker stop memgraph-ssl || echo "Memgraph container with SSL does not exist" docker wait memgraph-ssl || echo "Memgraph container with SSL does not exist" - for name in mg-coord1 mg-coord2 mg-coord3 mg-data1 mg-data2; do - docker stop "$name" || echo "$name does not exist" - done + bash mgclient/tool/ha_cluster.sh stop || echo "HA cluster does not exist" docker rmi pymgclient-builder:latest || echo "Image does not exist" docker network rm pymgclient-network || echo "Network does not exist" diff --git a/docs/source/module.rst b/docs/source/module.rst index 2a7e329e..7319e753 100644 --- a/docs/source/module.rst +++ b/docs/source/module.rst @@ -11,6 +11,54 @@ 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. + +``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, 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 conditions -- an instance briefly unreachable +during a failover, a replica still catching up, or a connection dropped +mid-request -- with a routing refresh and capped exponential backoff. Because +the work may run more than once, make it idempotent (e.g. ``MERGE`` rather than +``CREATE``) so a retry cannot duplicate a write. + +The classification used for retries is also exposed for building your own retry +loops: + +.. autofunction:: mgclient.is_transient_error + +For lower-level access to the routing table itself, see +:meth:`Connection.get_routing_table`. + ################ Module constants ################ @@ -54,6 +102,8 @@ through these exceptions or subclasses thereof: .. autoexception:: mgclient.OperationalError +.. autoexception:: mgclient.TransientError + .. autoexception:: mgclient.IntegrityError .. autoexception:: mgclient.InternalError @@ -64,8 +114,12 @@ through these exceptions or subclasses thereof: .. NOTE:: - In the current state, :exc:`OperationalError` is raised for all errors - obtained from the database. This will probably be improved in the future. + Most database errors are surfaced as :exc:`DatabaseError` (with + connection-related failures raised as :exc:`OperationalError`). Retryable + conditions such as a high-availability failover are raised as + :exc:`TransientError` (a subclass of :exc:`OperationalError`); see + :func:`is_transient_error`. + ################## Graph type objects diff --git a/docs/source/usage.rst b/docs/source/usage.rst index 1e4fd596..dd4984f2 100644 --- a/docs/source/usage.rst +++ b/docs/source/usage.rst @@ -124,3 +124,64 @@ 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("MERGE (: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 transient cluster conditions -- an instance briefly +unreachable during a failover, a replica still catching up, or a connection +dropped mid-request -- run your queries through the *managed* transaction +helpers, which retry them 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. diff --git a/mgclient b/mgclient index 6c6ba960..449ce85c 160000 --- a/mgclient +++ b/mgclient @@ -1 +1 @@ -Subproject commit 6c6ba9604060d8cecc6ddb472132fcf56007d6a1 +Subproject commit 449ce85c7a9bb81caec8590222f0272c631fb43b 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..6f3e6509 --- /dev/null +++ b/python/mgclient/__init__.py @@ -0,0 +1,29 @@ +# 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. +# 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`. +""" + +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. +from mgclient.routing import ( # noqa: F401,E402 + ACCESS_MODE_READ, + ACCESS_MODE_WRITE, + Router, + connect, + is_transient_error, +) diff --git a/python/mgclient/routing.py b/python/mgclient/routing.py new file mode 100644 index 00000000..70e552de --- /dev/null +++ b/python/mgclient/routing.py @@ -0,0 +1,222 @@ +# 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. +# 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 is a thin, Pythonic facade over libmgclient's routing engine, exposed by +the :mod:`mgclient._mgclient` C extension as ``_Router``. The engine -- routing +table fetch and caching (with TTL), read load-balancing, coordinator failover +and the managed-transaction retry loop -- lives in libmgclient and is shared +with every other driver built on it; this module only adapts it to an ergonomic +Python API. +""" + +from mgclient._mgclient import _Router +from mgclient._mgclient import connect as _base_connect +from mgclient._mgclient import TransientError + +#: 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) + +_DEFAULT_MAX_RETRIES = 8 +_DEFAULT_RETRY_BACKOFF = 1.0 +_DEFAULT_RETRY_BACKOFF_CAP = 15.0 + + +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. + + 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`` are treated as non-transient, like any other + client error. + """ + return isinstance(exc, TransientError) + + +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 + + +class Router: + """Client-side routing engine for a Memgraph high-availability cluster. + + A :class:`Router` is created against a seed coordinator and is meant to be + long-lived and reused. It fetches the cluster routing table, caches it until + its TTL expires, and hands out ordinary :class:`Connection` objects bound to + the appropriate data instance. + + All connection parameters other than ``host``/``address``/``port`` (for + example ``username``, ``password`` and the SSL options) are 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` forwarded to the coordinator's ``ROUTE`` + request. + + * :obj:`max_retries` / :obj:`retry_backoff` / :obj:`retry_backoff_cap` + + The managed-transaction retry budget (see :meth:`execute_read` and + :meth:`execute_write`). Backoff is capped exponential: + ``retry_backoff``, ``2 * retry_backoff``, ... up to + ``retry_backoff_cap`` seconds. + """ + + def __init__( + self, + *, + host=None, + address=None, + 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, + ): + # Forward only the parameters that were actually given; the C _Router + # applies its own defaults for anything omitted. + params = { + key: value + for key, value in dict( + connect_kwargs, + host=host, + address=address, + port=port, + resolver=resolver, + routing_context=routing_context, + ).items() + if value is not None + } + self._router = _Router( + max_retries=max_retries, + retry_backoff=retry_backoff, + retry_backoff_cap=retry_backoff_cap, + **params, + ) + + 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:`mgclient.TransientError` if no server for the requested access + mode can be reached even after refreshing the routing table (a + transient cluster condition, e.g. a failover in progress). + """ + if _normalize_access_mode(access_mode) == ACCESS_MODE_WRITE: + return self._router.connect_write() + return self._router.connect_read() + + 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; that 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._router.execute_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. A replication failure at + commit is surfaced as an error like any other -- including a SYNC + "committed on the main" failure, which is *not* treated as success: the + write is durable only on that main, so if the main is then lost before + an unreachable replica catches up the write is gone. As with any + retried write, make ``work`` idempotent (e.g. ``MERGE``) so a re-run + after such an error cannot duplicate it. + """ + return self._router.execute_write(work) + + def refresh(self): + """Force an immediate refresh of the cached routing table.""" + self._router.refresh() + + @property + def routing_table(self): + """A snapshot of the (refreshed if absent) routing table as a dict with + ``"ttl"``, ``"write"``, ``"read"`` and ``"route"`` entries.""" + return self._router.routing_table() + + +def connect( + *, + routing=False, + access_mode=ACCESS_MODE_WRITE, + resolver=None, + routing_context=None, + **kwargs, +): + """Open a connection to Memgraph. + + Takes only keyword arguments, exactly like the C extension's + :func:`connect`. + + With ``routing=False`` (the default) this is exactly the C extension's + :func:`connect`: it opens a direct connection to the given host/port. + + With ``routing=True`` it treats the target as a coordinator of a + high-availability cluster, fetches the routing table, and returns a + connection to a data instance serving ``access_mode`` (``"WRITE"`` -> the + main, ``"READ"`` -> a replica). A ``resolver`` may be supplied (see + :class:`Router`) for environments where advertised addresses are not + directly reachable. This is a one-shot convenience; construct a + :class:`Router` directly to reuse the cached routing table across + connections. + """ + if not routing: + return _base_connect(**kwargs) + + access_mode = _normalize_access_mode(access_mode) + router = Router( + resolver=resolver, routing_context=routing_context, **kwargs + ) + return router.connect(access_mode=access_mode) diff --git a/setup.py b/setup.py index 492b569b..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,7 +36,7 @@ # 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" +EXTENSION_NAME = "mgclient._mgclient" sources = [str(path) for path in Path("src").glob("*.c")] diff --git a/src/connection-int.c b/src/connection-int.c index 6ff9b720..388ed4d0 100644 --- a/src/connection-int.c +++ b/src/connection-int.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. @@ -36,7 +36,12 @@ 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)); + // A transient failure (a server-signalled TransientError, or a low-level + // transport failure worth retrying) surfaces as TransientError so callers can + // retry it; mgclient owns the classification (see mg_error_is_transient). + PyObject *exc = + mg_error_is_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..2b61367d 100644 --- a/src/connection.c +++ b/src/connection.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. @@ -21,10 +21,27 @@ #include "glue.h" static void connection_dealloc(ConnectionObject *conn) { - mg_session_destroy(conn->session); + if (conn->owns_session) { + mg_session_destroy(conn->session); + } Py_TYPE(conn)->tp_free(conn); } +PyObject *connection_wrap_session(mg_session *session, int owns_session, + int autocommit) { + ConnectionObject *conn = + (ConnectionObject *)ConnectionType.tp_alloc(&ConnectionType, 0); + if (!conn) { + return NULL; + } + conn->session = session; + conn->status = CONN_STATUS_READY; + conn->autocommit = autocommit ? 1 : 0; + conn->lazy = 0; + conn->owns_session = owns_session ? 1 : 0; + return (PyObject *)conn; +} + static int execute_trust_callback(const char *hostname, const char *ip_address, const char *key_type, const char *fingerprint, PyObject *pycallback) { @@ -113,9 +130,12 @@ 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)); + // A connection that failed for a transient reason (e.g. an instance was + // briefly unreachable during a failover) surfaces as TransientError so + // callers can retry it; mgclient owns the classification. + PyObject *exc = + mg_error_is_transient(status) ? TransientError : OperationalError; + PyErr_SetString(exc, mg_session_error(session)); mg_session_destroy(session); return -1; } @@ -125,6 +145,7 @@ static int connection_init(ConnectionObject *conn, PyObject *args, conn->status = CONN_STATUS_READY; conn->lazy = 0; conn->autocommit = 0; + conn->owns_session = 1; if (lazy) { conn->lazy = 1; @@ -178,8 +199,11 @@ static PyObject *connection_close(ConnectionObject *conn, PyObject *args) { } // No need to rollback, closing the connection will automatically - // rollback any open transactions. - mg_session_destroy(conn->session); + // rollback any open transactions. A borrowed session is left intact for its + // owner (the router); we only detach from it here. + if (conn->owns_session) { + mg_session_destroy(conn->session); + } conn->session = NULL; conn->status = CONN_STATUS_CLOSED; diff --git a/src/connection.h b/src/connection.h index 8cb78b67..07e203e4 100644 --- a/src/connection.h +++ b/src/connection.h @@ -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. @@ -35,11 +35,21 @@ typedef struct ConnectionObject { int status; int autocommit; int lazy; + // Whether closing/deallocating this connection destroys `session`. A routed + // managed transaction hands its work callback a *borrowed* connection over a + // session owned by the router, which must outlive the wrapper. + int owns_session; } ConnectionObject; // clang-format on extern PyTypeObject ConnectionType; +// Wraps an already-established `session` in a new Connection object. When +// `owns_session` is false the session is left untouched on close/dealloc (the +// caller keeps ownership). `autocommit` seeds the connection's autocommit mode. +PyObject *connection_wrap_session(mg_session *session, int owns_session, + int autocommit); + int connection_raise_if_bad_status(const ConnectionObject *conn); void connection_handle_error(ConnectionObject *conn, int error); diff --git a/src/exceptions.h b/src/exceptions.h index 1c8a5f81..764b75b9 100644 --- a/src/exceptions.h +++ b/src/exceptions.h @@ -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. @@ -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 4c97106d..e9e71dbc 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. @@ -26,6 +26,7 @@ #include "connection.h" #include "cursor.h" #include "glue.h" +#include "router.h" #include "types.h" PyObject *Warning; @@ -34,6 +35,7 @@ PyObject *InterfaceError; PyObject *DatabaseError; PyObject *DataError; PyObject *OperationalError; +PyObject *TransientError; PyObject *IntegrityError; PyObject *InternalError; PyObject *ProgrammingError; @@ -56,6 +58,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 +93,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, @@ -171,6 +180,7 @@ static struct { {"Node", &NodeType}, {"Relationship", &RelationshipType}, {"Path", &PathType}, + {"_Router", &RouterType}, {NULL, NULL}}; static int add_module_types(PyObject *module) { @@ -281,13 +291,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/src/router.c b/src/router.c new file mode 100644 index 00000000..d192d40a --- /dev/null +++ b/src/router.c @@ -0,0 +1,525 @@ +// 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. +// 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. + +#include "router.h" + +#include + +#include "connection.h" +#include "exceptions.h" +#include "glue.h" + +// clang-format off +typedef struct RouterObject { + PyObject_HEAD + + mg_router *router; + // The Python address resolver (or NULL for the identity mapping). Borrowed by + // `router` as its resolver_data, so it must outlive `router`. + PyObject *resolver; + // The most recent exception raised inside a callback (resolver or work), + // stashed while control is down in libmgclient's C code and re-raised once + // the top-level call returns. Everything here runs single-threaded with the + // GIL held, so a per-router slot is safe. + PyObject *exc_type; + PyObject *exc_value; + PyObject *exc_tb; +} RouterObject; +// clang-format on + +// -- callback exception stashing -------------------------------------------- + +static void router_clear_stashed(RouterObject *self) { + Py_CLEAR(self->exc_type); + Py_CLEAR(self->exc_value); + Py_CLEAR(self->exc_tb); +} + +// Move the currently-set Python exception into the router's stash (clearing the +// pending error so libmgclient's C code runs without one set). +static void router_stash_exception(RouterObject *self) { + router_clear_stashed(self); + PyErr_Fetch(&self->exc_type, &self->exc_value, &self->exc_tb); +} + +static int router_has_stashed(const RouterObject *self) { + return self->exc_type != NULL || self->exc_value != NULL; +} + +// Raise a Python exception for a failed router operation that returned `status`. +// A stashed callback exception takes precedence (it carries the real cause and +// traceback); otherwise the router's own message is used, classified as +// transient or not. +static void router_raise(RouterObject *self, int status) { + if (router_has_stashed(self)) { + PyErr_Restore(self->exc_type, self->exc_value, self->exc_tb); + self->exc_type = self->exc_value = self->exc_tb = NULL; + return; + } + PyObject *exc = + mg_error_is_transient(status) ? TransientError : OperationalError; + PyErr_SetString(exc, mg_router_error(self->router)); +} + +// -- resolver trampoline ----------------------------------------------------- + +// Bridges libmgclient's `mg_resolver_fn` to the Python resolver callable, which +// maps an advertised "host:port" to an iterable of "host:port" targets. +static int router_resolver_trampoline(const char *advertised, + mg_resolver_result *result, void *data) { + RouterObject *self = (RouterObject *)data; + + PyObject *targets = PyObject_CallFunction(self->resolver, "s", advertised); + if (!targets) { + router_stash_exception(self); + return MG_ERROR_CLIENT_ERROR; + } + PyObject *seq = + PySequence_Fast(targets, "resolver must return an iterable of addresses"); + Py_DECREF(targets); + if (!seq) { + router_stash_exception(self); + return MG_ERROR_CLIENT_ERROR; + } + + int rc = 0; + Py_ssize_t size = PySequence_Fast_GET_SIZE(seq); + for (Py_ssize_t i = 0; i < size; ++i) { + PyObject *item = PySequence_Fast_GET_ITEM(seq, i); // borrowed + const char *target = PyUnicode_AsUTF8(item); + if (!target) { + router_stash_exception(self); + rc = MG_ERROR_CLIENT_ERROR; + break; + } + if (mg_resolver_result_add(result, target) != 0) { + PyErr_NoMemory(); + router_stash_exception(self); + rc = MG_ERROR_OOM; + break; + } + } + Py_DECREF(seq); + return rc; +} + +// -- lifecycle --------------------------------------------------------------- + +static int router_init(RouterObject *self, PyObject *args, PyObject *kwargs) { + static char *kwlist[] = {"host", + "address", + "port", + "username", + "password", + "client_name", + "sslmode", + "sslcert", + "sslkey", + "resolver", + "routing_context", + "max_retries", + "retry_backoff", + "retry_backoff_cap", + NULL}; + + const char *host = NULL; + const char *address = NULL; + int port = -1; + const char *username = NULL; + const char *password = NULL; + const char *client_name = NULL; + int sslmode_int = MG_SSLMODE_DISABLE; + const char *sslcert = NULL; + const char *sslkey = NULL; + PyObject *resolver = NULL; + PyObject *routing_context = NULL; + unsigned int max_retries = 8; + double retry_backoff = 1.0; + double retry_backoff_cap = 15.0; + + if (!PyArg_ParseTupleAndKeywords( + args, kwargs, "|$zzizzzizzOOIdd", kwlist, &host, &address, &port, + &username, &password, &client_name, &sslmode_int, &sslcert, &sslkey, + &resolver, &routing_context, &max_retries, &retry_backoff, + &retry_backoff_cap)) { + return -1; + } + + if (port < 0 || port > 65535) { + PyErr_SetString(PyExc_ValueError, "port out of range"); + return -1; + } + + enum mg_sslmode sslmode; + switch (sslmode_int) { + case MG_SSLMODE_DISABLE: + case MG_SSLMODE_REQUIRE: + sslmode = sslmode_int; + break; + default: + PyErr_SetString(PyExc_ValueError, "invalid sslmode"); + return -1; + } + + if (resolver && resolver != Py_None && !PyCallable_Check(resolver)) { + PyErr_SetString(PyExc_TypeError, "resolver argument must be callable"); + return -1; + } + if (routing_context && routing_context != Py_None && + !PyDict_Check(routing_context)) { + PyErr_SetString(PyExc_TypeError, "routing_context must be a dict"); + return -1; + } + + mg_session_params *params = mg_session_params_make(); + if (!params) { + PyErr_SetString(PyExc_RuntimeError, + "couldn't allocate session parameters object"); + return -1; + } + mg_session_params_set_host(params, host); + mg_session_params_set_port(params, (uint16_t)port); + mg_session_params_set_address(params, address); + mg_session_params_set_username(params, username); + mg_session_params_set_password(params, password); + if (client_name) { + mg_session_params_set_user_agent(params, client_name); + } + mg_session_params_set_sslmode(params, sslmode); + mg_session_params_set_sslcert(params, sslcert); + mg_session_params_set_sslkey(params, sslkey); + + mg_map *mg_routing_context = NULL; + if (routing_context && routing_context != Py_None) { + mg_routing_context = py_dict_to_mg_map(routing_context); + if (!mg_routing_context) { + mg_session_params_destroy(params); + return -1; + } + } + + mg_router_config *config = mg_router_config_make(); + if (!config) { + mg_session_params_destroy(params); + mg_map_destroy(mg_routing_context); + PyErr_SetString(PyExc_RuntimeError, "couldn't allocate router config"); + return -1; + } + mg_router_config_set_session_params(config, params); + if (mg_routing_context) { + mg_router_config_set_routing_context(config, mg_routing_context); + } + mg_router_config_set_max_retries(config, (uint32_t)max_retries); + mg_router_config_set_retry_backoff(config, retry_backoff, retry_backoff_cap); + + PyObject *stored_resolver = NULL; + if (resolver && resolver != Py_None) { + stored_resolver = resolver; + Py_INCREF(stored_resolver); + mg_router_config_set_resolver(config, router_resolver_trampoline, self); + } + + mg_router *router = mg_router_make(config); + mg_session_params_destroy(params); + mg_map_destroy(mg_routing_context); + mg_router_config_destroy(config); + + if (!router) { + Py_XDECREF(stored_resolver); + PyErr_SetString(PyExc_RuntimeError, "couldn't create router"); + return -1; + } + + // Replace any previous state (in case __init__ is called twice). + mg_router_destroy(self->router); + Py_XDECREF(self->resolver); + router_clear_stashed(self); + self->router = router; + self->resolver = stored_resolver; + return 0; +} + +static PyObject *router_new(PyTypeObject *subtype, PyObject *args, + PyObject *kwargs) { + (void)args; + (void)kwargs; + RouterObject *self = (RouterObject *)subtype->tp_alloc(subtype, 0); + if (!self) { + return NULL; + } + self->router = NULL; + self->resolver = NULL; + self->exc_type = self->exc_value = self->exc_tb = NULL; + return (PyObject *)self; +} + +static void router_dealloc(RouterObject *self) { + // Destroy the router first: it borrows `resolver` as its resolver_data. + mg_router_destroy(self->router); + Py_XDECREF(self->resolver); + router_clear_stashed(self); + Py_TYPE(self)->tp_free(self); +} + +// -- connect ----------------------------------------------------------------- + +static PyObject *router_connect_role(RouterObject *self, int write) { + router_clear_stashed(self); + mg_session *session = NULL; + int status = write ? mg_router_connect_write(self->router, &session) + : mg_router_connect_read(self->router, &session); + if (status != 0) { + router_raise(self, status); + return NULL; + } + // The caller owns the returned connection; it owns and closes the session. + PyObject *conn = connection_wrap_session(session, /*owns_session=*/1, + /*autocommit=*/0); + if (!conn) { + mg_session_destroy(session); + return NULL; + } + return conn; +} + +PyDoc_STRVAR(router_connect_read_doc, + "connect_read()\n--\n\n" + "Open an owning Connection to a server that serves reads."); + +static PyObject *router_connect_read(RouterObject *self, PyObject *args) { + (void)args; + return router_connect_role(self, /*write=*/0); +} + +PyDoc_STRVAR(router_connect_write_doc, + "connect_write()\n--\n\n" + "Open an owning Connection to the server that serves writes."); + +static PyObject *router_connect_write(RouterObject *self, PyObject *args) { + (void)args; + return router_connect_role(self, /*write=*/1); +} + +// -- managed transactions ---------------------------------------------------- + +struct work_ctx { + RouterObject *self; + PyObject *work; // borrowed + PyObject *result; // owned; the value returned by the last successful work +}; + +// Bridges libmgclient's `mg_work_fn` to the Python work callable, which +// receives a Cursor and returns the caller's result. A raised exception is +// stashed and mapped to a transient/non-transient status so the retry loop in +// libmgclient can decide whether to try again. +static int router_work_trampoline(mg_session *session, void *data) { + struct work_ctx *ctx = (struct work_ctx *)data; + RouterObject *self = ctx->self; + + // A fresh attempt: drop any exception stashed by a previous one. + router_clear_stashed(self); + + // The work runs against a borrowed connection in autocommit mode: for a + // write, libmgclient owns the BEGIN/COMMIT around this callback, so the + // Python layer must not drive the transaction itself. + PyObject *conn = connection_wrap_session(session, /*owns_session=*/0, + /*autocommit=*/1); + if (!conn) { + router_stash_exception(self); + return MG_ERROR_CLIENT_ERROR; + } + PyObject *cursor = PyObject_CallMethod(conn, "cursor", NULL); + if (!cursor) { + router_stash_exception(self); + Py_DECREF(conn); + return MG_ERROR_CLIENT_ERROR; + } + + PyObject *result = PyObject_CallFunctionObjArgs(ctx->work, cursor, NULL); + Py_DECREF(cursor); + Py_DECREF(conn); + + if (!result) { + int transient = PyErr_ExceptionMatches(TransientError); + router_stash_exception(self); + return transient ? MG_ERROR_TRANSIENT_ERROR : MG_ERROR_CLIENT_ERROR; + } + Py_XSETREF(ctx->result, result); + return 0; +} + +static PyObject *router_execute_role(RouterObject *self, PyObject *work, + int write) { + if (!PyCallable_Check(work)) { + PyErr_SetString(PyExc_TypeError, "work argument must be callable"); + return NULL; + } + router_clear_stashed(self); + + struct work_ctx ctx = {.self = self, .work = work, .result = NULL}; + int status = + write ? mg_router_execute_write(self->router, router_work_trampoline, + &ctx) + : mg_router_execute_read(self->router, router_work_trampoline, + &ctx); + + if (status != 0) { + Py_XDECREF(ctx.result); + router_raise(self, status); + return NULL; + } + // Success: discard any exception stashed by an earlier, retried attempt. + router_clear_stashed(self); + if (ctx.result == NULL) { + Py_RETURN_NONE; + } + return ctx.result; // reference transferred to the caller +} + +PyDoc_STRVAR(router_execute_read_doc, + "execute_read(work)\n--\n\n" + "Run work(cursor) as a managed read against a replica."); + +static PyObject *router_execute_read(RouterObject *self, PyObject *work) { + return router_execute_role(self, work, /*write=*/0); +} + +PyDoc_STRVAR(router_execute_write_doc, + "execute_write(work)\n--\n\n" + "Run work(cursor) as a managed write against the main."); + +static PyObject *router_execute_write(RouterObject *self, PyObject *work) { + return router_execute_role(self, work, /*write=*/1); +} + +// -- refresh / routing table ------------------------------------------------- + +PyDoc_STRVAR(router_refresh_doc, + "refresh()\n--\n\n" + "Force an immediate refresh of the cached routing table."); + +static PyObject *router_refresh(RouterObject *self, PyObject *args) { + (void)args; + router_clear_stashed(self); + int status = mg_router_refresh(self->router); + if (status != 0) { + router_raise(self, status); + return NULL; + } + Py_RETURN_NONE; +} + +static PyObject *role_addresses(const mg_routing_table *table, + enum mg_routing_role role) { + uint32_t count = mg_routing_table_address_count(table, role); + PyObject *list = PyList_New(count); + if (!list) { + return NULL; + } + for (uint32_t i = 0; i < count; ++i) { + PyObject *item = + PyUnicode_FromString(mg_routing_table_address_at(table, role, i)); + if (!item) { + Py_DECREF(list); + return NULL; + } + PyList_SET_ITEM(list, i, item); // steals reference + } + return list; +} + +PyDoc_STRVAR(router_routing_table_doc, + "routing_table()\n--\n\n" + "Return the cached routing table (refreshing it if none is " + "cached yet) as a dict with 'ttl', 'write', 'read' and 'route'."); + +static PyObject *router_routing_table(RouterObject *self, PyObject *args) { + (void)args; + router_clear_stashed(self); + + if (mg_router_routing_table(self->router) == NULL) { + int status = mg_router_refresh(self->router); + if (status != 0) { + router_raise(self, status); + return NULL; + } + } + const mg_routing_table *table = mg_router_routing_table(self->router); + if (!table) { + PyErr_SetString(TransientError, "no routing table available"); + return NULL; + } + + PyObject *write = role_addresses(table, MG_ROUTING_ROLE_WRITE); + PyObject *read = role_addresses(table, MG_ROUTING_ROLE_READ); + PyObject *route = role_addresses(table, MG_ROUTING_ROLE_ROUTE); + PyObject *ttl = + PyLong_FromLongLong((long long)mg_routing_table_ttl(table)); + PyObject *dict = NULL; + if (write && read && route && ttl) { + dict = PyDict_New(); + } + if (!dict) { + Py_XDECREF(write); + Py_XDECREF(read); + Py_XDECREF(route); + Py_XDECREF(ttl); + return NULL; + } + int ok = PyDict_SetItemString(dict, "ttl", ttl) == 0 && + PyDict_SetItemString(dict, "write", write) == 0 && + PyDict_SetItemString(dict, "read", read) == 0 && + PyDict_SetItemString(dict, "route", route) == 0; + Py_DECREF(write); + Py_DECREF(read); + Py_DECREF(route); + Py_DECREF(ttl); + if (!ok) { + Py_DECREF(dict); + return NULL; + } + return dict; +} + +static PyMethodDef router_methods[] = { + {"connect_read", (PyCFunction)router_connect_read, METH_NOARGS, + router_connect_read_doc}, + {"connect_write", (PyCFunction)router_connect_write, METH_NOARGS, + router_connect_write_doc}, + {"execute_read", (PyCFunction)router_execute_read, METH_O, + router_execute_read_doc}, + {"execute_write", (PyCFunction)router_execute_write, METH_O, + router_execute_write_doc}, + {"refresh", (PyCFunction)router_refresh, METH_NOARGS, router_refresh_doc}, + {"routing_table", (PyCFunction)router_routing_table, METH_NOARGS, + router_routing_table_doc}, + {NULL, NULL, 0, NULL}}; + +PyDoc_STRVAR(RouterType_doc, + "Low-level client-side routing engine over libmgclient's " + "mg_router. Use mgclient.routing.Router instead."); + +// clang-format off +PyTypeObject RouterType = { + PyVarObject_HEAD_INIT(NULL, 0) + .tp_name = "mgclient._mgclient._Router", + .tp_doc = RouterType_doc, + .tp_basicsize = sizeof(RouterObject), + .tp_itemsize = 0, + .tp_flags = Py_TPFLAGS_DEFAULT, + .tp_dealloc = (destructor)router_dealloc, + .tp_methods = router_methods, + .tp_init = (initproc)router_init, + .tp_new = (newfunc)router_new}; +// clang-format on diff --git a/src/router.h b/src/router.h new file mode 100644 index 00000000..a19a792a --- /dev/null +++ b/src/router.h @@ -0,0 +1,25 @@ +// 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. +// 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. + +#ifndef PYMGCLIENT_ROUTER_H +#define PYMGCLIENT_ROUTER_H + +#include + +// A thin Python wrapper over libmgclient's `mg_router` (client-side routing +// engine). Not part of the public API surface directly; the ergonomic +// `mgclient.routing.Router` facade is built on top of it. +extern PyTypeObject RouterType; + +#endif // PYMGCLIENT_ROUTER_H diff --git a/test/test_connection.py b/test/test_connection.py index 282c1f2a..628f2499 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. @@ -14,17 +14,14 @@ import mgclient import pytest +import socket import tempfile -import time from common import ( start_memgraph, Memgraph, requires_ssl_enabled, requires_ssl_disabled, - requires_ha_cluster, - MEMGRAPH_HA_COORDINATOR_HOST, - MEMGRAPH_HA_COORDINATOR_PORT, ) from OpenSSL import crypto @@ -91,6 +88,29 @@ 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. + # + # Reserve a port with no listener -- bind a socket without calling listen() + # and keep it open -- so the connect is guaranteed to be refused rather than + # relying on a well-known port happening to be closed. + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + with pytest.raises(mgclient.TransientError): + mgclient.connect(host="127.0.0.1", port=port) + + +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 @@ -125,100 +145,6 @@ 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() - - -@requires_ha_cluster -def test_get_routing_table_ha(ha_cluster): - host, port = ha_cluster - conn = mgclient.connect(host=host, port=port) - - table = conn.get_routing_table() - - assert isinstance(table["ttl"], int) - assert table["servers"] - - # The cluster has a main (WRITE), a replica (READ) and coordinators (ROUTE), - # so all three roles must be present. - roles = {server["role"] for server in table["servers"]} - assert roles == {"READ", "WRITE", "ROUTE"} - - for server in table["servers"]: - assert isinstance(server["addresses"], list) - assert server["addresses"] - - @requires_ssl_disabled def test_connect_insecure_success(memgraph_server): host, port, sslmode, _ = memgraph_server @@ -228,6 +154,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..1ada928a --- /dev/null +++ b/test/test_routing.py @@ -0,0 +1,328 @@ +# 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. +# 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, ...)`` and +:class:`mgclient.Router`). + +These exercise the Python facade over libmgclient's routing engine. The engine +itself -- routing-table parsing and caching, read round-robin, coordinator +failover and the managed-transaction retry loop -- is unit-tested in +libmgclient's own suite (``tests/routing.cpp``); the cluster-gated tests here +confirm the facade wires it up correctly end to end. They are skipped unless +``MEMGRAPH_HA_COORDINATOR_HOST`` is set. The cluster's advertised addresses must +be directly reachable from the test runner (as they are on the shared CI Docker +network). +""" + +import mgclient +import pytest + +from common import ( + MEMGRAPH_HA_COORDINATOR_HOST, + MEMGRAPH_HA_COORDINATOR_PORT, + requires_ha_cluster, +) +from mgclient.routing import ( + Router, + is_transient_error, +) + + +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] + + +@pytest.fixture(scope="module") +def ha_cluster(): + """The coordinator ``(host, port)`` of a pre-provisioned HA cluster. + + The cluster is created, registered and converged out of band -- in CI by + the "Run Memgraph HA Cluster" step (``mgclient/tool/ha_cluster.sh``), or + locally by running that script -- so this fixture only checks it is up and + advertising all roles, then hands the coordinator address to the tests. + """ + host = MEMGRAPH_HA_COORDINATOR_HOST + port = MEMGRAPH_HA_COORDINATOR_PORT + + conn = mgclient.connect(host=host, port=port) + try: + roles = { + server["role"] for server in conn.get_routing_table()["servers"] + } + if not {"READ", "WRITE", "ROUTE"} <= roles: + raise RuntimeError( + f"HA cluster not ready; routing table advertised roles: {roles}" + ) + finally: + conn.close() + + return host, port + + +# --------------------------------------------------------------------------- +# Argument handling (no cluster needed). +# --------------------------------------------------------------------------- + + +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) + + +# --------------------------------------------------------------------------- +# Routing-table primitive: Connection.get_routing_table (cluster-gated). +# --------------------------------------------------------------------------- + + +@requires_ha_cluster +def test_get_routing_table_ha(ha_cluster): + host, port = ha_cluster + conn = mgclient.connect(host=host, port=port) + + table = conn.get_routing_table() + + assert isinstance(table["ttl"], int) + assert table["servers"] + + # The cluster has a main (WRITE), a replica (READ) and coordinators (ROUTE), + # so all three roles must be present. + roles = {server["role"] for server in table["servers"]} + assert roles == {"READ", "WRITE", "ROUTE"} + + for server in table["servers"]: + assert isinstance(server["addresses"], list) + assert server["addresses"] + + +# --------------------------------------------------------------------------- +# connect(routing=True): one-shot routed connections (cluster-gated). +# --------------------------------------------------------------------------- + + +@requires_ha_cluster +def test_connect_routing_write_reaches_main(ha_cluster): + host, port = ha_cluster + conn = mgclient.connect(host=host, port=port, routing=True, access_mode="WRITE") + 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): + host, port = ha_cluster + conn = mgclient.connect(host=host, port=port, routing=True, access_mode="READ") + 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): + host, port = ha_cluster + conn = mgclient.connect(host=host, port=port, routing=True) + assert _replication_role(conn) == "main" + conn.close() + + +@requires_ha_cluster +def test_connect_routing_access_mode_is_case_insensitive(ha_cluster): + host, port = ha_cluster + conn = mgclient.connect(host=host, port=port, routing=True, access_mode="read") + 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 [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 + + # Make the first replica the router selects unreachable (remap it to a dead + # address); every other replica resolves to itself. Routing must fail over + # to a reachable replica rather than giving up. Relies on the cluster having + # more than one replica. + killed = [] + + def resolver(address): + if not killed: + killed.append(address) + return ["127.0.0.1:1"] if address == killed[0] else [address] + + conn = mgclient.connect( + host=host, port=port, routing=True, access_mode="READ", resolver=resolver + ) + assert _replication_role(conn) == "replica" + 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"] + + # 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 + ) + + +# --------------------------------------------------------------------------- +# Router: a reusable engine (cluster-gated). +# --------------------------------------------------------------------------- + + +@requires_ha_cluster +def test_router_connect_reaches_correct_instances(ha_cluster): + host, port = ha_cluster + router = Router(host=host, port=port) + + 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_refreshes_when_targets_unreachable(ha_cluster): + host, port = ha_cluster + + # The coordinator (seed) is reachable, but the data instances are not, so + # routing refreshes and retries before giving up with a transient error. + def resolver(address): + return ["127.0.0.1:1"] + + router = Router(host=host, port=port, resolver=resolver) + + with pytest.raises(mgclient.TransientError): + router.connect(access_mode="WRITE") + + +@requires_ha_cluster +def test_router_routing_table_property(ha_cluster): + host, port = ha_cluster + router = Router(host=host, port=port) + + table = router.routing_table + assert table["write"] + assert table["read"] + assert table["route"] + assert isinstance(table["ttl"], int) + + +@requires_ha_cluster +def test_router_refresh_succeeds(ha_cluster): + host, port = ha_cluster + router = Router(host=host, port=port) + + # An explicit refresh works before and after any connection is opened. + router.refresh() + router.connect(access_mode="WRITE").close() + router.refresh() + + +# --------------------------------------------------------------------------- +# Managed transactions (cluster-gated). +# --------------------------------------------------------------------------- + + +@requires_ha_cluster +def test_execute_write_and_read_roundtrip(ha_cluster): + host, port = ha_cluster + router = Router(host=host, port=port) + + 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") + ) + + +# --------------------------------------------------------------------------- +# Error-classification helper (no cluster needed). +# --------------------------------------------------------------------------- + + +def test_is_transient_error_recognizes_transient_error_type(): + assert is_transient_error( + 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"))