fix: augment universe_domain handling (#1837)

shollyman · Linchin · web-flow · commit 53c2cbf98d29 · 2024-03-05T09:47:31.000-08:00
* fix: augment universe_domain handling

This PR revisits the universe resolution for the BQ client, and handles
new requirements like env-based specification and validation.

* lint

* skipif core too old

* deps

* add import

* no-cover in test helper

* lint

* ignore google.auth typing

* capitalization

* change to raise in test code

* reviewer feedback

* var fix

---------

Co-authored-by: Lingqing Gan &lt;lingqing.gan@gmail.com&gt;
diff --git a/google/cloud/bigquery/_helpers.py b/google/cloud/bigquery/_helpers.py
@@ -30,6 +30,8 @@
 from google.cloud._helpers import _RFC3339_MICROS
 from google.cloud._helpers import _RFC3339_NO_FRACTION
 from google.cloud._helpers import _to_bytes
+from google.auth import credentials as ga_credentials  # type: ignore
+from google.api_core import client_options as client_options_lib
 
 _RFC3339_MICROS_NO_ZULU = "%Y-%m-%dT%H:%M:%S.%f"
 _TIMEONLY_WO_MICROS = "%H:%M:%S"
@@ -55,9 +57,63 @@
 _DEFAULT_HOST = "https://bigquery.googleapis.com"
 """Default host for JSON API."""
 
+_DEFAULT_HOST_TEMPLATE = "https://bigquery.{UNIVERSE_DOMAIN}"
+""" Templatized endpoint format. """
+
 _DEFAULT_UNIVERSE = "googleapis.com"
 """Default universe for the JSON API."""
 
+_UNIVERSE_DOMAIN_ENV = "GOOGLE_CLOUD_UNIVERSE_DOMAIN"
+"""Environment variable for setting universe domain."""
+
+
+def _get_client_universe(
+    client_options: Optional[Union[client_options_lib.ClientOptions, dict]]
+) -> str:
+    """Retrieves the specified universe setting.
+
+    Args:
+        client_options: specified client options.
+    Returns:
+        str: resolved universe setting.
+
+    """
+    if isinstance(client_options, dict):
+        client_options = client_options_lib.from_dict(client_options)
+    universe = _DEFAULT_UNIVERSE
+    if hasattr(client_options, "universe_domain"):
+        options_universe = getattr(client_options, "universe_domain")
+        if options_universe is not None and len(options_universe) > 0:
+            universe = options_universe
+    else:
+        env_universe = os.getenv(_UNIVERSE_DOMAIN_ENV)
+        if isinstance(env_universe, str) and len(env_universe) > 0:
+            universe = env_universe
+    return universe
+
+
+def _validate_universe(client_universe: str, credentials: ga_credentials.Credentials):
+    """Validates that client provided universe and universe embedded in credentials match.
+
+    Args:
+        client_universe (str): The universe domain configured via the client options.
+        credentials (ga_credentials.Credentials): The credentials being used in the client.
+
+    Raises:
+        ValueError: when client_universe does not match the universe in credentials.
+    """
+    if hasattr(credentials, "universe_domain"):
+        cred_universe = getattr(credentials, "universe_domain")
+        if isinstance(cred_universe, str):
+            if client_universe != cred_universe:
+                raise ValueError(
+                    "The configured universe domain "
+                    f"({client_universe}) does not match the universe domain "
+                    f"found in the credentials ({cred_universe}). "
+                    "If you haven't configured the universe domain explicitly, "
+                    f"`{_DEFAULT_UNIVERSE}` is the default."
+                )
+
 
 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
@@ -78,7 +78,10 @@
 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_HOST_TEMPLATE
 from google.cloud.bigquery._helpers import _DEFAULT_UNIVERSE
+from google.cloud.bigquery._helpers import _validate_universe
+from google.cloud.bigquery._helpers import _get_client_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
@@ -245,6 +248,7 @@ def __init__(
         kw_args = {"client_info": client_info}
         bq_host = _get_bigquery_host()
         kw_args["api_endpoint"] = bq_host if bq_host != _DEFAULT_HOST else None
+        client_universe = None
         if client_options:
             if isinstance(client_options, dict):
                 client_options = google.api_core.client_options.from_dict(
@@ -253,14 +257,15 @@ 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
-                )
+            else:
+                client_universe = _get_client_universe(client_options)
+                if client_universe != _DEFAULT_UNIVERSE:
+                    kw_args["api_endpoint"] = _DEFAULT_HOST_TEMPLATE.replace(
+                        "{UNIVERSE_DOMAIN}", client_universe
+                    )
+        # Ensure credentials and universe are not in conflict.
+        if hasattr(self, "_credentials") and client_universe is not None:
+            _validate_universe(client_universe, self._credentials)
 
         self._connection = Connection(self, **kw_args)
         self._location = location
diff --git a/tests/unit/helpers.py b/tests/unit/helpers.py
@@ -43,6 +43,20 @@ def make_client(project="PROJECT", **kw):
     return google.cloud.bigquery.client.Client(project, credentials, **kw)
 
 
+def make_creds(creds_universe: None):
+    from google.auth import credentials
+
+    class TestingCreds(credentials.Credentials):
+        def refresh(self, request):  # pragma: NO COVER
+            raise NotImplementedError
+
+        @property
+        def universe_domain(self):
+            return creds_universe
+
+    return TestingCreds()
+
+
 def make_dataset_reference_string(project, ds_id):
     return f"{project}.{ds_id}"
 
diff --git a/tests/unit/test__helpers.py b/tests/unit/test__helpers.py
@@ -17,8 +17,86 @@
 import decimal
 import json
 import unittest
-
+import os
 import mock
+import pytest
+import packaging
+import google.api_core
+
+
+@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",
+)
+class Test_get_client_universe(unittest.TestCase):
+    def test_with_none(self):
+        from google.cloud.bigquery._helpers import _get_client_universe
+
+        self.assertEqual("googleapis.com", _get_client_universe(None))
+
+    def test_with_dict(self):
+        from google.cloud.bigquery._helpers import _get_client_universe
+
+        options = {"universe_domain": "foo.com"}
+        self.assertEqual("foo.com", _get_client_universe(options))
+
+    def test_with_dict_empty(self):
+        from google.cloud.bigquery._helpers import _get_client_universe
+
+        options = {"universe_domain": ""}
+        self.assertEqual("googleapis.com", _get_client_universe(options))
+
+    def test_with_client_options(self):
+        from google.cloud.bigquery._helpers import _get_client_universe
+        from google.api_core import client_options
+
+        options = client_options.from_dict({"universe_domain": "foo.com"})
+        self.assertEqual("foo.com", _get_client_universe(options))
+
+    @mock.patch.dict(os.environ, {"GOOGLE_CLOUD_UNIVERSE_DOMAIN": "foo.com"})
+    def test_with_environ(self):
+        from google.cloud.bigquery._helpers import _get_client_universe
+
+        self.assertEqual("foo.com", _get_client_universe(None))
+
+    @mock.patch.dict(os.environ, {"GOOGLE_CLOUD_UNIVERSE_DOMAIN": ""})
+    def test_with_environ_empty(self):
+        from google.cloud.bigquery._helpers import _get_client_universe
+
+        self.assertEqual("googleapis.com", _get_client_universe(None))
+
+
+class Test_validate_universe(unittest.TestCase):
+    def test_with_none(self):
+        from google.cloud.bigquery._helpers import _validate_universe
+
+        # should not raise
+        _validate_universe("googleapis.com", None)
+
+    def test_with_no_universe_creds(self):
+        from google.cloud.bigquery._helpers import _validate_universe
+        from .helpers import make_creds
+
+        creds = make_creds(None)
+        # should not raise
+        _validate_universe("googleapis.com", creds)
+
+    def test_with_matched_universe_creds(self):
+        from google.cloud.bigquery._helpers import _validate_universe
+        from .helpers import make_creds
+
+        creds = make_creds("googleapis.com")
+        # should not raise
+        _validate_universe("googleapis.com", creds)
+
+    def test_with_mismatched_universe_creds(self):
+        from google.cloud.bigquery._helpers import _validate_universe
+        from .helpers import make_creds
+
+        creds = make_creds("foo.com")
+        with self.assertRaises(ValueError):
+            _validate_universe("googleapis.com", creds)
 
 
 class Test_not_null(unittest.TestCase):
