Skip to content

Commit

Permalink
fix: relax timeout expectations (#1645)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
shollyman authored Aug 21, 2023
1 parent 3e021a4 commit 1760e94
Show file tree
Hide file tree
Showing 2 changed files with 102 additions and 2 deletions.
10 changes: 8 additions & 2 deletions google/cloud/bigquery/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -1895,7 +1895,10 @@ def _get_query_results(
extra_params: Dict[str, Any] = {"maxResults": 0}

if timeout is not None:
timeout = max(timeout, _MIN_GET_QUERY_RESULTS_TIMEOUT)
if type(timeout) == object:
timeout = _MIN_GET_QUERY_RESULTS_TIMEOUT
else:
timeout = max(timeout, _MIN_GET_QUERY_RESULTS_TIMEOUT)

if project is None:
project = self.project
Expand Down Expand Up @@ -3924,7 +3927,10 @@ def _list_rows_from_query_results(
}

if timeout is not None:
timeout = max(timeout, _MIN_GET_QUERY_RESULTS_TIMEOUT)
if type(timeout) == object:
timeout = _MIN_GET_QUERY_RESULTS_TIMEOUT
else:
timeout = max(timeout, _MIN_GET_QUERY_RESULTS_TIMEOUT)

if start_index is not None:
params["startIndex"] = start_index
Expand Down
94 changes: 94 additions & 0 deletions tests/unit/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -395,6 +395,31 @@ def test__get_query_results_miss_w_short_timeout(self):
timeout=google.cloud.bigquery.client._MIN_GET_QUERY_RESULTS_TIMEOUT,
)

def test__get_query_results_miss_w_default_timeout(self):
import google.cloud.bigquery.client
from google.cloud.exceptions import NotFound

creds = _make_credentials()
client = self._make_one(self.PROJECT, creds)
conn = client._connection = make_connection()
path = "/projects/other-project/queries/nothere"
with self.assertRaises(NotFound):
client._get_query_results(
"nothere",
None,
project="other-project",
location=self.LOCATION,
timeout_ms=500,
timeout=object(), # the api core default timeout
)

conn.api_request.assert_called_once_with(
method="GET",
path=path,
query_params={"maxResults": 0, "timeoutMs": 500, "location": self.LOCATION},
timeout=google.cloud.bigquery.client._MIN_GET_QUERY_RESULTS_TIMEOUT,
)

def test__get_query_results_miss_w_client_location(self):
from google.cloud.exceptions import NotFound

Expand Down Expand Up @@ -438,6 +463,75 @@ def test__get_query_results_hit(self):
self.assertEqual(query_results.total_rows, 10)
self.assertTrue(query_results.complete)

def test__list_rows_from_query_results_w_none_timeout(self):
from google.cloud.exceptions import NotFound
from google.cloud.bigquery.schema import SchemaField

creds = _make_credentials()
client = self._make_one(self.PROJECT, creds)
conn = client._connection = make_connection()
path = "/projects/project/queries/nothere"
iterator = client._list_rows_from_query_results(
"nothere",
location=None,
project="project",
schema=[
SchemaField("f1", "STRING", mode="REQUIRED"),
SchemaField("f2", "INTEGER", mode="REQUIRED"),
],
timeout=None,
)

# trigger the iterator to request data
with self.assertRaises(NotFound):
iterator._get_next_page_response()

conn.api_request.assert_called_once_with(
method="GET",
path=path,
query_params={
"fields": "jobReference,totalRows,pageToken,rows",
"location": None,
"formatOptions.useInt64Timestamp": True,
},
timeout=None,
)

def test__list_rows_from_query_results_w_default_timeout(self):
import google.cloud.bigquery.client
from google.cloud.exceptions import NotFound
from google.cloud.bigquery.schema import SchemaField

creds = _make_credentials()
client = self._make_one(self.PROJECT, creds)
conn = client._connection = make_connection()
path = "/projects/project/queries/nothere"
iterator = client._list_rows_from_query_results(
"nothere",
location=None,
project="project",
schema=[
SchemaField("f1", "STRING", mode="REQUIRED"),
SchemaField("f2", "INTEGER", mode="REQUIRED"),
],
timeout=object(),
)

# trigger the iterator to request data
with self.assertRaises(NotFound):
iterator._get_next_page_response()

conn.api_request.assert_called_once_with(
method="GET",
path=path,
query_params={
"fields": "jobReference,totalRows,pageToken,rows",
"location": None,
"formatOptions.useInt64Timestamp": True,
},
timeout=google.cloud.bigquery.client._MIN_GET_QUERY_RESULTS_TIMEOUT,
)

def test_default_query_job_config(self):
from google.cloud.bigquery import QueryJobConfig

Expand Down

0 comments on commit 1760e94

Please sign in to comment.