# 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. import pytest from pandas import DataFrame, Series, Timestamp from pandas.testing import assert_frame_equal from pytest import fixture, mark # noqa: PT013 from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType from superset.common.query_context import QueryContext from superset.common.query_context_processor import QueryContextProcessor from superset.connectors.sqla.models import BaseDatasource from superset.constants import TimeGrain from superset.exceptions import QueryObjectValidationError from superset.models.helpers import ExploreMixin # Create processor and bind ExploreMixin methods to datasource processor = QueryContextProcessor( QueryContext( datasource=BaseDatasource(), queries=[], result_type=ChartDataResultType.COLUMNS, form_data={}, slice_=None, result_format=ChartDataResultFormat.CSV, cache_values={}, ) ) # Bind ExploreMixin methods to datasource for testing # Type annotation needed because _qc_datasource is typed as Explorable in protocol _datasource: BaseDatasource = processor._qc_datasource # type: ignore _datasource.add_offset_join_column = ExploreMixin.add_offset_join_column.__get__( _datasource ) _datasource.join_offset_dfs = ExploreMixin.join_offset_dfs.__get__(_datasource) _datasource.is_valid_date_range = ExploreMixin.is_valid_date_range.__get__(_datasource) _datasource._determine_join_keys = ExploreMixin._determine_join_keys.__get__( _datasource ) _datasource._perform_join = ExploreMixin._perform_join.__get__(_datasource) _datasource._apply_cleanup_logic = ExploreMixin._apply_cleanup_logic.__get__( _datasource ) _datasource._coalesce_offset_index = ExploreMixin._coalesce_offset_index.__get__( _datasource ) # Static methods don't need binding - assign directly _datasource.generate_join_column = ExploreMixin.generate_join_column _datasource.is_valid_date_range_static = ExploreMixin.is_valid_date_range_static # Convenience reference for backward compatibility in tests query_context_processor = _datasource @fixture def make_join_column_producer(): def join_column_producer(row: Series, column_index: int) -> str: return "CUSTOM_FORMAT" return join_column_producer @mark.parametrize( ("time_grain", "expected"), [ (TimeGrain.WEEK, "2020-W01"), (TimeGrain.MONTH, "2020-01"), (TimeGrain.QUARTER, "2020-Q1"), (TimeGrain.YEAR, "2020"), ], ) def test_join_column(time_grain: str, expected: str): df = DataFrame({"ds": [Timestamp("2020-01-07")]}) column_name = "join_column" query_context_processor.add_offset_join_column(df, column_name, time_grain) result = DataFrame({"ds": [Timestamp("2020-01-07")], column_name: [expected]}) assert_frame_equal(df, result) def test_join_column_producer(make_join_column_producer): df = DataFrame({"ds": [Timestamp("2020-01-07")]}) column_name = "join_column" query_context_processor.add_offset_join_column( df, column_name, TimeGrain.YEAR, None, make_join_column_producer ) result = DataFrame( {"ds": [Timestamp("2020-01-07")], column_name: ["CUSTOM_FORMAT"]} ) assert_frame_equal(df, result) def test_join_offset_dfs_no_offsets(): df = DataFrame({"A": ["2021-01-01", "2021-02-01", "2021-03-01"]}) offset_dfs = {} time_grain = "YEAR" join_keys = ["A"] result = query_context_processor.join_offset_dfs( df, offset_dfs, time_grain, join_keys ) assert_frame_equal(df, result) def test_join_offset_dfs_with_offsets(): df = DataFrame({"A": ["2021-01-01", "2021-02-01", "2021-03-01"]}) offset_df = DataFrame( {"A": ["2021-02-01", "2021-03-01", "2021-04-01"], "B": [5, 6, 7]} ) offset_dfs = {"1_YEAR": offset_df} time_grain = "YEAR" join_keys = ["A"] expected = DataFrame( {"A": ["2021-01-01", "2021-02-01", "2021-03-01"], "B": [None, 5, 6]} ) result = query_context_processor.join_offset_dfs( df, offset_dfs, time_grain, join_keys ) assert_frame_equal(expected, result) def test_join_offset_dfs_with_multiple_offsets(): df = DataFrame({"A": ["2021-01-01", "2021-02-01", "2021-03-01"]}) offset_df1 = DataFrame( {"A": ["2021-02-01", "2021-03-01", "2021-04-01"], "B": [5, 6, 7]} ) offset_df2 = DataFrame( {"A": ["2021-03-01", "2021-04-01", "2021-05-01"], "C": [8, 9, 10]} ) offset_dfs = {"1_YEAR": offset_df1, "2_YEAR": offset_df2} time_grain = "YEAR" join_keys = ["A"] expected = DataFrame( { "A": ["2021-01-01", "2021-02-01", "2021-03-01"], "B": [None, 5, 6], "C": [None, None, 8], } ) result = query_context_processor.join_offset_dfs( df, offset_dfs, time_grain, join_keys ) assert_frame_equal(expected, result) def test_join_offset_dfs_with_month_granularity(): df = DataFrame( { "A": [ "2021-01-01", "2021-01-15", "2021-02-01", "2021-02-15", "2021-03-01", "2021-03-15", ], "D": [1, 2, 3, 4, 5, 6], } ) offset_df = DataFrame( { "A": [ "2021-02-01", "2021-02-15", "2021-03-01", "2021-03-15", "2021-04-01", "2021-04-15", ], "B": [5, 6, 7, 8, 9, 10], } ) offset_dfs = {"1_MONTH": offset_df} time_grain = "MONTH" join_keys = ["A"] expected = DataFrame( { "A": [ "2021-01-01", "2021-01-15", "2021-02-01", "2021-02-15", "2021-03-01", "2021-03-15", ], "D": [1, 2, 3, 4, 5, 6], "B": [None, None, 5, 6, 7, 8], } ) result = query_context_processor.join_offset_dfs( df, offset_dfs, time_grain, join_keys ) assert_frame_equal(expected, result) def test_join_offset_dfs_full_range_keeps_historical_tail(): """ With full_range=True the offset (historical) series keeps its full time range even when the main series ends earlier. Simulates "today so far" (main, ends at 01:00) compared against "1 day ago" (a complete prior day, runs to 02:00). The 02:00 historical point must survive and be aligned onto today's axis, with the main metric left null there. """ # Main series: today, only two hours of data so far. df = DataFrame( { "A": [Timestamp("2021-01-02 00:00"), Timestamp("2021-01-02 01:00")], "V": [1.0, 2.0], } ) # Offset series: the full prior day (already renamed metric column "B"). offset_df = DataFrame( { "A": [ Timestamp("2021-01-01 00:00"), Timestamp("2021-01-01 01:00"), Timestamp("2021-01-01 02:00"), ], "B": [10.0, 20.0, 30.0], } ) offset_dfs = {"1 day ago": offset_df} time_grain = TimeGrain.HOUR join_keys = ["A"] expected = DataFrame( { "A": [ Timestamp("2021-01-02 00:00"), Timestamp("2021-01-02 01:00"), Timestamp("2021-01-02 02:00"), ], "V": [1.0, 2.0, None], "B": [10.0, 20.0, 30.0], } ) result = query_context_processor.join_offset_dfs( df, offset_dfs, time_grain, join_keys, full_range=True ) assert_frame_equal(expected, result) def test_join_offset_dfs_full_range_disabled_truncates_historical(): """The default (full_range=False) left join drops the historical 02:00 point.""" df = DataFrame( { "A": [Timestamp("2021-01-02 00:00"), Timestamp("2021-01-02 01:00")], "V": [1.0, 2.0], } ) offset_df = DataFrame( { "A": [ Timestamp("2021-01-01 00:00"), Timestamp("2021-01-01 01:00"), Timestamp("2021-01-01 02:00"), ], "B": [10.0, 20.0, 30.0], } ) offset_dfs = {"1 day ago": offset_df} expected = DataFrame( { "A": [Timestamp("2021-01-02 00:00"), Timestamp("2021-01-02 01:00")], "V": [1.0, 2.0], "B": [10.0, 20.0], } ) result = query_context_processor.join_offset_dfs( df, offset_dfs, TimeGrain.HOUR, ["A"], full_range=False ) assert_frame_equal(expected, result) def test_join_offset_dfs_totals_query_no_dimensions(): """ Test time offset join for totals query with no dimension columns. This simulates a table chart totals query where: - columns=[] (no dimensions, only metrics) - time_offsets=["1 month ago"] - The dataframes only contain metric columns (no datetime column) The join should use the __temp_join_key__ fallback mechanism to properly join the offset data. """ # Main totals query result - only has metric column, no datetime df = DataFrame({"Total Cost": [54211.76]}) # Offset query result - renamed metric column offset_df = DataFrame({"Total Cost__1 month ago": [48000.50]}) offset_dfs = {"1 month ago": offset_df} time_grain = "P1D" # Daily grain from extras join_keys = [] # No dimension columns for totals query expected = DataFrame( {"Total Cost": [54211.76], "Total Cost__1 month ago": [48000.50]} ) result = query_context_processor.join_offset_dfs( df, offset_dfs, time_grain, join_keys ) assert_frame_equal(expected, result) def test_join_offset_dfs_raises_without_time_grain(): """Time comparison with relative offsets requires a time grain.""" df = DataFrame({"ds": [Timestamp("2021-01-01")], "D": [1]}) offset_df = DataFrame({"ds": [Timestamp("2021-02-01")], "B": [5]}) offset_dfs = {"1 year ago": offset_df} with pytest.raises( QueryObjectValidationError, match="Time Grain must be specified" ): query_context_processor.join_offset_dfs( df, offset_dfs, time_grain=None, join_keys=["ds"] ) def test_join_offset_dfs_allows_non_temporal_join_without_time_grain(): """Time comparison without time grain is valid when join keys are non-temporal.""" df = DataFrame({"country": ["US", "UK"], "metric": [10, 20]}) offset_df = DataFrame({"country": ["US", "UK"], "metric__1 year ago": [8, 15]}) offset_dfs = {"1 year ago": offset_df} result = query_context_processor.join_offset_dfs( df, offset_dfs, time_grain=None, join_keys=["country"] ) assert "metric__1 year ago" in result.columns def test_join_offset_dfs_raises_when_temporal_key_not_first(): """Temporal join key detection works even when it's not the first key.""" df = DataFrame( { "country": ["US", "UK"], "ds": [Timestamp("2021-01-01"), Timestamp("2021-02-01")], "D": [1, 2], } ) offset_df = DataFrame( { "country": ["US", "UK"], "ds": [Timestamp("2021-03-01"), Timestamp("2021-04-01")], "B": [5, 6], } ) offset_dfs = {"1 year ago": offset_df} with pytest.raises( QueryObjectValidationError, match="Time Grain must be specified" ): query_context_processor.join_offset_dfs( df, offset_dfs, time_grain=None, join_keys=["country", "ds"] )