Adds _connection object to bigquery magics context.

bigquery/google/cloud/bigquery/magics.py

@@ -159,6 +159,7 @@ class Context(object):
     def __init__(self):
         self._credentials = None
         self._project = None
+        self._connection = None
         self._use_bqstorage_api = None
 
     @property
@@ -363,6 +364,8 @@ def _cell_magic(line, query):
 
     project = args.project or context.project
     client = bigquery.Client(project=project, credentials=context.credentials)
+    if context._connection:
+        client._connection = context._connection
     bqstorage_client = _make_bqstorage_client(
         args.use_bqstorage_api or context.use_bqstorage_api, context.credentials
     )

bigquery/tests/unit/test__helpers.py

@@ -782,6 +782,17 @@ def _make_field(field_type, mode="NULLABLE", name="testing", fields=()):
     return SchemaField(name=name, field_type=field_type, mode=mode, fields=fields)
 
 
+def _make_connection(*responses):
+    import google.cloud.bigquery._http
+    import mock
+    from google.cloud.exceptions import NotFound
+
+    mock_conn = mock.create_autospec(google.cloud.bigquery._http.Connection)
+    mock_conn.user_agent = "testing 1.2.3"
+    mock_conn.api_request.side_effect = list(responses) + [NotFound("miss")]
+    return mock_conn
+
+
 class Test_scalar_field_to_json(unittest.TestCase):
     def _call_fut(self, field, value):
         from google.cloud.bigquery._helpers import _scalar_field_to_json

bigquery/tests/unit/test_client.py

@@ -40,6 +40,7 @@
 import google.api_core.exceptions
 from google.api_core.gapic_v1 import client_info
 import google.cloud._helpers
+from tests.unit.test__helpers import _make_connection
 from google.cloud.bigquery.dataset import DatasetReference
 
 
@@ -49,16 +50,6 @@ def _make_credentials():
     return mock.Mock(spec=google.auth.credentials.Credentials)
 
 
-def _make_connection(*responses):
-    import google.cloud.bigquery._http
-    from google.cloud.exceptions import NotFound
-
-    mock_conn = mock.create_autospec(google.cloud.bigquery._http.Connection)
-    mock_conn.user_agent = "testing 1.2.3"
-    mock_conn.api_request.side_effect = list(responses) + [NotFound("miss")]
-    return mock_conn
-
-
 def _make_list_partitons_meta_info(project, dataset_id, table_id, num_rows=0):
     return {
         "tableReference": {

bigquery/tests/unit/test_magics.py

@@ -14,6 +14,7 @@
 
 import re
 import mock
+import six
 from concurrent import futures
 
 import pytest
@@ -38,6 +39,7 @@
     bigquery_storage_v1beta1 = None
 from google.cloud.bigquery import table
 from google.cloud.bigquery import magics
+from tests.unit.test__helpers import _make_connection
 
 
 pytestmark = pytest.mark.skipif(IPython is None, reason="Requires `ipython`")
@@ -101,6 +103,58 @@ def test_context_credentials_and_project_can_be_set_explicitly():
     assert default_mock.call_count == 0
 
 
+@pytest.mark.usefixtures("ipython_interactive")
+def test_context_connection_can_be_overriden():
+    ip = IPython.get_ipython()
+    ip.extension_manager.load_extension("google.cloud.bigquery")
+    magics.context._project = None
+    magics.context._credentials = None
+
+    credentials_mock = mock.create_autospec(
+        google.auth.credentials.Credentials, instance=True
+    )
+    project = "project-123"
+    default_patch = mock.patch(
+        "google.auth.default", return_value=(credentials_mock, project)
+    )
+
+    query = "select * from persons"
+    job_reference = {"projectId": project, "jobId": "some-random-id"}
+    table = {"projectId": project, "datasetId": "ds", "tableId": "persons"}
+    resource = {
+        "jobReference": job_reference,
+        "configuration": {
+            "query": {
+                "destinationTable": table,
+                "query": query,
+                "queryParameters": [],
+                "useLegacySql": False,
+            }
+        },
+        "status": {"state": "DONE"},
+    }
+    data = {"jobReference": job_reference, "totalRows": 0, "rows": []}
+
+    conn = magics.context._connection = _make_connection(resource, data)
+    list_rows_patch = mock.patch(
+        "google.cloud.bigquery.client.Client.list_rows",
+        return_value=google.cloud.bigquery.table._EmptyRowIterator(),
+    )
+    with list_rows_patch as list_rows, default_patch:
+        ip.run_cell_magic("bigquery", "", query)
+
+    # Check that query actually starts the job.
+    list_rows.assert_called()
+    assert len(conn.api_request.call_args_list) == 2
+    _, req = conn.api_request.call_args_list[0]
+    assert req["method"] == "POST"
+    assert req["path"] == "/projects/{}/jobs".format(project)
+    sent = req["data"]
+    assert isinstance(sent["jobReference"]["jobId"], six.string_types)
+    sent_config = sent["configuration"]["query"]
+    assert sent_config["query"] == query
+
+
 def test__run_query():
     magics.context._credentials = None