aminghadersohi commented on code in PR #44604:
URL: https://github.com/apache/superset/pull/44604#discussion_r4141001725


##########
tests/unit_tests/db_engine_specs/test_bigquery.py:
##########
@@ -832,6 +869,43 @@ def test_fetch_data_fallback_on_exception(mocker: 
MockerFixture) -> None:
     assert flask_g.bq_memory_limited_row_count == 2
 
 
[email protected]("bounded", [False, True])
[email protected]("failure_stage", ["probe", "second_batch", 
"estimate"])
+def test_fetch_data_does_not_discard_sample_after_error(
+    mocker: MockerFixture, bounded: bool, failure_stage: str
+) -> None:
+    """A failure after consuming rows must not return only the cursor 
remainder."""
+    from superset.db_engine_specs.bigquery import BigQueryEngineSpec
+    from superset.db_engine_specs.exceptions import 
SupersetDBAPIConnectionError
+    from superset.sql.execution.executor import _LimitedCursor
+
+    _, app = _patch_bq_fetch_deps(mocker)
+    mocker.patch("superset.db_engine_specs.bigquery.has_app_context", 
return_value=True)
+    app.config = {"BQ_FETCH_MAX_MB": 1}
+    mocker.patch("superset.db_engine_specs.bigquery._BQ_INITIAL_SAMPLE_ROWS", 
2)
+    mocker.patch.object(
+        BigQueryEngineSpec,
+        "get_dbapi_exception_mapping",
+        return_value={OSError: SupersetDBAPIConnectionError},
+    )
+    row = ("x" * (600_000 if failure_stage == "probe" else 300_000),)

Review Comment:
   Fixed in 0c7cc2082c12c2c7019a6697cef5307316f87d01: named the probe and 
second-batch cell sizes and documented the 1 MiB budget math. Two roughly 600 
KB rows exceed the budget and trigger the EOF probe; roughly 300 KB rows leave 
room for the second batch. All bounded/unbounded probe, second-batch, and 
estimation cases remain. The requested BigQuery and SQL execution suites pass 
(363 tests); all applicable pre-commit hooks pass, including mypy.



##########
tests/unit_tests/db_engine_specs/test_bigquery.py:
##########
@@ -832,6 +869,43 @@ def test_fetch_data_fallback_on_exception(mocker: 
MockerFixture) -> None:
     assert flask_g.bq_memory_limited_row_count == 2
 
 
[email protected]("bounded", [False, True])
[email protected]("failure_stage", ["probe", "second_batch", 
"estimate"])
+def test_fetch_data_does_not_discard_sample_after_error(
+    mocker: MockerFixture, bounded: bool, failure_stage: str
+) -> None:
+    """A failure after consuming rows must not return only the cursor 
remainder."""
+    from superset.db_engine_specs.bigquery import BigQueryEngineSpec
+    from superset.db_engine_specs.exceptions import 
SupersetDBAPIConnectionError
+    from superset.sql.execution.executor import _LimitedCursor
+
+    _, app = _patch_bq_fetch_deps(mocker)
+    mocker.patch("superset.db_engine_specs.bigquery.has_app_context", 
return_value=True)
+    app.config = {"BQ_FETCH_MAX_MB": 1}
+    mocker.patch("superset.db_engine_specs.bigquery._BQ_INITIAL_SAMPLE_ROWS", 
2)
+    mocker.patch.object(
+        BigQueryEngineSpec,
+        "get_dbapi_exception_mapping",
+        return_value={OSError: SupersetDBAPIConnectionError},
+    )
+    row = ("x" * (600_000 if failure_stage == "probe" else 300_000),)
+    cursor = mock.MagicMock()
+    cursor.description = [("n", "STRING", None, None, None, None, None)]
+    cursor.fetchmany.side_effect = [[row, row], OSError("read failed"), [row]]

Review Comment:
   Fixed in 0c7cc2082c12c2c7019a6697cef5307316f87d01: derive fetchmany 
side_effect from failure_stage. Estimation consumes only the initial sample 
before getsizeof raises; probe and second-batch cases consume the sample and 
then raise on the second fetch. Removed both unreachable entries while 
retaining all six error cases and their fetch-count and no-fetchall assertions. 
The requested BigQuery and SQL execution suites pass (363 tests); all 
applicable pre-commit hooks pass, including mypy.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to