Connect to $DATABASE_URL if no connection string given

Author: catherinedevlin

README.rst

@@ -64,6 +64,10 @@ an existing connection by username@database
     ======================
     Poet       733
 
+If no connect string is supplied, ``5sql`` will provide a list of existing connections;
+however, if no connections have yet been made and the environment variable ``DATABASE_URL``
+is available, that will be used.
+
 For secure access, you may dynamically access your credentials (e.g. from your system environment or `getpass.getpass`) to avoid storing your password in the notebook itself. Use the `$` before any variable to access it in your `%sql` command.
 
 .. code-block:: python
@@ -250,10 +254,10 @@ specified) or in a file of the given name.
 PostgreSQL features
 -------------------
 
-``psql``-style "backslash" `meta-commands`_ commands (``\d``, ``\dt``, etc.) 
+``psql``-style "backslash" `meta-commands`_ commands (``\d``, ``\dt``, etc.)
 are provided by `PGSpecial`_.
 
-.. _PGSpecial: https://pypi.python.org/pypi/pgspecial 
+.. _PGSpecial: https://pypi.python.org/pypi/pgspecial
 
 .. _meta-commands: https://www.postgresql.org/docs/9.6/static/app-psql.html#APP-PSQL-META-COMMANDS
 

src/sql/connection.py

@@ -1,4 +1,9 @@
 import sqlalchemy
+import os
+
+class ConnectionError(Exception):
+    pass
+
 
 class Connection(object):
     current = None
@@ -23,27 +28,45 @@ def __init__(self, connect_str=None):
         self.connections[self.name] = self
         self.connections[str(self.metadata.bind.url)] = self
         Connection.current = self
+
     @classmethod
-    def get(cls, descriptor):
-        if isinstance(descriptor, Connection):
-            cls.current = descriptor
-        elif descriptor:
-            conn = cls.connections.get(descriptor) or \
-                   cls.connections.get(descriptor.lower())
-            if conn:
-                cls.current = conn
+    def set(cls, descriptor):
+        "Sets the current database connection"
+
+        if descriptor:
+            if isinstance(descriptor, Connection):
+                cls.current = descriptor
             else:
-                cls.current = Connection(descriptor)
-        if cls.current:
-            return cls.current
+                existing = cls.connections.get(descriptor) or \
+                           cls.connections.get(descriptor.lower())
+            cls.current = existing or Connection(descriptor)
         else:
-            raise Exception(cls.tell_format())
+            if cls.connections:
+                print(cls.connection_list())
+            else:
+                if os.getenv('DATABASE_URL'):
+                    cls.current = Connection(os.getenv('DATABASE_URL'))
+                else:
+                    raise ConnectionError('Environment variable $DATABASE_URL not set, and no connect string given.')
+        return cls.current
+
     @classmethod
     def assign_name(cls, engine):
-        core_name = '%s@%s' % (engine.url.username, engine.url.database)
+        core_name = '%s@%s' % (engine.url.username or '', engine.url.database)
         incrementer = 1
         name = core_name
         while name in cls.connections:
             name = '%s_%d' % (core_name, incrementer)
             incrementer += 1
         return name
+
+    @classmethod
+    def connection_list(cls):
+        result = []
+        for key in sorted(cls.connections):
+            if cls.connections[key] == cls.current:
+                template = ' * {}'
+            else:
+                template = '   {}'
+            result.append(template.format(key))
+        return '\n'.join(result)

src/sql/magic.py

@@ -80,9 +80,14 @@ def execute(self, line, cell='', local_ns={}):
 
         parsed = sql.parse.parse('%s\n%s' % (line, cell), self)
         flags = parsed['flags']
-        conn = sql.connection.Connection.get(parsed['connection'])
+        try:
+            conn = sql.connection.Connection.set(parsed['connection'])
+        except Exception as e:
+            print(e)
+            print(sql.connection.Connection.tell_format())
+            return None
 
-        if flags['persist']:
+        if flags.get('persist'):
             return self._persist_dataframe(parsed['sql'], conn, user_ns)
 
         try:
@@ -107,7 +112,7 @@ def execute(self, line, cell='', local_ns={}):
                 return None
             else:
 
-                if flags['result_var']:
+                if flags.get('result_var'):
                     result_var = flags['result_var']
                     print("Returning data to local variable {}".format(result_var))
                     self.shell.user_ns.update({result_var: result})

src/sql/parse.py

@@ -9,7 +9,7 @@ def parse(cell, config):
 
     parts = [part.strip() for part in cell.split(None, 1)]
     if not parts:
-        return {'connection': '', 'sql': ''}
+        return {'connection': '', 'sql': '', 'flags': {}}
     parts[0] = expandvars(parts[0])  # for environment variables
 
     if parts[0].startswith('[') and parts[0].endswith(']'):