Skip to content

Commit

Permalink
feat(apis): Pass query source to snuba - phase 2
Browse files Browse the repository at this point in the history
  • Loading branch information
sentaur-athena committed Jun 28, 2024
1 parent 54c9857 commit d3154bf
Show file tree
Hide file tree
Showing 17 changed files with 167 additions and 57 deletions.
10 changes: 10 additions & 0 deletions src/sentry/api/endpoints/organization_events_meta.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,9 @@
from sentry.api.serializers import serialize
from sentry.api.serializers.models.group import GroupSerializer
from sentry.api.utils import handle_query_errors
from sentry.middleware import is_frontend_request
from sentry.snuba import spans_indexed, spans_metrics
from sentry.snuba.query_sources import QuerySource
from sentry.snuba.referrer import Referrer


Expand All @@ -39,6 +41,10 @@ def get(self, request: Request, organization) -> Response:
params=params,
query=request.query_params.get("query"),
referrer=Referrer.API_ORGANIZATION_EVENTS_META.value,
# TODO: @athena - add query_source when all datasets support it
# query_source=(
# QuerySource.FRONTEND if is_frontend_request(request) else QuerySource.API
# ),
)

return Response({"count": result["data"][0]["count"]})
Expand Down Expand Up @@ -118,6 +124,7 @@ class OrganizationSpansSamplesEndpoint(OrganizationEventsEndpointBase):
snuba_methods = ["GET"]

def get(self, request: Request, organization) -> Response:
is_frontend = is_frontend_request(request)
try:
params = self.get_snuba_params(request, organization)
except NoProjects:
Expand Down Expand Up @@ -148,6 +155,7 @@ def get(self, request: Request, organization) -> Response:
params=params,
query=request.query_params.get("query"),
referrer=Referrer.API_SPAN_SAMPLE_GET_BOUNDS.value,
query_source=(QuerySource.FRONTEND if is_frontend else QuerySource.API),
)
if len(bound_results["data"]) != 1:
raise ParseError("Could not find bounds")
Expand All @@ -169,6 +177,7 @@ def get(self, request: Request, organization) -> Response:
params=params,
query=request.query_params.get("query"),
referrer=Referrer.API_SPAN_SAMPLE_GET_SPAN_IDS.value,
query_source=(QuerySource.FRONTEND if is_frontend else QuerySource.API),
)
span_ids = []
for row in result["data"]:
Expand All @@ -192,5 +201,6 @@ def get(self, request: Request, organization) -> Response:
query=query,
limit=9,
referrer=Referrer.API_SPAN_SAMPLE_GET_SPAN_DATA.value,
query_source=(QuerySource.FRONTEND if is_frontend else QuerySource.API),
)
return Response({"data": result["data"]})
18 changes: 18 additions & 0 deletions src/sentry/api/endpoints/organization_events_stats.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from sentry.api.base import region_silo_endpoint
from sentry.api.bases import OrganizationEventsV2EndpointBase
from sentry.constants import MAX_TOP_EVENTS
from sentry.middleware import is_frontend_request
from sentry.models.dashboard_widget import DashboardWidget, DashboardWidgetTypes
from sentry.models.organization import Organization
from sentry.snuba import (
Expand All @@ -23,6 +24,7 @@
spans_metrics,
)
from sentry.snuba.metrics.extraction import MetricSpecType
from sentry.snuba.query_sources import QuerySource
from sentry.snuba.referrer import Referrer
from sentry.utils.snuba import SnubaError, SnubaTSResult

Expand Down Expand Up @@ -73,6 +75,9 @@
ALLOWED_EVENTS_STATS_REFERRERS: set[str] = {
Referrer.API_ALERTS_ALERT_RULE_CHART.value,
Referrer.API_ALERTS_CHARTCUTERIE.value,
Referrer.API_ENDPOINT_REGRESSION_ALERT_CHARTCUTERIE.value,
Referrer.API_FUNCTION_REGRESSION_ALERT_CHARTCUTERIE.value,
Referrer.DISCOVER_SLACK_UNFURL.value,
Referrer.API_DASHBOARDS_WIDGET_AREA_CHART.value,
Referrer.API_DASHBOARDS_WIDGET_BAR_CHART.value,
Referrer.API_DASHBOARDS_WIDGET_LINE_CHART.value,
Expand Down Expand Up @@ -106,6 +111,14 @@
}


SENTRY_BACKEND_REFERRERS = [
Referrer.API_ALERTS_CHARTCUTERIE.value,
Referrer.API_ENDPOINT_REGRESSION_ALERT_CHARTCUTERIE.value,
Referrer.API_FUNCTION_REGRESSION_ALERT_CHARTCUTERIE.value,
Referrer.DISCOVER_SLACK_UNFURL.value,
]


@region_silo_endpoint
class OrganizationEventsStatsEndpoint(OrganizationEventsV2EndpointBase):
publish_status = {
Expand Down Expand Up @@ -163,6 +176,7 @@ def check_if_results_have_data(self, results: SnubaTSResult | dict[str, SnubaTSR
return has_data

def get(self, request: Request, organization: Organization) -> Response:
query_source = QuerySource.FRONTEND if is_frontend_request(request) else QuerySource.API
with sentry_sdk.start_span(op="discover.endpoint", description="filter_params") as span:
span.set_data("organization", organization)

Expand Down Expand Up @@ -201,6 +215,8 @@ def get(self, request: Request, organization: Organization) -> Response:
if referrer in ALLOWED_EVENTS_STATS_REFERRERS.union(METRICS_ENHANCED_REFERRERS)
else Referrer.API_ORGANIZATION_EVENT_STATS.value
)
if referrer in SENTRY_BACKEND_REFERRERS:
query_source = QuerySource.SENTRY_BACKEND
batch_features = self.get_features(organization, request)
has_chart_interpolation = batch_features.get(
"organizations:performance-chart-interpolation", False
Expand Down Expand Up @@ -275,6 +291,7 @@ def _get_event_stats(
on_demand_metrics_enabled=use_on_demand_metrics,
on_demand_metrics_type=on_demand_metrics_type,
include_other=include_other,
query_source=query_source,
)

return scoped_dataset.timeseries_query(
Expand All @@ -300,6 +317,7 @@ def _get_event_stats(
)
),
on_demand_metrics_type=on_demand_metrics_type,
query_source=query_source,
)

def get_event_stats_factory(scoped_dataset):
Expand Down
2 changes: 2 additions & 0 deletions src/sentry/integrations/slack/unfurl/discover.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
from sentry.models.organization import Organization
from sentry.models.user import User
from sentry.search.events.filter import to_list
from sentry.snuba.referrer import Referrer
from sentry.utils.dates import (
get_interval_from_range,
parse_stats_period,
Expand Down Expand Up @@ -235,6 +236,7 @@ def unfurl_discover(
params.setlist("statsPeriod", [stats_period])

endpoint = "events-stats/"
params["referrer"] = Referrer.DISCOVER_SLACK_UNFURL.value

try:
resp = client.get(
Expand Down
7 changes: 5 additions & 2 deletions src/sentry/search/events/builder/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@
from sentry.services.hybrid_cloud.user.service import user_service
from sentry.snuba.dataset import Dataset
from sentry.snuba.metrics.utils import MetricMeta
from sentry.snuba.query_sources import QuerySource
from sentry.utils.dates import outside_retention_with_modified_start
from sentry.utils.env import in_test_environment
from sentry.utils.snuba import (
Expand Down Expand Up @@ -1454,10 +1455,12 @@ def handle_invalid_float(cls, value: float | None) -> float | None:
return None
return value

def run_query(self, referrer: str, use_cache: bool = False) -> Any:
def run_query(
self, referrer: str, use_cache: bool = False, query_source: QuerySource | None = None
) -> Any:
if not referrer:
InvalidSearchQuery("Query missing referrer.")
return raw_snql_query(self.get_snql_query(), referrer, use_cache)
return raw_snql_query(self.get_snql_query(), referrer, use_cache, query_source)

def process_results(self, results: Any) -> EventsResponse:
with sentry_sdk.start_span(op="QueryBuilder", description="process_results") as span:
Expand Down
14 changes: 10 additions & 4 deletions src/sentry/search/events/builder/metrics.py
Original file line number Diff line number Diff line change
Expand Up @@ -1704,7 +1704,9 @@ def get_snql_query(self) -> list[Request]:

return queries

def run_query(self, referrer: str, use_cache: bool = False) -> Any:
def run_query(
self, referrer: str, use_cache: bool = False, query_source: QuerySource | None = None
) -> Any:
if self.use_metrics_layer or self.use_on_demand:
from sentry.snuba.metrics.datasource import get_series
from sentry.snuba.metrics.mqb_query_transformer import (
Expand Down Expand Up @@ -1752,7 +1754,7 @@ def run_query(self, referrer: str, use_cache: bool = False) -> Any:

queries = self.get_snql_query()
if queries:
results = bulk_snuba_queries(queries, referrer, use_cache)
results = bulk_snuba_queries(queries, referrer, use_cache, query_source=query_source)
else:
results = []

Expand Down Expand Up @@ -1960,7 +1962,9 @@ def resolve_top_event_conditions(

return final_condition

def run_query(self, referrer: str, use_cache: bool = False) -> Any:
def run_query(
self, referrer: str, use_cache: bool = False, query_source: QuerySource | None = None
) -> Any:
result = {}
results = []
meta_dict = {}
Expand Down Expand Up @@ -2028,7 +2032,9 @@ def run_query(self, referrer: str, use_cache: bool = False) -> Any:
else:
queries = self.get_snql_query()
if queries:
results = bulk_snuba_queries(queries, referrer, use_cache)
results = bulk_snuba_queries(
queries, referrer, use_cache, query_source=query_source
)

time_map: dict[str, dict[str, Any]] = defaultdict(dict)
for current_result in results:
Expand Down
18 changes: 15 additions & 3 deletions src/sentry/snuba/discover.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
)
from sentry.search.events.types import HistogramParams, ParamsType, QueryBuilderConfig
from sentry.snuba.dataset import Dataset
from sentry.snuba.query_sources import QuerySource
from sentry.tagstore.base import TOP_VALUES_DEFAULT_LIMIT
from sentry.utils.math import nice_int
from sentry.utils.snuba import (
Expand Down Expand Up @@ -200,6 +201,7 @@ def _query(
on_demand_metrics_enabled=False,
on_demand_metrics_type=None,
dataset=Dataset.Discover,
query_source: QuerySource | None = None,
) -> EventsResponse:
if not selected_columns:
raise InvalidSearchQuery("No columns selected")
Expand Down Expand Up @@ -233,7 +235,9 @@ def _query(
if extra_columns is not None:
builder.columns.extend(extra_columns)

result = builder.process_results(builder.run_query(referrer))
result = builder.process_results(
builder.run_query(referrer=referrer, query_source=query_source)
)
result["meta"]["tips"] = transform_tips(builder.tips)
return result

Expand All @@ -253,6 +257,7 @@ def _timeseries_query(
on_demand_metrics_enabled=False,
on_demand_metrics_type=None,
dataset=Dataset.Discover,
query_source: QuerySource | None = None,
):
assert dataset in [Dataset.Discover, Dataset.Transactions]
with sentry_sdk.start_span(op="discover.discover", description="timeseries.filter_transform"):
Expand Down Expand Up @@ -287,7 +292,7 @@ def _timeseries_query(
query_list.append(comparison_builder)

query_results = bulk_snuba_queries(
[query.get_snql_query() for query in query_list], referrer
[query.get_snql_query() for query in query_list], referrer, query_source=query_source
)

with sentry_sdk.start_span(op="discover.discover", description="timeseries.transform_results"):
Expand Down Expand Up @@ -438,6 +443,7 @@ def timeseries_query(
use_metrics_layer=False,
on_demand_metrics_enabled=False,
on_demand_metrics_type=None,
query_source: QuerySource | None = None,
):
"""
High-level API for doing arbitrary user timeseries queries against events.
Expand Down Expand Up @@ -477,6 +483,7 @@ def timeseries_query(
on_demand_metrics_enabled=on_demand_metrics_enabled,
on_demand_metrics_type=on_demand_metrics_type,
dataset=Dataset.Discover,
query_source=query_source,
)


Expand Down Expand Up @@ -525,6 +532,7 @@ def _top_events_timeseries(
on_demand_metrics_enabled: bool = False,
on_demand_metrics_type=None,
dataset=Dataset.Discover,
query_source: QuerySource | None = None,
):
assert dataset in [Dataset.Discover, Dataset.Transactions]
if top_events is None:
Expand All @@ -542,6 +550,7 @@ def _top_events_timeseries(
include_equation_fields=True,
skip_tag_resolution=True,
dataset=dataset,
query_source=query_source,
)

top_events_builder = TopEventsQueryBuilder(
Expand Down Expand Up @@ -574,9 +583,10 @@ def _top_events_timeseries(
result, other_result = bulk_snuba_queries(
[top_events_builder.get_snql_query(), other_events_builder.get_snql_query()],
referrer=referrer,
query_source=query_source,
)
else:
result = top_events_builder.run_query(referrer)
result = top_events_builder.run_query(referrer=referrer, query_source=query_source)
other_result = {"data": []}
if (
not allow_empty
Expand Down Expand Up @@ -664,6 +674,7 @@ def top_events_timeseries(
functions_acl=None,
on_demand_metrics_enabled: bool = False,
on_demand_metrics_type=None,
query_source: QuerySource | None = None,
):
"""
High-level API for doing arbitrary user timeseries queries for a limited number of top events
Expand Down Expand Up @@ -706,6 +717,7 @@ def top_events_timeseries(
on_demand_metrics_enabled=on_demand_metrics_enabled,
on_demand_metrics_type=on_demand_metrics_type,
dataset=Dataset.Discover,
query_source=query_source,
)


Expand Down
4 changes: 3 additions & 1 deletion src/sentry/snuba/errors.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
from sentry.snuba.dataset import Dataset
from sentry.snuba.discover import EventsResponse, transform_tips, zerofill
from sentry.snuba.metrics.extraction import MetricSpecType
from sentry.snuba.query_sources import QuerySource
from sentry.utils.snuba import SnubaTSResult, bulk_snuba_queries

is_filter_translation = {}
Expand Down Expand Up @@ -94,6 +95,7 @@ def timeseries_query(
use_metrics_layer=False,
on_demand_metrics_enabled=False,
on_demand_metrics_type: MetricSpecType | None = None,
query_source: QuerySource | None = None,
):
with sentry_sdk.start_span(op="errors", description="timeseries.filter_transform"):
equations, columns = categorize_columns(selected_columns)
Expand Down Expand Up @@ -129,7 +131,7 @@ def timeseries_query(
query_list.append(comparison_builder)

query_results = bulk_snuba_queries(
[query.get_snql_query() for query in query_list], referrer
[query.get_snql_query() for query in query_list], referrer, query_source=query_source
)

with sentry_sdk.start_span(op="errors", description="timeseries.transform_results"):
Expand Down
Loading

0 comments on commit d3154bf

Please sign in to comment.