feat: add client_endpoints_override to bq options (#1167)

* feat: add client_endpoints_override to bq options

* fix wording

Author: GarrettWu

bigframes/_config/bigquery_options.py

@@ -91,6 +91,7 @@ def __init__(
         skip_bq_connection_check: bool = False,
         *,
         ordering_mode: Literal["strict", "partial"] = "strict",
+        client_endpoints_override: dict = {},
     ):
         self._credentials = credentials
         self._project = project
@@ -103,6 +104,7 @@ def __init__(
         self._session_started = False
         # Determines the ordering strictness for the session.
         self._ordering_mode = _validate_ordering_mode(ordering_mode)
+        self._client_endpoints_override = client_endpoints_override
 
     @property
     def application_name(self) -> Optional[str]:
@@ -317,3 +319,21 @@ def ordering_mode(self) -> Literal["strict", "partial"]:
     @ordering_mode.setter
     def ordering_mode(self, ordering_mode: Literal["strict", "partial"]) -> None:
         self._ordering_mode = _validate_ordering_mode(ordering_mode)
+
+    @property
+    def client_endpoints_override(self) -> dict:
+        """Option that sets the BQ client endpoints addresses directly as a dict. Possible keys are "bqclient", "bqconnectionclient", "bqstoragereadclient"."""
+        return self._client_endpoints_override
+
+    @client_endpoints_override.setter
+    def client_endpoints_override(self, value: dict):
+        warnings.warn(
+            "This is an advanced configuration option for directly setting endpoints. Incorrect use may lead to unexpected behavior or system instability. Proceed only if you fully understand its implications."
+        )
+
+        if self._session_started and self._client_endpoints_override != value:
+            raise ValueError(
+                SESSION_STARTED_MESSAGE.format(attribute="client_endpoints_override")
+            )
+
+        self._client_endpoints_override = value

bigframes/_config/experiment_options.py

@@ -21,8 +21,8 @@ class ExperimentOptions:
     """
 
     def __init__(self):
-        self._semantic_operators = False
-        self._blob = False
+        self._semantic_operators: bool = False
+        self._blob: bool = False
 
     @property
     def semantic_operators(self) -> bool:

bigframes/pandas/io/api.py

@@ -332,6 +332,7 @@ def _set_default_session_location_if_possible(query):
         credentials=config.options.bigquery.credentials,
         application_name=config.options.bigquery.application_name,
         bq_kms_key_name=config.options.bigquery.kms_key_name,
+        client_endpoints_override=config.options.bigquery.client_endpoints_override,
     )
 
     bqclient = clients_provider.bqclient

bigframes/session/__init__.py

@@ -186,6 +186,7 @@ def __init__(
                 credentials=context.credentials,
                 application_name=context.application_name,
                 bq_kms_key_name=self._bq_kms_key_name,
+                client_endpoints_override=context.client_endpoints_override,
             )
 
         # TODO(shobs): Remove this logic after https://github.com/ibis-project/ibis/issues/8494

bigframes/session/clients.py

@@ -67,6 +67,7 @@ def __init__(
         credentials: Optional[google.auth.credentials.Credentials] = None,
         application_name: Optional[str] = None,
         bq_kms_key_name: Optional[str] = None,
+        client_endpoints_override: dict = {},
     ):
         credentials_project = None
         if credentials is None:
@@ -98,6 +99,7 @@ def __init__(
         self._use_regional_endpoints = use_regional_endpoints
         self._credentials = credentials
         self._bq_kms_key_name = bq_kms_key_name
+        self._client_endpoints_override = client_endpoints_override
 
         # cloud clients initialized for lazy load
         self._bqclient = None
@@ -126,6 +128,11 @@ def _create_bigquery_client(self):
                     else _BIGQUERY_LOCATIONAL_ENDPOINT
                 ).format(location=self._location),
             )
+        if "bqclient" in self._client_endpoints_override:
+            bq_options = google.api_core.client_options.ClientOptions(
+                api_endpoint=self._client_endpoints_override["bqclient"]
+            )
+
         bq_info = google.api_core.client_info.ClientInfo(
             user_agent=self._application_name
         )
@@ -172,6 +179,11 @@ def bqconnectionclient(self):
                         location=self._location
                     )
                 )
+            if "bqconnectionclient" in self._client_endpoints_override:
+                bqconnection_options = google.api_core.client_options.ClientOptions(
+                    api_endpoint=self._client_endpoints_override["bqconnectionclient"]
+                )
+
             bqconnection_info = google.api_core.gapic_v1.client_info.ClientInfo(
                 user_agent=self._application_name
             )
@@ -199,6 +211,11 @@ def bqstoragereadclient(self):
                         else _BIGQUERYSTORAGE_LOCATIONAL_ENDPOINT
                     ).format(location=self._location),
                 )
+
+            if "bqstoragereadclient" in self._client_endpoints_override:
+                bqstorage_options = google.api_core.client_options.ClientOptions(
+                    api_endpoint=self._client_endpoints_override["bqstoragereadclient"]
+                )
             bqstorage_info = google.api_core.gapic_v1.client_info.ClientInfo(
                 user_agent=self._application_name
             )

tests/unit/_config/test_bigquery_options.py

@@ -34,6 +34,7 @@
         ("use_regional_endpoints", False, True),
         ("kms_key_name", "kms/key/name/1", "kms/key/name/2"),
         ("skip_bq_connection_check", False, True),
+        ("client_endpoints_override", {}, {"bqclient": "endpoint_address"}),
     ],
 )
 def test_setter_raises_if_session_started(attribute, original_value, new_value):
@@ -67,6 +68,7 @@ def test_setter_raises_if_session_started(attribute, original_value, new_value):
             "bq_connection",
             "use_regional_endpoints",
             "bq_kms_key_name",
+            "client_endpoints_override",
         ]
     ],
 )
@@ -152,3 +154,10 @@ def set_location_property():
             ),
         ):
             op()
+
+
+def test_client_endpoints_override_set_shows_warning():
+    options = bigquery_options.BigQueryOptions()
+
+    with pytest.warns(UserWarning):
+        options.client_endpoints_override = {"bqclient": "endpoint_address"}