QT/superset2
1
0
Fork 0
mirror of https://github.com/apache/superset.git synced 2026-05-21 15:55:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(bigquery): address deferred review follow-ups

Browse Source 
Cache persistence: persist bq_memory_limited flag in QueryCacheManager
so truncation warnings survive cache hits. The flag is now saved into
the cache value dict in set_query_result (reading from g and resetting
it there), and restored from cache in get(). get_df_payload now reads
cache.bq_memory_limited instead of g, removing the need for g entirely
in query_context_processor.py.

Frontend test: add two tests to chartActions.test.ts verifying that
addWarningToast is dispatched (with noDuplicate: true) when a query
response carries a warning field, and not dispatched when it doesn't.

Updated existing test_query_context_processor tests to set the flag on
mock_cache directly rather than patching g.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Claude Code
2026-04-07 13:43:11 -07:00
parent 1894b77b44
commit bbbd1490ef
4 changed files with 79 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
superset-frontend/src/components/Chart/chartActions.test.ts
View File

@@ -31,6 +31,7 @@ import {
AnnotationSourceType,
AnnotationStyle,
} from '@superset-ui/core';
import * as toastActions from 'src/components/MessageToasts/actions';
import { LOG_EVENT } from 'src/logger/actions';
import * as exploreUtils from 'src/explore/exploreUtils';
import * as actions from 'src/components/Chart/chartAction';
@@ -340,6 +341,56 @@ describe('chart actions', () => {
);
expect(result).toEqual([1, 2, 3]);
});
test('dispatches addWarningToast when a query response includes a warning', async () => {
const warningMessage =
'Results truncated to 1,000 rows due to memory constraints.';
fetchMock.removeRoute(MOCK_URL);
fetchMock.post(
`glob:*${MOCK_URL}*`,
{ result: [{ warning: warningMessage }] },
{ name: MOCK_URL },
);
const addWarningToastSpy = jest.spyOn(toastActions, 'addWarningToast');
const actionThunk = actions.postChartFormData(
{ viz_type: 'my_viz' } as QueryFormData,
false,
undefined,
undefined,
);
await actionThunk(
dispatch as unknown as actions.ChartThunkDispatch,
mockGetState as unknown as () => actions.RootState,
undefined,
);
expect(addWarningToastSpy).toHaveBeenCalledWith(warningMessage, {
noDuplicate: true,
});
addWarningToastSpy.mockRestore();
fetchMock.removeRoute(MOCK_URL);
setupDefaultFetchMock();
});
test('does not dispatch addWarningToast when no query response has a warning', async () => {
const addWarningToastSpy = jest.spyOn(toastActions, 'addWarningToast');
const actionThunk = actions.postChartFormData(
{ viz_type: 'my_viz' } as QueryFormData,
false,
undefined,
undefined,
);
await actionThunk(
dispatch as unknown as actions.ChartThunkDispatch,
mockGetState as unknown as () => actions.RootState,
undefined,
);
expect(addWarningToastSpy).not.toHaveBeenCalled();
addWarningToastSpy.mockRestore();
});
});
// eslint-disable-next-line no-restricted-globals -- TODO: Migrate from describe blocks

10
superset/common/query_context_processor.py
View File

@@ -21,7 +21,7 @@ import re
from typing import Any, cast, ClassVar, Sequence, TYPE_CHECKING
import pandas as pd
from flask import current_app, g
from flask import current_app
from flask_babel import gettext as _
from superset.common.chart_data import ChartDataResultFormat
@@ -191,12 +191,8 @@ class QueryContextProcessor:
cache.df.columns = [unescape_separator(col) for col in cache.df.columns.values]
warning: str | None = None
if getattr(g, "bq_memory_limited", False):
row_count = getattr(g, "bq_memory_limited_row_count", len(cache.df))
# Reset flags immediately so subsequent queries in the same request
# don't inherit this warning
g.bq_memory_limited = False
g.bq_memory_limited_row_count = 0
if cache.bq_memory_limited:
row_count = cache.bq_memory_limited_row_count
chart_id = (self._query_context.form_data or {}).get("slice_id", "")
prefix = f"Chart {chart_id}: " if chart_id else ""
warning = _(

21
superset/common/utils/query_cache_manager.py
View File

@@ -20,7 +20,7 @@ import logging
from datetime import datetime, timezone
from typing import Any
from flask import current_app
from flask import current_app, g, has_request_context
from flask_caching import Cache
from pandas import DataFrame
@@ -86,6 +86,8 @@ class QueryCacheManager:
self.cache_value = cache_value
self.sql_rowcount = sql_rowcount
self.queried_dttm = queried_dttm
self.bq_memory_limited: bool = False
self.bq_memory_limited_row_count: int = 0
# pylint: disable=too-many-arguments
def set_query_result(
@@ -123,6 +125,15 @@ class QueryCacheManager:
)
self.is_loaded = True
# Capture BigQuery memory-limit flag so it survives cache hits
if has_request_context():
self.bq_memory_limited = getattr(g, "bq_memory_limited", False)
self.bq_memory_limited_row_count = getattr(
g, "bq_memory_limited_row_count", 0
)
g.bq_memory_limited = False
g.bq_memory_limited_row_count = 0
value = {
"df": self.df,
"query": self.query,
@@ -133,6 +144,8 @@ class QueryCacheManager:
"sql_rowcount": self.sql_rowcount,
"queried_dttm": self.queried_dttm,
"dttm": self.queried_dttm, # Backwards compatibility
"bq_memory_limited": self.bq_memory_limited,
"bq_memory_limited_row_count": self.bq_memory_limited_row_count,
}
if self.is_loaded and key and self.status != QueryStatus.FAILED:
self.set(
@@ -193,6 +206,12 @@ class QueryCacheManager:
"queried_dttm", cache_value.get("dttm")
)
query_cache.cache_value = cache_value
query_cache.bq_memory_limited = cache_value.get(
"bq_memory_limited", False
)
query_cache.bq_memory_limited_row_count = cache_value.get(
"bq_memory_limited_row_count", 0
)
current_app.config["STATS_LOGGER"].incr("loaded_from_cache")
except KeyError as ex:
logger.exception(ex)

16
tests/unit_tests/common/test_query_context_processor.py
View File

@@ -1426,17 +1426,14 @@ def test_get_df_payload_bq_memory_limited_warning():
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):
# Simulate BigQuery memory-limited flag being set on Flask g
with patch("superset.common.query_context_processor.g") as mock_g:
mock_g.bq_memory_limited = True
mock_g.bq_memory_limited_row_count = 5000
result = processor.get_df_payload(query_obj, force_cached=False)
result = processor.get_df_payload(query_obj, force_cached=False)
assert result["warning"] is not None
assert "Chart 42" in result["warning"]
@@ -1489,15 +1486,12 @@ def test_get_df_payload_no_warning_when_not_memory_limited():
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):
# g.bq_memory_limited is not set (default)
with patch("superset.common.query_context_processor.g") as mock_g:
mock_g.bq_memory_limited = False
result = processor.get_df_payload(query_obj, force_cached=False)
result = processor.get_df_payload(query_obj, force_cached=False)
assert result["warning"] is None