Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 4 commits into from
Jul 10, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 18 additions & 4 deletions google/cloud/bigquery/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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"
)
81 changes: 17 additions & 64 deletions google/cloud/bigquery/magics/magics.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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__)

Expand Down Expand Up @@ -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):
Expand Down
37 changes: 37 additions & 0 deletions samples/magics/noxfile_config.py
Original file line number Diff line number Diff line change
@@ -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": {},
}
2 changes: 1 addition & 1 deletion samples/magics/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
2 changes: 1 addition & 1 deletion samples/magics/query_params_scalars.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
1 change: 1 addition & 0 deletions samples/magics/requirements.txt
Original file line number Diff line number Diff line change
@@ -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
Expand Down
2 changes: 1 addition & 1 deletion samples/notebooks/jupyter_tutorial_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
1 change: 1 addition & 0 deletions samples/notebooks/requirements.txt
Original file line number Diff line number Diff line change
@@ -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
Expand Down
3 changes: 1 addition & 2 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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": [
Expand Down
1 change: 1 addition & 0 deletions testing/constraints-3.7.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
5 changes: 4 additions & 1 deletion tests/system/test_magics.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
Loading