Read project_id from JSON key file.

orangain · orangain · commit 7ce159a4002f · 2016-04-23T15:21:17.000+09:00
A JSON key file provided by Google contains project_id. Now project_id
argument of get_client() is optional and read from the JSON key file if
json_key or json_key_file is provided. I believe this improve usability of
get_client().
diff --git a/README.md b/README.md
@@ -31,7 +31,7 @@ client = get_client(project_id, service_account=service_account,
 # JSON key provided by Google
 json_key = 'key.json'
  
-client = get_client(project_id, json_key_file=json_key, readonly=True)
+client = get_client(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
@@ -45,7 +45,7 @@
 logger = getLogger(__name__)
 
 
-def get_client(project_id, credentials=None,
+def get_client(project_id=None, credentials=None,
                service_url=None, service_account=None,
                private_key=None, private_key_file=None,
                json_key=None, json_key_file=None,
@@ -56,8 +56,8 @@ def get_client(project_id, credentials=None,
 
     Parameters
     ----------
-    project_id : str
-        The BigQuery project id
+    project_id : str, optional
+        The BigQuery project id, required unless json_key or json_key_file is provided.
     credentials : oauth2client.client.SignedJwtAssertionCredentials, optional
         AssertionCredentials instance to authenticate requests to BigQuery (optional,
         must provide `service_account` and (`private_key` or `private_key_file`) or
@@ -96,6 +96,10 @@ def get_client(project_id, credentials=None,
             json_key or json_key_file), \
             'Must provide AssertionCredentials or service account and P12 key or JSON key'
 
+    if not project_id:
+        assert json_key or json_key_file, \
+            'Must provide project_id unless json_key or json_key_file is provided'
+
     if service_url is None:
         service_url = DISCOVERY_URI
 
@@ -119,12 +123,14 @@ def get_client(project_id, credentials=None,
             scopes=scope)
 
     if json_key_file:
-        credentials = _credentials().from_json_keyfile_name(json_key_file,
-                                                            scopes=scope)
+        with open(json_key_file, 'r') as key_file:
+            json_key = json.load(key_file)
 
     if json_key:
         credentials = _credentials().from_json_keyfile_dict(json_key,
                                                             scopes=scope)
+        if not project_id:
+            project_id = json_key['project_id']
 
     bq_service = _get_bq_service(credentials=credentials,
                                  service_url=service_url)
diff --git a/bigquery/tests/test_client.py b/bigquery/tests/test_client.py
@@ -143,19 +143,23 @@ def test_initialize_key_file(self, mock_build, mock_return_cred):
 
     @mock.patch('bigquery.client._credentials')
     @mock.patch('bigquery.client.build')
-    def test_initialize_json_key_file(self, mock_build, mock_return_cred):
+    @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_service_url = mock.Mock()
-        mock_cred.from_json_keyfile_name.return_value.authorize.return_value = mock_http
+        mock_cred.from_json_keyfile_dict.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
 
@@ -164,15 +168,51 @@ def test_initialize_json_key_file(self, mock_build, mock_return_cred):
             json_key_file=json_key_file, readonly=False)
 
         mock_return_cred.assert_called_once_with()
-        mock_cred.from_json_keyfile_name.assert_called_once_with(json_key_file,
+        mock_cred.from_json_keyfile_dict.assert_called_once_with(json_key,
                                                                  scopes=BIGQUERY_SCOPE)
         self.assertTrue(
-            mock_cred.from_json_keyfile_name.return_value.authorize.called)
+            mock_cred.from_json_keyfile_dict.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_without_project_id(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 without project_id.
+        """
+        from bigquery.client import BIGQUERY_SCOPE
+        import json
+
+        mock_cred = mock.Mock()
+        mock_http = mock.Mock()
+        mock_service_url = mock.Mock()
+        mock_cred.from_json_keyfile_dict.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', 'project_id': 'project'}
+        mock_open.return_value.__enter__.return_value.read.return_value = json.dumps(json_key)
+        mock_return_cred.return_value = mock_cred
+
+        bq_client = client.get_client(
+            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.from_json_keyfile_dict.assert_called_once_with(json_key,
+                                                                 scopes=BIGQUERY_SCOPE)
+        self.assertTrue(
+            mock_cred.from_json_keyfile_dict.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(json_key['project_id'], bq_client.project_id)
+
 
 class TestQuery(unittest.TestCase):
 
