ref(querybuilder): Clean up the way we load configs (#71710)

- Instead of a giant if statement for every dataset in the base query
builder, set the config in a load config function in each builder

src/sentry/search/events/builder/discover.py

@@ -51,17 +51,6 @@
 from sentry.search.events import filter as event_filter
 from sentry.search.events.datasets.base import DatasetConfig
 from sentry.search.events.datasets.discover import DiscoverDatasetConfig
-from sentry.search.events.datasets.metrics import MetricsDatasetConfig
-from sentry.search.events.datasets.metrics_layer import MetricsLayerDatasetConfig
-from sentry.search.events.datasets.metrics_summaries import MetricsSummariesDatasetConfig
-from sentry.search.events.datasets.profile_functions import ProfileFunctionsDatasetConfig
-from sentry.search.events.datasets.profile_functions_metrics import (
-    ProfileFunctionsMetricsDatasetConfig,
-)
-from sentry.search.events.datasets.profiles import ProfilesDatasetConfig
-from sentry.search.events.datasets.sessions import SessionsDatasetConfig
-from sentry.search.events.datasets.spans_indexed import SpansIndexedDatasetConfig
-from sentry.search.events.datasets.spans_metrics import SpansMetricsDatasetConfig
 from sentry.search.events.types import (
     EventsResponse,
     HistogramParams,
@@ -99,6 +88,7 @@ class BaseQueryBuilder:
     spans_metrics_builder = False
     profile_functions_metrics_builder = False
     entity: Entity | None = None
+    config_class: type[DatasetConfig] | None
 
     def get_middle(self):
         """Get the middle for comparison functions"""
@@ -274,12 +264,8 @@ def __init__(
         self.turbo = turbo
         self.sample_rate = sample_rate
 
-        (
-            self.field_alias_converter,
-            self.function_converter,
-            self.search_filter_converter,
-            self.orderby_converter,
-        ) = self.load_config()
+        self.config = self.load_config()
+        self.parse_config()
 
         self.start: datetime | None = None
         self.end: datetime | None = None
@@ -344,54 +330,16 @@ def resolve_query(
             with sentry_sdk.start_span(op="QueryBuilder", description="resolve_groupby"):
                 self.groupby = self.resolve_groupby(groupby_columns)
 
-    def load_config(
-        self,
-    ) -> tuple[
-        Mapping[str, Callable[[str], SelectType]],
-        Mapping[str, fields.SnQLFunction],
-        Mapping[str, Callable[[event_search.SearchFilter], WhereType | None]],
-        Mapping[str, Callable[[Direction], OrderBy]],
-    ]:
-        self.config: DatasetConfig
-        if self.dataset in [
-            Dataset.Discover,
-            Dataset.Transactions,
-            Dataset.Events,
-            Dataset.IssuePlatform,
-        ]:
-            self.config = DiscoverDatasetConfig(self)
-        elif self.dataset == Dataset.Sessions:
-            self.config = SessionsDatasetConfig(self)
-        elif self.dataset in [Dataset.Metrics, Dataset.PerformanceMetrics]:
-            if self.spans_metrics_builder:
-                # For now, we won't support the metrics layer for spans since it needs some work,
-                # but once the work will be done, we will have to add:
-                # if self.builder_config.use_metrics_layer:
-                #     self.config = SpansMetricsLayerDatasetConfig(self)
-                self.config = SpansMetricsDatasetConfig(self)
-            elif self.profile_functions_metrics_builder:
-                self.config = ProfileFunctionsMetricsDatasetConfig(self)
-            elif self.builder_config.use_metrics_layer:
-                self.config = MetricsLayerDatasetConfig(self)
-            else:
-                self.config = MetricsDatasetConfig(self)
-        elif self.dataset == Dataset.Profiles:
-            self.config = ProfilesDatasetConfig(self)
-        elif self.dataset == Dataset.Functions:
-            self.config = ProfileFunctionsDatasetConfig(self)
-        elif self.dataset == Dataset.SpansIndexed:
-            self.config = SpansIndexedDatasetConfig(self)
-        elif self.dataset == Dataset.MetricsSummaries:
-            self.config = MetricsSummariesDatasetConfig(self)
-        else:
-            raise NotImplementedError(f"Data Set configuration not found for {self.dataset}.")
-
-        field_alias_converter = self.config.field_alias_converter
-        function_converter = self.config.function_converter
-        search_filter_converter = self.config.search_filter_converter
-        orderby_converter = self.config.orderby_converter
+    def parse_config(self) -> None:
+        if not hasattr(self, "config") or self.config is None:
+            raise Exception("Setup failed, dataset config was not loaded")
+        self.field_alias_converter = self.config.field_alias_converter
+        self.function_converter = self.config.function_converter
+        self.search_filter_converter = self.config.search_filter_converter
+        self.orderby_converter = self.config.orderby_converter
 
-        return field_alias_converter, function_converter, search_filter_converter, orderby_converter
+    def load_config(self) -> DatasetConfig:
+        return self.config_class(self)
 
     def resolve_limit(self, limit: int | None) -> Limit | None:
         return None if limit is None else Limit(limit)
@@ -1549,6 +1497,24 @@ def get_row(row: dict[str, Any]) -> dict[str, Any]:
 class QueryBuilder(BaseQueryBuilder):
     """Builds a discover query"""
 
+    def load_config(
+        self,
+    ) -> DatasetConfig:
+        # Necessary until more classes inherit from BaseQueryBuilder instead
+        if hasattr(self, "config_class") and self.config_class is not None:
+            return super().load_config()
+
+        self.config: DatasetConfig
+        if self.dataset in [
+            Dataset.Discover,
+            Dataset.Transactions,
+            Dataset.Events,
+            Dataset.IssuePlatform,
+        ]:
+            return DiscoverDatasetConfig(self)
+        else:
+            raise NotImplementedError(f"Data Set configuration not found for {self.dataset}.")
+
     def resolve_field(self, raw_field: str, alias: bool = False) -> Column:
         tag_match = constants.TAG_KEY_RE.search(raw_field)
         field = tag_match.group("tag") if tag_match else raw_field

src/sentry/search/events/builder/metrics.py

@@ -40,6 +40,9 @@
     remove_hours,
     remove_minutes,
 )
+from sentry.search.events.datasets.base import DatasetConfig
+from sentry.search.events.datasets.metrics import MetricsDatasetConfig
+from sentry.search.events.datasets.metrics_layer import MetricsLayerDatasetConfig
 from sentry.search.events.fields import get_function_alias
 from sentry.search.events.filter import ParsedTerms
 from sentry.search.events.types import (
@@ -141,6 +144,18 @@ def __init__(
         sentry_sdk.set_tag("on_demand_metrics.enabled", config.on_demand_metrics_enabled)
         self.organization_id: int = org_id
 
+    def load_config(self) -> DatasetConfig:
+        if hasattr(self, "config_class") and self.config_class is not None:
+            return super().load_config()
+
+        if self.dataset in [Dataset.Metrics, Dataset.PerformanceMetrics]:
+            if self.builder_config.use_metrics_layer:
+                return MetricsLayerDatasetConfig(self)
+            else:
+                return MetricsDatasetConfig(self)
+        else:
+            raise NotImplementedError(f"Data Set configuration not found for {self.dataset}.")
+
     @property
     def use_default_tags(self) -> bool:
         if self._use_default_tags is None:

src/sentry/search/events/builder/metrics_summaries.py

@@ -1,11 +1,13 @@
 from snuba_sdk import Entity, Flags, Query, Request
 
 from sentry.search.events.builder import QueryBuilder
+from sentry.search.events.datasets.metrics_summaries import MetricsSummariesDatasetConfig
 from sentry.snuba.dataset import Dataset
 
 
 class MetricsSummariesQueryBuilder(QueryBuilder):
     requires_organization_condition = False
+    config_class = MetricsSummariesDatasetConfig
 
     def get_field_type(self, field: str) -> str | None:
         if field in ["min_metric", "max_metric", "sum_metric", "count_metric"]:

src/sentry/search/events/builder/profile_functions.py

@@ -45,12 +45,14 @@ def get_field_type(self: ProfileFunctionsQueryBuilderProtocol, field: str) -> st
 
 class ProfileFunctionsQueryBuilder(ProfileFunctionsQueryBuilderMixin, QueryBuilder):
     function_alias_prefix = "sentry_"
+    config_class = ProfileFunctionsDatasetConfig
 
 
 class ProfileFunctionsTimeseriesQueryBuilder(
     ProfileFunctionsQueryBuilderMixin, TimeseriesQueryBuilder
 ):
     function_alias_prefix = "sentry_"
+    config_class = ProfileFunctionsDatasetConfig
 
     def strip_alias_prefix(self, result):
         alias_mappings = {
@@ -75,6 +77,8 @@ def time_column(self) -> SelectType:
 
 
 class ProfileTopFunctionsTimeseriesQueryBuilder(ProfileFunctionsTimeseriesQueryBuilder):
+    config_class = ProfileFunctionsDatasetConfig
+
     def __init__(
         self,
         dataset: Dataset,

src/sentry/search/events/builder/profile_functions_metrics.py

@@ -3,12 +3,16 @@
     TimeseriesMetricQueryBuilder,
     TopMetricsQueryBuilder,
 )
+from sentry.search.events.datasets.profile_functions_metrics import (
+    ProfileFunctionsMetricsDatasetConfig,
+)
 from sentry.search.events.types import SelectType
 
 
 class ProfileFunctionsMetricsQueryBuilder(MetricsQueryBuilder):
     requires_organization_condition = True
     profile_functions_metrics_builder = True
+    config_class = ProfileFunctionsMetricsDatasetConfig
 
     column_remapping = {
         # We want to remap `message` to `name` for the free

src/sentry/search/events/builder/profiles.py

@@ -36,8 +36,8 @@ def get_field_type(self: ProfilesQueryBuilderProtocol, field: str) -> str | None
 
 
 class ProfilesQueryBuilder(ProfilesQueryBuilderMixin, QueryBuilder):
-    pass
+    config_class = ProfilesDatasetConfig
 
 
 class ProfilesTimeseriesQueryBuilder(ProfilesQueryBuilderMixin, TimeseriesQueryBuilder):
-    pass
+    config_class = ProfilesDatasetConfig

src/sentry/search/events/builder/sessions.py

@@ -6,13 +6,15 @@
 from sentry.api.event_search import SearchFilter
 from sentry.exceptions import InvalidSearchQuery
 from sentry.search.events.builder import QueryBuilder
+from sentry.search.events.datasets.sessions import SessionsDatasetConfig
 from sentry.search.events.types import SelectType, WhereType
 
 
 class SessionsV2QueryBuilder(QueryBuilder):
     filter_allowlist_fields = {"project", "project_id", "environment", "release"}
     requires_organization_condition = True
     organization_column: str = "org_id"
+    config_class = SessionsDatasetConfig
 
     def __init__(
         self,

src/sentry/search/events/builder/spans_indexed.py

@@ -3,6 +3,7 @@
 
 from sentry.search.events import constants
 from sentry.search.events.builder import QueryBuilder, TimeseriesQueryBuilder, TopEventsQueryBuilder
+from sentry.search.events.datasets.spans_indexed import SpansIndexedDatasetConfig
 from sentry.search.events.fields import custom_time_processor
 from sentry.search.events.types import SelectType
 
@@ -22,6 +23,7 @@ def get_field_type(self, field: str) -> str | None:
 class SpansIndexedQueryBuilder(SpansIndexedQueryBuilderMixin, QueryBuilder):
     requires_organization_condition = False
     uuid_fields = {"transaction.id", "replay.id", "profile.id", "trace"}
+    config_class = SpansIndexedDatasetConfig
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -31,6 +33,8 @@ def __init__(self, *args, **kwargs):
 
 
 class TimeseriesSpanIndexedQueryBuilder(SpansIndexedQueryBuilderMixin, TimeseriesQueryBuilder):
+    config_class = SpansIndexedDatasetConfig
+
     @property
     def time_column(self) -> SelectType:
         return custom_time_processor(
@@ -39,6 +43,8 @@ def time_column(self) -> SelectType:
 
 
 class TopEventsSpanIndexedQueryBuilder(SpansIndexedQueryBuilderMixin, TopEventsQueryBuilder):
+    config_class = SpansIndexedDatasetConfig
+
     @property
     def time_column(self) -> SelectType:
         return custom_time_processor(

src/sentry/search/events/builder/spans_metrics.py

@@ -6,13 +6,15 @@
     TimeseriesMetricQueryBuilder,
     TopMetricsQueryBuilder,
 )
+from sentry.search.events.datasets.spans_metrics import SpansMetricsDatasetConfig
 from sentry.search.events.types import SelectType
 
 
 class SpansMetricsQueryBuilder(MetricsQueryBuilder):
     requires_organization_condition = True
     spans_metrics_builder = True
     has_transaction = False
+    config_class = SpansMetricsDatasetConfig
 
     column_remapping = {
         # We want to remap `message` to `span.description` for the free

src/sentry/snuba/metrics/query_builder.py

@@ -37,6 +37,7 @@
 from sentry.exceptions import InvalidParams, InvalidSearchQuery
 from sentry.models.project import Project
 from sentry.search.events.builder import UnresolvedQuery
+from sentry.search.events.datasets.sessions import SessionsDatasetConfig
 from sentry.search.events.types import QueryBuilderConfig, WhereType
 from sentry.sentry_metrics.use_case_id_registry import UseCaseID
 from sentry.sentry_metrics.utils import (
@@ -469,6 +470,8 @@ def parse_conditions(
 
 
 class ReleaseHealthQueryBuilder(UnresolvedQuery):
+    config_class = SessionsDatasetConfig
+
     def _contains_wildcard_in_query(self, query: str | None) -> bool:
         parsed_terms = self.parse_query(query)
         for parsed_term in parsed_terms: