Files
superset2/tests/unit_tests/queries/test_query_object_prequery.py
Maxime Beauchemin 72e33ba811 more refactoring
2025-08-03 13:39:10 -07:00

116 lines
3.8 KiB
Python

# 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
from superset.common.query_object import QueryObject
def test_get_series_limit_prequery_obj():
"""
Test get_series_limit_prequery_obj method
"""
# Create a QueryObject with series limit settings
query_object = QueryObject(
columns=["country", "year"],
metrics=["sum__sales"],
series_limit=10,
from_dttm=datetime(2020, 1, 1),
to_dttm=datetime(2021, 1, 1),
filters=[{"col": "region", "op": "IN", "val": ["US", "EU"]}],
extras={"time_grain_sqla": "P1D"},
order_desc=False,
)
# Test basic prequery object creation
prequery_obj = query_object.get_series_limit_prequery_obj(
granularity="ds",
inner_from_dttm=None,
inner_to_dttm=None,
)
assert prequery_obj["is_timeseries"] is False
assert prequery_obj["row_limit"] == 10
assert prequery_obj["metrics"] == ["sum__sales"]
assert prequery_obj["granularity"] == "ds"
assert prequery_obj["groupby"] == ["country", "year"]
assert prequery_obj["from_dttm"] == datetime(2020, 1, 1)
assert prequery_obj["to_dttm"] == datetime(2021, 1, 1)
assert prequery_obj["filter"] == [
{"col": "region", "op": "IN", "val": ["US", "EU"]}
]
assert prequery_obj["orderby"] == []
assert prequery_obj["extras"] == {"time_grain_sqla": "P1D"}
assert prequery_obj["order_desc"] is True # Always True for prequery
def test_get_series_limit_prequery_obj_with_overrides():
"""
Test get_series_limit_prequery_obj with inner dates and orderby override
"""
query_object = QueryObject(
columns=["country"],
metrics=["count"],
series_limit=5,
from_dttm=datetime(2020, 1, 1),
to_dttm=datetime(2021, 1, 1),
)
# Test with inner dates and custom orderby
inner_from = datetime(2020, 6, 1)
inner_to = datetime(2020, 12, 31)
custom_orderby = [("sum__revenue", False)]
prequery_obj = query_object.get_series_limit_prequery_obj(
granularity="date_col",
inner_from_dttm=inner_from,
inner_to_dttm=inner_to,
orderby=custom_orderby,
)
assert prequery_obj["from_dttm"] == inner_from
assert prequery_obj["to_dttm"] == inner_to
assert prequery_obj["orderby"] == custom_orderby
def test_get_series_limit_prequery_obj_base_axis_filtering():
"""
Test that base axis columns are filtered out in prequery
"""
# Mock the x-axis column with proper structure for base axis
query_object = QueryObject(
columns=[
{
"label": "__timestamp",
"sqlExpression": "__timestamp",
"columnType": "BASE_AXIS",
},
"country",
"city",
],
metrics=["revenue"],
series_limit=20,
)
prequery_obj = query_object.get_series_limit_prequery_obj(
granularity=None,
inner_from_dttm=None,
inner_to_dttm=None,
)
# The columns in prequery should exclude the base axis column
assert prequery_obj["columns"] == ["country", "city"]