Extract raise_for_http_status() to net.util

neptune-ai · Jan 14, 2025 · 02dd56f · 02dd56f
1 parent 259f789
commit 02dd56f
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 20 deletions.
diff --git a/src/neptune_scale/net/util.py b/src/neptune_scale/net/util.py
@@ -13,10 +13,37 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from neptune_scale.exceptions import (
+    NeptuneConnectionLostError,
+    NeptuneInternalServerError,
+    NeptuneTooManyRequestsResponseError,
+    NeptuneUnauthorizedError,
+    NeptuneUnexpectedResponseError,
+)
+from neptune_scale.util import get_logger
+
+logger = get_logger()
+
 
 def escape_nql_criterion(criterion: str) -> str:
     """
     Escape backslash and (double-)quotes in the string, to match what the NQL engine expects.
     """
 
     return criterion.replace("\\", r"\\").replace('"', r"\"")
+
+
+def raise_for_http_status(status_code: int) -> None:
+    assert status_code >= 400, f"Status code {status_code} is not an error"
+
+    logger.error("HTTP response error: %s", status_code)
+    if status_code == 403:
+        raise NeptuneUnauthorizedError()
+    elif status_code == 408:
+        raise NeptuneConnectionLostError()
+    elif status_code == 429:
+        raise NeptuneTooManyRequestsResponseError()
+    elif status_code // 100 == 5:
+        raise NeptuneInternalServerError()
+    else:
+        raise NeptuneUnexpectedResponseError()
diff --git a/src/neptune_scale/sync/sync_process.py b/src/neptune_scale/sync/sync_process.py
@@ -54,7 +54,6 @@
     NeptuneAttributeTypeUnsupported,
     NeptuneConnectionLostError,
     NeptuneFloatValueNanInfUnsupported,
-    NeptuneInternalServerError,
     NeptuneProjectInvalidName,
     NeptuneProjectNotFound,
     NeptuneRetryableError,
@@ -70,16 +69,14 @@
     NeptuneStringSetExceedsSizeLimit,
     NeptuneStringValueExceedsSizeLimit,
     NeptuneSynchronizationStopped,
-    NeptuneTooManyRequestsResponseError,
-    NeptuneUnauthorizedError,
     NeptuneUnexpectedError,
-    NeptuneUnexpectedResponseError,
 )
 from neptune_scale.net.api_client import (
     ApiClient,
     backend_factory,
     with_api_errors_handling,
 )
+from neptune_scale.net.util import raise_for_http_status
 from neptune_scale.sync.aggregating_queue import AggregatingQueue
 from neptune_scale.sync.errors_tracking import ErrorsQueue
 from neptune_scale.sync.parameters import (
@@ -438,7 +435,7 @@ def submit(self, *, operation: RunOperation) -> Optional[SubmitResponse]:
 
         status_code = response.status_code
         if status_code != 200:
-            _raise_exception(status_code)
+            raise_for_http_status(status_code)
 
         return response.parsed
 
@@ -482,20 +479,6 @@ def work(self) -> None:
             raise NeptuneSynchronizationStopped() from e
 
 
-def _raise_exception(status_code: int) -> None:
-    logger.error("HTTP response error: %s", status_code)
-    if status_code == 403:
-        raise NeptuneUnauthorizedError()
-    elif status_code == 408:
-        raise NeptuneConnectionLostError()
-    elif status_code == 429:
-        raise NeptuneTooManyRequestsResponseError()
-    elif status_code // 100 == 5:
-        raise NeptuneInternalServerError()
-    else:
-        raise NeptuneUnexpectedResponseError()
-
-
 class StatusTrackingThread(Daemon, WithResources):
     def __init__(
         self,
@@ -542,7 +525,7 @@ def check_batch(self, *, request_ids: list[str]) -> Optional[BulkRequestStatus]:
         status_code = response.status_code
 
         if status_code != 200:
-            _raise_exception(status_code)
+            raise_for_http_status(status_code)
 
         return response.parsed