From 457cd3487ded7f09c6b260141dc6c2bc25045e4c Mon Sep 17 00:00:00 2001 From: Yuriy Krasilnikov <30294585+YuriyKrasilnikov@users.noreply.github.com> Date: Tue, 4 Aug 2026 03:27:57 +0300 Subject: [PATCH] fix(api): include query lifecycle timing in /api/v1/chart/data response (#37516) --- .../configuration/configuring-superset.mdx | 44 ++ .../src/query/types/QueryResponse.ts | 17 + superset/charts/data/api.py | 67 ++- superset/charts/schemas.py | 32 ++ .../commands/chart/data/get_data_command.py | 28 ++ superset/common/chart_data_timing.py | 114 +++++ superset/common/query_actions.py | 145 ++++-- superset/common/query_context.py | 23 + superset/common/query_context_processor.py | 74 +++- superset/config.py | 4 + .../charts/data/api_tests.py | 108 +++-- .../commands/data/test_get_data_command.py | 52 +++ .../unit_tests/charts/test_chart_data_api.py | 416 +++++++++++++++++- tests/unit_tests/charts/test_schemas.py | 43 ++ tests/unit_tests/common/test_query_actions.py | 318 +++++++++---- .../common/test_query_actions_drill_detail.py | 53 ++- .../common/test_query_context_processor.py | 30 +- .../test_query_context_processor_timing.py | 266 +++++++++++ 18 files changed, 1618 insertions(+), 216 deletions(-) create mode 100644 superset/common/chart_data_timing.py create mode 100644 tests/unit_tests/common/test_query_context_processor_timing.py diff --git a/docs/admin_docs/configuration/configuring-superset.mdx b/docs/admin_docs/configuration/configuring-superset.mdx index fe4be60e663..eaa5e99f0ed 100644 --- a/docs/admin_docs/configuration/configuring-superset.mdx +++ b/docs/admin_docs/configuration/configuring-superset.mdx @@ -97,6 +97,50 @@ for more information on how to configure it. At the very least, you'll want to change `SECRET_KEY` and `SQLALCHEMY_DATABASE_URI`. Continue reading for more about each of these. +## Chart-data query timing + +Set `CHART_DATA_INCLUDE_TIMING = True` to add an optional versioned timing object +to every successful JSON query result returned by the chart-data API. The setting +is `False` by default, so enabling it is an explicit API-contract choice for an +operator. File exports, streaming responses, and HTTP error responses do not +include this object. + +```json +{ + "timing": { + "version": 1, + "query": { + "query_planning_ms": 1.23, + "cache_resolution_ms": 0.45, + "data_acquisition_ms": null, + "payload_assembly_ms": 0.67, + "total_ms": 2.98 + } + } +} +``` + +Durations are milliseconds rounded to two decimal places. A numeric `0.0` +means that the corresponding stage ran but rounded below that precision; +`null` means it did not apply. For example, `data_acquisition_ms` is null for +a normal dataframe cache hit, while metadata-only results have null phase values +and a numeric total. + +The phases have fixed ownership: `query_planning_ms` includes Jinja rendering, +row-level-security transformation, and cache identity; `cache_resolution_ms` +includes cache lookup, compatibility policy, deserialization, and rehydration; +`data_acquisition_ms` includes database work and annotation dependencies; and +`payload_assembly_ms` includes response shaping and AUTO-currency fallback. + +`total_ms` is measured over one continuous per-query execution interval. It +contains the exposed stages and unattributed work such as query-result cache +persistence, so it is not the sum of the phase values. Query-context cache +persistence is excluded. The total also excludes request parsing, +authorization, contribution-total work before the per-query loop, client +post-processing, JSON serialization, network transfer, and background async +producer execution. Additive optional fields can preserve version 1; a field +rename, removal, or semantic change requires a new version. + ## Specifying a SECRET_KEY ### Adding an initial SECRET_KEY diff --git a/superset-frontend/packages/superset-ui-core/src/query/types/QueryResponse.ts b/superset-frontend/packages/superset-ui-core/src/query/types/QueryResponse.ts index 23c13d3c7c1..42433550bdf 100644 --- a/superset-frontend/packages/superset-ui-core/src/query/types/QueryResponse.ts +++ b/superset-frontend/packages/superset-ui-core/src/query/types/QueryResponse.ts @@ -29,6 +29,19 @@ export interface DataRecord { [key: string]: DataRecordValue; } +export interface ChartDataQueryTiming { + query_planning_ms: number | null; + cache_resolution_ms: number | null; + data_acquisition_ms: number | null; + payload_assembly_ms: number | null; + total_ms: number; +} + +export interface ChartDataTiming { + version: 1; + query: ChartDataQueryTiming; +} + /** * Queried data for charts. The `queries` field from `POST /chart/data`. * See superset/charts/schemas.py for the class of the same name. @@ -84,6 +97,10 @@ export interface ChartDataResponseResult { * or null if multiple currencies are present. */ detected_currency?: string | null; + /** + * Versioned query lifecycle timing in milliseconds. + */ + timing?: ChartDataTiming; } export interface TimeseriesChartDataResponseResult extends ChartDataResponseResult { diff --git a/superset/charts/data/api.py b/superset/charts/data/api.py index 4f45131b014..dfbba0467b0 100644 --- a/superset/charts/data/api.py +++ b/superset/charts/data/api.py @@ -51,6 +51,7 @@ from superset.commands.chart.exceptions import ( ChartDataQueryFailedError, ) from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType +from superset.common.chart_data_timing import ChartDataExecutionResult from superset.connectors.sqla.models import BaseDatasource from superset.constants import CACHE_DISABLED_TIMEOUT from superset.daos.exceptions import DatasourceNotFound @@ -435,14 +436,16 @@ class ChartDataRestApi(ChartRestApi): # First, look for the chart query results in the cache, # but only if we're not forcing a refresh. if not form_data.get("force"): - with contextlib.suppress(ChartDataCacheLoadError): - result = command.run(force_cached=True) + try: + result = command.execute(force_cached=True) if result is not None: # Log is_cached if extra payload callback is provided. # This indicates no async job was triggered - data was already # cached and a synchronous response is being returned immediately. - self._log_is_cached(result, add_extra_log_payload) + self._log_is_cached(result.materialize(), add_extra_log_payload) return self._send_chart_response(result) + except ChartDataCacheLoadError: + pass # Otherwise, kick off a background job to run the chart query. # Clients will either poll or be notified of query completion, # at which point they will call the /data/ endpoint @@ -453,26 +456,37 @@ class ChartDataRestApi(ChartRestApi): except AsyncQueryTokenException: return self.response_401() - result = async_command.run(form_data, get_user_id()) - return self.response(202, **result) + async_result = async_command.run(form_data, get_user_id()) + return self.response(202, **async_result) def _send_chart_response( # noqa: C901 self, - result: dict[Any, Any], + result: dict[Any, Any] | ChartDataExecutionResult, form_data: dict[str, Any] | None = None, datasource: BaseDatasource | Query | None = None, filename: str | None = None, expected_rows: int | None = None, dashboard_filter_context: DashboardFilterContext | None = None, ) -> Response: - result_type = result["query_context"].result_type - result_format = result["query_context"].result_format + if isinstance(result, ChartDataExecutionResult): + execution_result: ChartDataExecutionResult | None = result + materialized_result = result.materialize() + else: + execution_result = None + materialized_result = result + + result_type = materialized_result["query_context"].result_type + result_format = materialized_result["query_context"].result_format # Post-process the data so it matches the data presented in the chart. # This is needed for sending reports based on text charts that do the # post-processing of data, eg, the pivot table. if result_type == ChartDataResultType.POST_PROCESSED: - result = apply_client_processing(result, form_data, datasource) + materialized_result = apply_client_processing( + materialized_result, + form_data, + datasource, + ) if result_format in ChartDataResultFormat.table_like(): # Verify user has permission to export file @@ -485,15 +499,21 @@ class ChartDataRestApi(ChartRestApi): if not has_export_perm: return self.response_403() - if not result["queries"]: + if not materialized_result["queries"]: return self.response_400(_("Empty query result")) is_csv_format = result_format == ChartDataResultFormat.CSV # Check if we should use streaming for large datasets - if is_csv_format and self._should_use_streaming(result, form_data): + if is_csv_format and self._should_use_streaming( + materialized_result, + form_data, + ): return self._create_streaming_csv_response( - result, form_data, filename=filename, expected_rows=expected_rows + materialized_result, + form_data, + filename=filename, + expected_rows=expected_rows, ) export_filename = filename or self._get_default_export_filename(form_data) @@ -504,9 +524,9 @@ class ChartDataRestApi(ChartRestApi): r"\.(csv|xlsx|zip)$", "", export_filename, flags=re.IGNORECASE ) - if len(result["queries"]) == 1: + if len(materialized_result["queries"]) == 1: # return single query results - data = result["queries"][0]["data"] + data = materialized_result["queries"][0]["data"] if is_csv_format: return CsvResponse( data, headers=generate_download_headers("csv", export_filename) @@ -528,7 +548,7 @@ class ChartDataRestApi(ChartRestApi): files = { f"query_{idx + 1}.{result_format}": _process_data(query["data"]) - for idx, query in enumerate(result["queries"]) + for idx, query in enumerate(materialized_result["queries"]) } return Response( create_zip(files), @@ -537,7 +557,13 @@ class ChartDataRestApi(ChartRestApi): ) if result_format == ChartDataResultFormat.JSON: - queries = result["queries"] + queries = materialized_result["queries"] + if execution_result and app.config.get("CHART_DATA_INCLUDE_TIMING"): + for query, query_result in zip( + queries, execution_result.queries, strict=True + ): + query["timing"] = query_result.timing.as_public_dict() + if security_manager.is_guest_user(): for query in queries: query.pop("query", None) @@ -611,15 +637,18 @@ class ChartDataRestApi(ChartRestApi): ) -> Response: """Get data response and optionally log is_cached information.""" try: - result = command.run(force_cached=force_cached) + result = command.execute(force_cached=force_cached) except ChartDataCacheLoadError as exc: return self.response_422(message=exc.message) except ChartDataQueryFailedError as exc: return self.response_400(message=exc.message) # Log is_cached if extra payload callback is provided - if add_extra_log_payload and result and "queries" in result: - is_cached_values = [query.get("is_cached") for query in result["queries"]] + materialized_result = result.materialize() + if add_extra_log_payload and materialized_result.get("queries"): + is_cached_values = [ + query.get("is_cached") for query in materialized_result["queries"] + ] add_extra_log_payload(is_cached=is_cached_values) return self._send_chart_response( diff --git a/superset/charts/schemas.py b/superset/charts/schemas.py index d50c8f4b4e4..b97310fd4b2 100644 --- a/superset/charts/schemas.py +++ b/superset/charts/schemas.py @@ -27,6 +27,7 @@ from marshmallow.validate import Length, Range from marshmallow_union import Union from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType +from superset.common.chart_data_timing import CHART_DATA_TIMING_VERSION from superset.db_engine_specs.base import builtin_time_grains from superset.subjects.schemas import SubjectResponseSchema from superset.tags.models import TagType @@ -1565,6 +1566,26 @@ class AnnotationDataSchema(Schema): ) +class ChartDataQueryTimingSchema(Schema): + """Schema for the versioned per-query timing phases.""" + + query_planning_ms = fields.Float(required=True, allow_none=True) + cache_resolution_ms = fields.Float(required=True, allow_none=True) + data_acquisition_ms = fields.Float(required=True, allow_none=True) + payload_assembly_ms = fields.Float(required=True, allow_none=True) + total_ms = fields.Float(required=True) + + +class ChartDataTimingSchema(Schema): + """Schema for the versioned query lifecycle timing breakdown.""" + + version = fields.Integer( + required=True, + validate=validate.Equal(CHART_DATA_TIMING_VERSION), + ) + query = fields.Nested(ChartDataQueryTimingSchema, required=True) + + class ChartDataResponseResult(Schema): annotation_data = fields.List( fields.Dict( @@ -1676,6 +1697,17 @@ class ChartDataResponseResult(Schema): metadata={"description": "Warning message when results were truncated"}, allow_none=True, ) + timing = fields.Nested( + ChartDataTimingSchema, + metadata={ + "description": ( + "Optional versioned query lifecycle timing breakdown in milliseconds. " + "Present only when CHART_DATA_INCLUDE_TIMING is enabled; disabled by " + "default." + ) + }, + required=False, + ) class DashboardFilterInfoSchema(Schema): diff --git a/superset/commands/chart/data/get_data_command.py b/superset/commands/chart/data/get_data_command.py index eeaa860aadb..946811269bf 100644 --- a/superset/commands/chart/data/get_data_command.py +++ b/superset/commands/chart/data/get_data_command.py @@ -25,6 +25,7 @@ from superset.commands.chart.exceptions import ( ChartDataQueryFailedError, ) from superset.common.chart_data import ChartDataResultType +from superset.common.chart_data_timing import ChartDataExecutionResult from superset.common.query_context import QueryContext from superset.exceptions import CacheLoadError @@ -69,5 +70,32 @@ class ChartDataCommand(BaseCommand): return return_value + def execute(self, **kwargs: Any) -> ChartDataExecutionResult: + """Execute and return timing as a typed sidecar.""" + cache_query_context = kwargs.get("cache", False) + force_cached = kwargs.get("force_cached", False) + try: + result = self._query_context.get_payload_result( + cache_query_context=cache_query_context, + force_cached=force_cached, + ) + except CacheLoadError as ex: + raise ChartDataCacheLoadError(ex.message) from ex + + for query in result.queries: + if ( + query.payload.get("error") + and self._query_context.result_type != ChartDataResultType.QUERY + ): + raise ChartDataQueryFailedError( + _("Error: %(error)s", error=query.payload["error"]) + ) + + return ChartDataExecutionResult( + query_context=self._query_context, + queries=result.queries, + cache_key=result.cache_key, + ) + def validate(self) -> None: self._query_context.raise_for_access() diff --git a/superset/common/chart_data_timing.py b/superset/common/chart_data_timing.py new file mode 100644 index 00000000000..767981cb17d --- /dev/null +++ b/superset/common/chart_data_timing.py @@ -0,0 +1,114 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, TYPE_CHECKING + +if TYPE_CHECKING: + from superset.common.query_context import QueryContext + +NANOSECONDS_PER_MILLISECOND: int = 1_000_000 +CHART_DATA_TIMING_VERSION: int = 1 + + +def to_ms(value_ns: int | None) -> float | None: + """Convert nanoseconds to rounded milliseconds for public output.""" + if value_ns is None: + return None + return round(value_ns / NANOSECONDS_PER_MILLISECOND, 2) + + +@dataclass(frozen=True) +class QueryAcquisitionTiming: + """Timing captured by the dataframe payload owner.""" + + query_planning_ns: int + cache_resolution_ns: int + data_acquisition_ns: int | None + payload_assembly_ns: int + + +@dataclass(frozen=True) +class QueryTiming: + """Completed timing for one query in a chart-data execution.""" + + query_planning_ns: int | None + cache_resolution_ns: int | None + data_acquisition_ns: int | None + payload_assembly_ns: int | None + total_ns: int + + def as_public_dict(self) -> dict[str, Any]: + """Return the versioned chart-data API representation.""" + return { + "version": CHART_DATA_TIMING_VERSION, + "query": { + "query_planning_ms": to_ms(self.query_planning_ns), + "cache_resolution_ms": to_ms(self.cache_resolution_ns), + "data_acquisition_ms": to_ms(self.data_acquisition_ns), + "payload_assembly_ms": to_ms(self.payload_assembly_ns), + "total_ms": to_ms(self.total_ns), + }, + } + + +@dataclass(frozen=True) +class QueryAcquisitionResult: + """A dataframe payload paired with acquisition timing.""" + + payload: dict[str, Any] + timing: QueryAcquisitionTiming + + +@dataclass(frozen=True) +class QueryDataResult: + """A query payload paired with completed timing.""" + + payload: dict[str, Any] + timing: QueryTiming + + +@dataclass(frozen=True) +class QueryContextExecutionResult: + """Typed query-context result with timing outside query payloads.""" + + queries: tuple[QueryDataResult, ...] + cache_key: str | None = None + + +@dataclass(frozen=True) +class ChartDataExecutionResult: + """Typed result of executing a chart-data command.""" + + query_context: QueryContext + queries: tuple[QueryDataResult, ...] + cache_key: str | None = None + + def materialize(self) -> dict[str, Any]: + """Return the historical command payload shape.""" + queries: list[dict[str, Any]] = [] + for query_result in self.queries: + queries.append(dict(query_result.payload)) + + result: dict[str, Any] = { + "query_context": self.query_context, + "queries": queries, + } + if self.cache_key is not None: + result["cache_key"] = self.cache_key + return result diff --git a/superset/common/query_actions.py b/superset/common/query_actions.py index 3178ae07510..7668ebfe0c9 100644 --- a/superset/common/query_actions.py +++ b/superset/common/query_actions.py @@ -17,12 +17,17 @@ from __future__ import annotations import copy -import logging +import time from typing import Any, Callable, TYPE_CHECKING from flask_babel import _ from superset.common.chart_data import ChartDataResultType +from superset.common.chart_data_timing import ( + QueryAcquisitionTiming, + QueryDataResult, + QueryTiming, +) from superset.common.db_query_status import QueryStatus from superset.exceptions import QueryObjectValidationError, SupersetParseError from superset.explorables.base import Explorable @@ -43,8 +48,6 @@ if TYPE_CHECKING: from superset.common.query_context import QueryContext from superset.common.query_object import QueryObject -logger = logging.getLogger(__name__) - def _get_datasource(query_context: QueryContext, query_obj: QueryObject) -> Explorable: return query_obj.datasource or query_context.datasource @@ -150,14 +153,25 @@ def _detect_currency( ) -def _get_full( +def _get_full_with_timing( query_context: QueryContext, query_obj: QueryObject, force_cached: bool | None = False, +) -> tuple[dict[str, Any], QueryAcquisitionTiming, int]: + acquired = query_context.get_df_payload_result(query_obj, force_cached=force_cached) + payload_assembly_start_ns = time.perf_counter_ns() + payload = _materialize_full_payload(query_context, query_obj, acquired.payload) + payload_assembly_ns = max(0, time.perf_counter_ns() - payload_assembly_start_ns) + return payload, acquired.timing, payload_assembly_ns + + +def _materialize_full_payload( + query_context: QueryContext, + query_obj: QueryObject, + payload: dict[str, Any], ) -> dict[str, Any]: datasource = _get_datasource(query_context, query_obj) result_type = query_obj.result_type or query_context.result_type - payload = query_context.get_df_payload(query_obj, force_cached=force_cached) df = payload["df"] status = payload["status"] if status != QueryStatus.FAILED: @@ -202,9 +216,10 @@ def _get_full( return payload -def _get_samples( - query_context: QueryContext, query_obj: QueryObject, force_cached: bool = False -) -> dict[str, Any]: +def _prepare_samples_query( + query_context: QueryContext, + query_obj: QueryObject, +) -> QueryObject: datasource = _get_datasource(query_context, query_obj) query_obj = copy.copy(query_obj) query_obj.is_timeseries = False @@ -226,12 +241,13 @@ def _get_samples( query_obj.columns = qry_obj_cols query_obj.from_dttm = None query_obj.to_dttm = None - return _get_full(query_context, query_obj, force_cached) + return query_obj -def _get_drill_detail( - query_context: QueryContext, query_obj: QueryObject, force_cached: bool = False -) -> dict[str, Any]: +def _prepare_drill_detail_query( + query_context: QueryContext, + query_obj: QueryObject, +) -> QueryObject: # todo(yongjie): Remove this function, # when determining whether samples should be applied to the time filter. datasource = _get_datasource(query_context, query_obj) @@ -257,41 +273,52 @@ def _get_drill_detail( qry_obj_cols.append(o.column_name) query_obj.columns = qry_obj_cols query_obj.orderby = [(query_obj.columns[0], True)] - return _get_full(query_context, query_obj, force_cached) + return query_obj -def _get_results( - query_context: QueryContext, query_obj: QueryObject, force_cached: bool = False -) -> dict[str, Any]: - payload = _get_full(query_context, query_obj, force_cached) - return payload - - -_result_type_functions: dict[ +_metadata_result_type_functions: dict[ ChartDataResultType, Callable[[QueryContext, QueryObject, bool], dict[str, Any]] ] = { ChartDataResultType.COLUMNS: _get_columns, ChartDataResultType.TIMEGRAINS: _get_timegrains, ChartDataResultType.QUERY: _get_query, - ChartDataResultType.SAMPLES: _get_samples, - ChartDataResultType.FULL: _get_full, - ChartDataResultType.RESULTS: _get_results, - # for requests for post-processed data we return the full results, - # and post-process it later where we have the chart context, since - # post-processing is unique to each visualization type - ChartDataResultType.POST_PROCESSED: _get_full, - ChartDataResultType.DRILL_DETAIL: _get_drill_detail, } +_data_result_type_preparers: dict[ + ChartDataResultType, + Callable[[QueryContext, QueryObject], QueryObject] | None, +] = { + ChartDataResultType.SAMPLES: _prepare_samples_query, + ChartDataResultType.FULL: None, + ChartDataResultType.RESULTS: None, + # Post-processing is visualization-specific, so full data is returned and + # transformed later with the chart context. + ChartDataResultType.POST_PROCESSED: None, + ChartDataResultType.DRILL_DETAIL: _prepare_drill_detail_query, +} + + +def _metadata_timing(total_ns: int) -> QueryTiming: + return QueryTiming( + query_planning_ns=None, + cache_resolution_ns=None, + data_acquisition_ns=None, + payload_assembly_ns=None, + total_ns=total_ns, + ) + + def get_query_results( result_type: ChartDataResultType, query_context: QueryContext, query_obj: QueryObject, force_cached: bool, ) -> dict[str, Any]: - """ - Return result payload for a chart data request. + """Return the legacy payload-only view of a chart-data result. + + This compatibility wrapper deliberately discards the typed timing sidecar + returned by :func:`get_query_results_with_timing`. :param result_type: the type of result to return :param query_context: query context to which the query object belongs @@ -300,8 +327,56 @@ def get_query_results( :raises QueryObjectValidationError: if an unsupported result type is requested :return: JSON serializable result payload """ - if result_func := _result_type_functions.get(result_type): - return result_func(query_context, query_obj, force_cached) - raise QueryObjectValidationError( - _("Invalid result type: %(result_type)s", result_type=result_type) + return get_query_results_with_timing( + result_type, + query_context, + query_obj, + force_cached, + ).payload + + +def get_query_results_with_timing( + result_type: ChartDataResultType, + query_context: QueryContext, + query_obj: QueryObject, + force_cached: bool, +) -> QueryDataResult: + """ + Return result payload and timing without storing timing in the payload. + + The total interval begins before result-family preparation and dispatch. + Metadata result types do not acquire dataframe state, so their phase values + are null while the measured total remains available. + """ + started_ns = time.perf_counter_ns() + if result_func := _metadata_result_type_functions.get(result_type): + payload = result_func(query_context, query_obj, force_cached) + total_ns = max(0, time.perf_counter_ns() - started_ns) + return QueryDataResult(payload=payload, timing=_metadata_timing(total_ns)) + + if result_type not in _data_result_type_preparers: + raise QueryObjectValidationError( + _("Invalid result type: %(result_type)s", result_type=result_type) + ) + + if preparer := _data_result_type_preparers[result_type]: + query_obj = preparer(query_context, query_obj) + + payload, acquisition_timing, action_assembly_ns = _get_full_with_timing( + query_context, + query_obj, + force_cached, + ) + total_ns = max(0, time.perf_counter_ns() - started_ns) + return QueryDataResult( + payload=payload, + timing=QueryTiming( + query_planning_ns=acquisition_timing.query_planning_ns, + cache_resolution_ns=acquisition_timing.cache_resolution_ns, + data_acquisition_ns=acquisition_timing.data_acquisition_ns, + payload_assembly_ns=( + acquisition_timing.payload_assembly_ns + action_assembly_ns + ), + total_ns=total_ns, + ), ) diff --git a/superset/common/query_context.py b/superset/common/query_context.py index bef870035c0..8a5819e6aa6 100644 --- a/superset/common/query_context.py +++ b/superset/common/query_context.py @@ -22,6 +22,10 @@ from typing import Any, ClassVar, TYPE_CHECKING import pandas as pd from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType +from superset.common.chart_data_timing import ( + QueryAcquisitionResult, + QueryContextExecutionResult, +) from superset.common.query_context_processor import QueryContextProcessor from superset.common.query_object import QueryObject from superset.explorables.base import Explorable @@ -98,6 +102,14 @@ class QueryContext: """Returns the query results with both metadata and data""" return self._processor.get_payload(cache_query_context, force_cached) + def get_payload_result( + self, + cache_query_context: bool | None = False, + force_cached: bool = False, + ) -> QueryContextExecutionResult: + """Return query results with timing kept outside query payloads.""" + return self._processor.get_payload_result(cache_query_context, force_cached) + def get_cache_timeout(self) -> int | None: """ Get the cache timeout for this query context. @@ -131,6 +143,17 @@ class QueryContext: force_cached=force_cached, ) + def get_df_payload_result( + self, + query_obj: QueryObject, + force_cached: bool | None = False, + ) -> QueryAcquisitionResult: + """Return dataframe payload with timing kept outside the payload.""" + return self._processor.get_df_payload_result( + query_obj=query_obj, + force_cached=force_cached, + ) + def get_query_result(self, query_object: QueryObject) -> QueryResult: return self._processor.get_query_result(query_object) diff --git a/superset/common/query_context_processor.py b/superset/common/query_context_processor.py index f90f57021d5..87bb12d694f 100644 --- a/superset/common/query_context_processor.py +++ b/superset/common/query_context_processor.py @@ -19,6 +19,7 @@ from __future__ import annotations import copy import logging import re +import time from typing import Any, cast, ClassVar, Sequence, TYPE_CHECKING import pandas as pd @@ -26,9 +27,14 @@ from flask import current_app from flask_babel import gettext as _ from superset.common.chart_data import ChartDataResultFormat +from superset.common.chart_data_timing import ( + QueryAcquisitionResult, + QueryAcquisitionTiming, + QueryContextExecutionResult, +) from superset.common.db_query_status import QueryStatus from superset.common.grouping_sets import grouping_marker_label -from superset.common.query_actions import get_query_results +from superset.common.query_actions import get_query_results_with_timing from superset.common.utils.query_cache_manager import QueryCacheManager from superset.common.utils.time_range_utils import get_since_until_from_time_range from superset.constants import CACHE_DISABLED_TIMEOUT, CacheRegion @@ -86,7 +92,14 @@ class QueryContextProcessor: def get_df_payload( self, query_obj: QueryObject, force_cached: bool | None = False ) -> dict[str, Any]: - """Handles caching around the df payload retrieval""" + """Return the historical dataframe payload without timing metadata.""" + return self.get_df_payload_result(query_obj, force_cached).payload + + def get_df_payload_result( + self, query_obj: QueryObject, force_cached: bool | None = False + ) -> QueryAcquisitionResult: + """Acquire a dataframe and return timing as a typed sidecar.""" + query_planning_start_ns = time.perf_counter_ns() if query_obj: # Always validate the query object before generating cache key # This ensures sanitize_clause() is called and extras are normalized @@ -95,6 +108,9 @@ class QueryContextProcessor: cache_key = self.query_cache_key(query_obj) timeout = self.get_cache_timeout() force_query = self._query_context.force or timeout == CACHE_DISABLED_TIMEOUT + query_planning_ns = max(0, time.perf_counter_ns() - query_planning_start_ns) + + cache_resolution_start_ns = time.perf_counter_ns() cache = QueryCacheManager.get( key=cache_key, region=CacheRegion.DATA, @@ -114,7 +130,11 @@ class QueryContextProcessor: ): cache.is_loaded = False + cache_resolution_ns = max(0, time.perf_counter_ns() - cache_resolution_start_ns) + + data_acquisition_ns: int | None = None if query_obj and cache_key and not cache.is_loaded: + data_acquisition_start_ns = time.perf_counter_ns() try: if invalid_columns := [ col @@ -134,6 +154,15 @@ class QueryContextProcessor: query_result = self.get_query_result(query_obj) annotation_data = self.get_annotation_data(query_obj) + except QueryObjectValidationError as ex: + cache.error_message = str(ex) + cache.status = QueryStatus.FAILED + finally: + data_acquisition_ns = max( + 0, time.perf_counter_ns() - data_acquisition_start_ns + ) + + if cache.status != QueryStatus.FAILED: cache.set_query_result( key=cache_key, query_result=query_result, @@ -143,10 +172,8 @@ class QueryContextProcessor: datasource_uid=self._qc_datasource.uid, region=CacheRegion.DATA, ) - except QueryObjectValidationError as ex: - cache.error_message = str(ex) - cache.status = QueryStatus.FAILED + payload_assembly_start_ns = time.perf_counter_ns() # the N-dimensional DataFrame has converted into flat DataFrame # by `flatten operator`, "comma" in the column is escaped by `escape_separator` # the result DataFrame columns should be unescaped @@ -206,7 +233,7 @@ class QueryContextProcessor: row_count=f"{row_count:,}", ) - return { + payload = { "cache_key": cache_key, "cached_dttm": cache.cache_dttm, "queried_dttm": cache.queried_dttm, @@ -228,6 +255,15 @@ class QueryContextProcessor: "label_map": label_map, "warning": warning, } + timing = QueryAcquisitionTiming( + query_planning_ns=query_planning_ns, + cache_resolution_ns=cache_resolution_ns, + data_acquisition_ns=data_acquisition_ns, + payload_assembly_ns=max( + 0, time.perf_counter_ns() - payload_assembly_start_ns + ), + ) + return QueryAcquisitionResult(payload=payload, timing=timing) def query_cache_key(self, query_obj: QueryObject, **kwargs: Any) -> str | None: """ @@ -425,6 +461,20 @@ class QueryContextProcessor: force_cached: bool = False, ) -> dict[str, Any]: """Returns the query results with both metadata and data""" + result = self.get_payload_result(cache_query_context, force_cached) + return_value: dict[str, Any] = { + "queries": [query.payload for query in result.queries], + } + if result.cache_key is not None: + return_value["cache_key"] = result.cache_key + return return_value + + def get_payload_result( + self, + cache_query_context: bool | None = False, + force_cached: bool = False, + ) -> QueryContextExecutionResult: + """Return query results with timing kept outside query payloads.""" queries_needing_totals, totals_idx = self._prepare_contribution_totals() @@ -447,18 +497,17 @@ class QueryContextProcessor: ) ] - query_results = [ - get_query_results( + query_results = tuple( + get_query_results_with_timing( query_obj.result_type or self._query_context.result_type, self._query_context, query_obj, force_cached, ) for query_obj in self._query_context.queries - ] - - return_value = {"queries": query_results} + ) + cache_key = None if cache_query_context: cache_key = self.cache_key() set_and_log_cache( @@ -475,9 +524,8 @@ class QueryContextProcessor: }, self.get_cache_timeout(), ) - return_value["cache_key"] = cache_key # type: ignore - return return_value + return QueryContextExecutionResult(queries=query_results, cache_key=cache_key) def get_cache_timeout(self) -> int: """ diff --git a/superset/config.py b/superset/config.py index 82f14935754..cb31e5f793f 100644 --- a/superset/config.py +++ b/superset/config.py @@ -1365,6 +1365,10 @@ DATA_CACHE_CONFIG: CacheConfig = {"CACHE_TYPE": "NullCache"} # 10 * 1024 * 1024 for a 10 MB limit. DATA_CACHE_MAX_VALUE_SIZE: int | None = None +# Include per-query lifecycle timing in /api/v1/chart/data JSON responses. +# The default keeps the public response contract unchanged. +CHART_DATA_INCLUDE_TIMING: bool = False + # Cache for dashboard filter state. `CACHE_TYPE` defaults to `SupersetMetastoreCache` # that stores the values in the key-value table in the Superset metastore, as it's # required for Superset to operate correctly, but can be replaced by any diff --git a/tests/integration_tests/charts/data/api_tests.py b/tests/integration_tests/charts/data/api_tests.py index eebee483a46..18a30c543fd 100644 --- a/tests/integration_tests/charts/data/api_tests.py +++ b/tests/integration_tests/charts/data/api_tests.py @@ -37,6 +37,11 @@ from flask.ctx import AppContext from superset.charts.data.api import ChartDataRestApi from superset.commands.chart.data.get_data_command import ChartDataCommand from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType +from superset.common.chart_data_timing import ( + ChartDataExecutionResult, + QueryDataResult, + QueryTiming, +) from superset.connectors.sqla.models import SqlaTable, TableColumn from superset.constants import CACHE_DISABLED_TIMEOUT from superset.errors import SupersetErrorType @@ -93,6 +98,16 @@ INCOMPATIBLE_ADHOC_COLUMN_FIXTURE: AdhocColumn = { } +def _query_timing() -> QueryTiming: + return QueryTiming( + query_planning_ns=0, + cache_resolution_ns=0, + data_acquisition_ns=None, + payload_assembly_ns=0, + total_ns=0, + ) + + @pytest.fixture(autouse=True) def _skip_by_backend(app_context: AppContext): if backend() == "hive": @@ -777,23 +792,33 @@ class TestPostChartDataApi(BaseTestChartDataApi): result_format = ChartDataResultFormat.JSON result_type = ChartDataResultType.FULL - cmd_run_val = { - "query_context": QueryContext(), - "queries": [{"query": "select * from foo", "is_cached": True}], - } + cmd_execute_val = ChartDataExecutionResult( + query_context=QueryContext(), + queries=( + QueryDataResult( + payload={"query": "select * from foo", "is_cached": True}, + timing=_query_timing(), + ), + ), + ) with mock.patch.object( - ChartDataCommand, "run", return_value=cmd_run_val - ) as patched_run: + ChartDataCommand, "execute", return_value=cmd_execute_val + ) as patched_execute: self.query_context_payload["result_type"] = ChartDataResultType.FULL rv = self.post_assert_metric( CHART_DATA_URI, self.query_context_payload, "data" ) assert rv.status_code == 200 data = json.loads(rv.data.decode("utf-8")) - patched_run.assert_called_once_with(force_cached=True) + patched_execute.assert_called_once_with(force_cached=True) assert data == { - "result": [{"query": "select * from foo", "is_cached": True}] + "result": [ + { + "query": "select * from foo", + "is_cached": True, + } + ] } # Verify that is_cached was logged to event logger @@ -839,40 +864,45 @@ class TestPostChartDataApi(BaseTestChartDataApi): @with_feature_flags(GLOBAL_ASYNC_QUERIES=True) @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") - @mock.patch("superset.charts.data.api.ChartDataCommand.run") - def test_chart_data_async_force_refresh(self, mock_run): + @mock.patch("superset.charts.data.api.ChartDataCommand.execute") + def test_chart_data_async_force_refresh(self, mock_execute): """ Chart data API: Test that force=true skips cache and triggers async job """ app._got_first_request = False async_query_manager_factory.init_app(app) - # Mock the command.run to return cached data + # Mock the command execution to return cached data class QueryContext: result_format = ChartDataResultFormat.JSON result_type = ChartDataResultType.FULL - mock_run.return_value = { - "query_context": QueryContext(), - "queries": [{"query": "select * from foo", "is_cached": True}], - } + mock_execute.return_value = ChartDataExecutionResult( + query_context=QueryContext(), + queries=( + QueryDataResult( + payload={"query": "select * from foo", "is_cached": True}, + timing=_query_timing(), + ), + ), + ) # Test without force - should return cached data synchronously self.query_context_payload["result_type"] = ChartDataResultType.FULL rv = self.post_assert_metric(CHART_DATA_URI, self.query_context_payload, "data") assert rv.status_code == 200 - mock_run.assert_called_once_with(force_cached=True) + mock_execute.assert_called_once_with(force_cached=True) # Reset the mock - mock_run.reset_mock() + mock_execute.reset_mock() # Test with force=true - should skip cache and return async response self.query_context_payload["force"] = True rv = self.post_assert_metric(CHART_DATA_URI, self.query_context_payload, "data") assert rv.status_code == 202 - # When force=true, command.run should not be called at all in _run_async + # When force=true, command execution should not be called at all in _run_async # since we skip the cache check entirely - mock_run.assert_not_called() + mock_execute.assert_not_called() data = json.loads(rv.data.decode("utf-8")) keys = list(data.keys()) self.assertCountEqual( # noqa: PT009 @@ -1416,14 +1446,14 @@ class TestGetChartDataApi(BaseTestChartDataApi): app._got_first_request = False async_query_manager_factory.init_app(app) cache_loader.load.return_value = self.query_context_payload - orig_run = ChartDataCommand.run + orig_execute = ChartDataCommand.execute - def mock_run(self, **kwargs): + def mock_execute(self, **kwargs): assert kwargs["force_cached"] is True # noqa: E712 # override force_cached to get result from DB - return orig_run(self, force_cached=False) + return orig_execute(self, force_cached=False) - with mock.patch.object(ChartDataCommand, "run", new=mock_run): + with mock.patch.object(ChartDataCommand, "execute", new=mock_execute): rv = self.get_assert_metric( f"{CHART_DATA_URI}/test-cache-key", "data_from_cache" ) @@ -1465,14 +1495,14 @@ class TestGetChartDataApi(BaseTestChartDataApi): async_query_manager_factory.init_app(app) self.logout() cache_loader.load.return_value = self.query_context_payload - orig_run = ChartDataCommand.run + orig_execute = ChartDataCommand.execute - def mock_run(self, **kwargs): + def mock_execute(self, **kwargs): assert kwargs["force_cached"] is True # noqa: E712 # override force_cached to get result from DB - return orig_run(self, force_cached=False) + return orig_execute(self, force_cached=False) - with mock.patch.object(ChartDataCommand, "run", new=mock_run): + with mock.patch.object(ChartDataCommand, "execute", new=mock_execute): rv = self.client.get( f"{CHART_DATA_URI}/test-cache-key", ) @@ -2118,6 +2148,30 @@ class TestGetChartDataWithDashboardFilter(BaseTestChartDataApi): assert "dashboard_filters" not in data mock_get_filter_ctx.assert_not_called() + @with_config({"CHART_DATA_INCLUDE_TIMING": False}) + @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") + def test_get_data_excludes_timing_by_default(self): + """GET chart data preserves its default response contract.""" + chart = self._setup_chart_with_query_context() + + rv = self.get_assert_metric(f"api/v1/chart/{chart.id}/data/", "get_data") + data = json.loads(rv.data.decode("utf-8")) + + assert rv.status_code == 200 + assert "timing" not in data["result"][0] + + @with_config({"CHART_DATA_INCLUDE_TIMING": True}) + @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") + def test_get_data_projects_opt_in_timing(self): + """GET chart data projects the public timing object only when enabled.""" + chart = self._setup_chart_with_query_context() + + rv = self.get_assert_metric(f"api/v1/chart/{chart.id}/data/", "get_data") + data = json.loads(rv.data.decode("utf-8")) + + assert rv.status_code == 200 + assert data["result"][0]["timing"]["version"] == 1 + @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") def test_get_data_invalid_filters_dashboard_id_returns_400(self): """ diff --git a/tests/unit_tests/charts/commands/data/test_get_data_command.py b/tests/unit_tests/charts/commands/data/test_get_data_command.py index 378e127f216..c4d7c9281a1 100644 --- a/tests/unit_tests/charts/commands/data/test_get_data_command.py +++ b/tests/unit_tests/charts/commands/data/test_get_data_command.py @@ -22,9 +22,24 @@ import pytest from superset.commands.chart.data.get_data_command import ChartDataCommand from superset.commands.chart.exceptions import ChartDataQueryFailedError from superset.common.chart_data import ChartDataResultType +from superset.common.chart_data_timing import ( + QueryContextExecutionResult, + QueryDataResult, + QueryTiming, +) from superset.common.query_context import QueryContext +def _query_timing() -> QueryTiming: + return QueryTiming( + query_planning_ns=1_000_000, + cache_resolution_ns=2_000_000, + data_acquisition_ns=3_000_000, + payload_assembly_ns=4_000_000, + total_ns=10_000_000, + ) + + def test_query_result_type_allows_validation_error_payload() -> None: """ Regression test: Ensure result_type='query' with error payload returns @@ -148,6 +163,43 @@ def test_full_result_type_returns_successful_data() -> None: assert "error" not in result["queries"][0] +def test_execute_returns_timing_sidecar_without_mutating_payload() -> None: + mock_query_context = Mock(spec=QueryContext) + mock_query_context.result_type = ChartDataResultType.FULL + query_payload = {"data": [{"col1": "value1"}], "colnames": ["col1"]} + mock_query_context.get_payload_result.return_value = QueryContextExecutionResult( + queries=(QueryDataResult(query_payload, _query_timing()),), + cache_key="cache-key", + ) + + command = ChartDataCommand(mock_query_context) + + result = command.execute(cache=True) + + assert result.queries[0].payload is query_payload + assert "timing" not in query_payload + + materialized = result.materialize() + assert materialized["cache_key"] == "cache-key" + assert "timing" not in materialized["queries"][0] + assert "timing" not in query_payload + + +def test_execute_raises_on_error_payload_for_data_results() -> None: + mock_query_context = Mock(spec=QueryContext) + mock_query_context.result_type = ChartDataResultType.FULL + mock_query_context.get_payload_result.return_value = QueryContextExecutionResult( + queries=(QueryDataResult({"error": "Invalid column name"}, _query_timing()),), + ) + + command = ChartDataCommand(mock_query_context) + + with pytest.raises(ChartDataQueryFailedError) as exc_info: + command.execute() + + assert "Invalid column name" in str(exc_info.value) + + def test_query_result_type_with_multiple_queries_and_mixed_results() -> None: """ Test that result_type='query' handles multiple queries with mixed results. diff --git a/tests/unit_tests/charts/test_chart_data_api.py b/tests/unit_tests/charts/test_chart_data_api.py index ea999cc07f5..43dac1ad69b 100644 --- a/tests/unit_tests/charts/test_chart_data_api.py +++ b/tests/unit_tests/charts/test_chart_data_api.py @@ -19,11 +19,24 @@ from __future__ import annotations from typing import Any, TYPE_CHECKING from unittest.mock import MagicMock, patch +import pytest from flask import Flask, g +from superset.charts.data.api import ChartDataRestApi from superset.charts.data.dashboard_filter_context import ( apply_dashboard_filter_context, ) +from superset.charts.schemas import ChartDataTimingSchema +from superset.commands.chart.exceptions import ( + ChartDataCacheLoadError, + ChartDataQueryFailedError, +) +from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType +from superset.common.chart_data_timing import ( + ChartDataExecutionResult, + QueryDataResult, + QueryTiming, +) from superset.jinja_context import ExtraCache from superset.utils import json @@ -31,6 +44,29 @@ if TYPE_CHECKING: from superset.app import SupersetApp +def _query_timing() -> QueryTiming: + return QueryTiming( + query_planning_ns=1_000_000, + cache_resolution_ns=2_000_000, + data_acquisition_ns=3_000_000, + payload_assembly_ns=4_000_000, + total_ns=10_000_000, + ) + + +def _json_execution_result( + query_payload: dict[str, Any], + result_type: ChartDataResultType = ChartDataResultType.FULL, +) -> ChartDataExecutionResult: + query_context = MagicMock() + query_context.result_type = result_type + query_context.result_format = ChartDataResultFormat.JSON + return ChartDataExecutionResult( + query_context=query_context, + queries=(QueryDataResult(payload=query_payload, timing=_query_timing()),), + ) + + def test_get_data_sets_g_form_data_without_dashboard_filter() -> None: """ Regression test: GET /api/v1/chart//data/ must populate g.form_data @@ -209,10 +245,374 @@ def test_apply_dashboard_filter_context_keeps_grain_when_no_grain_filter() -> No assert query["columns"][0]["timeGrain"] == "P1M" +def test_send_chart_response_excludes_timing_by_default(app: SupersetApp) -> None: + query_payload = {"data": [{"col1": 1}], "query": "SELECT 1"} + result = _json_execution_result(query_payload) + + original = app.config.get("CHART_DATA_INCLUDE_TIMING") + try: + app.config["CHART_DATA_INCLUDE_TIMING"] = False + api = ChartDataRestApi() + with ( + app.test_request_context("/api/v1/chart/data"), + patch( + "superset.charts.data.api.security_manager.is_guest_user", + return_value=False, + ), + ): + response = api._send_chart_response(result) + finally: + app.config["CHART_DATA_INCLUDE_TIMING"] = original + + body = json.loads(response.get_data(as_text=True)) + assert body == {"result": [{"data": [{"col1": 1}], "query": "SELECT 1"}]} + assert "timing" not in query_payload + + +def test_send_chart_response_includes_opt_in_timing(app: SupersetApp) -> None: + query_payload = {"data": [{"col1": 1}], "query": "SELECT 1"} + result = _json_execution_result(query_payload) + + original = app.config.get("CHART_DATA_INCLUDE_TIMING") + try: + app.config["CHART_DATA_INCLUDE_TIMING"] = True + api = ChartDataRestApi() + with ( + app.test_request_context("/api/v1/chart/data"), + patch( + "superset.charts.data.api.security_manager.is_guest_user", + return_value=False, + ), + ): + response = api._send_chart_response(result) + finally: + app.config["CHART_DATA_INCLUDE_TIMING"] = original + + body = json.loads(response.get_data(as_text=True)) + projected_timing = body["result"][0]["timing"] + assert body["result"][0]["timing"] == { + "version": 1, + "query": { + "query_planning_ms": 1.0, + "cache_resolution_ms": 2.0, + "data_acquisition_ms": 3.0, + "payload_assembly_ms": 4.0, + "total_ms": 10.0, + }, + } + assert ChartDataTimingSchema().load(projected_timing) == projected_timing + assert "timing" not in query_payload + + +def test_send_chart_response_strips_guest_query_after_timing_projection( + app: SupersetApp, +) -> None: + query_payload = {"data": [{"col1": 1}], "query": "SELECT 1"} + result = _json_execution_result(query_payload) + + original = app.config.get("CHART_DATA_INCLUDE_TIMING") + try: + app.config["CHART_DATA_INCLUDE_TIMING"] = True + api = ChartDataRestApi() + with ( + app.test_request_context("/api/v1/chart/data"), + patch( + "superset.charts.data.api.security_manager.is_guest_user", + return_value=True, + ), + ): + response = api._send_chart_response(result) + finally: + app.config["CHART_DATA_INCLUDE_TIMING"] = original + + query = json.loads(response.get_data(as_text=True))["result"][0] + assert "query" not in query + assert "timing" in query + assert "query" in query_payload + + +def test_send_chart_response_pairs_each_timing_with_its_query( + app: SupersetApp, +) -> None: + first_payload = {"data": [{"col1": 1}], "query": "SELECT 1"} + second_payload = {"data": [{"col2": 2}], "query": "SELECT 2"} + query_context = MagicMock() + query_context.result_type = ChartDataResultType.FULL + query_context.result_format = ChartDataResultFormat.JSON + result = ChartDataExecutionResult( + query_context=query_context, + queries=( + QueryDataResult(payload=first_payload, timing=_query_timing()), + QueryDataResult( + payload=second_payload, + timing=QueryTiming( + query_planning_ns=5_000_000, + cache_resolution_ns=6_000_000, + data_acquisition_ns=None, + payload_assembly_ns=7_000_000, + total_ns=18_000_000, + ), + ), + ), + ) + + original = app.config.get("CHART_DATA_INCLUDE_TIMING") + try: + app.config["CHART_DATA_INCLUDE_TIMING"] = True + api = ChartDataRestApi() + with ( + app.test_request_context("/api/v1/chart/data"), + patch( + "superset.charts.data.api.security_manager.is_guest_user", + return_value=False, + ), + ): + response = api._send_chart_response(result) + finally: + app.config["CHART_DATA_INCLUDE_TIMING"] = original + + queries = json.loads(response.get_data(as_text=True))["result"] + assert queries[0]["timing"]["query"]["total_ms"] == 10.0 + assert queries[1]["timing"]["query"] == { + "query_planning_ms": 5.0, + "cache_resolution_ms": 6.0, + "data_acquisition_ms": None, + "payload_assembly_ms": 7.0, + "total_ms": 18.0, + } + assert "timing" not in first_payload + assert "timing" not in second_payload + + +def test_send_chart_response_refuses_mismatched_query_and_timing_counts( + app: SupersetApp, +) -> None: + result = _json_execution_result({"data": [{"col1": 1}]}) + materialized_result = { + "query_context": result.query_context, + "queries": [{"data": [{"col1": 1}]}, {"data": [{"col2": 2}]}], + } + + original = app.config.get("CHART_DATA_INCLUDE_TIMING") + try: + app.config["CHART_DATA_INCLUDE_TIMING"] = True + api = ChartDataRestApi() + with ( + app.test_request_context("/api/v1/chart/data"), + patch.object( + ChartDataExecutionResult, + "materialize", + return_value=materialized_result, + ), + patch( + "superset.charts.data.api.security_manager.is_guest_user", + return_value=False, + ), + pytest.raises( + ValueError, + match=r"zip\(\)", + ), + ): + api._send_chart_response(result) + finally: + app.config["CHART_DATA_INCLUDE_TIMING"] = original + + +def test_send_chart_response_projects_timing_for_query_preview_error( + app: SupersetApp, +) -> None: + result = _json_execution_result( + {"error": "Invalid column", "query": "SELECT invalid"}, + result_type=ChartDataResultType.QUERY, + ) + + original = app.config.get("CHART_DATA_INCLUDE_TIMING") + try: + app.config["CHART_DATA_INCLUDE_TIMING"] = True + api = ChartDataRestApi() + with ( + app.test_request_context("/api/v1/chart/data"), + patch( + "superset.charts.data.api.security_manager.is_guest_user", + return_value=False, + ), + ): + response = api._send_chart_response(result) + finally: + app.config["CHART_DATA_INCLUDE_TIMING"] = original + + query = json.loads(response.get_data(as_text=True))["result"][0] + assert response.status_code == 200 + assert query["error"] == "Invalid column" + assert query["timing"]["version"] == 1 + + +@pytest.mark.parametrize( + ("exception", "status_code"), + [ + (ChartDataCacheLoadError("cache unavailable"), 422), + (ChartDataQueryFailedError("query failed"), 400), + ], +) +def test_get_data_response_excludes_timing_from_http_errors( + app: SupersetApp, + exception: Exception, + status_code: int, +) -> None: + command = MagicMock() + command.execute.side_effect = exception + api = ChartDataRestApi() + + original = app.config.get("CHART_DATA_INCLUDE_TIMING") + try: + app.config["CHART_DATA_INCLUDE_TIMING"] = True + with app.test_request_context("/api/v1/chart/data"): + response = api._get_data_response(command) + finally: + app.config["CHART_DATA_INCLUDE_TIMING"] = original + + assert response.status_code == status_code + assert "timing" not in json.loads(response.get_data(as_text=True)) + + +def test_run_async_does_not_project_timing_onto_a_job_response( + app: SupersetApp, +) -> None: + command = MagicMock() + command.execute.side_effect = ChartDataCacheLoadError("cache miss") + async_command = MagicMock() + async_command.run.return_value = { + "channel_id": "channel", + "job_id": "job", + "user_id": 1, + "status": "pending", + "errors": [], + "result_url": "/api/v1/chart/data/job", + } + api = ChartDataRestApi() + + original = app.config.get("CHART_DATA_INCLUDE_TIMING") + try: + app.config["CHART_DATA_INCLUDE_TIMING"] = True + with ( + app.test_request_context("/api/v1/chart/data", method="POST"), + patch( + "superset.charts.data.api.CreateAsyncChartDataJobCommand", + return_value=async_command, + ), + patch("superset.charts.data.api.get_user_id", return_value=1), + ): + response = api._run_async({"force": False}, command) + finally: + app.config["CHART_DATA_INCLUDE_TIMING"] = original + + assert response.status_code == 202 + assert "timing" not in json.loads(response.get_data(as_text=True)) + async_command.validate.assert_called_once() + + +def test_run_async_projects_opt_in_timing_for_a_cached_result( + app: SupersetApp, +) -> None: + query_payload = {"data": [{"col1": 1}], "query": "SELECT 1"} + command = MagicMock() + command.execute.return_value = _json_execution_result(query_payload) + api = ChartDataRestApi() + + original = app.config.get("CHART_DATA_INCLUDE_TIMING") + try: + app.config["CHART_DATA_INCLUDE_TIMING"] = True + with ( + app.test_request_context("/api/v1/chart/data"), + patch( + "superset.charts.data.api.security_manager.is_guest_user", + return_value=False, + ), + ): + response = api._run_async({"force": False}, command) + finally: + app.config["CHART_DATA_INCLUDE_TIMING"] = original + + command.execute.assert_called_once_with(force_cached=True) + result = json.loads(response.get_data(as_text=True))["result"][0] + assert result["timing"]["version"] == 1 + assert "timing" not in query_payload + + +def test_send_chart_response_projects_timing_after_client_processing( + app: SupersetApp, +) -> None: + query_payload = {"data": [{"col1": 1}], "query": "SELECT 1"} + result = _json_execution_result( + query_payload, + result_type=ChartDataResultType.POST_PROCESSED, + ) + + def process( + materialized_result: dict[Any, Any], + _form_data: dict[str, Any] | None, + _datasource: Any, + ) -> dict[Any, Any]: + materialized_result["queries"][0]["data"] = [{"col1": 2}] + return materialized_result + + original = app.config.get("CHART_DATA_INCLUDE_TIMING") + try: + app.config["CHART_DATA_INCLUDE_TIMING"] = True + api = ChartDataRestApi() + with ( + app.test_request_context("/api/v1/chart/data"), + patch( + "superset.charts.data.api.apply_client_processing", + side_effect=process, + ), + patch( + "superset.charts.data.api.security_manager.is_guest_user", + return_value=False, + ), + ): + response = api._send_chart_response(result, form_data={"viz_type": "table"}) + finally: + app.config["CHART_DATA_INCLUDE_TIMING"] = original + + query = json.loads(response.get_data(as_text=True))["result"][0] + assert query["data"] == [{"col1": 2}] + assert query["timing"]["version"] == 1 + assert query_payload["data"] == [{"col1": 1}] + + +def test_send_chart_response_does_not_project_timing_for_csv( + app: SupersetApp, +) -> None: + query_payload = {"data": "col_a\n1\n"} + query_context = MagicMock() + query_context.result_type = ChartDataResultType.FULL + query_context.result_format = ChartDataResultFormat.CSV + result = ChartDataExecutionResult( + query_context=query_context, + queries=(QueryDataResult(query_payload, _query_timing()),), + ) + api = ChartDataRestApi() + + original = app.config.get("CHART_DATA_INCLUDE_TIMING") + try: + app.config["CHART_DATA_INCLUDE_TIMING"] = True + with ( + app.test_request_context("/api/v1/chart/data"), + patch("superset.charts.data.api.security_manager") as security_manager, + patch("superset.charts.data.api.is_feature_enabled", return_value=False), + ): + security_manager.can_access.return_value = True + response = api._send_chart_response(result) + finally: + app.config["CHART_DATA_INCLUDE_TIMING"] = original + + assert response.status_code == 200 + assert "timing" not in query_payload + + def _extract_filename(form_value: str) -> str | None: """Run _extract_export_params_from_request with a form filename value.""" - from superset.charts.data.api import ChartDataRestApi - app = Flask(__name__) with app.test_request_context("/", method="POST", data={"filename": form_value}): filename, _ = ChartDataRestApi._extract_export_params_from_request(MagicMock()) @@ -245,9 +645,6 @@ def test_send_chart_response_uses_chart_name_for_csv_filename() -> None: Content-Disposition header, not just a bare timestamp, mirroring the streaming CSV export path. """ - from superset.charts.data.api import ChartDataRestApi - from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType - query_context = MagicMock() query_context.result_type = ChartDataResultType.FULL query_context.result_format = ChartDataResultFormat.CSV @@ -273,9 +670,6 @@ def test_send_chart_response_uses_chart_name_for_csv_filename() -> None: def test_send_chart_response_uses_chart_name_for_xlsx_filename() -> None: """Same regression as above, for the XLSX export branch.""" - from superset.charts.data.api import ChartDataRestApi - from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType - query_context = MagicMock() query_context.result_type = ChartDataResultType.FULL query_context.result_format = ChartDataResultFormat.XLSX @@ -301,9 +695,6 @@ def test_send_chart_response_uses_chart_name_for_xlsx_filename() -> None: def test_send_chart_response_uses_chart_name_for_zip_filename() -> None: """Same regression as above, for the multi-query zip export branch.""" - from superset.charts.data.api import ChartDataRestApi - from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType - query_context = MagicMock() query_context.result_type = ChartDataResultType.FULL query_context.result_format = ChartDataResultFormat.CSV @@ -336,9 +727,6 @@ def test_send_chart_response_does_not_double_extension_for_csv_filename() -> Non ``.csv`` extension must not be doubled (e.g. ``export.csv.csv``) by the non-streaming CSV export branch of _send_chart_response. """ - from superset.charts.data.api import ChartDataRestApi - from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType - query_context = MagicMock() query_context.result_type = ChartDataResultType.FULL query_context.result_format = ChartDataResultFormat.CSV diff --git a/tests/unit_tests/charts/test_schemas.py b/tests/unit_tests/charts/test_schemas.py index 65dd1313274..73d5fd419b6 100644 --- a/tests/unit_tests/charts/test_schemas.py +++ b/tests/unit_tests/charts/test_schemas.py @@ -23,7 +23,9 @@ from superset.charts.schemas import ( ChartDataExtrasSchema, ChartDataProphetOptionsSchema, ChartDataQueryObjectSchema, + ChartDataResponseResult, ChartDataRollingOptionsSchema, + ChartDataTimingSchema, ChartPostSchema, ChartPutSchema, DEFAULT_MAX_PROPHET_PERIODS, @@ -61,6 +63,47 @@ def test_get_time_grain_choices(app_context: None) -> None: current_app.config["TIME_GRAIN_ADDONS"] = original_addons +def test_chart_data_timing_schema_validates_version(app_context: None) -> None: + schema = ChartDataTimingSchema() + payload = { + "version": 1, + "query": { + "query_planning_ms": 1.0, + "cache_resolution_ms": 2.0, + "data_acquisition_ms": None, + "payload_assembly_ms": 4.0, + "total_ms": 10.0, + }, + } + + assert schema.load(payload)["version"] == 1 + + with pytest.raises(ValidationError) as exc_info: + schema.load({**payload, "version": 2}) + assert "version" in exc_info.value.messages + + +def test_chart_data_response_timing_is_optional_but_never_null( + app_context: None, +) -> None: + timing_field = ChartDataResponseResult().fields["timing"] + timing_payload = { + "version": 1, + "query": { + "query_planning_ms": 1.0, + "cache_resolution_ms": 2.0, + "data_acquisition_ms": None, + "payload_assembly_ms": 4.0, + "total_ms": 10.0, + }, + } + + assert timing_field.required is False + assert timing_field.deserialize(timing_payload) == timing_payload + with pytest.raises(ValidationError): + timing_field.deserialize(None) + + def test_chart_data_prophet_options_schema_time_grain_validation( app_context: None, ) -> None: diff --git a/tests/unit_tests/common/test_query_actions.py b/tests/unit_tests/common/test_query_actions.py index 77dd962b3d9..0cb68b7b1f2 100644 --- a/tests/unit_tests/common/test_query_actions.py +++ b/tests/unit_tests/common/test_query_actions.py @@ -14,114 +14,270 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from typing import Any +from typing import cast from unittest.mock import MagicMock, patch +import pytest + +from superset.common import query_actions from superset.common.chart_data import ChartDataResultType -from superset.common.query_actions import _get_drill_detail +from superset.common.chart_data_timing import ( + QueryAcquisitionResult, + QueryAcquisitionTiming, +) +from superset.common.query_actions import ( + _prepare_drill_detail_query, + _prepare_samples_query, + get_query_results, + get_query_results_with_timing, +) from superset.common.query_object import QueryObject +from superset.exceptions import QueryObjectValidationError from superset.utils.core import QueryObjectFilterClause -@patch("superset.common.query_actions._get_full") -def test_get_drill_detail_does_not_strip_filters( - mock_get_full: MagicMock, -) -> None: - """ - Characterization test for ``_get_drill_detail`` (superset/common/query_actions.py), - the backend transform behind the "Drill to Detail by" samples query. - - ``_get_drill_detail`` shallow-copies the incoming ``QueryObject`` and rewrites - ``is_timeseries``, ``metrics``, ``post_processing``, ``columns`` and ``orderby``; - it never reads or reassigns ``QueryObject.filter``. This test pins that - behavior down: filters present on the ``QueryObject`` handed to - ``_get_drill_detail`` are still present, unmodified, on the object it hands off - to ``_get_full``. - - This is deliberately narrow and does NOT reproduce or close #28562 ("Drill to - Detail by" ignoring a dashboard's applied filters) -- that report describes - filters missing from the query *before* it reaches this function, which points - at the payload assembled upstream of ``_get_drill_detail`` (dashboard native - filter propagation into the chart's form data / the drill payload built on the - frontend), not at this transform. If that assembly is ever changed to drop - filters before calling into this code path, this test would not catch it; it - only guards against a regression introduced inside ``_get_drill_detail`` itself. - """ +def test_prepare_drill_detail_query_does_not_strip_filters() -> None: + """Drill preparation keeps filters while rewriting the row query shape.""" applied_filter: QueryObjectFilterClause = { "col": "region", "op": "==", "val": "USA", } - - query_obj: QueryObject = QueryObject( + query_obj = QueryObject( columns=["region", "sales"], metrics=["count"], filters=[applied_filter], ) - col_region: MagicMock = MagicMock() - col_region.column_name = "region" - col_sales: MagicMock = MagicMock() - col_sales.column_name = "sales" - datasource = MagicMock() - datasource.columns = [col_region, col_sales] - - query_context: MagicMock = MagicMock() + datasource.columns = [ + MagicMock(column_name="region"), + MagicMock(column_name="sales"), + ] + query_context = MagicMock() query_context.datasource = datasource - query_context.result_type = ChartDataResultType.DRILL_DETAIL - captured: dict[str, QueryObject] = {} + prepared = _prepare_drill_detail_query(query_context, query_obj) - def _capture(_ctx: MagicMock, obj: QueryObject, _force: bool) -> dict[str, Any]: - captured["query_obj"] = obj - return {} - - mock_get_full.side_effect = _capture - - _get_drill_detail(query_context, query_obj) - - executed: QueryObject = captured["query_obj"] - assert applied_filter in executed.filter, ( - "_get_drill_detail unexpectedly stripped a filter it never touches; " - "this guards against a regression introduced in that function, not #28562." - ) + assert applied_filter in prepared.filter -@patch("superset.common.query_actions._get_full") -def test_get_samples_marks_query_as_system_sampling( - mock_get_full: MagicMock, -) -> None: - """ - ``_get_samples`` marks the copied query object as system-authored sampling - (so query generation may apply the engine's bounded-read override) without - mutating the caller's query object. - """ - from superset.common.query_actions import _get_samples - - query_obj: QueryObject = QueryObject(columns=["region"], metrics=["count"]) +def test_prepare_samples_query_marks_query_as_system_sampling() -> None: + """Sample preparation is isolated from the caller's query object.""" + query_obj = QueryObject(columns=["region"], metrics=["count"]) original_extras = query_obj.extras - - col_region: MagicMock = MagicMock() - col_region.column_name = "region" - datasource = MagicMock() - datasource.columns = [col_region] - - query_context: MagicMock = MagicMock() + datasource.columns = [MagicMock(column_name="region")] + query_context = MagicMock() query_context.datasource = datasource - query_context.result_type = ChartDataResultType.SAMPLES - captured: dict[str, QueryObject] = {} + prepared = _prepare_samples_query(query_context, query_obj) - def _capture(_ctx: MagicMock, obj: QueryObject, _force: bool) -> dict[str, Any]: - captured["query_obj"] = obj - return {} - - mock_get_full.side_effect = _capture - _get_samples(query_context, query_obj) - - assert captured["query_obj"].extras.get("system_sampling") is True - # the caller's query object is untouched (shallow copy must not leak) + assert prepared.extras.get("system_sampling") is True assert "system_sampling" not in query_obj.extras assert query_obj.extras is original_extras + + +def test_timed_dataframe_result_uses_sidecar_and_continuous_total() -> None: + query_context = MagicMock() + query_obj = MagicMock() + acquisition_timing = QueryAcquisitionTiming( + query_planning_ns=1, + cache_resolution_ns=2, + data_acquisition_ns=3, + payload_assembly_ns=4, + ) + query_context.get_df_payload_result.return_value = QueryAcquisitionResult( + payload={"df": "frame"}, + timing=acquisition_timing, + ) + + with ( + patch( + "superset.common.query_actions._materialize_full_payload", + return_value={"data": [{"col1": 1}]}, + ) as materialize_full_payload, + patch( + "superset.common.query_actions.time.perf_counter_ns", + side_effect=[100, 110, 120, 150], + ), + ): + result = get_query_results_with_timing( + ChartDataResultType.FULL, + query_context, + query_obj, + force_cached=False, + ) + + query_context.get_df_payload_result.assert_called_once_with( + query_obj, + force_cached=False, + ) + query_context.get_df_payload.assert_not_called() + materialize_full_payload.assert_called_once_with( + query_context, + query_obj, + {"df": "frame"}, + ) + assert result.payload == {"data": [{"col1": 1}]} + assert result.timing.query_planning_ns == 1 + assert result.timing.cache_resolution_ns == 2 + assert result.timing.data_acquisition_ns == 3 + assert result.timing.payload_assembly_ns == 14 + assert result.timing.total_ns == 50 + + +def test_metadata_result_has_null_phases_and_numeric_total() -> None: + query_context = MagicMock() + query_obj = MagicMock() + result_func = MagicMock(return_value={"language": "sql", "query": "SELECT 1"}) + + with ( + patch.dict( + "superset.common.query_actions._metadata_result_type_functions", + {ChartDataResultType.QUERY: result_func}, + clear=True, + ), + patch( + "superset.common.query_actions.time.perf_counter_ns", + side_effect=[100, 125], + ), + ): + result = get_query_results_with_timing( + ChartDataResultType.QUERY, + query_context, + query_obj, + force_cached=True, + ) + + result_func.assert_called_once_with(query_context, query_obj, True) + assert result.payload == {"language": "sql", "query": "SELECT 1"} + assert result.timing.query_planning_ns is None + assert result.timing.cache_resolution_ns is None + assert result.timing.data_acquisition_ns is None + assert result.timing.payload_assembly_ns is None + assert result.timing.total_ns == 25 + + +def test_result_type_dispatchers_are_complete_and_disjoint() -> None: + """Every result type is owned by exactly one timing execution path.""" + metadata_types = set(query_actions._metadata_result_type_functions) + data_types = set(query_actions._data_result_type_preparers) + + assert metadata_types.isdisjoint(data_types) + assert metadata_types | data_types == set(ChartDataResultType) + + +def test_timed_result_refuses_a_result_type_without_a_dispatch_owner() -> None: + with pytest.raises(QueryObjectValidationError, match="Invalid result type"): + get_query_results_with_timing( + cast("ChartDataResultType", "unowned"), + MagicMock(), + MagicMock(), + force_cached=False, + ) + + +@pytest.mark.parametrize( + "result_type", + [ChartDataResultType.SAMPLES, ChartDataResultType.DRILL_DETAIL], +) +def test_data_preparation_is_inside_the_continuous_total( + result_type: ChartDataResultType, +) -> None: + query_context = MagicMock() + query_obj = MagicMock() + preparation_started_ns: list[int] = [] + acquisition_timing = QueryAcquisitionTiming( + query_planning_ns=1, + cache_resolution_ns=2, + data_acquisition_ns=3, + payload_assembly_ns=4, + ) + + def preparer(_query_context: MagicMock, prepared_query: MagicMock) -> MagicMock: + preparation_started_ns.append(query_actions.time.perf_counter_ns()) + return prepared_query + + with ( + patch.dict( + "superset.common.query_actions._data_result_type_preparers", + {result_type: preparer}, + ), + patch( + "superset.common.query_actions._get_full_with_timing", + return_value=( + {"data": []}, + acquisition_timing, + 5, + ), + ), + patch( + "superset.common.query_actions.time.perf_counter_ns", + side_effect=[100, 125, 150], + ), + ): + result = get_query_results_with_timing( + result_type, + query_context, + query_obj, + force_cached=False, + ) + + assert preparation_started_ns == [125] + assert result.timing.total_ns == 50 + + +def test_timed_drill_detail_keeps_capability_refusal() -> None: + query_context = MagicMock() + query_context.datasource = MagicMock(supports_drill_to_detail=False) + query_obj = MagicMock() + query_obj.datasource = None + + with pytest.raises(QueryObjectValidationError): + get_query_results_with_timing( + ChartDataResultType.DRILL_DETAIL, + query_context, + query_obj, + force_cached=False, + ) + + +def test_legacy_result_wrapper_keeps_drill_detail_capability_refusal() -> None: + query_context = MagicMock() + query_context.datasource = MagicMock(supports_drill_to_detail=False) + query_obj = MagicMock() + query_obj.datasource = None + + with pytest.raises(QueryObjectValidationError): + get_query_results( + ChartDataResultType.DRILL_DETAIL, + query_context, + query_obj, + force_cached=False, + ) + + +def test_legacy_result_wrapper_delegates_to_timed_resolver() -> None: + query_context = MagicMock() + query_obj = MagicMock() + + with patch( + "superset.common.query_actions.get_query_results_with_timing" + ) as timed_resolver: + timed_resolver.return_value.payload = {"data": []} + + result = get_query_results( + ChartDataResultType.FULL, + query_context, + query_obj, + force_cached=False, + ) + + assert result == {"data": []} + timed_resolver.assert_called_once_with( + ChartDataResultType.FULL, + query_context, + query_obj, + False, + ) diff --git a/tests/unit_tests/common/test_query_actions_drill_detail.py b/tests/unit_tests/common/test_query_actions_drill_detail.py index 0bd95aa8225..cf3a299ec56 100644 --- a/tests/unit_tests/common/test_query_actions_drill_detail.py +++ b/tests/unit_tests/common/test_query_actions_drill_detail.py @@ -14,19 +14,24 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from unittest.mock import MagicMock, patch +from types import SimpleNamespace +from typing import cast, TYPE_CHECKING +from unittest.mock import MagicMock import pytest -from superset.common.query_actions import _get_drill_detail +from superset.common.query_actions import _prepare_drill_detail_query from superset.exceptions import QueryObjectValidationError +if TYPE_CHECKING: + from superset.common.query_object import QueryObject -def test_get_drill_detail_refuses_datasource_that_opts_out() -> None: + +def test_prepare_drill_detail_query_refuses_datasource_that_opts_out() -> None: """ A datasource with ``supports_drill_to_detail = False`` (e.g. semantic views) must be hard-blocked on the server. Without this gate the request - would fall through to ``_get_full`` and fail with an opaque error, and + would fall through to dataframe acquisition and fail with an opaque error, and the flag would only be enforced by the frontend menu — leaving the chart-data API endpoint accepting drill-detail requests it shouldn't. """ @@ -42,29 +47,35 @@ def test_get_drill_detail_refuses_datasource_that_opts_out() -> None: QueryObjectValidationError, match="Drill to detail is not available", ): - _get_drill_detail(query_context, query_obj) + _prepare_drill_detail_query(query_context, query_obj) -def test_get_drill_detail_allows_datasource_without_flag() -> None: +def test_prepare_drill_detail_query_allows_datasource_without_flag() -> None: """ Datasources that don't declare the flag (e.g. legacy ``SqlaTable`` subclasses via ``getattr`` default) must continue to work — the gate only fires when the flag is explicitly ``False``. """ - datasource = MagicMock(spec=["columns"]) - column = MagicMock() - column.column_name = "id" - datasource.columns = [column] - - query_obj = MagicMock() - query_obj.datasource = datasource - query_obj.columns = [] - + datasource = SimpleNamespace(columns=[SimpleNamespace(column_name="id")]) + query_obj = SimpleNamespace( + columns=[], + datasource=datasource, + is_timeseries=True, + metrics=["count"], + orderby=[("name", False)], + post_processing=[{"operation": "pivot"}], + ) query_context = MagicMock() - expected_payload: dict[str, list[dict[str, str]]] = {"data": []} - with patch( - "superset.common.query_actions._get_full", return_value=expected_payload - ) as mock_get_full: - assert _get_drill_detail(query_context, query_obj) is expected_payload - mock_get_full.assert_called_once() + prepared_query = _prepare_drill_detail_query( + query_context, + cast("QueryObject", query_obj), + ) + + assert prepared_query is not query_obj + assert prepared_query.is_timeseries is False + assert prepared_query.metrics is None + assert prepared_query.post_processing == [] + assert prepared_query.columns == ["id"] + assert prepared_query.orderby == [("id", True)] + assert query_obj.columns == [] diff --git a/tests/unit_tests/common/test_query_context_processor.py b/tests/unit_tests/common/test_query_context_processor.py index 88c786ed1da..d81dd495e3a 100644 --- a/tests/unit_tests/common/test_query_context_processor.py +++ b/tests/unit_tests/common/test_query_context_processor.py @@ -24,6 +24,7 @@ import pandas as pd import pytest from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType +from superset.common.chart_data_timing import QueryDataResult, QueryTiming from superset.common.db_query_status import QueryStatus from superset.common.query_context_processor import QueryContextProcessor from superset.utils.core import GenericDataType @@ -38,6 +39,16 @@ def mock_query_context(): yield mock_query_context_processor +def _query_timing() -> QueryTiming: + return QueryTiming( + query_planning_ns=0, + cache_resolution_ns=0, + data_acquisition_ns=None, + payload_assembly_ns=0, + total_ns=0, + ) + + @pytest.fixture def processor(mock_query_context): from superset.models.helpers import ExploreMixin @@ -1426,8 +1437,8 @@ def test_ensure_totals_available_updates_cache_values(): # Now call get_payload which should update cache_values with patch( - "superset.common.query_context_processor.get_query_results" - ) as mock_get_query_results: + "superset.common.query_context_processor.get_query_results_with_timing" + ) as mock_get_query_results_with_timing: # Mock the query results mock_query_results_response = [ { @@ -1435,7 +1446,10 @@ def test_ensure_totals_available_updates_cache_values(): "query": "SELECT ...", } ] - mock_get_query_results.return_value = mock_query_results_response + mock_get_query_results_with_timing.return_value = QueryDataResult( + payload=mock_query_results_response[0], + timing=_query_timing(), + ) # Mock cache manager to avoid actual caching with patch( @@ -1650,15 +1664,18 @@ def test_cache_values_sync_after_ensure_totals_available(): # Mock the query results with patch( - "superset.common.query_context_processor.get_query_results" - ) as mock_get_query_results: + "superset.common.query_context_processor.get_query_results_with_timing" + ) as mock_get_query_results_with_timing: mock_query_results_response = [ { "data": [{"region": "North", "sales": 100}], "query": "SELECT region, SUM(sales) FROM table GROUP BY region", } ] - mock_get_query_results.return_value = mock_query_results_response + mock_get_query_results_with_timing.return_value = QueryDataResult( + payload=mock_query_results_response[0], + timing=_query_timing(), + ) # Call get_payload - this internally calls ensure_totals_available() # and then should update cache_values @@ -1896,6 +1913,7 @@ def test_force_cached_normalizes_totals_query_row_limit(): processor = QueryContextProcessor(mock_query_context) processor._qc_datasource = mock_datasource mock_query_context.get_df_payload = processor.get_df_payload + mock_query_context.get_df_payload_result = processor.get_df_payload_result mock_query_context.get_data = processor.get_data with patch( diff --git a/tests/unit_tests/common/test_query_context_processor_timing.py b/tests/unit_tests/common/test_query_context_processor_timing.py new file mode 100644 index 00000000000..cef8365aa51 --- /dev/null +++ b/tests/unit_tests/common/test_query_context_processor_timing.py @@ -0,0 +1,266 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from unittest.mock import MagicMock, patch + +import pandas as pd + +from superset.common.chart_data_timing import QueryDataResult, QueryTiming +from superset.common.db_query_status import QueryStatus +from superset.common.query_context_processor import QueryContextProcessor + + +def _query_timing() -> QueryTiming: + return QueryTiming( + query_planning_ns=1_000_000, + cache_resolution_ns=2_000_000, + data_acquisition_ns=3_000_000, + payload_assembly_ns=4_000_000, + total_ns=12_000_000, + ) + + +def _query_obj() -> MagicMock: + query_obj = MagicMock() + query_obj.columns = ["col1"] + query_obj.column_names = ["col1"] + query_obj.metrics = [] + query_obj.metric_names = [] + query_obj.from_dttm = None + query_obj.to_dttm = None + query_obj.annotation_layers = [] + query_obj.filter = [] + return query_obj + + +def _processor() -> QueryContextProcessor: + query_context = MagicMock() + query_context.force = False + query_context.form_data = {} + query_context.cache_values = {"queries": [{}]} + query_context.queries = [_query_obj()] + + processor = QueryContextProcessor.__new__(QueryContextProcessor) + processor._query_context = query_context + processor._qc_datasource = MagicMock() + processor._qc_datasource.uid = "test_uid" + processor._qc_datasource.column_names = ["col1"] + processor._qc_datasource.data = {} + return processor + + +def test_public_projection_is_explicit_and_versioned() -> None: + assert _query_timing().as_public_dict() == { + "version": 1, + "query": { + "query_planning_ms": 1.0, + "cache_resolution_ms": 2.0, + "data_acquisition_ms": 3.0, + "payload_assembly_ms": 4.0, + "total_ms": 12.0, + }, + } + + +def test_public_projection_uses_null_for_non_applicable_phases() -> None: + timing = QueryTiming( + query_planning_ns=None, + cache_resolution_ns=None, + data_acquisition_ns=None, + payload_assembly_ns=None, + total_ns=1, + ) + + assert timing.as_public_dict()["query"] == { + "query_planning_ms": None, + "cache_resolution_ms": None, + "data_acquisition_ms": None, + "payload_assembly_ms": None, + "total_ms": 0.0, + } + + +@patch("superset.common.query_context_processor.QueryCacheManager") +def test_dataframe_payload_result_keeps_timing_outside_payload( + cache_manager: MagicMock, +) -> None: + cache = MagicMock() + cache.is_loaded = True + cache.is_cached = True + cache.df = pd.DataFrame({"col1": [1]}) + cache.cache_dttm = "2026-01-01T00:00:00" + cache.queried_dttm = "2026-01-01T00:00:00" + cache.applied_template_filters = [] + cache.applied_filter_columns = [] + cache.rejected_filter_columns = [] + cache.annotation_data = {} + cache.error_message = None + cache.query = "SELECT 1" + cache.status = "success" + cache.stacktrace = None + cache.sql_rowcount = 1 + cache.bq_memory_limited = False + cache.bq_memory_limited_row_count = 0 + cache_manager.get.return_value = cache + + processor = _processor() + with ( + patch.object(processor, "query_cache_key", return_value="key"), + patch.object(processor, "get_cache_timeout", return_value=300), + ): + result = processor.get_df_payload_result(_query_obj()) + + assert "timing" not in result.payload + assert result.payload["query"] == "SELECT 1" + assert result.timing.query_planning_ns >= 0 + assert result.timing.cache_resolution_ns >= 0 + assert result.timing.data_acquisition_ns is None + assert result.timing.payload_assembly_ns >= 0 + + +@patch("superset.common.query_context_processor.QueryCacheManager") +def test_dataframe_payload_result_measures_non_overlapping_stages( + cache_manager: MagicMock, +) -> None: + cache = MagicMock() + cache.is_loaded = False + cache.df = pd.DataFrame({"col1": [1]}) + cache.cache_dttm = None + cache.queried_dttm = "2026-01-01T00:00:00" + cache.applied_template_filters = [] + cache.applied_filter_columns = [] + cache.rejected_filter_columns = [] + cache.annotation_data = {} + cache.error_message = None + cache.query = "SELECT 1" + cache.status = QueryStatus.SUCCESS + cache.stacktrace = None + cache.sql_rowcount = 1 + cache.bq_memory_limited = False + cache.bq_memory_limited_row_count = 0 + cache_manager.get.return_value = cache + + processor = _processor() + with ( + patch.object(processor, "query_cache_key", return_value="key"), + patch.object(processor, "get_cache_timeout", return_value=300), + patch.object(processor, "get_query_result", return_value=MagicMock()), + patch.object(processor, "get_annotation_data", return_value={}), + patch( + "superset.common.query_context_processor.time.perf_counter_ns", + side_effect=[100, 110, 120, 130, 140, 170, 180, 195], + ), + ): + result = processor.get_df_payload_result(_query_obj()) + + assert result.timing.query_planning_ns == 10 + assert result.timing.cache_resolution_ns == 10 + assert result.timing.data_acquisition_ns == 30 + assert result.timing.payload_assembly_ns == 15 + cache.set_query_result.assert_called_once() + + +@patch("superset.common.query_context_processor.QueryCacheManager") +def test_cache_resolution_includes_loaded_cache_compatibility_policy( + cache_manager: MagicMock, +) -> None: + """A legacy cached value becomes a miss before cache resolution completes.""" + cache = MagicMock() + cache.is_loaded = True + cache.is_cached = True + cache.df = pd.DataFrame({"col1": [1]}) + cache.cache_dttm = "2026-01-01T00:00:00" + cache.queried_dttm = "2026-01-01T00:00:00" + cache.applied_template_filters = [] + cache.applied_filter_columns = [] + cache.rejected_filter_columns = [] + cache.annotation_data = {} + cache.error_message = None + cache.query = "SELECT 1" + cache.status = QueryStatus.SUCCESS + cache.stacktrace = None + cache.sql_rowcount = 1 + cache.bq_memory_limited = False + cache.bq_memory_limited_row_count = 0 + cache_manager.get.return_value = cache + + query_obj = _query_obj() + query_obj.filter = [{"col": "col1", "op": "==", "val": 1}] + processor = _processor() + clock_values = iter((100, 110, 120, 130, 140, 170, 180, 195)) + clock_calls = 0 + + def perf_counter_ns() -> int: + nonlocal clock_calls + clock_calls += 1 + if clock_calls == 4: + assert cache.is_loaded is False + return next(clock_values) + + with ( + patch.object(processor, "query_cache_key", return_value="key"), + patch.object(processor, "get_cache_timeout", return_value=300), + patch.object(processor, "get_query_result", return_value=MagicMock()), + patch.object(processor, "get_annotation_data", return_value={}), + patch( + "superset.common.query_context_processor.time.perf_counter_ns", + side_effect=perf_counter_ns, + ), + ): + result = processor.get_df_payload_result(query_obj) + + assert result.timing.cache_resolution_ns == 10 + assert result.timing.data_acquisition_ns == 30 + cache.set_query_result.assert_called_once() + + +def test_get_payload_preserves_legacy_shape_without_timing() -> None: + processor = _processor() + query_payload = {"data": [{"col1": 1}]} + + with ( + patch.object( + processor, "_prepare_contribution_totals", return_value=([], None) + ), + patch( + "superset.common.query_context_processor.get_query_results_with_timing", + return_value=QueryDataResult(query_payload, _query_timing()), + ), + ): + result = processor.get_payload() + + assert result == {"queries": [query_payload]} + assert "timing" not in result["queries"][0] + + +def test_get_payload_result_keeps_timing_sidecar() -> None: + processor = _processor() + query_payload = {"data": [{"col1": 1}]} + + with ( + patch.object( + processor, "_prepare_contribution_totals", return_value=([], None) + ), + patch( + "superset.common.query_context_processor.get_query_results_with_timing", + return_value=QueryDataResult(query_payload, _query_timing()), + ), + ): + result = processor.get_payload_result() + + assert result.queries[0].payload == query_payload + assert result.queries[0].timing.total_ns == 12_000_000 + assert "timing" not in result.queries[0].payload