Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement filter offset for attribute keys API #6618

Merged
merged 3 commits into from
Dec 11, 2024
+93 −34
Merged

Implement filter offset for attribute keys API #6618

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3216448
c
Dec 5, 2024
5e8ede7
filter offset is default to avoid scanning
Dec 6, 2024
dc7db20
unnecessary
Dec 9, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 58 additions & 10 deletions snuba/web/rpc/v1/endpoint_trace_item_attribute_names.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,11 @@
TraceItemAttributeNamesResponse,
)
from sentry_protos.snuba.v1.request_common_pb2 import PageToken
from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey
from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey, AttributeValue
from sentry_protos.snuba.v1.trace_item_filter_pb2 import (
ComparisonFilter,
TraceItemFilter,
)

from snuba.attribution.appid import AppID
from snuba.attribution.attribution_info import AttributionInfo
Expand All @@ -17,7 +21,8 @@
from snuba.query import OrderBy, OrderByDirection, SelectedExpression
from snuba.query.data_source.simple import Entity
from snuba.query.dsl import Functions as f
from snuba.query.dsl import column
from snuba.query.dsl import column, literal
from snuba.query.expressions import Expression
from snuba.query.logical import Query
from snuba.query.query_settings import HTTPQuerySettings
from snuba.reader import Row
Expand All @@ -33,6 +38,27 @@
MAX_REQUEST_LIMIT = 1000


def _convert_filter_offset(filter_offset: TraceItemFilter) -> Expression:
if not filter_offset.HasField("comparison_filter"):
raise TypeError("filter_offset needs to be a comparison filter")
if filter_offset.comparison_filter.op != ComparisonFilter.OP_GREATER_THAN:
raise TypeError("filter_offset must use the greater than comparison")

k_expression = column(filter_offset.comparison_filter.key.name)
v = filter_offset.comparison_filter.value
value_type = v.WhichOneof("value")
if value_type is None:
raise BadSnubaRPCRequestException("comparison does not have a right hand side")

v_expression = {
"val_bool": literal(v.val_bool),
"val_str": literal(v.val_str),
"val_float": literal(v.val_float),
"val_int": literal(v.val_int),
}[value_type]
return f.greater(k_expression, v_expression)


def convert_to_snuba_request(req: TraceItemAttributeNamesRequest) -> SnubaRequest:
if req.limit > MAX_REQUEST_LIMIT:
raise BadSnubaRPCRequestException(
Expand Down Expand Up @@ -60,6 +86,18 @@ def convert_to_snuba_request(req: TraceItemAttributeNamesRequest) -> SnubaReques
f"Attribute type '{req.type}' is not supported. Supported types are: TYPE_STRING, TYPE_FLOAT, TYPE_INT, TYPE_BOOLEAN"
)

condition = (
base_conditions_and(
req.meta,
f.like(column("attr_key"), f"%{req.value_substring_match}%"),
_convert_filter_offset(req.page_token.filter_offset),
)
if req.page_token.HasField("filter_offset")
else base_conditions_and(
req.meta, f.like(column("attr_key"), f"%{req.value_substring_match}%")
)
)

query = Query(
from_clause=entity,
selected_columns=[
Expand All @@ -68,9 +106,7 @@ def convert_to_snuba_request(req: TraceItemAttributeNamesRequest) -> SnubaReques
expression=column("attr_key"),
),
],
condition=base_conditions_and(
req.meta, f.like(column("attr_key"), f"%{req.value_substring_match}%")
),
condition=condition,
order_by=[
OrderBy(direction=OrderByDirection.ASC, expression=column("attr_key")),
],
Expand Down Expand Up @@ -138,9 +174,24 @@ def _build_response(
res: QueryResult,
) -> TraceItemAttributeNamesResponse:
attributes = convert_to_attributes(res, req.type)
page_token = (
PageToken(offset=req.page_token.offset + len(attributes))
if req.page_token.HasField("offset")
else PageToken(
filter_offset=TraceItemFilter(
comparison_filter=ComparisonFilter(
key=AttributeKey(
type=AttributeKey.TYPE_STRING, name="attr_key"
),
op=ComparisonFilter.OP_GREATER_THAN,
value=AttributeValue(val_str=attributes[-1].name),
)
)
)
)
return TraceItemAttributeNamesResponse(
attributes=attributes,
page_token=PageToken(offset=req.page_token.offset + len(attributes)),
page_token=page_token,
meta=extract_response_meta(
req.meta.request_id, req.meta.debug, [res], [self._timer]
),
Expand All @@ -151,10 +202,7 @@ def _execute(
) -> TraceItemAttributeNamesResponse:
if not req.HasField("page_token"):
req.page_token.offset = 0
if req.page_token.HasField("filter_offset"):
raise NotImplementedError(
"EndpointTraceItemAttributeNames does not currently support page_token.filter_offset, please use page_token.offset instead."
)

if not req.meta.request_id:
req.meta.request_id = str(uuid.uuid4())

Expand Down
72 changes: 51 additions & 21 deletions tests/web/rpc/v1/test_endpoint_trace_item_attribute_names.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,11 @@
TraceItemAttributeNamesResponse,
)
from sentry_protos.snuba.v1.request_common_pb2 import PageToken, RequestMeta
from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey
from sentry_protos.snuba.v1.trace_item_filter_pb2 import TraceItemFilter
from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey, AttributeValue
from sentry_protos.snuba.v1.trace_item_filter_pb2 import (
ComparisonFilter,
TraceItemFilter,
)

from snuba.datasets.storages.factory import get_storage
from snuba.datasets.storages.storage_key import StorageKey
Expand Down Expand Up @@ -212,26 +215,53 @@ def test_with_page_token(self) -> None:
assert expected_attributes == []

def test_page_token_offset_filter(self) -> None:
req = TraceItemAttributeNamesRequest(
meta=RequestMeta(
project_ids=[1, 2, 3],
organization_id=1,
cogs_category="something",
referrer="something",
start_timestamp=Timestamp(
seconds=int((BASE_TIME - timedelta(days=1)).timestamp())
),
end_timestamp=Timestamp(
seconds=int((BASE_TIME + timedelta(days=1)).timestamp())
),
),
limit=10,
type=AttributeKey.Type.TYPE_STRING,
value_substring_match="",
page_token=PageToken(filter_offset=TraceItemFilter()),

expected_attributes = []
for i in range(TOTAL_GENERATED_ATTR_PER_TYPE):
expected_attributes.append(
TraceItemAttributeNamesResponse.Attribute(
name=f"a_tag_{str(i).zfill(3)}",
type=AttributeKey.Type.TYPE_STRING,
)
)
# grab 10 at a time until we get them all
done = 0
page_token = PageToken(
filter_offset=TraceItemFilter(
comparison_filter=ComparisonFilter(
key=AttributeKey(type=AttributeKey.TYPE_STRING, name="attr_key"),
op=ComparisonFilter.OP_GREATER_THAN,
value=AttributeValue(val_str=""),
)
)
xurui-c marked this conversation as resolved.
Show resolved Hide resolved
)
with pytest.raises(NotImplementedError):
EndpointTraceItemAttributeNames().execute(req)
at_a_time = 10

while done < TOTAL_GENERATED_ATTR_PER_TYPE:
req = TraceItemAttributeNamesRequest(
meta=RequestMeta(
project_ids=[1, 2, 3],
organization_id=1,
cogs_category="something",
referrer="something",
start_timestamp=Timestamp(
seconds=int((BASE_TIME - timedelta(days=1)).timestamp())
),
end_timestamp=Timestamp(
seconds=int((BASE_TIME + timedelta(days=1)).timestamp())
),
),
limit=at_a_time,
type=AttributeKey.Type.TYPE_STRING,
value_substring_match="",
page_token=page_token,
)
res = EndpointTraceItemAttributeNames().execute(req)
page_token = res.page_token
assert res.attributes == expected_attributes[:at_a_time]
expected_attributes = expected_attributes[at_a_time:]
done += at_a_time
assert expected_attributes == []

def test_response_metadata(self) -> None:
# debug must be true in RequestMeta for it to return query_info in the response
Expand Down
Loading