diff --git a/bigframes/_config/bigquery_options.py b/bigframes/_config/bigquery_options.py index 90f03c83cd..9962d4286e 100644 --- a/bigframes/_config/bigquery_options.py +++ b/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 diff --git a/bigframes/_config/experiment_options.py b/bigframes/_config/experiment_options.py index b0e1cbdd18..6b79dcf748 100644 --- a/bigframes/_config/experiment_options.py +++ b/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: diff --git a/bigframes/pandas/io/api.py b/bigframes/pandas/io/api.py index 4e08b3ef5e..19ff5e0a05 100644 --- a/bigframes/pandas/io/api.py +++ b/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 diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 21789835e7..26ec410637 100644 --- a/bigframes/session/__init__.py +++ b/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 diff --git a/bigframes/session/clients.py b/bigframes/session/clients.py index 04cd1a2ff0..ffc5fb05d9 100644 --- a/bigframes/session/clients.py +++ b/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 ) diff --git a/tests/unit/_config/test_bigquery_options.py b/tests/unit/_config/test_bigquery_options.py index 44cf024219..784e92af40 100644 --- a/tests/unit/_config/test_bigquery_options.py +++ b/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"}