From 55dc9a2b6980b2d0c0774c699f89c5f7c9e6da62 Mon Sep 17 00:00:00 2001 From: abhishekacharya200 Date: Tue, 20 Oct 2026 03:48:30 +1100 Subject: [PATCH] fix(bigquery): pass timeout to the lookup after an exists_ok conflict create_dataset, create_table and create_routine with exists_ok=True fall back to get_dataset, get_table or get_routine when the insert returns 409. The fallback passed retry but not timeout, and BigQuery's default timeout is None, so that GET could wait forever even when the caller set a timeout. Forward the caller's timeout to the lookup too. --- .../google/cloud/bigquery/client.py | 6 +- .../tests/unit/test_client.py | 58 +++++++++++++++++++ .../tests/unit/test_create_dataset.py | 25 ++++++++ 3 files changed, 86 insertions(+), 3 deletions(-) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py index 214982dc2a23..04f96545e120 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py @@ -775,7 +775,7 @@ def create_dataset( except core_exceptions.Conflict: if not exists_ok: raise - return self.get_dataset(dataset.reference, retry=retry) + return self.get_dataset(dataset.reference, retry=retry, timeout=timeout) def create_routine( self, @@ -830,7 +830,7 @@ def create_routine( except core_exceptions.Conflict: if not exists_ok: raise - return self.get_routine(routine.reference, retry=retry) + return self.get_routine(routine.reference, retry=retry, timeout=timeout) def create_table( self, @@ -891,7 +891,7 @@ def create_table( except core_exceptions.Conflict: if not exists_ok: raise - return self.get_table(table.reference, retry=retry) + return self.get_table(table.reference, retry=retry, timeout=timeout) def _call_api( self, diff --git a/packages/google-cloud-bigquery/tests/unit/test_client.py b/packages/google-cloud-bigquery/tests/unit/test_client.py index 88629500a007..b32b79248b51 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_client.py +++ b/packages/google-cloud-bigquery/tests/unit/test_client.py @@ -1189,6 +1189,38 @@ def test_create_routine_w_conflict_exists_ok(self): ] ) + def test_create_routine_w_conflict_exists_ok_w_timeout(self): + from google.cloud.bigquery.routine import Routine + + creds = _make_credentials() + client = self._make_one(project=self.PROJECT, credentials=creds) + resource = { + "routineReference": { + "projectId": "test-routine-project", + "datasetId": "test_routines", + "routineId": "minimal_routine", + } + } + path = "/projects/test-routine-project/datasets/test_routines/routines" + conn = client._connection = make_connection( + google.api_core.exceptions.AlreadyExists("routine already exists"), resource + ) + routine = Routine("test-routine-project.test_routines.minimal_routine") + + client.create_routine(routine, exists_ok=True, timeout=7.5) + + # The lookup after the conflict must keep the caller's timeout. + conn.api_request.assert_has_calls( + [ + mock.call(method="POST", path=path, data=resource, timeout=7.5), + mock.call( + method="GET", + path="/projects/test-routine-project/datasets/test_routines/routines/minimal_routine", + timeout=7.5, + ), + ] + ) + def test_create_table_w_day_partition(self): from google.cloud.bigquery.table import Table, TimePartitioning @@ -1642,6 +1674,32 @@ def test_create_table_alreadyexists_w_exists_ok_true(self): ] ) + def test_create_table_alreadyexists_w_exists_ok_true_w_timeout(self): + post_path = "/projects/{}/datasets/{}/tables".format(self.PROJECT, self.DS_ID) + get_path = "/projects/{}/datasets/{}/tables/{}".format( + self.PROJECT, self.DS_ID, self.TABLE_ID + ) + resource = self._make_table_resource() + creds = _make_credentials() + client = self._make_one( + project=self.PROJECT, credentials=creds, location=self.LOCATION + ) + conn = client._connection = make_connection( + google.api_core.exceptions.AlreadyExists("table already exists"), resource + ) + + client.create_table( + "{}.{}".format(self.DS_ID, self.TABLE_ID), exists_ok=True, timeout=7.5 + ) + + # The lookup after the conflict must keep the caller's timeout. + conn.api_request.assert_has_calls( + [ + mock.call(method="POST", path=post_path, data=mock.ANY, timeout=7.5), + mock.call(method="GET", path=get_path, timeout=7.5), + ] + ) + def test_close(self): creds = _make_credentials() http = mock.Mock() diff --git a/packages/google-cloud-bigquery/tests/unit/test_create_dataset.py b/packages/google-cloud-bigquery/tests/unit/test_create_dataset.py index b144471ca503..190ca28c2fe7 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_create_dataset.py +++ b/packages/google-cloud-bigquery/tests/unit/test_create_dataset.py @@ -382,6 +382,31 @@ def test_create_dataset_alreadyexists_w_exists_ok_true(PROJECT, DS_ID, LOCATION) ) +def test_create_dataset_alreadyexists_w_exists_ok_true_w_timeout( + PROJECT, DS_ID, LOCATION +): + post_path = "/projects/{}/datasets".format(PROJECT) + get_path = "/projects/{}/datasets/{}".format(PROJECT, DS_ID) + resource = { + "datasetReference": {"projectId": PROJECT, "datasetId": DS_ID}, + "location": LOCATION, + } + client = make_client(location=LOCATION) + conn = client._connection = make_connection( + google.api_core.exceptions.AlreadyExists("dataset already exists"), resource + ) + + client.create_dataset(DS_ID, exists_ok=True, timeout=7.5) + + # The lookup after the conflict must keep the caller's timeout. + conn.api_request.assert_has_calls( + [ + mock.call(method="POST", path=post_path, data=mock.ANY, timeout=7.5), + mock.call(method="GET", path=get_path, timeout=7.5, query_params={}), + ] + ) + + def test_create_dataset_with_default_rounding_mode_if_value_is_none( PROJECT, DS_ID, LOCATION ):