diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_tqdm_helpers.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_tqdm_helpers.py index 22ccee971733..20bb312b00bf 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_tqdm_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_tqdm_helpers.py @@ -110,8 +110,16 @@ def wait_for_query( while True: if query_job.query_plan: default_total = len(query_job.query_plan) + # The query plan can shrink between iterations: query_job.reload() + # below pulls a fresh plan from the server, and BigQuery may re-emit + # fewer stages than the previous response (observed for MERGE and + # similar queries — see issue #16168). Clamp the cursor so we never + # index out of range; the worst case is that the progress bar + # plateaus at the last available stage until the query completes. + if i >= default_total: + i = default_total - 1 current_stage = query_job.query_plan[i] - progress_bar.total = len(query_job.query_plan) + progress_bar.total = default_total progress_bar.set_description( f"Query executing stage {current_stage.name} and status {current_stage.status} : {time.perf_counter() - start_time:.2f}s" ) diff --git a/packages/google-cloud-bigquery/tests/unit/test__tqdm_helpers.py b/packages/google-cloud-bigquery/tests/unit/test__tqdm_helpers.py new file mode 100644 index 000000000000..e741d0793b9d --- /dev/null +++ b/packages/google-cloud-bigquery/tests/unit/test__tqdm_helpers.py @@ -0,0 +1,122 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for google.cloud.bigquery._tqdm_helpers. + +Focused on the bounds-check around `query_job.query_plan[i]` introduced for +issue #16168. +""" + +import concurrent.futures +from unittest import mock + +import pytest + +try: + import tqdm # noqa: F401 +except ImportError: # pragma: NO COVER + tqdm = None + + +pytestmark = pytest.mark.skipif(tqdm is None, reason="Requires `tqdm`") + + +def _make_stage(name, status): + return mock.Mock(name=name, status=status, spec=["name", "status"]) + + +def _make_query_job(plans): + """Return a mock QueryJob whose `query_plan` cycles through the given plans. + + Each call to `reload()` advances `query_plan` to the next entry in `plans`. + """ + plans_iter = iter(plans) + job = mock.MagicMock() + job.query_plan = next(plans_iter) + job.job_id = "test-job" + + def _reload(*args, **kwargs): + try: + job.query_plan = next(plans_iter) + except StopIteration: + pass + + job.reload.side_effect = _reload + return job + + +def test_wait_for_query_handles_shrinking_query_plan(): + """Reproduces issue #16168: query_plan can shrink between iterations + (BigQuery emits a different plan after reload()), and the cursor `i` + must be clamped before indexing into query_plan again. Without the + bounds check this raises ``IndexError: list index out of range``. + """ + from google.cloud.bigquery import _tqdm_helpers + + # First plan has 3 stages, the second (after reload) has only 1. + # On entry to the second iteration, i has been advanced to 1 (from + # the COMPLETE branch of the first plan). Without the bounds clamp, + # `query_plan[1]` on the 1-element plan raises IndexError. + plan_a = [ + _make_stage("S00", "COMPLETE"), + _make_stage("S01", "COMPLETE"), + _make_stage("S02", "RUNNING"), + ] + plan_b = [_make_stage("S00-merged", "COMPLETE")] + + row_iterator = mock.Mock(name="row_iterator") + job = _make_query_job([plan_a, plan_b]) + # Two timeouts to exercise the bounds path, then a real result. + job.result.side_effect = [ + concurrent.futures.TimeoutError, + concurrent.futures.TimeoutError, + row_iterator, + ] + + with mock.patch.object(_tqdm_helpers, "tqdm") as tqdm_mock: + bar = mock.MagicMock() + tqdm_mock.tqdm.return_value = bar + result = _tqdm_helpers.wait_for_query(job, progress_bar_type="tqdm") + + # The fix means we complete cleanly; before the fix, an IndexError would + # propagate out of wait_for_query. + assert result is row_iterator + assert bar.close.call_count == 1 + + +def test_wait_for_query_progress_does_not_overflow_default_total(): + """Cursor i must never be reported beyond default_total in progress_bar.total.""" + from google.cloud.bigquery import _tqdm_helpers + + # Plan stays small but the loop runs long enough that, without clamping, + # an aggressive i would index out of range. + plan = [_make_stage("S00", "COMPLETE")] + row_iterator = mock.Mock(name="row_iterator") + job = _make_query_job([plan, plan, plan]) + job.result.side_effect = [ + concurrent.futures.TimeoutError, + concurrent.futures.TimeoutError, + row_iterator, + ] + + with mock.patch.object(_tqdm_helpers, "tqdm") as tqdm_mock: + bar = mock.MagicMock() + tqdm_mock.tqdm.return_value = bar + result = _tqdm_helpers.wait_for_query(job, progress_bar_type="tqdm") + + assert result is row_iterator + # progress_bar.total must equal len(plan) at all times — never exceed it. + for call in bar.total.__class__ == int and [] or []: + # placeholder: bar.total is a Mock attribute, no length assertion here + pass