Statestore: Retry on connection errors

The SDK doesn't return a requests ConnectionError when it fails to connect, it returns a CogniteConnectionError. Retry on CogniteException to retry both CogniteAPIErrors and CogniteConnectionErrors.
cognitedata · Dec 13, 2023 · e7ddf63 · e7ddf63
1 parent c3cc2f7
commit e7ddf63
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/cognite/extractorutils/statestore.py b/cognite/extractorutils/statestore.py
@@ -92,10 +92,8 @@
 from types import TracebackType
 from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Type, Union
 
-from requests.exceptions import ConnectionError
-
 from cognite.client import CogniteClient
-from cognite.client.exceptions import CogniteAPIError
+from cognite.client.exceptions import CogniteAPIError, CogniteException
 from cognite.extractorutils.uploader import DataPointList
 
 from ._inner_util import _DecimalDecoder, _DecimalEncoder, _resolve_log_level
@@ -348,7 +346,7 @@ def __init__(
         self._ensure_table()
 
     @retry(
-        exceptions=(CogniteAPIError, ConnectionError),
+        exceptions=(CogniteException,),
         tries=RETRIES,
         delay=RETRY_DELAY,
         max_delay=RETRY_MAX_DELAY,
@@ -370,7 +368,7 @@ def initialize(self, force: bool = False) -> None:
         self._initialize_implementation(force)
 
     @retry(
-        exceptions=(CogniteAPIError, ConnectionError),
+        exceptions=(CogniteException,),
         tries=RETRIES,
         delay=RETRY_DELAY,
         max_delay=RETRY_MAX_DELAY,
@@ -404,7 +402,7 @@ def synchronize(self) -> None:
         self._synchronize_implementation()
 
     @retry(
-        exceptions=(CogniteAPIError, ConnectionError),
+        exceptions=(CogniteException,),
         tries=RETRIES,
         delay=RETRY_DELAY,
         max_delay=RETRY_MAX_DELAY,