Visitar URL original
fix(bigquery): pass universe_domain to the BigQuery Storage client (#… · googleapis/google-cloud-python@6ab4063 · GitHub
Skip to content

Commit 6ab4063

Browse files
authored
fix(bigquery): pass universe_domain to the BigQuery Storage client (#18783)
Points the BigQuery Storage client at the same universe as the BigQuery client. Before, `_ensure_bqstorage_client()` created it without `client_options`, so in a non-default universe `to_arrow()`, `to_dataframe()` and DB-API reads went to `bigquerystorage.googleapis.com`. - Stores the resolved universe on `Client`. - Passes it to the new Storage client when the caller gives no `client_options`. - Leaves caller-provided `client_options` unchanged. - Adds unit tests for the default universe, a custom universe, and caller-provided options.
1 parent 8166325 commit 6ab4063

2 files changed

Lines changed: 61 additions & 1 deletion

File tree

‎packages/google-cloud-bigquery/google/cloud/bigquery/client.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -296,6 +296,7 @@ def __init__(
296296
_validate_universe(client_universe, self._credentials)
297297

298298
self._connection = Connection(self, **kw_args)
299+
self._universe_domain = _get_client_universe(client_options)
299300
self._location = location
300301
self._default_load_job_config = copy.deepcopy(default_load_job_config)
301302
self.default_job_creation_mode = default_job_creation_mode
@@ -609,7 +610,8 @@ def _ensure_bqstorage_client(
609610
instance is created and returned.
610611
client_options:
611612
Custom options used with a new BigQuery Storage client instance
612-
if one is created.
613+
if one is created. If ``None``, the new client uses this
614+
client's universe domain.
613615
client_info:
614616
The client info used with a new BigQuery Storage client
615617
instance if one is created.
@@ -676,6 +678,11 @@ def _ensure_bqstorage_client(
676678
)
677679

678680
if bqstorage_client is None:
681+
if client_options is None and self._universe_domain != _DEFAULT_UNIVERSE:
682+
# Without this, the BigQuery Storage client uses googleapis.com.
683+
client_options = google.api_core.client_options.ClientOptions(
684+
universe_domain=self._universe_domain
685+
)
679686
bqstorage_client = bigquery_storage.BigQueryReadClient(
680687
credentials=self._credentials,
681688
client_options=client_options,

‎packages/google-cloud-bigquery/tests/unit/test_client.py‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -844,6 +844,59 @@ def test_ensure_bqstorage_client_creating_new_instance(self):
844844
f"legacy-gcb/{version.__version__}", kwargs["client_info"].user_agent
845845
)
846846

847+
def test_ensure_bqstorage_client_default_universe(self):
848+
bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage")
849+
850+
mock_client = mock.create_autospec(bigquery_storage.BigQueryReadClient)
851+
creds = _make_credentials()
852+
client = self._make_one(project=self.PROJECT, credentials=creds)
853+
854+
with mock.patch(
855+
"google.cloud.bigquery_storage.BigQueryReadClient", mock_client
856+
):
857+
client._ensure_bqstorage_client()
858+
859+
_, kwargs = mock_client.call_args
860+
self.assertIsNone(kwargs["client_options"])
861+
862+
def test_ensure_bqstorage_client_uses_client_universe(self):
863+
bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage")
864+
865+
mock_client = mock.create_autospec(bigquery_storage.BigQueryReadClient)
866+
creds = _make_credentials()
867+
client = self._make_one(
868+
project=self.PROJECT,
869+
credentials=creds,
870+
client_options={"universe_domain": "foo.com"},
871+
)
872+
873+
with mock.patch(
874+
"google.cloud.bigquery_storage.BigQueryReadClient", mock_client
875+
):
876+
client._ensure_bqstorage_client()
877+
878+
_, kwargs = mock_client.call_args
879+
self.assertEqual(kwargs["client_options"].universe_domain, "foo.com")
880+
881+
def test_ensure_bqstorage_client_keeps_given_client_options(self):
882+
bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage")
883+
884+
mock_client = mock.create_autospec(bigquery_storage.BigQueryReadClient)
885+
creds = _make_credentials()
886+
client = self._make_one(
887+
project=self.PROJECT,
888+
credentials=creds,
889+
client_options={"universe_domain": "foo.com"},
890+
)
891+
892+
with mock.patch(
893+
"google.cloud.bigquery_storage.BigQueryReadClient", mock_client
894+
):
895+
client._ensure_bqstorage_client(client_options=mock.sentinel.client_options)
896+
897+
_, kwargs = mock_client.call_args
898+
self.assertIs(kwargs["client_options"], mock.sentinel.client_options)
899+
847900
def test_ensure_bqstorage_client_pandas_gbq_installed(self):
848901
bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage")
849902
import sys

0 commit comments

Comments
 (0)