Skip to content

Commit ecbf77d

Browse files
authored
chore: show FutureWarning for LEP endpoints usage with use_regional_endpoints=True (#1432)
* chore: show `DeprecationWarning` for LEP endpoints usage * fix mypy * use custom bigframes warning as the default DeprecationWarning is disabled by default * Revert "use custom bigframes warning as the default DeprecationWarning is disabled by default" This reverts commit faa3050. * revert custom warning, use FutureWarning instead for LEP deprecation
1 parent 7c37c7d commit ecbf77d

File tree

3 files changed

+93
-32
lines changed

3 files changed

+93
-32
lines changed

bigframes/constants.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
"""
1919

2020
import datetime
21+
import textwrap
2122

2223
DEFAULT_EXPIRATION = datetime.timedelta(days=7)
2324

@@ -100,6 +101,25 @@
100101
ALL_BIGQUERY_LOCATIONS - REP_ENABLED_BIGQUERY_LOCATIONS
101102
)
102103

104+
LEP_DEPRECATION_WARNING_MESSAGE = textwrap.dedent(
105+
"""
106+
Support for regional endpoints is not yet available in the location
107+
{location} for BigQuery and BigQuery Storage APIs. For the supported
108+
locations and APIs see https://cloud.google.com/bigquery/docs/regional-endpoints.
109+
For other locations and APIs, currently an older, now deprecated locational
110+
endpoints are being used, which requires your project to be allowlisted. In
111+
future version 2.0 onwards the locational endpoints will no longer be
112+
supported automatically when you enable regional endpoints. However, if you
113+
still need them, you will be able to override the endpoints directly by
114+
doing the following:
115+
bigframes.pandas.options.bigquery.client_endpoints_override = {{
116+
"bqclient": "https://{location}-bigquery.googleapis.com",
117+
"bqconnectionclient": "{location}-bigqueryconnection.googleapis.com",
118+
"bqstoragereadclient": "{location}-bigquerystorage.googleapis.com"
119+
}}
120+
"""
121+
).strip()
122+
103123
# BigQuery default is 10000, leave 100 for overhead
104124
MAX_COLUMNS = 9900
105125

bigframes/session/clients.py

Lines changed: 44 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import os
1818
import typing
1919
from typing import Optional
20+
import warnings
2021

2122
import google.api_core.client_info
2223
import google.api_core.client_options
@@ -94,8 +95,22 @@ def __init__(
9495
else _APPLICATION_NAME
9596
)
9697
self._project = project
98+
99+
if (
100+
use_regional_endpoints
101+
and location is not None
102+
and location.lower()
103+
not in bigframes.constants.REP_ENABLED_BIGQUERY_LOCATIONS
104+
):
105+
warnings.warn(
106+
bigframes.constants.LEP_DEPRECATION_WARNING_MESSAGE.format(
107+
location=location
108+
),
109+
category=FutureWarning,
110+
)
97111
self._location = location
98112
self._use_regional_endpoints = use_regional_endpoints
113+
99114
self._credentials = credentials
100115
self._bq_kms_key_name = bq_kms_key_name
101116
self._client_endpoints_override = client_endpoints_override
@@ -117,20 +132,22 @@ def __init__(
117132

118133
def _create_bigquery_client(self):
119134
bq_options = None
120-
if self._use_regional_endpoints:
121-
bq_options = google.api_core.client_options.ClientOptions(
122-
api_endpoint=(
123-
_BIGQUERY_REGIONAL_ENDPOINT
124-
if self._location is not None
125-
and self._location.lower()
126-
in bigframes.constants.REP_ENABLED_BIGQUERY_LOCATIONS
127-
else _BIGQUERY_LOCATIONAL_ENDPOINT
128-
).format(location=self._location),
129-
)
130135
if "bqclient" in self._client_endpoints_override:
131136
bq_options = google.api_core.client_options.ClientOptions(
132137
api_endpoint=self._client_endpoints_override["bqclient"]
133138
)
139+
elif self._use_regional_endpoints:
140+
endpoint_template = _BIGQUERY_REGIONAL_ENDPOINT
141+
if (
142+
self._location is not None
143+
and self._location.lower()
144+
not in bigframes.constants.REP_ENABLED_BIGQUERY_LOCATIONS
145+
):
146+
endpoint_template = _BIGQUERY_LOCATIONAL_ENDPOINT
147+
148+
bq_options = google.api_core.client_options.ClientOptions(
149+
api_endpoint=endpoint_template.format(location=self._location)
150+
)
134151

135152
bq_info = google.api_core.client_info.ClientInfo(
136153
user_agent=self._application_name
@@ -172,16 +189,16 @@ def bqclient(self):
172189
def bqconnectionclient(self):
173190
if not self._bqconnectionclient:
174191
bqconnection_options = None
175-
if self._use_regional_endpoints:
192+
if "bqconnectionclient" in self._client_endpoints_override:
193+
bqconnection_options = google.api_core.client_options.ClientOptions(
194+
api_endpoint=self._client_endpoints_override["bqconnectionclient"]
195+
)
196+
elif self._use_regional_endpoints:
176197
bqconnection_options = google.api_core.client_options.ClientOptions(
177198
api_endpoint=_BIGQUERYCONNECTION_LOCATIONAL_ENDPOINT.format(
178199
location=self._location
179200
)
180201
)
181-
if "bqconnectionclient" in self._client_endpoints_override:
182-
bqconnection_options = google.api_core.client_options.ClientOptions(
183-
api_endpoint=self._client_endpoints_override["bqconnectionclient"]
184-
)
185202

186203
bqconnection_info = google.api_core.gapic_v1.client_info.ClientInfo(
187204
user_agent=self._application_name
@@ -200,21 +217,23 @@ def bqconnectionclient(self):
200217
def bqstoragereadclient(self):
201218
if not self._bqstoragereadclient:
202219
bqstorage_options = None
203-
if self._use_regional_endpoints:
220+
if "bqstoragereadclient" in self._client_endpoints_override:
204221
bqstorage_options = google.api_core.client_options.ClientOptions(
205-
api_endpoint=(
206-
_BIGQUERYSTORAGE_REGIONAL_ENDPOINT
207-
if self._location is not None
208-
and self._location.lower()
209-
in bigframes.constants.REP_ENABLED_BIGQUERY_LOCATIONS
210-
else _BIGQUERYSTORAGE_LOCATIONAL_ENDPOINT
211-
).format(location=self._location),
222+
api_endpoint=self._client_endpoints_override["bqstoragereadclient"]
212223
)
224+
elif self._use_regional_endpoints:
225+
endpoint_template = _BIGQUERYSTORAGE_REGIONAL_ENDPOINT
226+
if (
227+
self._location is not None
228+
and self._location.lower()
229+
not in bigframes.constants.REP_ENABLED_BIGQUERY_LOCATIONS
230+
):
231+
endpoint_template = _BIGQUERYSTORAGE_LOCATIONAL_ENDPOINT
213232

214-
if "bqstoragereadclient" in self._client_endpoints_override:
215233
bqstorage_options = google.api_core.client_options.ClientOptions(
216-
api_endpoint=self._client_endpoints_override["bqstoragereadclient"]
234+
api_endpoint=endpoint_template.format(location=self._location)
217235
)
236+
218237
bqstorage_info = google.api_core.gapic_v1.client_info.ClientInfo(
219238
user_agent=self._application_name
220239
)

tests/system/large/test_location.py

Lines changed: 29 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
# limitations under the License.
1414

1515
import typing
16+
import warnings
1617

1718
from google.cloud import bigquery
1819
import pytest
@@ -118,12 +119,22 @@ def test_bq_location_non_canonical(set_location, resolved_location):
118119
sorted(bigframes.constants.REP_ENABLED_BIGQUERY_LOCATIONS),
119120
)
120121
def test_bq_rep_endpoints(bigquery_location):
121-
session = bigframes.Session(
122-
context=bigframes.BigQueryOptions(
123-
location=bigquery_location, use_regional_endpoints=True
122+
with warnings.catch_warnings(record=True) as record:
123+
warnings.simplefilter("always")
124+
session = bigframes.Session(
125+
context=bigframes.BigQueryOptions(
126+
location=bigquery_location, use_regional_endpoints=True
127+
)
128+
)
129+
assert (
130+
len([warn for warn in record if isinstance(warn.message, FutureWarning)])
131+
== 0
124132
)
125-
)
126133

134+
# Verify that location and endpoints are correctly set for the BigQuery API
135+
# client
136+
# TODO(shobs): Figure out if the same can be verified for the other API
137+
# clients.
127138
assert session.bqclient.location == bigquery_location
128139
assert (
129140
session.bqclient._connection.API_BASE_URL
@@ -147,10 +158,21 @@ def test_bq_lep_endpoints(bigquery_location):
147158
# allowlisted for LEP access. We could hardcode one project which is
148159
# allowlisted but then not every open source developer will have access to
149160
# that. Let's rely on just creating the clients for LEP.
150-
clients_provider = bigframes.session.clients.ClientsProvider(
151-
location=bigquery_location, use_regional_endpoints=True
152-
)
161+
with pytest.warns(FutureWarning) as record:
162+
clients_provider = bigframes.session.clients.ClientsProvider(
163+
location=bigquery_location, use_regional_endpoints=True
164+
)
165+
assert len(record) == 1
166+
assert typing.cast(Warning, record[0].message).args[
167+
0
168+
] == bigframes.constants.LEP_DEPRECATION_WARNING_MESSAGE.format(
169+
location=bigquery_location
170+
)
153171

172+
# Verify that location and endpoints are correctly set for the BigQuery API
173+
# client
174+
# TODO(shobs): Figure out if the same can be verified for the other API
175+
# clients.
154176
assert clients_provider.bqclient.location == bigquery_location
155177
assert (
156178
clients_provider.bqclient._connection.API_BASE_URL

0 commit comments

Comments
 (0)