feat: use bigquery-magics package for the %%bigquery magic (#1965)

* feat: use `bigquery-magics` package for the `%%bigquery` magic

* ignore types on bigquery-magics package

* Update samples/magics/noxfile_config.py

Co-authored-by: Chalmer Lowe <[email protected]>

---------

Co-authored-by: Chalmer Lowe <[email protected]>

Author: tswast

google/cloud/bigquery/__init__.py

@@ -27,6 +27,7 @@
 - :class:`~google.cloud.bigquery.table.Table` represents a single "relation".
 """
 
+import warnings
 
 from google.cloud.bigquery import version as bigquery_version
 
@@ -114,6 +115,11 @@
 from google.cloud.bigquery.table import TimePartitioning
 from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration
 
+try:
+    import bigquery_magics  # type: ignore
+except ImportError:
+    bigquery_magics = None
+
 __all__ = [
     "__version__",
     "Client",
@@ -214,8 +220,16 @@
 
 def load_ipython_extension(ipython):
     """Called by IPython when this module is loaded as an IPython extension."""
-    from google.cloud.bigquery.magics.magics import _cell_magic
-
-    ipython.register_magic_function(
-        _cell_magic, magic_kind="cell", magic_name="bigquery"
+    warnings.warn(
+        "%load_ext google.cloud.bigquery is deprecated. Install bigquery-magics package and use `%load_ext bigquery_magics`, instead.",
+        category=FutureWarning,
     )
+
+    if bigquery_magics is not None:
+        bigquery_magics.load_ipython_extension(ipython)
+    else:
+        from google.cloud.bigquery.magics.magics import _cell_magic
+
+        ipython.register_magic_function(
+            _cell_magic, magic_kind="cell", magic_name="bigquery"
+        )

google/cloud/bigquery/magics/magics.py

@@ -14,70 +14,11 @@
 
 """IPython Magics
 
-.. function:: %%bigquery
-
-    IPython cell magic to run a query and display the result as a DataFrame
-
-    .. code-block:: python
-
-        %%bigquery [<destination_var>] [--project <project>] [--use_legacy_sql]
-                   [--verbose] [--params <params>]
-        <query>
-
-    Parameters:
-
-    * ``<destination_var>`` (Optional[line argument]):
-        variable to store the query results. The results are not displayed if
-        this parameter is used. If an error occurs during the query execution,
-        the corresponding ``QueryJob`` instance (if available) is stored in
-        the variable instead.
-    * ``--destination_table`` (Optional[line argument]):
-        A dataset and table to store the query results. If table does not exists,
-        it will be created. If table already exists, its data will be overwritten.
-        Variable should be in a format <dataset_id>.<table_id>.
-    * ``--no_query_cache`` (Optional[line argument]):
-        Do not use cached query results.
-    * ``--project <project>`` (Optional[line argument]):
-        Project to use for running the query. Defaults to the context
-        :attr:`~google.cloud.bigquery.magics.Context.project`.
-    * ``--use_bqstorage_api`` (Optional[line argument]):
-        [Deprecated] Not used anymore, as BigQuery Storage API is used by default.
-    * ``--use_rest_api`` (Optional[line argument]):
-        Use the BigQuery REST API instead of the Storage API.
-    * ``--use_legacy_sql`` (Optional[line argument]):
-        Runs the query using Legacy SQL syntax. Defaults to Standard SQL if
-        this argument not used.
-    * ``--verbose`` (Optional[line argument]):
-        If this flag is used, information including the query job ID and the
-        amount of time for the query to complete will not be cleared after the
-        query is finished. By default, this information will be displayed but
-        will be cleared after the query is finished.
-    * ``--params <params>`` (Optional[line argument]):
-        If present, the argument following the ``--params`` flag must be
-        either:
-
-        * :class:`str` - A JSON string representation of a dictionary in the
-          format ``{"param_name": "param_value"}`` (ex. ``{"num": 17}``). Use
-          of the parameter in the query should be indicated with
-          ``@param_name``. See ``In[5]`` in the Examples section below.
-
-        * :class:`dict` reference - A reference to a ``dict`` in the format
-          ``{"param_name": "param_value"}``, where the value types must be JSON
-          serializable. The variable reference is indicated by a ``$`` before
-          the variable name (ex. ``$my_dict_var``). See ``In[6]`` and ``In[7]``
-          in the Examples section below.
-
-    * ``<query>`` (required, cell argument):
-        SQL query to run. If the query does not contain any whitespace (aside
-        from leading and trailing whitespace), it is assumed to represent a
-        fully-qualified table ID, and the latter's data will be fetched.
+Install ``bigquery-magics`` and call ``%load_ext bigquery_magics`` to use the
+``%%bigquery`` cell magic.
 
-    Returns:
-        A :class:`pandas.DataFrame` with the query results.
-
-    .. note::
-        All queries run using this magic will run using the context
-        :attr:`~google.cloud.bigquery.magics.Context.credentials`.
+See the `BigQuery Magics reference documentation
+<https://googleapis.dev/python/bigquery-magics/latest/>`_.
 """
 
 from __future__ import print_function
@@ -109,6 +50,11 @@
 from google.cloud.bigquery.dbapi import _helpers
 from google.cloud.bigquery.magics import line_arg_parser as lap
 
+try:
+    import bigquery_magics  # type: ignore
+except ImportError:
+    bigquery_magics = None
+
 
 IPYTHON_USER_AGENT = "ipython-{}".format(IPython.__version__)
 
@@ -280,7 +226,14 @@ def progress_bar_type(self, value):
         self._progress_bar_type = value
 
 
-context = Context()
+# If bigquery_magics is available, we load that extension rather than this one.
+# Ensure google.cloud.bigquery.magics.context setters are on the correct magics
+# implementation in case the user has installed the package but hasn't updated
+# their code.
+if bigquery_magics is not None:
+    context = bigquery_magics.context
+else:
+    context = Context()
 
 
 def _handle_error(error, destination_var=None):

samples/magics/noxfile_config.py

@@ -0,0 +1,37 @@
+# Copyright 2020 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.
+
+# Default TEST_CONFIG_OVERRIDE for python repos.
+
+# You can copy this file into your directory, then it will be inported from
+# the noxfile.py.
+
+# The source of truth:
+# https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/noxfile_config.py
+
+TEST_CONFIG_OVERRIDE = {
+    # You can opt out from the test for specific Python versions.
+    "ignored_versions": [
+        "2.7",
+    ],
+    # An envvar key for determining the project id to use. Change it
+    # to 'BUILD_SPECIFIC_GCLOUD_PROJECT' if you want to opt in using a
+    # build specific Cloud project. You can also use your own string
+    # to use your own Cloud project.
+    "gcloud_project_env": "GOOGLE_CLOUD_PROJECT",
+    # "gcloud_project_env": "BUILD_SPECIFIC_GCLOUD_PROJECT",
+    # A dictionary you want to inject into your test. Don't put any
+    # secrets here. These values will override predefined values.
+    "envs": {},
+}

samples/magics/query.py

@@ -24,7 +24,7 @@
 
 def query() -> "pandas.DataFrame":
     ip = IPython.get_ipython()
-    ip.extension_manager.load_extension("google.cloud.bigquery")
+    ip.extension_manager.load_extension("bigquery_magics")
 
     sample = """
     # [START bigquery_jupyter_query]

samples/magics/query_params_scalars.py

@@ -24,7 +24,7 @@
 
 def query_with_parameters() -> "pandas.DataFrame":
     ip = IPython.get_ipython()
-    ip.extension_manager.load_extension("google.cloud.bigquery")
+    ip.extension_manager.load_extension("bigquery_magics")
 
     sample = """
     # [START bigquery_jupyter_query_params_scalars]

samples/magics/requirements.txt

@@ -1,3 +1,4 @@
+bigquery_magics==0.1.0
 db-dtypes==1.2.0
 google.cloud.bigquery==3.25.0
 google-cloud-bigquery-storage==2.25.0

samples/notebooks/jupyter_tutorial_test.py

@@ -60,7 +60,7 @@ def _strip_region_tags(sample_text: str) -> str:
 def test_jupyter_tutorial(ipython: "TerminalInteractiveShell") -> None:
     matplotlib.use("agg")
     ip = IPython.get_ipython()
-    ip.extension_manager.load_extension("google.cloud.bigquery")
+    ip.extension_manager.load_extension("bigquery_magics")
 
     sample = """
     # [START bigquery_jupyter_magic_gender_by_year]

samples/notebooks/requirements.txt

@@ -1,3 +1,4 @@
+bigquery-magics==0.1.0
 db-dtypes==1.2.0
 google-cloud-bigquery==3.25.0
 google-cloud-bigquery-storage==2.25.0

setup.py

@@ -66,8 +66,7 @@
     ],
     "geopandas": ["geopandas>=0.9.0, <1.0dev", "Shapely>=1.8.4, <3.0.0dev"],
     "ipython": [
-        "ipython>=7.23.1,!=8.1.0",
-        "ipykernel>=6.0.0",
+        "bigquery-magics >= 0.1.0",
     ],
     "tqdm": ["tqdm >= 4.7.4, <5.0.0dev"],
     "opentelemetry": [

testing/constraints-3.7.txt

@@ -5,6 +5,7 @@
 #
 # e.g., if setup.py has "foo >= 1.14.0, < 2.0.0dev",
 # Then this file should have foo==1.14.0
+bigquery-magics==0.1.0
 db-dtypes==0.3.0
 geopandas==0.9.0
 google-api-core==2.11.1

tests/system/test_magics.py

@@ -50,7 +50,10 @@ def test_bigquery_magic(ipython_interactive):
     current_process = psutil.Process()
     conn_count_start = len(current_process.connections())
 
-    ip.extension_manager.load_extension("google.cloud.bigquery")
+    # Deprecated, but should still work in google-cloud-bigquery 3.x.
+    with pytest.warns(FutureWarning, match="bigquery_magics"):
+        ip.extension_manager.load_extension("google.cloud.bigquery")
+
     sql = """
         SELECT
             CONCAT(