Skip to content

Commit

Permalink
feat: support universe resolution (#1774)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
shollyman authored Jan 16, 2024
1 parent cf920f4 commit 0b5c1d5
Show file tree
Hide file tree
Showing 3 changed files with 29 additions and 0 deletions.
3 changes: 3 additions & 0 deletions google/cloud/bigquery/_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
9 changes: 9 additions & 0 deletions google/cloud/bigquery/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
17 changes: 17 additions & 0 deletions tests/unit/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down

0 comments on commit 0b5c1d5

Please sign in to comment.