Adds _connection object to bigquery magics context. (#8192)

* Adds _connection object to bigquery magics context.

* Moves make_connection to helpers.py and adds unit test for no connection in 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/helpers.py

@@ -0,0 +1,24 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+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