add TransientError and Memgraph vendor

Author: mattkjames7

gqlalchemy/__init__.py

@@ -37,7 +37,12 @@
     wait_for_docker_container,
     wait_for_port,
 )
-from gqlalchemy.exceptions import GQLAlchemyError, GQLAlchemyWarning  # noqa F401
+from gqlalchemy.exceptions import (  # noqa F401
+    GQLAlchemyError,
+    GQLAlchemyWarning,
+    GQLAlchemyDatabaseError,
+    GQLAlchemyTransientError,
+)
 
 from gqlalchemy.query_builders import (  # noqa F401
     neo4j_query_builder,

gqlalchemy/exceptions.py

@@ -15,6 +15,8 @@
 from enum import Enum
 import time
 
+import mgclient
+
 DATABASE_MISSING_IN_FIELD_ERROR_MESSAGE = """
 Can't have an index on a property without providing the database `db` object.
 Define your property as:
@@ -172,6 +174,12 @@ def __init__(self, message):
         self.message = message
 
 
+class GQLAlchemyTransientError(GQLAlchemyDatabaseError):
+    """A database error that is worth retrying."""
+
+    pass
+
+
 class GQLAlchemyOperatorTypeError(GQLAlchemyError):
     def __init__(self, clause) -> None:
         self.message = OPERATOR_TYPE_ERROR.format(clause=clause)
@@ -212,6 +220,8 @@ def database_error_handler(func):
     def inner_function(*args, **kwargs):
         try:
             return func(*args, **kwargs)
+        except mgclient.TransientError as e:
+            raise GQLAlchemyTransientError(e) from e
         except Exception as e:
             raise GQLAlchemyDatabaseError(e) from e
 

gqlalchemy/vendors/memgraph.py

@@ -122,11 +122,17 @@ def __init__(
         encrypted: bool = mg_consts.MG_ENCRYPTED,
         client_name: str = mg_consts.MG_CLIENT_NAME,
         lazy: bool = mg_consts.MG_LAZY,
+        routing: bool = False,
+        access_mode: Optional[str] = None,
+        resolver: Optional[object] = None,
     ):
         super().__init__(
             host=host, port=port, username=username, password=password, encrypted=encrypted, client_name=client_name
         )
         self._lazy = lazy
+        self._routing = routing
+        self._access_mode = access_mode
+        self._resolver = resolver
         self._on_disk_db = None
 
     @staticmethod
@@ -228,6 +234,9 @@ def new_connection(self) -> Connection:
             password=self._password,
             encrypted=self._encrypted,
             client_name=self._client_name,
+            routing=self._routing,
+            access_mode=self._access_mode,
+            resolver=self._resolver,
         )
         return MemgraphConnection(**args)
 
@@ -301,6 +310,9 @@ def _new_connection(self) -> Connection:
             password=self._password,
             encrypted=self._encrypted,
             client_name=self._client_name,
+            routing=self._routing,
+            access_mode=self._access_mode,
+            resolver=self._resolver,
         )
         return MemgraphConnection(**args)
 

tests/integration/test_routing.py

@@ -17,15 +17,15 @@
 These tests need a running HA cluster whose coordinator is advertised through
 ``MEMGRAPH_HA_COORDINATOR_HOST`` / ``MEMGRAPH_HA_COORDINATOR_PORT`` (started in
 CI by the "Run Memgraph HA Cluster" step, ``scripts/ha_cluster.sh``, or locally
-by running that script). They are skipped otherwise, and also when the
-installed pymgclient predates client-side routing.
+by running that script). They are skipped otherwise.
 """
 
 import os
 
 import mgclient
 import pytest
 
+from gqlalchemy import Memgraph
 from gqlalchemy.connection import MemgraphConnection
 
 HA_HOST = os.environ.get("MEMGRAPH_HA_COORDINATOR_HOST")
@@ -37,10 +37,6 @@
         not (HA_HOST and HA_PORT),
         reason="requires a Memgraph HA cluster (set MEMGRAPH_HA_COORDINATOR_HOST/PORT)",
     ),
-    pytest.mark.skipif(
-        not hasattr(mgclient, "ACCESS_MODE_WRITE"),
-        reason="requires a routing-capable pymgclient",
-    ),
 ]
 
 
@@ -84,3 +80,27 @@ def test_routed_write_is_readable_from_main():
     connection.execute("MERGE (n:RoutingTest {id: 1}) SET n.value = 'ok'")
     result = list(connection.execute_and_fetch("MATCH (n:RoutingTest {id: 1}) RETURN n.value AS value"))
     assert result[0]["value"] == "ok"
+
+
+# The Memgraph vendor client (routing plumbed through new_connection).
+
+
+def _routing_memgraph(access_mode=None):
+    return Memgraph(host=HA_HOST, port=int(HA_PORT), routing=True, access_mode=access_mode)
+
+
+def _vendor_role(db):
+    row = list(db.execute_and_fetch("SHOW REPLICATION ROLE"))[0]
+    return next(iter(row.values()))
+
+
+def test_vendor_write_client_targets_main():
+    assert _vendor_role(_routing_memgraph(mgclient.ACCESS_MODE_WRITE)) == "main"
+
+
+def test_vendor_read_client_targets_replica():
+    assert _vendor_role(_routing_memgraph(mgclient.ACCESS_MODE_READ)) == "replica"
+
+
+def test_vendor_default_access_mode_targets_main():
+    assert _vendor_role(_routing_memgraph()) == "main"

tests/test_exceptions.py

@@ -0,0 +1,57 @@
+# 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.
+
+import mgclient
+import pytest
+
+from gqlalchemy.exceptions import (
+    GQLAlchemyDatabaseError,
+    GQLAlchemyTransientError,
+    database_error_handler,
+)
+
+
+def test_transient_error_is_a_database_error():
+    # Backwards compatible: code catching GQLAlchemyDatabaseError still catches
+    # the transient subclass.
+    assert issubclass(GQLAlchemyTransientError, GQLAlchemyDatabaseError)
+
+
+def test_transient_error_is_mapped_to_gqlalchemy_transient_error():
+    @database_error_handler
+    def boom():
+        raise mgclient.TransientError("instance briefly unreachable during a failover")
+
+    with pytest.raises(GQLAlchemyTransientError):
+        boom()
+
+
+def test_non_transient_database_error_is_not_transient():
+    @database_error_handler
+    def boom():
+        raise mgclient.DatabaseError("syntax error")
+
+    with pytest.raises(GQLAlchemyDatabaseError) as exc_info:
+        boom()
+    assert not isinstance(exc_info.value, GQLAlchemyTransientError)
+
+
+def test_plain_exception_is_mapped_to_database_error():
+    @database_error_handler
+    def boom():
+        raise ValueError("not a database error")
+
+    with pytest.raises(GQLAlchemyDatabaseError) as exc_info:
+        boom()
+    assert not isinstance(exc_info.value, GQLAlchemyTransientError)