diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index 11cceea42..2712b0c83 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -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 @@ -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 diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index cf0aa4028..faa073dce 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -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 @@ -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