tests/unit_tests/common/test_query_context_processor.py

# 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 datetime import datetime, timedelta
from typing import Any
from unittest.mock import MagicMock, patch

import numpy as np
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.exceptions import QueryObjectValidationError
from superset.utils.core import GenericDataType
from superset.utils.date_parser import get_past_or_future


@pytest.fixture
def mock_query_context():
    with patch(
        "superset.common.query_context_processor.QueryContextProcessor"
    ) as mock_query_context_processor:
        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

    mock_query_context.datasource.data = MagicMock()
    mock_query_context.datasource.data.get.return_value = {
        "col1": "Column 1",
        "col2": "Column 2",
    }

    # Create a processor instance
    processor = QueryContextProcessor(mock_query_context)

    # Setup datasource methods from ExploreMixin to be real methods
    # by binding them to the mock datasource
    processor._qc_datasource.is_valid_date_range = (
        ExploreMixin.is_valid_date_range.__get__(processor._qc_datasource)
    )
    processor._qc_datasource.is_valid_date = ExploreMixin.is_valid_date.__get__(
        processor._qc_datasource
    )
    processor._qc_datasource.get_offset_custom_or_inherit = (
        ExploreMixin.get_offset_custom_or_inherit.__get__(processor._qc_datasource)
    )
    processor._qc_datasource._get_temporal_column_for_filter = (
        ExploreMixin._get_temporal_column_for_filter.__get__(processor._qc_datasource)
    )
    processor._qc_datasource.join_offset_dfs = ExploreMixin.join_offset_dfs.__get__(
        processor._qc_datasource
    )
    processor._qc_datasource._determine_join_keys = (
        ExploreMixin._determine_join_keys.__get__(processor._qc_datasource)
    )
    processor._qc_datasource._process_date_range_offset = (
        ExploreMixin._process_date_range_offset.__get__(processor._qc_datasource)
    )
    processor._qc_datasource._perform_join = ExploreMixin._perform_join.__get__(
        processor._qc_datasource
    )
    processor._qc_datasource._apply_cleanup_logic = (
        ExploreMixin._apply_cleanup_logic.__get__(processor._qc_datasource)
    )
    processor._qc_datasource.add_offset_join_column = (
        ExploreMixin.add_offset_join_column.__get__(processor._qc_datasource)
    )

    return processor


def test_query_cache_key_binds_annotation_data_to_requesting_user(processor):
    """The cache key for annotated queries must differ per requesting user."""
    query_obj = MagicMock()
    query_obj.annotation_layers = [{"sourceType": "NATIVE", "name": "a", "value": 1}]
    with (
        patch(
            "superset.common.query_context_processor.get_user_id",
            side_effect=[1, 2],
        ),
        patch("superset.common.query_context_processor.security_manager"),
    ):
        processor.query_cache_key(query_obj)
        processor.query_cache_key(query_obj)
    contexts = [
        call.kwargs["annotation_context"] for call in query_obj.cache_key.call_args_list
    ]
    assert contexts[0] != contexts[1]


def test_get_data_table_like(processor, mock_query_context):
    df = pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]})
    coltypes = [GenericDataType.NUMERIC, GenericDataType.STRING]
    mock_query_context.result_format = ChartDataResultFormat.JSON

    result = processor.get_data(df, coltypes)
    expected = [
        {"col1": 1, "col2": "a"},
        {"col1": 2, "col2": "b"},
        {"col1": 3, "col2": "c"},
    ]
    assert result == expected


@patch("superset.common.query_context_processor.csv.df_to_escaped_csv")
def test_get_data_csv(mock_df_to_escaped_csv, processor, mock_query_context):
    df = pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]})
    coltypes = [GenericDataType.NUMERIC, GenericDataType.STRING]
    mock_query_context.result_format = ChartDataResultFormat.CSV

    mock_df_to_escaped_csv.return_value = "col1,col2\n1,a\n2,b\n3,c\n"
    result = processor.get_data(df, coltypes)
    # CSV output is encoded to bytes using the CSV_EXPORT encoding so dashboard
    # chart exports honor the configured encoding (e.g. the utf-8-sig BOM).
    assert result == "col1,col2\n1,a\n2,b\n3,c\n".encode("utf-8-sig")
    mock_df_to_escaped_csv.assert_called_once_with(
        df, index=False, encoding="utf-8-sig"
    )


@patch("superset.common.query_context_processor.excel.df_to_excel")
@patch("superset.common.query_context_processor.excel.apply_column_types")
def test_get_data_xlsx(
    mock_apply_column_types, mock_df_to_excel, processor, mock_query_context
):
    df = pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]})
    coltypes = [GenericDataType.NUMERIC, GenericDataType.STRING]
    mock_query_context.result_format = ChartDataResultFormat.XLSX

    mock_df_to_excel.return_value = b"binary data"
    result = processor.get_data(df, coltypes)
    assert result == b"binary data"
    mock_apply_column_types.assert_called_once_with(df, coltypes)
    mock_df_to_excel.assert_called_once_with(df, index=False)


def test_get_data_json(processor, mock_query_context):
    df = pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]})
    coltypes = [GenericDataType.NUMERIC, GenericDataType.STRING]
    mock_query_context.result_format = ChartDataResultFormat.JSON

    result = processor.get_data(df, coltypes)
    expected = [
        {"col1": 1, "col2": "a"},
        {"col1": 2, "col2": "b"},
        {"col1": 3, "col2": "c"},
    ]
    assert result == expected


def test_get_data_invalid_dataframe(processor, mock_query_context):
    df = pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]})
    coltypes = [GenericDataType.NUMERIC, GenericDataType.STRING]
    mock_query_context.result_format = ChartDataResultFormat.JSON

    with patch.object(df, "to_dict", side_effect=ValueError("Invalid DataFrame")):
        with pytest.raises(ValueError, match="Invalid DataFrame"):
            processor.get_data(df, coltypes)


def test_get_data_non_unique_columns(processor, mock_query_context):
    data = [[1, "a"], [2, "b"], [3, "c"]]
    df = pd.DataFrame(data, columns=["col1", "col1"])
    coltypes = [GenericDataType.NUMERIC, GenericDataType.STRING]
    mock_query_context.result_format = ChartDataResultFormat.JSON

    with pytest.warns(
        UserWarning,
        match="DataFrame columns are not unique, some columns will be omitted",
    ):
        processor.get_data(df, coltypes)


def test_get_data_empty_dataframe_json(processor, mock_query_context):
    df = pd.DataFrame(columns=["col1", "col2"])
    coltypes = [GenericDataType.NUMERIC, GenericDataType.STRING]
    mock_query_context.result_format = ChartDataResultFormat.JSON
    result = processor.get_data(df, coltypes)
    assert result == []


@patch("superset.common.query_context_processor.csv.df_to_escaped_csv")
def test_get_data_empty_dataframe_csv(
    mock_df_to_escaped_csv, processor, mock_query_context
):
    df = pd.DataFrame(columns=["col1", "col2"])
    coltypes = [GenericDataType.NUMERIC, GenericDataType.STRING]
    mock_query_context.result_format = ChartDataResultFormat.CSV
    mock_df_to_escaped_csv.return_value = "col1,col2\n"
    result = processor.get_data(df, coltypes)
    assert result == "col1,col2\n".encode("utf-8-sig")
    mock_df_to_escaped_csv.assert_called_once_with(
        df, index=False, encoding="utf-8-sig"
    )


@patch("superset.common.query_context_processor.excel.df_to_excel")
@patch("superset.common.query_context_processor.excel.apply_column_types")
def test_get_data_empty_dataframe_xlsx(
    mock_apply_column_types, mock_df_to_excel, processor, mock_query_context
):
    df = pd.DataFrame(columns=["col1", "col2"])
    coltypes = [GenericDataType.NUMERIC, GenericDataType.STRING]
    mock_query_context.result_format = ChartDataResultFormat.XLSX
    mock_df_to_excel.return_value = b"binary data empty"
    result = processor.get_data(df, coltypes)
    assert result == b"binary data empty"
    mock_apply_column_types.assert_called_once_with(df, coltypes)
    mock_df_to_excel.assert_called_once_with(df, index=False)


def test_get_data_nan_values_json(processor, mock_query_context):
    df = pd.DataFrame({"col1": [1, np.nan, 3], "col2": ["a", "b", "c"]})
    coltypes = [GenericDataType.NUMERIC, GenericDataType.STRING]
    mock_query_context.result_format = ChartDataResultFormat.JSON
    result = processor.get_data(df, coltypes)
    assert result[0]["col1"] == 1
    assert pd.isna(result[1]["col1"])
    assert result[2]["col1"] == 3


def test_get_data_invalid_input(processor, mock_query_context):
    df = "not a dataframe"
    coltypes = [GenericDataType.NUMERIC, GenericDataType.STRING]
    mock_query_context.result_format = ChartDataResultFormat.JSON
    with pytest.raises(AttributeError):
        processor.get_data(df, coltypes)


def test_get_data_default_format_when_result_format_is_none(
    processor, mock_query_context
):
    df = pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]})
    coltypes = [GenericDataType.NUMERIC, GenericDataType.STRING]
    mock_query_context.result_format = None
    result = processor.get_data(df, coltypes)
    expected = [
        {"col1": 1, "col2": "a"},
        {"col1": 2, "col2": "b"},
        {"col1": 3, "col2": "c"},
    ]
    assert result == expected


def fake_apply_column_types(df, coltypes):
    if len(coltypes) != len(df.columns):
        raise ValueError("Mismatch between column types and dataframe columns")
    return df


@patch("superset.common.query_context_processor.excel.df_to_excel")
@patch(
    "superset.common.query_context_processor.excel.apply_column_types",
    side_effect=fake_apply_column_types,
)
def test_get_data_invalid_coltypes_length_xlsx(
    mock_apply_column_types, mock_df_to_excel, processor, mock_query_context
):
    df = pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]})
    coltypes = [GenericDataType.NUMERIC]  # Mismatched length
    mock_query_context.result_format = ChartDataResultFormat.XLSX
    with pytest.raises(
        ValueError, match="Mismatch between column types and dataframe columns"
    ):
        processor.get_data(df, coltypes)


def test_get_data_does_not_mutate_dataframe(processor, mock_query_context):
    df = pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]})
    original_df = df.copy(deep=True)
    coltypes = [GenericDataType.NUMERIC, GenericDataType.STRING]
    mock_query_context.result_format = ChartDataResultFormat.JSON
    _ = processor.get_data(df, coltypes)
    pd.testing.assert_frame_equal(df, original_df)


@patch(
    "superset.common.query_context_processor.excel.apply_column_types",
    side_effect=ValueError("Conversion error"),
)
def test_get_data_xlsx_apply_column_types_error(
    mock_apply_column_types, processor, mock_query_context
):
    df = pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]})
    coltypes = [GenericDataType.NUMERIC, GenericDataType.STRING]
    mock_query_context.result_format = ChartDataResultFormat.XLSX
    with pytest.raises(ValueError, match="Conversion error"):
        processor.get_data(df, coltypes)


def test_is_valid_date_range_format(processor):
    """Test that date range format validation works correctly."""
    # Should return True for valid date range format
    assert (
        processor._qc_datasource.is_valid_date_range("2023-01-01 : 2023-01-31") is True
    )
    assert (
        processor._qc_datasource.is_valid_date_range("2020-12-25 : 2020-12-31") is True
    )

    # Should return False for invalid format
    assert processor._qc_datasource.is_valid_date_range("1 day ago") is False
    assert processor._qc_datasource.is_valid_date_range("2023-01-01") is False
    assert processor._qc_datasource.is_valid_date_range("invalid") is False


def test_is_valid_date_range_static_format():
    """Test that static date range format validation works correctly."""
    from superset.models.helpers import ExploreMixin

    # Should return True for valid date range format
    assert ExploreMixin.is_valid_date_range_static("2023-01-01 : 2023-01-31") is True
    assert ExploreMixin.is_valid_date_range_static("2020-12-25 : 2020-12-31") is True

    # Should return False for invalid format
    assert ExploreMixin.is_valid_date_range_static("1 day ago") is False
    assert ExploreMixin.is_valid_date_range_static("2023-01-01") is False
    assert ExploreMixin.is_valid_date_range_static("invalid") is False


def test_processing_time_offsets_date_range_logic(processor):
    """Test that date range timeshift logic works correctly with feature flag checks."""
    from superset.models.helpers import ExploreMixin

    # Test that the date range validation works
    assert (
        processor._qc_datasource.is_valid_date_range("2023-01-01 : 2023-01-31") is True
    )
    assert processor._qc_datasource.is_valid_date_range("1 year ago") is False

    # Test that static method also works
    assert ExploreMixin.is_valid_date_range_static("2023-01-01 : 2023-01-31") is True
    assert ExploreMixin.is_valid_date_range_static("1 year ago") is False


def test_feature_flag_validation_logic():
    """Test that feature flag validation logic works as expected."""
    from superset.extensions import feature_flag_manager

    # This tests the concept - actual feature flag value depends on config
    # The important thing is that the code checks for DATE_RANGE_TIMESHIFTS_ENABLED
    flag_name = "DATE_RANGE_TIMESHIFTS_ENABLED"

    # Test that the feature flag is being checked
    # (This will vary based on actual config but tests the mechanism)
    result = feature_flag_manager.is_feature_enabled(flag_name)
    assert isinstance(result, bool)  # Should return a boolean


def test_join_offset_dfs_date_range_basic(processor):
    """Test basic join logic for date range offsets."""
    # Create simple test data
    main_df = pd.DataFrame({"dim1": ["A", "B", "C"], "metric1": [10, 20, 30]})

    offset_df = pd.DataFrame({"dim1": ["A", "B", "C"], "metric1": [5, 10, 15]})

    # Mock query context
    mock_query = MagicMock()
    mock_query.granularity = "date_col"
    processor._query_context.queries = [mock_query]

    # Test basic join with date range offset
    offset_dfs = {"2023-01-01 : 2023-01-31": offset_df}
    join_keys = ["dim1"]

    with patch("superset.models.helpers.feature_flag_manager") as mock_ff:
        mock_ff.is_feature_enabled.return_value = True
        with patch("superset.common.utils.dataframe_utils.left_join_df") as mock_join:
            mock_join.return_value = pd.DataFrame(
                {
                    "dim1": ["A", "B", "C"],
                    "metric1": [10, 20, 30],
                    "metric1 2023-01-01 : 2023-01-31": [5, 10, 15],
                }
            )

            result_df = processor._qc_datasource.join_offset_dfs(
                main_df, offset_dfs, time_grain=None, join_keys=join_keys
            )

            # Verify join was called
            mock_join.assert_called_once()
            assert len(result_df) == 3


def test_get_offset_custom_or_inherit_with_inherit(processor):
    """Test get_offset_custom_or_inherit with 'inherit' option."""
    from_dttm = pd.Timestamp("2024-01-01")
    to_dttm = pd.Timestamp("2024-01-10")

    result = processor._qc_datasource.get_offset_custom_or_inherit(
        "inherit", from_dttm, to_dttm
    )

    # Should shift back by the length of the range: 9 days, in seconds
    assert result == "777600 seconds ago"


@pytest.mark.parametrize(
    "hours",
    [
        # Sub-day ranges truncated to "0 days ago" and compared the range
        # against itself.
        1,
        12,
        # Ranges that are not a whole number of days truncated downwards and
        # overlapped the range they were compared against.
        36,
        # Whole-day ranges were already correct and must stay so.
        24,
        48,
    ],
)
def test_get_offset_custom_or_inherit_shifts_by_full_range(
    processor, hours: int
) -> None:
    """
    'inherit' compares a range against the period immediately preceding it,
    for any range length. The offset must not be truncated to whole days: it
    has to reproduce the range's exact duration, so that the comparison window
    neither overlaps the range nor leaves a gap before it.
    """
    from_dttm = datetime(2024, 3, 10, 6, 0)
    to_dttm = from_dttm + timedelta(hours=hours)

    offset = processor._qc_datasource.get_offset_custom_or_inherit(
        "inherit", from_dttm, to_dttm
    )

    shifted_from = get_past_or_future(offset, from_dttm)
    shifted_to = get_past_or_future(offset, to_dttm)
    assert shifted_to == from_dttm
    assert shifted_to - shifted_from == to_dttm - from_dttm


def test_get_offset_custom_or_inherit_with_date(processor):
    """Test get_offset_custom_or_inherit with specific date."""
    from_dttm = pd.Timestamp("2024-01-10")
    to_dttm = pd.Timestamp("2024-01-20")

    result = processor._qc_datasource.get_offset_custom_or_inherit(
        "2024-01-05", from_dttm, to_dttm
    )

    # Should return difference between from_dttm and the specified date
    assert result == "5 days ago"


def test_get_offset_custom_or_inherit_with_invalid_date(processor):
    """Test get_offset_custom_or_inherit with invalid date format."""
    from_dttm = pd.Timestamp("2024-01-10")
    to_dttm = pd.Timestamp("2024-01-20")

    result = processor._qc_datasource.get_offset_custom_or_inherit(
        "invalid-date", from_dttm, to_dttm
    )

    # Should return empty string for invalid format
    assert result == ""


def test_get_temporal_column_for_filter_with_granularity(processor):
    """Test _get_temporal_column_for_filter returns granularity when available."""
    query_object = MagicMock()
    query_object.granularity = "date_column"

    result = processor._qc_datasource._get_temporal_column_for_filter(
        query_object, "x_axis_col"
    )

    assert result == "date_column"


def test_get_temporal_column_for_filter_with_x_axis_fallback(processor):
    """Test _get_temporal_column_for_filter falls back to x_axis_label."""
    query_object = MagicMock()
    query_object.granularity = None

    result = processor._qc_datasource._get_temporal_column_for_filter(
        query_object, "x_axis_col"
    )

    assert result == "x_axis_col"


def test_get_temporal_column_for_filter_with_datasource_columns(processor):
    """Test _get_temporal_column_for_filter
    returns None when no clear temporal column."""
    query_object = MagicMock()
    query_object.granularity = None
    query_object.filter = []

    mock_datetime_col = MagicMock()
    mock_datetime_col.is_dttm = True
    mock_datetime_col.column_name = "created_at"

    mock_regular_col = MagicMock()
    mock_regular_col.is_dttm = False
    mock_regular_col.column_name = "name"

    processor._qc_datasource.columns = [mock_regular_col, mock_datetime_col]

    result = processor._qc_datasource._get_temporal_column_for_filter(
        query_object, None
    )

    assert result is None


def test_get_temporal_column_for_filter_prefers_granularity(processor):
    """Test _get_temporal_column_for_filter uses granularity when available."""
    query_object = MagicMock()
    query_object.granularity = "timestamp_col"
    query_object.filter = []

    mock_datetime_col = MagicMock()
    mock_datetime_col.is_dttm = True
    mock_datetime_col.name = "other_col"

    processor._qc_datasource.columns = [mock_datetime_col]

    result = processor._qc_datasource._get_temporal_column_for_filter(
        query_object, None
    )

    assert result == "timestamp_col"


def test_get_temporal_column_for_filter_no_columns_found(processor):
    """Test _get_temporal_column_for_filter
    returns None when no temporal column found."""
    query_object = MagicMock()
    query_object.granularity = None

    # Mock datasource with no datetime columns
    mock_regular_col = MagicMock()
    mock_regular_col.is_dttm = False
    mock_regular_col.column_name = "name"

    processor._qc_datasource.columns = [mock_regular_col]

    result = processor._qc_datasource._get_temporal_column_for_filter(
        query_object, None
    )

    assert result is None


def test_get_temporal_column_for_filter_no_datasource_columns(processor):
    """Test _get_temporal_column_for_filter handles datasource
    without columns attribute."""
    query_object = MagicMock()
    query_object.granularity = None

    # Remove columns attribute from datasource
    if hasattr(processor._qc_datasource, "columns"):
        delattr(processor._qc_datasource, "columns")

    result = processor._qc_datasource._get_temporal_column_for_filter(
        query_object, None
    )

    assert result is None


def test_processing_time_offsets_temporal_column_error(processor):
    """Test processing_time_offsets raises QueryObjectValidationError
    when temporal column can't be determined."""
    from superset.common.query_object import QueryObject
    from superset.exceptions import QueryObjectValidationError

    # Create a dataframe for testing
    df = pd.DataFrame({"dim1": ["A", "B", "C"], "metric1": [10, 20, 30]})

    # Create query object with date range offset and proper time range
    query_object = QueryObject(
        datasource=MagicMock(),
        granularity=None,  # No granularity set
        columns=[],
        is_timeseries=True,
        time_offsets=["2023-01-01 : 2023-01-31"],
        filter=[
            {
                "col": "some_date_col",
                "op": "TEMPORAL_RANGE",
                "val": "2024-01-01 : 2024-01-31",
            }
        ],
    )

    # Mock get_since_until_from_query_object to return valid dates
    with patch(
        "superset.common.utils.time_range_utils.get_since_until_from_query_object"
    ) as mock_dates:
        mock_dates.return_value = (
            pd.Timestamp("2024-01-01"),
            pd.Timestamp("2024-01-31"),
        )

        # Mock feature flag to be enabled
        with patch("superset.models.helpers.feature_flag_manager") as mock_ff:
            mock_ff.is_feature_enabled.return_value = True

            # Mock _get_temporal_column_for_filter to return None
            # (no temporal column found)
            with patch.object(
                processor._qc_datasource,
                "_get_temporal_column_for_filter",
                return_value=None,
            ):
                # Mock the datasource's processing_time_offsets to raise the error
                def raise_error(*args, **kwargs):
                    raise QueryObjectValidationError(
                        "Unable to identify temporal column for date "
                        "range time comparison."
                    )

                with patch.object(
                    processor._qc_datasource,
                    "processing_time_offsets",
                    side_effect=raise_error,
                ):
                    with pytest.raises(
                        QueryObjectValidationError,
                        match="Unable to identify temporal column",
                    ):
                        processor._qc_datasource.processing_time_offsets(
                            df, query_object, None, None, False
                        )


def test_processing_time_offsets_date_range_enabled(processor):
    """Test processing_time_offsets correctly handles
    date range offsets when enabled."""
    from superset.common.query_object import QueryObject

    # Create a dataframe for testing
    df = pd.DataFrame(
        {
            "dim1": ["A", "B", "C"],
            "metric1": [10, 20, 30],
            "__timestamp": pd.date_range("2023-01-01", periods=3, freq="D"),
        }
    )

    # Create a properly mocked datasource
    mock_datasource = MagicMock()
    mock_datasource.id = 123
    mock_datasource.uid = "abc123"
    mock_datasource.cache_timeout = None
    mock_datasource.changed_on = pd.Timestamp("2023-01-01")
    mock_datasource.get_extra_cache_keys.return_value = {}

    # Create query object with date range offset
    query_object = QueryObject(
        datasource=mock_datasource,
        granularity="date_col",
        columns=[],
        is_timeseries=True,
        time_offsets=["2022-01-01 : 2022-01-31"],
        filter=[],
    )

    # Mock the query context and its methods
    processor._query_context.queries = [query_object]

    with patch("superset.models.helpers.feature_flag_manager") as mock_ff:
        mock_ff.is_feature_enabled.return_value = True

        with patch(
            "superset.utils.core.get_base_axis_labels",
            return_value=["__timestamp"],
        ):
            with patch(
                "superset.common.utils.time_range_utils.get_since_until_from_query_object"
            ) as mock_dates:
                mock_dates.return_value = (
                    pd.Timestamp("2023-01-01"),
                    pd.Timestamp("2023-01-03"),
                )

                with patch(
                    "superset.common.utils.time_range_utils.get_since_until_from_time_range"
                ) as mock_time_range:
                    mock_time_range.return_value = (
                        pd.Timestamp("2022-01-01"),
                        pd.Timestamp("2022-01-31"),
                    )

                    with patch.object(
                        processor, "get_query_result"
                    ) as mock_query_result:
                        mock_result = MagicMock()
                        mock_result.df = pd.DataFrame(
                            {
                                "dim1": ["A", "B"],
                                "metric1": [5, 10],
                                "__timestamp": pd.date_range(
                                    "2022-01-01", periods=2, freq="D"
                                ),
                            }
                        )
                        mock_result.query = "SELECT * FROM table"
                        mock_result.cache_key = "offset_cache_key"
                        mock_query_result.return_value = mock_result

                        # Mock the datasource's processing_time_offsets to
                        # return a proper result
                        mock_cached_result = {
                            "df": pd.DataFrame(
                                {
                                    "dim1": ["A", "B", "C"],
                                    "metric1": [10, 20, 30],
                                    "metric1 2022-01-01 : 2022-01-31": [5, 10, 15],
                                    "__timestamp": pd.date_range(
                                        "2023-01-01", periods=3, freq="D"
                                    ),
                                }
                            ),
                            "queries": ["SELECT * FROM table"],
                            "cache_keys": ["mock_cache_key"],
                        }

                        with patch.object(
                            processor._qc_datasource,
                            "processing_time_offsets",
                            return_value=mock_cached_result,
                        ):
                            # Test the method (call datasource method directly)
                            result = processor._qc_datasource.processing_time_offsets(
                                df, query_object, None, None, False
                            )

                            # Verify that the method completes successfully
                            assert "df" in result
                            assert "queries" in result
                            assert "cache_keys" in result

                            # Verify the result has the expected structure
                            assert isinstance(result["df"], pd.DataFrame)
                            assert isinstance(result["queries"], list)
                            assert isinstance(result["cache_keys"], list)


def test_processing_time_offsets_uses_chart_row_limit(processor):
    """Offset subquery inherits the chart's row_limit when one is set."""
    from superset.common.query_object import QueryObject
    from superset.models.helpers import ExploreMixin

    processor._qc_datasource.processing_time_offsets = (
        ExploreMixin.processing_time_offsets.__get__(processor._qc_datasource)
    )

    df = pd.DataFrame({"__timestamp": ["1990-01-01"], "sum__num": [100]})

    query_object = QueryObject(
        datasource=MagicMock(),
        granularity="ds",
        columns=[],
        metrics=["sum__num"],
        is_timeseries=True,
        row_limit=100,
        row_offset=0,
        time_offsets=["1 year ago"],
        filters=[
            {
                "col": "ds",
                "op": "TEMPORAL_RANGE",
                "val": "1990-01-01 : 1991-01-01",
            }
        ],
    )

    captured: list[dict[str, Any]] = []

    def fake_query(dct: dict[str, Any]) -> MagicMock:
        captured.append(dct)
        result = MagicMock()
        result.df = pd.DataFrame()
        result.query = "SELECT 1"
        return result

    processor._qc_datasource.query = fake_query
    processor._qc_datasource.normalize_df = MagicMock(return_value=pd.DataFrame())

    with (
        patch(
            "superset.models.helpers.get_since_until_from_query_object",
            return_value=(pd.Timestamp("1990-01-01"), pd.Timestamp("1991-01-01")),
        ),
        patch(
            "superset.common.utils.query_cache_manager.QueryCacheManager"
        ) as mock_cache_manager,
        patch.object(
            processor._qc_datasource,
            "get_time_grain",
            return_value=None,
        ),
        patch.object(
            processor._qc_datasource,
            "join_offset_dfs",
            return_value=df,
        ),
    ):
        mock_cache = MagicMock()
        mock_cache.is_loaded = False
        mock_cache_manager.get.return_value = mock_cache

        processor._qc_datasource.processing_time_offsets(
            df, query_object, None, None, False
        )

    assert len(captured) == 1
    assert captured[0]["row_limit"] == 100
    assert captured[0]["row_offset"] == 0


def test_processing_time_offsets_row_offset_extends_window(processor):
    """Offset subquery limit covers the main query's window (row_limit + row_offset).

    When the chart has pagination (row_offset > 0), fetching only row_limit rows
    in the offset period would likely miss the dimensions present in the main
    query's page, yielding null comparison values. The subquery instead drops
    row_offset and widens row_limit to cover the full window.
    """
    from superset.common.query_object import QueryObject
    from superset.models.helpers import ExploreMixin

    processor._qc_datasource.processing_time_offsets = (
        ExploreMixin.processing_time_offsets.__get__(processor._qc_datasource)
    )

    df = pd.DataFrame({"__timestamp": ["1990-01-01"], "sum__num": [100]})

    query_object = QueryObject(
        datasource=MagicMock(),
        granularity="ds",
        columns=[],
        metrics=["sum__num"],
        is_timeseries=True,
        row_limit=100,
        row_offset=10,
        time_offsets=["1 year ago"],
        filters=[
            {
                "col": "ds",
                "op": "TEMPORAL_RANGE",
                "val": "1990-01-01 : 1991-01-01",
            }
        ],
    )

    captured: list[dict[str, Any]] = []

    def fake_query(dct: dict[str, Any]) -> MagicMock:
        captured.append(dct)
        result = MagicMock()
        result.df = pd.DataFrame()
        result.query = "SELECT 1"
        return result

    processor._qc_datasource.query = fake_query
    processor._qc_datasource.normalize_df = MagicMock(return_value=pd.DataFrame())

    with (
        patch(
            "superset.models.helpers.get_since_until_from_query_object",
            return_value=(pd.Timestamp("1990-01-01"), pd.Timestamp("1991-01-01")),
        ),
        patch(
            "superset.common.utils.query_cache_manager.QueryCacheManager"
        ) as mock_cache_manager,
        patch.object(
            processor._qc_datasource,
            "get_time_grain",
            return_value=None,
        ),
        patch.object(
            processor._qc_datasource,
            "join_offset_dfs",
            return_value=df,
        ),
    ):
        mock_cache = MagicMock()
        mock_cache.is_loaded = False
        mock_cache_manager.get.return_value = mock_cache

        processor._qc_datasource.processing_time_offsets(
            df, query_object, None, None, False
        )

    assert len(captured) == 1
    assert captured[0]["row_limit"] == 110
    assert captured[0]["row_offset"] == 0


def test_processing_time_offsets_falls_back_to_config_row_limit(processor):
    """Offset subquery uses app config ROW_LIMIT when chart has offset but no limit."""
    from superset.common.query_object import QueryObject
    from superset.models.helpers import ExploreMixin

    processor._qc_datasource.processing_time_offsets = (
        ExploreMixin.processing_time_offsets.__get__(processor._qc_datasource)
    )

    df = pd.DataFrame({"__timestamp": ["1990-01-01"], "sum__num": [100]})

    query_object = QueryObject(
        datasource=MagicMock(),
        granularity="ds",
        columns=[],
        metrics=["sum__num"],
        is_timeseries=True,
        row_limit=None,
        row_offset=10,
        time_offsets=["1 year ago"],
        filters=[
            {
                "col": "ds",
                "op": "TEMPORAL_RANGE",
                "val": "1990-01-01 : 1991-01-01",
            }
        ],
    )

    captured: list[dict[str, Any]] = []

    def fake_query(dct: dict[str, Any]) -> MagicMock:
        captured.append(dct)
        result = MagicMock()
        result.df = pd.DataFrame()
        result.query = "SELECT 1"
        return result

    processor._qc_datasource.query = fake_query
    processor._qc_datasource.normalize_df = MagicMock(return_value=pd.DataFrame())

    with (
        patch(
            "superset.models.helpers.get_since_until_from_query_object",
            return_value=(pd.Timestamp("1990-01-01"), pd.Timestamp("1991-01-01")),
        ),
        patch(
            "superset.common.utils.query_cache_manager.QueryCacheManager"
        ) as mock_cache_manager,
        patch.object(
            processor._qc_datasource,
            "get_time_grain",
            return_value=None,
        ),
        patch.object(
            processor._qc_datasource,
            "join_offset_dfs",
            return_value=df,
        ),
        patch("superset.models.helpers.app") as mock_app,
    ):
        mock_app.config = {"ROW_LIMIT": 4242}
        mock_cache = MagicMock()
        mock_cache.is_loaded = False
        mock_cache_manager.get.return_value = mock_cache

        processor._qc_datasource.processing_time_offsets(
            df, query_object, None, None, False
        )

    assert len(captured) == 1
    assert captured[0]["row_limit"] == 4242
    assert captured[0]["row_offset"] == 0


def test_processing_time_offsets_updates_temporal_filter_with_adhoc_x_axis(processor):
    """Offset query's TEMPORAL_RANGE filter must be shifted when the X-axis is
    an adhoc Custom SQL column whose label differs from the underlying time
    column. Previously the filter was matched against the X-axis label, which
    never equals the dataset column, so the filter stayed at the original
    range and the offset query AND'd both ranges together (empty intersection).
    """
    from superset.common.query_object import QueryObject
    from superset.models.helpers import ExploreMixin

    processor._qc_datasource.processing_time_offsets = (
        ExploreMixin.processing_time_offsets.__get__(processor._qc_datasource)
    )

    df = pd.DataFrame(
        {
            "wedding_date_cast": pd.to_datetime(["2025-01-01", "2025-02-01"]),
            "SUM(revenue)": [110, 120],
        }
    )

    adhoc_x_axis = {
        "label": "wedding_date_cast",
        "sqlExpression": "CAST(wedding_date AS TIMESTAMP)",
        "expressionType": "SQL",
        "columnType": "BASE_AXIS",
        "timeGrain": "P1M",
    }

    query_object = QueryObject(
        datasource=MagicMock(),
        granularity=None,
        columns=[adhoc_x_axis],
        metrics=["SUM(revenue)"],
        is_timeseries=True,
        row_limit=10000,
        time_offsets=["1 year ago"],
        filters=[
            {
                "col": "wedding_date",
                "op": "TEMPORAL_RANGE",
                "val": "2025-01-01 : 2026-06-01",
            }
        ],
    )

    captured: list[dict[str, Any]] = []

    def fake_query(dct: dict[str, Any]) -> MagicMock:
        captured.append(dct)
        result = MagicMock()
        result.df = pd.DataFrame()
        result.query = "SELECT 1"
        return result

    processor._qc_datasource.query = fake_query
    processor._qc_datasource.normalize_df = MagicMock(return_value=pd.DataFrame())

    with (
        patch(
            "superset.models.helpers.get_since_until_from_query_object",
            return_value=(pd.Timestamp("2025-01-01"), pd.Timestamp("2026-06-01")),
        ),
        patch(
            "superset.common.utils.query_cache_manager.QueryCacheManager"
        ) as mock_cache_manager,
        patch.object(
            processor._qc_datasource,
            "get_time_grain",
            return_value="P1M",
        ),
        patch.object(
            processor._qc_datasource,
            "join_offset_dfs",
            return_value=df,
        ),
    ):
        mock_cache = MagicMock()
        mock_cache.is_loaded = False
        mock_cache_manager.get.return_value = mock_cache

        processor._qc_datasource.processing_time_offsets(
            df, query_object, None, None, False
        )

    assert len(captured) == 1
    temporal_filters = [
        flt for flt in captured[0]["filter"] if flt.get("op") == "TEMPORAL_RANGE"
    ]
    assert len(temporal_filters) == 1
    val = temporal_filters[0]["val"]
    assert "2024-01-01" in val, f"Expected shifted-from-dttm in val, got: {val!r}"
    assert "2025-06-01" in val, f"Expected shifted-to-dttm in val, got: {val!r}"


def test_processing_time_offsets_quarter_offset_shifts_query_window(
    processor: QueryContextProcessor,
) -> None:
    """A quarter offset must shift the offset query's window, not just the
    join keys. parsedatetime does not understand "1 quarter ago", so without
    rewriting quarters to months the offset subquery silently runs against
    the current period and the comparison series joins to nulls. The fake
    query below derives its rows from the requested window, so the join only
    yields the expected values when the window was actually shifted.
    """
    from superset.common.query_object import QueryObject
    from superset.models.helpers import ExploreMixin

    # The fixture's datasource is a MagicMock, not a real Explorable
    datasource: Any = processor._qc_datasource

    for method in (
        "processing_time_offsets",
        "_align_offset_without_time_grain",
        "_coalesce_offset_index",
    ):
        setattr(
            datasource,
            method,
            getattr(ExploreMixin, method).__get__(datasource),
        )

    df = pd.DataFrame(
        {
            "__timestamp": pd.to_datetime(["2024-04-01", "2024-05-01", "2024-06-01"]),
            "sum__num": [100, 200, 300],
        }
    )

    query_object = QueryObject(
        datasource=MagicMock(),
        granularity="ds",
        columns=[],
        metrics=["sum__num"],
        is_timeseries=True,
        time_offsets=["1 quarter ago"],
        filters=[
            {
                "col": "ds",
                "op": "TEMPORAL_RANGE",
                "val": "2024-04-01 : 2024-07-01",
            }
        ],
    )

    captured: list[dict[str, Any]] = []

    def fake_query(dct: dict[str, Any]) -> MagicMock:
        captured.append(dct)
        result = MagicMock()
        result.df = pd.DataFrame(
            {
                "__timestamp": pd.date_range(
                    start=dct["from_dttm"], periods=3, freq="MS"
                ),
                "sum__num": [1.0, 2.0, 3.0],
            }
        )
        result.query = "SELECT 1"
        return result

    datasource.query = fake_query
    datasource.normalize_df = MagicMock(
        side_effect=lambda offset_df, _query_object: offset_df
    )

    with (
        patch(
            "superset.models.helpers.get_since_until_from_query_object",
            return_value=(pd.Timestamp("2024-04-01"), pd.Timestamp("2024-07-01")),
        ),
        patch(
            "superset.common.utils.query_cache_manager.QueryCacheManager"
        ) as mock_cache_manager,
        patch.object(
            datasource,
            "get_time_grain",
            return_value=None,
        ),
    ):
        mock_cache = MagicMock()
        mock_cache.is_loaded = False
        mock_cache_manager.get.return_value = mock_cache

        result = datasource.processing_time_offsets(df, query_object, None, None, False)

    assert len(captured) == 1
    assert captured[0]["from_dttm"] == pd.Timestamp("2024-01-01")
    assert captured[0]["to_dttm"] == pd.Timestamp("2024-04-01")
    assert result["df"]["sum__num__1 quarter ago"].tolist() == [1.0, 2.0, 3.0]


def test_processing_time_offsets_accepts_zero_shift_offset(
    processor: QueryContextProcessor,
) -> None:
    """An offset that legitimately parses to no shift (e.g. "0 days ago")
    must render a self-comparison instead of being rejected as
    uninterpretable: a zero delta is not a parse failure.
    """
    from superset.common.query_object import QueryObject
    from superset.models.helpers import ExploreMixin

    # The fixture's datasource is a MagicMock, not a real Explorable
    datasource: Any = processor._qc_datasource

    for method in (
        "processing_time_offsets",
        "_align_offset_without_time_grain",
        "_coalesce_offset_index",
    ):
        setattr(
            datasource,
            method,
            getattr(ExploreMixin, method).__get__(datasource),
        )

    df = pd.DataFrame(
        {
            "__timestamp": pd.to_datetime(["2024-04-01", "2024-05-01", "2024-06-01"]),
            "sum__num": [100, 200, 300],
        }
    )

    query_object = QueryObject(
        datasource=MagicMock(),
        granularity="ds",
        columns=[],
        metrics=["sum__num"],
        is_timeseries=True,
        time_offsets=["0 days ago"],
        filters=[
            {
                "col": "ds",
                "op": "TEMPORAL_RANGE",
                "val": "2024-04-01 : 2024-07-01",
            }
        ],
    )

    captured: list[dict[str, Any]] = []

    def fake_query(dct: dict[str, Any]) -> MagicMock:
        captured.append(dct)
        result = MagicMock()
        result.df = pd.DataFrame(
            {
                "__timestamp": pd.date_range(
                    start=dct["from_dttm"], periods=3, freq="MS"
                ),
                "sum__num": [1.0, 2.0, 3.0],
            }
        )
        result.query = "SELECT 1"
        return result

    datasource.query = fake_query
    datasource.normalize_df = MagicMock(
        side_effect=lambda offset_df, _query_object: offset_df
    )

    with (
        patch(
            "superset.models.helpers.get_since_until_from_query_object",
            return_value=(pd.Timestamp("2024-04-01"), pd.Timestamp("2024-07-01")),
        ),
        patch(
            "superset.common.utils.query_cache_manager.QueryCacheManager"
        ) as mock_cache_manager,
        patch.object(
            datasource,
            "get_time_grain",
            return_value=None,
        ),
    ):
        mock_cache = MagicMock()
        mock_cache.is_loaded = False
        mock_cache_manager.get.return_value = mock_cache

        result = datasource.processing_time_offsets(df, query_object, None, None, False)

    # The offset query runs against the unshifted window and its rows join
    # back onto the main series one-to-one
    assert len(captured) == 1
    assert captured[0]["from_dttm"] == pd.Timestamp("2024-04-01")
    assert captured[0]["to_dttm"] == pd.Timestamp("2024-07-01")
    assert result["df"]["sum__num__0 days ago"].tolist() == [1.0, 2.0, 3.0]


def test_processing_time_offsets_rejects_unparseable_offset(
    processor: QueryContextProcessor,
) -> None:
    """An offset no parser understands must fail with a validation error
    instead of querying an unshifted window and presenting the current
    period's rows as the comparison series.
    """
    from superset.common.query_object import QueryObject
    from superset.exceptions import QueryObjectValidationError
    from superset.models.helpers import ExploreMixin

    # The fixture's datasource is a MagicMock, not a real Explorable
    datasource: Any = processor._qc_datasource

    datasource.processing_time_offsets = ExploreMixin.processing_time_offsets.__get__(
        datasource
    )

    df = pd.DataFrame(
        {
            "__timestamp": pd.to_datetime(["2024-04-01"]),
            "sum__num": [100],
        }
    )

    query_object = QueryObject(
        datasource=MagicMock(),
        granularity="ds",
        columns=[],
        metrics=["sum__num"],
        is_timeseries=True,
        time_offsets=["not a real offset"],
        filters=[
            {
                "col": "ds",
                "op": "TEMPORAL_RANGE",
                "val": "2024-04-01 : 2024-07-01",
            }
        ],
    )

    datasource.query = MagicMock()

    with (
        patch(
            "superset.models.helpers.get_since_until_from_query_object",
            return_value=(pd.Timestamp("2024-04-01"), pd.Timestamp("2024-07-01")),
        ),
        patch.object(
            datasource,
            "get_time_grain",
            return_value=None,
        ),
    ):
        with pytest.raises(
            QueryObjectValidationError, match="Unable to interpret the time offset"
        ):
            datasource.processing_time_offsets(df, query_object, None, None, False)

    datasource.query.assert_not_called()


def test_ensure_totals_available_updates_cache_values():
    """
    Test that ensure_totals_available() updates the query objects AND
    cache_values to keep them in sync.

    The issue was that ensure_totals_available() modified QueryObject instances
    (e.g., setting row_limit=None on totals queries and adding contribution_totals
    to post_processing), but cache_values still contained the original queries.
    This caused cache key mismatches between worker execution and cache fetch.
    """
    import pandas as pd

    from superset.common.query_object import QueryObject

    # Create a mock datasource
    mock_datasource = MagicMock()
    mock_datasource.uid = "test_datasource"
    mock_datasource.database.db_engine_spec.engine = "postgresql"
    mock_datasource.cache_timeout = None
    mock_datasource.changed_on = None

    # Create QueryObjects that would trigger ensure_totals_available logic
    # Query 1: Main query with contribution post-processing (needs totals)
    main_query = QueryObject(
        datasource=mock_datasource,
        columns=["brokerage"],
        metrics=["Net Amount In", "Amount Out", "Amount In"],
        row_limit=50000,
        orderby=[["Net Amount In", False]],
        post_processing=[
            {
                "operation": "contribution",
                "options": {
                    "columns": ["Amount In", "Amount Out"],
                    "rename_columns": ["%Amount In", "%Amount Out"],
                },
            }
        ],
    )

    # Query 2: Totals query (no columns, has metrics, no post-processing)
    totals_query = QueryObject(
        datasource=mock_datasource,
        columns=[],  # No columns = totals query
        metrics=["Net Amount In", "Amount Out", "Amount In"],
        row_limit=50000,
        post_processing=[],  # No post-processing
    )

    # Create mock query context
    mock_query_context = MagicMock()
    mock_query_context.force = False
    mock_query_context.datasource = mock_datasource
    mock_query_context.queries = [main_query, totals_query]
    mock_query_context.result_type = "full"
    mock_query_context.cache_values = {
        "datasource": {"type": "table", "id": 1},
        "queries": [
            # These are the original queries as they would be stored in cache_values
            {
                "columns": ["brokerage"],
                "metrics": ["Net Amount In", "Amount Out", "Amount In"],
                "row_limit": 50000,
                "orderby": [("Net Amount In", False)],
                "post_processing": [
                    {
                        "operation": "contribution",
                        "options": {
                            "columns": ["Amount In", "Amount Out"],
                            "rename_columns": ["%Amount In", "%Amount Out"],
                        },
                    }
                ],
            },
            {
                "columns": [],
                "metrics": ["Net Amount In", "Amount Out", "Amount In"],
                "row_limit": 50000,
                "post_processing": [],
            },
        ],
        "result_type": "full",
        "result_format": "json",
    }

    # Create processor
    processor = QueryContextProcessor(mock_query_context)
    processor._qc_datasource = mock_datasource

    # Mock the query execution result for totals query
    mock_query_result = MagicMock()
    mock_df = pd.DataFrame(
        {
            "Net Amount In": [20228060486.838825],
            "Amount Out": [-20543489614.980007],
            "Amount In": [40771550101.81883],
        }
    )
    mock_query_result.df = mock_df

    with patch.object(
        mock_query_context, "get_query_result", return_value=mock_query_result
    ):
        # Call ensure_totals_available
        processor.ensure_totals_available()

        # Now call get_payload which should update cache_values
        with patch(
            "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 = [
                {
                    "data": [{"brokerage": "Test", "Net Amount In": 100}],
                    "query": "SELECT ...",
                }
            ]
            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(
                "superset.common.query_context_processor.QueryCacheManager"
            ) as mock_cache_manager:
                mock_cache = MagicMock()
                mock_cache.is_loaded = True
                mock_cache.df = pd.DataFrame(
                    {"brokerage": ["Test"], "Net Amount In": [100]}
                )
                mock_cache.query = "SELECT ..."
                mock_cache.error_message = None
                mock_cache.status = "success"
                mock_cache_manager.get.return_value = mock_cache

                # This should update cache_values to match the modified queries
                processor.get_payload(cache_query_context=False)

    # Verify that cache_values has been updated to reflect the modifications
    updated_cache_queries = mock_query_context.cache_values["queries"]

    # Check that totals query has row_limit=None (modified by ensure_totals_available)
    assert updated_cache_queries[1]["row_limit"] is None, (
        "Expected totals query to have row_limit=None after ensure_totals_available, "
        f"but got: {updated_cache_queries[1]['row_limit']}"
    )

    # Check that the main query has contribution_totals in post_processing
    assert (
        "contribution_totals"
        in updated_cache_queries[0]["post_processing"][0]["options"]
    ), "Expected main query post_processing to have contribution_totals added"

    # Verify the contribution_totals match what we mocked
    expected_totals = {
        "Net Amount In": 20228060486.838825,
        "Amount Out": -20543489614.980007,
        "Amount In": 40771550101.81883,
    }
    assert (
        updated_cache_queries[0]["post_processing"][0]["options"]["contribution_totals"]
        == expected_totals
    )


def test_get_df_payload_validates_before_cache_key_generation():
    """
    Test that get_df_payload calls validate() before generating cache key.
    """
    from superset.common.query_object import QueryObject

    # Create a mock query context
    mock_query_context = MagicMock()
    mock_query_context.force = False
    mock_query_context.result_type = "full"

    # Create a mock datasource
    mock_datasource = MagicMock()
    mock_datasource.id = 123
    mock_datasource.uid = "test_datasource"
    mock_datasource.cache_timeout = None
    mock_datasource.database.db_engine_spec.engine = "postgresql"
    mock_datasource.database.extra = "{}"
    mock_datasource.get_extra_cache_keys.return_value = []
    mock_datasource.changed_on = None

    # Create processor
    processor = QueryContextProcessor(mock_query_context)
    processor._qc_datasource = mock_datasource

    # Create a query object with unsanitized where clause
    query_obj = QueryObject(
        datasource=mock_datasource,
        columns=["col1"],
        metrics=[],
        extras={"where": "(\n  col1 > 0\n)"},  # Unsanitized with newlines
    )

    # Track the order of calls
    call_order = []

    original_validate = query_obj.validate

    def mock_validate(*args, **kwargs):
        call_order.append("validate")
        # Update extras to simulate sanitization
        query_obj.extras["where"] = "(col1 > 0)"  # Sanitized, compact format
        return original_validate(*args, **kwargs)

    original_cache_key = query_obj.cache_key

    def mock_cache_key(*args, **kwargs):
        call_order.append("cache_key")
        # Verify that extras have been sanitized at this point
        assert query_obj.extras["where"] == "(col1 > 0)", (
            f"Expected sanitized clause in cache_key, got: {query_obj.extras['where']}"
        )
        return original_cache_key(*args, **kwargs)

    with patch.object(query_obj, "validate", side_effect=mock_validate):
        with patch.object(query_obj, "cache_key", side_effect=mock_cache_key):
            with patch(
                "superset.common.query_context_processor.QueryCacheManager"
            ) as mock_cache_manager:
                mock_cache = MagicMock()
                mock_cache.is_loaded = True
                mock_cache.df = pd.DataFrame({"col1": [1, 2, 3]})
                mock_cache.query = "SELECT * FROM table"
                mock_cache.error_message = None
                mock_cache.status = "success"
                mock_cache.bq_memory_limited = False
                mock_cache_manager.get.return_value = mock_cache

                # Call get_df_payload
                processor.get_df_payload(query_obj, force_cached=False)

    # Verify validate was called before cache_key
    assert call_order == ["validate", "cache_key"], (
        f"Expected validate to be called before cache_key, "
        f"but got call order: {call_order}"
    )


def test_cache_values_sync_after_ensure_totals_available():
    """
    Test that cache_values is synchronized with QueryObject modifications
    after ensure_totals_available() runs.

    This is a focused regression test for the cache key mismatch issue.
    It verifies that when ensure_totals_available() modifies QueryObject
    instances, those changes are reflected in cache_values before the
    QueryContext cache key is generated.
    """
    import pandas as pd

    from superset.common.query_object import QueryObject

    # Create a mock datasource
    mock_datasource = MagicMock()
    mock_datasource.uid = "test_datasource_456"
    mock_datasource.database.db_engine_spec.engine = "pinot"
    mock_datasource.cache_timeout = None
    mock_datasource.changed_on = None

    # Create two queries: one totals query and one main query with contribution
    totals_query = QueryObject(
        datasource=mock_datasource,
        columns=[],
        metrics=["sales"],
        row_limit=1000,
        post_processing=[],
    )

    main_query = QueryObject(
        datasource=mock_datasource,
        columns=["region"],
        metrics=["sales"],
        row_limit=1000,
        post_processing=[{"operation": "contribution", "options": {}}],
    )

    # Create mock query context with initial cache_values
    mock_query_context = MagicMock()
    mock_query_context.force = False
    mock_query_context.datasource = mock_datasource
    mock_query_context.queries = [main_query, totals_query]
    mock_query_context.result_type = "full"
    mock_query_context.cache_values = {
        "datasource": {"type": "table", "id": 20},
        "queries": [
            {
                "columns": ["region"],
                "metrics": ["sales"],
                "row_limit": 1000,
                "post_processing": [{"operation": "contribution", "options": {}}],
            },
            {
                "columns": [],
                "metrics": ["sales"],
                "row_limit": 1000,
                "post_processing": [],
            },
        ],
        "result_type": "full",
        "result_format": "json",
    }

    # Create processor
    processor = QueryContextProcessor(mock_query_context)
    processor._qc_datasource = mock_datasource

    # Mock query execution result (totals query execution)
    mock_query_result = MagicMock()
    mock_df = pd.DataFrame({"sales": [1000.0]})
    mock_query_result.df = mock_df

    # Patch methods to isolate the test
    with patch.object(
        mock_query_context, "get_query_result", return_value=mock_query_result
    ):
        # Mock cache management to prevent actual caching
        with patch(
            "superset.common.query_context_processor.QueryCacheManager"
        ) as mock_cache_manager:
            mock_cache = MagicMock()
            mock_cache.is_loaded = True
            mock_cache.df = pd.DataFrame({"region": ["North"], "sales": [100]})
            mock_cache.query = "SELECT region, SUM(sales) FROM table GROUP BY region"
            mock_cache.error_message = None
            mock_cache.status = "success"
            mock_cache_manager.get.return_value = mock_cache

            # Mock the query results
            with patch(
                "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_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
                processor.get_payload(cache_query_context=False)

    # Verify the fix: cache_values should now reflect the modifications
    updated_cache_queries = mock_query_context.cache_values["queries"]
    updated_totals_row_limit = updated_cache_queries[1]["row_limit"]

    # Before the fix: row_limit would remain 1000 in cache_values
    # After the fix: row_limit should be None (modified by
    # ensure_totals_available)
    assert updated_totals_row_limit is None, (
        "Expected row_limit to be None after ensure_totals_available, "
        f"but got: {updated_totals_row_limit}"
    )

    # Verify that contribution_totals was added to the main query
    assert (
        "contribution_totals"
        in updated_cache_queries[0]["post_processing"][0]["options"]
    )

    # Verify that the main query row_limit is still 1000 (only totals query
    # should be modified)
    assert updated_cache_queries[0]["row_limit"] == 1000


def test_cache_key_excludes_contribution_totals():
    """
    Test that cache_key() excludes contribution_totals from post_processing.

    contribution_totals is computed at runtime by ensure_totals_available() and
    varies per request. Including it in the cache key would cause mismatches
    between workers that compute different totals for the same query.
    """
    from superset.common.query_object import QueryObject

    mock_datasource = MagicMock()
    mock_datasource.uid = "test_datasource"
    mock_datasource.database.extra = "{}"
    mock_datasource.get_extra_cache_keys.return_value = []

    # Create query with contribution post-processing that includes contribution_totals
    query_with_totals = QueryObject(
        datasource=mock_datasource,
        columns=["region"],
        metrics=["sales", "profit"],
        post_processing=[
            {
                "operation": "contribution",
                "options": {
                    "columns": ["sales", "profit"],
                    "rename_columns": ["%sales", "%profit"],
                    "contribution_totals": {"sales": 1000.0, "profit": 200.0},
                },
            }
        ],
    )

    # Create identical query without contribution_totals
    query_without_totals = QueryObject(
        datasource=mock_datasource,
        columns=["region"],
        metrics=["sales", "profit"],
        post_processing=[
            {
                "operation": "contribution",
                "options": {
                    "columns": ["sales", "profit"],
                    "rename_columns": ["%sales", "%profit"],
                },
            }
        ],
    )

    # Cache keys should be identical since contribution_totals is excluded
    cache_key_with = query_with_totals.cache_key()
    cache_key_without = query_without_totals.cache_key()

    assert cache_key_with == cache_key_without, (
        "Cache keys should match regardless of contribution_totals. "
        f"With totals: {cache_key_with}, Without totals: {cache_key_without}"
    )


def test_cache_key_preserves_other_post_processing_options():
    """
    Test that cache_key() only excludes contribution_totals, not other options.
    """
    from superset.common.query_object import QueryObject

    mock_datasource = MagicMock()
    mock_datasource.uid = "test_datasource"
    mock_datasource.database.extra = "{}"
    mock_datasource.get_extra_cache_keys.return_value = []

    # Create query with contribution post-processing
    query1 = QueryObject(
        datasource=mock_datasource,
        columns=["region"],
        metrics=["sales"],
        post_processing=[
            {
                "operation": "contribution",
                "options": {
                    "columns": ["sales"],
                    "rename_columns": ["%sales"],
                    "contribution_totals": {"sales": 1000.0},
                },
            }
        ],
    )

    # Create query with different rename_columns
    query2 = QueryObject(
        datasource=mock_datasource,
        columns=["region"],
        metrics=["sales"],
        post_processing=[
            {
                "operation": "contribution",
                "options": {
                    "columns": ["sales"],
                    "rename_columns": ["%sales_pct"],  # Different!
                    "contribution_totals": {"sales": 1000.0},
                },
            }
        ],
    )

    # Cache keys should differ because rename_columns is different
    assert query1.cache_key() != query2.cache_key(), (
        "Cache keys should differ when other post_processing options differ"
    )


def test_cache_key_non_contribution_post_processing_unchanged():
    """
    Test that non-contribution post_processing operations are unchanged in cache key.
    """
    from superset.common.query_object import QueryObject

    mock_datasource = MagicMock()
    mock_datasource.uid = "test_datasource"
    mock_datasource.database.extra = "{}"
    mock_datasource.get_extra_cache_keys.return_value = []

    # Create query with non-contribution post-processing
    query1 = QueryObject(
        datasource=mock_datasource,
        columns=["region"],
        metrics=["sales"],
        post_processing=[
            {
                "operation": "pivot",
                "options": {"columns": ["region"], "aggregates": {"sales": "sum"}},
            }
        ],
    )

    query2 = QueryObject(
        datasource=mock_datasource,
        columns=["region"],
        metrics=["sales"],
        post_processing=[
            {
                "operation": "pivot",
                "options": {"columns": ["region"], "aggregates": {"sales": "mean"}},
            }
        ],
    )

    # Cache keys should differ because aggregates option is different
    assert query1.cache_key() != query2.cache_key(), (
        "Cache keys should differ for different non-contribution post_processing"
    )


def test_force_cached_normalizes_totals_query_row_limit():
    """
    When fetching from cache (force_cached=True), the totals query should still be
    normalized so its cache key matches the cached entry, but the totals query should
    not be executed.
    """
    from superset.common.query_object import QueryObject

    mock_datasource = MagicMock()
    mock_datasource.uid = "test_datasource"
    mock_datasource.column_names = ["region", "sales"]
    mock_datasource.cache_timeout = None
    mock_datasource.changed_on = None
    mock_datasource.get_extra_cache_keys.return_value = []
    mock_datasource.database.extra = "{}"
    mock_datasource.database.impersonate_user = False
    mock_datasource.database.db_engine_spec.get_impersonation_key.return_value = None

    totals_query = QueryObject(
        datasource=mock_datasource,
        columns=[],
        metrics=["sales"],
        row_limit=1000,
    )
    main_query = QueryObject(
        datasource=mock_datasource,
        columns=["region"],
        metrics=["sales"],
        row_limit=1000,
        post_processing=[{"operation": "contribution", "options": {}}],
    )

    totals_query.validate = MagicMock()
    main_query.validate = MagicMock()

    captured_limits: list[int | None] = []

    def totals_cache_key(**kwargs: Any) -> str:
        captured_limits.append(totals_query.row_limit)
        return "totals-cache-key"

    totals_query.cache_key = totals_cache_key
    main_query.cache_key = lambda **kwargs: "main-cache-key"

    mock_query_context = MagicMock()
    mock_query_context.force = False
    mock_query_context.datasource = mock_datasource
    mock_query_context.queries = [main_query, totals_query]
    mock_query_context.result_type = ChartDataResultType.FULL
    mock_query_context.result_format = ChartDataResultFormat.JSON
    mock_query_context.cache_values = {
        "queries": [main_query.to_dict(), totals_query.to_dict()]
    }
    mock_query_context.get_query_result = MagicMock()

    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(
        "superset.common.query_context_processor.security_manager"
    ) as mock_security_manager:
        mock_security_manager.get_rls_cache_key.return_value = None

        with patch(
            "superset.common.query_context_processor.QueryCacheManager"
        ) as mock_cache_manager:

            def cache_get(*args: Any, **kwargs: Any) -> Any:
                df = pd.DataFrame({"region": ["North"], "sales": [100]})
                cache = MagicMock()
                cache.is_loaded = True
                cache.df = df
                cache.query = "SELECT 1"
                cache.error_message = None
                cache.status = QueryStatus.SUCCESS
                cache.applied_template_filters = []
                cache.applied_filter_columns = []
                cache.rejected_filter_columns = []
                cache.annotation_data = {}
                cache.is_cached = True
                cache.sql_rowcount = len(df)
                cache.cache_dttm = "2024-01-01T00:00:00"
                cache.bq_memory_limited = False
                return cache

            mock_cache_manager.get.side_effect = cache_get

            processor.get_payload(cache_query_context=False, force_cached=True)

    assert captured_limits == [None], "Totals query should be normalized before caching"
    mock_query_context.get_query_result.assert_not_called()


def test_get_df_payload_invalidates_cache_missing_applied_filter_columns():
    """
    Test that get_df_payload invalidates cache when cache is loaded but missing
    applied_filter_columns and query has filters.

    This ensures that old cache entries without applied_filter_columns are
    invalidated and fresh queries are executed to populate the field correctly.
    """
    from superset.common.query_object import QueryObject

    # Minimal setup
    mock_query_context = MagicMock()
    mock_query_context.force = False
    mock_datasource = MagicMock()
    mock_datasource.column_names = ["col1"]

    processor = QueryContextProcessor(mock_query_context)
    processor._qc_datasource = mock_datasource

    # Create query object with filters (note: `filters` kwarg, not `filter`)
    query_obj = QueryObject(
        datasource=mock_datasource,
        columns=["col1"],
        filters=[{"col": "col1", "op": "IN", "val": ["value1"]}],
    )

    # Simple cache class that tracks is_loaded changes
    class MockCache:
        def __init__(self):
            self.is_loaded = True
            self.applied_filter_columns = []  # Empty = missing
            self.df = pd.DataFrame()
            self.query = ""
            self.status = "success"
            self.cache_dttm = "2024-01-01T00:00:00"
            self.queried_dttm = "2024-01-01T00:00:00"
            self.stacktrace = None
            self.error_message = None
            self.is_cached = True
            self.sql_rowcount = 0
            self.cache_value = None
            self.cache_timeout = 3600
            self.datasource_uid = "test_datasource"
            self.applied_template_filters = []
            self.rejected_filter_columns = []
            self.annotation_data = {}
            self.bq_memory_limited = False
            self.bq_memory_limited_row_count = 0
            self.set_query_result = MagicMock()

    mock_cache = MockCache()

    with patch(
        "superset.common.query_context_processor.QueryCacheManager"
    ) as mock_cache_manager:
        mock_cache_manager.get.return_value = mock_cache

        # Prevent validate from doing any heavy work; it shouldn't modify filters
        with patch.object(query_obj, "validate", return_value=None):
            with patch.object(processor, "query_cache_key", return_value="key"):
                with patch.object(processor, "get_cache_timeout", return_value=3600):
                    # Call get_df_payload - should invalidate cache
                    processor.get_df_payload(query_obj, force_cached=False)

    # Verify cache was invalidated
    assert mock_cache.is_loaded is False, (
        "Cache should be inv when no applied_filter_columns and query has filters"
    )


def test_get_df_payload_bq_memory_limited_warning() -> None:
    """
    Test that get_df_payload includes a warning when BigQuery results are
    truncated due to the memory limit (g.bq_memory_limited is set).
    """
    from superset.common.query_object import QueryObject

    mock_query_context = MagicMock()
    mock_query_context.force = False
    mock_query_context.form_data = {"slice_id": 42}

    mock_datasource = MagicMock()
    mock_datasource.column_names = ["col1"]
    mock_datasource.uid = "test_ds"
    mock_datasource.cache_timeout = None
    mock_datasource.changed_on = None
    mock_datasource.get_extra_cache_keys.return_value = []
    mock_datasource.data = MagicMock()
    mock_datasource.data.get.return_value = {}

    processor = QueryContextProcessor(mock_query_context)
    processor._qc_datasource = mock_datasource

    query_obj = QueryObject(
        datasource=mock_datasource,
        columns=["col1"],
    )

    with patch(
        "superset.common.query_context_processor.QueryCacheManager"
    ) as mock_cache_manager:
        mock_cache = MagicMock()
        mock_cache.is_loaded = True
        mock_cache.df = pd.DataFrame({"col1": [1, 2, 3]})
        mock_cache.query = "SELECT col1 FROM table"
        mock_cache.error_message = None
        mock_cache.status = "success"
        mock_cache.applied_filter_columns = ["col1"]
        mock_cache.applied_template_filters = []
        mock_cache.rejected_filter_columns = []
        mock_cache.annotation_data = {}
        mock_cache.is_cached = True
        mock_cache.sql_rowcount = 3
        mock_cache.cache_dttm = "2024-01-01T00:00:00"
        mock_cache.queried_dttm = "2024-01-01T00:00:00"
        mock_cache.bq_memory_limited = True
        mock_cache.bq_memory_limited_row_count = 5000
        mock_cache_manager.get.return_value = mock_cache

        with patch.object(query_obj, "validate", return_value=None):
            with patch.object(processor, "query_cache_key", return_value="key"):
                with patch.object(processor, "get_cache_timeout", return_value=3600):
                    result = processor.get_df_payload(query_obj, force_cached=False)

    assert result["warning"] is not None
    assert "Chart 42" in result["warning"]
    assert "5,000 rows" in result["warning"]
    assert "memory constraints" in result["warning"]


def test_get_df_payload_no_warning_when_not_memory_limited() -> None:
    """
    Test that get_df_payload does not include a warning when BigQuery
    results were not truncated.
    """
    from superset.common.query_object import QueryObject

    mock_query_context = MagicMock()
    mock_query_context.force = False
    mock_query_context.form_data = {}

    mock_datasource = MagicMock()
    mock_datasource.column_names = ["col1"]
    mock_datasource.uid = "test_ds"
    mock_datasource.cache_timeout = None
    mock_datasource.changed_on = None
    mock_datasource.get_extra_cache_keys.return_value = []
    mock_datasource.data = MagicMock()
    mock_datasource.data.get.return_value = {}

    processor = QueryContextProcessor(mock_query_context)
    processor._qc_datasource = mock_datasource

    query_obj = QueryObject(
        datasource=mock_datasource,
        columns=["col1"],
    )

    with patch(
        "superset.common.query_context_processor.QueryCacheManager"
    ) as mock_cache_manager:
        mock_cache = MagicMock()
        mock_cache.is_loaded = True
        mock_cache.df = pd.DataFrame({"col1": [1, 2]})
        mock_cache.query = "SELECT col1 FROM table"
        mock_cache.error_message = None
        mock_cache.status = "success"
        mock_cache.applied_filter_columns = ["col1"]
        mock_cache.applied_template_filters = []
        mock_cache.rejected_filter_columns = []
        mock_cache.annotation_data = {}
        mock_cache.is_cached = True
        mock_cache.sql_rowcount = 2
        mock_cache.cache_dttm = "2024-01-01T00:00:00"
        mock_cache.queried_dttm = "2024-01-01T00:00:00"
        mock_cache.bq_memory_limited = False
        mock_cache_manager.get.return_value = mock_cache

        with patch.object(query_obj, "validate", return_value=None):
            with patch.object(processor, "query_cache_key", return_value="key"):
                with patch.object(processor, "get_cache_timeout", return_value=3600):
                    result = processor.get_df_payload(query_obj, force_cached=False)

    assert result["warning"] is None


def test_raise_for_access_evaluates_access_before_validate():
    """
    Access must be evaluated before the queries are validated, because query
    validation renders the request's filter expressions. When access is denied,
    no query is validated (so caller-supplied input is never rendered).
    """
    from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
    from superset.exceptions import SupersetSecurityException
    from superset.utils.core import DatasourceType

    query = MagicMock()
    query_context = MagicMock()
    query_context.queries = [query]
    query_context.datasource.type = DatasourceType.TABLE

    processor = QueryContextProcessor(query_context)

    denied = SupersetSecurityException(
        SupersetError(
            message="denied",
            error_type=SupersetErrorType.DATASOURCE_SECURITY_ACCESS_ERROR,
            level=ErrorLevel.ERROR,
        )
    )
    with patch(
        "superset.common.query_context_processor.security_manager.raise_for_access",
        side_effect=denied,
    ):
        with pytest.raises(SupersetSecurityException):
            processor.raise_for_access()

    query.validate.assert_not_called()


def test_grouping_sets_fallback_handles_adhoc_and_physical_columns() -> None:
    """
    The fallback used on engines without native GROUPING SETS support must
    build ``label_to_column`` from the same labels used to run each level's
    subquery, for both physical (string) and adhoc (dict) columns, and in the
    original column order. Otherwise an adhoc column ahead of a physical one
    in ``query_object.columns`` misaligns the mapping, and adhoc groupby
    columns referenced in ``grouping_sets`` are silently dropped.
    """
    import copy
    from datetime import timedelta

    from superset.common.query_object import QueryObject
    from superset.models.helpers import QueryResult
    from superset.superset_typing import AdhocColumn

    adhoc_col: AdhocColumn = {
        "sqlExpression": "DATE_TRUNC('month', dttm)",
        "label": "month",
    }
    mock_datasource = MagicMock()

    query_obj = QueryObject(
        datasource=mock_datasource,
        columns=[adhoc_col, "state"],
        grouping_sets=[["month", "state"], ["month"], []],
    )

    processor = QueryContextProcessor(MagicMock())
    processor._qc_datasource = mock_datasource

    captured_columns: list[list[Any]] = []

    def fake_get_query_result(sub_query: QueryObject) -> QueryResult:
        captured_columns.append(copy.copy(sub_query.columns))
        return QueryResult(
            df=pd.DataFrame({"metric": [1]}),
            query="SELECT 1",
            duration=timedelta(seconds=0),
        )

    mock_datasource.get_query_result.side_effect = fake_get_query_result

    processor._grouping_sets_fallback(query_obj)

    # Level ["month", "state"] must include both the adhoc and physical column,
    # each mapped to its own definition (not swapped).
    assert captured_columns[0] == [adhoc_col, "state"]
    # Level ["month"] must still include the adhoc column, not drop it.
    assert captured_columns[1] == [adhoc_col]
    # Grand total level has no groupby columns.
    assert captured_columns[2] == []


def test_grouping_sets_fallback_applies_row_offset_once_globally() -> None:
    """
    The native GROUPING SETS path applies `row_offset` exactly once, to the
    combined multi-level result (see the unconditional `qry.offset()` call in
    `models/helpers.py`). The fallback must match that: it must not apply the
    same offset independently to every per-level subquery, since that would
    apply it once per level (and can drop an entire low-row-count level, e.g.
    a single grand-total row, outright).
    """
    from datetime import timedelta

    from superset.common.query_object import QueryObject
    from superset.models.helpers import QueryResult

    mock_datasource = MagicMock()

    query_obj = QueryObject(
        datasource=mock_datasource,
        columns=["state"],
        grouping_sets=[["state"], []],
        row_offset=1,
    )

    processor = QueryContextProcessor(MagicMock())
    processor._qc_datasource = mock_datasource

    captured_offsets: list[int] = []

    # Each level returns 2 rows regardless of the (should-be-ignored) offset,
    # emulating a real datasource that would otherwise apply row_offset itself.
    def fake_get_query_result(sub_query: QueryObject) -> QueryResult:
        captured_offsets.append(sub_query.row_offset)
        return QueryResult(
            df=pd.DataFrame({"state": ["CA", "NY"]})
            if sub_query.columns
            else pd.DataFrame({"state": ["total"]}),
            query="SELECT 1",
            duration=timedelta(seconds=0),
        )

    mock_datasource.get_query_result.side_effect = fake_get_query_result

    result = processor._grouping_sets_fallback(query_obj)

    # Each per-level subquery must run unoffset...
    assert captured_offsets == [0, 0]
    # ...and the requested offset is applied exactly once, to the combined
    # result: 2 + 1 = 3 total rows in, minus an offset of 1 = 2 rows out.
    assert len(result.df) == 2


def test_relative_offset_preserves_inner_bounds(
    processor: QueryContextProcessor,
) -> None:
    """
    Regression test for #40501: Relative time comparison offset should
    preserve inner bounds as the original (unshifted) period, not the shifted one.

    When comparing 2026-05-01 : 2026-05-28 with offset "365 days ago":
    - inner_from_dttm should be 2026-05-01 (original), NOT 2025-05-01 (shifted)
    - inner_to_dttm should be 2026-05-28 (original), NOT 2025-05-28 (shifted)
    """
    from superset.common.query_object import QueryObject
    from superset.models.helpers import ExploreMixin

    datasource: Any = processor._qc_datasource

    for method in (
        "processing_time_offsets",
        "_align_offset_without_time_grain",
        "_coalesce_offset_index",
    ):
        setattr(
            datasource,
            method,
            getattr(ExploreMixin, method).__get__(datasource),
        )

    df = pd.DataFrame(
        {
            "__timestamp": pd.to_datetime(["2026-05-01", "2026-05-15", "2026-05-28"]),
            "sum__num": [100, 200, 300],
        }
    )

    query_object = QueryObject(
        datasource=MagicMock(),
        granularity="ds",
        columns=[],
        metrics=["sum__num"],
        is_timeseries=True,
        time_offsets=["365 days ago"],
        filters=[
            {
                "col": "ds",
                "op": "TEMPORAL_RANGE",
                "val": "2026-05-01 : 2026-05-28",
            }
        ],
    )

    captured: list[dict[str, Any]] = []

    def fake_query(dct: dict[str, Any]) -> MagicMock:
        captured.append(dct)
        result = MagicMock()
        result.df = pd.DataFrame(
            {
                "__timestamp": pd.date_range(
                    start=dct["from_dttm"], periods=3, freq="14D"
                ),
                "sum__num": [1.0, 2.0, 3.0],
            }
        )
        result.query = "SELECT 1"
        return result

    datasource.query = fake_query
    datasource.normalize_df = MagicMock(
        side_effect=lambda offset_df, _query_object: offset_df
    )

    with (
        patch(
            "superset.models.helpers.get_since_until_from_query_object",
            return_value=(pd.Timestamp("2026-05-01"), pd.Timestamp("2026-05-28")),
        ),
        patch(
            "superset.common.utils.query_cache_manager.QueryCacheManager"
        ) as mock_cache_manager,
        patch.object(
            datasource,
            "get_time_grain",
            return_value=None,
        ),
    ):
        mock_cache = MagicMock()
        mock_cache.is_loaded = False
        mock_cache_manager.get.return_value = mock_cache

        datasource.processing_time_offsets(df, query_object, None, None, False)

    # The offset query should use shifted dates for the main window
    assert len(captured) == 1
    assert captured[0]["from_dttm"] == pd.Timestamp("2025-05-01")
    assert captured[0]["to_dttm"] == pd.Timestamp("2025-05-28")

    # The inner bounds (used for series-limit subquery) should be the
    # ORIGINAL unshifted dates, not the shifted ones — this is the fix
    # for #40501. Without the fix, inner_from/to_dttm == shifted dates.
    assert captured[0]["inner_from_dttm"] == pd.Timestamp("2026-05-01")
    assert captured[0]["inner_to_dttm"] == pd.Timestamp("2026-05-28")


def test_get_native_annotation_data_requires_annotation_read_access():
    """Native annotation layers are only served to users who can read them."""
    query_obj = MagicMock()
    query_obj.annotation_layers = [{"sourceType": "NATIVE", "name": "a", "value": 1}]
    with (
        patch(
            "superset.common.query_context_processor.security_manager"
        ) as security_manager_mock,
        patch(
            "superset.common.query_context_processor.AnnotationLayerDAO.find_by_ids",
            return_value=[],
        ) as find_by_ids_mock,
    ):
        # ``can_access`` is synchronous; force a plain Mock so the patched
        # manager doesn't hand back a truthy coroutine that slips past the
        # ``not can_access(...)`` guard.
        security_manager_mock.can_access = MagicMock(return_value=False)
        with pytest.raises(QueryObjectValidationError):
            QueryContextProcessor.get_native_annotation_data(query_obj)
    security_manager_mock.can_access.assert_called_once_with("can_read", "Annotation")
    find_by_ids_mock.assert_not_called()