Add request timeout, error handling, and retry to KGClient

fairgraph/client.py

@@ -22,6 +22,8 @@
 import os
 import logging
 import re
+from functools import wraps
+from time import sleep
 from typing import Any, Dict, Iterable, List, Optional, Union, TYPE_CHECKING
 from uuid import uuid4, UUID
 
@@ -36,7 +38,7 @@
 
 from openminds.registry import lookup_type
 
-from .errors import AuthenticationError, AuthorizationError, ResourceExistsError
+from .errors import AuthenticationError, AuthorizationError, KGConnectionError, ResourceExistsError
 from .utility import handle_scope_keyword
 from .base import OPENMINDS_VERSION
 
@@ -52,14 +54,116 @@
 logger = logging.getLogger("fairgraph")
 
 
+def translate_network_errors(method):
+    """Turn a network-level `requests` exception escaping a KGClient method into a
+    clear KGConnectionError.
+    """
+
+    @wraps(method)
+    def wrapper(self, *args, **kwargs):
+        try:
+            return method(self, *args, **kwargs)
+        except KGConnectionError:
+            raise
+        except Exception as err:
+            if have_kg_core and isinstance(err, requests.exceptions.RequestException):
+                raise KGConnectionError(f"Request to the KG failed: {err}") from err
+            raise
+
+    return wrapper
+
+
+def retry_on_connection_error(method):
+    """Retry a KGClient method a few times if it raises KGConnectionError.
+
+    Only apply this to methods that are read-only or execute a query without persisting
+    anything. Retrying a request that modifies the KG risks duplicating or corrupting
+    data if the original request actually succeeded but the response was lost. Reads
+    `self._max_retries` / `self._retry_backoff`, set in KGClient.__init__.
+    """
+
+    @wraps(method)
+    def wrapper(self, *args, **kwargs):
+        for attempt in range(self._max_retries + 1):
+            try:
+                return method(self, *args, **kwargs)
+            except KGConnectionError as err:
+                if attempt < self._max_retries:
+                    wait = (2**attempt) * self._retry_backoff
+                    logger.warning(
+                        "Retrying KG request after connection error (attempt %d/%d), waiting %.1fs: %s",
+                        attempt + 1,
+                        self._max_retries,
+                        wait,
+                        err,
+                    )
+                    sleep(wait)
+                else:
+                    raise
+
+    return wrapper
+
+
 if have_kg_core:
+    import requests
+
     STAGE_MAP = {
         "released": Stage.RELEASED,
         "latest": Stage.IN_PROGRESS,
         "in progress": Stage.IN_PROGRESS,
     }
     default_response_configuration = ExtendedResponseConfiguration(return_embedded=True)
 
+    # kg_core sets no timeout on its outbound HTTP requests, and exposes no supported way to
+    # configure one (KGConfig carries no timeout field, and neither ClientBuilder.build() nor
+    # .build_admin() offer a hook for one). A request that loses its connection mid-flight -
+    # e.g. after the machine sleeps and resumes on a different network - then hangs forever
+    # instead of failing. This is a temporary workaround, patching a private (and therefore
+    # unstable) kg_core module, pending a fix in kg_core itself; remove it once that exists.
+    #
+    # While we're in there, we also give the two most common network failures a more specific
+    # KGConnectionError message than the generic one from translate_network_errors above:
+    #  - requests.Timeout: no response within KG_REQUEST_TIMEOUT (our own client-side limit).
+    #  - requests.ConnectionError: the connection was lost before a response arrived, e.g. a
+    #    "Connection aborted ... RemoteDisconnected" when the KG's own gateway drops a
+    #    request that has been running too long (observed around 50s against core.kg-ppd,
+    #    though slow-but-successful `File` queries on core.kg have taken much longer than
+    #    that, so this is not a fixed value we can rely on). If this patch can't be applied,
+    #    translate_network_errors above still turns both into a (less specific) KGConnectionError.
+    KG_REQUEST_TIMEOUT = 300  # seconds
+    try:
+        from kg_core.__communication import RequestsWithTokenHandler
+
+        _kg_core_do_request = RequestsWithTokenHandler._do_request
+
+        def _do_request_with_timeout(self, args, payload):
+            args.setdefault("timeout", KG_REQUEST_TIMEOUT)
+            request_description = f"{args.get('method')} {args.get('url')}"
+            try:
+                return _kg_core_do_request(self, args, payload)
+            except requests.exceptions.Timeout as err:
+                raise KGConnectionError(
+                    f"No response from the KG within {KG_REQUEST_TIMEOUT}s ({request_description}). "
+                    "The KG may still be processing the request (this is more likely for queries "
+                    "against types with very many instances, such as File); consider passing a "
+                    "larger request_timeout to KGClient, or retrying."
+                ) from err
+            except requests.exceptions.ConnectionError as err:
+                raise KGConnectionError(
+                    f"Lost connection to the KG before a response arrived ({request_description}): "
+                    f"{err}. This can happen when the network connection changes (e.g. switching "
+                    "wifi networks) or when the KG's own infrastructure drops a long-running "
+                    "request; retrying the request may succeed."
+                ) from err
+
+        RequestsWithTokenHandler._do_request = _do_request_with_timeout
+    except (ImportError, AttributeError) as err:
+        logger.warning(
+            "Could not patch kg_core to apply a request timeout (%s). "
+            "Requests to the KG may hang indefinitely instead of timing out.",
+            err,
+        )
+
 
 BARE_UUID = re.compile(r"[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}")
 
@@ -125,6 +229,20 @@ class KGClient(object):
         openminds_version (str, default "v4"): the openMINDS schema version that responses should be
             deserialized into. Must be one of "v4" or "v5". v4 is the default so existing code is
             unaffected; pass "v5" when connecting to a KG instance that has been migrated to v5.
+        request_timeout (int, optional): how many seconds to wait for a response from the KG
+            before giving up, since kg_core sets no timeout of its own (see KG_REQUEST_TIMEOUT).
+            This is a temporary workaround applied process-wide, not per client: passing it here
+            changes the timeout for every KGClient in the process, not just this instance.
+            Leave unset to keep whatever timeout is currently in effect.
+        max_retries (int, default 2): how many times to retry a request that fails with
+            KGConnectionError (a lost connection, or no response within request_timeout),
+            using exponential backoff. Only applied to read-only or query-only methods
+            (e.g. query, list, instance_from_full_uri, user_info), never to methods that
+            modify the KG (e.g. create_new_instance, update_instance, delete_instance,
+            release), since retrying a write risks duplicating or corrupting data if the
+            original request actually succeeded but its response was lost.
+        retry_backoff (float, default 5.0): base number of seconds to wait before the first
+            retry; each subsequent retry doubles the wait (5s, 10s, 20s, ...).
 
     Raises:
         ImportError: If the kg_core package is not installed.
@@ -140,6 +258,9 @@ def __init__(
         client_secret: Optional[str] = None,
         allow_interactive: bool = True,
         openminds_version: str = OPENMINDS_VERSION,
+        request_timeout: Optional[int] = None,
+        max_retries: int = 2,
+        retry_backoff: float = 5.0,
     ):
         if openminds_version not in ("v4", "v5"):
             raise ValueError(
@@ -148,6 +269,11 @@ def __init__(
         self.openminds_version = openminds_version
         if not have_kg_core:
             raise ImportError("Please install the ebrains-kg-core package")
+        if request_timeout is not None:
+            global KG_REQUEST_TIMEOUT
+            KG_REQUEST_TIMEOUT = request_timeout
+        self._max_retries = max_retries
+        self._retry_backoff = retry_backoff
         if client_id and client_secret:
             self._kg_client_builder = kg(host).with_credentials(client_id, client_secret)
             self._auth_method = "credentials"
@@ -206,6 +332,8 @@ def _kg_admin_client(self):
         return self.__kg_admin_client
 
     @property
+    @retry_on_connection_error
+    @translate_network_errors
     def token(self) -> Optional[str]:
         return self._kg_client.instances._kg_config.token_handler._fetch_token()
 
@@ -250,6 +378,8 @@ def _check_response(
             expand_bare_uuids(response.data, self._kg_client.instances._kg_config.id_namespace)
             return response
 
+    @retry_on_connection_error
+    @translate_network_errors
     def query(
         self,
         query: Dict[str, Any],
@@ -359,6 +489,8 @@ def _query(release_status, from_index, size):
             response = _query(release_status, from_index, size)
         return response
 
+    @retry_on_connection_error
+    @translate_network_errors
     def list(
         self,
         target_type: str,
@@ -416,6 +548,8 @@ def _list(release_status, from_index, size):
         else:
             return _list(release_status, from_index, size)
 
+    @retry_on_connection_error
+    @translate_network_errors
     def instance_from_full_uri(
         self,
         uri: str,
@@ -493,6 +627,7 @@ def _get_instance(release_status):
                 self.cache[uri] = data
         return data
 
+    @translate_network_errors
     def create_new_instance(
         self, data: JsonLdDocument, space: str, instance_id: Optional[str] = None
     ) -> JsonLdDocument:
@@ -525,6 +660,7 @@ def create_new_instance(
         error_context = f"create_new_instance(data={data}, space={space}, instance_id={instance_id})"
         return self._check_response(response, error_context=error_context).data
 
+    @translate_network_errors
     def update_instance(self, instance_id: str, data: JsonLdDocument) -> JsonLdDocument:
         """
         Update an existing KG instance using the data provided.
@@ -549,6 +685,7 @@ def update_instance(self, instance_id: str, data: JsonLdDocument) -> JsonLdDocum
         self.cache.pop(self.uri_from_uuid(instance_id), None)
         return response_data
 
+    @translate_network_errors
     def replace_instance(self, instance_id: str, data: JsonLdDocument) -> JsonLdDocument:
         """
         Replace an existing KG instance using the data provided.
@@ -569,6 +706,7 @@ def replace_instance(self, instance_id: str, data: JsonLdDocument) -> JsonLdDocu
         self.cache.pop(self.uri_from_uuid(instance_id), None)
         return response_data
 
+    @translate_network_errors
     def delete_instance(self, instance_id: str, ignore_not_found: bool = True, ignore_errors: bool = True):
         """
         Delete a KG instance.
@@ -597,6 +735,7 @@ def uuid_from_uri(self, uri: str) -> UUID:
         assert uri.startswith(namespace)
         return UUID(uri[len(namespace) :])
 
+    @translate_network_errors
     def store_query(self, query_label: str, query_definition: Dict[str, Any], space: str):
         """
         Store a query definition in the KG.
@@ -625,6 +764,8 @@ def store_query(self, query_label: str, query_definition: Dict[str, Any], space:
 
         query_definition["@id"] = self.uri_from_uuid(query_id)
 
+    @retry_on_connection_error
+    @translate_network_errors
     def retrieve_query(self, query_label: str) -> Dict[str, Any]:
         """
         Retrieve a stored query definition from the KG.
@@ -652,6 +793,8 @@ def retrieve_query(self, query_label: str) -> Dict[str, Any]:
             self._query_cache[query_label] = query_definition
         return self._query_cache[query_label]
 
+    @retry_on_connection_error
+    @translate_network_errors
     def user_info(self) -> Dict[str, Any]:
         """
         Returns information about the current user.
@@ -670,6 +813,8 @@ def user_info(self) -> Dict[str, Any]:
                 raise Exception(response.error)
         return self._user_info
 
+    @retry_on_connection_error
+    @translate_network_errors
     def spaces(
         self, permissions: Optional[Iterable[str] | bool] = None, names_only: bool = False
     ) -> Union[List[str], List[SpaceInformation]]:
@@ -717,6 +862,7 @@ def _private_space(self) -> str:
         # temporary workaround
         return f"private-{self.user_info().identifiers[0]}"
 
+    @translate_network_errors
     def configure_space(self, space_name: Optional[str] = None, types: Optional[List[KGObject]] = None) -> str:
         """
         Creates and configures a Knowledge Graph (KG) space with the specified name and types.
@@ -761,6 +907,7 @@ def configure_space(self, space_name: Optional[str] = None, types: Optional[List
                 raise Exception(f"Unable to assign {cls.__name__} to space {space_name}: {result}")
         return space_name
 
+    @translate_network_errors
     def move_to_space(self, uri: str, destination_space: str):
         """
         Move a KG instance from one space to another.
@@ -772,6 +919,8 @@ def move_to_space(self, uri: str, destination_space: str):
         if response.error:
             raise Exception(response.error)
 
+    @retry_on_connection_error
+    @translate_network_errors
     def space_info(
         self,
         space_name: str,
@@ -880,6 +1029,8 @@ def move_all_to_space(self, source_space: str, destination_space: str):
         else:
             print(f"The space '{source_space}' is empty, nothing to move.")
 
+    @retry_on_connection_error
+    @translate_network_errors
     def is_released(self, uri: str, with_children: bool = False) -> bool:
         """
         Release status of a KG instance identified by its URI.
@@ -907,12 +1058,14 @@ def is_released(self, uri: str, with_children: bool = False) -> bool:
         else:
             raise AuthorizationError("You are not able to access the release status")
 
+    @translate_network_errors
     def release(self, uri: str):
         """Release the instance with the given uri"""
         response = self._kg_client.instances.release(self.uuid_from_uri(uri))
         if response:
             raise Exception(f"Can't release instance with id {uri}. Error message: {response}")
 
+    @translate_network_errors
     def unrelease(self, uri: str):
         """Unrelease the instance with the given uri"""
         response = self._kg_client.instances.unrelease(self.uuid_from_uri(uri))

fairgraph/errors.py

@@ -45,3 +45,11 @@ class CannotBuildExistenceQuery(Exception):
     """Raised when it is not possible to build an existence query"""
 
     pass
+
+
+class KGConnectionError(Exception):
+    """Raised when a request to the KG fails for network reasons (lost connection,
+    or no response within the configured timeout) rather than because the KG itself
+    reported an error"""
+
+    pass