let sql-magic handle identifying connections by uri or alias string, and persist pgcli object through multiple cli sessions

1 files changed, 25 insertions, 8 deletions

diff --git a/pgcli/magic.py b/pgcli/magic.py
index 112a6e9b..c7b558e6 100644
--- a/pgcli/magic.py
+++ b/pgcli/magic.py
@@ -1,4 +1,9 @@
-from pgcli.main import PGCli
+from .main import PGCli
+import sql.parse
+import sql.connection
+import logging
+
+_logger = logging.getLogger(__name__)
 
 def load_ipython_extension(ipython):
 
@@ -9,15 +14,27 @@ def load_ipython_extension(ipython):
         ipython.run_line_magic('load_ext', 'sql')
 
     #register our own magic
-    ipython.register_magic_function(pgcli_line_magic, 'line','pgcli')
+    ipython.register_magic_function(pgcli_line_magic, 'line', 'pgcli')
 
 def pgcli_line_magic(line):
+    _logger.debug('pgcli magic called: %r', line)
+    parsed = sql.parse.parse(line, {})
+    conn = sql.connection.Connection.get(parsed['connection'])
 
-    #for now, assume line is connection string e.g. postgres://localhost
-    uri = line
-
-    pgcli = PGCli()
-    pgcli.connect_uri(uri)
+    try:
+        #A corresponding pgcli object already exists
+        pgcli = conn._pgcli
+        _logger.debug('Reusing existing pgcli')
+    except AttributeError:
+        #I can't figure out how to get the underylying psycopg2 connection
+        #from the sqlalchemy connection, so just grab the url and make a
+        #new connection
+        pgcli = PGCli()
+        u = conn.session.engine.url
+        _logger.debug('New pgcli: %r', str(u))
+
+        pgcli.connect(u.database, u.host, u.username, u.port, u.password)
+        conn._pgcli = pgcli
 
     try:
         pgcli.run_cli()
@@ -30,5 +47,5 @@ def pgcli_line_magic(line):
     q = pgcli.query_history[-1]
     if q.successful:
         ipython = get_ipython()
-        return ipython.run_cell_magic('sql', uri, q.query)
+        return ipython.run_cell_magic('sql', line, q.query)