Improved connection handling

pgcli_sublime.py

@@ -20,7 +20,7 @@
 completers = {}  # Dict mapping urls to pgcompleter objects
 completer_lock = Lock()
 
-executors = {}  # Dict mapping buffer ids to pgexecutor objects
+executors = {}  # Dict mapping view ids to pgexecutor objects
 executor_lock = Lock()
 
 recent_urls = []
@@ -89,6 +89,11 @@ def plugin_unloaded():
 
 
 class PgcliPlugin(sublime_plugin.EventListener):
+    def on_close(self, view):
+        executor = executors.pop(view.id(), None)
+        if executor:
+            executor.conn.close()
+
     def on_post_save_async(self, view):
         check_pgcli(view)
 
@@ -372,8 +377,8 @@ def check_pgcli(view):
         return
 
     with executor_lock:
-        buffer_id = view.buffer_id()
-        if buffer_id not in executors:
+        view_id = view.id()
+        if view_id not in executors:
             url = get(view, 'pgcli_url')
 
             if not url:
@@ -394,7 +399,7 @@ def check_pgcli(view):
                     status = 'ERROR CONNECTING TO {}'.format(url)
                     view.set_status('pgcli', status)
 
-                executors[buffer_id] = executor
+                executors[view_id] = executor
 
                 # Make sure we have a completer for the corresponding url
                 with completer_lock:
@@ -461,28 +466,35 @@ def run_sqls_async(view, sqls):
 
 
 def run_sql_async(view, sql, panel):
-    executor = executors[view.buffer_id()]
+    executor = executors[view.id()]
     logger.debug('Command: PgcliExecute: %r', sql)
     save_mode = get(view, 'pgcli_save_on_run_query_mode')
-
     # Make sure the output panel is visiblle
     sublime.active_window().run_command('pgcli_show_output_panel')
     # Put a leading datetime
     datestr = str(datetime.datetime.now()) + '\n\n'
     panel.run_command('append', {'characters': datestr, 'pos': 0})
-    results = executor.run(sql, pgspecial=special)
-    try:
-        for (title, cur, headers, status, _, _) in results:
-            fmt = format_output(title, cur, headers, status, 'psql')
-            out = ('\n'.join(fmt)
-                   + '\n\n' + str(datetime.datetime.now()) + '\n\n')
+    for attempts in range(2):
+        results = executor.run(sql, pgspecial=special)
+        try:
+            for (title, cur, headers, status, _, _) in results:
+                fmt = format_output(title, cur, headers, status, 'psql')
+                out = ('\n'.join(fmt)
+                       + '\n\n' + str(datetime.datetime.now()) + '\n\n')
+                panel.run_command('append', {'characters': out})
+        except psycopg2.DatabaseError as e:
+            success = False
+            out = str(e) + '\n\n' + str(datetime.datetime.now()) + '\n\n'
             panel.run_command('append', {'characters': out})
-    except psycopg2.DatabaseError as e:
-        success = False
-        out = str(e) + '\n\n' + str(datetime.datetime.now()) + '\n\n'
-        panel.run_command('append', {'characters': out})
-    else:
-        success = True
+            if executor.conn.closed and attempts == 0:
+                panel.run_command('append', {
+                    'characters': 'Reconnecting ...\n\n'})
+                executor.connect()
+            else:
+                break
+        else:
+            success = True
+            break
 
     if (view.file_name()
             and ((save_mode == 'always')