diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py index ccbab8b5eb44..8ea8860a7ac2 100644 --- a/bigquery/google/cloud/bigquery/job.py +++ b/bigquery/google/cloud/bigquery/job.py @@ -2903,6 +2903,7 @@ def _begin(self, client=None, retry=DEFAULT_RETRY): super(QueryJob, self)._begin(client=client, retry=retry) except exceptions.GoogleCloudError as exc: exc.message += self._format_for_exception(self.query, self.job_id) + exc.query_job = self raise def result( @@ -2945,6 +2946,7 @@ def result( ) except exceptions.GoogleCloudError as exc: exc.message += self._format_for_exception(self.query, self.job_id) + exc.query_job = self raise # If the query job is complete but there are no query results, this was diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index c238bb50317a..2a174cefeea3 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -28,7 +28,9 @@ * ```` (optional, line argument): variable to store the query results. The results are not displayed if - this parameter is used. + this parameter is used. If an error occurs during the query execution, + the corresponding ``QueryJob`` instance (if available) is stored in + the variable instead. * ``--project `` (optional, line argument): Project to use for running the query. Defaults to the context :attr:`~google.cloud.bigquery.magics.Context.project`. @@ -267,13 +269,29 @@ def default_query_job_config(self, value): context = Context() -def _print_error(error, destination_var=None): +def _handle_error(error, destination_var=None): + """Process a query execution error. + + Args: + error (Exception): + An exception that ocurred during the query exectution. + destination_var (Optional[str]): + The name of the IPython session variable to store the query job. + """ if destination_var: - print( - "Could not save output to variable '{}'.".format(destination_var), - file=sys.stderr, - ) - print("\nERROR:\n", error, file=sys.stderr) + query_job = getattr(error, "query_job", None) + + if query_job is not None: + IPython.get_ipython().push({destination_var: query_job}) + else: + # this is the case when previewing table rows by providing just + # table ID to cell magic + print( + "Could not save output to variable '{}'.".format(destination_var), + file=sys.stderr, + ) + + print("\nERROR:\n", str(error), file=sys.stderr) def _run_query(client, query, job_config=None): @@ -452,7 +470,7 @@ def _cell_magic(line, query): try: rows = client.list_rows(query, max_results=max_results) except Exception as ex: - _print_error(str(ex), args.destination_var) + _handle_error(ex, args.destination_var) return result = rows.to_dataframe(bqstorage_client=bqstorage_client) @@ -476,7 +494,7 @@ def _cell_magic(line, query): try: query_job = _run_query(client, query, job_config=job_config) except Exception as ex: - _print_error(str(ex), args.destination_var) + _handle_error(ex, args.destination_var) return if not args.verbose: diff --git a/bigquery/tests/unit/test_job.py b/bigquery/tests/unit/test_job.py index b34184f00cd9..07c3ba88da5b 100644 --- a/bigquery/tests/unit/test_job.py +++ b/bigquery/tests/unit/test_job.py @@ -4337,8 +4337,10 @@ def test_result_error(self): self.assertIsInstance(exc_info.exception, exceptions.GoogleCloudError) self.assertEqual(exc_info.exception.code, http_client.BAD_REQUEST) - full_text = str(exc_info.exception) + exc_job_instance = getattr(exc_info.exception, "query_job", None) + self.assertIs(exc_job_instance, job) + full_text = str(exc_info.exception) assert job.job_id in full_text assert "Query Job SQL Follows" in full_text @@ -4370,8 +4372,10 @@ def test__begin_error(self): self.assertIsInstance(exc_info.exception, exceptions.GoogleCloudError) self.assertEqual(exc_info.exception.code, http_client.BAD_REQUEST) - full_text = str(exc_info.exception) + exc_job_instance = getattr(exc_info.exception, "query_job", None) + self.assertIs(exc_job_instance, job) + full_text = str(exc_info.exception) assert job.job_id in full_text assert "Query Job SQL Follows" in full_text diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index ec642ff384e1..ed253636c468 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -902,6 +902,37 @@ def test_bigquery_magic_dryrun_option_saves_query_job_to_variable(): assert isinstance(q_job, job.QueryJob) +@pytest.mark.usefixtures("ipython_interactive") +def test_bigquery_magic_saves_query_job_to_variable_on_error(): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + magics.context.credentials = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + + client_query_patch = mock.patch( + "google.cloud.bigquery.client.Client.query", autospec=True + ) + + query_job = mock.create_autospec(job.QueryJob, instance=True) + exception = Exception("Unexpected SELECT") + exception.query_job = query_job + query_job.result.side_effect = exception + + sql = "SELECT SELECT 17 AS num" + + assert "result" not in ip.user_ns + + with client_query_patch as client_query_mock: + client_query_mock.return_value = query_job + return_value = ip.run_cell_magic("bigquery", "result", sql) + + assert return_value is None + assert "result" in ip.user_ns + result = ip.user_ns["result"] + assert isinstance(result, job.QueryJob) + + @pytest.mark.usefixtures("ipython_interactive") def test_bigquery_magic_w_maximum_bytes_billed_invalid(): ip = IPython.get_ipython()