Repository: ploomber/jupysql (python)

Issue:
detecting when a db driver is missing
In https://github.com/ploomber/jupysql/pull/120, we improved the error message displayed when failing to open a connection to a database. However, we're ignoring the error message content and printing a generic message. This leads to a confusing message when the connection string is valid, and the failure is due to another reason.

There are X reasons a connection might fail, but we should at least cover the most common ones. For example, missing a driver ([see this](https://github.com/ploomber/jupysql/pull/120#pullrequestreview-1289598402)). In such case, we should display a more appropriate message like "fix it by installing X package and try again"

I ran a few tests with some sample [URIs](https://docs.sqlalchemy.org/en/20/core/engines.html#postgresql) and the error is always the same `ModuleNotFoundError: No module named 'name'`.

The second case we must cover is when the plugin is missing. For example, if we run:

```
create_engine("duckdb://")
```

with duckdb-engine installed but duckdb *uninstalled*, we get: `ModuleNotFoundError: No module named 'duckdb'` (same as the use case described above).

However, if we uninstall duckdb-engine (regardless of duckdb installed or not), we get:

```
NoSuchModuleError: Can't load plugin: sqlalchemy.dialects:duckdb
```

`NoSuchModuleError` could be either a typo or an actual plugin missing (in this case, the duckdb-engine plugin is missing). So we should keep a list of plugins to distinguish between the two.



Pull request description:
Fix: detecting when a db driver is missing
## Describe your changes

* Provide the suggestion message when the driver package is not found
  - Correct driver name, but not install in the user environment
  - Closely matching to correct driver name, but not quite correct
  - Neither the above cases and incorrect driver name

* Test cases to cover those 

### Case 1 - Missing driver

Correct driver name `pymysql`, no `pymysql` package installed

1. 
```
%sql mysql+pymysql://database_url
```

2.
```
UsageError: No module named 'pymysql'

To fix it, try to install package: pymysql

For technical support: https://ploomber.io/community
Documentation: https://jupysql.ploomber.io/en/latest/connecting.html
```

### Case 2 - Closely correct but incorrect driver name `pymysqll`

1. 
```
%sql mysql+pymysqll://
```

2.
```
UsageError: Can't load plugin: sqlalchemy.dialects:mysql.pymysqll

To fix it, perhaps you meant to use driver name: pymysql

For technical support: https://ploomber.io/community
Documentation: https://jupysql.ploomber.io/en/latest/connecting.html
```

### Case 3 - Neither close or correct driver name `something`

1. 
```
%sql mysql+something://
```

2.
```
UsageError: Can't load plugin: sqlalchemy.dialects:mysql.something

To fix it, make sure you are using correct driver name:
Ref: https://docs.sqlalchemy.org/en/20/core/engines.html#database-urls

For technical support: https://ploomber.io/community
Documentation: https://jupysql.ploomber.io/en/latest/connecting.html
```

## Issue ticket number and link
Closes #124

## Checklist before requesting a review

- [ ] I have performed a self-review of my code
- [ ] 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--127.org.readthedocs.build/en/127/

<!-- readthedocs-preview jupysql end -->

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b669fa0..0fc0d2a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ## 0.5.6dev
 
+* [Feature] Shows missing driver package suggestion message ([#124](https://github.com/ploomber/jupysql/issues/124))
+
 ## 0.5.5 (2023-02-08)
 
 * [Fix] Clearer error message on connection failure ([#120](https://github.com/ploomber/jupysql/issues/120))
diff --git a/src/sql/connection.py b/src/sql/connection.py
index 4ba0fd1..79be479 100644
--- a/src/sql/connection.py
+++ b/src/sql/connection.py
@@ -3,7 +3,89 @@ from difflib import get_close_matches
 
 import sqlalchemy
 from sqlalchemy.engine import Engine
+from sqlalchemy.exc import NoSuchModuleError
 from IPython.core.error import UsageError
+import difflib
+
+PLOOMBER_SUPPORT_LINK_STR = (
+    "For technical support: https://ploomber.io/community"
+    "\nDocumentation: https://jupysql.ploomber.io/en/latest/connecting.html"
+)
+
+# Check Full List: https://docs.sqlalchemy.org/en/20/dialects
+MISSING_PACKAGE_LIST_EXCEPT_MATCHERS = {
+    # SQLite
+    "sqlite": "sqlite",
+    "pysqlcipher3": "pysqlcipher3",
+    # DuckDB
+    "duckdb": "duckdb-engine",
+    # MySQL + MariaDB
+    "pymysql": "pymysql",
+    "mysqldb": "mysqlclient",
+    "mariadb": "mariadb",
+    "mysql": "mysql-connector-python",
+    "asyncmy": "asyncmy",
+    "aiomysql": "aiomysql",
+    "cymysql": "cymysql",
+    "pyodbc": "pyodbc",
+    # PostgreSQL
+    "psycopg2": "psycopg2",
+    "psycopg": "psycopg",
+    "pg8000": "pg8000",
+    "asyncpg": "asyncpg",
+    "psycopg2cffi": "psycopg2cffi",
+    # Oracle
+    "cx_oracle": "cx_oracle",
+    "oracledb": "oracledb",
+    # MSSQL
+    "pyodbc": "pyodbc",
+    "pymssql": "pymssql",
+}
+
+
+def extract_module_name_from_ModuleNotFoundError(e):
+    return e.name
+
+
+def extract_module_name_from_NoSuchModuleError(e):
+    return str(e).split(":")[-1].split(".")[-1]
+
+
+"""
+When there is ModuleNotFoundError or NoSuchModuleError case
+Three types of suggestions will be shown when the missing module name is:
+1. Excepted in the pre-defined map, suggest the user to install the driver pkg
+2. Closely matched to the pre-defined map, suggest the user to type correct driver name
+3. Not found in the pre-defined map, suggest user to use valid driver pkg
+"""
+
+
+def get_missing_package_suggestion_str(e):
+    suggestion_prefix = "To fix it, "
+    module_name = None
+    if isinstance(e, ModuleNotFoundError):
+        module_name = extract_module_name_from_ModuleNotFoundError(e)
+    elif isinstance(e, NoSuchModuleError):
+        module_name = extract_module_name_from_NoSuchModuleError(e)
+
+    module_name = module_name.lower()
+    # Excepted
+    for matcher, suggested_package in MISSING_PACKAGE_LIST_EXCEPT_MATCHERS.items():
+        if matcher == module_name:
+            return suggestion_prefix + "try to install package: " + suggested_package
+    # Closely matched
+    close_matches = difflib.get_close_matches(
+        module_name, MISSING_PACKAGE_LIST_EXCEPT_MATCHERS.keys()
+    )
+    if close_matches:
+        return "Perhaps you meant to use driver the dialect: \"{}\"".format(
+            close_matches[0]
+        )
+    # Not found
+    return (
+        suggestion_prefix + "make sure you are using correct driver name:\n"
+        "Ref: https://docs.sqlalchemy.org/en/20/core/engines.html#database-urls"
+    )
 
 
 def rough_dict_get(dct, sought, default=None):
@@ -15,7 +97,7 @@ def rough_dict_get(dct, sought, default=None):
     """
 
     sought = sought.split("@")
-    for (key, val) in dct.items():
+    for key, val in dct.items():
         if not any(s.lower() not in key.lower() for s in sought):
             return val
     return default
@@ -36,6 +118,16 @@ class Connection:
     # all connections
     connections = {}
 
+    @classmethod
+    def _suggest_fix_no_module_found(module_name):
+        DEFAULT_PREFIX = "\n\n"
+
+        prefix = DEFAULT_PREFIX
+        suffix = "To fix it:"
+        suggest_str = "Install X package and try again"
+        options = [f"{prefix}{suffix}", suggest_str]
+        return "\n\n".join(options)
+
     @classmethod
     def _suggest_fix(cls, env_var, connect_str=None):
         """
@@ -82,10 +174,7 @@ class Connection:
         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"
-        )
+        options.append(PLOOMBER_SUPPORT_LINK_STR)
 
         return "\n\n".join(options)
 
@@ -101,6 +190,10 @@ class Connection:
             f"{e}.{cls._suggest_fix(env_var=False, connect_str=connect_str)}"
         )
 
+    @classmethod
+    def _error_module_not_found(cls, e):
+        return ModuleNotFoundError("test")
+
     def __init__(self, engine, alias=None):
         self.dialect = engine.url.get_dialect()
         self.metadata = sqlalchemy.MetaData(bind=engine)
@@ -130,6 +223,17 @@ class Connection:
                     connect_str,
                     connect_args=connect_args,
                 )
+        except (ModuleNotFoundError, NoSuchModuleError) as e:
+            suggestion_str = get_missing_package_suggestion_str(e)
+            raise UsageError(
+                "\n\n".join(
+                    [
+                        str(e),
+                        suggestion_str,
+                        PLOOMBER_SUPPORT_LINK_STR,
+                    ]
+                )
+            ) from e
         except Exception as e:
             raise cls._error_invalid_connection_info(e, connect_str) from e
 
@@ -169,7 +273,6 @@ class Connection:
                 )
 
         else:
-
             if cls.connections:
                 if displaycon:
                     # display list of connections


Test patch:
diff --git a/src/tests/mock_pymysql.py b/src/tests/mock_pymysql.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/tests/test_connection.py b/src/tests/test_connection.py
index 6422e86..1e99d01 100644
--- a/src/tests/test_connection.py
+++ b/src/tests/test_connection.py
@@ -1,10 +1,9 @@
 import sys
-from unittest.mock import Mock
-
+from unittest.mock import Mock, patch
 import pytest
 from sqlalchemy.engine import Engine
-
 from sql.connection import Connection
+from IPython.core.error import UsageError
 
 
 @pytest.fixture
@@ -35,3 +34,52 @@ def test_alias(cleanup):
     Connection.from_connect_str("sqlite://", alias="some-alias")
 
     assert list(Connection.connections) == ["some-alias"]
+
+
+# Mock the missing package
+# Ref: https://stackoverflow.com/a/28361013
+def test_missing_duckdb_dependencies(cleanup, monkeypatch):
+    with patch.dict(sys.modules):
+        sys.modules["duckdb"] = None
+        sys.modules["duckdb-engine"] = None
+
+        with pytest.raises(UsageError) as error:
+            Connection.from_connect_str("duckdb://")
+        assert "try to install package: duckdb-engine" + str(error.value)
+
+
+@pytest.mark.parametrize(
+    "missing_pkg, except_missing_pkg_suggestion, connect_str",
+    [
+        # MySQL + MariaDB
+        ["pymysql", "pymysql", "mysql+pymysql://"],
+        ["mysqlclient", "mysqlclient", "mysql+mysqldb://"],
+        ["mariadb", "mariadb", "mariadb+mariadbconnector://"],
+        ["mysql-connector-python", "mysql-connector-python", "mysql+mysqlconnector://"],
+        ["asyncmy", "asyncmy", "mysql+asyncmy://"],
+        ["aiomysql", "aiomysql", "mysql+aiomysql://"],
+        ["cymysql", "cymysql", "mysql+cymysql://"],
+        ["pyodbc", "pyodbc", "mysql+pyodbc://"],
+        # PostgreSQL
+        ["psycopg2", "psycopg2", "postgresql+psycopg2://"],
+        ["psycopg", "psycopg", "postgresql+psycopg://"],
+        ["pg8000", "pg8000", "postgresql+pg8000://"],
+        ["asyncpg", "asyncpg", "postgresql+asyncpg://"],
+        ["psycopg2cffi", "psycopg2cffi", "postgresql+psycopg2cffi://"],
+        # Oracle
+        ["cx_oracle", "cx_oracle", "oracle+cx_oracle://"],
+        ["oracledb", "oracledb", "oracle+oracledb://"],
+        # MSSQL
+        ["pyodbc", "pyodbc", "mssql+pyodbc://"],
+        ["pymssql", "pymssql", "mssql+pymssql://"],
+    ],
+)
+def test_missing_driver(
+    missing_pkg, except_missing_pkg_suggestion, connect_str, monkeypatch
+):
+    with patch.dict(sys.modules):
+        sys.modules[missing_pkg] = None
+        with pytest.raises(UsageError) as error:
+            Connection.from_connect_str(connect_str)
+
+        assert "try to install package: " + missing_pkg in str(error.value)
diff --git a/src/tests/test_magic.py b/src/tests/test_magic.py
index 88d7d5c..dddc184 100644
--- a/src/tests/test_magic.py
+++ b/src/tests/test_magic.py
@@ -489,12 +489,10 @@ def test_error_on_invalid_connection_string(ip_empty, clean_conns):
 
 
 invalid_connection_string_format = """\
-An error happened while creating the connection: Can't load plugin: sqlalchemy.dialects:something.
+Can't load plugin: sqlalchemy.dialects:something
 
-To fix it:
-
-Pass a valid connection string:
-    Example: %sql postgresql://username:password@hostname/dbname
+To fix it, make sure you are using correct driver name:
+Ref: https://docs.sqlalchemy.org/en/20/core/engines.html#database-urls
 
 For technical support: https://ploomber.io/community
 Documentation: https://jupysql.ploomber.io/en/latest/connecting.html
@@ -509,17 +507,10 @@ def test_error_on_invalid_connection_string_format(ip_empty, clean_conns):
 
 
 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
+Can't load plugin: sqlalchemy.dialects:something
 
-OR
-
-Pass a connection key (one of: 'sqlite://')
-    Example: %sql 'sqlite://'
+To fix it, make sure you are using correct driver name:
+Ref: https://docs.sqlalchemy.org/en/20/core/engines.html#database-urls
 
 For technical support: https://ploomber.io/community
 Documentation: https://jupysql.ploomber.io/en/latest/connecting.html
@@ -535,19 +526,9 @@ def test_error_on_invalid_connection_string_with_existing_conns(ip_empty, clean_
 
 
 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
+Can't load plugin: sqlalchemy.dialects:sqlit
 
-Pass a connection key (one of: 'sqlite://')
-    Example: %sql 'sqlite://'
+Perhaps you meant to use driver the dialect: "sqlite"
 
 For technical support: https://ploomber.io/community
 Documentation: https://jupysql.ploomber.io/en/latest/connecting.html