feat: store QueryJob to destination var on error

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

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`.
@@ -450,16 +452,22 @@ def _cell_magic(line, query):
         query_job = _run_query(client, query, job_config=job_config)
     except Exception as ex:
         error = str(ex)
+        query_job = getattr(ex, "query_job", None)
 
     if not args.verbose:
         display.clear_output()
 
     if error:
         if args.destination_var:
-            print(
-                "Could not save output to variable '{}'.".format(args.destination_var),
-                file=sys.stderr,
-            )
+            if query_job is not None:
+                IPython.get_ipython().push({args.destination_var: query_job})
+            else:
+                print(
+                    "Could not save output to variable '{}'.".format(
+                        args.destination_var
+                    ),
+                    file=sys.stderr,
+                )
         print("\nERROR:\n", error, file=sys.stderr)
         return
 

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
 

bigquery/tests/unit/test_magics.py

@@ -794,6 +794,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()