Skip to content

Commit

Permalink
feat: Enable regional endpoints for me-central2 (#386)
Browse files Browse the repository at this point in the history
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a [bug/issue](https://togithub.com/googleapis/python-bigquery-dataframes/issues/new/choose) before writing your code!  That way we can discuss the change, evaluate designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)

Fixes internal issue 312304785 🦕
  • Loading branch information
shobsi authored Feb 24, 2024
1 parent 1040dff commit 469674d
Show file tree
Hide file tree
Showing 3 changed files with 27 additions and 17 deletions.
10 changes: 4 additions & 6 deletions bigframes/_config/bigquery_options.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
2 changes: 1 addition & 1 deletion bigframes/session/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -152,7 +152,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.",
Expand Down
32 changes: 22 additions & 10 deletions bigframes/session/clients.py
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down Expand Up @@ -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_REGIONAL_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
Expand All @@ -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
)
)
Expand All @@ -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_REGIONAL_ENDPOINT
if self._location.lower() in _REP_SUPPORTED_REGIONS
else _BIGQUERYSTORAGE_LOCATIONAL_ENDPOINT
).format(location=self._location),
)
bqstorage_info = google.api_core.gapic_v1.client_info.ClientInfo(
user_agent=self._application_name
Expand Down

0 comments on commit 469674d

Please sign in to comment.