Repository: ploomber/jupysql
Pull request #446 (2023-04-25 19:11:35)

fixes --alias when passing an existing engine
## Describe your changes


Fixed a bug when using existing engine objects and passing `--alias`

## Issue number

Closes #X

## Checklist before requesting a review

- [ ] Performed a self-review of my code
- [ ] Formatted my code with [`pkgmt format`](https://ploomber-contributing.readthedocs.io/en/latest/contributing/pr.html#linting-formatting)
- [ ] Added [tests](https://ploomber-contributing.readthedocs.io/en/latest/contributing/pr.html#testing) (when necessary).
- [ ] 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--446.org.readthedocs.build/en/446/

<!-- readthedocs-preview jupysql end -->

edublancas: > looks good, missing integration tests?

This change is vendor-independent. since passing an alias is just setting a key in an internal mapping (I think we should have some guidelines on where to add integration tests since I've seen many cases where they're redundant)

> Any docs needed besides the existing one?

No, this doesn't modify anything. It fixes something that should work as the docs mention it.
yafimvo: lgtm! just rebase

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 143cf53..a4f5955 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## 0.7.3dev
 * [Doc] Tutorial on querying excel files with pandas and jupysql ([#423](https://github.com/ploomber/jupysql/pull/423))
+* [Fix] Fix `--alias` when passing an existing engine
 
 ## 0.7.2 (2023-04-25)
 
diff --git a/src/sql/command.py b/src/sql/command.py
index c633075..62b43b0 100644
--- a/src/sql/command.py
+++ b/src/sql/command.py
@@ -49,6 +49,7 @@ class SQLCommand:
             add_alias = True
         else:
             add_alias = False
+
         self.command_text = " ".join(line_for_command) + "\n" + cell
 
         if self.args.file:
diff --git a/src/sql/connection.py b/src/sql/connection.py
index d9aeb8f..dba05ed 100644
--- a/src/sql/connection.py
+++ b/src/sql/connection.py
@@ -270,7 +270,7 @@ class Connection:
             if isinstance(descriptor, Connection):
                 cls.current = descriptor
             elif isinstance(descriptor, Engine):
-                cls.current = Connection(descriptor)
+                cls.current = Connection(descriptor, alias=alias)
             elif is_custom_connection_:
                 cls.current = CustomConnection(descriptor, alias=alias)
             else:


Test patch:
diff --git a/src/tests/test_magic.py b/src/tests/test_magic.py
index 75e8bb9..a98360e 100644
--- a/src/tests/test_magic.py
+++ b/src/tests/test_magic.py
@@ -1,3 +1,4 @@
+import sqlite3
 import platform
 from pathlib import Path
 import os.path
@@ -504,6 +505,23 @@ def test_close_connection(ip, tmp_empty):
     assert "sqlite:///two.db" not in Connection.connections
 
 
+def test_alias(clean_conns, ip_empty, tmp_empty):
+    ip_empty.run_cell("%sql sqlite:///one.db --alias one")
+    assert {"one"} == set(Connection.connections)
+
+
+def test_alias_existing_engine(clean_conns, ip_empty, tmp_empty):
+    ip_empty.user_global_ns["first"] = create_engine("sqlite:///first.db")
+    ip_empty.run_cell("%sql first --alias one")
+    assert {"one"} == set(Connection.connections)
+
+
+def test_alias_custom_connection(clean_conns, ip_empty, tmp_empty):
+    ip_empty.user_global_ns["first"] = sqlite3.connect(":memory:")
+    ip_empty.run_cell("%sql first --alias one")
+    assert {"one"} == set(Connection.connections)
+
+
 def test_close_connection_with_alias(ip, tmp_empty):
     # open two connections
     ip.run_cell("%sql sqlite:///one.db --alias one")
@@ -519,6 +537,42 @@ def test_close_connection_with_alias(ip, tmp_empty):
     assert "two" not in Connection.connections
 
 
+def test_close_connection_with_existing_engine_and_alias(ip, tmp_empty):
+    ip.user_global_ns["first"] = create_engine("sqlite:///first.db")
+    ip.user_global_ns["second"] = create_engine("sqlite:///second.db")
+
+    # open two connections
+    ip.run_cell("%sql first --alias one")
+    ip.run_cell("%sql second --alias two")
+
+    # close them
+    ip.run_cell("%sql -x one")
+    ip.run_cell("%sql --close two")
+
+    assert "sqlite:///first.db" not in Connection.connections
+    assert "sqlite:///second.db" not in Connection.connections
+    assert "first" not in Connection.connections
+    assert "second" not in Connection.connections
+
+
+def test_close_connection_with_custom_connection_and_alias(ip, tmp_empty):
+    ip.user_global_ns["first"] = sqlite3.connect("first.db")
+    ip.user_global_ns["second"] = sqlite3.connect("second.db")
+
+    # open two connections
+    ip.run_cell("%sql first --alias one")
+    ip.run_cell("%sql second --alias two")
+
+    # close them
+    ip.run_cell("%sql -x one")
+    ip.run_cell("%sql --close two")
+
+    assert "sqlite:///first.db" not in Connection.connections
+    assert "sqlite:///second.db" not in Connection.connections
+    assert "first" not in Connection.connections
+    assert "second" not in Connection.connections
+
+
 def test_column_names_visible(ip, tmp_empty):
     res = ip.run_line_magic("sql", "SELECT * FROM empty_table")