Repository: ploomber/jupysql (python)

Issue:
Better error messages when passing invalid connection string argument
When passing an invalid connection argument, the current error message could be clearer. This is an example we got from a user.

Input:

```
%sql mysql+pymysql://userid:pwd@hostname/dbname
```

Output:

```
Connection info needed in SQLAlchemy format, example:
               postgresql://username:password@hostname/dbname
               or an existing connection: dict_keys(['duckdb://'])
No module named 'pymysql'
```

Note that the actual error is displayed in the last line: `No module named 'pymysql'`

We should detect these issues and raises shorter error messages. In this case, it'd be better to say that we could not connect using the passed connection string since we're missing the pymysql package. Since we rely on SQLAlchemy for db connections we should check if SQLAlchemy is already doing some validation and forward the error message (as opposed to writing validation logic ourselves).

In this case, the connection string is valid, but the connection cannot be established. But another error can be that the connection string is invalid. Furthermore, we should also consider the case where a user wants to connect to an existing [engine](https://jupysql.readthedocs.io/en/latest/connecting.html#using-an-existing-sqlalchemy-engine-engine), so the error message should provide some feedback. A typical scenario can be a typo, for that, we can use [difflib](https://docs.python.org/3/library/difflib.html#difflib.get_close_matches).

There might be other use cases as well. Since this is the entry point for using JupySQL, we should ensure a robust validation logic so we cover most scenarios with accurate error messages.








Pull request:
Connection str error
## Describe your changes

## Issue ticket number and link
Closes #70

## 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).
- [x] 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--120.org.readthedocs.build/en/120/

<!-- readthedocs-preview jupysql end -->

Patch:
diff --git a/.gitignore b/.gitignore
index d8ebb04..ab1b296 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
+*.jsonl
 *.db
 *.csv
 *.parquet
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 600b474..471ae51 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ## 0.5.5dev
 
+* [Doc] Adds tutorial on querying JSON data
+
 ## 0.5.4 (2023-02-06)
 
 * [Feature] Adds `%jupysql`/`%%jupysql` as alias for `%sql`/`%%sql`
diff --git a/doc/_toc.yml b/doc/_toc.yml
index 425d74c..eaceeac 100644
--- a/doc/_toc.yml
+++ b/doc/_toc.yml
@@ -33,6 +33,7 @@ parts:
     - file: howto
     - file: howto/postgres-install
     - file: howto/postgres-connect
+    - file: howto/json
 
   - caption: Community
     chapters:
diff --git a/doc/howto/json.md b/doc/howto/json.md
new file mode 100644
index 0000000..fb48aaa
--- /dev/null
+++ b/doc/howto/json.md
@@ -0,0 +1,187 @@
+---
+jupytext:
+  text_representation:
+    extension: .md
+    format_name: myst
+    format_version: 0.13
+    jupytext_version: 1.14.4
+kernelspec:
+  display_name: Python 3 (ipykernel)
+  language: python
+  name: python3
+---
+
+# Run SQL on JSON files
+
+In this tutorial, we'll show you how to query JSON with JupySQL and DuckDB.
+
+
+First, let's install the required dependencies:
+
+```{code-cell} ipython3
+:tags: [remove-cell]
+
+# this cell won't be visible in the docs
+from pathlib import Path
+
+paths = ["people.jsonl", "people.csv"]
+
+for path in paths:
+    path = Path(p)
+
+    if path.exists():
+        print(f"Deleting {path}")
+        path.unlink()
+```
+
+```{code-cell} ipython3
+:tags: [hide-output]
+
+%pip install jupysql duckdb duckdb-engine rich --quiet
+```
+
+Now, let's generate some data. Note that DuckDB expects your data to contain *one JSON object per line*; this format is called [JSON Lines](https://jsonlines.org/), and it often comes with the `.json`, `.jsonl.gz`, or `.jsonl.bz2` extension.
+
+Our sample data contains four rows:
+
+```{code-cell} ipython3
+from pathlib import Path
+import json
+
+data = [
+    {"name": "John", "age": 25,  "friends": ["Jake", "Kelly"], "likes": {"pizza": True, "tacos": True}},
+    {"name": "Jake", "age": 20,  "friends": ["John"], "likes": {"pizza": False, "tacos": True}},
+    {"name": "Kelly", "age": 21,  "friends": ["John", "Sam"], "likes": {"pizza": True, "tacos": True}},
+    {"name": "Sam", "age": 22,  "friends": ["Kelly"], "likes": {"pizza": False, "tacos": True}},
+]
+
+lines = ""
+
+for d in data:
+    lines += json.dumps(d) + "\n"
+
+_ = Path("people.jsonl").write_text(lines)
+```
+
+```{code-cell} ipython3
+print(lines)
+```
+
+## Query
+
+```{note}
+Documentation for DuckDB's JSON capabilities is available [here](https://duckdb.org/docs/extensions/json.html).
+```
+
+Load the extension and start a DuckDB in-memory database:
+
+```{code-cell} ipython3
+%load_ext sql
+%sql duckdb://
+```
+
+Read the JSON data:
+
+```{code-cell} ipython3
+%%sql
+SELECT *
+FROM read_json_objects('people.jsonl')
+```
+
+## Extract fields
+
+Extract fields from each JSON record:
+
+```{code-cell} ipython3
+%%sql
+SELECT
+    json ->> '$.name' AS name,
+    json ->> '$.friends[0]' AS first_friend,
+    json ->> '$.likes.pizza' AS likes_pizza,
+    json ->> '$.likes.tacos' AS likes_tacos,
+FROM read_json_objects('people.jsonl')
+```
+
+Looks like everybody likes tacos!
+
++++
+
+## Extract schema
+
+Infer the JSON schema:
+
+```{code-cell} ipython3
+%%sql
+SELECT
+    json_structure(json),
+    json_structure(json ->> '$.likes'),
+FROM read_json_objects('people.jsonl')
+```
+
+```{code-cell} ipython3
+%%sql schema <<
+SELECT
+    json_structure(json) AS schema_all,
+    json_structure(json ->> '$.likes') AS schema_likes,
+FROM read_json_objects('people.jsonl')
+```
+
+Pretty print the inferred schema:
+
+```{code-cell} ipython3
+from rich import print_json
+
+row = schema.DataFrame().iloc[0]
+
+print("Schema:")
+print_json(row.schema_all)
+
+print("\n\nSchema (likes):")
+print_json(row.schema_likes)
+```
+
+## Store snippets
+
+You can use JupySQL's `--save` feature to store a SQL snippet so you can keep your queries succint:
+
+```{code-cell} ipython3
+%%sql --save clean_data
+SELECT
+    json ->> '$.name' AS name,
+    json ->> '$.friends[0]' AS first_friend,
+    json ->> '$.likes.pizza' AS likes_pizza,
+    json ->> '$.likes.tacos' AS likes_tacos,
+FROM read_json_objects('people.jsonl')
+```
+
+```{code-cell} ipython3
+%%sql --with clean_data
+SELECT * FROM clean_data
+```
+
+## Export to CSV
+
+```{note}
+Using `--with` isn't supported when exporting to CSV.
+```
+
+To export to CSV:
+
+```{code-cell} ipython3
+%%sql
+COPY (
+    SELECT
+    json ->> '$.name' AS name,
+    json ->> '$.friends[0]' AS first_friend,
+    json ->> '$.likes.pizza' AS likes_pizza,
+    json ->> '$.likes.tacos' AS likes_tacos,
+    FROM read_json_objects('people.jsonl')
+)
+
+TO 'people.csv' (HEADER, DELIMITER ',');
+```
+
+```{code-cell} ipython3
+%%sql
+SELECT * FROM 'people.csv'
+```
diff --git a/src/sql/connection.py b/src/sql/connection.py
index 49da5b2..4ba0fd1 100644
--- a/src/sql/connection.py
+++ b/src/sql/connection.py
@@ -1,11 +1,9 @@
 import os
+from difflib import get_close_matches
 
 import sqlalchemy
 from sqlalchemy.engine import Engine
-
-
-class ConnectionError(Exception):
-    pass
+from IPython.core.error import UsageError
 
 
 def rough_dict_get(dct, sought, default=None):
@@ -39,15 +37,68 @@ class Connection:
     connections = {}
 
     @classmethod
-    def tell_format(cls):
+    def _suggest_fix(cls, env_var, connect_str=None):
         """
         Returns an error message that we can display to the user
         to tell them how to pass the connection string
         """
-        return """Connection info needed in SQLAlchemy format, example:
-               postgresql://username:password@hostname/dbname
-               or an existing connection: %s""" % str(
-            cls.connections.keys()
+        DEFAULT_PREFIX = "\n\n"
+
+        if connect_str:
+            matches = get_close_matches(connect_str, list(cls.connections), n=1)
+
+            if matches:
+                prefix = (
+                    "\n\nPerhaps you meant to use the existing "
+                    f"connection: %sql {matches[0]!r}?\n\n"
+                )
+
+            else:
+                prefix = DEFAULT_PREFIX
+        else:
+            matches = None
+            prefix = DEFAULT_PREFIX
+
+        connection_string = (
+            "Pass a valid connection string:\n    "
+            "Example: %sql postgresql://username:password@hostname/dbname"
+        )
+
+        suffix = "To fix it:" if not matches else "Otherwise, try the following:"
+        options = [f"{prefix}{suffix}", connection_string]
+
+        keys = list(cls.connections.keys())
+
+        if keys:
+            keys_ = ",".join(repr(k) for k in keys)
+            options.append(
+                f"Pass a connection key (one of: {keys_})"
+                f"\n    Example: %sql {keys[0]!r}"
+            )
+
+        if env_var:
+            options.append("Set the environment variable $DATABASE_URL")
+
+        if len(options) >= 3:
+            options.insert(-1, "OR")
+
+        options.append(
+            "For technical support: https://ploomber.io/community"
+            "\nDocumentation: https://jupysql.ploomber.io/en/latest/connecting.html"
+        )
+
+        return "\n\n".join(options)
+
+    @classmethod
+    def _error_no_connection(cls):
+        """Error when there isn't any connection"""
+        return UsageError("No active connection." + cls._suggest_fix(env_var=True))
+
+    @classmethod
+    def _error_invalid_connection_info(cls, e, connect_str):
+        return UsageError(
+            "An error happened while creating the connection: "
+            f"{e}.{cls._suggest_fix(env_var=False, connect_str=connect_str)}"
         )
 
     def __init__(self, engine, alias=None):
@@ -79,9 +130,8 @@ class Connection:
                     connect_str,
                     connect_args=connect_args,
                 )
-        except Exception:
-            print(cls.tell_format())
-            raise
+        except Exception as e:
+            raise cls._error_invalid_connection_info(e, connect_str) from e
 
         connection = cls(engine, alias=alias)
         connection.connect_args = connect_args
@@ -91,9 +141,10 @@ class Connection:
     @classmethod
     def set(cls, descriptor, displaycon, connect_args=None, creator=None, alias=None):
         """
-        Sets the current database connection
+        Set the current database connection. This method is called from the magic to
+        determine which connection to use (either use an existing one or open a new one)
         """
-        connect_args = connect_args or connect_args
+        connect_args = connect_args or {}
 
         if descriptor:
             if isinstance(descriptor, Connection):
@@ -116,24 +167,23 @@ class Connection:
                     creator=creator,
                     alias=alias,
                 )
+
         else:
 
             if cls.connections:
                 if displaycon:
+                    # display list of connections
                     print(cls.connection_list())
+            elif os.getenv("DATABASE_URL"):
+                cls.current = Connection.from_connect_str(
+                    connect_str=os.getenv("DATABASE_URL"),
+                    connect_args=connect_args,
+                    creator=creator,
+                    alias=alias,
+                )
             else:
-                if os.getenv("DATABASE_URL"):
-                    cls.current = Connection.from_connect_str(
-                        connect_str=os.getenv("DATABASE_URL"),
-                        connect_args=connect_args,
-                        creator=creator,
-                        alias=alias,
-                    )
-                else:
-                    raise ConnectionError(
-                        "Environment variable $DATABASE_URL "
-                        "not set, and no connect string given."
-                    )
+                raise cls._error_no_connection()
+
         return cls.current
 
     @classmethod
@@ -143,6 +193,7 @@ class Connection:
 
     @classmethod
     def connection_list(cls):
+        """Returns the list of connections, appending '*' to the current one"""
         result = []
         for key in sorted(cls.connections):
             conn = cls.connections[key]
diff --git a/src/sql/magic.py b/src/sql/magic.py
index 8b26463..d281a8c 100644
--- a/src/sql/magic.py
+++ b/src/sql/magic.py
@@ -266,20 +266,15 @@ class SqlMagic(Magics, Configurable):
         if args.creator:
             args.creator = user_ns[args.creator]
 
-        try:
-            # this creates a new connection or use an existing one
-            # depending on the connect_arg value
-            conn = sql.connection.Connection.set(
-                connect_arg,
-                displaycon=self.displaycon,
-                connect_args=args.connection_arguments,
-                creator=args.creator,
-                alias=args.alias,
-            )
-        except Exception as e:
-            print(e)
-            print(sql.connection.Connection.tell_format())
-            return None
+        # this creates a new connection or use an existing one
+        # depending on the connect_arg value
+        conn = sql.connection.Connection.set(
+            connect_arg,
+            displaycon=self.displaycon,
+            connect_args=args.connection_arguments,
+            creator=args.creator,
+            alias=args.alias,
+        )
 
         if args.persist:
             return self._persist_dataframe(


Test patch:
diff --git a/src/tests/conftest.py b/src/tests/conftest.py
index 4ee30a4..a3934a1 100644
--- a/src/tests/conftest.py
+++ b/src/tests/conftest.py
@@ -8,6 +8,7 @@ from IPython.core.interactiveshell import InteractiveShell
 from sql.magic import SqlMagic, RenderMagic
 from sql.magic_plot import SqlPlotMagic
 from sql.magic_cmd import SqlCmdMagic
+from sql.connection import Connection
 
 PATH_TO_TESTS = Path(__file__).absolute().parent
 PATH_TO_TMP_ASSETS = PATH_TO_TESTS / "tmp"
@@ -41,17 +42,30 @@ def runsql(ip_session, statements):
 
 
 @pytest.fixture
-def ip():
-    """Provides an IPython session in which tables have been created"""
+def clean_conns():
+    Connection.current = None
+    Connection.connections = dict()
+    yield
+
+
+@pytest.fixture
+def ip_empty():
+
     ip_session = InteractiveShell()
     ip_session.register_magics(SqlMagic)
     ip_session.register_magics(RenderMagic)
     ip_session.register_magics(SqlPlotMagic)
     ip_session.register_magics(SqlCmdMagic)
+    yield ip_session
+
+
+@pytest.fixture
+def ip(ip_empty):
+    """Provides an IPython session in which tables have been created"""
 
     # runsql creates an inmemory sqlitedatabase
     runsql(
-        ip_session,
+        ip_empty,
         [
             "CREATE TABLE test (n INT, name TEXT)",
             "INSERT INTO test VALUES (1, 'foo')",
@@ -62,9 +76,9 @@ def ip():
             "CREATE TABLE empty_table (column INT, another INT)",
         ],
     )
-    yield ip_session
-    runsql(ip_session, "DROP TABLE test")
-    runsql(ip_session, "DROP TABLE author")
+    yield ip_empty
+    runsql(ip_empty, "DROP TABLE test")
+    runsql(ip_empty, "DROP TABLE author")
 
 
 @pytest.fixture
diff --git a/src/tests/test_magic.py b/src/tests/test_magic.py
index f86059a..88d7d5c 100644
--- a/src/tests/test_magic.py
+++ b/src/tests/test_magic.py
@@ -7,6 +7,7 @@ from textwrap import dedent
 
 import pytest
 from sqlalchemy import create_engine
+from IPython.core.error import UsageError
 
 from sql.connection import Connection
 from sql.magic import SqlMagic
@@ -463,6 +464,106 @@ def test_autolimit(ip):
     assert len(result) == 1
 
 
+invalid_connection_string = """
+No active connection.
+
+To fix it:
+
+Pass a valid connection string:
+    Example: %sql postgresql://username:password@hostname/dbname
+
+OR
+
+Set the environment variable $DATABASE_URL
+
+For technical support: https://ploomber.io/community
+Documentation: https://jupysql.ploomber.io/en/latest/connecting.html
+"""
+
+
+def test_error_on_invalid_connection_string(ip_empty, clean_conns):
+    result = ip_empty.run_cell("%sql some invalid connection string")
+
+    assert invalid_connection_string.strip() == str(result.error_in_exec)
+    assert isinstance(result.error_in_exec, UsageError)
+
+
+invalid_connection_string_format = """\
+An error happened while creating the connection: Can't load plugin: sqlalchemy.dialects:something.
+
+To fix it:
+
+Pass a valid connection string:
+    Example: %sql postgresql://username:password@hostname/dbname
+
+For technical support: https://ploomber.io/community
+Documentation: https://jupysql.ploomber.io/en/latest/connecting.html
+"""  # noqa
+
+
+def test_error_on_invalid_connection_string_format(ip_empty, clean_conns):
+    result = ip_empty.run_cell("%sql something://")
+
+    assert invalid_connection_string_format.strip() == str(result.error_in_exec)
+    assert isinstance(result.error_in_exec, UsageError)
+
+
+invalid_connection_string_existing_conns = """
+An error happened while creating the connection: Can't load plugin: sqlalchemy.dialects:something.
+
+To fix it:
+
+Pass a valid connection string:
+    Example: %sql postgresql://username:password@hostname/dbname
+
+OR
+
+Pass a connection key (one of: 'sqlite://')
+    Example: %sql 'sqlite://'
+
+For technical support: https://ploomber.io/community
+Documentation: https://jupysql.ploomber.io/en/latest/connecting.html
+"""  # noqa
+
+
+def test_error_on_invalid_connection_string_with_existing_conns(ip_empty, clean_conns):
+    ip_empty.run_cell("%sql sqlite://")
+    result = ip_empty.run_cell("%sql something://")
+
+    assert invalid_connection_string_existing_conns.strip() == str(result.error_in_exec)
+    assert isinstance(result.error_in_exec, UsageError)
+
+
+invalid_connection_string_with_possible_typo = """
+An error happened while creating the connection: Can't load plugin: sqlalchemy.dialects:sqlit.
+
+Perhaps you meant to use the existing connection: %sql 'sqlite://'?
+
+Otherwise, try the following:
+
+Pass a valid connection string:
+    Example: %sql postgresql://username:password@hostname/dbname
+
+OR
+
+Pass a connection key (one of: 'sqlite://')
+    Example: %sql 'sqlite://'
+
+For technical support: https://ploomber.io/community
+Documentation: https://jupysql.ploomber.io/en/latest/connecting.html
+"""  # noqa
+
+
+def test_error_on_invalid_connection_string_with_possible_typo(ip_empty, clean_conns):
+    ip_empty.run_cell("%sql sqlite://")
+    result = ip_empty.run_cell("%sql sqlit://")
+
+    assert invalid_connection_string_with_possible_typo.strip() == str(
+        result.error_in_exec
+    )
+    assert isinstance(result.error_in_exec, UsageError)
+
+
 def test_jupysql_alias():
     assert SqlMagic.magics == {
         "line": {"jupysql": "execute", "sql": "execute"},