Skip to content

fix: use AuthMetadataPlugin with context manager #460

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
busunkim96 wants to merge 2 commits into from
Closed

fix: use AuthMetadataPlugin with context manager #460

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
673b145
fix: use AuthMetadataPlugin with context manager
busunkim96 Mar 16, 2020
259b78d
Merge branch 'master' into use-context-manager
busunkim96 Mar 16, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 16 additions & 14 deletions google/auth/transport/grpc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,12 @@ class AuthMetadataPlugin(grpc.AuthMetadataPlugin):
object used to refresh credentials as needed.
"""

def __enter__(self):
return self

def __exit__(self, exc_type, exc_value, traceback):
self._pool.shutdown(wait=False)

def __init__(self, credentials, request):
# pylint: disable=no-value-for-parameter
# pylint doesn't realize that the super method takes no arguments
Expand Down Expand Up @@ -87,9 +93,6 @@ def __call__(self, context, callback):
future = self._pool.submit(self._get_authorization_headers, context)
future.add_done_callback(self._callback_wrapper(callback))

def __del__(self):
self._pool.shutdown(wait=False)


def secure_authorized_channel(
credentials, request, target, ssl_credentials=None, **kwargs
Expand Down Expand Up @@ -135,17 +138,16 @@ def secure_authorized_channel(
grpc.Channel: The created gRPC channel.
"""
# Create the metadata plugin for inserting the authorization header.
metadata_plugin = AuthMetadataPlugin(credentials, request)

# Create a set of grpc.CallCredentials using the metadata plugin.
google_auth_credentials = grpc.metadata_call_credentials(metadata_plugin)
with AuthMetadataPlugin(credentials, request) as metadata_plugin:
# Create a set of grpc.CallCredentials using the metadata plugin.
google_auth_credentials = grpc.metadata_call_credentials(metadata_plugin)

if ssl_credentials is None:
ssl_credentials = grpc.ssl_channel_credentials()
if ssl_credentials is None:
ssl_credentials = grpc.ssl_channel_credentials()

# Combine the ssl credentials and the authorization credentials.
composite_credentials = grpc.composite_channel_credentials(
ssl_credentials, google_auth_credentials
)
# Combine the ssl credentials and the authorization credentials.
composite_credentials = grpc.composite_channel_credentials(
ssl_credentials, google_auth_credentials
)

return grpc.secure_channel(target, composite_credentials, **kwargs)
return grpc.secure_channel(target, composite_credentials, **kwargs)