From 0b5c1d597cdec3a05a16fb935595f773c5840bd4 Mon Sep 17 00:00:00 2001
From: shollyman <shollyman@google.com>
Date: Tue, 16 Jan 2024 11:09:40 -0800
Subject: [PATCH] feat: support universe resolution (#1774)

* feat: support universe resolution

This PR wires up consumption of the universe_domain client option for
resolving the endpoint for constructing the BQ client.

Testing universes is not yet something we want to in this repo, so
validation was done out of band.

* formatting and testing

* conditionals for stale core

* formatting

* unused import
---
 google/cloud/bigquery/_helpers.py |  3 +++
 google/cloud/bigquery/client.py   |  9 +++++++++
 tests/unit/test_client.py         | 17 +++++++++++++++++
 3 files changed, 29 insertions(+)

diff --git a/google/cloud/bigquery/_helpers.py b/google/cloud/bigquery/_helpers.py
index 4cf6dddac..905d4aee1 100644
--- a/google/cloud/bigquery/_helpers.py
+++ b/google/cloud/bigquery/_helpers.py
@@ -55,6 +55,9 @@
 _DEFAULT_HOST = "https://bigquery.googleapis.com"
 """Default host for JSON API."""
 
+_DEFAULT_UNIVERSE = "googleapis.com"
+"""Default universe for the JSON API."""
+
 
 def _get_bigquery_host():
     return os.environ.get(BIGQUERY_EMULATOR_HOST, _DEFAULT_HOST)
diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py
index 182319646..b2ea130c4 100644
--- a/google/cloud/bigquery/client.py
+++ b/google/cloud/bigquery/client.py
@@ -78,6 +78,7 @@
 from google.cloud.bigquery._helpers import _verify_job_config_type
 from google.cloud.bigquery._helpers import _get_bigquery_host
 from google.cloud.bigquery._helpers import _DEFAULT_HOST
+from google.cloud.bigquery._helpers import _DEFAULT_UNIVERSE
 from google.cloud.bigquery._job_helpers import make_job_id as _make_job_id
 from google.cloud.bigquery.dataset import Dataset
 from google.cloud.bigquery.dataset import DatasetListItem
@@ -252,6 +253,14 @@ def __init__(
             if client_options.api_endpoint:
                 api_endpoint = client_options.api_endpoint
                 kw_args["api_endpoint"] = api_endpoint
+            elif (
+                hasattr(client_options, "universe_domain")
+                and client_options.universe_domain
+                and client_options.universe_domain is not _DEFAULT_UNIVERSE
+            ):
+                kw_args["api_endpoint"] = _DEFAULT_HOST.replace(
+                    _DEFAULT_UNIVERSE, client_options.universe_domain
+                )
 
         self._connection = Connection(self, **kw_args)
         self._location = location
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
index ad22e0ddb..56bdbad5e 100644
--- a/tests/unit/test_client.py
+++ b/tests/unit/test_client.py
@@ -201,6 +201,23 @@ def test_ctor_w_client_options_object(self):
             client._connection.API_BASE_URL, "https://www.foo-googleapis.com"
         )
 
+    @pytest.mark.skipif(
+        packaging.version.parse(getattr(google.api_core, "__version__", "0.0.0"))
+        < packaging.version.Version("2.15.0"),
+        reason="universe_domain not supported with google-api-core < 2.15.0",
+    )
+    def test_ctor_w_client_options_universe(self):
+        creds = _make_credentials()
+        http = object()
+        client_options = {"universe_domain": "foo.com"}
+        client = self._make_one(
+            project=self.PROJECT,
+            credentials=creds,
+            _http=http,
+            client_options=client_options,
+        )
+        self.assertEqual(client._connection.API_BASE_URL, "https://bigquery.foo.com")
+
     def test_ctor_w_location(self):
         from google.cloud.bigquery._http import Connection