Repository: googleapis/google-cloud-python

Issue:
BigQuery: [magics] if `destination_var` is set and an error occurs, save the QueryJob for further inspection
 **Is your feature request related to a problem? Please describe.**

When an error occurs, it's useful to have a QueryJob object to inspect the errors properties, as well as check for other details about the job that may have caused the failure. At the very least, a job ID is needed if one wants to file a support ticket about the failure.

 **Describe the solution you'd like**

As identified in https://github.com/googleapis/google-cloud-python/pull/9067#pullrequestreview-278337340, when an error occurs with a `%%bigquery` magic, it would be useful to save the query job object for further inspection of the error and the rest of the job.

A problem with this is that in successful cases `destination_var` will be a pandas DataFrame, but in failure cases `destination_var` will be a QueryJob object. Since the magics are primarily used interactively, this probably isn't that much of a problem.

 **Describe alternatives you've considered**

Print out the job ID and have the user call `client.get_job()` in order to inspect the object. This does not work in the case of dry run queries, as no job resource is saved in the backend.



Discussion:

Resolving patch:
diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py
--- 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
--- a/bigquery/google/cloud/bigquery/magics.py
+++ b/bigquery/google/cloud/bigquery/magics.py
@@ -28,7 +28,9 @@
 
     * ``<destination_var>`` (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 <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:


Test patch:
diff --git a/bigquery/tests/unit/test_job.py b/bigquery/tests/unit/test_job.py
--- 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
--- 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()