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

Commit30f8326

Browse files
committed
Add default_query_job_config property.
Test API params get set.
1 parent629d34b commit30f8326

File tree

2 files changed

+120
-98
lines changed

2 files changed

+120
-98
lines changed

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

Lines changed: 24 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -161,7 +161,7 @@ def __init__(self):
161161
self._project=None
162162
self._connection=None
163163
self._use_bqstorage_api=None
164-
self._maximum_bytes_billed=0
164+
self._default_query_job_config=bigquery.QueryJobConfig()
165165

166166
@property
167167
defcredentials(self):
@@ -239,30 +239,26 @@ def use_bqstorage_api(self, value):
239239
self._use_bqstorage_api=value
240240

241241
@property
242-
defmaximum_bytes_billed(self):
243-
"""int: Maximum bytes to be billed for this job or :data:`None` if not set.
242+
defdefault_query_job_config(self):
243+
"""google.cloud.bigquery.job.QueryJobConfig: Default job
244+
configuration for queries.
245+
246+
The context's :class:`~google.cloud.bigquery.job.QueryJobConfig` is
247+
used for queries. Some properties can be overridden with arguments to
248+
the magics.
244249
245250
Example:
246-
Manually setting the context maximum_bytes_billed:
251+
Manually setting the default value for ``maximum_bytes_billed``
252+
to 100 MB:
247253
248254
>>> from google.cloud.bigquery import magics
249-
>>> magics.context.maximum_bytes_billed = '123'
250-
251-
Raises:
252-
ValueError: If the parameters are invalid.
255+
>>> magics.context.default_query_job_config.maximum_bytes_billed = 100000000
253256
"""
254-
returnself._maximum_bytes_billed
257+
returnself._default_query_job_config
255258

256-
@maximum_bytes_billed.setter
257-
defmaximum_bytes_billed(self,value):
258-
try:
259-
value=int(value)
260-
self._maximum_bytes_billed=value
261-
job_config=bigquery.job.QueryJobConfig()
262-
job_config.maximum_bytes_billed=self._maximum_bytes_billed
263-
264-
exceptException:
265-
raiseValueError("value is not a valid integer.")
259+
@default_query_job_config.setter
260+
defdefault_query_job_config(self,value):
261+
self._default_query_job_config=value
266262

267263

268264
context=Context()
@@ -322,7 +318,8 @@ def _run_query(client, query, job_config=None):
322318
"--maximum_bytes_billed",
323319
default=None,
324320
help=(
325-
"maximum_bytes_billed to use for executing this query. Defaults to the context maximum_bytes_billed."
321+
"maximum_bytes_billed to use for executing this query. Defaults to "
322+
"the context default_query_job_config.maximum_bytes_billed."
326323
),
327324
)
328325
@magic_arguments.argument(
@@ -397,7 +394,11 @@ def _cell_magic(line, query):
397394
)
398395

399396
project=args.projectorcontext.project
400-
client=bigquery.Client(project=project,credentials=context.credentials)
397+
client=bigquery.Client(
398+
project=project,
399+
credentials=context.credentials,
400+
default_query_job_config=context.default_query_job_config,
401+
)
401402
ifcontext._connection:
402403
client._connection=context._connection
403404
bqstorage_client=_make_bqstorage_client(
@@ -410,13 +411,8 @@ def _cell_magic(line, query):
410411
ifargs.maximum_bytes_billed=="None":
411412
job_config.maximum_bytes_billed=0
412413
elifargs.maximum_bytes_billedisnotNone:
413-
try:
414-
value=int(args.maximum_bytes_billed)
415-
job_config.maximum_bytes_billed=value
416-
exceptException:
417-
raiseValueError("value is not a valid integer.")
418-
else:
419-
job_config.maximum_bytes_billed=context.maximum_bytes_billed
414+
value=int(args.maximum_bytes_billed)
415+
job_config.maximum_bytes_billed=value
420416
query_job=_run_query(client,query,job_config)
421417

422418
ifnotargs.verbose:

‎bigquery/tests/unit/test_magics.py‎

Lines changed: 96 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,13 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
importcopy
1516
importre
16-
importmock
17-
importsix
1817
fromconcurrentimportfutures
1918

19+
importmock
2020
importpytest
21+
importsix
2122

2223
try:
2324
importpandas
@@ -63,6 +64,26 @@ def ipython_interactive(request, ipython):
6364
yieldipython
6465

6566

67+
JOB_REFERENCE_RESOURCE= {"projectId":"its-a-project-eh","jobId":"some-random-id"}
68+
TABLE_REFERENCE_RESOURCE= {
69+
"projectId":"its-a-project-eh",
70+
"datasetId":"ds",
71+
"tableId":"persons",
72+
}
73+
QUERY_RESOURCE= {
74+
"jobReference":JOB_REFERENCE_RESOURCE,
75+
"configuration": {
76+
"query": {
77+
"destinationTable":TABLE_REFERENCE_RESOURCE,
78+
"query":"SELECT 42 FROM `life.the_universe.and_everything`;",
79+
"queryParameters": [],
80+
"useLegacySql":False,
81+
}
82+
},
83+
"status": {"state":"DONE"},
84+
}
85+
86+
6687
deftest_context_credentials_auto_set_w_application_default_credentials():
6788
"""When Application Default Credentials are set, the context credentials
6889
will be created the first time it is called
@@ -117,22 +138,13 @@ def test_context_connection_can_be_overriden():
117138
default_patch=mock.patch(
118139
"google.auth.default",return_value=(credentials_mock,project)
119140
)
141+
job_reference=copy.deepcopy(JOB_REFERENCE_RESOURCE)
142+
job_reference["projectId"]=project
120143

121144
query="select * from persons"
122-
job_reference= {"projectId":project,"jobId":"some-random-id"}
123-
table= {"projectId":project,"datasetId":"ds","tableId":"persons"}
124-
resource= {
125-
"jobReference":job_reference,
126-
"configuration": {
127-
"query": {
128-
"destinationTable":table,
129-
"query":query,
130-
"queryParameters": [],
131-
"useLegacySql":False,
132-
}
133-
},
134-
"status": {"state":"DONE"},
135-
}
145+
resource=copy.deepcopy(QUERY_RESOURCE)
146+
resource["jobReference"]=job_reference
147+
resource["configuration"]["query"]["query"]=query
136148
data= {"jobReference":job_reference,"totalRows":0,"rows": []}
137149

138150
conn=magics.context._connection=make_connection(resource,data)
@@ -170,22 +182,13 @@ def test_context_no_connection():
170182
default_patch=mock.patch(
171183
"google.auth.default",return_value=(credentials_mock,project)
172184
)
185+
job_reference=copy.deepcopy(JOB_REFERENCE_RESOURCE)
186+
job_reference["projectId"]=project
173187

174188
query="select * from persons"
175-
job_reference= {"projectId":project,"jobId":"some-random-id"}
176-
table= {"projectId":project,"datasetId":"ds","tableId":"persons"}
177-
resource= {
178-
"jobReference":job_reference,
179-
"configuration": {
180-
"query": {
181-
"destinationTable":table,
182-
"query":query,
183-
"queryParameters": [],
184-
"useLegacySql":False,
185-
}
186-
},
187-
"status": {"state":"DONE"},
188-
}
189+
resource=copy.deepcopy(QUERY_RESOURCE)
190+
resource["jobReference"]=job_reference
191+
resource["configuration"]["query"]["query"]=query
189192
data= {"jobReference":job_reference,"totalRows":0,"rows": []}
190193

191194
conn_mock=make_connection(resource,data,data,data)
@@ -549,7 +552,7 @@ def test_bigquery_magic_without_bqstorage(monkeypatch):
549552

550553

551554
@pytest.mark.usefixtures("ipython_interactive")
552-
deftest_maximum_bytes_billed_w_int_magic():
555+
deftest_bigquery_magic_w_maximum_bytes_billed():
553556
ip=IPython.get_ipython()
554557
ip.extension_manager.load_extension("google.cloud.bigquery")
555558
magics.context._project=None
@@ -576,15 +579,17 @@ def test_maximum_bytes_billed_w_int_magic():
576579
query_job_mock.to_dataframe.return_value=result
577580
withrun_query_patchasrun_query_mock,default_patch:
578581
run_query_mock.return_value=query_job_mock
579-
return_value=ip.run_cell_magic("bigquery","--maximum_bytes_billed=123456789",sql)
582+
return_value=ip.run_cell_magic(
583+
"bigquery","--maximum_bytes_billed=123456789",sql
584+
)
580585

581586
bqstorage_mock.assert_not_called()
582587
query_job_mock.to_dataframe.assert_called_once_with(bqstorage_client=None)
583588
assertisinstance(return_value,pandas.DataFrame)
584589

585590

586591
@pytest.mark.usefixtures("ipython_interactive")
587-
deftest_maximum_bytes_billed_w_string_params():
592+
deftest_bigquery_magic_w_maximum_bytes_billed_invalid():
588593
ip=IPython.get_ipython()
589594
ip.extension_manager.load_extension("google.cloud.bigquery")
590595
magics.context._project=None
@@ -595,39 +600,80 @@ def test_maximum_bytes_billed_w_string_params():
595600
ip.run_cell_magic("bigquery","--maximum_bytes_billed=abc",sql)
596601

597602

603+
@pytest.mark.parametrize(
604+
"param_value,expected", [("987654321","987654321"), ("None","0")]
605+
)
598606
@pytest.mark.usefixtures("ipython_interactive")
599-
deftest_maximum_bytes_billed_w_none__magic():
607+
deftest_bigquery_magic_w_maximum_bytes_billed_overrides_context(param_value,expected):
600608
ip=IPython.get_ipython()
601609
ip.extension_manager.load_extension("google.cloud.bigquery")
602610
magics.context._project=None
603611

604-
bqstorage_mock=mock.create_autospec(
605-
bigquery_storage_v1beta1.BigQueryStorageClient
606-
)
612+
# Set the default maximum bytes billed, so we know it's overridable by the param.
613+
magics.context.default_query_job_config.maximum_bytes_billed=1234567
607614

615+
project="test-project"
616+
job_reference=copy.deepcopy(JOB_REFERENCE_RESOURCE)
617+
job_reference["projectId"]=project
618+
query="SELECT 17 AS num"
619+
resource=copy.deepcopy(QUERY_RESOURCE)
620+
resource["jobReference"]=job_reference
621+
resource["configuration"]["query"]["query"]=query
622+
data= {"jobReference":job_reference,"totalRows":0,"rows": []}
608623
credentials_mock=mock.create_autospec(
609624
google.auth.credentials.Credentials,instance=True
610625
)
611626
default_patch=mock.patch(
612627
"google.auth.default",return_value=(credentials_mock,"general-project")
613628
)
614-
run_query_patch=mock.patch(
615-
"google.cloud.bigquery.magics._run_query",autospec=True
629+
conn=magics.context._connection=make_connection(resource,data)
630+
list_rows_patch=mock.patch(
631+
"google.cloud.bigquery.client.Client.list_rows",
632+
return_value=google.cloud.bigquery.table._EmptyRowIterator(),
616633
)
634+
withlist_rows_patch,default_patch:
635+
ip.run_cell_magic(
636+
"bigquery","--maximum_bytes_billed={}".format(param_value),query
637+
)
617638

618-
sql="SELECT 17 AS num"
619-
result=pandas.DataFrame([17],columns=["num"])
620-
query_job_mock=mock.create_autospec(
621-
google.cloud.bigquery.job.QueryJob,instance=True
639+
_,req=conn.api_request.call_args_list[0]
640+
sent_config=req["data"]["configuration"]["query"]
641+
assertsent_config["maximumBytesBilled"]==expected
642+
643+
644+
@pytest.mark.usefixtures("ipython_interactive")
645+
deftest_bigquery_magic_w_maximum_bytes_billed_w_context():
646+
ip=IPython.get_ipython()
647+
ip.extension_manager.load_extension("google.cloud.bigquery")
648+
magics.context._project=None
649+
650+
magics.context.default_query_job_config.maximum_bytes_billed=1234567
651+
652+
project="test-project"
653+
job_reference=copy.deepcopy(JOB_REFERENCE_RESOURCE)
654+
job_reference["projectId"]=project
655+
query="SELECT 17 AS num"
656+
resource=copy.deepcopy(QUERY_RESOURCE)
657+
resource["jobReference"]=job_reference
658+
resource["configuration"]["query"]["query"]=query
659+
data= {"jobReference":job_reference,"totalRows":0,"rows": []}
660+
credentials_mock=mock.create_autospec(
661+
google.auth.credentials.Credentials,instance=True
622662
)
623-
query_job_mock.to_dataframe.return_value=result
624-
withrun_query_patchasrun_query_mock,default_patch:
625-
run_query_mock.return_value=query_job_mock
626-
return_value=ip.run_cell_magic("bigquery","--maximum_bytes_billed=None",sql)
663+
default_patch=mock.patch(
664+
"google.auth.default",return_value=(credentials_mock,"general-project")
665+
)
666+
conn=magics.context._connection=make_connection(resource,data)
667+
list_rows_patch=mock.patch(
668+
"google.cloud.bigquery.client.Client.list_rows",
669+
return_value=google.cloud.bigquery.table._EmptyRowIterator(),
670+
)
671+
withlist_rows_patch,default_patch:
672+
ip.run_cell_magic("bigquery","",query)
627673

628-
bqstorage_mock.assert_not_called()
629-
query_job_mock.to_dataframe.assert_called_once_with(bqstorage_client=None)
630-
assertisinstance(return_value,pandas.DataFrame)
674+
_,req=conn.api_request.call_args_list[0]
675+
sent_config=req["data"]["configuration"]["query"]
676+
assertsent_config["maximumBytesBilled"]=="1234567"
631677

632678

633679
@pytest.mark.usefixtures("ipython_interactive")
@@ -734,23 +780,3 @@ def test_bigquery_magic_with_improperly_formatted_params():
734780

735781
withpytest.raises(SyntaxError):
736782
ip.run_cell_magic("bigquery","--params {17}",sql)
737-
738-
739-
deftest_maximum_bytes_billed_set_value():
740-
"""When Application Default Credentials are set, the context credentials
741-
will be created the first time it is called
742-
"""
743-
744-
fromgoogle.cloud.bigqueryimportQueryJobConfig
745-
job_config=QueryJobConfig()
746-
magics.context.maximum_bytes_billed=1234567489
747-
assertjob_config.maximum_bytes_billed==magics.context.maximum_bytes_billed
748-
749-
750-
deftest_maximum_bytes_billed_set_string():
751-
"""When Application Default Credentials are set, the context credentials
752-
will be created the first time it is called
753-
"""
754-
withpytest.raises(ValueError):
755-
magics.context.maximum_bytes_billed="abc"
756-

0 commit comments

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