From 7918033b15c717a3d939622b8bca4180d7cbdf87 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Jankowski?= Date: Thu, 11 Apr 2024 15:22:37 +0200 Subject: [PATCH] Added `Accept` and `Accept-Encoding` headers to protocol buffer requests (#1736) --- CHANGELOG.md | 1 + src/neptune/api/searching_entries.py | 26 ++++++----------- .../internal/backends/hosted_client.py | 12 ++++++++ .../backends/hosted_neptune_backend.py | 29 ++++++++++++++++--- .../neptune/new/api/test_searching_entries.py | 9 +++--- 5 files changed, 51 insertions(+), 26 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f17895b73..8dc53d2bb 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/neptune/api/searching_entries.py b/src/neptune/api/searching_entries.py index 0f7bc9109..3d5da5a7a 100644 --- a/src/neptune/api/searching_entries.py +++ b/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": diff --git a/src/neptune/internal/backends/hosted_client.py b/src/neptune/internal/backends/hosted_client.py index fb233f07c..69c8a4843 100644 --- a/src/neptune/internal/backends/hosted_client.py +++ b/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: diff --git a/src/neptune/internal/backends/hosted_neptune_backend.py b/src/neptune/internal/backends/hosted_neptune_backend.py index 4313923a7..4d1079123 100644 --- a/src/neptune/internal/backends/hosted_neptune_backend.py +++ b/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: diff --git a/tests/unit/neptune/new/api/test_searching_entries.py b/tests/unit/neptune/new/api/test_searching_entries.py index 01e081871..0fe95bb60 100644 --- a/tests/unit/neptune/new/api/test_searching_entries.py +++ b/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, )