Merge pull request #91 from jasonab/google-credentials

Upgrade to latest OAuth library(2.0.2) and Google Python library (1.5.0)

Author: tylertreat

bigquery/client.py

@@ -4,19 +4,19 @@
 from collections import defaultdict
 from datetime import datetime, timedelta
 from hashlib import sha256
+from io import StringIO
 from time import sleep, time
 
-import httplib2
 import six
-from apiclient.discovery import build, DISCOVERY_URI
-from apiclient.errors import HttpError
-
 from bigquery.errors import (BigQueryTimeoutException, JobExecutingException,
                              JobInsertException, UnfinishedQueryException)
 from bigquery.schema_builder import schema_from_record
+from googleapiclient.discovery import build, DISCOVERY_URI
+from googleapiclient.errors import HttpError
+from httplib2 import Http
 
-BIGQUERY_SCOPE = 'https://www.googleapis.com/auth/bigquery'
-BIGQUERY_SCOPE_READ_ONLY = 'https://www.googleapis.com/auth/bigquery.readonly'
+BIGQUERY_SCOPE = ['https://www.googleapis.com/auth/bigquery']
+BIGQUERY_SCOPE_READ_ONLY = ['https://www.googleapis.com/auth/bigquery.readonly']
 
 CACHE_TIMEOUT = timedelta(seconds=30)
 
@@ -90,56 +90,63 @@ def get_client(project_id, credentials=None,
     """
 
     if not credentials:
-        assert (service_account and (private_key or private_key_file)) or (json_key or json_key_file), \
+        assert (service_account and (private_key or private_key_file)) or (
+            json_key or json_key_file), \
             'Must provide AssertionCredentials or service account and P12 key or JSON key'
 
     if service_url is None:
         service_url = DISCOVERY_URI
 
+    scope = BIGQUERY_SCOPE_READ_ONLY if readonly else BIGQUERY_SCOPE
+
     if private_key_file:
-        with open(private_key_file, 'rb') as key_file:
-            private_key = key_file.read()
+        credentials = _credentials().from_p12_keyfile(service_account,
+                                                      private_key_file,
+                                                      scopes=scope)
+
+    if private_key:
+        try:
+            if isinstance(private_key, basestring):
+                private_key = private_key.decode('utf-8')
+        except NameError:
+            # python3 -- private_key is already unicode
+            pass
+        credentials = _credentials().from_p12_keyfile_buffer(
+            service_account,
+            StringIO(private_key),
+            scopes=scope)
 
     if json_key_file:
-        with open(json_key_file, 'r') as key_file:
-            json_key = json.load(key_file)
+        credentials = _credentials().from_json_keyfile_name(json_key_file,
+                                                            scopes=scope)
 
     if json_key:
-        service_account = json_key['client_email']
-        private_key = json_key['private_key']
+        credentials = _credentials().from_json_keyfile_dict(json_key,
+                                                            scopes=scope)
 
     bq_service = _get_bq_service(credentials=credentials,
-                                 service_url=service_url,
-                                 service_account=service_account,
-                                 private_key=private_key,
-                                 readonly=readonly)
+                                 service_url=service_url)
 
     return BigQueryClient(bq_service, project_id, swallow_results)
 
 
-def _get_bq_service(credentials=None, service_url=None, service_account=None, private_key=None,
-                    readonly=True):
+def _get_bq_service(credentials=None, service_url=None):
     """Construct an authorized BigQuery service object."""
 
-    assert credentials or (service_account and private_key), \
-        'Must provide AssertionCredentials or service account and key'
-
-    if not credentials:
-        scope = BIGQUERY_SCOPE_READ_ONLY if readonly else BIGQUERY_SCOPE
-        credentials = _credentials()(service_account, private_key, scope=scope)
+    assert credentials, 'Must provide ServiceAccountCredentials'
 
-    http = httplib2.Http()
-    http = credentials.authorize(http)
-    service = build('bigquery', 'v2', http=http, discoveryServiceUrl=service_url)
+    http = credentials.authorize(Http())
+    service = build('bigquery', 'v2', http=http,
+                    discoveryServiceUrl=service_url)
 
     return service
 
 
 def _credentials():
     """Import and return SignedJwtAssertionCredentials class"""
-    from oauth2client.client import SignedJwtAssertionCredentials
+    from oauth2client.service_account import ServiceAccountCredentials
 
-    return SignedJwtAssertionCredentials
+    return ServiceAccountCredentials
 
 
 class BigQueryClient(object):

bigquery/tests/test_client.py

@@ -2,18 +2,16 @@
 
 import mock
 import six
-from nose.tools import raises
-
-from apiclient.errors import HttpError
 from bigquery import client
 from bigquery.errors import (
     JobInsertException, JobExecutingException,
     BigQueryTimeoutException
 )
+from googleapiclient.errors import HttpError
+from nose.tools import raises
 
 
 class HttpResponse(object):
-
     def __init__(self, status, reason='There was an error'):
         """
         Args:
@@ -24,7 +22,6 @@ def __init__(self, status, reason='There was an error'):
 
 
 class TestGetClient(unittest.TestCase):
-
     def setUp(self):
         client._bq_client = None
 
@@ -51,7 +48,7 @@ def test_initialize_readonly(self, mock_build, mock_return_cred):
         mock_cred = mock.Mock()
         mock_http = mock.Mock()
         mock_service_url = mock.Mock()
-        mock_cred.return_value.authorize.return_value = mock_http
+        mock_cred.from_p12_keyfile_buffer.return_value.authorize.return_value = mock_http
         mock_bq = mock.Mock()
         mock_build.return_value = mock_bq
         key = 'key'
@@ -65,9 +62,11 @@ def test_initialize_readonly(self, mock_build, mock_return_cred):
             readonly=True)
 
         mock_return_cred.assert_called_once_with()
-        mock_cred.assert_called_once_with(service_account, key,
-                                          scope=BIGQUERY_SCOPE_READ_ONLY)
-        self.assertTrue(mock_cred.return_value.authorize.called)
+        mock_cred.from_p12_keyfile_buffer.assert_called_once_with(
+            service_account, mock.ANY,
+            scopes=BIGQUERY_SCOPE_READ_ONLY)
+        self.assertTrue(
+            mock_cred.from_p12_keyfile_buffer.return_value.authorize.called)
         mock_build.assert_called_once_with('bigquery', 'v2', http=mock_http,
                                            discoveryServiceUrl=mock_service_url)
         self.assertEquals(mock_bq, bq_client.bigquery)
@@ -84,7 +83,7 @@ def test_initialize_read_write(self, mock_build, mock_return_cred):
         mock_cred = mock.Mock()
         mock_http = mock.Mock()
         mock_service_url = mock.Mock()
-        mock_cred.return_value.authorize.return_value = mock_http
+        mock_cred.from_p12_keyfile_buffer.return_value.authorize.return_value = mock_http
         mock_bq = mock.Mock()
         mock_build.return_value = mock_bq
         key = 'key'
@@ -98,19 +97,18 @@ def test_initialize_read_write(self, mock_build, mock_return_cred):
             readonly=False)
 
         mock_return_cred.assert_called_once_with()
-        mock_cred.assert_called_once_with(service_account, key,
-                                          scope=BIGQUERY_SCOPE)
-        self.assertTrue(mock_cred.return_value.authorize.called)
+        mock_cred.from_p12_keyfile_buffer.assert_called_once_with(
+            service_account, mock.ANY, scopes=BIGQUERY_SCOPE)
+        self.assertTrue(
+            mock_cred.from_p12_keyfile_buffer.return_value.authorize.called)
         mock_build.assert_called_once_with('bigquery', 'v2', http=mock_http,
                                            discoveryServiceUrl=mock_service_url)
         self.assertEquals(mock_bq, bq_client.bigquery)
         self.assertEquals(project_id, bq_client.project_id)
 
     @mock.patch('bigquery.client._credentials')
     @mock.patch('bigquery.client.build')
-    @mock.patch('__builtin__.open' if six.PY2 else 'builtins.open')
-    def test_initialize_key_file(self, mock_open, mock_build,
-                                 mock_return_cred):
+    def test_initialize_key_file(self, mock_build, mock_return_cred):
         """Ensure that a BigQueryClient is initialized and returned with
         read/write permissions using a private key file.
         """
@@ -119,12 +117,10 @@ def test_initialize_key_file(self, mock_open, mock_build,
         mock_cred = mock.Mock()
         mock_http = mock.Mock()
         mock_service_url = mock.Mock()
-        mock_cred.return_value.authorize.return_value = mock_http
+        mock_cred.from_p12_keyfile.return_value.authorize.return_value = mock_http
         mock_bq = mock.Mock()
         mock_build.return_value = mock_bq
         key_file = 'key.pem'
-        key = 'key'
-        mock_open.return_value.__enter__.return_value.read.return_value = key
         service_account = 'account'
         project_id = 'project'
         mock_return_cred.return_value = mock_cred
@@ -134,46 +130,46 @@ def test_initialize_key_file(self, mock_open, mock_build,
             service_account=service_account,
             private_key_file=key_file, readonly=False)
 
-        mock_open.assert_called_once_with(key_file, 'rb')
         mock_return_cred.assert_called_once_with()
-        mock_cred.assert_called_once_with(service_account, key,
-                                          scope=BIGQUERY_SCOPE)
-        self.assertTrue(mock_cred.return_value.authorize.called)
+        mock_cred.from_p12_keyfile.assert_called_once_with(service_account,
+                                                           key_file,
+                                                           scopes=BIGQUERY_SCOPE)
+        self.assertTrue(
+            mock_cred.from_p12_keyfile.return_value.authorize.called)
         mock_build.assert_called_once_with('bigquery', 'v2', http=mock_http,
                                            discoveryServiceUrl=mock_service_url)
         self.assertEquals(mock_bq, bq_client.bigquery)
         self.assertEquals(project_id, bq_client.project_id)
 
     @mock.patch('bigquery.client._credentials')
     @mock.patch('bigquery.client.build')
-    @mock.patch('__builtin__.open' if six.PY2 else 'builtins.open')
-    def test_initialize_json_key_file(self, mock_open, mock_build, mock_return_cred):
+    def test_initialize_json_key_file(self, mock_build, mock_return_cred):
         """Ensure that a BigQueryClient is initialized and returned with
         read/write permissions using a JSON key file.
         """
         from bigquery.client import BIGQUERY_SCOPE
-        import json
 
         mock_cred = mock.Mock()
         mock_http = mock.Mock()
         mock_service_url = mock.Mock()
-        mock_cred.return_value.authorize.return_value = mock_http
+        mock_cred.from_json_keyfile_name.return_value.authorize.return_value = mock_http
         mock_bq = mock.Mock()
         mock_build.return_value = mock_bq
         json_key_file = 'key.json'
-        json_key = {'client_email': 'mail', 'private_key': 'pkey'}
-        mock_open.return_value.__enter__.return_value.read.return_value = json.dumps(json_key)
         project_id = 'project'
         mock_return_cred.return_value = mock_cred
 
         bq_client = client.get_client(
-            project_id, service_url=mock_service_url, json_key_file=json_key_file, readonly=False)
+            project_id, service_url=mock_service_url,
+            json_key_file=json_key_file, readonly=False)
 
-        mock_open.assert_called_once_with(json_key_file, 'r')
         mock_return_cred.assert_called_once_with()
-        mock_cred.assert_called_once_with(json_key['client_email'], json_key['private_key'], scope=BIGQUERY_SCOPE)
-        self.assertTrue(mock_cred.return_value.authorize.called)
-        mock_build.assert_called_once_with('bigquery', 'v2', http=mock_http, discoveryServiceUrl=mock_service_url)
+        mock_cred.from_json_keyfile_name.assert_called_once_with(json_key_file,
+                                                                 scopes=BIGQUERY_SCOPE)
+        self.assertTrue(
+            mock_cred.from_json_keyfile_name.return_value.authorize.called)
+        mock_build.assert_called_once_with('bigquery', 'v2', http=mock_http,
+                                           discoveryServiceUrl=mock_service_url)
         self.assertEquals(mock_bq, bq_client.bigquery)
         self.assertEquals(project_id, bq_client.project_id)