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

docs/en_US/release_notes_4_3.rst

@@ -9,4 +9,5 @@ This release contains a number of fixes reported since the release of pgAdmin4 4
 Bug fixes
 *********
 
+| `Bug #3942 <https://redmine.postgresql.org/issues/3942>`_ - Close connections gracefully when the user logs out of pgAdmin.
 | `Bug #3963 <https://redmine.postgresql.org/issues/3963>`_ - Fix alignment of import/export toggle switch.

web/pgadmin/__init__.py

@@ -12,6 +12,7 @@ such as setup of logging, dynamic loading of modules etc."""
 import logging
 import os
 import sys
+from types import MethodType
 from collections import defaultdict
 from importlib import import_module
 
@@ -56,6 +57,8 @@ class PgAdmin(Flask):
             extensions=['jinja2.ext.autoescape', 'jinja2.ext.with_'],
             loader=VersionedTemplateLoader(self)
         )
+        self.logout_hooks = []
+
         super(PgAdmin, self).__init__(*args, **kwargs)
 
     def find_submodules(self, basemodule):
@@ -161,6 +164,11 @@ class PgAdmin(Flask):
                           for key, value in menu_items.items())
         return menu_items
 
+    def register_logout_hook(self, module):
+        if hasattr(module, 'on_logout') and \
+                type(getattr(module, 'on_logout')) == MethodType:
+            self.logout_hooks.append(module)
+
 
 def _find_blueprint():
     if request.blueprint:
@@ -556,9 +564,17 @@ def create_app(app_name=None):
         session.force_write = True
 
     @user_logged_out.connect_via(app)
-    def clear_current_user_connections(app, user):
+    def current_user_cleanup(app, user):
         from config import PG_DEFAULT_DRIVER
         from pgadmin.utils.driver import get_driver
+        from flask import current_app
+
+        for mdl in current_app.logout_hooks:
+            try:
+                mdl.on_logout(user)
+            except Exception as e:
+                current_app.logger.exception(e)
+
         _driver = get_driver(PG_DEFAULT_DRIVER)
         _driver.gc_own()
 
@@ -568,6 +584,7 @@ def create_app(app_name=None):
     for module in app.find_submodules('pgadmin'):
         app.logger.info('Registering blueprint module: %s' % module)
         app.register_blueprint(module)
+        app.register_logout_hook(module)
 
     ##########################################################################
     # Handle the desktop login

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,30 +409,16 @@ 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']
+    with query_tool_close_session_lock:
-    # 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:
         try:
-            manager = get_driver(
+            close_query_tool_session(trans_id)
-                PG_DEFAULT_DRIVER).connection_manager(cmd_obj.sid)
-            conn = manager.connection(
-                did=cmd_obj.did, conn_id=cmd_obj.conn_id)
-        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
+        except Exception as e:
+            app.logger.error(e)
+            return internal_server_error(errormsg=str(e))
 
     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,23 +844,13 @@ 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)]
+    with debugger_close_session_lock:
         try:
-        manager = get_driver(
+            close_debugger_session(trans_id)
-            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)]
@@ -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'])

web/pgadmin/utils/__init__.py

@@ -66,6 +66,7 @@ class PgAdminModule(Blueprint):
             if first_registration:
                 module.parentmodules.append(self)
             app.register_blueprint(module)
+            app.register_logout_hook(module)
 
     def get_own_stylesheets(self):
         """