From 8c7e2894df308dc19bf642904176dd1f5805ee05 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Anita=20=C5=9Aled=C5=BA?= <anita.gaja.sledz@gmail.com>
Date: Thu, 5 Oct 2023 13:20:12 +0200
Subject: [PATCH] add code review

---
 src/neptune/common/backends/utils.py             | 16 +++++++++++-----
 src/neptune/common/utils.py                      | 13 -------------
 .../api_clients/hosted_api_clients/utils.py      |  2 +-
 3 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/src/neptune/common/backends/utils.py b/src/neptune/common/backends/utils.py
index 5a8a2615a..29323e0d3 100644
--- a/src/neptune/common/backends/utils.py
+++ b/src/neptune/common/backends/utils.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__all__ = ["with_api_exceptions_handler"]
+__all__ = ["with_api_exceptions_handler", "get_retry_from_headers_or_default"]
 
 import itertools
 import logging
@@ -47,10 +47,7 @@
     NeptuneSSLVerificationError,
     Unauthorized,
 )
-from neptune.common.utils import (
-    get_retry_from_headers_or_default,
-    reset_internal_ssl_state,
-)
+from neptune.common.utils import reset_internal_ssl_state
 
 _logger = logging.getLogger(__name__)
 
@@ -59,6 +56,15 @@
 retries_timeout = int(os.getenv(NEPTUNE_RETRIES_TIMEOUT_ENV, "60"))
 
 
+def get_retry_from_headers_or_default(headers, retry_count):
+    try:
+        return (
+            int(headers["retry-after"][0]) if "retry-after" in headers else 2 ** min(MAX_RETRY_MULTIPLIER, retry_count)
+        )
+    except Exception:
+        return min(2 ** min(MAX_RETRY_MULTIPLIER, retry_count), MAX_RETRY_TIME)
+
+
 def with_api_exceptions_handler(func):
     def wrapper(*args, **kwargs):
         ssl_error_occurred = False
diff --git a/src/neptune/common/utils.py b/src/neptune/common/utils.py
index 077ce6a7b..5883e233a 100644
--- a/src/neptune/common/utils.py
+++ b/src/neptune/common/utils.py
@@ -268,16 +268,3 @@ def __enter__(self):
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         pass
-
-
-MAX_RETRY_TIME = 30
-MAX_RETRY_MULTIPLIER = 10
-
-
-def get_retry_from_headers_or_default(headers, retry_count):
-    try:
-        return (
-            int(headers["retry-after"][0]) if "retry-after" in headers else 2 ** min(MAX_RETRY_MULTIPLIER, retry_count)
-        )
-    except Exception:
-        return min(2 ** min(MAX_RETRY_MULTIPLIER, retry_count), MAX_RETRY_TIME)
diff --git a/src/neptune/legacy/internal/api_clients/hosted_api_clients/utils.py b/src/neptune/legacy/internal/api_clients/hosted_api_clients/utils.py
index 4700b9469..b59a6d188 100644
--- a/src/neptune/legacy/internal/api_clients/hosted_api_clients/utils.py
+++ b/src/neptune/legacy/internal/api_clients/hosted_api_clients/utils.py
@@ -32,7 +32,7 @@
 )
 from urllib3.exceptions import NewConnectionError
 
-from neptune.common.utils import get_retry_from_headers_or_default
+from neptune.common.backends.utils import get_retry_from_headers_or_default
 from neptune.legacy.api_exceptions import (
     ConnectionLost,
     Forbidden,