Skip to content

Commit 75fec36

Browse files
ZylphrexZylphrex
authored
feat(eap): Support more search types in eap aggregations (#83129)
Since these additional size and duration units are valid search types, we should allow them in aggregations.
1 parent 06a3f1e commit 75fec36

File tree

1 file changed

+22
-17
lines changed

1 file changed

+22
-17
lines changed

‎src/sentry/search/eap/columns.py

+22-17
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515
from sentry.exceptions import InvalidSearchQuery
1616
from sentry.search.eap import constants
17-
from sentry.search.events.constants import SPAN_MODULE_CATEGORY_VALUES
17+
from sentry.search.events.constants import DURATION_UNITS, SIZE_UNITS, SPAN_MODULE_CATEGORY_VALUES
1818
from sentry.search.events.types import SnubaParams
1919
from sentry.search.utils import DEVICE_CLASS
2020
from sentry.utils.validators import is_event_id, is_span_id
@@ -74,9 +74,14 @@ def proto_definition(self) -> AttributeKey:
7474
)
7575

7676

77+
SIZE_TYPE: set[constants.SearchType] = set(SIZE_UNITS.keys())
78+
79+
DURATION_TYPE: set[constants.SearchType] = set(DURATION_UNITS.keys())
80+
81+
7782
@dataclass
7883
class ArgumentDefinition:
79-
argument_types: list[constants.SearchType] | None = None
84+
argument_types: set[constants.SearchType] | None = None
8085
# The public alias for the default arg, the SearchResolver will resolve this value
8186
default_arg: str | None = None
8287
# Whether this argument is completely ignored, used for `count()`
@@ -505,7 +510,7 @@ def module_context_constructor(params: SnubaParams) -> VirtualColumnContext:
505510
default_search_type="duration",
506511
arguments=[
507512
ArgumentDefinition(
508-
argument_types=["byte", "duration", "millisecond", "number"],
513+
argument_types={"duration", "number", *SIZE_TYPE, *DURATION_TYPE},
509514
default_arg="span.duration",
510515
)
511516
],
@@ -515,7 +520,7 @@ def module_context_constructor(params: SnubaParams) -> VirtualColumnContext:
515520
default_search_type="duration",
516521
arguments=[
517522
ArgumentDefinition(
518-
argument_types=["byte", "duration", "millisecond", "number", "percentage"],
523+
argument_types={"duration", "number", "percentage", *SIZE_TYPE, *DURATION_TYPE},
519524
default_arg="span.duration",
520525
)
521526
],
@@ -525,7 +530,7 @@ def module_context_constructor(params: SnubaParams) -> VirtualColumnContext:
525530
default_search_type="duration",
526531
arguments=[
527532
ArgumentDefinition(
528-
argument_types=["byte", "duration", "millisecond", "number", "percentage"],
533+
argument_types={"duration", "number", "percentage", *SIZE_TYPE, *DURATION_TYPE},
529534
default_arg="span.duration",
530535
)
531536
],
@@ -537,7 +542,7 @@ def module_context_constructor(params: SnubaParams) -> VirtualColumnContext:
537542
default_search_type="integer",
538543
arguments=[
539544
ArgumentDefinition(
540-
argument_types=["byte", "duration", "millisecond", "number"],
545+
argument_types={"duration", "number", *SIZE_TYPE, *DURATION_TYPE},
541546
default_arg="span.duration",
542547
)
543548
],
@@ -548,7 +553,7 @@ def module_context_constructor(params: SnubaParams) -> VirtualColumnContext:
548553
default_search_type="integer",
549554
arguments=[
550555
ArgumentDefinition(
551-
argument_types=["byte", "duration", "millisecond", "number"],
556+
argument_types={"duration", "number", *SIZE_TYPE, *DURATION_TYPE},
552557
default_arg="span.duration",
553558
)
554559
],
@@ -559,7 +564,7 @@ def module_context_constructor(params: SnubaParams) -> VirtualColumnContext:
559564
default_search_type="duration",
560565
arguments=[
561566
ArgumentDefinition(
562-
argument_types=["byte", "duration", "millisecond", "number"],
567+
argument_types={"duration", "number", *SIZE_TYPE, *DURATION_TYPE},
563568
default_arg="span.duration",
564569
)
565570
],
@@ -569,7 +574,7 @@ def module_context_constructor(params: SnubaParams) -> VirtualColumnContext:
569574
default_search_type="duration",
570575
arguments=[
571576
ArgumentDefinition(
572-
argument_types=["byte", "duration", "millisecond", "number"],
577+
argument_types={"duration", "number", *SIZE_TYPE, *DURATION_TYPE},
573578
default_arg="span.duration",
574579
)
575580
],
@@ -580,7 +585,7 @@ def module_context_constructor(params: SnubaParams) -> VirtualColumnContext:
580585
default_search_type="duration",
581586
arguments=[
582587
ArgumentDefinition(
583-
argument_types=["byte", "duration", "millisecond", "number"],
588+
argument_types={"duration", "number", *SIZE_TYPE, *DURATION_TYPE},
584589
default_arg="span.duration",
585590
)
586591
],
@@ -590,7 +595,7 @@ def module_context_constructor(params: SnubaParams) -> VirtualColumnContext:
590595
default_search_type="duration",
591596
arguments=[
592597
ArgumentDefinition(
593-
argument_types=["byte", "duration", "millisecond", "number"],
598+
argument_types={"duration", "number", *SIZE_TYPE, *DURATION_TYPE},
594599
default_arg="span.duration",
595600
)
596601
],
@@ -600,7 +605,7 @@ def module_context_constructor(params: SnubaParams) -> VirtualColumnContext:
600605
default_search_type="duration",
601606
arguments=[
602607
ArgumentDefinition(
603-
argument_types=["byte", "duration", "millisecond", "number"],
608+
argument_types={"duration", "number", *SIZE_TYPE, *DURATION_TYPE},
604609
default_arg="span.duration",
605610
)
606611
],
@@ -610,7 +615,7 @@ def module_context_constructor(params: SnubaParams) -> VirtualColumnContext:
610615
default_search_type="duration",
611616
arguments=[
612617
ArgumentDefinition(
613-
argument_types=["byte", "duration", "millisecond", "number"],
618+
argument_types={"duration", "number", *SIZE_TYPE, *DURATION_TYPE},
614619
default_arg="span.duration",
615620
)
616621
],
@@ -620,7 +625,7 @@ def module_context_constructor(params: SnubaParams) -> VirtualColumnContext:
620625
default_search_type="duration",
621626
arguments=[
622627
ArgumentDefinition(
623-
argument_types=["byte", "duration", "millisecond", "number"],
628+
argument_types={"duration", "number", *SIZE_TYPE, *DURATION_TYPE},
624629
default_arg="span.duration",
625630
)
626631
],
@@ -630,7 +635,7 @@ def module_context_constructor(params: SnubaParams) -> VirtualColumnContext:
630635
default_search_type="duration",
631636
arguments=[
632637
ArgumentDefinition(
633-
argument_types=["byte", "duration", "millisecond", "number", "percentage"],
638+
argument_types={"duration", "number", "percentage", *SIZE_TYPE, *DURATION_TYPE},
634639
default_arg="span.duration",
635640
)
636641
],
@@ -640,7 +645,7 @@ def module_context_constructor(params: SnubaParams) -> VirtualColumnContext:
640645
default_search_type="duration",
641646
arguments=[
642647
ArgumentDefinition(
643-
argument_types=["byte", "duration", "millisecond", "number", "percentage"],
648+
argument_types={"duration", "number", "percentage", *SIZE_TYPE, *DURATION_TYPE},
644649
default_arg="span.duration",
645650
)
646651
],
@@ -650,7 +655,7 @@ def module_context_constructor(params: SnubaParams) -> VirtualColumnContext:
650655
default_search_type="number",
651656
arguments=[
652657
ArgumentDefinition(
653-
argument_types=["string"],
658+
argument_types={"string"},
654659
)
655660
],
656661
),

0 commit comments

Comments
 (0)