Repository: ploomber/jupysql (python)

Issue:
Support passing existing connections
The initialization locks the context to a specific engine. [See example](https://github.com/catherinedevlin/ipython-sql/issues/216)

We'd like to allow users to connect to an engine directly by reference (similar to SQL alchemy).

Pull request description:
allow to pass existing `sqlalchemy.engine.Engine`
Closes #16

I also added more unit tests and created `SQLCommand` to encapsulate the input parsing logic (which allows us to test it more easily)

Patch:
diff --git a/doc/connecting.md b/doc/connecting.md
index d08a9b1..09f6b62 100644
--- a/doc/connecting.md
+++ b/doc/connecting.md
@@ -6,7 +6,7 @@ jupytext:
     extension: .md
     format_name: myst
     format_version: 0.13
-    jupytext_version: 1.14.0
+    jupytext_version: 1.14.4
 kernelspec:
   display_name: Python 3 (ipykernel)
   language: python
@@ -74,6 +74,35 @@ then you can:
 %sql --section DB_CONFIG_1 
 ```
 
++++
+
+## Using an existing `sqlalchemy.engine.Engine`
+
+Use an existing `Engine` by passing the variable name to `%sql`.
+
 ```{code-cell} ipython3
+import pandas as pd
+from sqlalchemy.engine import create_engine
+```
 
+```{code-cell} ipython3
+engine = create_engine("sqlite://")
+```
+
+```{code-cell} ipython3
+df = pd.DataFrame({"x": range(5)})
+df.to_sql("numbers", engine)
+```
+
+```{code-cell} ipython3
+%load_ext sql
+```
+
+```{code-cell} ipython3
+%sql engine
+```
+
+```{code-cell} ipython3
+%%sql
+SELECT * FROM numbers
 ```
diff --git a/src/sql/command.py b/src/sql/command.py
new file mode 100644
index 0000000..2b06e1c
--- /dev/null
+++ b/src/sql/command.py
@@ -0,0 +1,66 @@
+from sqlalchemy.engine import Engine
+
+from sql import parse
+from sql.store import store
+
+
+class SQLCommand:
+    """
+    Encapsulates the parsing logic (arguments, SQL code, connection string, etc.)
+
+    """
+
+    def __init__(self, magic, user_ns, line, cell) -> None:
+        self.args = parse.magic_args(magic.execute, line)
+
+        if (
+            len(self.args.line) == 1
+            and self.args.line[0] in user_ns
+            and isinstance(user_ns[self.args.line[0]], Engine)
+        ):
+            line_for_command = []
+            add_conn = True
+        else:
+            line_for_command = self.args.line
+            add_conn = False
+
+        self.command_text = " ".join(line_for_command) + "\n" + cell
+
+        if self.args.file:
+            with open(self.args.file, "r") as infile:
+                self.command_text = infile.read() + "\n" + self.command_text
+
+        self.parsed = parse.parse(self.command_text, magic)
+
+        self.parsed["sql_original"] = self.parsed["sql"]
+
+        if add_conn:
+            self.parsed["connection"] = user_ns[self.args.line[0]]
+
+        if self.args.with_:
+            final = store.render(self.parsed["sql"], with_=self.args.with_)
+            self.parsed["sql"] = str(final)
+
+    @property
+    def sql(self):
+        """
+        Returns the SQL query to execute, without any other options or arguments
+        """
+        return self.parsed["sql"]
+
+    @property
+    def sql_original(self):
+        """
+        Returns the raw SQL query. Might be different from `sql` if using --with
+        """
+        return self.parsed["sql_original"]
+
+    @property
+    def connection(self):
+        """Returns the connection string"""
+        return self.parsed["connection"]
+
+    @property
+    def result_var(self):
+        """Returns the result_var"""
+        return self.parsed["result_var"]
diff --git a/src/sql/connection.py b/src/sql/connection.py
index e84a8d0..d64e27e 100644
--- a/src/sql/connection.py
+++ b/src/sql/connection.py
@@ -1,6 +1,7 @@
 import os
 
 import sqlalchemy
+from sqlalchemy.engine import Engine
 
 
 class ConnectionError(Exception):
@@ -22,19 +23,47 @@ def rough_dict_get(dct, sought, default=None):
     return default
 
 
-class Connection(object):
+class Connection:
+    """Manages connections to databases
+
+    Parameters
+    ----------
+    engine: sqlalchemy.engine.Engine
+        The SQLAlchemy engine to use
+    """
+
+    # the active connection
     current = None
+
+    # all connections
     connections = {}
 
     @classmethod
     def tell_format(cls):
+        """
+        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()
         )
 
-    def __init__(self, connect_str=None, connect_args={}, creator=None):
+    def __init__(self, engine):
+        self.dialect = engine.url.get_dialect()
+        self.metadata = sqlalchemy.MetaData(bind=engine)
+        self.name = self.assign_name(engine)
+        self.session = engine.connect()
+        self.connections[repr(self.metadata.bind.url)] = self
+        self.connect_args = None
+        Connection.current = self
+
+    @classmethod
+    def from_connect_str(cls, connect_str=None, connect_args=None, creator=None):
+        """Creates a new connection from a connection string"""
+        connect_args = connect_args or {}
+
         try:
             if creator:
                 engine = sqlalchemy.create_engine(
@@ -45,27 +74,39 @@ class Connection(object):
                     connect_str, connect_args=connect_args
                 )
         except Exception:
-            print(self.tell_format())
+            print(cls.tell_format())
             raise
-        self.dialect = engine.url.get_dialect()
-        self.metadata = sqlalchemy.MetaData(bind=engine)
-        self.name = self.assign_name(engine)
-        self.session = engine.connect()
-        self.connections[repr(self.metadata.bind.url)] = self
-        self.connect_args = connect_args
-        Connection.current = self
+
+        connection = cls(engine)
+        connection.connect_args = connect_args
+
+        return connection
 
     @classmethod
-    def set(cls, descriptor, displaycon, connect_args={}, creator=None):
-        "Sets the current database connection"
+    def set(cls, descriptor, displaycon, connect_args=None, creator=None):
+        """
+        Sets the current database connection
+        """
+        connect_args = connect_args or connect_args
 
         if descriptor:
             if isinstance(descriptor, Connection):
                 cls.current = descriptor
+            elif isinstance(descriptor, Engine):
+                cls.current = Connection(descriptor)
             else:
                 existing = rough_dict_get(cls.connections, descriptor)
-            # http://docs.sqlalchemy.org/en/rel_0_9/core/engines.html#custom-dbapi-connect-arguments # noqa
-            cls.current = existing or Connection(descriptor, connect_args, creator)
+
+                # NOTE: I added one indentation level, otherwise
+                # the "existing" variable would not exist if
+                # passing an engine object as descriptor.
+                # Since I never saw this breaking, my guess
+                # is that we're missing some unit tests
+                # when descriptor is a connection object
+                # http://docs.sqlalchemy.org/en/rel_0_9/core/engines.html#custom-dbapi-connect-arguments # noqa
+                cls.current = existing or Connection.from_connect_str(
+                    descriptor, connect_args, creator
+                )
         else:
 
             if cls.connections:
@@ -73,7 +114,7 @@ class Connection(object):
                     print(cls.connection_list())
             else:
                 if os.getenv("DATABASE_URL"):
-                    cls.current = Connection(
+                    cls.current = Connection.from_connect_str(
                         os.getenv("DATABASE_URL"), connect_args, creator
                     )
                 else:
diff --git a/src/sql/magic.py b/src/sql/magic.py
index eaf10d5..4daae54 100644
--- a/src/sql/magic.py
+++ b/src/sql/magic.py
@@ -20,6 +20,7 @@ import sql.connection
 import sql.parse
 import sql.run
 from sql.store import store
+from sql.command import SQLCommand
 
 try:
     from traitlets.config.configurable import Configurable
@@ -217,39 +218,37 @@ class SqlMagic(Magics, Configurable):
           mysql+pymysql://me:mypw@localhost/mydb
 
         """
+        # line is the text after the magic, cell is the cell's body
 
-        # Parse variables (words wrapped in {}) for %%sql magic
-        # (for %sql this is done automatically)
-        cell = self.shell.var_expand(cell)
-        line = sql.parse.without_sql_comment(parser=self.execute.parser, line=line)
-        args = parse_argstring(self.execute, line)
+        # Examples
 
-        if args.connections:
-            return sql.connection.Connection.connections
-        elif args.close:
-            return sql.connection.Connection._close(args.close)
+        # %sql {line}
+        # note that line magic has no body
+
+        # %%sql {line}
+        # {cell}
 
         # save globals and locals so they can be referenced in bind vars
         user_ns = self.shell.user_ns.copy()
         user_ns.update(local_ns)
 
-        command_text = " ".join(args.line) + "\n" + cell
-
-        if args.file:
-            with open(args.file, "r") as infile:
-                command_text = infile.read() + "\n" + command_text
+        # Parse variables (words wrapped in {}) for %%sql magic
+        # (for %sql this is done automatically)
+        cell = self.shell.var_expand(cell)
 
-        parsed = sql.parse.parse(command_text, self)
+        command = SQLCommand(self, user_ns, line, cell)
+        # args.line: contains the line after the magic with all options removed
+        args = command.args
 
-        original = parsed["sql"]
+        if args.connections:
+            return sql.connection.Connection.connections
+        elif args.close:
+            return sql.connection.Connection._close(args.close)
 
-        if args.with_:
-            final = self._store.render(original, with_=args.with_)
-            parsed["sql"] = str(final)
+        connect_arg = command.connection
 
-        connect_str = parsed["connection"]
         if args.section:
-            connect_str = sql.parse.connection_from_dsn_section(args.section, self)
+            connect_arg = sql.parse.connection_from_dsn_section(args.section, self)
 
         if args.connection_arguments:
             try:
@@ -272,7 +271,7 @@ class SqlMagic(Magics, Configurable):
 
         try:
             conn = sql.connection.Connection.set(
-                connect_str,
+                connect_arg,
                 displaycon=self.displaycon,
                 connect_args=args.connection_arguments,
                 creator=args.creator,
@@ -284,27 +283,27 @@ class SqlMagic(Magics, Configurable):
 
         if args.persist:
             return self._persist_dataframe(
-                parsed["sql"], conn, user_ns, append=False, index=not args.no_index
+                command.sql, conn, user_ns, append=False, index=not args.no_index
             )
 
         if args.append:
             return self._persist_dataframe(
-                parsed["sql"], conn, user_ns, append=True, index=not args.no_index
+                command.sql, conn, user_ns, append=True, index=not args.no_index
             )
 
-        if not parsed["sql"]:
+        if not command.sql:
             return
 
         # store the query if needed
         if args.save:
-            self._store.store(args.save, original, with_=args.with_)
+            self._store.store(args.save, command.sql_original, with_=args.with_)
 
         if args.no_execute:
             print("Skipping execution...")
             return
 
         try:
-            result = sql.run.run(conn, parsed["sql"], self, user_ns)
+            result = sql.run.run(conn, command.sql, self, user_ns)
 
             if (
                 result is not None
@@ -330,9 +329,8 @@ class SqlMagic(Magics, Configurable):
                 return None
             else:
 
-                if parsed["result_var"]:
-                    result_var = parsed["result_var"]
-                    self.shell.user_ns.update({result_var: result})
+                if command.result_var:
+                    self.shell.user_ns.update({command.result_var: result})
                     return None
 
                 # Return results into the default ipython _ variable
diff --git a/src/sql/parse.py b/src/sql/parse.py
index 336efc5..6eb11a1 100644
--- a/src/sql/parse.py
+++ b/src/sql/parse.py
@@ -4,6 +4,7 @@ from os.path import expandvars
 
 from six.moves import configparser as CP
 from sqlalchemy.engine.url import URL
+from IPython.core.magic_arguments import parse_argstring
 
 
 def connection_from_dsn_section(section, config):
@@ -90,3 +91,8 @@ def without_sql_comment(parser, line):
         shlex.split(line, posix=False),
     )
     return " ".join(result)
+
+
+def magic_args(magic_execute, line):
+    line = without_sql_comment(parser=magic_execute.parser, line=line)
+    return parse_argstring(magic_execute, line)


Test patch:
diff --git a/src/tests/conftest.py b/src/tests/conftest.py
index 13b9ee3..18480da 100644
--- a/src/tests/conftest.py
+++ b/src/tests/conftest.py
@@ -1,3 +1,4 @@
+import os
 import urllib.request
 from pathlib import Path
 
@@ -44,6 +45,7 @@ def ip():
     ip_session.register_magics(SqlMagic)
     ip_session.register_magics(RenderMagic)
 
+    # runsql creates an inmemory sqlitedatabase
     runsql(
         ip_session,
         [
@@ -58,3 +60,15 @@ def ip():
     yield ip_session
     runsql(ip_session, "DROP TABLE test")
     runsql(ip_session, "DROP TABLE author")
+
+
+@pytest.fixture
+def tmp_empty(tmp_path):
+    """
+    Create temporary path using pytest native fixture,
+    them move it, yield, and restore the original path
+    """
+    old = os.getcwd()
+    os.chdir(str(tmp_path))
+    yield str(Path(tmp_path).resolve())
+    os.chdir(old)
diff --git a/src/tests/test_command.py b/src/tests/test_command.py
new file mode 100644
index 0000000..052859f
--- /dev/null
+++ b/src/tests/test_command.py
@@ -0,0 +1,167 @@
+from pathlib import Path
+
+import pytest
+from sqlalchemy import create_engine
+
+from sql.command import SQLCommand
+
+
+@pytest.fixture
+def sql_magic(ip):
+    return ip.magics_manager.lsmagic()["line"]["sql"].__self__
+
+
+@pytest.mark.parametrize(
+    "line, cell, parsed_sql, parsed_connection, parsed_result_var",
+    [
+        ("something --no-execute", "", "something\n", "", None),
+        ("sqlite://", "", "", "sqlite://", None),
+        ("SELECT * FROM TABLE", "", "SELECT * FROM TABLE\n", "", None),
+        ("SELECT * FROM", "TABLE", "SELECT * FROM\nTABLE", "", None),
+        ("my_var << SELECT *", "FROM table", "SELECT *\nFROM table", "", "my_var"),
+        ("[db]", "", "", "sqlite://", None),
+    ],
+    ids=[
+        "arg-with-option",
+        "connection-string",
+        "sql-query",
+        "sql-query-in-line-and-cell",
+        "parsed-var",
+        "config",
+    ],
+)
+def test_parsed(
+    ip,
+    sql_magic,
+    line,
+    cell,
+    parsed_sql,
+    parsed_connection,
+    parsed_result_var,
+    tmp_empty,
+):
+    # needed for the last test case
+    Path("odbc.ini").write_text(
+        """
+[db]
+drivername = sqlite
+"""
+    )
+
+    cmd = SQLCommand(sql_magic, ip.user_ns, line, cell)
+
+    assert cmd.parsed == {
+        "connection": parsed_connection,
+        "result_var": parsed_result_var,
+        "sql": parsed_sql,
+        "sql_original": parsed_sql,
+    }
+
+    assert cmd.connection == parsed_connection
+    assert cmd.sql == parsed_sql
+    assert cmd.sql_original == parsed_sql
+
+
+def test_parsed_sql_when_using_with(ip, sql_magic):
+    ip.run_cell_magic(
+        "sql",
+        "--save author_one",
+        """
+        SELECT * FROM author LIMIT 1
+        """,
+    )
+
+    cmd = SQLCommand(
+        sql_magic, ip.user_ns, line="--with author_one", cell="SELECT * FROM author_one"
+    )
+
+    sql = (
+        "WITH author_one AS (\n    \n\n        "
+        "SELECT * FROM author LIMIT 1\n        \n)"
+        "\n\nSELECT * FROM author_one"
+    )
+
+    sql_original = "\nSELECT * FROM author_one"
+
+    assert cmd.parsed == {
+        "connection": "",
+        "result_var": None,
+        "sql": sql,
+        "sql_original": sql_original,
+    }
+
+    assert cmd.connection == ""
+    assert cmd.sql == sql
+    assert cmd.sql_original == sql_original
+
+
+def test_parsed_sql_when_using_file(ip, sql_magic, tmp_empty):
+    Path("query.sql").write_text("SELECT * FROM author")
+    cmd = SQLCommand(sql_magic, ip.user_ns, "--file query.sql", "")
+
+    assert cmd.parsed == {
+        "connection": "",
+        "result_var": None,
+        "sql": "SELECT * FROM author\n\n",
+        "sql_original": "SELECT * FROM author\n\n",
+    }
+
+    assert cmd.connection == ""
+    assert cmd.sql == "SELECT * FROM author\n\n"
+    assert cmd.sql_original == "SELECT * FROM author\n\n"
+
+
+def test_args(ip, sql_magic):
+    ip.run_cell_magic(
+        "sql",
+        "--save author_one",
+        """
+        SELECT * FROM author LIMIT 1
+        """,
+    )
+
+    cmd = SQLCommand(sql_magic, ip.user_ns, line="--with author_one", cell="")
+
+    assert cmd.args.__dict__ == {
+        "line": "",
+        "connections": False,
+        "close": None,
+        "creator": None,
+        "section": None,
+        "persist": False,
+        "no_index": False,
+        "append": False,
+        "connection_arguments": None,
+        "file": None,
+        "save": None,
+        "with_": ["author_one"],
+        "no_execute": False,
+    }
+
+
+@pytest.mark.parametrize(
+    "line",
+    [
+        "my_engine",
+        " my_engine",
+        "my_engine ",
+    ],
+)
+def test_parse_sql_when_passing_engine(ip, sql_magic, tmp_empty, line):
+    engine = create_engine("sqlite:///my.db")
+    ip.user_global_ns["my_engine"] = engine
+
+    cmd = SQLCommand(sql_magic, ip.user_ns, line, cell="SELECT * FROM author")
+
+    sql_expected = "\nSELECT * FROM author"
+
+    assert cmd.parsed == {
+        "connection": engine,
+        "result_var": None,
+        "sql": sql_expected,
+        "sql_original": sql_expected,
+    }
+
+    assert cmd.connection is engine
+    assert cmd.sql == sql_expected
+    assert cmd.sql_original == sql_expected
diff --git a/src/tests/test_magic.py b/src/tests/test_magic.py
index fe7e235..c412465 100644
--- a/src/tests/test_magic.py
+++ b/src/tests/test_magic.py
@@ -4,6 +4,7 @@ import tempfile
 from textwrap import dedent
 
 import pytest
+from sqlalchemy import create_engine
 
 from conftest import runsql
 
@@ -375,6 +376,24 @@ def test_close_connection(ip):
     assert connection_name not in connections_afterward
 
 
+def test_pass_existing_engine(ip, tmp_empty):
+    ip.user_global_ns["my_engine"] = create_engine("sqlite:///my.db")
+    ip.run_line_magic("sql", "  my_engine ")
+
+    runsql(
+        ip,
+        [
+            "CREATE TABLE some_data (n INT, name TEXT)",
+            "INSERT INTO some_data VALUES (10, 'foo')",
+            "INSERT INTO some_data VALUES (20, 'bar')",
+        ],
+    )
+
+    result = ip.run_line_magic("sql", "SELECT * FROM some_data")
+
+    assert result == [(10, "foo"), (20, "bar")]
+
+
 # theres some weird shared state with this one, moving it to the end
 def test_autolimit(ip):
     ip.run_line_magic("config", "SqlMagic.autolimit = 0")
diff --git a/src/tests/test_parse.py b/src/tests/test_parse.py
index 35a37d6..5c4df62 100644
--- a/src/tests/test_parse.py
+++ b/src/tests/test_parse.py
@@ -2,7 +2,14 @@ import os
 from pathlib import Path
 
 
-from sql.parse import connection_from_dsn_section, parse, without_sql_comment
+import pytest
+
+from sql.parse import (
+    connection_from_dsn_section,
+    parse,
+    without_sql_comment,
+    magic_args,
+)
 
 try:
     from traitlets.config.configurable import Configurable
@@ -176,3 +183,49 @@ def test_without_sql_persist():
     line = "--persist my_table --uff da"
     expected = "--persist my_table"
     assert without_sql_comment(parser=parser_stub, line=line) == expected
+
+
+def complete_with_defaults(mapping):
+    defaults = {
+        "line": ["some-argument"],
+        "connections": False,
+        "close": None,
+        "creator": None,
+        "section": None,
+        "persist": False,
+        "no_index": False,
+        "append": False,
+        "connection_arguments": None,
+        "file": None,
+        "save": None,
+        "with_": None,
+        "no_execute": False,
+    }
+
+    return {**defaults, **mapping}
+
+
+@pytest.mark.parametrize(
+    "line, expected",
+    [
+        (
+            "some-argument",
+            {"line": ["some-argument"]},
+        ),
+        (
+            "a b c",
+            {"line": ["a", "b", "c"]},
+        ),
+        (
+            "a b c --file query.sql",
+            {"line": ["a", "b", "c"], "file": "query.sql"},
+        ),
+    ],
+)
+def test_magic_args(ip, line, expected):
+
+    sql_line = ip.magics_manager.lsmagic()["line"]["sql"]
+
+    args = magic_args(sql_line, line)
+
+    assert args.__dict__ == complete_with_defaults(expected)