Skip to content

Commit

Permalink
Order by start_timestamp only
Browse files Browse the repository at this point in the history
  • Loading branch information
@phacops
phacops committed Jan 3, 2025
1 parent 293f71a commit 01651ca
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 46 deletions.
30 changes: 10 additions & 20 deletions snuba/web/rpc/v1/endpoint_get_traces.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import uuid
from collections import defaultdict
from typing import Any, Callable, Dict, Iterable, Sequence, Type
from typing import Any, Callable, Dict, Iterable, Type

from google.protobuf.internal.containers import RepeatedCompositeFieldContainer
from google.protobuf.json_format import MessageToDict
Expand Down Expand Up @@ -42,6 +42,7 @@
from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException

_DEFAULT_ROW_LIMIT = 10_000
_BUFFER_WINDOW = 2 * 3600 # 2 hours

_ATTRIBUTES: dict[
TraceAttribute.Key.ValueType,
Expand Down Expand Up @@ -119,22 +120,6 @@ def _attribute_to_expression(
)


def _convert_order_by(
order_by: Sequence[GetTracesRequest.OrderBy],
) -> Sequence[OrderBy]:
res: list[OrderBy] = []
for x in order_by:
res.append(
OrderBy(
direction=(
OrderByDirection.DESC if x.descending else OrderByDirection.ASC
),
expression=_attribute_to_expression(TraceAttribute(key=x.key)),
)
)
return res


def _build_snuba_request(request: GetTracesRequest, query: Query) -> SnubaRequest:
query_settings = (
setup_trace_query_settings() if request.meta.debug else HTTPQuerySettings()
Expand Down Expand Up @@ -383,8 +368,8 @@ def _get_metadata_for_traces(
condition=and_cond(
project_id_and_org_conditions(request.meta),
timestamp_in_range_condition(
min(timestamps) - 2 * 3600,
max(timestamps) + 2 * 3600,
min(timestamps) - _BUFFER_WINDOW,
max(timestamps) + _BUFFER_WINDOW,
),
in_cond(
f.cast(
Expand All @@ -404,7 +389,12 @@ def _get_metadata_for_traces(
),
),
],
order_by=_convert_order_by(request.order_by),
order_by=[
OrderBy(
direction=OrderByDirection.DESC,
expression=column("trace_start_timestamp"),
),
],
)

treeify_or_and_conditions(query)
Expand Down
42 changes: 16 additions & 26 deletions tests/web/rpc/v1/test_endpoint_get_traces.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -183,7 +183,7 @@ def test_without_data(self) -> None:
error_proto.ParseFromString(response.data)
assert response.status_code == 200, error_proto

def test_with_data_and_order_by(self, setup_teardown: Any) -> None:
def test_with_data(self, setup_teardown: Any) -> None:
ts = Timestamp(seconds=int(_BASE_TIME.timestamp()))
three_hours_later = int((_BASE_TIME + timedelta(hours=3)).timestamp())
message = GetTracesRequest(
Expand All @@ -201,11 +201,6 @@ def test_with_data_and_order_by(self, setup_teardown: Any) -> None:
key=TraceAttribute.Key.KEY_TRACE_ID,
),
],
order_by=[
GetTracesRequest.OrderBy(
key=TraceAttribute.Key.KEY_TRACE_ID,
),
],
)
response = EndpointGetTraces().execute(message)
expected_response = GetTracesResponse(
Expand All @@ -228,7 +223,7 @@ def test_with_data_and_order_by(self, setup_teardown: Any) -> None:
)
assert MessageToDict(response) == MessageToDict(expected_response)

def test_with_data_order_by_and_limit(self, setup_teardown: Any) -> None:
def test_with_data_and_limit(self, setup_teardown: Any) -> None:
ts = Timestamp(seconds=int(_BASE_TIME.timestamp()))
three_hours_later = int((_BASE_TIME + timedelta(hours=3)).timestamp())
message = GetTracesRequest(
Expand All @@ -246,11 +241,6 @@ def test_with_data_order_by_and_limit(self, setup_teardown: Any) -> None:
key=TraceAttribute.Key.KEY_TRACE_ID,
),
],
order_by=[
GetTracesRequest.OrderBy(
key=TraceAttribute.Key.KEY_TRACE_ID,
),
],
limit=1,
)
response = EndpointGetTraces().execute(message)
Expand Down Expand Up @@ -333,9 +323,7 @@ def test_with_data_and_filter(self, setup_teardown: Any) -> None:
)
assert MessageToDict(response) == MessageToDict(expected_response)

def test_with_data_order_by_and_aggregated_fields(
self, setup_teardown: Any
) -> None:
def test_with_data_and_aggregated_fields(self, setup_teardown: Any) -> None:
ts = Timestamp(seconds=int(_BASE_TIME.timestamp()))
three_hours_later = int((_BASE_TIME + timedelta(hours=3)).timestamp())
start_timestamp_per_trace_id: dict[str, float] = defaultdict(lambda: 2 * 1e10)
Expand All @@ -344,6 +332,10 @@ def test_with_data_order_by_and_aggregated_fields(
start_timestamp_per_trace_id[s["trace_id"]],
s["start_timestamp_precise"],
)
trace_id_per_start_timestamp: dict[float, str] = {
timestamp: trace_id
for trace_id, timestamp in start_timestamp_per_trace_id.items()
}
message = GetTracesRequest(
meta=RequestMeta(
project_ids=[1, 2, 3],
Expand Down Expand Up @@ -391,14 +383,6 @@ def test_with_data_order_by_and_aggregated_fields(
),
),
],
order_by=[
GetTracesRequest.OrderBy(
key=TraceAttribute.Key.KEY_TRACE_ID,
),
GetTracesRequest.OrderBy(
key=TraceAttribute.Key.KEY_START_TIMESTAMP,
),
],
)
response = EndpointGetTraces().execute(message)
expected_response = GetTracesResponse(
Expand All @@ -409,14 +393,16 @@ def test_with_data_order_by_and_aggregated_fields(
key=TraceAttribute.Key.KEY_TRACE_ID,
type=AttributeKey.TYPE_STRING,
value=AttributeValue(
val_str=trace_id,
val_str=trace_id_per_start_timestamp[start_timestamp],
),
),
TraceAttribute(
key=TraceAttribute.Key.KEY_START_TIMESTAMP,
type=AttributeKey.TYPE_FLOAT,
value=AttributeValue(
val_float=start_timestamp_per_trace_id[trace_id],
val_float=start_timestamp_per_trace_id[
trace_id_per_start_timestamp[start_timestamp]
],
),
),
TraceAttribute(
Expand All @@ -442,9 +428,13 @@ def test_with_data_order_by_and_aggregated_fields(
),
],
)
for trace_id in sorted(_TRACE_IDS)
for start_timestamp in reversed(
sorted(trace_id_per_start_timestamp.keys())
)
],
page_token=PageToken(offset=len(_TRACE_IDS)),
meta=ResponseMeta(request_id=_REQUEST_ID),
)
for start_timestamp in reversed(sorted(trace_id_per_start_timestamp.keys())):
print(start_timestamp, trace_id_per_start_timestamp[start_timestamp])
assert MessageToDict(response) == MessageToDict(expected_response)

0 comments on commit 01651ca

Please sign in to comment.