Skip to content

Commit 1760e94

Browse files
authored
fix: relax timeout expectations (#1645)
* fix: relax timeout expectations Changes to python-api-core can in certain cases cause timeout to be represented as a literal python base object type. This CL adjusts logic that selects from multiple timeout values to better handle this case, which previously assumed either a None or scalar value being present. Fixes: #1612 * augment testing * blacken and lint fixes * unused import
1 parent 3e021a4 commit 1760e94

File tree

2 files changed

+102
-2
lines changed

2 files changed

+102
-2
lines changed

google/cloud/bigquery/client.py

+8-2
Original file line numberDiff line numberDiff line change
@@ -1895,7 +1895,10 @@ def _get_query_results(
18951895
extra_params: Dict[str, Any] = {"maxResults": 0}
18961896

18971897
if timeout is not None:
1898-
timeout = max(timeout, _MIN_GET_QUERY_RESULTS_TIMEOUT)
1898+
if type(timeout) == object:
1899+
timeout = _MIN_GET_QUERY_RESULTS_TIMEOUT
1900+
else:
1901+
timeout = max(timeout, _MIN_GET_QUERY_RESULTS_TIMEOUT)
18991902

19001903
if project is None:
19011904
project = self.project
@@ -3924,7 +3927,10 @@ def _list_rows_from_query_results(
39243927
}
39253928

39263929
if timeout is not None:
3927-
timeout = max(timeout, _MIN_GET_QUERY_RESULTS_TIMEOUT)
3930+
if type(timeout) == object:
3931+
timeout = _MIN_GET_QUERY_RESULTS_TIMEOUT
3932+
else:
3933+
timeout = max(timeout, _MIN_GET_QUERY_RESULTS_TIMEOUT)
39283934

39293935
if start_index is not None:
39303936
params["startIndex"] = start_index

tests/unit/test_client.py

+94
Original file line numberDiff line numberDiff line change
@@ -395,6 +395,31 @@ def test__get_query_results_miss_w_short_timeout(self):
395395
timeout=google.cloud.bigquery.client._MIN_GET_QUERY_RESULTS_TIMEOUT,
396396
)
397397

398+
def test__get_query_results_miss_w_default_timeout(self):
399+
import google.cloud.bigquery.client
400+
from google.cloud.exceptions import NotFound
401+
402+
creds = _make_credentials()
403+
client = self._make_one(self.PROJECT, creds)
404+
conn = client._connection = make_connection()
405+
path = "/projects/other-project/queries/nothere"
406+
with self.assertRaises(NotFound):
407+
client._get_query_results(
408+
"nothere",
409+
None,
410+
project="other-project",
411+
location=self.LOCATION,
412+
timeout_ms=500,
413+
timeout=object(), # the api core default timeout
414+
)
415+
416+
conn.api_request.assert_called_once_with(
417+
method="GET",
418+
path=path,
419+
query_params={"maxResults": 0, "timeoutMs": 500, "location": self.LOCATION},
420+
timeout=google.cloud.bigquery.client._MIN_GET_QUERY_RESULTS_TIMEOUT,
421+
)
422+
398423
def test__get_query_results_miss_w_client_location(self):
399424
from google.cloud.exceptions import NotFound
400425

@@ -438,6 +463,75 @@ def test__get_query_results_hit(self):
438463
self.assertEqual(query_results.total_rows, 10)
439464
self.assertTrue(query_results.complete)
440465

466+
def test__list_rows_from_query_results_w_none_timeout(self):
467+
from google.cloud.exceptions import NotFound
468+
from google.cloud.bigquery.schema import SchemaField
469+
470+
creds = _make_credentials()
471+
client = self._make_one(self.PROJECT, creds)
472+
conn = client._connection = make_connection()
473+
path = "/projects/project/queries/nothere"
474+
iterator = client._list_rows_from_query_results(
475+
"nothere",
476+
location=None,
477+
project="project",
478+
schema=[
479+
SchemaField("f1", "STRING", mode="REQUIRED"),
480+
SchemaField("f2", "INTEGER", mode="REQUIRED"),
481+
],
482+
timeout=None,
483+
)
484+
485+
# trigger the iterator to request data
486+
with self.assertRaises(NotFound):
487+
iterator._get_next_page_response()
488+
489+
conn.api_request.assert_called_once_with(
490+
method="GET",
491+
path=path,
492+
query_params={
493+
"fields": "jobReference,totalRows,pageToken,rows",
494+
"location": None,
495+
"formatOptions.useInt64Timestamp": True,
496+
},
497+
timeout=None,
498+
)
499+
500+
def test__list_rows_from_query_results_w_default_timeout(self):
501+
import google.cloud.bigquery.client
502+
from google.cloud.exceptions import NotFound
503+
from google.cloud.bigquery.schema import SchemaField
504+
505+
creds = _make_credentials()
506+
client = self._make_one(self.PROJECT, creds)
507+
conn = client._connection = make_connection()
508+
path = "/projects/project/queries/nothere"
509+
iterator = client._list_rows_from_query_results(
510+
"nothere",
511+
location=None,
512+
project="project",
513+
schema=[
514+
SchemaField("f1", "STRING", mode="REQUIRED"),
515+
SchemaField("f2", "INTEGER", mode="REQUIRED"),
516+
],
517+
timeout=object(),
518+
)
519+
520+
# trigger the iterator to request data
521+
with self.assertRaises(NotFound):
522+
iterator._get_next_page_response()
523+
524+
conn.api_request.assert_called_once_with(
525+
method="GET",
526+
path=path,
527+
query_params={
528+
"fields": "jobReference,totalRows,pageToken,rows",
529+
"location": None,
530+
"formatOptions.useInt64Timestamp": True,
531+
},
532+
timeout=google.cloud.bigquery.client._MIN_GET_QUERY_RESULTS_TIMEOUT,
533+
)
534+
441535
def test_default_query_job_config(self):
442536
from google.cloud.bigquery import QueryJobConfig
443537

0 commit comments

Comments
 (0)