Skip to content

Commit c245182

Browse files
author
chenyumic
authored
fix: make gRPC auth plugin non-blocking + add default timeout value for requests transport (#390)
This commit includes the following changes: - `transport.grpc.AuthMetadataPlugin` is now non-blocking as gRPC requires - `transport.requests.Request` now has a default timeout value of 120 seconds so that token refreshing will not be stuck Resolves: #351
1 parent 7ea1fdd commit c245182

File tree

3 files changed

+21
-2
lines changed

3 files changed

+21
-2
lines changed

packages/google-auth/google/auth/transport/grpc.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616

1717
from __future__ import absolute_import
1818

19+
from concurrent import futures
20+
1921
import six
2022

2123
try:
@@ -51,6 +53,7 @@ def __init__(self, credentials, request):
5153
super(AuthMetadataPlugin, self).__init__()
5254
self._credentials = credentials
5355
self._request = request
56+
self._pool = futures.ThreadPoolExecutor(max_workers=1)
5457

5558
def _get_authorization_headers(self, context):
5659
"""Gets the authorization headers for a request.
@@ -66,6 +69,13 @@ def _get_authorization_headers(self, context):
6669

6770
return list(six.iteritems(headers))
6871

72+
@staticmethod
73+
def _callback_wrapper(callback):
74+
def wrapped(future):
75+
callback(future.result(), None)
76+
77+
return wrapped
78+
6979
def __call__(self, context, callback):
7080
"""Passes authorization metadata into the given callback.
7181
@@ -74,7 +84,11 @@ def __call__(self, context, callback):
7484
callback (grpc.AuthMetadataPluginCallback): The callback that will
7585
be invoked to pass in the authorization metadata.
7686
"""
77-
callback(self._get_authorization_headers(context), None)
87+
future = self._pool.submit(self._get_authorization_headers, context)
88+
future.add_done_callback(self._callback_wrapper(callback))
89+
90+
def __del__(self):
91+
self._pool.shutdown(wait=False)
7892

7993

8094
def secure_authorized_channel(

packages/google-auth/google/auth/transport/requests.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@ def __init__(self, session=None):
9595
self.session = session
9696

9797
def __call__(
98-
self, url, method="GET", body=None, headers=None, timeout=None, **kwargs
98+
self, url, method="GET", body=None, headers=None, timeout=120, **kwargs
9999
):
100100
"""Make an HTTP request using requests.
101101

packages/google-auth/tests/transport/test_grpc.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
# limitations under the License.
1414

1515
import datetime
16+
import time
1617

1718
import mock
1819
import pytest
@@ -58,6 +59,8 @@ def test_call_no_refresh(self):
5859

5960
plugin(context, callback)
6061

62+
time.sleep(2)
63+
6164
callback.assert_called_once_with(
6265
[(u"authorization", u"Bearer {}".format(credentials.token))], None
6366
)
@@ -76,6 +79,8 @@ def test_call_refresh(self):
7679

7780
plugin(context, callback)
7881

82+
time.sleep(2)
83+
7984
assert credentials.token == "token1"
8085
callback.assert_called_once_with(
8186
[(u"authorization", u"Bearer {}".format(credentials.token))], None

0 commit comments

Comments
 (0)