diff --git a/fairgraph/client.py b/fairgraph/client.py index 3a21a3b8..08c681a0 100644 --- a/fairgraph/client.py +++ b/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,7 +54,59 @@ 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, @@ -60,6 +114,56 @@ } 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)) diff --git a/fairgraph/errors.py b/fairgraph/errors.py index 6c7830a4..6ddfc795 100644 --- a/fairgraph/errors.py +++ b/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 diff --git a/test/test_client.py b/test/test_client.py index c5e621b0..b7bae622 100644 --- a/test/test_client.py +++ b/test/test_client.py @@ -1,12 +1,13 @@ import copy import os import pytest +import requests from kg_core.request import Stage, Pagination from kg_core.response import Error as KGError from fairgraph.kgobject import KGObject from fairgraph.queries import Query, QueryProperty, Filter -from fairgraph.errors import AuthenticationError, AuthorizationError, ResourceExistsError +from fairgraph.errors import AuthenticationError, AuthorizationError, KGConnectionError, ResourceExistsError from fairgraph.base import OPENMINDS_VERSION from fairgraph.client import KGClient from .utils import ( @@ -367,6 +368,46 @@ def offline_kg_client(mocker): return client +class TestRetryOnConnectionError: + """Read-only/query-only KGClient methods retry on KGConnectionError; methods that + modify the KG never do, since retrying a write risks duplicating or corrupting + data if the original request actually succeeded but its response was lost.""" + + def test_retryable_method_recovers_after_connection_error(self, offline_kg_client, mocker): + mocker.patch("fairgraph.client.sleep") + mocker.patch.object( + offline_kg_client._kg_client.instances, + "list", + side_effect=[requests.exceptions.ConnectionError("boom"), MockKGResponse([])], + ) + response = offline_kg_client.list("https://openminds.om-i.org/types/File") + assert response.data == [] + assert offline_kg_client._kg_client.instances.list.call_count == 2 + + def test_retryable_method_exhausts_retries(self, offline_kg_client, mocker): + mocker.patch("fairgraph.client.sleep") + mocker.patch.object( + offline_kg_client._kg_client.instances, + "list", + side_effect=requests.exceptions.ConnectionError("boom"), + ) + with pytest.raises(KGConnectionError): + offline_kg_client.list("https://openminds.om-i.org/types/File") + assert offline_kg_client._kg_client.instances.list.call_count == offline_kg_client._max_retries + 1 + + def test_mutating_method_does_not_retry(self, offline_kg_client, mocker): + mocker.patch("fairgraph.client.sleep") + mocker.patch.object( + offline_kg_client._kg_client.instances, + "delete", + side_effect=requests.exceptions.ConnectionError("boom"), + ) + fake_id = "00000000-0000-0000-0000-000000000000" + with pytest.raises(KGConnectionError): + offline_kg_client.delete_instance(fake_id) + offline_kg_client._kg_client.instances.delete.assert_called_once() + + class TestCacheInvalidationOnWrite: """Regression tests for the bug where writes left stale entries in `client.cache`, causing subsequent `from_id(use_cache=True)` calls to diff --git a/test/utils.py b/test/utils.py index f77fc447..5b13a3f9 100644 --- a/test/utils.py +++ b/test/utils.py @@ -69,6 +69,8 @@ def __init__(self, data, error=None): class MockKGClient: _private_space = "myspace_1234" + _max_retries = 2 + _retry_backoff = 5.0 def __init__(self, openminds_version: str = OPENMINDS_VERSION): if openminds_version not in ("v4", "v5"):