googleapis · Linchin · Mar 5, 2024 · Feb 29, 2024 · Feb 29, 2024 · Feb 29, 2024
@@ -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)

@@ -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

@@ -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}"
 

@@ -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):