fix(eap-api): return normalized fields for GetTraceEndpoint (#6856)

Currently, if the user doesn't specify attributes to the
GetTraceEndpoint, we won't return normalized fields; however, we should

getsentry/eap-planning#187

---------

Co-authored-by: Rachel Chen <[email protected]>
Co-authored-by: Rachel Chen <[email protected]>

Author: 3 people

snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_get_trace.py

@@ -1,4 +1,5 @@
 import uuid
+from datetime import datetime
 from operator import attrgetter
 from typing import Any, Dict, Iterable
 
@@ -41,6 +42,21 @@
 _BUCKET_COUNT = 20
 
 
+NORMALIZED_COLUMNS_TO_INCLUDE = [
+    col.name
+    for col in get_entity(EntityKey("eap_spans")).get_data_model().columns
+    if col.name
+    not in [
+        "retention_days",
+        "sign",
+        "attr_str",
+        "attr_num",
+        "span_id",
+        "timestamp",
+    ]
+]
+
+
 def _build_query(request: GetTraceRequest) -> Query:
     selected_columns: list[SelectedExpression] = [
         SelectedExpression(
@@ -86,6 +102,14 @@ def _build_query(request: GetTraceRequest) -> Query:
                 ),
             ),
         ]
+        selected_columns.extend(
+            [
+                SelectedExpression(
+                    name=col_name, expression=column(col_name, alias=col_name)
+                )
+                for col_name in NORMALIZED_COLUMNS_TO_INCLUDE
+            ]
+        )
 
     entity = Entity(
         key=EntityKey("eap_spans"),
@@ -178,6 +202,16 @@ def _value_to_attribute(key: str, value: Any) -> tuple[AttributeKey, AttributeVa
                 val_str=value,
             ),
         )
+    elif isinstance(value, datetime):
+        return (
+            AttributeKey(
+                name=key,
+                type=AttributeKey.Type.TYPE_DOUBLE,
+            ),
+            AttributeValue(
+                val_double=value.timestamp(),
+            ),
+        )
     else:
         raise BadSnubaRPCRequestException(f"data type unknown: {type(value)}")
 

tests/web/rpc/v1/test_endpoint_get_trace.py

@@ -22,6 +22,9 @@
 from snuba.datasets.storages.factory import get_storage
 from snuba.datasets.storages.storage_key import StorageKey
 from snuba.web.rpc.v1.endpoint_get_trace import EndpointGetTrace
+from snuba.web.rpc.v1.resolvers.R_eap_spans.resolver_get_trace import (
+    NORMALIZED_COLUMNS_TO_INCLUDE,
+)
 from tests.base import BaseApiTest
 from tests.helpers import write_raw_unprocessed_events
 
@@ -266,7 +269,21 @@ def test_with_data(self, setup_teardown: Any) -> None:
                 ),
             ],
         )
-        assert MessageToDict(response) == MessageToDict(expected_response)
+
+        assert list(
+            [
+                attribute
+                for attribute in response.item_groups[0].items[0].attributes
+                if attribute.key.name not in NORMALIZED_COLUMNS_TO_INCLUDE
+            ]
+        ) == list(expected_response.item_groups[0].items[0].attributes)
+        assert set(
+            [
+                attribute.key.name
+                for attribute in response.item_groups[0].items[0].attributes
+                if attribute.key.name in NORMALIZED_COLUMNS_TO_INCLUDE
+            ]
+        ) == set(NORMALIZED_COLUMNS_TO_INCLUDE)
 
     def test_with_specific_attributes(self, setup_teardown: Any) -> None:
         ts = Timestamp(seconds=int(_BASE_TIME.timestamp()))