Add service_url parameter to get_client and _get_bq_service

andrewgardners · andrewgardners · commit d45009a074f3 · 2016-02-11T15:20:59.000Z
diff --git a/bigquery/client.py b/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
 
