From 45b672a9a6359ec8c4755d94e63e5ae77a39754b Mon Sep 17 00:00:00 2001
From: Shobhit Singh <shobs@google.com>
Date: Tue, 8 Oct 2024 15:36:29 -0700
Subject: [PATCH] feat: support regional endpoints for more bigquery locations
 (#1061)

* feat: Support bigquery regional endpoints for more locations

* typo plural

* correct the arg documentation
---
 bigframes/_config/bigquery_options.py | 27 +++++++++++++++++++--------
 bigframes/constants.py                |  5 +++--
 bigframes/session/clients.py          | 16 +++++-----------
 3 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/bigframes/_config/bigquery_options.py b/bigframes/_config/bigquery_options.py
index afb0f00a27..2fdd7d6feb 100644
--- a/bigframes/_config/bigquery_options.py
+++ b/bigframes/_config/bigquery_options.py
@@ -233,18 +233,29 @@ def skip_bq_connection_check(self, value: bool):
     def use_regional_endpoints(self) -> bool:
         """Flag to connect to regional API endpoints.
 
-        .. deprecated:: 0.13.0
-            Use of regional endpoints is a feature in Preview and
-            available only in selected regions and projects.
+        .. note::
+            Use of regional endpoints is a feature in Preview and available only
+            in regions "europe-west3", "europe-west9", "europe-west8",
+            "me-central2", "us-east4" and "us-west1".
 
-        Requires that ``location`` is set. For example, to connect to
-        asia-northeast1-bigquery.googleapis.com, specify
-        ``location='asia-northeast1'`` and ``use_regional_endpoints=True``.
+        .. deprecated:: 0.13.0
+            Use of locational endpoints is available only in selected projects.
+
+        Requires that ``location`` is set. For supported regions, for example
+        ``europe-west3``, you need to specify ``location='europe-west3'`` and
+        ``use_regional_endpoints=True``, and then BigQuery DataFrames would
+        connect to the BigQuery endpoint ``bigquery.europe-west3.rep.googleapis.com``.
+        For not supported regions, for example ``asia-northeast1``, when you
+        specify ``location='asia-northeast1'`` and ``use_regional_endpoints=True``,
+        a different endpoint (called locational endpoint, now deprecated, used
+        to provide weaker promise on the request remaining within the location
+        during transit) ``europe-west3-bigquery.googleapis.com`` would be used.
 
         Returns:
             bool:
-              A boolean value, where True indicates that a location is set;
-              otherwise False.
+              A boolean value, where True indicates that regional endpoints
+              would be used for BigQuery and BigQuery storage APIs; otherwise
+              global endpoints would be used.
         """
         return self._use_regional_endpoints
 
diff --git a/bigframes/constants.py b/bigframes/constants.py
index e0c8305079..13636a4484 100644
--- a/bigframes/constants.py
+++ b/bigframes/constants.py
@@ -80,9 +80,10 @@
 # https://cloud.google.com/storage/docs/regional-endpoints
 REP_ENABLED_BIGQUERY_LOCATIONS = frozenset(
     {
-        "me-central2",
-        "europe-west9",
         "europe-west3",
+        "europe-west9",
+        "europe-west8",
+        "me-central2",
         "us-east4",
         "us-west1",
     }
diff --git a/bigframes/session/clients.py b/bigframes/session/clients.py
index 7b53d40f74..04cd1a2ff0 100644
--- a/bigframes/session/clients.py
+++ b/bigframes/session/clients.py
@@ -31,21 +31,13 @@
 import ibis
 import pydata_google_auth
 
+import bigframes.constants
 import bigframes.version
 
 _ENV_DEFAULT_PROJECT = "GOOGLE_CLOUD_PROJECT"
 _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",
-    "europe-west9",
-    "europe-west3",
-    "us-east4",
-    "us-west1",
-}
-
 
 # BigQuery is a REST API, which requires the protocol as part of the URL.
 _BIGQUERY_LOCATIONAL_ENDPOINT = "https://{location}-bigquery.googleapis.com"
@@ -129,7 +121,8 @@ def _create_bigquery_client(self):
                 api_endpoint=(
                     _BIGQUERY_REGIONAL_ENDPOINT
                     if self._location is not None
-                    and self._location.lower() in _REP_SUPPORTED_REGIONS
+                    and self._location.lower()
+                    in bigframes.constants.REP_ENABLED_BIGQUERY_LOCATIONS
                     else _BIGQUERY_LOCATIONAL_ENDPOINT
                 ).format(location=self._location),
             )
@@ -201,7 +194,8 @@ def bqstoragereadclient(self):
                     api_endpoint=(
                         _BIGQUERYSTORAGE_REGIONAL_ENDPOINT
                         if self._location is not None
-                        and self._location.lower() in _REP_SUPPORTED_REGIONS
+                        and self._location.lower()
+                        in bigframes.constants.REP_ENABLED_BIGQUERY_LOCATIONS
                         else _BIGQUERYSTORAGE_LOCATIONAL_ENDPOINT
                     ).format(location=self._location),
                 )