From fa0511ef9e6ebdc2ea06d0a25e01b51cde5426ae Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Tue, 20 Feb 2024 23:14:01 +0000 Subject: [PATCH 1/2] feat: Enable regional endpoints for me-central2 --- bigframes/_config/bigquery_options.py | 10 ++++----- bigframes/session/__init__.py | 2 +- bigframes/session/clients.py | 32 ++++++++++++++++++--------- 3 files changed, 27 insertions(+), 17 deletions(-) diff --git a/bigframes/_config/bigquery_options.py b/bigframes/_config/bigquery_options.py index 2875a11de3..74b83429d0 100644 --- a/bigframes/_config/bigquery_options.py +++ b/bigframes/_config/bigquery_options.py @@ -125,9 +125,8 @@ def use_regional_endpoints(self) -> bool: """Flag to connect to regional API endpoints. .. deprecated:: 0.13.0 - BigQuery regional endpoints is a feature in preview and - available only to selected projects. - Enable it only if your project has regional endpoints access. + Use of regional endpoints is a feature in preview and + available only in selected regions and projects. Requires ``location`` to also be set. For example, set ``location='asia-northeast1'`` and ``use_regional_endpoints=True`` to @@ -144,9 +143,8 @@ def use_regional_endpoints(self, value: bool): if value: warnings.warn( - "BigQuery regional endpoints is a feature in preview and " - "available only to selected projects. " - "Enable it only if your project has regional endpoints access." + "Use of regional endpoints is a feature in preview and " + "available only in selected regions and projects. " ) self._use_regional_endpoints = value diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index df0cd6e947..f100effccf 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -148,7 +148,7 @@ def __init__( context = bigquery_options.BigQueryOptions() # TODO(swast): Get location from the environment. - if context is None or context.location is None: + if context.location is None: self._location = "US" warnings.warn( f"No explicit location is set, so using location {self._location} for the session.", diff --git a/bigframes/session/clients.py b/bigframes/session/clients.py index e33413002f..e731f2b13c 100644 --- a/bigframes/session/clients.py +++ b/bigframes/session/clients.py @@ -37,13 +37,21 @@ _APPLICATION_NAME = f"bigframes/{bigframes.version.__version__} ibis/{ibis.__version__}" _SCOPES = ["https://www.googleapis.com/auth/cloud-platform"] +# Regions for which Regional Endpoints (REPs) are supported +_REP_SUPPORTED_REGIONS = {"me-central2"} + + # BigQuery is a REST API, which requires the protocol as part of the URL. -_BIGQUERY_REGIONAL_ENDPOINT = "https://{location}-bigquery.googleapis.com" +_BIGQUERY_LOCATIONAL_ENDPOINT = "https://{location}-bigquery.googleapis.com" +_BIGQUERY_REGIONAL_ENDPOINT = "https://bigquery.{location}.rep.googleapis.com" # BigQuery Connection and Storage are gRPC APIs, which don't support the # https:// protocol in the API endpoint URL. -_BIGQUERYCONNECTION_REGIONAL_ENDPOINT = "{location}-bigqueryconnection.googleapis.com" -_BIGQUERYSTORAGE_REGIONAL_ENDPOINT = "{location}-bigquerystorage.googleapis.com" +_BIGQUERYCONNECTION_LOCATIONAL_ENDPOINT = "{location}-bigqueryconnection.googleapis.com" +_BIGQUERYSTORAGE_LOCATIONAL_ENDPOINT = "{location}-bigquerystorage.googleapis.com" +_BIGQUERYSTORAGE_REGIONAL_ENDPOINT = ( + "https://bigquerystorage.{location}.rep.googleapis.com" +) def _get_default_credentials_with_project(): @@ -104,9 +112,11 @@ def bqclient(self): bq_options = None if self._use_regional_endpoints: bq_options = google.api_core.client_options.ClientOptions( - api_endpoint=_BIGQUERY_REGIONAL_ENDPOINT.format( - location=self._location - ), + api_endpoint=( + _BIGQUERY_LOCATIONAL_ENDPOINT + if self._location.lower() in _REP_SUPPORTED_REGIONS + else _BIGQUERY_LOCATIONAL_ENDPOINT + ).format(location=self._location), ) bq_info = google.api_core.client_info.ClientInfo( user_agent=self._application_name @@ -127,7 +137,7 @@ def bqconnectionclient(self): bqconnection_options = None if self._use_regional_endpoints: bqconnection_options = google.api_core.client_options.ClientOptions( - api_endpoint=_BIGQUERYCONNECTION_REGIONAL_ENDPOINT.format( + api_endpoint=_BIGQUERYCONNECTION_LOCATIONAL_ENDPOINT.format( location=self._location ) ) @@ -150,9 +160,11 @@ def bqstoragereadclient(self): bqstorage_options = None if self._use_regional_endpoints: bqstorage_options = google.api_core.client_options.ClientOptions( - api_endpoint=_BIGQUERYSTORAGE_REGIONAL_ENDPOINT.format( - location=self._location - ) + api_endpoint=( + _BIGQUERYSTORAGE_LOCATIONAL_ENDPOINT + if self._location.lower() in _REP_SUPPORTED_REGIONS + else _BIGQUERYSTORAGE_REGIONAL_ENDPOINT + ).format(location=self._location), ) bqstorage_info = google.api_core.gapic_v1.client_info.ClientInfo( user_agent=self._application_name From 635ad8ed36d6648a74b88fd7e8ecc7a4dfd690e6 Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Tue, 20 Feb 2024 23:16:44 +0000 Subject: [PATCH 2/2] use correct endpoints --- bigframes/session/clients.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/bigframes/session/clients.py b/bigframes/session/clients.py index e731f2b13c..627c9258a6 100644 --- a/bigframes/session/clients.py +++ b/bigframes/session/clients.py @@ -113,7 +113,7 @@ def bqclient(self): if self._use_regional_endpoints: bq_options = google.api_core.client_options.ClientOptions( api_endpoint=( - _BIGQUERY_LOCATIONAL_ENDPOINT + _BIGQUERY_REGIONAL_ENDPOINT if self._location.lower() in _REP_SUPPORTED_REGIONS else _BIGQUERY_LOCATIONAL_ENDPOINT ).format(location=self._location), @@ -161,9 +161,9 @@ def bqstoragereadclient(self): if self._use_regional_endpoints: bqstorage_options = google.api_core.client_options.ClientOptions( api_endpoint=( - _BIGQUERYSTORAGE_LOCATIONAL_ENDPOINT + _BIGQUERYSTORAGE_REGIONAL_ENDPOINT if self._location.lower() in _REP_SUPPORTED_REGIONS - else _BIGQUERYSTORAGE_REGIONAL_ENDPOINT + else _BIGQUERYSTORAGE_LOCATIONAL_ENDPOINT ).format(location=self._location), ) bqstorage_info = google.api_core.gapic_v1.client_info.ClientInfo(