diff --git a/snuba/admin/static/api_client.tsx b/snuba/admin/static/api_client.tsx index bbf05c24d4..eafe5c226b 100644 --- a/snuba/admin/static/api_client.tsx +++ b/snuba/admin/static/api_client.tsx @@ -25,7 +25,10 @@ import { SnQLRequest, SnQLResult, SnubaDatasetName } from "./snql_to_sql/types"; import { KafkaTopicData } from "./kafka/types"; import { QuerylogRequest, QuerylogResult } from "./querylog/types"; -import { CardinalityQueryRequest, CardinalityQueryResult } from "./cardinality_analyzer/types"; +import { + CardinalityQueryRequest, + CardinalityQueryResult, +} from "./cardinality_analyzer/types"; import { AllocationPolicy } from "./capacity_management/types"; @@ -57,7 +60,9 @@ interface Client { getPredefinedQuerylogOptions: () => Promise<[PredefinedQuery]>; getQuerylogSchema: () => Promise; executeQuerylogQuery: (req: QuerylogRequest) => Promise; - executeCardinalityQuery: (req: CardinalityQueryRequest) => Promise; + executeCardinalityQuery: ( + req: CardinalityQueryRequest + ) => Promise; getAllMigrationGroups: () => Promise; runMigration: (req: RunMigrationRequest) => Promise; getAllowedTools: () => Promise; @@ -183,7 +188,7 @@ function Client() { }, convertSnQLQuery: (query: SnQLRequest) => { - const url = baseUrl + "snql_to_sql"; + const url = baseUrl + "snuba_debug"; return fetch(url, { headers: { "Content-Type": "application/json" }, method: "POST", diff --git a/snuba/admin/tool_policies.py b/snuba/admin/tool_policies.py index 319b581dbe..eeda63e2d7 100644 --- a/snuba/admin/tool_policies.py +++ b/snuba/admin/tool_policies.py @@ -28,6 +28,7 @@ class AdminTools(Enum): KAFKA = "kafka" CAPACITY_MANAGEMENT = "capacity-management" CARDINALITY_ANALYZER = "cardinality-analyzer" + SNUBA_EXPLAIN = "snuba_explain" DEVELOPER_TOOLS: set[AdminTools] = {AdminTools.SNQL_TO_SQL, AdminTools.QUERY_TRACING} diff --git a/snuba/admin/views.py b/snuba/admin/views.py index acea1966b9..b5890069b3 100644 --- a/snuba/admin/views.py +++ b/snuba/admin/views.py @@ -64,6 +64,7 @@ from snuba.migrations.groups import MigrationGroup from snuba.migrations.runner import MigrationKey, Runner from snuba.query.exceptions import InvalidQueryException +from snuba.state.explain_meta import explain_cleanup, get_explain_meta from snuba.utils.metrics.timer import Timer from snuba.web.views import dataset_query @@ -721,15 +722,23 @@ def snuba_datasets() -> Response: ) -@application.route("/snql_to_sql", methods=["POST"]) -@check_tool_perms(tools=[AdminTools.SNQL_TO_SQL]) -def snql_to_sql() -> Response: +@application.route("/snuba_debug", methods=["POST"]) +@check_tool_perms(tools=[AdminTools.SNQL_TO_SQL, AdminTools.SNUBA_EXPLAIN]) +def snuba_debug() -> Response: body = json.loads(request.data) body["debug"] = True body["dry_run"] = True try: dataset = get_dataset(body.pop("dataset")) - return dataset_query(dataset, body, Timer("admin")) + response = dataset_query(dataset, body, Timer("admin")) + data = response.get_json() + assert isinstance(data, dict) + + meta = get_explain_meta() + if meta: + data["explain"] = asdict(meta) + response.data = json.dumps(data) + return response except InvalidQueryException as exception: return Response( json.dumps({"error": {"message": str(exception)}}, indent=4), @@ -742,6 +751,8 @@ def snql_to_sql() -> Response: 400, {"Content-Type": "application/json"}, ) + finally: + explain_cleanup() @application.route("/storages_with_allocation_policies") diff --git a/snuba/pipeline/processors.py b/snuba/pipeline/processors.py index 581f75e69b..3c48106255 100644 --- a/snuba/pipeline/processors.py +++ b/snuba/pipeline/processors.py @@ -2,6 +2,7 @@ import sentry_sdk +from snuba.state import explain_meta from snuba.datasets.entities.factory import get_entity from snuba.datasets.plans.query_plan import ClickhouseQueryPlan from snuba.query.logical import Query as LogicalQuery @@ -74,4 +75,7 @@ def execute_entity_processors(query: LogicalQuery, settings: QuerySettings) -> N with sentry_sdk.start_span( description=type(processor).__name__, op="processor" ): + if settings.get_dry_run(): + explain_meta.add_step("entity_processor", type(processor).__name__) + processor.process_query(query, settings) diff --git a/snuba/state/explain_meta.py b/snuba/state/explain_meta.py new file mode 100644 index 0000000000..b00cf0c123 --- /dev/null +++ b/snuba/state/explain_meta.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, cast + +from flask import g + + +@dataclass +class ExplainStep: + category: str # The type of step e.g. "processor" + name: str # The specific name for the step e.g. "TimeSeriesProcessor" + data: dict[str, Any] = field( + default_factory=dict + ) # Any extra information about the step + + +@dataclass +class ExplainMeta: + steps: list[ExplainStep] = field(default_factory=list) + + def add_step(self, step: ExplainStep) -> None: + self.steps.append(step) + + +def add_step(category: str, name: str, data: dict[str, Any] | None = None) -> None: + try: + if data is None: + data = {} + step = ExplainStep(category, name, data) + + if not hasattr(g, "explain_meta"): + g.explain_meta = ExplainMeta() + + g.explain_meta.add_step(step) + except RuntimeError: + # Code is executing outside of a flask context + return + + +def get_explain_meta() -> ExplainMeta | None: + try: + if hasattr(g, "explain_meta"): + return cast(ExplainMeta, g.explain_meta) + return None + except RuntimeError: + # Code is executing outside of a flask context + return None + + +def explain_cleanup() -> None: + try: + if hasattr(g, "explain_meta"): + g.pop("explain_meta") + except RuntimeError: + # Code is executing outside of a flask context + pass diff --git a/tests/admin/test_api.py b/tests/admin/test_api.py index 5f1c4d4bb3..e21e58e081 100644 --- a/tests/admin/test_api.py +++ b/tests/admin/test_api.py @@ -320,9 +320,9 @@ def test_get_snuba_datasets(admin_api: FlaskClient) -> None: assert set(data) == set(get_enabled_dataset_names()) -def test_convert_SnQL_to_SQL_invalid_dataset(admin_api: FlaskClient) -> None: +def test_snuba_debug_invalid_dataset(admin_api: FlaskClient) -> None: response = admin_api.post( - "/snql_to_sql", data=json.dumps({"dataset": "", "query": ""}) + "/snuba_debug", data=json.dumps({"dataset": "", "query": ""}) ) assert response.status_code == 400 data = json.loads(response.data) @@ -330,9 +330,9 @@ def test_convert_SnQL_to_SQL_invalid_dataset(admin_api: FlaskClient) -> None: @pytest.mark.redis_db -def test_convert_SnQL_to_SQL_invalid_query(admin_api: FlaskClient) -> None: +def test_snuba_debug_invalid_query(admin_api: FlaskClient) -> None: response = admin_api.post( - "/snql_to_sql", data=json.dumps({"dataset": "sessions", "query": ""}) + "/snuba_debug", data=json.dumps({"dataset": "sessions", "query": ""}) ) assert response.status_code == 400 data = json.loads(response.data) @@ -344,7 +344,7 @@ def test_convert_SnQL_to_SQL_invalid_query(admin_api: FlaskClient) -> None: @pytest.mark.redis_db @pytest.mark.clickhouse_db -def test_convert_SnQL_to_SQL_valid_query(admin_api: FlaskClient) -> None: +def test_snuba_debug_valid_query(admin_api: FlaskClient) -> None: snql_query = """ MATCH (sessions) SELECT sessions_crashed @@ -354,11 +354,17 @@ def test_convert_SnQL_to_SQL_valid_query(admin_api: FlaskClient) -> None: AND started < toDateTime('2022-02-01 00:00:00') """ response = admin_api.post( - "/snql_to_sql", data=json.dumps({"dataset": "sessions", "query": snql_query}) + "/snuba_debug", data=json.dumps({"dataset": "sessions", "query": snql_query}) ) assert response.status_code == 200 data = json.loads(response.data) assert data["sql"] != "" + assert len(data["explain"]["steps"]) > 0 + assert { + "category": "entity_processor", + "name": "BasicFunctionsProcessor", + "data": {}, + } in data["explain"]["steps"] @pytest.mark.redis_db