diff --git a/web/pgadmin/__init__.py b/web/pgadmin/__init__.py index ff465d7..49e76d0 100644 --- a/web/pgadmin/__init__.py +++ b/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 diff --git a/web/pgadmin/tools/datagrid/__init__.py b/web/pgadmin/tools/datagrid/__init__.py index 85ab5da..3736e38 100644 --- a/web/pgadmin/tools/datagrid/__init__.py +++ b/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) diff --git a/web/pgadmin/tools/debugger/__init__.py b/web/pgadmin/tools/debugger/__init__.py index 42fa12f..6ea816c 100644 --- a/web/pgadmin/tools/debugger/__init__.py +++ b/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']) diff --git a/web/pgadmin/utils/__init__.py b/web/pgadmin/utils/__init__.py index 277fcd3..653a607 100644 --- a/web/pgadmin/utils/__init__.py +++ b/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): """