Merge pull request tylertreat#71 from kamatama41/json-key

tylertreat · tylertreat · commit b0893974fa2d · 2015-11-25T12:26:44.000-06:00
Add support for JSON key authorization
diff --git a/README.md b/README.md
@@ -26,6 +26,11 @@ key = 'key.pem'
 client = get_client(project_id, service_account=service_account,
                     private_key_file=key, readonly=True)
 
+# JSON key provided by Google
+json_key = 'key.json'
+ 
+client = get_client(project_id, json_key_file=json_key, readonly=True)
+
 # Submit an async query.
 job_id, _results = client.query('SELECT * FROM dataset.my_table LIMIT 1000')
 
diff --git a/bigquery/client.py b/bigquery/client.py
@@ -44,8 +44,9 @@
 
 
 def get_client(project_id, credentials=None, service_account=None,
-               private_key=None, private_key_file=None, readonly=True,
-               swallow_results=True):
+               private_key=None, private_key_file=None,
+               json_key=None, json_key_file=None,
+               readonly=True, swallow_results=True):
     """Return a singleton instance of BigQueryClient. Either
     AssertionCredentials or a service account and private key combination need
     to be provided in order to authenticate requests to BigQuery.
@@ -60,6 +61,9 @@ def get_client(project_id, credentials=None, service_account=None,
         private_key_file: the name of the file containing the private key
                           associated with the service account in PKCS12 or PEM
                           format.
+        json_key: the JSON key associated with the service account
+        json_key_file: the name of the JSON key file associated with
+                       the service account
         readonly: bool indicating if BigQuery access is read-only. Has no
                   effect if credentials are provided.
         swallow_results: If set to false then return the actual response value
@@ -70,13 +74,21 @@ def get_client(project_id, credentials=None, service_account=None,
     """
 
     if not credentials:
-        assert service_account and (private_key or private_key_file), \
-            'Must provide AssertionCredentials or service account and key'
+        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 private_key_file:
         with open(private_key_file, 'rb') as key_file:
             private_key = key_file.read()
 
+    if json_key_file:
+        with open(json_key_file, 'rb') as key_file:
+            json_key = json.loads(key_file.read())
+
+    if json_key:
+        service_account = json_key['client_email']
+        private_key = json_key['private_key']
+
     bq_service = _get_bq_service(credentials=credentials,
                                  service_account=service_account,
                                  private_key=private_key,
diff --git a/bigquery/tests/test_client.py b/bigquery/tests/test_client.py
@@ -135,6 +135,37 @@ def test_initialize_key_file(self, mock_open, mock_build,
         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):
+        """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_cred.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, json_key_file=json_key_file, readonly=False)
+
+        mock_open.assert_called_once_with(json_key_file, 'rb')
+        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)
+        self.assertEquals(mock_bq, bq_client.bigquery)
+        self.assertEquals(project_id, bq_client.project_id)
+
 
 class TestQuery(unittest.TestCase):
 
