From 03fe1d544ef22865c07c680873f980c64bbc7abc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tim=20Swe=C3=B1a=20=28Swast=29?= <swast@google.com>
Date: Tue, 17 Dec 2024 16:32:31 -0600
Subject: [PATCH] feat: add `%%bqsql` as an alias to the `%%bigquery` magic
 (#72)

---
 README.rst                  | 12 ++++++++++++
 bigquery_magics/__init__.py |  1 +
 bigquery_magics/bigquery.py |  9 +++++----
 docs/magics.rst             | 10 ++++++++--
 tests/unit/test_bigquery.py | 11 +++++++++--
 5 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/README.rst b/README.rst
index c531e3c..33738d7 100644
--- a/README.rst
+++ b/README.rst
@@ -101,3 +101,15 @@ Perform a query
     GROUP BY name
     ORDER BY count DESC
     LIMIT 3
+
+Since BigQuery supports Python via BigQuery DataFrames, `%%bqsql` is offered as
+an alias to clarify the language of these cells.
+
+.. code:: python
+
+    %%bqsql
+    SELECT name, SUM(number) as count
+    FROM 'bigquery-public-data.usa_names.usa_1910_current'
+    GROUP BY name
+    ORDER BY count DESC
+    LIMIT 3
diff --git a/bigquery_magics/__init__.py b/bigquery_magics/__init__.py
index 627ca98..61a7c76 100644
--- a/bigquery_magics/__init__.py
+++ b/bigquery_magics/__init__.py
@@ -27,6 +27,7 @@ def load_ipython_extension(ipython):
     ipython.register_magic_function(
         _cell_magic, magic_kind="cell", magic_name="bigquery"
     )
+    ipython.register_magic_function(_cell_magic, magic_kind="cell", magic_name="bqsql")
 
 
 __all__ = (
diff --git a/bigquery_magics/bigquery.py b/bigquery_magics/bigquery.py
index 94632d9..5e99883 100644
--- a/bigquery_magics/bigquery.py
+++ b/bigquery_magics/bigquery.py
@@ -14,13 +14,13 @@
 
 """IPython Magics
 
-.. function:: %%bigquery
+.. function:: ``%%bigquery`` or ``%%bqsql``
 
     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]
+        %%bqsql [<destination_var>] [--project <project>] [--use_legacy_sql]
                    [--verbose] [--params <params>]
         <query>
 
@@ -74,11 +74,12 @@
         fully-qualified table ID, and the latter's data will be fetched.
 
     Returns:
-        A :class:`pandas.DataFrame` with the query results.
+        A :class:`pandas.DataFrame` or :class:`bigframes.pandas.DataFrame`
+        with the query results, depending on the ``engine`` chosen.
 
     .. note::
         All queries run using this magic will run using the context
-        :attr:`~google.cloud.bigquery.magics.Context.credentials`.
+        :attr:`~bigquery_magics.config.Context.credentials`.
 """
 
 from __future__ import print_function
diff --git a/docs/magics.rst b/docs/magics.rst
index 8fec049..2a029e4 100644
--- a/docs/magics.rst
+++ b/docs/magics.rst
@@ -8,7 +8,7 @@ in a Jupyter notebook cell.
 
     %load_ext bigquery_magics
 
-This makes the ``%%bigquery`` magic available.
+This makes the ``%%bigquery`` and ``%%bqsql`` magics available.
 
 Code Samples
 ------------
@@ -30,5 +30,11 @@ Running a parameterized query:
 API Reference
 -------------
 
-.. automodule:: google.cloud.bigquery.magics.magics
+.. automodule:: bigquery_magics.bigquery
+    :members:
+
+Configuration
+~~~~~~~~~~~~~
+
+.. automodule:: bigquery_magics.config
     :members:
diff --git a/tests/unit/test_bigquery.py b/tests/unit/test_bigquery.py
index d927345..61ad744 100644
--- a/tests/unit/test_bigquery.py
+++ b/tests/unit/test_bigquery.py
@@ -395,13 +395,20 @@ def test__create_dataset_if_necessary_not_exist():
         client.create_dataset.assert_called_once()
 
 
+@pytest.mark.parametrize(
+    ("magic_name",),
+    (
+        ("bigquery",),
+        ("bqsql",),
+    ),
+)
 @pytest.mark.usefixtures("ipython_interactive")
-def test_extension_load():
+def test_extension_load(magic_name):
     ip = IPython.get_ipython()
     ip.extension_manager.load_extension("bigquery_magics")
 
     # verify that the magic is registered and has the correct source
-    magic = ip.magics_manager.magics["cell"].get("bigquery")
+    magic = ip.magics_manager.magics["cell"].get(magic_name)
     assert magic.__module__ == "bigquery_magics.bigquery"