From 691a19cd2a83447e0307a2933471a583fd57e043 Mon Sep 17 00:00:00 2001 From: Volo Kluev Date: Thu, 19 Sep 2024 15:09:19 -0700 Subject: [PATCH] make tests pass --- snuba/web/rpc/__init__.py | 51 +++++++++++++++++++ snuba/web/rpc/{v1alpha => }/common.py | 0 snuba/web/rpc/{v1alpha => }/exceptions.py | 0 snuba/web/rpc/v1alpha/__init__.py | 32 ------------ snuba/web/views.py | 8 +-- tests/web/rpc/test_span_samples.py | 4 +- tests/web/rpc/test_timeseries_api.py | 4 +- .../web/rpc/test_trace_item_attribute_list.py | 4 +- .../rpc/test_trace_item_attribute_values.py | 4 +- 9 files changed, 65 insertions(+), 42 deletions(-) create mode 100644 snuba/web/rpc/__init__.py rename snuba/web/rpc/{v1alpha => }/common.py (100%) rename snuba/web/rpc/{v1alpha => }/exceptions.py (100%) diff --git a/snuba/web/rpc/__init__.py b/snuba/web/rpc/__init__.py new file mode 100644 index 0000000000..63cb374a3b --- /dev/null +++ b/snuba/web/rpc/__init__.py @@ -0,0 +1,51 @@ +from typing import Any, Callable, Mapping, Tuple + +from google.protobuf.message import Message as ProtobufMessage +from sentry_protos.snuba.v1alpha.endpoint_aggregate_bucket_pb2 import ( + AggregateBucketRequest, +) +from sentry_protos.snuba.v1alpha.endpoint_span_samples_pb2 import SpanSamplesRequest +from sentry_protos.snuba.v1alpha.endpoint_tags_list_pb2 import ( + AttributeValuesRequest, + TraceItemAttributesRequest, +) + +from snuba.utils.metrics.timer import Timer +from snuba.web.rpc.v1alpha.span_samples import span_samples_query +from snuba.web.rpc.v1alpha.timeseries import timeseries_query +from snuba.web.rpc.v1alpha.trace_item_attribute_list import ( + trace_item_attribute_list_query, +) +from snuba.web.rpc.v1alpha.trace_item_attribute_values import ( + trace_item_attribute_values_query, +) + +Version = str +EndpointName = str + +ALL_RPCS: Mapping[ + Version, + Mapping[ + EndpointName, + Tuple[Callable[[Any, Timer], ProtobufMessage], type[ProtobufMessage]], + ], +] = { + "v1alpha": { + "AggregateBucketRequest": (timeseries_query, AggregateBucketRequest), + "SpanSamplesRequest": (span_samples_query, SpanSamplesRequest), + "TraceItemAttributesRequest": ( + trace_item_attribute_list_query, + TraceItemAttributesRequest, + ), + "AttributeValuesRequest": ( + trace_item_attribute_values_query, + AttributeValuesRequest, + ), + } +} + + +def get_rpc_endpoint( + name: str, version: str +) -> Tuple[Callable[[Any, Timer], ProtobufMessage], type[ProtobufMessage]]: + return ALL_RPCS[version][name] diff --git a/snuba/web/rpc/v1alpha/common.py b/snuba/web/rpc/common.py similarity index 100% rename from snuba/web/rpc/v1alpha/common.py rename to snuba/web/rpc/common.py diff --git a/snuba/web/rpc/v1alpha/exceptions.py b/snuba/web/rpc/exceptions.py similarity index 100% rename from snuba/web/rpc/v1alpha/exceptions.py rename to snuba/web/rpc/exceptions.py diff --git a/snuba/web/rpc/v1alpha/__init__.py b/snuba/web/rpc/v1alpha/__init__.py index 6034257542..e69de29bb2 100644 --- a/snuba/web/rpc/v1alpha/__init__.py +++ b/snuba/web/rpc/v1alpha/__init__.py @@ -1,32 +0,0 @@ -from typing import Any, Callable, Mapping, Tuple - -from google.protobuf.message import Message as ProtobufMessage -from sentry_protos.snuba.v1alpha.endpoint_aggregate_bucket_pb2 import ( - AggregateBucketRequest, -) -from sentry_protos.snuba.v1alpha.endpoint_span_samples_pb2 import SpanSamplesRequest -from sentry_protos.snuba.v1alpha.endpoint_tags_list_pb2 import ( - AttributeValuesRequest, - TraceItemAttributesRequest, -) - -from snuba.utils.metrics.timer import Timer -from snuba.web.rpc.span_samples import span_samples_query -from snuba.web.rpc.timeseries import timeseries_query -from snuba.web.rpc.trace_item_attribute_list import trace_item_attribute_list_query -from snuba.web.rpc.trace_item_attribute_values import trace_item_attribute_values_query - -ALL_RPCS: Mapping[ - str, Tuple[Callable[[Any, Timer], ProtobufMessage], type[ProtobufMessage]] -] = { - "AggregateBucketRequest": (timeseries_query, AggregateBucketRequest), - "SpanSamplesRequest": (span_samples_query, SpanSamplesRequest), - "TraceItemAttributesRequest": ( - trace_item_attribute_list_query, - TraceItemAttributesRequest, - ), - "AttributeValuesRequest": ( - trace_item_attribute_values_query, - AttributeValuesRequest, - ), -} diff --git a/snuba/web/views.py b/snuba/web/views.py index ce50b5f6e9..35ba953629 100644 --- a/snuba/web/views.py +++ b/snuba/web/views.py @@ -75,7 +75,7 @@ from snuba.web.converters import DatasetConverter, EntityConverter, StorageConverter from snuba.web.delete_query import DeletesNotEnabledError, delete_from_storage from snuba.web.query import parse_and_run_query -from snuba.web.rpc import ALL_RPCS +from snuba.web.rpc import get_rpc_endpoint from snuba.web.rpc.exceptions import BadSnubaRPCRequestException from snuba.writer import BatchWriterEncoderWrapper, WriterTableRow @@ -272,11 +272,11 @@ def unqualified_query_view(*, timer: Timer) -> Union[Response, str, WerkzeugResp assert False, "unexpected fallthrough" -@application.route("/rpc//", methods=["POST"]) +@application.route("/rpc//", methods=["POST"]) @util.time_request("rpc") -def rpc(*, name: str, timer: Timer) -> Response: +def rpc(*, name: str, version: str, timer: Timer) -> Response: try: - endpoint, req_class = ALL_RPCS[name] + endpoint, req_class = get_rpc_endpoint(name, version) req = req_class() req.ParseFromString(http_request.data) diff --git a/tests/web/rpc/test_span_samples.py b/tests/web/rpc/test_span_samples.py index 1e70c69ceb..903958ee34 100644 --- a/tests/web/rpc/test_span_samples.py +++ b/tests/web/rpc/test_span_samples.py @@ -21,7 +21,7 @@ from snuba.datasets.storages.factory import get_storage from snuba.datasets.storages.storage_key import StorageKey -from snuba.web.rpc.span_samples import span_samples_query +from snuba.web.rpc.v1alpha.span_samples import span_samples_query from tests.base import BaseApiTest from tests.helpers import write_raw_unprocessed_events @@ -141,7 +141,7 @@ def test_basic(self) -> None: limit=10, ) response = self.app.post( - "/rpc/SpanSamplesRequest", data=message.SerializeToString() + "/rpc/SpanSamplesRequest/v1alpha", data=message.SerializeToString() ) assert response.status_code == 200, response.text diff --git a/tests/web/rpc/test_timeseries_api.py b/tests/web/rpc/test_timeseries_api.py index d5093c9641..23519b6ede 100644 --- a/tests/web/rpc/test_timeseries_api.py +++ b/tests/web/rpc/test_timeseries_api.py @@ -13,7 +13,7 @@ from snuba.datasets.storages.factory import get_storage from snuba.datasets.storages.storage_key import StorageKey -from snuba.web.rpc.timeseries import timeseries_query +from snuba.web.rpc.v1alpha.timeseries import timeseries_query from tests.base import BaseApiTest from tests.helpers import write_raw_unprocessed_events @@ -127,7 +127,7 @@ def test_basic(self) -> None: granularity_secs=60, ) response = self.app.post( - "/rpc/AggregateBucketRequest", data=message.SerializeToString() + "/rpc/AggregateBucketRequest/v1alpha", data=message.SerializeToString() ) assert response.status_code == 200 diff --git a/tests/web/rpc/test_trace_item_attribute_list.py b/tests/web/rpc/test_trace_item_attribute_list.py index ecb0846e70..e349f1c981 100644 --- a/tests/web/rpc/test_trace_item_attribute_list.py +++ b/tests/web/rpc/test_trace_item_attribute_list.py @@ -13,7 +13,9 @@ from snuba.datasets.storages.factory import get_storage from snuba.datasets.storages.storage_key import StorageKey -from snuba.web.rpc.trace_item_attribute_list import trace_item_attribute_list_query +from snuba.web.rpc.v1alpha.trace_item_attribute_list import ( + trace_item_attribute_list_query, +) from tests.base import BaseApiTest from tests.helpers import write_raw_unprocessed_events diff --git a/tests/web/rpc/test_trace_item_attribute_values.py b/tests/web/rpc/test_trace_item_attribute_values.py index 3f400219be..fc2f1a4fae 100644 --- a/tests/web/rpc/test_trace_item_attribute_values.py +++ b/tests/web/rpc/test_trace_item_attribute_values.py @@ -9,7 +9,9 @@ from snuba.datasets.storages.factory import get_storage from snuba.datasets.storages.storage_key import StorageKey -from snuba.web.rpc.trace_item_attribute_values import trace_item_attribute_values_query +from snuba.web.rpc.v1alpha.trace_item_attribute_values import ( + trace_item_attribute_values_query, +) from tests.base import BaseApiTest from tests.helpers import write_raw_unprocessed_events