Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commit5178b55

Browse files
authored
feat: allow client options to be set in magics context (#322)
* feat: allow client options to be set in magics context* add separate client options for storage client
1 parentfb401bd commit5178b55

File tree

2 files changed

+188
-8
lines changed

2 files changed

+188
-8
lines changed

‎google/cloud/bigquery/magics/magics.py‎

Lines changed: 94 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,7 @@
139139

140140
importre
141141
importast
142+
importcopy
142143
importfunctools
143144
importsys
144145
importtime
@@ -155,6 +156,7 @@
155156
importsix
156157

157158
fromgoogle.api_coreimportclient_info
159+
fromgoogle.api_coreimportclient_options
158160
fromgoogle.api_core.exceptionsimportNotFound
159161
importgoogle.auth
160162
fromgoogle.cloudimportbigquery
@@ -178,11 +180,13 @@ def __init__(self):
178180
self._project=None
179181
self._connection=None
180182
self._default_query_job_config=bigquery.QueryJobConfig()
183+
self._bigquery_client_options=client_options.ClientOptions()
184+
self._bqstorage_client_options=client_options.ClientOptions()
181185

182186
@property
183187
defcredentials(self):
184188
"""google.auth.credentials.Credentials: Credentials to use for queries
185-
performed through IPython magics
189+
performed through IPython magics.
186190
187191
Note:
188192
These credentials do not need to be explicitly defined if you are
@@ -217,7 +221,7 @@ def credentials(self, value):
217221
@property
218222
defproject(self):
219223
"""str: Default project to use for queries performed through IPython
220-
magics
224+
magics.
221225
222226
Note:
223227
The project does not need to be explicitly defined if you have an
@@ -239,6 +243,54 @@ def project(self):
239243
defproject(self,value):
240244
self._project=value
241245

246+
@property
247+
defbigquery_client_options(self):
248+
"""google.api_core.client_options.ClientOptions: client options to be
249+
used through IPython magics.
250+
251+
Note::
252+
The client options do not need to be explicitly defined if no
253+
special network connections are required. Normally you would be
254+
using the https://bigquery.googleapis.com/ end point.
255+
256+
Example:
257+
Manually setting the endpoint:
258+
259+
>>> from google.cloud.bigquery import magics
260+
>>> client_options = {}
261+
>>> client_options['api_endpoint'] = "https://some.special.url"
262+
>>> magics.context.bigquery_client_options = client_options
263+
"""
264+
returnself._bigquery_client_options
265+
266+
@bigquery_client_options.setter
267+
defbigquery_client_options(self,value):
268+
self._bigquery_client_options=value
269+
270+
@property
271+
defbqstorage_client_options(self):
272+
"""google.api_core.client_options.ClientOptions: client options to be
273+
used through IPython magics for the storage client.
274+
275+
Note::
276+
The client options do not need to be explicitly defined if no
277+
special network connections are required. Normally you would be
278+
using the https://bigquerystorage.googleapis.com/ end point.
279+
280+
Example:
281+
Manually setting the endpoint:
282+
283+
>>> from google.cloud.bigquery import magics
284+
>>> client_options = {}
285+
>>> client_options['api_endpoint'] = "https://some.special.url"
286+
>>> magics.context.bqstorage_client_options = client_options
287+
"""
288+
returnself._bqstorage_client_options
289+
290+
@bqstorage_client_options.setter
291+
defbqstorage_client_options(self,value):
292+
self._bqstorage_client_options=value
293+
242294
@property
243295
defdefault_query_job_config(self):
244296
"""google.cloud.bigquery.job.QueryJobConfig: Default job
@@ -410,6 +462,24 @@ def _create_dataset_if_necessary(client, dataset_id):
410462
"Standard SQL if this argument is not used."
411463
),
412464
)
465+
@magic_arguments.argument(
466+
"--bigquery_api_endpoint",
467+
type=str,
468+
default=None,
469+
help=(
470+
"The desired API endpoint, e.g., bigquery.googlepis.com. Defaults to this "
471+
"option's value in the context bigquery_client_options."
472+
),
473+
)
474+
@magic_arguments.argument(
475+
"--bqstorage_api_endpoint",
476+
type=str,
477+
default=None,
478+
help=(
479+
"The desired API endpoint, e.g., bigquerystorage.googlepis.com. Defaults to "
480+
"this option's value in the context bqstorage_client_options."
481+
),
482+
)
413483
@magic_arguments.argument(
414484
"--use_bqstorage_api",
415485
action="store_true",
@@ -511,15 +581,34 @@ def _cell_magic(line, query):
511581
params=_helpers.to_query_parameters(ast.literal_eval(params_option_value))
512582

513583
project=args.projectorcontext.project
584+
585+
bigquery_client_options=copy.deepcopy(context.bigquery_client_options)
586+
ifargs.bigquery_api_endpoint:
587+
ifisinstance(bigquery_client_options,dict):
588+
bigquery_client_options["api_endpoint"]=args.bigquery_api_endpoint
589+
else:
590+
bigquery_client_options.api_endpoint=args.bigquery_api_endpoint
591+
514592
client=bigquery.Client(
515593
project=project,
516594
credentials=context.credentials,
517595
default_query_job_config=context.default_query_job_config,
518596
client_info=client_info.ClientInfo(user_agent=IPYTHON_USER_AGENT),
597+
client_options=bigquery_client_options,
519598
)
520599
ifcontext._connection:
521600
client._connection=context._connection
522-
bqstorage_client=_make_bqstorage_client(use_bqstorage_api,context.credentials)
601+
602+
bqstorage_client_options=copy.deepcopy(context.bqstorage_client_options)
603+
ifargs.bqstorage_api_endpoint:
604+
ifisinstance(bqstorage_client_options,dict):
605+
bqstorage_client_options["api_endpoint"]=args.bqstorage_api_endpoint
606+
else:
607+
bqstorage_client_options.api_endpoint=args.bqstorage_api_endpoint
608+
609+
bqstorage_client=_make_bqstorage_client(
610+
use_bqstorage_api,context.credentials,bqstorage_client_options,
611+
)
523612

524613
close_transports=functools.partial(_close_transports,client,bqstorage_client)
525614

@@ -632,7 +721,7 @@ def _split_args_line(line):
632721
returnparams_option_value,rest_of_args
633722

634723

635-
def_make_bqstorage_client(use_bqstorage_api,credentials):
724+
def_make_bqstorage_client(use_bqstorage_api,credentials,client_options):
636725
ifnotuse_bqstorage_api:
637726
returnNone
638727

@@ -658,6 +747,7 @@ def _make_bqstorage_client(use_bqstorage_api, credentials):
658747
returnbigquery_storage.BigQueryReadClient(
659748
credentials=credentials,
660749
client_info=gapic_client_info.ClientInfo(user_agent=IPYTHON_USER_AGENT),
750+
client_options=client_options,
661751
)
662752

663753

‎tests/unit/test_magics.py‎

Lines changed: 94 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -309,7 +309,7 @@ def test__make_bqstorage_client_false():
309309
credentials_mock=mock.create_autospec(
310310
google.auth.credentials.Credentials,instance=True
311311
)
312-
got=magics._make_bqstorage_client(False,credentials_mock)
312+
got=magics._make_bqstorage_client(False,credentials_mock, {})
313313
assertgotisNone
314314

315315

@@ -320,7 +320,7 @@ def test__make_bqstorage_client_true():
320320
credentials_mock=mock.create_autospec(
321321
google.auth.credentials.Credentials,instance=True
322322
)
323-
got=magics._make_bqstorage_client(True,credentials_mock)
323+
got=magics._make_bqstorage_client(True,credentials_mock, {})
324324
assertisinstance(got,bigquery_storage.BigQueryReadClient)
325325

326326

@@ -330,7 +330,7 @@ def test__make_bqstorage_client_true_raises_import_error(missing_bq_storage):
330330
)
331331

332332
withpytest.raises(ImportError)asexc_context,missing_bq_storage:
333-
magics._make_bqstorage_client(True,credentials_mock)
333+
magics._make_bqstorage_client(True,credentials_mock, {})
334334

335335
error_msg=str(exc_context.value)
336336
assert"google-cloud-bigquery-storage"inerror_msg
@@ -347,7 +347,7 @@ def test__make_bqstorage_client_true_missing_gapic(missing_grpcio_lib):
347347
)
348348

349349
withpytest.raises(ImportError)asexc_context,missing_grpcio_lib:
350-
magics._make_bqstorage_client(True,credentials_mock)
350+
magics._make_bqstorage_client(True,credentials_mock, {})
351351

352352
assert"grpcio"instr(exc_context.value)
353353

@@ -1180,6 +1180,96 @@ def test_bigquery_magic_with_project():
11801180
assertmagics.context.project=="general-project"
11811181

11821182

1183+
@pytest.mark.usefixtures("ipython_interactive")
1184+
deftest_bigquery_magic_with_bigquery_api_endpoint(ipython_ns_cleanup):
1185+
ip=IPython.get_ipython()
1186+
ip.extension_manager.load_extension("google.cloud.bigquery")
1187+
magics.context._connection=None
1188+
1189+
run_query_patch=mock.patch(
1190+
"google.cloud.bigquery.magics.magics._run_query",autospec=True
1191+
)
1192+
withrun_query_patchasrun_query_mock:
1193+
ip.run_cell_magic(
1194+
"bigquery",
1195+
"--bigquery_api_endpoint=https://bigquery_api.endpoint.com",
1196+
"SELECT 17 as num",
1197+
)
1198+
1199+
connection_used=run_query_mock.call_args_list[0][0][0]._connection
1200+
assertconnection_used.API_BASE_URL=="https://bigquery_api.endpoint.com"
1201+
# context client options should not change
1202+
assertmagics.context.bigquery_client_options.api_endpointisNone
1203+
1204+
1205+
@pytest.mark.usefixtures("ipython_interactive")
1206+
deftest_bigquery_magic_with_bigquery_api_endpoint_context_dict():
1207+
ip=IPython.get_ipython()
1208+
ip.extension_manager.load_extension("google.cloud.bigquery")
1209+
magics.context._connection=None
1210+
magics.context.bigquery_client_options= {}
1211+
1212+
run_query_patch=mock.patch(
1213+
"google.cloud.bigquery.magics.magics._run_query",autospec=True
1214+
)
1215+
withrun_query_patchasrun_query_mock:
1216+
ip.run_cell_magic(
1217+
"bigquery",
1218+
"--bigquery_api_endpoint=https://bigquery_api.endpoint.com",
1219+
"SELECT 17 as num",
1220+
)
1221+
1222+
connection_used=run_query_mock.call_args_list[0][0][0]._connection
1223+
assertconnection_used.API_BASE_URL=="https://bigquery_api.endpoint.com"
1224+
# context client options should not change
1225+
assertmagics.context.bigquery_client_options== {}
1226+
1227+
1228+
@pytest.mark.usefixtures("ipython_interactive")
1229+
deftest_bigquery_magic_with_bqstorage_api_endpoint(ipython_ns_cleanup):
1230+
ip=IPython.get_ipython()
1231+
ip.extension_manager.load_extension("google.cloud.bigquery")
1232+
magics.context._connection=None
1233+
1234+
run_query_patch=mock.patch(
1235+
"google.cloud.bigquery.magics.magics._run_query",autospec=True
1236+
)
1237+
withrun_query_patchasrun_query_mock:
1238+
ip.run_cell_magic(
1239+
"bigquery",
1240+
"--bqstorage_api_endpoint=https://bqstorage_api.endpoint.com",
1241+
"SELECT 17 as num",
1242+
)
1243+
1244+
client_used=run_query_mock.mock_calls[1][2]["bqstorage_client"]
1245+
assertclient_used._transport._host=="https://bqstorage_api.endpoint.com"
1246+
# context client options should not change
1247+
assertmagics.context.bqstorage_client_options.api_endpointisNone
1248+
1249+
1250+
@pytest.mark.usefixtures("ipython_interactive")
1251+
deftest_bigquery_magic_with_bqstorage_api_endpoint_context_dict():
1252+
ip=IPython.get_ipython()
1253+
ip.extension_manager.load_extension("google.cloud.bigquery")
1254+
magics.context._connection=None
1255+
magics.context.bqstorage_client_options= {}
1256+
1257+
run_query_patch=mock.patch(
1258+
"google.cloud.bigquery.magics.magics._run_query",autospec=True
1259+
)
1260+
withrun_query_patchasrun_query_mock:
1261+
ip.run_cell_magic(
1262+
"bigquery",
1263+
"--bqstorage_api_endpoint=https://bqstorage_api.endpoint.com",
1264+
"SELECT 17 as num",
1265+
)
1266+
1267+
client_used=run_query_mock.mock_calls[1][2]["bqstorage_client"]
1268+
assertclient_used._transport._host=="https://bqstorage_api.endpoint.com"
1269+
# context client options should not change
1270+
assertmagics.context.bqstorage_client_options== {}
1271+
1272+
11831273
@pytest.mark.usefixtures("ipython_interactive")
11841274
deftest_bigquery_magic_with_multiple_options():
11851275
ip=IPython.get_ipython()

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp