Merge branch 'master' of https://github.com/catherinedevlin/ipython-sql

Author: catherinedevlin

.gitignore

@@ -33,3 +33,6 @@ nosetests.xml
 .mr.developer.cfg
 .project
 .pydevproject
+
+# Pycharm
+/.idea

src/sql/connection.py

@@ -3,20 +3,23 @@
 class Connection(object):
     current = None
     connections = {}
+
     @classmethod
     def tell_format(cls):
-        return "Format: (postgresql|mysql)://username:password@hostname/dbname, or one of %s" \
-               % str(cls.connections.keys())
+        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):
         try:
             engine = sqlalchemy.create_engine(connect_str)
         except: # TODO: bare except; but what's an ArgumentError?
             print(self.tell_format())
-            raise 
+            raise
         self.dialect = engine.url.get_dialect()
         self.metadata = sqlalchemy.MetaData(bind=engine)
         self.name = self.assign_name(engine)
-        self.session = engine.connect() 
+        self.session = engine.connect()
         self.connections[self.name] = self
         self.connections[str(self.metadata.bind.url)] = self
         Connection.current = self
@@ -26,7 +29,7 @@ def get(cls, descriptor):
             cls.current = descriptor
         elif descriptor:
             conn = cls.connections.get(descriptor) or \
-                   cls.connections.get(descriptor.lower()) 
+                   cls.connections.get(descriptor.lower())
             if conn:
                 cls.current = conn
             else:

src/sql/magic.py

@@ -78,9 +78,10 @@ def execute(self, line, cell='', local_ns={}):
         user_ns.update(local_ns)
 
         parsed = sql.parse.parse('%s\n%s' % (line, cell), self)
+        flags = parsed['flags']
         conn = sql.connection.Connection.get(parsed['connection'])
-        first_word = parsed['sql'].split(None, 1)[:1]
-        if first_word and first_word[0].lower() == 'persist':
+
+        if flags['persist']:
             return self._persist_dataframe(parsed['sql'], conn, user_ns)
 
         try:
@@ -104,6 +105,13 @@ def execute(self, line, cell='', local_ns={}):
 
                 return None
             else:
+
+                if flags['result_var']:
+                    result_var = flags['result_var']
+                    print("Returning data to local variable {}".format(result_var))
+                    self.shell.user_ns.update({result_var: result})
+                    return None
+
                 #Return results into the default ipython _ variable
                 return result
 
@@ -120,13 +128,11 @@ def _persist_dataframe(self, raw, conn, user_ns):
         if not DataFrame:
             raise ImportError("Must `pip install pandas` to use DataFrames")
 
-        # Parse input to get name of DataFrame
-        pieces = raw.split()
-        if len(pieces) != 2:
-            raise SyntaxError("Format: %sql [connection] persist <DataFrameName>")
-        frame_name = pieces[1].strip(';')
+        frame_name = raw.strip(';')
 
-       # Get the DataFrame from the user namespace
+        # Get the DataFrame from the user namespace
+        if not frame_name:
+            raise SyntaxError('Syntax: %sql PERSIST <name_of_data_frame>')
         frame = eval(frame_name, user_ns)
         if not isinstance(frame, DataFrame) and not isinstance(frame, Series):
             raise TypeError('%s is not a Pandas DataFrame or Series' % frame_name)

src/sql/parse.py

@@ -27,5 +27,27 @@ def parse(cell, config):
     else:
         connection = ''
         sql = cell
+    flags, sql = parse_sql_flags(sql.strip())
     return {'connection': connection.strip(),
-            'sql': sql.strip()}
+            'sql': sql,
+            'flags': flags}
+
+
+def parse_sql_flags(sql):
+    words = sql.split()
+    flags = {
+        'persist': False,
+        'result_var': None
+    }
+    if not words:
+        return (flags, "")
+    num_words = len(words)
+    trimmed_sql = sql
+    if words[0].lower() == 'persist':
+        print("Persist parsed to True")
+        flags['persist'] = True
+        trimmed_sql =  " ".join(words[1:])
+    elif num_words >= 2 and words[1] == '<<':
+        flags['result_var'] = words[0]
+        trimmed_sql = " ".join(words[2:])
+    return (flags, trimmed_sql.strip())

src/tests/test_magic.py

@@ -1,4 +1,5 @@
 from nose import with_setup
+from nose.tools import raises
 from sql.magic import SqlMagic
 from textwrap import dedent
 import os.path
@@ -57,6 +58,16 @@ def test_multi_sql():
         """)
     assert 'Shakespeare' in str(result) and 'Brecht' in str(result)
 
+@with_setup(_setup_writer, _teardown_writer)
+def test_result_var():
+    ip.run_cell_magic('sql', '', """
+        sqlite://
+        x <<
+        SELECT last_name FROM writer;
+        """)
+    result = ip.user_global_ns['x']
+    assert 'Shakespeare' in str(result) and 'Brecht' in str(result)
+
 @with_setup(_setup_writer, _teardown_writer)
 def test_access_results_by_keys():
     assert ip.run_line_magic('sql', "sqlite:// SELECT * FROM writer;")['William'] == (u'William', u'Shakespeare', 1616)
@@ -78,7 +89,6 @@ def test_autolimit():
     result = ip.run_line_magic('sql',  "sqlite:// SELECT * FROM test;")
     assert len(result) == 1
 
-
 @with_setup(_setup, _teardown)
 def test_persist():
     ip.run_cell("results = %sql SELECT * FROM test;")
@@ -87,13 +97,37 @@ def test_persist():
     persisted = ip.run_line_magic('sql', 'SELECT * FROM results_dframe')
     assert 'foo' in str(persisted)
 
+@raises(NameError)
+def test_persist_nonexistent_raises():
+    ip.run_line_magic('sql',  "sqlite://")
+    ip.run_line_magic('sql', 'PERSIST no_such_dataframe')
+
+@raises(TypeError)
+def test_persist_non_frame_raises():
+    ip.run_cell("not_a_dataframe = 22")
+    ip.run_line_magic('sql', "sqlite://")
+    ip.run_line_magic('sql', 'PERSIST not_a_dataframe')
+
+@raises(SyntaxError)
+def test_persist_bare():
+    ip.run_line_magic('sql', "sqlite://")
+    ip.run_line_magic('sql', 'PERSIST')
+
 @with_setup(_setup_writer, _teardown_writer)
-def test_unnamed_persist():
+def test_persist_frame_at_its_creation():
     ip.run_cell("results = %sql SELECT * FROM writer;")
     ip.run_line_magic('sql', 'PERSIST results.DataFrame()')
     persisted = ip.run_line_magic('sql', 'SELECT * FROM results')
     assert 'Shakespeare' in str(persisted)
 
+# TODO: support 
+# @with_setup(_setup_writer, _teardown_writer)
+# def test_persist_with_connection_info():
+#     ip.run_cell("results = %sql SELECT * FROM writer;")
+#     ip.run_line_magic('sql', 'sqlite:// PERSIST results.DataFrame()')
+#     persisted = ip.run_line_magic('sql', 'SELECT * FROM results')
+#     assert 'Shakespeare' in str(persisted)
+
 @with_setup(_setup_writer, _teardown_writer)
 def test_displaylimit():
     ip.run_line_magic('config',  "SqlMagic.autolimit = 0")

src/tests/test_parse.py

@@ -6,24 +6,28 @@
     from IPython.config.configurable import Configurable
 
 empty_config = Configurable()
-
+default_flags = {'persist': False, 'result_var': None}
 def test_parse_no_sql():
     assert parse("will:longliveliz@localhost/shakes", empty_config) == \
            {'connection': "will:longliveliz@localhost/shakes",
-            'sql': ''}
+            'sql': '',
+            'flags': default_flags}
 
 def test_parse_with_sql():
     assert parse("postgresql://will:longliveliz@localhost/shakes SELECT * FROM work", 
                  empty_config) == \
            {'connection': "postgresql://will:longliveliz@localhost/shakes",
-            'sql': 'SELECT * FROM work'}    
+            'sql': 'SELECT * FROM work',
+            'flags': default_flags}    
 
 def test_parse_sql_only():
     assert parse("SELECT * FROM work", empty_config) == \
            {'connection': "",
-            'sql': 'SELECT * FROM work'} 
+            'sql': 'SELECT * FROM work',
+            'flags': default_flags} 
 
 def test_parse_postgresql_socket_connection():
     assert parse("postgresql:///shakes SELECT * FROM work", empty_config) == \
            {'connection': "postgresql:///shakes",
-            'sql': 'SELECT * FROM work'}            
+            'sql': 'SELECT * FROM work',
+            'flags': default_flags}