Visitar URL original
feat: store QueryJob to destination var on error · googleapis/google-cloud-python@02e17ea · GitHub
Skip to content

Commit 02e17ea

Browse files
committed
feat: store QueryJob to destination var on error
1 parent b4ea85c commit 02e17ea

4 files changed

Lines changed: 52 additions & 7 deletions

File tree

‎bigquery/google/cloud/bigquery/job.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2903,6 +2903,7 @@ def _begin(self, client=None, retry=DEFAULT_RETRY):
29032903
super(QueryJob, self)._begin(client=client, retry=retry)
29042904
except exceptions.GoogleCloudError as exc:
29052905
exc.message += self._format_for_exception(self.query, self.job_id)
2906+
exc.query_job = self
29062907
raise
29072908

29082909
def result(
@@ -2945,6 +2946,7 @@ def result(
29452946
)
29462947
except exceptions.GoogleCloudError as exc:
29472948
exc.message += self._format_for_exception(self.query, self.job_id)
2949+
exc.query_job = self
29482950
raise
29492951

29502952
# If the query job is complete but there are no query results, this was

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

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,9 @@
2828
2929
* ``<destination_var>`` (optional, line argument):
3030
variable to store the query results. The results are not displayed if
31-
this parameter is used.
31+
this parameter is used. If an error occurs during the query execution,
32+
the corresponding ``QueryJob`` instance (if available) is stored in
33+
the variable instead.
3234
* ``--project <project>`` (optional, line argument):
3335
Project to use for running the query. Defaults to the context
3436
:attr:`~google.cloud.bigquery.magics.Context.project`.
@@ -450,16 +452,22 @@ def _cell_magic(line, query):
450452
query_job = _run_query(client, query, job_config=job_config)
451453
except Exception as ex:
452454
error = str(ex)
455+
query_job = getattr(ex, "query_job", None)
453456

454457
if not args.verbose:
455458
display.clear_output()
456459

457460
if error:
458461
if args.destination_var:
459-
print(
460-
"Could not save output to variable '{}'.".format(args.destination_var),
461-
file=sys.stderr,
462-
)
462+
if query_job is not None:
463+
IPython.get_ipython().push({args.destination_var: query_job})
464+
else:
465+
print(
466+
"Could not save output to variable '{}'.".format(
467+
args.destination_var
468+
),
469+
file=sys.stderr,
470+
)
463471
print("\nERROR:\n", error, file=sys.stderr)
464472
return
465473

‎bigquery/tests/unit/test_job.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4337,8 +4337,10 @@ def test_result_error(self):
43374337
self.assertIsInstance(exc_info.exception, exceptions.GoogleCloudError)
43384338
self.assertEqual(exc_info.exception.code, http_client.BAD_REQUEST)
43394339

4340-
full_text = str(exc_info.exception)
4340+
exc_job_instance = getattr(exc_info.exception, "query_job", None)
4341+
self.assertIs(exc_job_instance, job)
43414342

4343+
full_text = str(exc_info.exception)
43424344
assert job.job_id in full_text
43434345
assert "Query Job SQL Follows" in full_text
43444346

@@ -4370,8 +4372,10 @@ def test__begin_error(self):
43704372
self.assertIsInstance(exc_info.exception, exceptions.GoogleCloudError)
43714373
self.assertEqual(exc_info.exception.code, http_client.BAD_REQUEST)
43724374

4373-
full_text = str(exc_info.exception)
4375+
exc_job_instance = getattr(exc_info.exception, "query_job", None)
4376+
self.assertIs(exc_job_instance, job)
43744377

4378+
full_text = str(exc_info.exception)
43754379
assert job.job_id in full_text
43764380
assert "Query Job SQL Follows" in full_text
43774381

‎bigquery/tests/unit/test_magics.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -794,6 +794,37 @@ def test_bigquery_magic_dryrun_option_saves_query_job_to_variable():
794794
assert isinstance(q_job, job.QueryJob)
795795

796796

797+
@pytest.mark.usefixtures("ipython_interactive")
798+
def test_bigquery_magic_saves_query_job_to_variable_on_error():
799+
ip = IPython.get_ipython()
800+
ip.extension_manager.load_extension("google.cloud.bigquery")
801+
magics.context.credentials = mock.create_autospec(
802+
google.auth.credentials.Credentials, instance=True
803+
)
804+
805+
client_query_patch = mock.patch(
806+
"google.cloud.bigquery.client.Client.query", autospec=True
807+
)
808+
809+
query_job = mock.create_autospec(job.QueryJob, instance=True)
810+
exception = Exception("Unexpected SELECT")
811+
exception.query_job = query_job
812+
query_job.result.side_effect = exception
813+
814+
sql = "SELECT SELECT 17 AS num"
815+
816+
assert "result" not in ip.user_ns
817+
818+
with client_query_patch as client_query_mock:
819+
client_query_mock.return_value = query_job
820+
return_value = ip.run_cell_magic("bigquery", "result", sql)
821+
822+
assert return_value is None
823+
assert "result" in ip.user_ns
824+
result = ip.user_ns["result"]
825+
assert isinstance(result, job.QueryJob)
826+
827+
797828
@pytest.mark.usefixtures("ipython_interactive")
798829
def test_bigquery_magic_w_maximum_bytes_billed_invalid():
799830
ip = IPython.get_ipython()

0 commit comments

Comments
 (0)