Close connections gracefully when the user logs out of pgAdmin. Fixes #3942

web/pgadmin/tools/datagrid/__init__.py

@@ -14,6 +14,7 @@ import simplejson as json
 import pickle
 import random
 
+from threading import Lock
 from flask import Response, url_for, session, request, make_response
 from werkzeug.useragents import UserAgent
 from flask import current_app as app
@@ -28,6 +29,8 @@ from pgadmin.model import Server
 from pgadmin.utils.driver import get_driver
 from pgadmin.utils.exception import ConnectionLost, SSHTunnelConnectionLost
 
+query_tool_close_session_lock = Lock()
+
 
 class DataGridModule(PgAdminModule):
     """
@@ -66,6 +69,20 @@ class DataGridModule(PgAdminModule):
             'datagrid.close'
         ]
 
+    def on_logout(self, user):
+        """
+        This is a callback function when user logout from pgAdmin
+        :param user:
+        :return:
+        """
+        with query_tool_close_session_lock:
+            if 'gridData' in session:
+                for trans_id in session['gridData']:
+                    close_query_tool_session(trans_id)
+
+                # Delete all grid data from session variable
+                del session['gridData']
+
 
 blueprint = DataGridModule(MODULE_NAME, __name__, static_url_path='/static')
 
@@ -392,31 +409,17 @@ def close(trans_id):
     if str(trans_id) not in grid_data:
         return make_json_response(data={'status': True})
 
-    cmd_obj_str = grid_data[str(trans_id)]['command_obj']
-    # Use pickle.loads function to get the command object
-    cmd_obj = pickle.loads(cmd_obj_str)
-
-    # if connection id is None then no need to release the connection
-    if cmd_obj.conn_id is not None:
+    with query_tool_close_session_lock:
         try:
-            manager = get_driver(
-                PG_DEFAULT_DRIVER).connection_manager(cmd_obj.sid)
-            conn = manager.connection(
-                did=cmd_obj.did, conn_id=cmd_obj.conn_id)
+            close_query_tool_session(trans_id)
+            # Remove the information of unique transaction id from the
+            # session variable.
+            grid_data.pop(str(trans_id), None)
+            session['gridData'] = grid_data
         except Exception as e:
             app.logger.error(e)
             return internal_server_error(errormsg=str(e))
 
-        # Release the connection
-        if conn.connected():
-            conn.cancel_transaction(cmd_obj.conn_id, cmd_obj.did)
-            manager.release(did=cmd_obj.did, conn_id=cmd_obj.conn_id)
-
-        # Remove the information of unique transaction id from the
-        # session variable.
-        grid_data.pop(str(trans_id), None)
-        session['gridData'] = grid_data
-
     return make_json_response(data={'status': True})
 
 
@@ -461,3 +464,29 @@ def script():
         status=200,
         mimetype="application/javascript"
     )
+
+
+def close_query_tool_session(trans_id):
+    """
+    This function is used to cancel the transaction and release the connection.
+
+    :param trans_id: Transaction id
+    :return:
+    """
+
+    cmd_obj_str = session['gridData'][str(trans_id)]['command_obj']
+    # Use pickle.loads function to get the command object
+    cmd_obj = pickle.loads(cmd_obj_str)
+
+    # if connection id is None then no need to release the connection
+    if cmd_obj.conn_id is not None:
+        manager = get_driver(
+            PG_DEFAULT_DRIVER).connection_manager(cmd_obj.sid)
+        if manager is not None:
+            conn = manager.connection(
+                did=cmd_obj.did, conn_id=cmd_obj.conn_id)
+
+            # Release the connection
+            if conn.connected():
+                conn.cancel_transaction(cmd_obj.conn_id, cmd_obj.did)
+                manager.release(did=cmd_obj.did, conn_id=cmd_obj.conn_id)

web/pgadmin/tools/debugger/__init__.py

@@ -15,6 +15,7 @@ import simplejson as json
 import random
 import re
 
+from threading import Lock
 from flask import url_for, Response, render_template, request, session, \
     current_app
 from flask_babelex import gettext
@@ -31,10 +32,10 @@ from pgadmin.utils.driver import get_driver
 
 from config import PG_DEFAULT_DRIVER
 from pgadmin.model import db, DebuggerFunctionArguments
-from pgadmin.utils.preferences import Preferences
 
 # Constants
 ASYNC_OK = 1
+debugger_close_session_lock = Lock()
 
 
 class DebuggerModule(PgAdminModule):
@@ -226,6 +227,21 @@ class DebuggerModule(PgAdminModule):
                 'debugger.poll_end_execution_result', 'debugger.poll_result'
                 ]
 
+    def on_logout(self, user):
+        """
+        This is a callback function when user logout from pgAdmin
+        :param user:
+        :return:
+        """
+        with debugger_close_session_lock:
+            if 'debuggerData' in session:
+                for trans_id in session['debuggerData']:
+                    close_debugger_session(trans_id)
+
+                # Delete the all debugger data from session variable
+                del session['debuggerData']
+                del session['functionData']
+
 
 blueprint = DebuggerModule(MODULE_NAME, __name__)
 
@@ -828,29 +844,19 @@ def close(trans_id):
     if 'debuggerData' not in session:
         return make_json_response(data={'status': True})
 
-    debugger_data = session['debuggerData']
     # Return from the function if transaction id not found
-    if str(trans_id) not in debugger_data:
+    if str(trans_id) not in session['debuggerData']:
         return make_json_response(data={'status': True})
 
-    obj = debugger_data[str(trans_id)]
-    try:
-        manager = get_driver(
-            PG_DEFAULT_DRIVER).connection_manager(obj['server_id'])
-        conn = manager.connection(
-            did=obj['database_id'], conn_id=obj['conn_id'])
-        conn.cancel_transaction(obj['conn_id'], obj['database_id'])
-        conn = manager.connection(
-            did=obj['database_id'], conn_id=obj['exe_conn_id'])
-        conn.cancel_transaction(obj['exe_conn_id'], obj['database_id'])
-        manager.release(conn_id=obj['conn_id'])
-        manager.release(conn_id=obj['exe_conn_id'])
-        # Delete the existing debugger data in session variable
-        del session['debuggerData'][str(trans_id)]
-        del session['functionData'][str(trans_id)]
-        return make_json_response(data={'status': True})
-    except Exception as e:
-        return internal_server_error(errormsg=str(e))
+    with debugger_close_session_lock:
+        try:
+            close_debugger_session(trans_id)
+            # Delete the existing debugger data in session variable
+            del session['debuggerData'][str(trans_id)]
+            del session['functionData'][str(trans_id)]
+            return make_json_response(data={'status': True})
+        except Exception as e:
+            return internal_server_error(errormsg=str(e))
 
 
 @blueprint.route(
@@ -2105,3 +2111,31 @@ def poll_result(trans_id):
             'result': result
         }
     )
+
+
+def close_debugger_session(trans_id):
+    """
+    This function is used to cancel the debugger transaction and
+    release the connection.
+
+    :param trans_id: Transaction id
+    :return:
+    """
+    dbg_obj = session['debuggerData'][str(trans_id)]
+
+    manager = get_driver(
+        PG_DEFAULT_DRIVER).connection_manager(dbg_obj['server_id'])
+
+    if manager is not None:
+        conn = manager.connection(
+            did=dbg_obj['database_id'], conn_id=dbg_obj['conn_id'])
+        if conn.connected():
+            conn.cancel_transaction(dbg_obj['conn_id'],
+                                    dbg_obj['database_id'])
+        conn = manager.connection(
+            did=dbg_obj['database_id'], conn_id=dbg_obj['exe_conn_id'])
+        if conn.connected():
+            conn.cancel_transaction(dbg_obj['exe_conn_id'],
+                                    dbg_obj['database_id'])
+        manager.release(conn_id=dbg_obj['conn_id'])
+        manager.release(conn_id=dbg_obj['exe_conn_id'])