feat(bigquery): store QueryJob to destination var on error (googleapis#9245)

plamut · tswast · commit 6e6cd392d736 · 2019-09-26T12:37:57.000-07:00
diff --git 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
@@ -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:
diff --git 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
@@ -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()
