parse connect_args from DATABASE_URL env var

Author: danykaufman

src/sql/connection.py

@@ -6,6 +6,8 @@
 from sqlalchemy.exc import NoSuchModuleError
 from IPython.core.error import UsageError
 import difflib
+import urllib.parse
+import ast
 
 PLOOMBER_SUPPORT_LINK_STR = (
     "For technical support: https://ploomber.io/community"
@@ -278,12 +280,23 @@ def set(cls, descriptor, displaycon, connect_args=None, creator=None, alias=None
                     # display list of connections
                     print(cls.connection_list())
             elif os.getenv("DATABASE_URL"):
-                cls.current = Connection.from_connect_str(
-                    connect_str=os.getenv("DATABASE_URL"),
-                    connect_args=connect_args,
-                    creator=creator,
-                    alias=alias,
-                )
+                # try and extract connect_args from DATABASE_URL
+                if not connect_args and '&' in os.getenv("DATABASE_URL"):
+                    cls.current = Connection.from_connect_str(
+                        connect_str=os.getenv("DATABASE_URL").split('&')[0],
+                        connect_args=ast.literal_eval(
+                            '{' + urllib.parse.unquote_plus(os.getenv('DATABASE_URL').split('&')[1]).replace(
+                                'requests_kwargs=','"requests_kwargs":') + '}'),
+                        creator=creator,
+                        alias=alias,
+                    )
+                else:
+                    cls.current = Connection.from_connect_str(
+                        connect_str=os.getenv("DATABASE_URL"),
+                        connect_args=connect_args,
+                        creator=creator,
+                        alias=alias,
+                    )
             else:
                 raise cls._error_no_connection()
 

src/sql/run.py

@@ -385,6 +385,7 @@ def _commit(conn, config):
 
 def run(conn, sql, config, user_namespace):
     if sql.strip():
+        conn.dialect.supports_statement_cache = False # disable SQL compilation caching warning
         for statement in sqlparse.split(sql):
             first_word = sql.strip().split()[0].lower()
             if first_word == "begin":