Repository: googleapis/google-auth-library-python

Issue:
Add timeout parameter to AuthorizedSession.request()
As a Google Cloud client library developer, I would like an easy way of specifying a timeout when making API calls using the `requests.AuthorizedSession` transport, for example:

```py
transport = AuthorizedSession(...)
transport.request(..., timeout=42)
```

**Motivation:**
Sometimes the BigQuery client gets stuck at the transport layer. Having a timeout parameter that can be passed to the transport is necessary for fixing several issue, for example: https://github.com/googleapis/google-cloud-python/pull/9875#discussion_r351504480


Discussion:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 300de2f..c9e6364 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,13 @@
 
 [1]: https://pypi.org/project/google-auth/#history
 
+### [1.8.2](https://www.github.com/googleapis/google-auth-library-python/compare/v1.8.1...v1.8.2) (2019-12-11)
+
+
+### Bug Fixes
+
+* revert "feat: send quota project id in x-goog-user-project header for OAuth2 credentials ([#400](https://www.github.com/googleapis/google-auth-library-python/issues/400))" ([#407](https://www.github.com/googleapis/google-auth-library-python/issues/407)) ([25ea942](https://www.github.com/googleapis/google-auth-library-python/commit/25ea942cef4378ff22adf235dd1baf1ca0d595f8))
+
 ### [1.8.1](https://www.github.com/googleapis/google-auth-library-python/compare/v1.8.0...v1.8.1) (2019-12-09)
 
 
diff --git a/docs/reference/google.auth.crypt.base.rst b/docs/reference/google.auth.crypt.base.rst
deleted file mode 100644
index a899650..0000000
--- a/docs/reference/google.auth.crypt.base.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-google.auth.crypt.base module
-=============================
-
-.. automodule:: google.auth.crypt.base
-   :members:
-   :inherited-members:
-   :show-inheritance:
diff --git a/docs/reference/google.auth.crypt.rsa.rst b/docs/reference/google.auth.crypt.rsa.rst
deleted file mode 100644
index 7060b03..0000000
--- a/docs/reference/google.auth.crypt.rsa.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-google.auth.crypt.rsa module
-============================
-
-.. automodule:: google.auth.crypt.rsa
-   :members:
-   :inherited-members:
-   :show-inheritance:
diff --git a/google/auth/transport/requests.py b/google/auth/transport/requests.py
index d1971cd..1caec0b 100644
--- a/google/auth/transport/requests.py
+++ b/google/auth/transport/requests.py
@@ -18,6 +18,8 @@ from __future__ import absolute_import
 
 import functools
 import logging
+import numbers
+import time
 
 try:
     import requests
@@ -64,6 +66,50 @@ class _Response(transport.Response):
         return self._response.content
 
 
+class TimeoutGuard(object):
+    """A context manager raising an error if the suite execution took too long.
+
+    Args:
+        timeout ([Union[None, float, Tuple[float, float]]]):
+            The maximum number of seconds a suite can run without the context
+            manager raising a timeout exception on exit. If passed as a tuple,
+            the smaller of the values is taken as a timeout. If ``None``, a
+            timeout error is never raised.
+        timeout_error_type (Optional[Exception]):
+            The type of the error to raise on timeout. Defaults to
+            :class:`requests.exceptions.Timeout`.
+    """
+
+    def __init__(self, timeout, timeout_error_type=requests.exceptions.Timeout):
+        self._timeout = timeout
+        self.remaining_timeout = timeout
+        self._timeout_error_type = timeout_error_type
+
+    def __enter__(self):
+        self._start = time.time()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_value:
+            return  # let the error bubble up automatically
+
+        if self._timeout is None:
+            return  # nothing to do, the timeout was not specified
+
+        elapsed = time.time() - self._start
+        deadline_hit = False
+
+        if isinstance(self._timeout, numbers.Number):
+            self.remaining_timeout = self._timeout - elapsed
+            deadline_hit = self.remaining_timeout <= 0
+        else:
+            self.remaining_timeout = tuple(x - elapsed for x in self._timeout)
+            deadline_hit = min(self.remaining_timeout) <= 0
+
+        if deadline_hit:
+            raise self._timeout_error_type()
+
+
 class Request(transport.Request):
     """Requests request adapter.
 
@@ -193,8 +239,19 @@ class AuthorizedSession(requests.Session):
         # credentials.refresh).
         self._auth_request = auth_request
 
-    def request(self, method, url, data=None, headers=None, **kwargs):
-        """Implementation of Requests' request."""
+    def request(self, method, url, data=None, headers=None, timeout=None, **kwargs):
+        """Implementation of Requests' request.
+
+        Args:
+            timeout (Optional[Union[float, Tuple[float, float]]]): The number
+                of seconds to wait before raising a ``Timeout`` exception. If
+                multiple requests are made under the hood, ``timeout`` is
+                interpreted as the approximate total time of **all** requests.
+
+                If passed as a tuple ``(connect_timeout, read_timeout)``, the
+                smaller of the values is taken as the total allowed time across
+                all requests.
+        """
         # pylint: disable=arguments-differ
         # Requests has a ton of arguments to request, but only two
         # (method, url) are required. We pass through all of the other
@@ -208,13 +265,28 @@ class AuthorizedSession(requests.Session):
         # and we want to pass the original headers if we recurse.
         request_headers = headers.copy() if headers is not None else {}
 
-        self.credentials.before_request(
-            self._auth_request, method, url, request_headers
+        # Do not apply the timeout unconditionally in order to not override the
+        # _auth_request's default timeout.
+        auth_request = (
+            self._auth_request
+            if timeout is None
+            else functools.partial(self._auth_request, timeout=timeout)
         )
 
-        response = super(AuthorizedSession, self).request(
-            method, url, data=data, headers=request_headers, **kwargs
-        )
+        with TimeoutGuard(timeout) as guard:
+            self.credentials.before_request(auth_request, method, url, request_headers)
+        timeout = guard.remaining_timeout
+
+        with TimeoutGuard(timeout) as guard:
+            response = super(AuthorizedSession, self).request(
+                method,
+                url,
+                data=data,
+                headers=request_headers,
+                timeout=timeout,
+                **kwargs
+            )
+        timeout = guard.remaining_timeout
 
         # If the response indicated that the credentials needed to be
         # refreshed, then refresh the credentials and re-attempt the
@@ -233,17 +305,34 @@ class AuthorizedSession(requests.Session):
                 self._max_refresh_attempts,
             )
 
-            auth_request_with_timeout = functools.partial(
-                self._auth_request, timeout=self._refresh_timeout
+            if self._refresh_timeout is not None:
+                if timeout is None:
+                    timeout = self._refresh_timeout
+                elif isinstance(timeout, numbers.Number):
+                    timeout = min(timeout, self._refresh_timeout)
+                else:
+                    timeout = tuple(min(x, self._refresh_timeout) for x in timeout)
+
+            # Do not apply the timeout unconditionally in order to not override the
+            # _auth_request's default timeout.
+            auth_request = (
+                self._auth_request
+                if timeout is None
+                else functools.partial(self._auth_request, timeout=timeout)
             )
-            self.credentials.refresh(auth_request_with_timeout)
 
-            # Recurse. Pass in the original headers, not our modified set.
+            with TimeoutGuard(timeout) as guard:
+                self.credentials.refresh(auth_request)
+            timeout = guard.remaining_timeout
+
+            # Recurse. Pass in the original headers, not our modified set, but
+            # do pass the adjusted timeout (i.e. the remaining time).
             return self.request(
                 method,
                 url,
                 data=data,
                 headers=headers,
+                timeout=timeout,
                 _credential_refresh_attempt=_credential_refresh_attempt + 1,
                 **kwargs
             )
diff --git a/google/oauth2/credentials.py b/google/oauth2/credentials.py
index 422c8ab..3a32c06 100644
--- a/google/oauth2/credentials.py
+++ b/google/oauth2/credentials.py
@@ -58,7 +58,6 @@ class Credentials(credentials.ReadOnlyScoped, credentials.Credentials):
         client_id=None,
         client_secret=None,
         scopes=None,
-        quota_project_id=None,
     ):
         """
         Args:
@@ -82,9 +81,6 @@ class Credentials(credentials.ReadOnlyScoped, credentials.Credentials):
                 token if refresh information is provided (e.g. The refresh
                 token scopes are a superset of this or contain a wild card
                 scope like 'https://www.googleapis.com/auth/any-api').
-            quota_project_id (Optional[str]): The project ID used for quota and billing.
-                This project may be different from the project used to
-                create the credentials.
         """
         super(Credentials, self).__init__()
         self.token = token
@@ -94,7 +90,6 @@ class Credentials(credentials.ReadOnlyScoped, credentials.Credentials):
         self._token_uri = token_uri
         self._client_id = client_id
         self._client_secret = client_secret
-        self._quota_project_id = quota_project_id
 
     @property
     def refresh_token(self):
@@ -128,11 +123,6 @@ class Credentials(credentials.ReadOnlyScoped, credentials.Credentials):
         """Optional[str]: The OAuth 2.0 client secret."""
         return self._client_secret
 
-    @property
-    def quota_project_id(self):
-        """Optional[str]: The project to use for quota and billing purposes."""
-        return self._quota_project_id
-
     @property
     def requires_scopes(self):
         """False: OAuth 2.0 credentials have their scopes set when
@@ -179,12 +169,6 @@ class Credentials(credentials.ReadOnlyScoped, credentials.Credentials):
                     )
                 )
 
-    @_helpers.copy_docstring(credentials.Credentials)
-    def apply(self, headers, token=None):
-        super(Credentials, self).apply(headers, token=token)
-        if self.quota_project_id is not None:
-            headers["x-goog-user-project"] = self.quota_project_id
-
     @classmethod
     def from_authorized_user_info(cls, info, scopes=None):
         """Creates a Credentials instance from parsed authorized user info.
@@ -218,9 +202,6 @@ class Credentials(credentials.ReadOnlyScoped, credentials.Credentials):
             scopes=scopes,
             client_id=info["client_id"],
             client_secret=info["client_secret"],
-            quota_project_id=info.get(
-                "quota_project_id"
-            ),  # quota project may not exist
         )
 
     @classmethod
diff --git a/noxfile.py b/noxfile.py
index aaf1bc5..e170ee5 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -16,6 +16,7 @@ import nox
 
 TEST_DEPENDENCIES = [
     "flask",
+    "freezegun",
     "mock",
     "oauth2client",
     "pytest",
diff --git a/setup.py b/setup.py
index d8ba62c..d98e89c 100644
--- a/setup.py
+++ b/setup.py
@@ -30,7 +30,7 @@ DEPENDENCIES = (
 with io.open("README.rst", "r") as fh:
     long_description = fh.read()
 
-version = "1.8.1"
+version = "1.8.2"
 
 setup(
     name="google-auth",


Test patch:
diff --git a/tests/oauth2/test_credentials.py b/tests/oauth2/test_credentials.py
index 59031d7..bb70f15 100644
--- a/tests/oauth2/test_credentials.py
+++ b/tests/oauth2/test_credentials.py
@@ -294,33 +294,6 @@ class TestCredentials(object):
         # expired.)
         assert creds.valid
 
-    def test_apply_with_quota_project_id(self):
-        creds = credentials.Credentials(
-            token="token",
-            refresh_token=self.REFRESH_TOKEN,
-            token_uri=self.TOKEN_URI,
-            client_id=self.CLIENT_ID,
-            client_secret=self.CLIENT_SECRET,
-            quota_project_id="quota-project-123",
-        )
-
-        headers = {}
-        creds.apply(headers)
-        assert headers["x-goog-user-project"] == "quota-project-123"
-
-    def test_apply_with_no_quota_project_id(self):
-        creds = credentials.Credentials(
-            token="token",
-            refresh_token=self.REFRESH_TOKEN,
-            token_uri=self.TOKEN_URI,
-            client_id=self.CLIENT_ID,
-            client_secret=self.CLIENT_SECRET,
-        )
-
-        headers = {}
-        creds.apply(headers)
-        assert "x-goog-user-project" not in headers
-
     def test_from_authorized_user_info(self):
         info = AUTH_USER_INFO.copy()
 
diff --git a/tests/transport/test_requests.py b/tests/transport/test_requests.py
index 0e165ac..0026974 100644
--- a/tests/transport/test_requests.py
+++ b/tests/transport/test_requests.py
@@ -12,7 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import datetime
+import functools
+
+import freezegun
 import mock
+import pytest
 import requests
 import requests.adapters
 from six.moves import http_client
@@ -22,6 +27,12 @@ import google.auth.transport.requests
 from tests.transport import compliance
 
 
+@pytest.fixture
+def frozen_time():
+    with freezegun.freeze_time("1970-01-01 00:00:00", tick=False) as frozen:
+        yield frozen
+
+
 class TestRequestResponse(compliance.RequestResponseTests):
     def make_request(self):
         return google.auth.transport.requests.Request()
@@ -34,6 +45,52 @@ class TestRequestResponse(compliance.RequestResponseTests):
         assert http.request.call_args[1]["timeout"] == 5
 
 
+class TestTimeoutGuard(object):
+    def make_guard(self, *args, **kwargs):
+        return google.auth.transport.requests.TimeoutGuard(*args, **kwargs)
+
+    def test_tracks_elapsed_time_w_numeric_timeout(self, frozen_time):
+        with self.make_guard(timeout=10) as guard:
+            frozen_time.tick(delta=3.8)
+        assert guard.remaining_timeout == 6.2
+
+    def test_tracks_elapsed_time_w_tuple_timeout(self, frozen_time):
+        with self.make_guard(timeout=(16, 19)) as guard:
+            frozen_time.tick(delta=3.8)
+        assert guard.remaining_timeout == (12.2, 15.2)
+
+    def test_noop_if_no_timeout(self, frozen_time):
+        with self.make_guard(timeout=None) as guard:
+            frozen_time.tick(delta=datetime.timedelta(days=3650))
+        # NOTE: no timeout error raised, despite years have passed
+        assert guard.remaining_timeout is None
+
+    def test_timeout_error_w_numeric_timeout(self, frozen_time):
+        with pytest.raises(requests.exceptions.Timeout):
+            with self.make_guard(timeout=10) as guard:
+                frozen_time.tick(delta=10.001)
+        assert guard.remaining_timeout == pytest.approx(-0.001)
+
+    def test_timeout_error_w_tuple_timeout(self, frozen_time):
+        with pytest.raises(requests.exceptions.Timeout):
+            with self.make_guard(timeout=(11, 10)) as guard:
+                frozen_time.tick(delta=10.001)
+        assert guard.remaining_timeout == pytest.approx((0.999, -0.001))
+
+    def test_custom_timeout_error_type(self, frozen_time):
+        class FooError(Exception):
+            pass
+
+        with pytest.raises(FooError):
+            with self.make_guard(timeout=1, timeout_error_type=FooError):
+                frozen_time.tick(2)
+
+    def test_lets_suite_errors_bubble_up(self, frozen_time):
+        with pytest.raises(IndexError):
+            with self.make_guard(timeout=1):
+                [1, 2, 3][3]
+
+
 class CredentialsStub(google.auth.credentials.Credentials):
     def __init__(self, token="token"):
         super(CredentialsStub, self).__init__()
@@ -49,6 +106,18 @@ class CredentialsStub(google.auth.credentials.Credentials):
         self.token += "1"
 
 
+class TimeTickCredentialsStub(CredentialsStub):
+    """Credentials that spend some (mocked) time when refreshing a token."""
+
+    def __init__(self, time_tick, token="token"):
+        self._time_tick = time_tick
+        super(TimeTickCredentialsStub, self).__init__(token=token)
+
+    def refresh(self, request):
+        self._time_tick()
+        super(TimeTickCredentialsStub, self).refresh(requests)
+
+
 class AdapterStub(requests.adapters.BaseAdapter):
     def __init__(self, responses, headers=None):
         super(AdapterStub, self).__init__()
@@ -69,6 +138,18 @@ class AdapterStub(requests.adapters.BaseAdapter):
         return
 
 
+class TimeTickAdapterStub(AdapterStub):
+    """Adapter that spends some (mocked) time when making a request."""
+
+    def __init__(self, time_tick, responses, headers=None):
+        self._time_tick = time_tick
+        super(TimeTickAdapterStub, self).__init__(responses, headers=headers)
+
+    def send(self, request, **kwargs):
+        self._time_tick()
+        return super(TimeTickAdapterStub, self).send(request, **kwargs)
+
+
 def make_response(status=http_client.OK, data=None):
     response = requests.Response()
     response.status_code = status
@@ -121,7 +202,9 @@ class TestAuthorizedHttp(object):
             [make_response(status=http_client.UNAUTHORIZED), final_response]
         )
 
-        authed_session = google.auth.transport.requests.AuthorizedSession(credentials)
+        authed_session = google.auth.transport.requests.AuthorizedSession(
+            credentials, refresh_timeout=60
+        )
         authed_session.mount(self.TEST_URL, adapter)
 
         result = authed_session.request("GET", self.TEST_URL)
@@ -136,3 +219,72 @@ class TestAuthorizedHttp(object):
 
         assert adapter.requests[1].url == self.TEST_URL
         assert adapter.requests[1].headers["authorization"] == "token1"
+
+    def test_request_timeout(self, frozen_time):
+        tick_one_second = functools.partial(frozen_time.tick, delta=1.0)
+
+        credentials = mock.Mock(
+            wraps=TimeTickCredentialsStub(time_tick=tick_one_second)
+        )
+        adapter = TimeTickAdapterStub(
+            time_tick=tick_one_second,
+            responses=[
+                make_response(status=http_client.UNAUTHORIZED),
+                make_response(status=http_client.OK),
+            ],
+        )
+
+        authed_session = google.auth.transport.requests.AuthorizedSession(credentials)
+        authed_session.mount(self.TEST_URL, adapter)
+
+        # Because at least two requests have to be made, and each takes one
+        # second, the total timeout specified will be exceeded.
+        with pytest.raises(requests.exceptions.Timeout):
+            authed_session.request("GET", self.TEST_URL, timeout=1.9)
+
+    def test_request_timeout_w_refresh_timeout(self, frozen_time):
+        tick_one_second = functools.partial(frozen_time.tick, delta=1.0)
+
+        credentials = mock.Mock(
+            wraps=TimeTickCredentialsStub(time_tick=tick_one_second)
+        )
+        adapter = TimeTickAdapterStub(
+            time_tick=tick_one_second,
+            responses=[
+                make_response(status=http_client.UNAUTHORIZED),
+                make_response(status=http_client.OK),
+            ],
+        )
+
+        authed_session = google.auth.transport.requests.AuthorizedSession(
+            credentials, refresh_timeout=1.9
+        )
+        authed_session.mount(self.TEST_URL, adapter)
+
+        # The timeout is long, but the short refresh timeout will prevail.
+        with pytest.raises(requests.exceptions.Timeout):
+            authed_session.request("GET", self.TEST_URL, timeout=60)
+
+    def test_request_timeout_w_refresh_timeout_and_tuple_timeout(self, frozen_time):
+        tick_one_second = functools.partial(frozen_time.tick, delta=1.0)
+
+        credentials = mock.Mock(
+            wraps=TimeTickCredentialsStub(time_tick=tick_one_second)
+        )
+        adapter = TimeTickAdapterStub(
+            time_tick=tick_one_second,
+            responses=[
+                make_response(status=http_client.UNAUTHORIZED),
+                make_response(status=http_client.OK),
+            ],
+        )
+
+        authed_session = google.auth.transport.requests.AuthorizedSession(
+            credentials, refresh_timeout=100
+        )
+        authed_session.mount(self.TEST_URL, adapter)
+
+        # The shortest timeout will prevail and cause a Timeout error, despite
+        # other timeouts being quite long.
+        with pytest.raises(requests.exceptions.Timeout):
+            authed_session.request("GET", self.TEST_URL, timeout=(100, 2.9))