diff --git a/CHANGELOG.md b/CHANGELOG.md index e30aebb..3540d79 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,27 @@ tag releases both in lockstep, so entries below are keyed by the engine version. ### Fixed +- **`explore query` no longer fails inside its own confirmed budget on + BigQuery** ([#320]). The server-side `maximum_bytes_billed` cap was set to + this command's own reservation against the cumulative session ceiling + (sized to the dry-run estimate), not the wider per-command budget the + operator actually confirmed. BigQuery's own execution-time rounding of + bytes billed can exceed any dry-run estimate regardless of how accurate + that estimate was, so a multi-table statement confirmed at a budget six + times its estimate still failed with `bytesBilledLimitExceeded`, a + self-imposed cap the error message never named as the cause. + + BigQuery's own refusal already states the exact byte count it needed + (`"163595928. 164626432 or higher required."`), so a statement that hits + the cap now widens its charge to precisely that number and retries once, + rather than guessing at a margin or handing the warehouse a cap wider than + what this command actually reserved. A retry that still can't fit the + confirmed ceiling refuses on the real number, exactly as it would have + without the retry; the concurrency guarantee a cumulative session ceiling + depends on (two commands sharing one ceiling can never jointly overspend + it, [#159]) is unaffected, since the widening goes through the same locked + admission path an estimate drifting past its booking already used. + - **`get_dialect` now raises on an unrecognized connector instead of silently parsing every subsequent statement as DuckDB** ([#319]). A hyphenated BigQuery project id, the shape BigQuery itself hands out and diff --git a/packages/dex-core/src/exmergo_dex_core/adapters/bigquery.py b/packages/dex-core/src/exmergo_dex_core/adapters/bigquery.py index 14414a4..172913d 100644 --- a/packages/dex-core/src/exmergo_dex_core/adapters/bigquery.py +++ b/packages/dex-core/src/exmergo_dex_core/adapters/bigquery.py @@ -16,6 +16,7 @@ from __future__ import annotations +import re from collections.abc import Callable from dataclasses import dataclass from typing import Any @@ -66,6 +67,16 @@ # the math instead of letting the server fail the job after the fact. _MIN_BILLED_BYTES = 10 * 1024 * 1024 +# BigQuery's own refusal names the exact byte count it needed (issue #320): +# "Query exceeded limit for bytes billed: . or higher +# required." No dry run predicts this number, since it reflects the server's +# own execution-time rounding of bytes billed (observed per table scanned, +# not per query), which is exactly the gap between a dry-run-based estimate +# and what the job actually needed. Parsed rather than guessed at with a +# margin, so the retry below asks for precisely what BigQuery says it needs, +# whatever the underlying rounding rule turns out to be. +_BYTES_BILLED_REQUIRED_RE = re.compile(r"(\d+) or higher required") + # Field types whose values are nested or non-scalar: no approx-distinct, no # min/max, and non-null counting via COUNTIF (COUNT DISTINCT is invalid on # them and plain COUNT is not supported for every one of these types). @@ -705,7 +716,7 @@ def exact_distinct_counts( "not cover the extra scan; uniqueness verdicts stay approximate", ) return {} - _job, iterator = self._run(sql) + _job, iterator = self._run(sql, floored) rows = list(iterator) return {name: int(rows[0][f"d_{i}"]) for i, name in enumerate(columns)} @@ -737,7 +748,7 @@ def distinct_combination_counts( "cover the extra scan; grain stays unknown", ) return {} - _job, iterator = self._run(sql) + _job, iterator = self._run(sql, floored) rows = list(iterator) return { tuple(combo): int(rows[0][f"d_{i}"]) for i, combo in enumerate(combinations) @@ -779,7 +790,7 @@ def value_domain_counts( "cover the extra scan; no value domain reported", ) return {} - _job, iterator = self._run(sql) + _job, iterator = self._run(sql, floored) rows = list(iterator) return { name: ValueDomainSample( @@ -1054,19 +1065,30 @@ def _execute( """SELECT-only guard, free dry-run, gate charge, then the capped run.""" assert_select_only(sql, dialect=self.dialect) - self.cost_gate.charge(self._dry_run(sql)) - return self._run(sql, timeout_seconds=timeout_seconds, max_results=max_results) + estimate = self._dry_run(sql) + self.cost_gate.charge(estimate) + return self._run( + sql, estimate, timeout_seconds=timeout_seconds, max_results=max_results + ) def _run( self, sql: str, + dry_run_estimate: float, *, timeout_seconds: float | None = None, max_results: int | None = None, + _retried: bool = False, ) -> tuple[Any, Any]: """The single billed door past the gate: run with the server-side byte cap, wait for completion (bounded when a timeout is given), account the - actual billed bytes, and return (job, row iterator).""" + actual billed bytes, and return (job, row iterator). + + ``dry_run_estimate`` is what :meth:`_execute` already charged for this + exact statement; a bytes-billed refusal below widens the charge by the + gap between that estimate and what BigQuery says it actually needed, + rather than charging the full requirement a second time on top of it. + """ cap = self.cost_gate.remaining_for_statement() if cap is not None and cap < _MIN_BILLED_BYTES: @@ -1087,6 +1109,27 @@ def _run( iterator = job.result(timeout=timeout_seconds, max_results=max_results) except self._api_exceptions.BadRequest as exc: if "bytes billed" in str(exc) or "bytesBilledLimitExceeded" in str(exc): + required = _parse_bytes_billed_required(str(exc)) + if required is not None and not _retried: + # No dry run predicts BigQuery's own execution-time + # rounding of bytes billed (issue #320), so the estimate + # this command already charged for the statement can be + # a genuine underestimate even though nothing was wrong + # with it at dry-run time. BigQuery's own refusal names + # exactly what it needed; widen the charge by the gap and + # retry once with that as the new cap. Raises the same + # OverCeilingError/ConfirmationRequiredError this would + # raise anyway if the confirmed ceiling itself can't + # cover the real requirement, so a genuine over-budget + # query still refuses, correctly, on the real number. + self.cost_gate.charge(required - dry_run_estimate) + return self._run( + sql, + required, + timeout_seconds=timeout_seconds, + max_results=max_results, + _retried=True, + ) raise OverCeilingError( "the query would bill more than the remaining budget " "(server-side maximum_bytes_billed); raise --budget or " @@ -1155,3 +1198,12 @@ def _quote_ident(name: str) -> str: escaped = name.replace("`", "\\`") return f"`{escaped}`" + + +def _parse_bytes_billed_required(message: str) -> float | None: + """The byte count BigQuery's own bytes-billed refusal names as required, + or ``None`` if the message doesn't have the expected shape (a future + wording change should degrade to the old flat refusal, not a crash).""" + + match = _BYTES_BILLED_REQUIRED_RE.search(message) + return float(match.group(1)) if match else None diff --git a/packages/dex-core/tests/adapters/test_connect_bigquery.py b/packages/dex-core/tests/adapters/test_connect_bigquery.py index 21aee93..e0187cf 100644 --- a/packages/dex-core/tests/adapters/test_connect_bigquery.py +++ b/packages/dex-core/tests/adapters/test_connect_bigquery.py @@ -335,6 +335,81 @@ def test_server_side_cap_translates_when_the_estimate_drifts(fake_bq_client): assert "budget" in str(exc_info.value) +# --- issue #320: the server-side cap must not be pinned below the confirmed +# budget by BigQuery's own execution-time billing rounding, which no dry run +# predicts. The bug's actual trigger is a `session_ceiling`: without one, +# `remaining_for_statement()` already equals the confirmed ceiling directly +# (see `test_every_executed_job_carries_maximum_bytes_billed` above); with +# one, the per-statement cap is additionally bounded by this command's own +# reservation (booked at confirm time, sized to the dry-run estimate), which +# a dry run cannot widen for a gap execution alone reveals. + + +def test_billing_rounding_retries_once_within_the_confirmed_budget(fake_bq_client): + # Dry run is accurate to within 1%: nothing was wrong with the estimate at + # dry-run time, but the real bill still exceeds the tight reservation a + # session_ceiling books it against, the way BigQuery's own execution-time + # rounding can regardless of how accurate the dry run was. + fake_bq_client.dry_run_underestimate = 0.99 + fake_bq_client.tables["test-proj.shop.customers"].num_bytes = 160 * MB + adapter = make_adapter( + fake_bq_client, + ceiling=1_000 * MB, # roughly 6x the estimate, like the reported budget + session_ceiling=10_000 * MB, + ) + sql = "SELECT COUNT(*) FROM `test-proj`.`shop`.`customers`" + adapter.cost_gate.preflight_command( + adapter.query_estimate(sql) + ) # confirm handshake + + adapter.run_query(sql, max_rows=10, timeout_seconds=30) + + non_dry = [c for c in fake_bq_client.query_calls if not c.dry_run] + assert len(non_dry) == 2 # the too-tight attempt, then the retry + first_cap = non_dry[0].job_config.maximum_bytes_billed + second_cap = non_dry[1].job_config.maximum_bytes_billed + assert first_cap < 160 * MB # pinned to the reservation, not the ceiling + assert second_cap == 160 * MB # widened to exactly what BigQuery required + + +def test_billing_rounding_retry_still_refuses_past_the_confirmed_ceiling( + fake_bq_client, +): + # The real bill (200 MB) exceeds not just the reservation but the + # confirmed per-command ceiling itself (120 MB): the retry's own widening + # must refuse here, on the real number, rather than silently granting more + # than the operator confirmed. + fake_bq_client.dry_run_underestimate = 0.5 + fake_bq_client.tables["test-proj.shop.customers"].num_bytes = 200 * MB + adapter = make_adapter(fake_bq_client, ceiling=120 * MB, session_ceiling=1_000 * MB) + sql = "SELECT COUNT(*) FROM `test-proj`.`shop`.`customers`" + adapter.cost_gate.preflight_command(adapter.query_estimate(sql)) + + with pytest.raises(OverCeilingError) as exc_info: + adapter.run_query(sql, max_rows=10, timeout_seconds=30) + assert "budget" in str(exc_info.value) + # No second server-side attempt: the widening itself refused before a + # retry could even be issued. + non_dry = [c for c in fake_bq_client.query_calls if not c.dry_run] + assert len(non_dry) == 1 + + +def test_parse_bytes_billed_required_reads_bigquerys_own_number(): + from exmergo_dex_core.adapters.bigquery import _parse_bytes_billed_required + + message = ( + "Query exceeded limit for bytes billed: 163595928. " + "164626432 or higher required." + ) + assert _parse_bytes_billed_required(message) == 164626432.0 + + +def test_parse_bytes_billed_required_degrades_to_none_on_an_unexpected_message(): + from exmergo_dex_core.adapters.bigquery import _parse_bytes_billed_required + + assert _parse_bytes_billed_required("bytesBilledLimitExceeded") is None + + def test_timeout_cancels_the_job(fake_bq_client): fake_bq_client.result_error = TimeoutError("deadline") adapter = make_adapter(fake_bq_client)