You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by ru...@apache.org on 2026/05/23 04:07:52 UTC
(superset) branch master updated: fix(bigquery): limit result set size to prevent browser memory crashes (#38588)
This is an automated email from the ASF dual-hosted git repository.
rusackas pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new e57387098b4 fix(bigquery): limit result set size to prevent browser memory crashes (#38588)
e57387098b4 is described below
commit e57387098b4450d6077d4c06191598a6b324a38b
Author: Evan Rusackas <ev...@preset.io>
AuthorDate: Fri May 22 21:07:45 2026 -0700
fix(bigquery): limit result set size to prevent browser memory crashes (#38588)
Co-authored-by: Claude Code <no...@anthropic.com>
Co-authored-by: ethan-l-geotab <et...@geotab.com>
---
.../src/query/types/QueryResponse.ts | 1 +
.../src/components/Chart/chartAction.ts | 10 +-
.../src/components/Chart/chartActions.test.ts | 51 ++++++++
superset/charts/schemas.py | 4 +
superset/common/query_context_processor.py | 13 ++
superset/common/utils/query_cache_manager.py | 21 +++-
superset/config.py | 3 +
superset/db_engine_specs/bigquery.py | 109 +++++++++++++++--
.../common/test_query_context_processor.py | 120 +++++++++++++++++++
tests/unit_tests/db_engine_specs/test_bigquery.py | 132 ++++++++++++++++++++-
10 files changed, 454 insertions(+), 10 deletions(-)
diff --git a/superset-frontend/packages/superset-ui-core/src/query/types/QueryResponse.ts b/superset-frontend/packages/superset-ui-core/src/query/types/QueryResponse.ts
index 210685f3480..23c13d3c7c1 100644
--- a/superset-frontend/packages/superset-ui-core/src/query/types/QueryResponse.ts
+++ b/superset-frontend/packages/superset-ui-core/src/query/types/QueryResponse.ts
@@ -77,6 +77,7 @@ export interface ChartDataResponseResult {
// TODO(hainenber): define proper type for below attributes
rejected_filters?: any[];
applied_filters?: any[];
+ warning?: string | null;
/**
* Detected ISO 4217 currency code when AUTO mode is used.
* Returns the currency code if all filtered data contains a single currency,
diff --git a/superset-frontend/src/components/Chart/chartAction.ts b/superset-frontend/src/components/Chart/chartAction.ts
index 9ba8b8dfd55..c44c3ccd341 100644
--- a/superset-frontend/src/components/Chart/chartAction.ts
+++ b/superset-frontend/src/components/Chart/chartAction.ts
@@ -42,7 +42,10 @@ import {
getQuerySettings,
getChartDataUri,
} from 'src/explore/exploreUtils';
-import { addDangerToast } from 'src/components/MessageToasts/actions';
+import {
+ addDangerToast,
+ addWarningToast,
+} from 'src/components/MessageToasts/actions';
import { logEvent } from 'src/logger/actions';
import { Logger, LOG_ACTIONS_LOAD_CHART } from 'src/logger/LogUtils';
import { allowCrossDomain as domainShardingEnabled } from 'src/utils/hostNamesConfig';
@@ -813,6 +816,11 @@ export function exploreJSON(
}),
),
);
+ (queriesResponse as QueryData[]).forEach(response => {
+ if (response.warning) {
+ dispatch(addWarningToast(response.warning, { noDuplicate: true }));
+ }
+ });
return dispatch(
chartUpdateSucceeded(queriesResponse as QueryData[], key as number),
);
diff --git a/superset-frontend/src/components/Chart/chartActions.test.ts b/superset-frontend/src/components/Chart/chartActions.test.ts
index 6379326ad5e..30ba44359aa 100644
--- a/superset-frontend/src/components/Chart/chartActions.test.ts
+++ b/superset-frontend/src/components/Chart/chartActions.test.ts
@@ -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';
@@ -412,6 +413,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
diff --git a/superset/charts/schemas.py b/superset/charts/schemas.py
index e0cee7758c4..0a603cbb6ec 100644
--- a/superset/charts/schemas.py
+++ b/superset/charts/schemas.py
@@ -1561,6 +1561,10 @@ class ChartDataResponseResult(Schema):
required=False,
allow_none=True,
)
+ warning = fields.String(
+ metadata={"description": "Warning message when results were truncated"},
+ allow_none=True,
+ )
class DashboardFilterInfoSchema(Schema):
diff --git a/superset/common/query_context_processor.py b/superset/common/query_context_processor.py
index 52fc6d24f28..56c7ac4f5c8 100644
--- a/superset/common/query_context_processor.py
+++ b/superset/common/query_context_processor.py
@@ -190,6 +190,18 @@ class QueryContextProcessor:
)
cache.df.columns = [unescape_separator(col) for col in cache.df.columns.values]
+ warning: str | None = None
+ 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 = _(
+ "%(prefix)sResults truncated to %(row_count)s rows"
+ " due to memory constraints.",
+ prefix=prefix,
+ row_count=f"{row_count:,}",
+ )
+
return {
"cache_key": cache_key,
"cached_dttm": cache.cache_dttm,
@@ -210,6 +222,7 @@ class QueryContextProcessor:
"from_dttm": query_obj.from_dttm,
"to_dttm": query_obj.to_dttm,
"label_map": label_map,
+ "warning": warning,
}
def query_cache_key(self, query_obj: QueryObject, **kwargs: Any) -> str | None:
diff --git a/superset/common/utils/query_cache_manager.py b/superset/common/utils/query_cache_manager.py
index da2d668e8c9..a7a97e2b874 100644
--- a/superset/common/utils/query_cache_manager.py
+++ b/superset/common/utils/query_cache_manager.py
@@ -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)
diff --git a/superset/config.py b/superset/config.py
index 2a29ce87147..2f83c65e4d6 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -1452,6 +1452,9 @@ DEFAULT_DB_ID = None
# Timeout duration for SQL Lab synchronous queries
SQLLAB_TIMEOUT = int(timedelta(seconds=30).total_seconds())
+# BigQuery max fetch size in MB (limits memory usage when fetching large results)
+BQ_FETCH_MAX_MB = 200
+
# Timeout duration for SQL Lab query validation
SQLLAB_VALIDATION_TIMEOUT = int(timedelta(seconds=10).total_seconds())
diff --git a/superset/db_engine_specs/bigquery.py b/superset/db_engine_specs/bigquery.py
index 8a82623e164..b8bc4783145 100644
--- a/superset/db_engine_specs/bigquery.py
+++ b/superset/db_engine_specs/bigquery.py
@@ -19,6 +19,7 @@ from __future__ import annotations
import logging
import re
+import sys
import urllib
from datetime import datetime
from re import Pattern
@@ -27,6 +28,7 @@ from typing import Any, TYPE_CHECKING, TypedDict
import pandas as pd
from apispec import APISpec
from apispec.ext.marshmallow import MarshmallowPlugin
+from flask import current_app, g, has_app_context, has_request_context
from flask_babel import gettext as __
from marshmallow import fields, Schema
from marshmallow.exceptions import ValidationError
@@ -106,6 +108,11 @@ SYNTAX_ERROR_REGEX = re.compile(
ma_plugin = MarshmallowPlugin()
+# Initial sample size for the progressive fetch in ``fetch_data``. Reading a
+# small first batch lets us measure the row size before deciding how many
+# more rows fit within ``BQ_FETCH_MAX_MB``.
+_BQ_INITIAL_SAMPLE_ROWS = 1000
+
class BigQueryParametersSchema(Schema):
credentials_info = EncryptedString(
@@ -303,13 +310,101 @@ class BigQueryEngineSpec(BaseEngineSpec): # pylint: disable=too-many-public-met
return None
@classmethod
- def fetch_data(cls, cursor: Any, limit: int | None = None) -> list[tuple[Any, ...]]:
- data = super().fetch_data(cursor, limit)
- # Support type BigQuery Row, introduced here PR #4071
- # google.cloud.bigquery.table.Row
- if data and type(data[0]).__name__ == "Row":
- data = [r.values() for r in data] # type: ignore
- return data
+ def fetch_data(cls, cursor: Any, limit: int | None = None) -> list[tuple[Any, ...]]: # noqa: C901
+ """
+ Progressive fetch for BigQuery to prevent browser memory overload.
+
+ Samples a first batch to estimate row size, then extrapolates the
+ total number of rows that fit within ``BQ_FETCH_MAX_MB``.
+ Falls back to the parent implementation on any error.
+ """
+ # ``BQ_FETCH_MAX_MB`` has a default in ``config.py``, so use bracket
+ # access in-context — a missing key should surface as a loud KeyError
+ # rather than be silently masked by a duplicated default here. The
+ # 200 fallback is only used when running outside an app context
+ # (e.g., direct unit-test calls to ``fetch_data``).
+ max_mb: int = (
+ current_app.config["BQ_FETCH_MAX_MB"] if has_app_context() else 200
+ )
+ max_bytes = max_mb * 1024 * 1024
+
+ try:
+ initial_batch_size = (
+ min(_BQ_INITIAL_SAMPLE_ROWS, limit)
+ if limit
+ else _BQ_INITIAL_SAMPLE_ROWS
+ )
+ first_batch: list[Any] = cursor.fetchmany(initial_batch_size)
+
+ if not first_batch:
+ if has_request_context():
+ g.bq_memory_limited = False
+ g.bq_memory_limited_row_count = 0
+ return []
+
+ # Support BigQuery Row objects (PR #4071)
+ if type(first_batch[0]).__name__ == "Row":
+ first_batch = [r.values() for r in first_batch]
+
+ # Estimate how many rows fit in the memory budget.
+ # Sum container + element sizes (one level deep) for a better
+ # estimate. Most BigQuery cell values are primitives (str, int,
+ # float, date), so one level captures the dominant allocation.
+ first_batch_bytes = sum(
+ sys.getsizeof(row) + sum(sys.getsizeof(v) for v in row)
+ for row in first_batch
+ )
+ rows_fetched = len(first_batch)
+ avg_bytes_per_row = first_batch_bytes / rows_fetched
+ total_rows_for_target = int(max_bytes / avg_bytes_per_row)
+
+ if limit:
+ total_rows_for_target = min(limit, total_rows_for_target)
+
+ remaining_rows = total_rows_for_target - rows_fetched
+
+ # First batch already covers the budget or the result set
+ if rows_fetched < initial_batch_size or remaining_rows <= 0:
+ memory_limited = (
+ remaining_rows <= 0 and rows_fetched == initial_batch_size
+ )
+ if has_request_context():
+ g.bq_memory_limited = memory_limited
+ g.bq_memory_limited_row_count = len(first_batch)
+ return first_batch
+
+ # Fetch one extra row to confirm truncation without false positives
+ second_batch: list[Any] = cursor.fetchmany(remaining_rows + 1) or []
+ if second_batch and type(second_batch[0]).__name__ == "Row":
+ second_batch = [r.values() for r in second_batch]
+
+ # Truncation is confirmed only when more rows exist beyond the budget
+ memory_limited = len(second_batch) > remaining_rows
+ if memory_limited:
+ second_batch = second_batch[:remaining_rows]
+
+ data = first_batch + second_batch
+
+ if has_request_context():
+ g.bq_memory_limited = memory_limited
+ g.bq_memory_limited_row_count = len(data)
+ return data
+
+ except Exception: # pylint: disable=broad-except
+ # Broad catch on purpose: any failure in the size-estimation /
+ # progressive-fetch path (BigQuery DB-API errors, network or
+ # auth timeouts mid-fetch, ``sys.getsizeof`` raising on an
+ # unexpected cell type, or a future ``Row`` subclass we don't
+ # know how to unwrap) must degrade gracefully to the parent's
+ # straight fetch so the user still gets data.
+ # Fallback to parent implementation
+ data = super().fetch_data(cursor, limit)
+ if data and type(data[0]).__name__ == "Row":
+ data = [r.values() for r in data] # type: ignore
+ if has_request_context():
+ g.bq_memory_limited = False
+ g.bq_memory_limited_row_count = len(data) if data else 0
+ return data
@staticmethod
def _mutate_label(label: str) -> str:
diff --git a/tests/unit_tests/common/test_query_context_processor.py b/tests/unit_tests/common/test_query_context_processor.py
index 8e33b4d448f..2b6a1d4f0c5 100644
--- a/tests/unit_tests/common/test_query_context_processor.py
+++ b/tests/unit_tests/common/test_query_context_processor.py
@@ -1155,6 +1155,7 @@ def test_get_df_payload_validates_before_cache_key_generation():
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
@@ -1530,6 +1531,7 @@ def test_force_cached_normalizes_totals_query_row_limit():
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
@@ -1586,6 +1588,8 @@ def test_get_df_payload_invalidates_cache_missing_applied_filter_columns():
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()
@@ -1606,3 +1610,119 @@ def test_get_df_payload_invalidates_cache_missing_applied_filter_columns():
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
diff --git a/tests/unit_tests/db_engine_specs/test_bigquery.py b/tests/unit_tests/db_engine_specs/test_bigquery.py
index a7231508e8e..a8695008282 100644
--- a/tests/unit_tests/db_engine_specs/test_bigquery.py
+++ b/tests/unit_tests/db_engine_specs/test_bigquery.py
@@ -18,7 +18,7 @@
# pylint: disable=line-too-long, import-outside-toplevel, protected-access, invalid-name
from datetime import datetime
-from typing import Optional
+from typing import Any, Optional
from unittest import mock
import pytest
@@ -528,3 +528,133 @@ def test_get_view_names_excludes_materialized_views() -> None:
assert "table_type = 'VIEW'" in executed_query
# Ensure it's not querying for materialized views
assert "MATERIALIZED VIEW" not in executed_query
+
+
+def _patch_bq_fetch_deps(
+ mocker: MockerFixture, max_mb: int = 200
+) -> tuple[mock.MagicMock, mock.MagicMock]:
+ """Helper to patch Flask g and current_app for BigQuery fetch_data tests."""
+ flask_g = mocker.patch("superset.db_engine_specs.bigquery.g")
+ app = mocker.patch("superset.db_engine_specs.bigquery.current_app")
+ # Make current_app truthy and .config.get() return a plain int
+ app.__bool__ = mock.Mock(return_value=True)
+ app.config = mock.MagicMock()
+ app.config.get = mock.Mock(return_value=max_mb)
+ return flask_g, app
+
+
+def test_fetch_data_within_memory_limit(mocker: MockerFixture) -> None:
+ """
+ Test that fetch_data returns all rows when the result fits within the
+ configured memory limit.
+ """
+ from superset.db_engine_specs.bigquery import BigQueryEngineSpec
+
+ rows = [(1, "a"), (2, "b"), (3, "c")]
+
+ cursor = mock.MagicMock()
+ # First fetchmany returns all rows; the result set is smaller than limit
+ cursor.fetchmany.return_value = rows
+
+ flask_g, _ = _patch_bq_fetch_deps(mocker, max_mb=200)
+
+ result = BigQueryEngineSpec.fetch_data(cursor, limit=100)
+
+ assert result == rows
+ assert flask_g.bq_memory_limited is False
+ assert flask_g.bq_memory_limited_row_count == 3
+
+
+def test_fetch_data_truncated_by_memory_limit(mocker: MockerFixture) -> None:
+ """
+ Test that fetch_data truncates results and sets the memory_limited flag
+ when the memory budget is exceeded.
+
+ We use a very small budget (1 MB) so that after the first batch the
+ method computes ``remaining_rows <= 0``, hitting the truncation path.
+ """
+ from superset.db_engine_specs.bigquery import BigQueryEngineSpec
+
+ # 1000 rows of ~10KB each --> first batch ~10 MB >> 1 MB budget
+ first_batch = [(i, "x" * 10_000) for i in range(1000)]
+
+ cursor = mock.MagicMock()
+ cursor.fetchmany.return_value = first_batch
+
+ # 1 MB budget: first batch exceeds it, so remaining_rows <= 0
+ flask_g, _ = _patch_bq_fetch_deps(mocker, max_mb=1)
+
+ result = BigQueryEngineSpec.fetch_data(cursor, limit=None)
+
+ assert result == first_batch
+ assert flask_g.bq_memory_limited is True
+ assert flask_g.bq_memory_limited_row_count == len(first_batch)
+
+
+def test_fetch_data_empty_result(mocker: MockerFixture) -> None:
+ """
+ Test that fetch_data handles an empty result set gracefully.
+ """
+ from superset.db_engine_specs.bigquery import BigQueryEngineSpec
+
+ cursor = mock.MagicMock()
+ cursor.fetchmany.return_value = []
+
+ flask_g, _ = _patch_bq_fetch_deps(mocker, max_mb=200)
+
+ result = BigQueryEngineSpec.fetch_data(cursor, limit=100)
+
+ assert result == []
+ assert flask_g.bq_memory_limited is False
+ assert flask_g.bq_memory_limited_row_count == 0
+
+
+def test_fetch_data_fallback_on_exception(mocker: MockerFixture) -> None:
+ """
+ Test that fetch_data falls back to the parent implementation when the
+ progressive fetch raises an exception.
+ """
+ from superset.db_engine_specs.bigquery import BigQueryEngineSpec
+
+ cursor = mock.MagicMock()
+ cursor.fetchmany.side_effect = RuntimeError("cursor error")
+ cursor.fetchall.return_value = [(1, "a"), (2, "b")]
+ cursor.description = [("col1", None), ("col2", None)]
+
+ flask_g, _ = _patch_bq_fetch_deps(mocker, max_mb=200)
+
+ result = BigQueryEngineSpec.fetch_data(cursor, limit=None)
+
+ assert result == [(1, "a"), (2, "b")]
+ assert flask_g.bq_memory_limited is False
+ assert flask_g.bq_memory_limited_row_count == 2
+
+
+def test_fetch_data_converts_bigquery_row_objects(mocker: MockerFixture) -> None:
+ """
+ Test that BigQuery Row objects are converted to plain values.
+ """
+ from superset.db_engine_specs.bigquery import BigQueryEngineSpec
+
+ class FakeRow:
+ """Mimics google.cloud.bigquery.table.Row"""
+
+ def __init__(self, vals: tuple[Any, ...]) -> None:
+ self._vals = vals
+
+ def values(self) -> tuple[Any, ...]:
+ return self._vals
+
+ FakeRow.__name__ = "Row"
+
+ rows = [FakeRow((1, "a")), FakeRow((2, "b"))]
+
+ cursor = mock.MagicMock()
+ cursor.fetchmany.return_value = rows
+
+ flask_g, _ = _patch_bq_fetch_deps(mocker, max_mb=200)
+
+ result = BigQueryEngineSpec.fetch_data(cursor, limit=100)
+
+ assert result == [(1, "a"), (2, "b")]
+ assert flask_g.bq_memory_limited is False