Merge pull request #81 from andrewgardners/master

Add service_url parameter to get_client and _get_bq_service

Author: tylertreat

bigquery/client.py

@@ -8,7 +8,7 @@
 
 import httplib2
 import six
-from apiclient.discovery import build
+from apiclient.discovery import build, DISCOVERY_URI
 from apiclient.errors import HttpError
 
 from bigquery.errors import (BigQueryTimeoutException, JobExecutingException,
@@ -43,7 +43,8 @@
 JOB_DESTINATION_FORMAT_CSV = JOB_FORMAT_CSV
 
 
-def get_client(project_id, credentials=None, service_account=None,
+def get_client(project_id, credentials=None,
+               service_url=None, service_account=None,
                private_key=None, private_key_file=None,
                json_key=None, json_key_file=None,
                readonly=True, swallow_results=True):
@@ -55,6 +56,12 @@ def get_client(project_id, credentials=None, service_account=None,
         project_id: the BigQuery project id.
         credentials: an AssertionCredentials instance to authenticate requests
                      to BigQuery.
+        service_url: a URI string template pointing to the location of
+                     Google's API discovery service. Requires two parameters
+                     {api} and {apiVersion} that when filled in produce an
+                     absolute URI to the discovery document for that service.
+                     If not set then the default googleapiclient disovery URI
+                     is used.
         service_account: the Google API service account name.
         private_key: the private key associated with the service account in
                      PKCS12 or PEM format.
@@ -77,6 +84,9 @@ def get_client(project_id, credentials=None, service_account=None,
         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
+
     if private_key_file:
         with open(private_key_file, 'rb') as key_file:
             private_key = key_file.read()
@@ -90,14 +100,15 @@ def get_client(project_id, credentials=None, service_account=None,
         private_key = json_key['private_key']
 
     bq_service = _get_bq_service(credentials=credentials,
+                                 service_url=service_url,
                                  service_account=service_account,
                                  private_key=private_key,
                                  readonly=readonly)
 
     return BigQueryClient(bq_service, project_id, swallow_results)
 
 
-def _get_bq_service(credentials=None, service_account=None, private_key=None,
+def _get_bq_service(credentials=None, service_url=None, service_account=None, private_key=None,
                     readonly=True):
     """Construct an authorized BigQuery service object."""
 
@@ -110,7 +121,7 @@ def _get_bq_service(credentials=None, service_account=None, private_key=None,
 
     http = httplib2.Http()
     http = credentials.authorize(http)
-    service = build('bigquery', 'v2', http=http)
+    service = build('bigquery', 'v2', http=http, discoveryServiceUrl=service_url)
 
     return service
 

bigquery/tests/test_client.py

@@ -50,6 +50,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_bq = mock.Mock()
         mock_build.return_value = mock_bq
@@ -59,14 +60,16 @@ def test_initialize_readonly(self, mock_build, mock_return_cred):
         mock_return_cred.return_value = mock_cred
 
         bq_client = client.get_client(
-            project_id, service_account=service_account, private_key=key,
+            project_id, service_url=mock_service_url,
+            service_account=service_account, private_key=key,
             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_build.assert_called_once_with('bigquery', 'v2', http=mock_http)
+        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)
 
@@ -80,6 +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_bq = mock.Mock()
         mock_build.return_value = mock_bq
@@ -89,14 +93,16 @@ def test_initialize_read_write(self, mock_build, mock_return_cred):
         mock_return_cred.return_value = mock_cred
 
         bq_client = client.get_client(
-            project_id, service_account=service_account, private_key=key,
+            project_id, service_url=mock_service_url,
+            service_account=service_account, private_key=key,
             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_build.assert_called_once_with('bigquery', 'v2', http=mock_http)
+        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)
 
@@ -112,6 +118,7 @@ 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_bq = mock.Mock()
         mock_build.return_value = mock_bq
@@ -123,15 +130,17 @@ def test_initialize_key_file(self, mock_open, mock_build,
         mock_return_cred.return_value = mock_cred
 
         bq_client = client.get_client(
-            project_id, service_account=service_account,
+            project_id, service_url=mock_service_url,
+            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_build.assert_called_once_with('bigquery', 'v2', http=mock_http)
+        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)
 
@@ -147,6 +156,7 @@ def test_initialize_json_key_file(self, mock_open, 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_bq = mock.Mock()
         mock_build.return_value = mock_bq
@@ -156,13 +166,14 @@ def test_initialize_json_key_file(self, mock_open, mock_build, mock_return_cred)
         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)
+        bq_client = client.get_client(
+            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)
+        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)