Added Accept and Accept-Encoding headers to protocol buffer reque…

…sts (#1736)

CHANGELOG.md

@@ -28,6 +28,7 @@
 - Fixed `tqdm.notebook` import only in Notebook environment ([#1716](https://github.com/neptune-ai/neptune-client/pull/1716))
 - Added `setuptools` to dependencies for `python >= 3.12` ([#1721](https://github.com/neptune-ai/neptune-client/pull/1721))
 - Fixed compatibility checks with pre-release versions ([#1730](https://github.com/neptune-ai/neptune-client/pull/1730))
+- Added `Accept` and `Accept-Encoding` headers to protocol buffer requests ([#1736](https://github.com/neptune-ai/neptune-client/pull/1736))
 
 
 ## neptune 1.10.0

src/neptune/api/searching_entries.py

@@ -24,8 +24,6 @@
     Optional,
 )
 
-from bravado.client import construct_request  # type: ignore
-from bravado.config import RequestConfig  # type: ignore
 from bravado.exception import HTTPBadRequest  # type: ignore
 from typing_extensions import (
     Literal,
@@ -41,7 +39,10 @@
 )
 from neptune.api.proto.neptune_pb.api.model.leaderboard_entries_pb2 import ProtoLeaderboardEntriesSearchResultDTO
 from neptune.exceptions import NeptuneInvalidQueryException
-from neptune.internal.backends.hosted_client import DEFAULT_REQUEST_KWARGS
+from neptune.internal.backends.hosted_client import (
+    DEFAULT_PROTO_REQUEST_KWARGS,
+    DEFAULT_REQUEST_KWARGS,
+)
 from neptune.internal.backends.nql import (
     NQLAggregator,
     NQLAttributeOperator,
@@ -143,23 +144,14 @@ def get_single_page(
 
     try:
         if use_proto:
-            result = client.api.searchLeaderboardEntriesProto(**params).response().result
+            result = (
+                client.api.searchLeaderboardEntriesProto(**params, **DEFAULT_PROTO_REQUEST_KWARGS).response().result
+            )
             proto_data = ProtoLeaderboardEntriesSearchResultDTO.FromString(result)
             return LeaderboardEntriesSearchResult.from_proto(proto_data)
         else:
-            request_options = DEFAULT_REQUEST_KWARGS.get("_request_options", {})
-            request_config = RequestConfig(request_options, True)
-            request_params = construct_request(client.api.searchLeaderboardEntries, request_options, **params)
-
-            http_client = client.swagger_spec.http_client
-
-            json_data = (
-                http_client.request(request_params, operation=None, request_config=request_config)
-                .response()
-                .incoming_response.json()
-            )
-
-            return LeaderboardEntriesSearchResult.from_dict(json_data)
+            model_data = client.api.searchLeaderboardEntries(**params, **DEFAULT_REQUEST_KWARGS).response().result
+            return LeaderboardEntriesSearchResult.from_model(model_data)
     except HTTPBadRequest as e:
         title = e.response.json().get("title")
         if title == "Syntax error":

src/neptune/internal/backends/hosted_client.py

@@ -15,6 +15,7 @@
 #
 __all__ = [
     "DEFAULT_REQUEST_KWARGS",
+    "DEFAULT_PROTO_REQUEST_KWARGS",
     "create_http_client_with_auth",
     "create_backend_client",
     "create_leaderboard_client",
@@ -65,6 +66,17 @@
     }
 }
 
+DEFAULT_PROTO_REQUEST_KWARGS = {
+    "_request_options": {
+        **DEFAULT_REQUEST_KWARGS["_request_options"],
+        "headers": {
+            **DEFAULT_REQUEST_KWARGS["_request_options"]["headers"],
+            "Accept": "application/x-protobuf,application/json",
+            "Accept-Encoding": "gzip, deflate, br",
+        },
+    }
+}
+
 
 def _close_connections_on_fork(session: requests.Session):
     try:

src/neptune/internal/backends/hosted_neptune_backend.py

@@ -95,6 +95,7 @@
     track_to_new_artifact,
 )
 from neptune.internal.backends.hosted_client import (
+    DEFAULT_PROTO_REQUEST_KWARGS,
     DEFAULT_REQUEST_KWARGS,
     create_artifacts_client,
     create_backend_client,
@@ -1138,16 +1139,29 @@ def get_fields_definitions(
 
         params = {
             "experimentIdentifier": container_id,
-            **DEFAULT_REQUEST_KWARGS,
         }
 
         try:
             if use_proto:
-                result = self.leaderboard_client.api.queryAttributeDefinitionsProto(**params).response().result
+                result = (
+                    self.leaderboard_client.api.queryAttributeDefinitionsProto(
+                        **params,
+                        **DEFAULT_PROTO_REQUEST_KWARGS,
+                    )
+                    .response()
+                    .result
+                )
                 data = ProtoAttributesSearchResultDTO.FromString(result)
                 return [FieldDefinition.from_proto(field_def) for field_def in data.entries]
             else:
-                data = self.leaderboard_client.api.queryAttributeDefinitions(**params).response().result
+                data = (
+                    self.leaderboard_client.api.queryAttributeDefinitions(
+                        **params,
+                        **DEFAULT_REQUEST_KWARGS,
+                    )
+                    .response()
+                    .result
+                )
                 return [FieldDefinition.from_model(field_def) for field_def in data.entries]
         except HTTPNotFound as e:
             raise ContainerUUIDNotFound(
@@ -1171,7 +1185,14 @@ def get_fields_with_paths_filter(
 
         try:
             if use_proto:
-                result = self.leaderboard_client.api.getAttributesWithPathsFilterProto(**params).response().result
+                result = (
+                    self.leaderboard_client.api.getAttributesWithPathsFilterProto(
+                        **params,
+                        **DEFAULT_PROTO_REQUEST_KWARGS,
+                    )
+                    .response()
+                    .result
+                )
                 data = ProtoAttributesDTO.FromString(result)
                 return [Field.from_proto(field) for field in data.attributes]
             else:

tests/unit/neptune/new/api/test_searching_entries.py

@@ -259,14 +259,13 @@ def generate_leaderboard_entries(values: Sequence, experiment_id: str = "foo") -
     )
 
 
-@patch("neptune.api.searching_entries.construct_request")
-def test_get_single_page_error_handling(construct_request_mock):
+def test_get_single_page_error_handling():
     # given
     bravado_exception = HTTPBadRequest(response=Mock())
     bravado_exception.response.json.return_value = {"title": "Syntax error"}
 
-    failing_clinet = Mock()
-    failing_clinet.swagger_spec.http_client.request.side_effect = bravado_exception
+    failing_client = Mock()
+    failing_client.api.searchLeaderboardEntries.side_effect = bravado_exception
 
     # then
     with pytest.raises(NeptuneInvalidQueryException):
@@ -281,5 +280,5 @@ def test_get_single_page_error_handling(construct_request_mock):
             ascending=False,
             sort_by_column_type="string",
             searching_after=None,
-            client=failing_clinet,
+            client=failing_client,
         )