From 79c44bcad08015d3cfc869d94f235158b3d8fe42 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Fri, 19 Sep 2025 14:48:47 +0000 Subject: [PATCH 1/8] fix(job): honor custom retry in job.result() The `_AsyncJob.result()` method was not correctly passing the `retry` argument to the superclass's `result()` method when the `retry` object was the same as the default retry object. This caused the default retry settings to be ignored in some cases. This change modifies the `result()` method to always pass the `retry` argument to the superclass, ensuring that the provided retry settings are always honored. A new test case is added to verify that `job.result()` correctly handles both the default retry and a custom retry object. --- google/cloud/bigquery/job/base.py | 2 +- tests/unit/test_job_retry.py | 42 +++++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/google/cloud/bigquery/job/base.py b/google/cloud/bigquery/job/base.py index 9b7ddb82d..d11c6c52f 100644 --- a/google/cloud/bigquery/job/base.py +++ b/google/cloud/bigquery/job/base.py @@ -1044,7 +1044,7 @@ def result( # type: ignore # (incompatible with supertype) if self.state is None: self._begin(retry=retry, timeout=timeout) - kwargs = {} if retry is DEFAULT_RETRY else {"retry": retry} + kwargs = {"retry": retry} return super(_AsyncJob, self).result(timeout=timeout, **kwargs) def cancelled(self): diff --git a/tests/unit/test_job_retry.py b/tests/unit/test_job_retry.py index 7343fed3d..cedd722de 100644 --- a/tests/unit/test_job_retry.py +++ b/tests/unit/test_job_retry.py @@ -615,3 +615,45 @@ def test_query_and_wait_retries_job_for_DDL_queries(global_time_lock): _, kwargs = calls[3] assert kwargs["method"] == "POST" assert kwargs["path"] == query_request_path + + +@pytest.mark.parametrize( + "result_retry", + [ + pytest.param( + {}, + id="default retry", + ), + pytest.param( + {"retry": google.cloud.bigquery.retry.DEFAULT_RETRY.with_timeout(timeout=10.0)}, + id="custom retry object", + ), + ], +) +def test_retry_load_job_result(result_retry, PROJECT, DS_ID): + from google.cloud.bigquery.dataset import DatasetReference + from google.cloud.bigquery.job.load import LoadJob + + client = make_client() + conn = client._connection = make_connection( + dict( + status=dict(state="RUNNING"), + jobReference={"jobId": "id_1"}, + ), + google.api_core.exceptions.ServiceUnavailable("retry me"), + dict( + status=dict(state="DONE"), + jobReference={"jobId": "id_1"}, + statistics={"load": {"outputRows": 1}}, + ), + ) + + table_ref = DatasetReference(project=PROJECT, dataset_id=DS_ID).table("new_table") + job = LoadJob("id_1", source_uris=None, destination=table_ref, client=client) + result = job.result(**result_retry) + + assert job.state == "DONE" + assert result.output_rows == 1 + + # We made all the calls we expected to. + assert conn.api_request.call_count == 3 From d9cc9f8db8410fea494c3b8f3cab2762c39deb7b Mon Sep 17 00:00:00 2001 From: Chalmer Lowe Date: Fri, 19 Sep 2025 10:53:30 -0400 Subject: [PATCH 2/8] Update tests/unit/test_job_retry.py --- tests/unit/test_job_retry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_job_retry.py b/tests/unit/test_job_retry.py index cedd722de..7929e676b 100644 --- a/tests/unit/test_job_retry.py +++ b/tests/unit/test_job_retry.py @@ -622,7 +622,7 @@ def test_query_and_wait_retries_job_for_DDL_queries(global_time_lock): [ pytest.param( {}, - id="default retry", + id="default retry use case", ), pytest.param( {"retry": google.cloud.bigquery.retry.DEFAULT_RETRY.with_timeout(timeout=10.0)}, From 5476f0238cb3389c14d75bdb57cc7429c48584be Mon Sep 17 00:00:00 2001 From: Chalmer Lowe Date: Fri, 19 Sep 2025 10:53:51 -0400 Subject: [PATCH 3/8] Update tests/unit/test_job_retry.py --- tests/unit/test_job_retry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_job_retry.py b/tests/unit/test_job_retry.py index 7929e676b..f4f229fd1 100644 --- a/tests/unit/test_job_retry.py +++ b/tests/unit/test_job_retry.py @@ -626,7 +626,7 @@ def test_query_and_wait_retries_job_for_DDL_queries(global_time_lock): ), pytest.param( {"retry": google.cloud.bigquery.retry.DEFAULT_RETRY.with_timeout(timeout=10.0)}, - id="custom retry object", + id="custom retry object use case", ), ], ) From 71e31b3669181e904b0b73b0bacc4fd2d6b9834b Mon Sep 17 00:00:00 2001 From: chalmer lowe Date: Fri, 19 Sep 2025 14:12:26 -0400 Subject: [PATCH 4/8] blacken and lint --- tests/unit/test_job_retry.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_job_retry.py b/tests/unit/test_job_retry.py index f4f229fd1..590fb5429 100644 --- a/tests/unit/test_job_retry.py +++ b/tests/unit/test_job_retry.py @@ -625,7 +625,11 @@ def test_query_and_wait_retries_job_for_DDL_queries(global_time_lock): id="default retry use case", ), pytest.param( - {"retry": google.cloud.bigquery.retry.DEFAULT_RETRY.with_timeout(timeout=10.0)}, + { + "retry": google.cloud.bigquery.retry.DEFAULT_RETRY.with_timeout( + timeout=10.0 + ) + }, id="custom retry object use case", ), ], From e0b9ec715f6d2a54d962c888e6800897df6dc84a Mon Sep 17 00:00:00 2001 From: chalmer lowe Date: Tue, 30 Sep 2025 09:52:55 -0400 Subject: [PATCH 5/8] udpates retry handling and testing of retry handling --- google/cloud/bigquery/job/base.py | 3 +- tests/unit/test_job_retry.py | 109 ++++++++++++++++++++++++++++-- 2 files changed, 104 insertions(+), 8 deletions(-) diff --git a/google/cloud/bigquery/job/base.py b/google/cloud/bigquery/job/base.py index d11c6c52f..7576fc9aa 100644 --- a/google/cloud/bigquery/job/base.py +++ b/google/cloud/bigquery/job/base.py @@ -1044,8 +1044,7 @@ def result( # type: ignore # (incompatible with supertype) if self.state is None: self._begin(retry=retry, timeout=timeout) - kwargs = {"retry": retry} - return super(_AsyncJob, self).result(timeout=timeout, **kwargs) + return super(_AsyncJob, self).result(timeout=timeout, retry=retry) def cancelled(self): """Check if the job has been cancelled. diff --git a/tests/unit/test_job_retry.py b/tests/unit/test_job_retry.py index 590fb5429..ade16c7d4 100644 --- a/tests/unit/test_job_retry.py +++ b/tests/unit/test_job_retry.py @@ -618,11 +618,11 @@ def test_query_and_wait_retries_job_for_DDL_queries(global_time_lock): @pytest.mark.parametrize( - "result_retry", + "result_retry_param", [ pytest.param( {}, - id="default retry use case", + id="default retry {}", ), pytest.param( { @@ -630,13 +630,14 @@ def test_query_and_wait_retries_job_for_DDL_queries(global_time_lock): timeout=10.0 ) }, - id="custom retry object use case", + id="custom retry object with timeout 10.0", ), ], ) -def test_retry_load_job_result(result_retry, PROJECT, DS_ID): +def test_retry_load_job_result(result_retry_param, PROJECT, DS_ID): from google.cloud.bigquery.dataset import DatasetReference from google.cloud.bigquery.job.load import LoadJob + import google.cloud.bigquery.retry client = make_client() conn = client._connection = make_connection( @@ -654,10 +655,106 @@ def test_retry_load_job_result(result_retry, PROJECT, DS_ID): table_ref = DatasetReference(project=PROJECT, dataset_id=DS_ID).table("new_table") job = LoadJob("id_1", source_uris=None, destination=table_ref, client=client) - result = job.result(**result_retry) + + with mock.patch.object( + client, "_call_api", wraps=client._call_api + ) as wrapped_call_api: + result = job.result(**result_retry_param) assert job.state == "DONE" assert result.output_rows == 1 - # We made all the calls we expected to. + # Check that _call_api was called multiple times due to retry + assert wrapped_call_api.call_count > 1 + + # Verify the retry object used in the calls to _call_api + expected_retry = result_retry_param.get( + "retry", google.cloud.bigquery.retry.DEFAULT_RETRY + ) + + for call in wrapped_call_api.mock_calls: + name, args, kwargs = call + # The retry object is the first positional argument to _call_api + called_retry = args[0] + + # We only care about the calls made during the job.result() polling + if kwargs.get("method") == "GET" and "jobs/id_1" in kwargs.get("path", ""): + assert called_retry._predicate == expected_retry._predicate + assert called_retry._initial == expected_retry._initial + assert called_retry._maximum == expected_retry._maximum + assert called_retry._multiplier == expected_retry._multiplier + assert called_retry._deadline == expected_retry._deadline + if "retry" in result_retry_param: + # Specifically check the timeout for the custom retry case + assert called_retry._timeout == 10.0 + print(10.0) + else: + assert called_retry._timeout == expected_retry._timeout + print("not 10.0") + + # The number of api_request calls should still be 3 assert conn.api_request.call_count == 3 + + +# @pytest.mark.parametrize( +# "result_retry", +# [ +# pytest.param( +# {}, +# id="default retry use case", +# ), +# pytest.param( +# { +# "retry": google.cloud.bigquery.retry.DEFAULT_RETRY.with_timeout( +# timeout=10.0 +# ) +# }, +# id="custom retry object use case", +# ), +# ], +# ) +# def test_retry_load_job_result(result_retry, PROJECT, DS_ID): + +# from google.cloud.bigquery.dataset import DatasetReference +# from google.cloud.bigquery.job.load import LoadJob +# import google.cloud.bigquery.retry + +# client = make_client() +# conn = client._connection = make_connection( +# dict( +# status=dict(state="RUNNING"), +# jobReference={"jobId": "id_1"}, +# ), +# google.api_core.exceptions.ServiceUnavailable("retry me"), +# dict( +# status=dict(state="DONE"), +# jobReference={"jobId": "id_1"}, +# statistics={"load": {"outputRows": 1}}, +# ), +# ) + +# table_ref = DatasetReference(project=PROJECT, dataset_id=DS_ID).table("new_table") +# job = LoadJob("id_1", source_uris=None, destination=table_ref, client=client) +# result = job.result(**result_retry) + +# print(f"Test Case ID: {result_retry}") +# for call in conn.api_request.mock_calls: +# name, args, kwargs = call +# print(f" Call to api_request with kwargs: {kwargs}") + +# if result_retry: +# custom_retry_used = False +# for name, args, kwargs in conn.api_request.mock_calls: +# if 'retry' in kwargs: +# # Check if the retry object passed has the custom timeout +# if kwargs['retry']._timeout == 10.0: +# custom_retry_used = True +# assert kwargs['retry']._deadline == 10.0 # Deadline should also be updated +# # Optional: check other properties if needed +# assert custom_retry_used, "Custom retry object with timeout 10.0 was not used" + +# assert job.state == "DONE" +# assert result.output_rows == 1 + +# # We made all the calls we expected to. +# assert conn.api_request.call_count == 3 From 01483363a35978b0cdab753400f57c482f19c29d Mon Sep 17 00:00:00 2001 From: Chalmer Lowe Date: Tue, 30 Sep 2025 09:54:10 -0400 Subject: [PATCH 6/8] Update tests/unit/test_job_retry.py --- tests/unit/test_job_retry.py | 63 ------------------------------------ 1 file changed, 63 deletions(-) diff --git a/tests/unit/test_job_retry.py b/tests/unit/test_job_retry.py index ade16c7d4..8706742a4 100644 --- a/tests/unit/test_job_retry.py +++ b/tests/unit/test_job_retry.py @@ -695,66 +695,3 @@ def test_retry_load_job_result(result_retry_param, PROJECT, DS_ID): # The number of api_request calls should still be 3 assert conn.api_request.call_count == 3 - -# @pytest.mark.parametrize( -# "result_retry", -# [ -# pytest.param( -# {}, -# id="default retry use case", -# ), -# pytest.param( -# { -# "retry": google.cloud.bigquery.retry.DEFAULT_RETRY.with_timeout( -# timeout=10.0 -# ) -# }, -# id="custom retry object use case", -# ), -# ], -# ) -# def test_retry_load_job_result(result_retry, PROJECT, DS_ID): - -# from google.cloud.bigquery.dataset import DatasetReference -# from google.cloud.bigquery.job.load import LoadJob -# import google.cloud.bigquery.retry - -# client = make_client() -# conn = client._connection = make_connection( -# dict( -# status=dict(state="RUNNING"), -# jobReference={"jobId": "id_1"}, -# ), -# google.api_core.exceptions.ServiceUnavailable("retry me"), -# dict( -# status=dict(state="DONE"), -# jobReference={"jobId": "id_1"}, -# statistics={"load": {"outputRows": 1}}, -# ), -# ) - -# table_ref = DatasetReference(project=PROJECT, dataset_id=DS_ID).table("new_table") -# job = LoadJob("id_1", source_uris=None, destination=table_ref, client=client) -# result = job.result(**result_retry) - -# print(f"Test Case ID: {result_retry}") -# for call in conn.api_request.mock_calls: -# name, args, kwargs = call -# print(f" Call to api_request with kwargs: {kwargs}") - -# if result_retry: -# custom_retry_used = False -# for name, args, kwargs in conn.api_request.mock_calls: -# if 'retry' in kwargs: -# # Check if the retry object passed has the custom timeout -# if kwargs['retry']._timeout == 10.0: -# custom_retry_used = True -# assert kwargs['retry']._deadline == 10.0 # Deadline should also be updated -# # Optional: check other properties if needed -# assert custom_retry_used, "Custom retry object with timeout 10.0 was not used" - -# assert job.state == "DONE" -# assert result.output_rows == 1 - -# # We made all the calls we expected to. -# assert conn.api_request.call_count == 3 From 5b6012a0f2ba892f79596b8866b25cb066056305 Mon Sep 17 00:00:00 2001 From: Chalmer Lowe Date: Tue, 30 Sep 2025 09:54:55 -0400 Subject: [PATCH 7/8] Update tests/unit/test_job_retry.py --- tests/unit/test_job_retry.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/unit/test_job_retry.py b/tests/unit/test_job_retry.py index 8706742a4..0d7593a6b 100644 --- a/tests/unit/test_job_retry.py +++ b/tests/unit/test_job_retry.py @@ -687,10 +687,8 @@ def test_retry_load_job_result(result_retry_param, PROJECT, DS_ID): if "retry" in result_retry_param: # Specifically check the timeout for the custom retry case assert called_retry._timeout == 10.0 - print(10.0) else: assert called_retry._timeout == expected_retry._timeout - print("not 10.0") # The number of api_request calls should still be 3 assert conn.api_request.call_count == 3 From e8d5fb0e88b4e949b7c4dd2d1ecd1fd8149f2a0d Mon Sep 17 00:00:00 2001 From: Chalmer Lowe Date: Tue, 30 Sep 2025 10:31:36 -0400 Subject: [PATCH 8/8] Update tests/unit/test_job_retry.py --- tests/unit/test_job_retry.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unit/test_job_retry.py b/tests/unit/test_job_retry.py index 0d7593a6b..fa55e8f6a 100644 --- a/tests/unit/test_job_retry.py +++ b/tests/unit/test_job_retry.py @@ -692,4 +692,3 @@ def test_retry_load_job_result(result_retry_param, PROJECT, DS_ID): # The number of api_request calls should still be 3 assert conn.api_request.call_count == 3 -