From 6e228db2d2d6fac3a4104fa9218f3d39639c51a6 Mon Sep 17 00:00:00 2001
From: Thorkild Stray <thorkild.stray@cognitedata.com>
Date: Fri, 12 Jan 2024 10:55:19 +0100
Subject: [PATCH] Retry on requests.exceptions.ConnectionError (again) (#282)

In a previous PR, we switched from retrying on requests connection
errors to the internal Cognite exceptions that wrapped these. This
fixed some of the issues encountered by the customer (SLB).

But, in some corner cases, such as for authorization, you can still
get an unwrapped requests.exceptions.ConnectionError (as seen by
the customer).

This re-adds the retry on ConnectionError in addition to the internal
Cognite SDK classes, to be more resilient.
---
 CHANGELOG.md                         | 7 +++++++
 cognite/extractorutils/statestore.py | 8 +++++---
 pyproject.toml                       | 2 +-
 3 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8de8b5f9..f4c10c06 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,13 @@ Changes are grouped as follows
 - `Fixed` for any bug fixes.
 - `Security` in case of vulnerabilities.
 
+## [6.3.1]
+
+### Fixed
+
+ * Improved the state store retry behavior to handle both fundamental
+   and wrapped network connection errors.
+
 ## [6.3.0]
 
 ### Added
diff --git a/cognite/extractorutils/statestore.py b/cognite/extractorutils/statestore.py
index 14b959d8..810c87bb 100644
--- a/cognite/extractorutils/statestore.py
+++ b/cognite/extractorutils/statestore.py
@@ -92,6 +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, CogniteException
 from cognite.extractorutils.uploader import DataPointList
@@ -346,7 +348,7 @@ def __init__(
         self._ensure_table()
 
     @retry(
-        exceptions=(CogniteException,),
+        exceptions=(CogniteException, ConnectionError),
         tries=RETRIES,
         delay=RETRY_DELAY,
         max_delay=RETRY_MAX_DELAY,
@@ -368,7 +370,7 @@ def initialize(self, force: bool = False) -> None:
         self._initialize_implementation(force)
 
     @retry(
-        exceptions=(CogniteException,),
+        exceptions=(CogniteException, ConnectionError),
         tries=RETRIES,
         delay=RETRY_DELAY,
         max_delay=RETRY_MAX_DELAY,
@@ -402,7 +404,7 @@ def synchronize(self) -> None:
         self._synchronize_implementation()
 
     @retry(
-        exceptions=(CogniteException,),
+        exceptions=(CogniteException, ConnectionError),
         tries=RETRIES,
         delay=RETRY_DELAY,
         max_delay=RETRY_MAX_DELAY,
diff --git a/pyproject.toml b/pyproject.toml
index 56c4cc77..8bbf275c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognite-extractor-utils"
-version = "6.3.0"
+version = "6.3.1"
 description = "Utilities for easier development of extractors for CDF"
 authors = ["Mathias Lohne <mathias.lohne@cognite.com>"]
 license = "Apache-2.0"