beam-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From rober...@apache.org
Subject [1/2] incubator-beam git commit: auth: add application default credentials as fallback
Date Fri, 02 Dec 2016 19:30:05 GMT
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 <vikasrk@google.com>
Authored: Wed Nov 30 17:55:20 2016 -0800
Committer: Robert Bradshaw <robertwb@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):


Mime
View raw message