Repository: incubator-beam Updated Branches: refs/heads/python-sdk 9ded359da -> a463f000e
auth: add application default credentials as fallback Project: http://git-wip-us.apache.org/repos/asf/incubator-beam/repo Commit: http://git-wip-us.apache.org/repos/asf/incubator-beam/commit/01bddf29 Tree: http://git-wip-us.apache.org/repos/asf/incubator-beam/tree/01bddf29 Diff: http://git-wip-us.apache.org/repos/asf/incubator-beam/diff/01bddf29 Branch: refs/heads/python-sdk Commit: 01bddf296dfb84a70bc733add6a76c76cf6afaef Parents: 9ded359 Author: Vikas Kedigehalli <vika...@google.com> Authored: Wed Nov 30 17:55:20 2016 -0800 Committer: Robert Bradshaw <rober...@google.com> Committed: Fri Dec 2 11:29:41 2016 -0800 ---------------------------------------------------------------------- sdks/python/apache_beam/internal/auth.py | 37 +++++++++++++++----- .../apache_beam/io/datastore/v1/datastoreio.py | 10 ++++-- .../apache_beam/io/datastore/v1/helper.py | 6 ++-- 3 files changed, 37 insertions(+), 16 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/01bddf29/sdks/python/apache_beam/internal/auth.py ---------------------------------------------------------------------- diff --git a/sdks/python/apache_beam/internal/auth.py b/sdks/python/apache_beam/internal/auth.py index f324a2d..a043fcf 100644 --- a/sdks/python/apache_beam/internal/auth.py +++ b/sdks/python/apache_beam/internal/auth.py @@ -24,7 +24,7 @@ import os import sys import urllib2 - +from oauth2client.client import GoogleCredentials from oauth2client.client import OAuth2Credentials from apache_beam.utils import processes @@ -125,6 +125,14 @@ def get_service_credentials(): # them again. return GCEMetadataCredentials(user_agent=user_agent) else: + client_scopes = [ + 'https://www.googleapis.com/auth/bigquery', + 'https://www.googleapis.com/auth/cloud-platform', + 'https://www.googleapis.com/auth/devstorage.full_control', + 'https://www.googleapis.com/auth/userinfo.email', + 'https://www.googleapis.com/auth/datastore' + ] + # We are currently being run from the command line. google_cloud_options = PipelineOptions( sys.argv).view_as(GoogleCloudOptions) @@ -135,13 +143,6 @@ def get_service_credentials(): if not os.path.exists(google_cloud_options.service_account_key_file): raise AuthenticationException( 'Specified service account key file does not exist.') - client_scopes = [ - 'https://www.googleapis.com/auth/bigquery', - 'https://www.googleapis.com/auth/cloud-platform', - 'https://www.googleapis.com/auth/devstorage.full_control', - 'https://www.googleapis.com/auth/userinfo.email', - 'https://www.googleapis.com/auth/datastore' - ] # The following code uses oauth2client >=2.0.0 functionality and if this # is not available due to import errors will use 1.5.2 functionality. @@ -163,4 +164,22 @@ def get_service_credentials(): user_agent=user_agent) else: - return _GCloudWrapperCredentials(user_agent) + try: + credentials = _GCloudWrapperCredentials(user_agent) + # Check if we are able to get an access token. If not fallback to + # application default credentials. + credentials.get_access_token() + return credentials + except AuthenticationException: + logging.warning('Unable to find credentials from gcloud.') + + # Falling back to application default credentials. + try: + credentials = GoogleCredentials.get_application_default() + credentials = credentials.create_scoped(client_scopes) + logging.debug('Connecting using Google Application Default ' + 'Credentials.') + return credentials + except Exception: + logging.warning('Unable to find default credentials to use.') + raise http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/01bddf29/sdks/python/apache_beam/io/datastore/v1/datastoreio.py ---------------------------------------------------------------------- diff --git a/sdks/python/apache_beam/io/datastore/v1/datastoreio.py b/sdks/python/apache_beam/io/datastore/v1/datastoreio.py index 20466b9..054002f 100644 --- a/sdks/python/apache_beam/io/datastore/v1/datastoreio.py +++ b/sdks/python/apache_beam/io/datastore/v1/datastoreio.py @@ -22,6 +22,7 @@ import logging from google.datastore.v1 import datastore_pb2 from googledatastore import helper as datastore_helper +from apache_beam.internal import auth from apache_beam.io.datastore.v1 import helper from apache_beam.io.datastore.v1 import query_splitter from apache_beam.transforms import Create @@ -153,7 +154,8 @@ class ReadFromDatastore(PTransform): self._num_splits = num_splits def start_bundle(self, context): - self._datastore = helper.get_datastore(self._project) + self._datastore = helper.get_datastore(self._project, + auth.get_service_credentials()) def process(self, p_context, *args, **kwargs): # distinct key to be used to group query splits. @@ -208,7 +210,8 @@ class ReadFromDatastore(PTransform): self._datastore = None def start_bundle(self, context): - self._datastore = helper.get_datastore(self._project) + self._datastore = helper.get_datastore(self._project, + auth.get_service_credentials()) def process(self, p_context, *args, **kwargs): query = p_context.element @@ -338,7 +341,8 @@ class _Mutate(PTransform): def start_bundle(self, context): self._mutations = [] - self._datastore = helper.get_datastore(self._project) + self._datastore = helper.get_datastore(self._project, + auth.get_service_credentials()) def process(self, context): self._mutations.append(context.element) http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/01bddf29/sdks/python/apache_beam/io/datastore/v1/helper.py ---------------------------------------------------------------------- diff --git a/sdks/python/apache_beam/io/datastore/v1/helper.py b/sdks/python/apache_beam/io/datastore/v1/helper.py index 28cb123..720f30a 100644 --- a/sdks/python/apache_beam/io/datastore/v1/helper.py +++ b/sdks/python/apache_beam/io/datastore/v1/helper.py @@ -95,11 +95,9 @@ def str_compare(s1, s2): return 1 -def get_datastore(project): +def get_datastore(project, credentials): """Returns a Cloud Datastore client.""" - credentials = datastore_helper.get_credentials_from_env() - datastore = Datastore(project, credentials) - return datastore + return Datastore(project, credentials) def make_request(project, namespace, query):