Split timeout in multi-request methods

If a method makes multiple requests and is given a timeout, that
timeout should represent the total allowed time for all requests
combined.

Author: plamut

bigquery/google/cloud/bigquery/client.py

@@ -22,6 +22,7 @@
 except ImportError:  # Python 2.7
     import collections as collections_abc
 
+import concurrent.futures
 import copy
 import functools
 import gzip
@@ -47,6 +48,7 @@
 import google.api_core.client_options
 import google.api_core.exceptions
 from google.api_core import page_iterator
+from google.auth.transport.requests import TimeoutGuard
 import google.cloud._helpers
 from google.cloud import exceptions
 from google.cloud.client import ClientWithProject
@@ -2557,21 +2559,27 @@ def list_partitions(self, table, retry=DEFAULT_RETRY, timeout=None):
             timeout (Optional[float]):
                 The number of seconds to wait for the underlying HTTP transport
                 before using ``retry``.
+                If multiple requests are made under the hood, ``timeout`` is
+                interpreted as the approximate total time of **all** requests.
 
         Returns:
             List[str]:
                 A list of the partition ids present in the partitioned table
         """
-        # TODO: split timeout between all API calls in the method
         table = _table_arg_to_table_ref(table, default_project=self.project)
-        meta_table = self.get_table(
-            TableReference(
-                self.dataset(table.dataset_id, project=table.project),
-                "%s$__PARTITIONS_SUMMARY__" % table.table_id,
-            ),
-            retry=retry,
-            timeout=timeout,
-        )
+
+        with TimeoutGuard(
+            timeout, timeout_error_type=concurrent.futures.TimeoutError
+        ) as guard:
+            meta_table = self.get_table(
+                TableReference(
+                    self.dataset(table.dataset_id, project=table.project),
+                    "%s$__PARTITIONS_SUMMARY__" % table.table_id,
+                ),
+                retry=retry,
+                timeout=timeout,
+            )
+        timeout = guard.remaining_timeout
 
         subset = [col for col in meta_table.schema if col.name == "partition_id"]
         return [
@@ -2638,6 +2646,8 @@ def list_rows(
             timeout (Optional[float]):
                 The number of seconds to wait for the underlying HTTP transport
                 before using ``retry``.
+                If multiple requests are made under the hood, ``timeout`` is
+                interpreted as the approximate total time of **all** requests.
 
         Returns:
             google.cloud.bigquery.table.RowIterator:
@@ -2648,7 +2658,6 @@ def list_rows(
                 (this is distinct from the total number of rows in the
                 current page: ``iterator.page.num_items``).
         """
-        # TODO: split timeout between all internal API calls
         table = _table_arg_to_table(table, default_project=self.project)
 
         if not isinstance(table, Table):
@@ -2663,7 +2672,11 @@ def list_rows(
         # No schema, but no selected_fields. Assume the developer wants all
         # columns, so get the table resource for them rather than failing.
         elif len(schema) == 0:
-            table = self.get_table(table.reference, retry=retry, timeout=timeout)
+            with TimeoutGuard(
+                timeout, timeout_error_type=concurrent.futures.TimeoutError
+            ) as guard:
+                table = self.get_table(table.reference, retry=retry, timeout=timeout)
+            timeout = guard.remaining_timeout
             schema = table.schema
 
         params = {}

bigquery/google/cloud/bigquery/job.py

@@ -26,6 +26,7 @@
 from six.moves import http_client
 
 import google.api_core.future.polling
+from google.auth.transport.requests import TimeoutGuard
 from google.cloud import exceptions
 from google.cloud.exceptions import NotFound
 from google.cloud.bigquery.dataset import Dataset
@@ -793,6 +794,8 @@ def result(self, retry=DEFAULT_RETRY, timeout=None):
             timeout (Optional[float]):
                 The number of seconds to wait for the underlying HTTP transport
                 before using ``retry``.
+                If multiple requests are made under the hood, ``timeout`` is
+                interpreted as the approximate total time of **all** requests.
 
         Returns:
             _AsyncJob: This instance.
@@ -803,10 +806,12 @@ def result(self, retry=DEFAULT_RETRY, timeout=None):
             concurrent.futures.TimeoutError:
                 if the job did not complete in the given timeout.
         """
-        # TODO: combine _begin timeout with super().result() timeout!
-        # borrow timeout guard from google auth lib
         if self.state is None:
-            self._begin(retry=retry, timeout=timeout)
+            with TimeoutGuard(
+                timeout, timeout_error_type=concurrent.futures.TimeoutError
+            ) as guard:
+                self._begin(retry=retry, timeout=timeout)
+            timeout = guard.remaining_timeout
         # TODO: modify PollingFuture so it can pass a retry argument to done().
         return super(_AsyncJob, self).result(timeout=timeout)
 
@@ -3163,6 +3168,8 @@ def result(
             timeout (Optional[float]):
                 The number of seconds to wait for the underlying HTTP transport
                 before using ``retry``.
+                If multiple requests are made under the hood, ``timeout`` is
+                interpreted as the approximate total time of **all** requests.
 
         Returns:
             google.cloud.bigquery.table.RowIterator:
@@ -3180,16 +3187,27 @@ def result(
                 If the job did not complete in the given timeout.
         """
         try:
-            # TODO: combine timeout with timeouts passed to super().result()
-            # and _get_query_results (total timeout shared by both)
-            # borrow timeout guard from google auth lib
-            super(QueryJob, self).result(timeout=timeout)
+            guard = TimeoutGuard(
+                timeout, timeout_error_type=concurrent.futures.TimeoutError
+            )
+            with guard:
+                super(QueryJob, self).result(retry=retry, timeout=timeout)
+            timeout = guard.remaining_timeout
 
             # Return an iterator instead of returning the job.
             if not self._query_results:
-                self._query_results = self._client._get_query_results(
-                    self.job_id, retry, project=self.project, location=self.location
+                guard = TimeoutGuard(
+                    timeout, timeout_error_type=concurrent.futures.TimeoutError
                 )
+                with guard:
+                    self._query_results = self._client._get_query_results(
+                        self.job_id,
+                        retry,
+                        project=self.project,
+                        location=self.location,
+                        timeout=timeout,
+                    )
+                timeout = guard.remaining_timeout
         except exceptions.GoogleCloudError as exc:
             exc.message += self._format_for_exception(self.query, self.job_id)
             exc.query_job = self
@@ -3209,7 +3227,11 @@ def result(
         dest_table = Table(dest_table_ref, schema=schema)
         dest_table._properties["numRows"] = self._query_results.total_rows
         rows = self._client.list_rows(
-            dest_table, page_size=page_size, retry=retry, max_results=max_results
+            dest_table,
+            page_size=page_size,
+            max_results=max_results,
+            retry=retry,
+            timeout=timeout,
         )
         rows._preserve_order = _contains_order_by(self.query)
         return rows

bigquery/noxfile.py

@@ -34,7 +34,7 @@ def default(session):
     run the tests.
     """
     # Install all test dependencies, then install local packages in-place.
-    session.install("mock", "pytest", "pytest-cov")
+    session.install("mock", "pytest", "pytest-cov", "freezegun")
     for local_dep in LOCAL_DEPS:
         session.install("-e", local_dep)
 

bigquery/tests/unit/test_client.py

@@ -24,6 +24,7 @@
 import unittest
 import warnings
 
+import freezegun
 import mock
 import requests
 import six
@@ -5367,6 +5368,43 @@ def test_list_partitions_with_string_id(self):
 
         self.assertEqual(len(partition_list), 0)
 
+    def test_list_partitions_splitting_timout_between_requests(self):
+        from google.cloud.bigquery.table import Table
+
+        row_count = 2
+        meta_info = _make_list_partitons_meta_info(
+            self.PROJECT, self.DS_ID, self.TABLE_ID, row_count
+        )
+
+        data = {
+            "totalRows": str(row_count),
+            "rows": [{"f": [{"v": "20180101"}]}, {"f": [{"v": "20180102"}]},],
+        }
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
+        client._connection = make_connection(meta_info, data)
+        table = Table(self.TABLE_REF)
+
+        with freezegun.freeze_time("2019-01-01 00:00:00", tick=False) as frozen_time:
+
+            def delayed_get_table(*args, **kwargs):
+                frozen_time.tick(delta=1.4)
+                return orig_get_table(*args, **kwargs)
+
+            orig_get_table = client.get_table
+            client.get_table = mock.Mock(side_effect=delayed_get_table)
+
+            client.list_partitions(table, timeout=5.0)
+
+        client.get_table.assert_called_once()
+        _, kwargs = client.get_table.call_args
+        self.assertEqual(kwargs.get("timeout"), 5.0)
+
+        client._connection.api_request.assert_called()
+        _, kwargs = client._connection.api_request.call_args
+        self.assertAlmostEqual(kwargs.get("timeout"), 3.6, places=5)
+
     def test_list_rows(self):
         import datetime
         from google.cloud._helpers import UTC
@@ -5687,6 +5725,46 @@ def test_list_rows_with_missing_schema(self):
             self.assertEqual(rows[1].age, 31, msg=repr(table))
             self.assertIsNone(rows[2].age, msg=repr(table))
 
+    def test_list_rows_splitting_timout_between_requests(self):
+        from google.cloud.bigquery.schema import SchemaField
+        from google.cloud.bigquery.table import Table
+
+        response = {"totalRows": "0", "rows": []}
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
+        client._connection = make_connection(response, response)
+
+        table = Table(
+            self.TABLE_REF, schema=[SchemaField("field_x", "INTEGER", mode="NULLABLE")]
+        )
+
+        with freezegun.freeze_time("1970-01-01 00:00:00", tick=False) as frozen_time:
+
+            def delayed_get_table(*args, **kwargs):
+                frozen_time.tick(delta=1.4)
+                return table
+
+            client.get_table = mock.Mock(side_effect=delayed_get_table)
+
+            rows_iter = client.list_rows(
+                "{}.{}.{}".format(
+                    self.TABLE_REF.project,
+                    self.TABLE_REF.dataset_id,
+                    self.TABLE_REF.table_id,
+                ),
+                timeout=5.0,
+            )
+            six.next(rows_iter.pages)
+
+        client.get_table.assert_called_once()
+        _, kwargs = client.get_table.call_args
+        self.assertEqual(kwargs.get("timeout"), 5.0)
+
+        client._connection.api_request.assert_called_once()
+        _, kwargs = client._connection.api_request.call_args
+        self.assertAlmostEqual(kwargs.get("timeout"), 3.6)
+
     def test_list_rows_error(self):
         creds = _make_credentials()
         http = object()

bigquery/tests/unit/test_job.py

@@ -18,6 +18,7 @@
 import textwrap
 import unittest
 
+import freezegun
 import mock
 import pytest
 import requests
@@ -956,6 +957,24 @@ def test_result_explicit_w_state(self, result):
         begin.assert_not_called()
         result.assert_called_once_with(timeout=timeout)
 
+    @mock.patch("google.api_core.future.polling.PollingFuture.result")
+    def test_result_splitting_timout_between_requests(self, result):
+        client = _make_client(project=self.PROJECT)
+        job = self._make_one(self.JOB_ID, client)
+        begin = job._begin = mock.Mock()
+        retry = mock.Mock()
+
+        with freezegun.freeze_time("1970-01-01 00:00:00", tick=False) as frozen_time:
+
+            def delayed_begin(*args, **kwargs):
+                frozen_time.tick(delta=0.3)
+
+            begin.side_effect = delayed_begin
+            job.result(retry=retry, timeout=1.0)
+
+        begin.assert_called_once_with(retry=retry, timeout=1.0)
+        result.assert_called_once_with(timeout=0.7)
+
     def test_cancelled_wo_error_result(self):
         client = _make_client(project=self.PROJECT)
         job = self._make_one(self.JOB_ID, client)
@@ -4551,7 +4570,8 @@ def test_result_w_timeout(self):
         client = _make_client(project=self.PROJECT, connection=connection)
         job = self._make_one(self.JOB_ID, self.QUERY, client)
 
-        job.result(timeout=1.0)
+        with freezegun.freeze_time("1970-01-01 00:00:00", tick=False):
+            job.result(timeout=1.0)
 
         self.assertEqual(len(connection.api_request.call_args_list), 3)
         begin_request = connection.api_request.call_args_list[0]
@@ -4566,6 +4586,49 @@ def test_result_w_timeout(self):
         self.assertEqual(query_request[1]["query_params"]["timeoutMs"], 900)
         self.assertEqual(reload_request[1]["method"], "GET")
 
+    @mock.patch("google.api_core.future.polling.PollingFuture.result")
+    def test_result_splitting_timout_between_requests(self, polling_result):
+        begun_resource = self._make_resource()
+        query_resource = {
+            "jobComplete": True,
+            "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
+            "schema": {"fields": [{"name": "col1", "type": "STRING"}]},
+            "totalRows": "5",
+        }
+        done_resource = copy.deepcopy(begun_resource)
+        done_resource["status"] = {"state": "DONE"}
+
+        connection = _make_connection(begun_resource, query_resource, done_resource)
+        client = _make_client(project=self.PROJECT, connection=connection)
+        job = self._make_one(self.JOB_ID, self.QUERY, client)
+
+        client.list_rows = mock.Mock()
+
+        with freezegun.freeze_time("1970-01-01 00:00:00", tick=False) as frozen_time:
+
+            def delayed_result(*args, **kwargs):
+                frozen_time.tick(delta=0.8)
+
+            polling_result.side_effect = delayed_result
+
+            def delayed_get_results(*args, **kwargs):
+                frozen_time.tick(delta=0.5)
+                return orig_get_results(*args, **kwargs)
+
+            orig_get_results = client._get_query_results
+            client._get_query_results = mock.Mock(side_effect=delayed_get_results)
+            job.result(timeout=2.0)
+
+        polling_result.assert_called_once_with(timeout=2.0)
+
+        client._get_query_results.assert_called_once()
+        _, kwargs = client._get_query_results.call_args
+        self.assertAlmostEqual(kwargs.get("timeout"), 1.2)
+
+        client.list_rows.assert_called_once()
+        _, kwargs = client.list_rows.call_args
+        self.assertAlmostEqual(kwargs.get("timeout"), 0.7)
+
     def test_result_w_page_size(self):
         # Arrange
         query_results_resource = {