Repository: ploomber/jupysql (python)

Issue:
better error when missing CTE
when using the `--with` argument, we access [here](https://github.com/ploomber/jupysql/blob/fcbf525be3b4e52225a9e3ecc0e08b9d324c96e9/src/sql/store.py#L42); however, if the passed key does not exist, the error traceback is pretty long (see attached notebook).

Instead, we should raise an UsageError:

```python
from IPython.core.error import UsageError

raise UsageError(f"{key!r} is not a valid snippet identifier")
```

To provide a more helpful message, we can also use [difflib](https://docs.python.org/3/library/difflib.html#difflib.get_close_matches) to check if there are any closes matches:

> "subsett"  is not a valid snippet identifier. Did you mean "subset"?

If there are no close matches, we can display all available keys:

> "subsett"  is not a valid snippet identifier. Valid identifiers are "one", "another", and "something"


Pull request description:
fixed 257-better-error-when-missing-CTE with added test cases
## Describe your changes
fixed 257-better-error-when-missing-CTE with added test cases

## Issue ticket number and link
Closes #257 

## Checklist before requesting a review

- [x] I have performed a self-review of my code
- [x] I have added thorough [tests](https://github.com/ploomber/contributing/blob/main/CONTRIBUTING.md#testing) (when necessary).
- [ ] I have added the right documentation in the [docstring](https://github.com/ploomber/contributing/blob/main/CONTRIBUTING.md#documenting-changes-and-new-features) and [changelog](https://github.com/ploomber/contributing/blob/main/CONTRIBUTING.md#changelog) (when needed)



<!-- readthedocs-preview jupysql start -->
----
:books: Documentation preview :books:: https://jupysql--290.org.readthedocs.build/en/290/

<!-- readthedocs-preview jupysql end -->

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f06ce20..aff80e9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,7 +8,7 @@
 * [Feature] Adds `--interact` argument to `%%sql` to enable interactivity in parametrized SQL queries (#293)
 * [Feature] Results parse HTTP URLs to make them clickable (#230)
 * [Feature] Adds `ggplot` plotting API (histogram and boxplot)
-
+* [Fix] Clearer error when using `--with` with snippets that do not exist (#257)
 ## 0.6.6 (2023-03-16)
 
 * [Fix] Pinning SQLAlchemy 1.x
diff --git a/doc/user-guide/data-profiling.md b/doc/user-guide/data-profiling.md
index 716d84c..44c0323 100644
--- a/doc/user-guide/data-profiling.md
+++ b/doc/user-guide/data-profiling.md
@@ -116,6 +116,7 @@ To save the generated report as an HTML file, use the `--output`/`-o` attribute
 
 ```{code-cell} ipython3
 from IPython.display import HTML
+
 HTML("my-report.html")
 ```
 
diff --git a/src/sql/inspect.py b/src/sql/inspect.py
index a8970ae..3580a14 100644
--- a/src/sql/inspect.py
+++ b/src/sql/inspect.py
@@ -122,7 +122,7 @@ class TableDescription(DatabaseInspection):
                     Connection.current,
                     f"""SELECT DISTINCT {column} as top,
                     COUNT({column}) as frequency FROM {table_name}
-                    GROUP BY {column} ORDER BY Count({column}) Desc"""
+                    GROUP BY {column} ORDER BY Count({column}) Desc""",
                 ).fetchall()
 
                 table_stats[column]["freq"] = result_col_freq_values[0][1]
@@ -144,7 +144,7 @@ class TableDescription(DatabaseInspection):
                     COUNT({column}) AS count
                     FROM {table_name}
                     WHERE {column} IS NOT NULL
-                    """
+                    """,
                 ).fetchall()
 
                 table_stats[column]["min"] = result_value_values[0][0]
@@ -164,7 +164,7 @@ class TableDescription(DatabaseInspection):
                                 SELECT AVG({column}) AS avg
                                 FROM {table_name}
                                 WHERE {column} IS NOT NULL
-                                """
+                                """,
                 ).fetchall()
 
                 table_stats[column]["mean"] = float(results_avg[0][0])
@@ -190,7 +190,7 @@ class TableDescription(DatabaseInspection):
                         percentile_disc(0.75) WITHIN GROUP
                         (ORDER BY {column}) as key_75
                     FROM {table_name}
-                    """
+                    """,
                 ).fetchall()
 
                 for i, key in enumerate(special_numeric_keys):
diff --git a/src/sql/magic_cmd.py b/src/sql/magic_cmd.py
index 3f52873..f5521a8 100644
--- a/src/sql/magic_cmd.py
+++ b/src/sql/magic_cmd.py
@@ -2,11 +2,7 @@ import sys
 import argparse
 
 from IPython.utils.process import arg_split
-from IPython.core.magic import (
-    Magics,
-    line_magic,
-    magics_class
-)
+from IPython.core.magic import Magics, line_magic, magics_class
 from IPython.core.magic_arguments import argument, magic_arguments
 from IPython.core.error import UsageError
 from sqlglot import select, condition
@@ -135,7 +131,6 @@ class SqlCmdMagic(Magics, Configurable):
                 return True
 
         elif cmd_name == "profile":
-
             parser = CmdParser()
             parser.add_argument(
                 "-t", "--table", type=str, help="Table name", required=True
@@ -151,9 +146,7 @@ class SqlCmdMagic(Magics, Configurable):
 
             args = parser.parse_args(others)
 
-            report = inspect.get_table_statistics(
-                schema=args.schema, name=args.table
-            )
+            report = inspect.get_table_statistics(schema=args.schema, name=args.table)
 
             if args.output:
                 with open(args.output, "w") as f:
diff --git a/src/sql/run.py b/src/sql/run.py
index 2496995..944d0d5 100644
--- a/src/sql/run.py
+++ b/src/sql/run.py
@@ -118,9 +118,7 @@ class ResultSet(list, ColumnGuesserMixin):
             if isinstance(config.style, str):
                 _style = prettytable.__dict__[config.style.upper()]
 
-            self.pretty = PrettyTable(
-                self.field_names, style=_style
-            )
+            self.pretty = PrettyTable(self.field_names, style=_style)
         else:
             list.__init__(self, [])
             self.pretty = None
@@ -356,7 +354,7 @@ class FakeResultProxy(object):
         def fetchmany(size):
             pos = 0
             while pos < len(source_list):
-                yield source_list[pos: pos + size]
+                yield source_list[pos : pos + size]
                 pos += size
 
         self.fetchmany = fetchmany
diff --git a/src/sql/store.py b/src/sql/store.py
index e803b51..d98fcf6 100644
--- a/src/sql/store.py
+++ b/src/sql/store.py
@@ -4,6 +4,7 @@ from ploomber_core.exceptions import modify_exceptions
 from sqlglot import parse_one
 from IPython.core.error import UsageError
 import warnings
+import difflib
 
 
 class SQLStore(MutableMapping):
@@ -40,6 +41,16 @@ class SQLStore(MutableMapping):
         self._data[key] = value
 
     def __getitem__(self, key) -> str:
+        if not self._data:
+            raise UsageError("No saved SQL")
+        if key not in self._data:
+            matches = difflib.get_close_matches(key, self._data)
+            error = f'"{key}" is not a valid snippet identifier.'
+            if matches:
+                raise UsageError(error + f' Did you mean "{matches[0]}"?')
+            else:
+                valid = ", ".join(f'"{key}"' for key in self._data.keys())
+                raise UsageError(error + f" Valid identifiers are {valid}.")
         return self._data[key]
 
     def __iter__(self) -> Iterator[str]:


Test patch:
diff --git a/src/tests/test_compose.py b/src/tests/test_compose.py
index 4dcb491..098f092 100644
--- a/src/tests/test_compose.py
+++ b/src/tests/test_compose.py
@@ -12,7 +12,7 @@ def test_compose(ip):
     )
 
     result = ip.run_cell("%sqlrender final").result
-    expected = 'WITH author_sub AS (SELECT last_name FROM author\
- WHERE year_of_death > 1900) SELECT last_name FROM author_sub'
+    expected = "WITH author_sub AS (SELECT last_name FROM author\
+ WHERE year_of_death > 1900) SELECT last_name FROM author_sub"
 
     assert result == expected
diff --git a/src/tests/test_magic.py b/src/tests/test_magic.py
index 050e1a3..578c1b4 100644
--- a/src/tests/test_magic.py
+++ b/src/tests/test_magic.py
@@ -709,7 +709,7 @@ def test_save_with_non_existing_with(ip):
     out = ip.run_cell(
         "%sql --with non_existing_sub_query " "SELECT * FROM non_existing_sub_query"
     )
-    assert isinstance(out.error_in_exec, KeyError)
+    assert isinstance(out.error_in_exec, UsageError)
 
 
 def test_save_with_non_existing_table(ip, capsys):
diff --git a/src/tests/test_magic_cmd.py b/src/tests/test_magic_cmd.py
index 20242df..9134855 100644
--- a/src/tests/test_magic_cmd.py
+++ b/src/tests/test_magic_cmd.py
@@ -125,7 +125,6 @@ def test_table_profile(ip, tmp_empty):
 
 
 def test_table_schema_profile(ip, tmp_empty):
-
     with sqlite3.connect("a.db") as conn:
         conn.execute("CREATE TABLE t (n FLOAT)")
         conn.execute("INSERT INTO t VALUES (1)")
diff --git a/src/tests/test_store.py b/src/tests/test_store.py
index 60317c7..9522d7f 100644
--- a/src/tests/test_store.py
+++ b/src/tests/test_store.py
@@ -1,6 +1,6 @@
 import pytest
 from sql.connection import Connection
-
+from IPython.core.error import UsageError
 from sql.store import SQLStore
 
 
@@ -15,6 +15,58 @@ def test_sqlstore_setitem():
     assert store["a"] == "SELECT * FROM a"
 
 
+def test_sqlstore_getitem():
+    store = SQLStore()
+
+    # Test case 1: Test for a valid key
+    store["first"] = "SELECT * FROM a"
+    assert store.__getitem__("first") == "SELECT * FROM a"
+
+    # Test case 2: Test for an invalid key with no matches
+    with pytest.raises(UsageError) as err_info:
+        store.__getitem__("second")
+    assert (
+        str(err_info.value)
+        == '"second" is not a valid snippet identifier. Valid identifiers are "first".'
+    )
+
+    # Test case 3: Test for invalid key with close match
+    with pytest.raises(UsageError) as err_info:
+        store.__getitem__("firs")
+    assert (
+        str(err_info.value)
+        == '"firs" is not a valid snippet identifier. Did you mean "first"?'
+    )
+
+    # Test case 4: Test for multiple keys with close match
+    store["first2"] = "SELECT * FROM b"
+    with pytest.raises(UsageError) as err_info:
+        store.__getitem__("firs")
+    assert (
+        str(err_info.value)
+        == '"firs" is not a valid snippet identifier. Did you mean "first"?'
+    )
+
+    # Test case 5: Test for multiple keys with no close match
+    with pytest.raises(UsageError) as err_info:
+        store.__getitem__("second")
+    assert (
+        str(err_info.value)
+        == '"second" is not a valid snippet identifier. '
+        + 'Valid identifiers are "first", "first2".'
+    )
+
+    # Test case 6: Test for empty dictionary:
+    store2 = SQLStore()
+    with pytest.raises(UsageError) as err_info:
+        store2.__getitem__("second")
+    assert str(err_info.value) == "No saved SQL"
+
+    # Test case 7: Test for special character in key:
+    store["$%#"] = "SELECT * FROM a"
+    assert store.__getitem__("$%#") == "SELECT * FROM a"
+
+
 def test_key():
     store = SQLStore()
 
@@ -46,8 +98,7 @@ def test_serial(with_):
 
     result = store.render("SELECT * FROM third", with_=with_)
     assert (
-        str(result)
-        == "WITH first AS (SELECT * FROM a WHERE x > 10), "
+        str(result) == "WITH first AS (SELECT * FROM a WHERE x > 10), "
         "second AS (SELECT * FROM first WHERE x > 20), "
         "third AS (SELECT * FROM second WHERE x > 30) SELECT * FROM third"
     )
@@ -64,8 +115,7 @@ def test_branch_root():
 
     result = store.render("SELECT * FROM third", with_=["third_a", "first_b"])
     assert (
-        str(result)
-        == "WITH first_a AS (SELECT * FROM a WHERE x > 10), "
+        str(result) == "WITH first_a AS (SELECT * FROM a WHERE x > 10), "
         "second_a AS (SELECT * FROM first_a WHERE x > 20), "
         "third_a AS (SELECT * FROM second_a WHERE x > 30), "
         "first_b AS (SELECT * FROM b WHERE y > 10) SELECT * FROM third"
@@ -83,8 +133,7 @@ def test_branch_root_reverse_final_with():
 
     result = store.render("SELECT * FROM third", with_=["first_b", "third_a"])
     assert (
-        str(result)
-        == "WITH first_a AS (SELECT * FROM a WHERE x > 10), "
+        str(result) == "WITH first_a AS (SELECT * FROM a WHERE x > 10), "
         "second_a AS (SELECT * FROM first_a WHERE x > 20), "
         "first_b AS (SELECT * FROM b WHERE y > 10), "
         "third_a AS (SELECT * FROM second_a WHERE x > 30) SELECT * FROM third"