Skip to content

Commit

Permalink
ok this method is jank but its still functional
Browse files Browse the repository at this point in the history
  • Loading branch information
Rachel Chen authored and Rachel Chen committed Jan 16, 2025
1 parent ca2cf30 commit ecc5a6b
Show file tree
Hide file tree
Showing 3 changed files with 59 additions and 6 deletions.
6 changes: 6 additions & 0 deletions snuba/web/rpc/common/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -294,6 +294,8 @@ def trace_item_filters_to_expression(item_filter: TraceItemFilter) -> Expression
v_expression = literal(v.val_str)
case "val_float":
v_expression = literal(v.val_float)
case "val_double":
v_expression = literal(v.val_double)
case "val_int":
v_expression = literal(v.val_int)
case "val_null":
Expand All @@ -310,6 +312,10 @@ def trace_item_filters_to_expression(item_filter: TraceItemFilter) -> Expression
v_expression = literals_array(
None, list(map(lambda x: literal(x), v.val_float_array.values))
)
case "val_double_array":
v_expression = literals_array(
None, list(map(lambda x: literal(x), v.val_double_array.values))
)
case default:
raise NotImplementedError(
f"translation of AttributeValue type {default} is not implemented"
Expand Down
54 changes: 49 additions & 5 deletions snuba/web/rpc/v1/endpoint_get_traces.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,13 @@
_DEFAULT_ROW_LIMIT = 10_000
_BUFFER_WINDOW = 2 * 3600 # 2 hours


def _convert_key_to_support_doubles_and_floats_for_backward_compat(
key: TraceAttribute.Key.ValueType,
) -> TraceAttribute.Key.ValueType:
return TraceAttribute.Key.ValueType(-1 * key)


_ATTRIBUTES: dict[
TraceAttribute.Key.ValueType,
tuple[str, AttributeKey.Type.ValueType],
Expand All @@ -69,6 +76,16 @@
AttributeKey.Type.TYPE_STRING,
),
}
# for every AttributeKey of TYPE_FLOAT a user may add during the backward compat period, this adds the TYPE_DOUBLE equivalent
_attributes_backward_compat = dict()
for k in _ATTRIBUTES:
v = _ATTRIBUTES[k]
if v[1] == AttributeKey.Type.TYPE_FLOAT:
_attributes_backward_compat[
_convert_key_to_support_doubles_and_floats_for_backward_compat(k)
] = (v[0], AttributeKey.Type.TYPE_DOUBLE)
_ATTRIBUTES.update(_attributes_backward_compat)

_TYPES_TO_CLICKHOUSE: dict[
AttributeKey.Type.ValueType,
tuple[str, Callable[[Any], AttributeValue]],
Expand All @@ -85,6 +102,10 @@
"Float64",
lambda x: AttributeValue(val_float=float(x)),
),
AttributeKey.Type.TYPE_DOUBLE: (
"Float64",
lambda x: AttributeValue(val_double=float(x)),
),
}


Expand All @@ -102,11 +123,19 @@ def _attribute_to_expression(
alias=_ATTRIBUTES[trace_attribute.key][0],
)
if trace_attribute.key == TraceAttribute.Key.KEY_START_TIMESTAMP:
attribute = _ATTRIBUTES[trace_attribute.key]
attribute = (
_ATTRIBUTES[
_convert_key_to_support_doubles_and_floats_for_backward_compat(
trace_attribute.key
)
]
if trace_attribute.type == AttributeKey.Type.TYPE_DOUBLE
else _ATTRIBUTES[trace_attribute.key]
)
return f.cast(
f.min(column("start_timestamp")),
_TYPES_TO_CLICKHOUSE[attribute[1]][0],
alias=_ATTRIBUTES[trace_attribute.key][0],
alias=attribute[0],
)
if trace_attribute.key == TraceAttribute.Key.KEY_ROOT_SPAN_NAME:
# TODO: Change to return the root span name instead of the trace's first span's name.
Expand All @@ -116,7 +145,15 @@ def _attribute_to_expression(
alias=_ATTRIBUTES[trace_attribute.key][0],
)
if trace_attribute.key in _ATTRIBUTES:
attribute = _ATTRIBUTES[trace_attribute.key]
attribute = (
_ATTRIBUTES[
_convert_key_to_support_doubles_and_floats_for_backward_compat(
trace_attribute.key
)
]
if trace_attribute.type == AttributeKey.Type.TYPE_DOUBLE
else _ATTRIBUTES[trace_attribute.key]
)
return f.cast(
column(attribute[0]),
_TYPES_TO_CLICKHOUSE[attribute[1]][0],
Expand Down Expand Up @@ -165,8 +202,15 @@ def _convert_results(
TraceAttribute,
] = defaultdict(TraceAttribute)
for attribute in request.attributes:
value = row[_ATTRIBUTES[attribute.key][0]]
type = _ATTRIBUTES[attribute.key][1]
backward_compat_attribute_key = (
_convert_key_to_support_doubles_and_floats_for_backward_compat(
attribute.key
)
if attribute.type == AttributeKey.Type.TYPE_DOUBLE
else attribute.key
)
value = row[_ATTRIBUTES[backward_compat_attribute_key][0]]
type = _ATTRIBUTES[backward_compat_attribute_key][1]
values[attribute.key] = TraceAttribute(
key=attribute.key,
value=_TYPES_TO_CLICKHOUSE[type][1](value),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -196,7 +196,10 @@ def _convert_results(
elif column.key.type == AttributeKey.TYPE_DOUBLE:
converters[column.label] = lambda x: AttributeValue(val_double=float(x))
elif column.HasField("aggregation"):
converters[column.label] = lambda x: AttributeValue(val_float=float(x))
if column.key.type == AttributeKey.TYPE_FLOAT:
converters[column.label] = lambda x: AttributeValue(val_float=float(x))
if column.key.type == AttributeKey.TYPE_DOUBLE:
converters[column.label] = lambda x: AttributeValue(val_double=float(x))
else:
raise BadSnubaRPCRequestException(
"column is neither an attribute or aggregation"
Expand Down

0 comments on commit ecc5a6b

Please sign in to comment.