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

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,10 @@ 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_parsed: bool = False
+        self.load_config()
+        if not self.config_parsed:
+            raise Exception("Setup failed, dataset config was not parsed")
 
         self.start: datetime | None = None
         self.end: datetime | None = None
@@ -344,54 +332,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:
+        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
+        self.config_parsed = True
 
-        return field_alias_converter, function_converter, search_filter_converter, orderby_converter
+    def load_config(self) -> None:
+        self.config = self.config_class(self)
+        self.parse_config()
 
     def resolve_limit(self, limit: int | None) -> Limit | None:
         return None if limit is None else Limit(limit)
@@ -1547,6 +1497,27 @@ def get_row(row: dict[str, Any]) -> dict[str, Any]:
 class QueryBuilder(BaseQueryBuilder):
     """Builds a discover query"""
 
+    def load_config(
+        self,
+    ) -> None:
+        # Necessary until more classes inherit from BaseQueryBuilder instead
+        if hasattr(self, "config_class") and self.config_class is not None:
+            super().load_config()
+            return
+
+        self.config: DatasetConfig
+        if self.dataset in [
+            Dataset.Discover,
+            Dataset.Transactions,
+            Dataset.Events,
+            Dataset.IssuePlatform,
+        ]:
+            self.config = DiscoverDatasetConfig(self)
+        else:
+            raise NotImplementedError(f"Data Set configuration not found for {self.dataset}.")
+
+        self.parse_config()
+
     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,8 @@
     remove_hours,
     remove_minutes,
 )
+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 (
@@ -77,7 +79,6 @@
 
 class MetricsQueryBuilder(QueryBuilder):
     requires_organization_condition = True
-    is_alerts_query = False
 
     organization_column: str = "organization_id"
 
@@ -141,6 +142,21 @@ 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):
+        if hasattr(self, "config_class") and self.config_class is not None:
+            super().load_config()
+            return
+
+        if self.dataset in [Dataset.Metrics, Dataset.PerformanceMetrics]:
+            if self.builder_config.use_metrics_layer:
+                self.config = MetricsLayerDatasetConfig(self)
+            else:
+                self.config = MetricsDatasetConfig(self)
+        else:
+            raise NotImplementedError(f"Data Set configuration not found for {self.dataset}.")
+
+        self.parse_config()
+
     @property
     def use_default_tags(self) -> bool:
         if self._use_default_tags is None:
@@ -345,7 +361,7 @@ def _get_metrics_query_from_on_demand_spec(
             offset=self.offset,
             granularity=self.granularity,
             interval=interval,
-            is_alerts_query=True,
+            is_alerts_query=isinstance(self, AlertMetricsQueryBuilder),
             org_id=self.params.organization.id,
             project_ids=[p.id for p in self.params.projects],
             include_series=include_series,
@@ -1200,7 +1216,7 @@ def run_query(self, referrer: str, use_cache: bool = False) -> Any:
                                     self.get_metrics_layer_snql_query(
                                         query_details, extra_conditions
                                     ),
-                                    self.is_alerts_query,
+                                    isinstance(self, AlertMetricsQueryBuilder),
                                 )
                             )
                     metrics_data = []
@@ -1407,8 +1423,6 @@ def _get_metric_prefix(
 
 
 class AlertMetricsQueryBuilder(MetricsQueryBuilder):
-    is_alerts_query = True
-
     def __init__(
         self,
         *args: Any,
@@ -1449,7 +1463,7 @@ def get_snql_query(self) -> Request:
             else:
                 intermediate_query = self.get_metrics_layer_snql_query()
                 metrics_query = transform_mqb_query_to_metrics_query(
-                    intermediate_query, is_alerts_query=self.is_alerts_query
+                    intermediate_query, is_alerts_query=isinstance(self, AlertMetricsQueryBuilder)
                 )
 
             snuba_queries, _ = SnubaQueryBuilder(
@@ -1710,7 +1724,9 @@ def run_query(self, referrer: str, use_cache: bool = False) -> Any:
                     elif self.use_metrics_layer:
                         snuba_query = self.get_snql_query()[0].query
                         metrics_queries.append(
-                            transform_mqb_query_to_metrics_query(snuba_query, self.is_alerts_query)
+                            transform_mqb_query_to_metrics_query(
+                                snuba_query, isinstance(self, AlertMetricsQueryBuilder)
+                            )
                         )
                 metrics_data = []
                 with sentry_sdk.start_span(op="metric_layer", description="run_query"):
@@ -1981,7 +1997,9 @@ def run_query(self, referrer: str, use_cache: bool = False) -> Any:
                     elif self.use_metrics_layer:
                         snuba_query = self.get_snql_query()[0].query
                         metrics_queries.append(
-                            transform_mqb_query_to_metrics_query(snuba_query, self.is_alerts_query)
+                            transform_mqb_query_to_metrics_query(
+                                snuba_query, isinstance(self, AlertMetricsQueryBuilder)
+                            )
                         )
                 metrics_data = []
                 for metrics_query in metrics_queries:

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

@@ -6,6 +6,7 @@
 from sentry.api.event_search import SearchFilter, SearchKey, SearchValue
 from sentry.discover.arithmetic import categorize_columns
 from sentry.search.events.builder import QueryBuilder, TimeseriesQueryBuilder
+from sentry.search.events.datasets.base import DatasetConfig
 from sentry.search.events.datasets.profile_functions import ProfileFunctionsDatasetConfig
 from sentry.search.events.fields import custom_time_processor, get_function_alias
 from sentry.search.events.types import (
@@ -19,6 +20,8 @@
 
 
 class ProfileFunctionsQueryBuilderProtocol(Protocol):
+    config_class: type[DatasetConfig] = ProfileFunctionsDatasetConfig
+
     @property
     def config(self) -> ProfileFunctionsDatasetConfig:
         ...
@@ -45,12 +48,14 @@ def get_field_type(self: ProfileFunctionsQueryBuilderProtocol, field: str) -> st
 
 class ProfileFunctionsQueryBuilder(ProfileFunctionsQueryBuilderMixin, QueryBuilder):
     function_alias_prefix = "sentry_"
+    config_class: type[DatasetConfig] = ProfileFunctionsDatasetConfig
 
 
 class ProfileFunctionsTimeseriesQueryBuilder(
     ProfileFunctionsQueryBuilderMixin, TimeseriesQueryBuilder
 ):
     function_alias_prefix = "sentry_"
+    config_class: type[DatasetConfig] = ProfileFunctionsDatasetConfig
 
     def strip_alias_prefix(self, result):
         alias_mappings = {
@@ -75,6 +80,8 @@ def time_column(self) -> SelectType:
 
 
 class ProfileTopFunctionsTimeseriesQueryBuilder(ProfileFunctionsTimeseriesQueryBuilder):
+    config_class: type[DatasetConfig] = 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/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/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: