658 snippet support sqlcmd

CHANGELOG.md

@@ -13,9 +13,11 @@
 * [Fix] Refactored `ResultSet` to lazy loading (#470)
 * [Fix] Removed `WITH` when a snippet does not have a dependency (#657)
 * [Fix] Used display module when generating CTE (#649)
+* [Fix] Added support for `profile` and `explore` commands with saved snippets (#658)
 * [Doc] Re-organized sections. Adds section showing how to share notebooks via Ploomber Cloud
 * [Fix] Adding `--with` back because of issues with sqlglot query parser (#684)
 
+
 ## 0.7.9 (2023-06-19)
 
 * [Feature] Modified `histogram` command to support data with NULL values ([#176](https://github.com/ploomber/jupysql/issues/176))

src/sql/inspect.py

@@ -9,6 +9,7 @@
 from sql import util
 from IPython.core.display import HTML
 import uuid
+from sql.store import store
 
 
 def _get_inspector(conn):
@@ -173,7 +174,7 @@ class Columns(DatabaseInspection):
     """
 
     def __init__(self, name, schema, conn=None) -> None:
-        util.is_table_exists(name, schema)
+        util.is_saved_snippet_or_table_exists(name, schema)
 
         inspector = _get_inspector(conn)
 
@@ -231,14 +232,20 @@ class TableDescription(DatabaseInspection):
     """
 
     def __init__(self, table_name, schema=None) -> None:
-        util.is_table_exists(table_name, schema)
+        with_ = util.is_saved_snippet_or_table_exists(table_name, schema)
 
         if schema:
             table_name = f"{schema}.{table_name}"
 
+        if with_:
+            with_clause = str(store.render("", with_=with_))
+        else:
+            with_clause = ""
+
         columns_query_result = sql.run.raw_run(
-            Connection.current, f"SELECT * FROM {table_name} WHERE 1=0"
+            Connection.current, f"{with_clause} SELECT * FROM {table_name} WHERE 1=0"
         )
+
         if Connection.is_custom_connection():
             columns = [i[0] for i in columns_query_result.description]
         else:
@@ -255,7 +262,8 @@ def __init__(self, table_name, schema=None) -> None:
             # check the datatype of a column
             try:
                 result = sql.run.raw_run(
-                    Connection.current, f"""SELECT {column} FROM {table_name} LIMIT 1"""
+                    Connection.current,
+                    f"""{with_clause} SELECT {column} FROM {table_name} LIMIT 1""",
                 ).fetchone()
 
                 value = result[0]
@@ -272,7 +280,7 @@ def __init__(self, table_name, schema=None) -> None:
             try:
                 result_col_freq_values = sql.run.raw_run(
                     Connection.current,
-                    f"""SELECT DISTINCT {column} as top,
+                    f"""{with_clause} SELECT DISTINCT {column} as top,
                     COUNT({column}) as frequency FROM {table_name}
                     GROUP BY top ORDER BY frequency Desc""",
                 ).fetchall()
@@ -290,6 +298,7 @@ def __init__(self, table_name, schema=None) -> None:
                 result_value_values = sql.run.raw_run(
                     Connection.current,
                     f"""
+                    {with_clause}
                     SELECT MIN({column}) AS min,
                     MAX({column}) AS max,
                     COUNT({column}) AS count
@@ -314,6 +323,7 @@ def __init__(self, table_name, schema=None) -> None:
                 result_value_values = sql.run.raw_run(
                     Connection.current,
                     f"""
+                    {with_clause}
                     SELECT
                     COUNT(DISTINCT {column}) AS unique_count
                     FROM {table_name}
@@ -329,10 +339,11 @@ def __init__(self, table_name, schema=None) -> None:
                 results_avg = sql.run.raw_run(
                     Connection.current,
                     f"""
-                                SELECT AVG({column}) AS avg
-                                FROM {table_name}
-                                WHERE {column} IS NOT NULL
-                                """,
+                    {with_clause}
+                    SELECT AVG({column}) AS avg
+                    FROM {table_name}
+                    WHERE {column} IS NOT NULL
+                    """,
                 ).fetchall()
 
                 columns_to_include_in_report.update(["mean"])
@@ -349,6 +360,7 @@ def __init__(self, table_name, schema=None) -> None:
                 result = sql.run.raw_run(
                     Connection.current,
                     f"""
+                    {with_clause}
                     SELECT
                         stddev_pop({column}) as key_std,
                         percentile_disc(0.25) WITHIN GROUP

src/sql/magic_plot.py

@@ -92,7 +92,7 @@ def execute(self, line="", cell="", local_ns=None):
         if cmd.args.with_:
             with_ = cmd.args.with_
         else:
-            with_ = self._check_table_exists(table)
+            with_ = self.is_saved_snippet_or_table_exists(table)
 
         if cmd.args.line[0] in {"box", "boxplot"}:
             return plot.boxplot(
@@ -129,7 +129,7 @@ def execute(self, line="", cell="", local_ns=None):
             )
 
     @staticmethod
-    def _check_table_exists(table):
+    def is_saved_snippet_or_table_exists(table):
         with_ = None
         if util.is_saved_snippet(table):
             with_ = [table]

src/sql/util.py

@@ -255,7 +255,7 @@ def support_only_sql_alchemy_connection(command):
 
 
 def fetch_sql_with_pagination(
-    table, offset, n_rows, sort_column=None, sort_order=None
+    table, offset, n_rows, sort_column=None, sort_order=None, with_clause=None
 ) -> tuple:
     """
     Returns next n_rows and columns from table starting at the offset
@@ -277,19 +277,25 @@ def fetch_sql_with_pagination(
 
     sort_order : 'DESC' or 'ASC', default None
         Order list
+
+    with_clause : str, default None
+        Name of the snippet used to generate the table
     """
-    is_table_exists(table)
 
-    order_by = "" if not sort_column else f"ORDER BY {sort_column} {sort_order}"
+    if with_clause is None:
+        with_clause = ""
+        is_table_exists(table)
 
+    order_by = "" if not sort_column else f"ORDER BY {sort_column} {sort_order}"
     query = f"""
+    {with_clause}
     SELECT * FROM {table} {order_by}
-    OFFSET {offset} ROWS FETCH NEXT {n_rows} ROWS ONLY"""
+    LIMIT {n_rows} OFFSET {offset}"""
 
     rows = Connection.current.execute(query).fetchall()
 
     columns = sql.run.raw_run(
-        Connection.current, f"SELECT * FROM {table} WHERE 1=0"
+        Connection.current, f"{with_clause} SELECT * FROM {table} WHERE 1=0"
     ).keys()
 
     return rows, columns
@@ -363,3 +369,22 @@ def show_deprecation_warning():
         "raise an exception in the next major release so please remove it.",
         FutureWarning,
     )
+
+
+def is_saved_snippet_or_table_exists(table, schema=None):
+    """
+    Check if the referenced table is a snippet
+    Parameters
+    ----------
+    table : str, name of table
+    Returns
+    -------
+    with_ : str, default None
+        name of the snippet
+    """
+    with_ = None
+    if is_saved_snippet(table):
+        with_ = [table]
+    else:
+        is_table_exists(table, schema)
+    return with_

src/sql/widgets/table_widget/table_widget.py

@@ -6,8 +6,9 @@
 from sql.util import (
     fetch_sql_with_pagination,
     parse_sql_results_to_json,
-    is_table_exists,
+    is_saved_snippet_or_table_exists,
 )
+from sql.store import store
 
 from sql.widgets import utils
 from sql.telemetry import telemetry
@@ -33,7 +34,12 @@ def __init__(self, table):
 
         self.html = ""
 
-        is_table_exists(table)
+        self.with_ = is_saved_snippet_or_table_exists(table)
+
+        if self.with_:
+            self.with_clause = str(store.render("", with_=self.with_))
+        else:
+            self.with_clause = ""
 
         # load css
         html_style = utils.load_css(f"{BASE_DIR}/css/tableWidget.css")
@@ -58,10 +64,12 @@ def create_table(self, table):
         Creates an HTML table with default data
         """
         rows_per_page = 10
-        rows, columns = fetch_sql_with_pagination(table, 0, rows_per_page)
+        rows, columns = fetch_sql_with_pagination(
+            table, 0, rows_per_page, with_clause=self.with_clause
+        )
         rows = parse_sql_results_to_json(rows, columns)
 
-        query = f"SELECT count(*) FROM {table}"
+        query = f"{self.with_clause} SELECT count(*) FROM {table}"
         n_total = Connection.current.session.execute(
             sqlalchemy.sql.text(query)
         ).fetchone()[0]
@@ -187,6 +195,7 @@ def _recv(msg):
                     n_rows,
                     sort_column=sort_column,
                     sort_order=sort_order,
+                    with_clause=self.with_clause,
                 )
                 rows_json = parse_sql_results_to_json(rows, columns)
 

src/tests/test_magic_cmd.py

@@ -9,6 +9,7 @@
 from sql.store import store
 from sql.inspect import _is_numeric
 from sql.display import Table, Message
+from IPython.utils.io import capture_output
 
 
 VALID_COMMANDS_MESSAGE = (
@@ -391,6 +392,63 @@ def test_table_profile_store(ip, tmp_empty):
     assert report.is_file()
 
 
+def test_table_profile_with_snippets(ip, tmp_empty):
+    ip.run_cell(
+        """
+        %%sql duckdb://
+        CREATE TABLE people (name varchar(50), number int, country varchar(50));
+        INSERT INTO people VALUES ('joe', 82, 'usa');
+        INSERT INTO people VALUES ('paula', 93, 'uk');
+        INSERT INTO people VALUES ('sam', 77, 'canada');
+        INSERT INTO people VALUES ('emily', 65, 'usa');
+        INSERT INTO people VALUES ('michael', 89, 'usa');
+        INSERT INTO people VALUES ('sarah', 81, 'uk');
+        INSERT INTO people VALUES ('james', 76, 'usa');
+        INSERT INTO people VALUES ('angela', 88, 'usa');
+        INSERT INTO people VALUES ('robert', 82, 'usa');
+        INSERT INTO people VALUES ('lisa', 92, 'uk');
+        INSERT INTO people VALUES ('mark', 77, 'usa');
+        INSERT INTO people VALUES ('jennifer', 68, 'australia');
+        """
+    )
+    ip.run_cell(
+        """
+        %%sql --save citizen
+        select * from people where country = 'usa'
+        """
+    )
+    out = ip.run_cell("%sqlcmd profile -t citizen").result
+
+    expected = {
+        "count": [7, 7, 7],
+        "mean": [math.nan, "79.8571", math.nan],
+        "min": [math.nan, 65, math.nan],
+        "max": [math.nan, 89, math.nan],
+        "unique": [7, 6, 1],
+        "freq": [1, math.nan, 7],
+        "top": ["joe", math.nan, "usa"],
+        "std": [math.nan, "7.5862", math.nan],
+        "25%": [math.nan, "76.0000", math.nan],
+        "50%": [math.nan, "82.0000", math.nan],
+        "75%": [math.nan, "88.0000", math.nan],
+    }
+
+    stats_table = out._table
+
+    assert len(stats_table.rows) == len(expected)
+
+    for row in stats_table:
+        profile_metric = _get_row_string(row, " ")
+        name = _get_row_string(row, "name")
+        number = _get_row_string(row, "number")
+        country = _get_row_string(row, "country")
+
+        assert profile_metric in expected
+        assert name == str(expected[profile_metric][0])
+        assert number == str(expected[profile_metric][1])
+        assert country == str(expected[profile_metric][2])
+
+
 @pytest.mark.parametrize(
     "cell, error_type, error_message",
     [
@@ -585,3 +643,35 @@ def test_delete_invalid_snippet(arg, ip_snippets):
         str(out.error_in_exec) == "No such saved snippet found "
         ": non_existent_snippet"
     )
+
+
+def test_table_explore_with_snippets(ip, tmp_empty):
+    ip.run_cell(
+        """
+        %%sql duckdb://
+        CREATE TABLE people (name varchar(50), number int, country varchar(50));
+        INSERT INTO people VALUES ('joe', 82, 'usa');
+        INSERT INTO people VALUES ('paula', 93, 'uk');
+        INSERT INTO people VALUES ('sam', 77, 'canada');
+        INSERT INTO people VALUES ('emily', 65, 'usa');
+        INSERT INTO people VALUES ('michael', 89, 'usa');
+        INSERT INTO people VALUES ('sarah', 81, 'uk');
+        INSERT INTO people VALUES ('james', 76, 'usa');
+        INSERT INTO people VALUES ('angela', 88, 'usa');
+        INSERT INTO people VALUES ('robert', 82, 'usa');
+        INSERT INTO people VALUES ('lisa', 92, 'uk');
+        INSERT INTO people VALUES ('mark', 77, 'usa');
+        INSERT INTO people VALUES ('jennifer', 68, 'australia');
+        """
+    )
+    ip.run_cell(
+        """
+        %%sql --save citizen
+        select * from people where country = 'usa'
+        """
+    )
+    with capture_output() as captured:
+        ip.run_cell("%sqlcmd explore -t citizen")
+
+    assert "Plotting using saved snippet : citizen" in captured.stdout
+    assert "sql.widgets.table_widget.table_widget.TableWidget object" in captured.stdout