Repository: ploomber/jupysql (python)

Issue:
error when detecting is connection is custom
(Note that this error can only be reproduced with https://github.com/ploomber/jupysql/pull/431, which hasn't been merged yet)

while improving some error messages, I encountered this:

![image](https://user-images.githubusercontent.com/989250/233754280-328dda08-ec54-402d-98e9-5a17e8f2fc37.png)

Note that the error message doesn't accurately display the error in the `df` variable. Instead, if shows `''`

I tracked down the error to this method:

https://github.com/ploomber/jupysql/blob/fcc74b0f56c007d04f1f0997373c79cba342f23f/src/sql/connection.py#L362

Since `int` is not part of the list of checks:

https://github.com/ploomber/jupysql/blob/fcc74b0f56c007d04f1f0997373c79cba342f23f/src/sql/connection.py#L381

It fails. Which breaks the `SQLCommand` logic (the `is_custom_connection` is called there).

We need a more robust way to check whether a connection is custom or not. A better approach is to check if the connection object has all the attributes expected from the DB API: close, commit, rollback and cursor (https://peps.python.org/pep-0249/)

Pull request description:
checking custom connections according to PEP 249
## Describe your changes
Added PEP 249 compliance to custom connection. 
Checking whether the connection object has `close`, `commit` and `rollback` methods
## Issue number

Closes #433 

## Checklist before requesting a review

- [X] Performed a self-review of my code
- [X] Formatted my code with [`pkgmt format`](https://ploomber-contributing.readthedocs.io/en/latest/contributing/pr.html#linting-formatting)
- [X] Added [tests](https://ploomber-contributing.readthedocs.io/en/latest/contributing/pr.html#testing) (when necessary).
- [X] Added [docstring](https://ploomber-contributing.readthedocs.io/en/latest/contributing/pr.html#documenting-changes-and-new-features) documentation and update the [changelog](https://ploomber-contributing.readthedocs.io/en/latest/contributing/pr.html#changelog) (when needed)



<!-- readthedocs-preview jupysql start -->
----
:books: Documentation preview :books:: https://jupysql--517.org.readthedocs.build/en/517/

<!-- readthedocs-preview jupysql end -->

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2ec53e7..d8cc63e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,9 +2,11 @@
 
 ## 0.7.6dev
 
+* [Fix] Fix error when checking if custom connection was PEP 249 Compliant (#517)
 * [Feature] Support for printing capture variables using `=<<` syntax (by [@jorisroovers](https://github.com/jorisroovers))
 * [Feature] Adds `--persist-replace` argument to replace existing tables when persisting data frames (#440)
 
+
 ## 0.7.5 (2023-05-24)
 
 * [Feature] Using native DuckDB `.df()` method when using `autopandas`
diff --git a/src/sql/connection.py b/src/sql/connection.py
index 1df3a5f..166d0ec 100644
--- a/src/sql/connection.py
+++ b/src/sql/connection.py
@@ -123,6 +123,29 @@ def rough_dict_get(dct, sought, default=None):
     return default
 
 
+def is_pep249_compliant(conn):
+    """
+    Checks if given connection object complies with PEP 249
+    """
+    pep249_methods = [
+        "close",
+        "commit",
+        # "rollback",
+        # "cursor",
+        # PEP 249 doesn't require the connection object to have
+        # a cursor method strictly
+        # ref: https://peps.python.org/pep-0249/#id52
+    ]
+
+    for method_name in pep249_methods:
+        # Checking whether the connection object has the method
+        # and if it is callable
+        if not hasattr(conn, method_name) or not callable(getattr(conn, method_name)):
+            return False
+
+    return True
+
+
 class Connection:
     """Manages connections to databases
 
@@ -418,14 +441,9 @@ class Connection:
         if isinstance(conn, (CustomConnection, CustomSession)):
             is_custom_connection_ = True
         else:
-            # TODO: Better check when user passes a custom
-            # connection
-            if (
-                isinstance(
-                    conn, (sqlalchemy.engine.base.Connection, Connection, str, bool)
-                )
-                or conn.__class__.__name__ == "DataFrame"
-            ):
+            if isinstance(
+                conn, (sqlalchemy.engine.base.Connection, Connection)
+            ) or not (is_pep249_compliant(conn)):
                 is_custom_connection_ = False
             else:
                 is_custom_connection_ = True


Test patch:
diff --git a/src/tests/test_connection.py b/src/tests/test_connection.py
index 50f9bb7..52d1d0b 100644
--- a/src/tests/test_connection.py
+++ b/src/tests/test_connection.py
@@ -4,10 +4,11 @@ import pytest
 from sqlalchemy import create_engine
 from sqlalchemy.engine import Engine
 import sql.connection
-from sql.connection import Connection
+from sql.connection import Connection, CustomConnection
 from IPython.core.error import UsageError
 import sqlglot
 import sqlalchemy
+import sqlite3
 
 
 @pytest.fixture
@@ -266,3 +267,41 @@ def test_no_current_connection_and_get_info(monkeypatch, mock_database):
 
     monkeypatch.setattr(conn, "session", None)
     assert conn._get_curr_sqlalchemy_connection_info() is None
+
+
+class dummy_connection:
+    def __init__(self):
+        self.engine_name = "dummy_engine"
+
+    def close(self):
+        pass
+
+
+@pytest.mark.parametrize(
+    "conn, expected",
+    [
+        [sqlite3.connect(""), True],
+        [
+            CustomConnection(engine=sqlalchemy.create_engine("sqlite://")),
+            True,
+        ],
+        [
+            Connection(engine=sqlalchemy.create_engine("sqlite://")),
+            False,
+        ],
+        [dummy_connection(), False],
+        ["not_a_valid_connection", False],
+        [0, False],
+    ],
+    ids=[
+        "sqlite3_connection",
+        "custom_connection",
+        "normal_connection",
+        "dummy_connection",
+        "str",
+        "int",
+    ],
+)
+def test_custom_connection(conn, expected):
+    is_custom = Connection.is_custom_connection(conn)
+    assert is_custom == expected