Expand environment variable in connection string

Author: catherinedevlin

src/sql/parse.py

@@ -1,3 +1,4 @@
+from os.path import expandvars
 import six
 from six.moves import configparser as CP
 from sqlalchemy.engine.url import URL
@@ -9,6 +10,7 @@ def parse(cell, config):
     parts = [part.strip() for part in cell.split(None, 1)]
     if not parts:
         return {'connection': '', 'sql': ''}
+    parts[0] = expandvars(parts[0])  # for environment variables
 
     if parts[0].startswith('[') and parts[0].endswith(']'):
         section = parts[0].lstrip('[').rstrip(']')

src/tests/test_parse.py

@@ -1,3 +1,4 @@
+import os
 from sql.parse import parse
 from six.moves import configparser
 try:
@@ -12,22 +13,29 @@ def test_parse_no_sql():
            {'connection': "will:longliveliz@localhost/shakes",
             'sql': '',
             'flags': default_flags}
-    
+
 def test_parse_with_sql():
-    assert parse("postgresql://will:longliveliz@localhost/shakes SELECT * FROM work", 
+    assert parse("postgresql://will:longliveliz@localhost/shakes SELECT * FROM work",
                  empty_config) == \
            {'connection': "postgresql://will:longliveliz@localhost/shakes",
             'sql': 'SELECT * FROM work',
-            'flags': default_flags}    
-    
+            'flags': default_flags}
+
 def test_parse_sql_only():
     assert parse("SELECT * FROM work", empty_config) == \
            {'connection': "",
             'sql': 'SELECT * FROM work',
-            'flags': default_flags} 
-    
+            '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',
-            'flags': default_flags}            
+            'flags': default_flags}
+
+def test_expand_environment_variables_in_connection():
+    os.environ['DATABASE_URL'] = 'postgresql:///shakes'
+    assert parse("$DATABASE_URL SELECT * FROM work", empty_config) == \
+            {'connection': "postgresql:///shakes",
+            'sql': 'SELECT * FROM work',
+            'flags': default_flags}