Merge branch 'master' into external-pr-intermediate

cognitedata · Dec 13, 2023 · 96a3ceb · 96a3ceb
2 parents 00a8c8d + a21efcb
commit 96a3ceb
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 7 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -12,12 +12,19 @@ Changes are grouped as follows
 - `Fixed` for any bug fixes.
 - `Security` in case of vulnerabilities.
 
+
 ## [6.1.0]
 
 ### Added
 
  * Added ability to specify dataset under which metrics timeseries are created
 
+## [6.0.2]
+
+### Fixed
+
+ * Improved the state store retry behavior to handle connection errors
+
 ## [6.0.1]
 
 ### Fixed

diff --git a/cognite/extractorutils/__init__.py b/cognite/extractorutils/__init__.py
@@ -16,5 +16,5 @@
 Cognite extractor utils is a Python package that simplifies the development of new extractors.
 """
 
-__version__ = "6.1.0"
+__version__ = "6.1.0
 from .base import Extractor
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,