From 53c2cbf98d2961f553747514de273bcd5c117f0e Mon Sep 17 00:00:00 2001 From: shollyman Date: Tue, 5 Mar 2024 09:47:31 -0800 Subject: [PATCH] fix: augment universe_domain handling (#1837) * 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 --- google/cloud/bigquery/_helpers.py | 56 ++++++++++++++++++++++ google/cloud/bigquery/client.py | 21 ++++---- tests/unit/helpers.py | 14 ++++++ tests/unit/test__helpers.py | 80 ++++++++++++++++++++++++++++++- 4 files changed, 162 insertions(+), 9 deletions(-) diff --git a/google/cloud/bigquery/_helpers.py b/google/cloud/bigquery/_helpers.py index 905d4aee1..ec4ac9970 100644 --- 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 index a871dc003..cb4daa897 100644 --- 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 index 67aeaca35..bc92c0df6 100644 --- 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 index 87ab46669..019d2e7bd 100644 --- 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):