Add private_key_file kwarg on get_client

tylertreat · tylertreat · commit e8a77965baf0 · 2015-01-23T19:26:47.000-06:00
This kwarg allows you to pass the key file name and get_client handles
reading it. The actual key string can still be passed in using
private_key.
diff --git a/README.md b/README.md
@@ -21,9 +21,10 @@ project_id = 'project_id'
 service_account = 'my_id_123@developer.gserviceaccount.com'
 
 # PKCS12 or PEM key provided by Google.
-key = 'secret_key'
+key = 'key.pem'
 
-client = get_client(project_id, service_account=service_account, private_key=key, readonly=True)
+client = get_client(project_id, service_account=service_account,
+                    private_key_file=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
@@ -46,7 +46,8 @@
 
 
 def get_client(project_id, credentials=None, service_account=None,
-               private_key=None, readonly=True, swallow_results=True):
+               private_key=None, private_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.
@@ -58,6 +59,9 @@ def get_client(project_id, credentials=None, service_account=None,
         service_account: the Google API service account name.
         private_key: the private key associated with the service account in
                      PKCS12 or PEM format.
+        private_key_file: the name of the file containing the private key
+                          associated with the service account in PKCS12 or PEM
+                          format.
         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
@@ -67,9 +71,13 @@ def get_client(project_id, credentials=None, service_account=None,
         an instance of BigQueryClient.
     """
 
-    if not credentials and not (service_account and private_key):
-        raise Exception('AssertionCredentials or service account and private'
-                        'key need to be provided')
+    if not credentials:
+        assert service_account and (private_key or private_key_file), \
+            'Must provide AssertionCredentials or service account and key'
+
+    if private_key_file:
+        with open(private_key_file, 'rb') as key_file:
+            private_key = key_file.read()
 
     bq_service = _get_bq_service(credentials=credentials,
                                  service_account=service_account,
@@ -83,7 +91,8 @@ def _get_bq_service(credentials=None, service_account=None, private_key=None,
                     readonly=True):
     """Construct an authorized BigQuery service object."""
 
-    assert credentials or (service_account and private_key)
+    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
@@ -820,7 +829,7 @@ def _get_all_tables(self, dataset_id, cache=False):
                     projectId=self.project_id,
                     datasetId=dataset_id,
                     pageToken=page_token
-                    ).execute()
+                ).execute()
                 page_token = res.get('nextPageToken')
                 result['tables'] += res.get('tables', [])
             self.cache[dataset_id] = (datetime.now(), result)
diff --git a/bigquery/tests/test_client.py b/bigquery/tests/test_client.py
@@ -37,7 +37,7 @@ def setUp(self):
     def test_no_credentials(self):
         """Ensure an Exception is raised when no credentials are provided."""
 
-        self.assertRaises(Exception, client.get_client, 'foo', 'bar')
+        self.assertRaises(AssertionError, client.get_client, 'foo')
 
     @mock.patch('bigquery.client._credentials')
     @mock.patch('bigquery.client.build')
@@ -99,6 +99,41 @@ def test_initialize_read_write(self, mock_build, mock_return_cred):
         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')
+    def test_initialize_key_file(self, mock_open, mock_build,
+                                 mock_return_cred):
+        """Ensure that a BigQueryClient is initialized and returned with
+        read/write permissions using a private key file.
+        """
+        from bigquery.client import BIGQUERY_SCOPE
+
+        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
+        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
+
+        bq_client = client.get_client(
+            project_id, 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)
+        mock_cred.authorize.assert_called_once()
+        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):
 
