Gzip request body if server support it (#1476)

Co-authored-by: BartoszPrusak <[email protected]>

CHANGELOG.md

@@ -2,7 +2,7 @@
 
 ### Changes
 - Upgraded performance of sending series data to Neptune ([#1483](https://github.com/neptune-ai/neptune-client/pull/1483))
-
+- Compress (gzip) request to server, when server support it ([#1476](https://github.com/neptune-ai/neptune-client/pull/1476))
 
 ## neptune 1.8.1
 

src/neptune/internal/backends/api_model.py

@@ -132,6 +132,7 @@ class ClientConfig:
     _missing_features: FrozenSet[str]
     version_info: VersionInfo
     multipart_config: MultipartConfig
+    gzip_upload: bool
 
     def has_feature(self, feature_name: str) -> bool:
         return feature_name not in self._missing_features
@@ -162,6 +163,7 @@ def from_api_response(config) -> "ClientConfig":
             multipart_upload_config = MultipartConfig(
                 min_chunk_size, max_chunk_size, max_chunk_count, max_single_part_size
             )
+        gzip_upload = getattr(config, "gzipUpload", False)
 
         artifacts_config_obj = getattr(config, "artifacts", None)
         has_artifacts = getattr(artifacts_config_obj, "enabled", False)
@@ -179,6 +181,7 @@ def from_api_response(config) -> "ClientConfig":
             _missing_features=frozenset(missing_features),
             version_info=VersionInfo.build(min_recommended, min_compatible, max_compatible),
             multipart_config=multipart_upload_config,
+            gzip_upload=gzip_upload,
         )
 
 

src/neptune/internal/backends/hosted_client.py

@@ -23,6 +23,7 @@
 
 import os
 import platform
+import zlib
 from typing import (
     Dict,
     Tuple,
@@ -31,6 +32,11 @@
 import requests
 from bravado.http_client import HttpClient
 from bravado.requests_client import RequestsClient
+from requests import (
+    PreparedRequest,
+    Response,
+)
+from requests.adapters import HTTPAdapter
 
 from neptune.common.backends.utils import with_api_exceptions_handler
 from neptune.common.oauth import NeptuneAuthenticator
@@ -48,6 +54,7 @@
     verify_host_resolution,
 )
 from neptune.internal.credentials import Credentials
+from neptune.internal.utils.logger import logger
 from neptune.version import version as neptune_client_version
 
 BACKEND_SWAGGER_PATH = "/api/backend/swagger.json"
@@ -66,6 +73,22 @@
 }
 
 
+class GzipAdapter(HTTPAdapter):
+    def send(self, request: PreparedRequest, stream: bool = False, **kw) -> Response:
+        if request.body is not None and not stream and request.headers.get("Content-Type", None) == "application/json":
+            try:
+                request_body = request.body if isinstance(request.body, bytes) else bytes(request.body, "utf-8")
+                gzip_compress = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, zlib.MAX_WBITS | 16)
+                compressed = gzip_compress.compress(request_body) + gzip_compress.flush()
+                request.prepare_body(compressed, None)
+                request.headers["Content-Encoding"] = "gzip"
+            except zlib.error:
+                logger.warning("Error on compressing request")
+                pass
+
+        return super(GzipAdapter, self).send(request, stream, **kw)
+
+
 def _close_connections_on_fork(session: requests.Session):
     try:
         os.register_at_fork(before=session.close, after_in_child=session.close, after_in_parent=session.close)
@@ -174,7 +197,11 @@ def create_backend_client(client_config: ClientConfig, http_client: HttpClient)
 
 
 @cache
-def create_leaderboard_client(client_config: ClientConfig, http_client: HttpClient) -> SwaggerClientWrapper:
+def create_leaderboard_client(client_config: ClientConfig, http_client: RequestsClient) -> SwaggerClientWrapper:
+    if client_config.gzip_upload:
+        http_client.session.mount("http://", GzipAdapter())
+        http_client.session.mount("https://", GzipAdapter())
+
     return SwaggerClientWrapper(
         create_swagger_client(
             build_operation_url(client_config.api_url, LEADERBOARD_SWAGGER_PATH),