feat: add mtls feature

google/cloud/bigquery/_http.py

@@ -26,13 +26,18 @@ class Connection(_http.JSONConnection):
         client (google.cloud.bigquery.client.Client): The client that owns the current connection.
 
         client_info (Optional[google.api_core.client_info.ClientInfo]): Instance used to generate user agent.
+
+        api_endpoint (str): The api_endpoint to use. If None, the library will decide what endpoint to use.
     """
 
     DEFAULT_API_ENDPOINT = "https://bigquery.googleapis.com"
+    DEFAULT_API_MTLS_ENDPOINT = "https://bigquery.mtls.googleapis.com"
 
-    def __init__(self, client, client_info=None, api_endpoint=DEFAULT_API_ENDPOINT):
+    def __init__(self, client, client_info=None, api_endpoint=None):
         super(Connection, self).__init__(client, client_info)
-        self.API_BASE_URL = api_endpoint
+        self.API_BASE_URL = api_endpoint or self.DEFAULT_API_ENDPOINT
+        self.API_BASE_MTLS_URL = self.DEFAULT_API_MTLS_ENDPOINT
+        self.ALLOW_AUTO_SWITCH_TO_MTLS_URL = api_endpoint is None
         self._client_info.gapic_version = __version__
         self._client_info.client_library_version = __version__
 

google/cloud/bigquery/client.py

@@ -78,10 +78,7 @@
 _DEFAULT_CHUNKSIZE = 1048576  # 1024 * 1024 B = 1 MB
 _MAX_MULTIPART_SIZE = 5 * 1024 * 1024
 _DEFAULT_NUM_RETRIES = 6
-_BASE_UPLOAD_TEMPLATE = (
-    "https://bigquery.googleapis.com/upload/bigquery/v2/projects/"
-    "{project}/jobs?uploadType="
-)
+_BASE_UPLOAD_TEMPLATE = "{host}/upload/bigquery/v2/projects/{project}/jobs?uploadType="
 _MULTIPART_URL_TEMPLATE = _BASE_UPLOAD_TEMPLATE + "multipart"
 _RESUMABLE_URL_TEMPLATE = _BASE_UPLOAD_TEMPLATE + "resumable"
 _GENERIC_CONTENT_TYPE = "*/*"
@@ -2547,7 +2544,12 @@ def _initiate_resumable_upload(
 
         if project is None:
             project = self.project
-        upload_url = _RESUMABLE_URL_TEMPLATE.format(project=project)
+        hostname = (
+            self._connection.API_BASE_URL
+            if not hasattr(self._connection, "get_api_base_url_for_mtls")
+            else self._connection.get_api_base_url_for_mtls()
+        )
+        upload_url = _RESUMABLE_URL_TEMPLATE.format(host=hostname, project=project)
 
         # TODO: modify ResumableUpload to take a retry.Retry object
         # that it can use for the initial RPC.
@@ -2616,7 +2618,12 @@ def _do_multipart_upload(
         if project is None:
             project = self.project
 
-        upload_url = _MULTIPART_URL_TEMPLATE.format(project=project)
+        hostname = (
+            self._connection.API_BASE_URL
+            if not hasattr(self._connection, "get_api_base_url_for_mtls")
+            else self._connection.get_api_base_url_for_mtls()
+        )
+        upload_url = _MULTIPART_URL_TEMPLATE.format(host=hostname, project=project)
         upload = MultipartUpload(upload_url, headers=headers)
 
         if num_retries is not None:

tests/system/test_client.py

@@ -28,6 +28,7 @@
 import uuid
 
 import psutil
+import pytest
 import pytz
 import pkg_resources
 
@@ -132,6 +133,8 @@
 else:
     PYARROW_INSTALLED_VERSION = None
 
+MTLS_TESTING = os.getenv("GOOGLE_API_USE_CLIENT_CERTIFICATE") == "true"
+
 
 def _has_rows(result):
     return len(result) > 0
@@ -2651,6 +2654,9 @@ def test_insert_rows_nested_nested_dictionary(self):
         expected_rows = [("Some value", record)]
         self.assertEqual(row_tuples, expected_rows)
 
+    @pytest.mark.skipif(
+        MTLS_TESTING, reason="mTLS testing has no permission to the max-value.js file"
+    )
     def test_create_routine(self):
         routine_name = "test_routine"
         dataset = self.temp_dataset(_make_dataset_id("create_routine"))

tests/unit/helpers.py

@@ -21,6 +21,7 @@ def make_connection(*responses):
     mock_conn = mock.create_autospec(google.cloud.bigquery._http.Connection)
     mock_conn.user_agent = "testing 1.2.3"
     mock_conn.api_request.side_effect = list(responses) + [NotFound("miss")]
+    mock_conn.API_BASE_URL = "https://bigquery.googleapis.com"
     return mock_conn