aminghadersohi commented on code in PR #44604:
URL: https://github.com/apache/superset/pull/44604#discussion_r4161027088
##########
superset/sql/execution/executor.py:
##########
@@ -186,6 +187,65 @@ def build_statement_blocks(
return parsed_script, blocks
+class _LimitedCursor:
+ """Bound cumulative cursor reads without bypassing engine fetch
processing."""
+
+ def __init__(self, cursor: Any, limit: int) -> None:
+ """Wrap a cursor with a shared budget for all row-reading methods."""
+ self._cursor = cursor
+ self._remaining = limit
+
+ def __getattr__(self, name: str) -> Any:
+ """Delegate metadata and driver-specific methods to the real cursor."""
+ return getattr(self._cursor, name)
+
+ @property
+ def arraysize(self) -> int:
+ """Expose the driver's default fetch batch size."""
+ return self._cursor.arraysize
+
+ @arraysize.setter
+ def arraysize(self, value: int) -> None:
+ """Preserve engine-specific cursor batch-size configuration."""
+ self._cursor.arraysize = value
+
+ def fetchmany(self, size: int | None = None) -> list[Any]:
+ """Read no more than the remaining budget, including across batches."""
+ size = self.arraysize if size is None else size
+ size = max(0, min(size, self._remaining))
+ # Some drivers interpret zero as unbounded, so do not call them at all.
+ if not size:
+ return []
+ rows = self._cursor.fetchmany(size)
+ self._remaining -= len(rows)
+ return rows
+
+ def check_truncated(self, db_engine_spec: type[BaseEngineSpec]) -> bool:
+ """Probe one extra row using the engine's fetch and error handling."""
+ return self._remaining == 0 and bool(
+ db_engine_spec.fetch_data(_LimitedCursor(self._cursor, 1))
+ )
Review Comment:
Agreed, degrading is the better trade. Done in 5eee71d8c7: `check_truncated`
now catches a probe failure, logs a warning and returns `True`, so the rows
already fetched within the budget survive and the caller is told the result may
be partial. `test_execute_truncation_probe_engine_errors` now asserts the
`OSError` case returns both rows with `truncated is True` (Drill EOF still
returns `truncated is False`).
##########
superset/sql/execution/executor.py:
##########
@@ -186,6 +187,65 @@ def build_statement_blocks(
return parsed_script, blocks
+class _LimitedCursor:
+ """Bound cumulative cursor reads without bypassing engine fetch
processing."""
+
+ def __init__(self, cursor: Any, limit: int) -> None:
+ """Wrap a cursor with a shared budget for all row-reading methods."""
+ self._cursor = cursor
+ self._remaining = limit
+
+ def __getattr__(self, name: str) -> Any:
+ """Delegate metadata and driver-specific methods to the real cursor."""
+ return getattr(self._cursor, name)
+
+ @property
+ def arraysize(self) -> int:
+ """Expose the driver's default fetch batch size."""
+ return self._cursor.arraysize
+
+ @arraysize.setter
+ def arraysize(self, value: int) -> None:
+ """Preserve engine-specific cursor batch-size configuration."""
+ self._cursor.arraysize = value
+
+ def fetchmany(self, size: int | None = None) -> list[Any]:
+ """Read no more than the remaining budget, including across batches."""
+ size = self.arraysize if size is None else size
+ size = max(0, min(size, self._remaining))
+ # Some drivers interpret zero as unbounded, so do not call them at all.
+ if not size:
+ return []
+ rows = self._cursor.fetchmany(size)
+ self._remaining -= len(rows)
+ return rows
+
+ def check_truncated(self, db_engine_spec: type[BaseEngineSpec]) -> bool:
+ """Probe one extra row using the engine's fetch and error handling."""
+ return self._remaining == 0 and bool(
+ db_engine_spec.fetch_data(_LimitedCursor(self._cursor, 1))
+ )
+
+ def fetchall(self) -> list[Any]:
+ """Translate an unbounded read into a bounded driver fetch."""
+ return self.fetchmany(self._remaining)
Review Comment:
Closed in 5eee71d8c7: `_LimitedCursor.fetchall()` now loops
`fetchmany(remaining)` until the budget is spent or a batch comes back empty.
New `test_limited_cursor_fetchall_reads_through_short_batches` uses a cursor
that serves two rows per call with ten available and a budget of five:
`fetch_data` returns 5 rows, the driver sees `fetchmany(5), (3), (1)` plus the
one-row probe, and `check_truncated` reports `True`.
##########
superset/db_engine_specs/bigquery.py:
##########
@@ -488,16 +494,15 @@ def fetch_data(cls, cursor: Any, limit: int | None =
None) -> list[tuple[Any, ..
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
+ return FetchedRows(data, truncated=memory_limited)
+
+ except Exception as ex: # pylint: disable=broad-except
+ # A forward-only cursor cannot replay the consumed sample. Falling
+ # back after an estimation, second-batch, or EOF-probe failure
would
+ # silently discard it and return only the remaining rows as
success.
+ if first_batch:
+ raise cls.get_dbapi_mapped_exception(ex) from ex
Review Comment:
Added in 5eee71d8c7: `UPDATING.md` now says `BigQueryEngineSpec.fetch_data`
no longer falls back to a plain `fetchall()` after its initial sample, so a
mid-fetch driver error fails the query on every caller, including legacy SQL
Lab execution and dataset column discovery, while an error on the initial read
still falls back as before.
--
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]
