From 08e09f81cccc5a40836e641f4a87ac0297f9c562 Mon Sep 17 00:00:00 2001 From: Akshay Joshi Date: Tue, 17 Dec 2024 17:28:49 +0530 Subject: [PATCH] Close the connection and delete adhoc server if all it's Query tool and PSQL connections are closed. --- docs/en_US/release_notes_9_0.rst | 1 + .../browser/server_groups/servers/utils.py | 7 ++-- web/pgadmin/misc/workspaces/__init__.py | 33 ++++++++++++++++++- web/pgadmin/tools/psql/__init__.py | 33 ++++++++++++++++--- web/pgadmin/tools/sqleditor/__init__.py | 4 +++ 5 files changed, 71 insertions(+), 7 deletions(-) diff --git a/docs/en_US/release_notes_9_0.rst b/docs/en_US/release_notes_9_0.rst index 0e5c99a56cb..7580a8108ec 100644 --- a/docs/en_US/release_notes_9_0.rst +++ b/docs/en_US/release_notes_9_0.rst @@ -29,4 +29,5 @@ Housekeeping Bug fixes ********* + | `Issue #8142 `_ - Correct the documentation for the MFA configuration. | `Issue #8208 `_ - Allow deleting the entry while creating/adding new label to enumeration type. \ No newline at end of file diff --git a/web/pgadmin/browser/server_groups/servers/utils.py b/web/pgadmin/browser/server_groups/servers/utils.py index 3f2115da6fe..7d0cd78b14c 100644 --- a/web/pgadmin/browser/server_groups/servers/utils.py +++ b/web/pgadmin/browser/server_groups/servers/utils.py @@ -650,12 +650,15 @@ def disconnect_from_all_servers(): manager.release() -def delete_adhoc_servers(): +def delete_adhoc_servers(sid=None): """ This function will remove all the adhoc servers. """ try: - db.session.query(Server).filter(Server.is_adhoc == 1).delete() + if sid is not None: + db.session.query(Server).filter(Server.id == sid).delete() + else: + db.session.query(Server).filter(Server.is_adhoc == 1).delete() db.session.commit() # Reset the sequence again diff --git a/web/pgadmin/misc/workspaces/__init__.py b/web/pgadmin/misc/workspaces/__init__.py index e148941b998..546757f719e 100644 --- a/web/pgadmin/misc/workspaces/__init__.py +++ b/web/pgadmin/misc/workspaces/__init__.py @@ -10,18 +10,21 @@ """A blueprint module implementing the workspace.""" import json import config +from config import PG_DEFAULT_DRIVER from flask import request, current_app from pgadmin.user_login_check import pga_login_required from flask_babel import gettext from flask_security import current_user from pgadmin.utils import PgAdminModule from pgadmin.model import db, Server +from pgadmin.utils.driver import get_driver from pgadmin.utils.ajax import bad_request, make_json_response from pgadmin.browser.server_groups.servers.utils import ( is_valid_ipaddress, convert_connection_parameter, check_ssl_fields) from pgadmin.tools.schema_diff.node_registry import SchemaDiffRegistry from pgadmin.browser.server_groups.servers.utils import ( disconnect_from_all_servers, delete_adhoc_servers) +from pgadmin.tools.psql import get_open_psql_connections MODULE_NAME = 'workspace' @@ -34,7 +37,8 @@ def get_exposed_url_endpoints(self): list: URL endpoints for Workspace module """ return [ - 'workspace.adhoc_connect_server' + 'workspace.adhoc_connect_server', + 'workspace.layout_changed' ] @@ -188,3 +192,30 @@ def layout_changed(): delete_adhoc_servers() return make_json_response(status=200) + + +def check_and_delete_adhoc_server(sid): + """ + This function is used to check for adhoc server and if all Query Tool + and PSQL connections are closed then delete that server. + """ + server = Server.query.filter_by(id=sid).first() + if server.is_adhoc: + # Check PSQL connections. If more connections are open for + # the given sid return from the function. + psql_connections = get_open_psql_connections() + if sid in psql_connections.values(): + return + + # Check Query Tool connections for the given sid + manager = get_driver(PG_DEFAULT_DRIVER).connection_manager(sid) + for key, value in manager.connections.items(): + if key.startswith('CONN') and value.connected(): + return + + # Assumption at this point all the Query Tool and PSQL connections + # is closed, so now we can release the manager + manager.release() + + # Delete the adhoc server from the pgadmin database + delete_adhoc_servers(sid) diff --git a/web/pgadmin/tools/psql/__init__.py b/web/pgadmin/tools/psql/__init__.py index e8495d76e6b..575fa8a18ab 100644 --- a/web/pgadmin/tools/psql/__init__.py +++ b/web/pgadmin/tools/psql/__init__.py @@ -45,6 +45,7 @@ session_input = dict() pdata = dict() cdata = dict() +open_psql_connections = dict() class PSQLModule(PgAdminModule): @@ -175,7 +176,7 @@ def get_user_env(): return env -def create_pty_terminal(connection_data): +def create_pty_terminal(connection_data, server_id): # Create the pty terminal process, parent and fd are file descriptors # for parent and child. parent, fd = pty.openpty() @@ -194,6 +195,7 @@ def create_pty_terminal(connection_data): app.config['sessions'][request.sid] = parent pdata[request.sid] = p cdata[request.sid] = fd + open_psql_connections[request.sid] = server_id else: app.config['sessions'][request.sid] = parent cdata[request.sid] = fd @@ -234,7 +236,7 @@ def read_stdout(process, sid, max_read_bytes, win_emit_output=True): sio.sleep(0) -def windows_platform(connection_data, sid, max_read_bytes): +def windows_platform(connection_data, sid, max_read_bytes, server_id): process = PtyProcess.spawn('cmd.exe', env=get_user_env()) process.write(r'"{0}" "{1}" 2>>&1'.format(connection_data[0], @@ -242,6 +244,7 @@ def windows_platform(connection_data, sid, max_read_bytes): process.write("\r\n") app.config['sessions'][request.sid] = process pdata[request.sid] = process + open_psql_connections[request.sid] = server_id set_term_size(process, 50, 50) while True: @@ -278,9 +281,10 @@ def non_windows_platform(parent, p, fd, data, max_read_bytes, sid): def pty_handel_io(connection_data, data, sid): max_read_bytes = 1024 * 20 if _platform == 'win32': - windows_platform(connection_data, sid, max_read_bytes) + windows_platform(connection_data, sid, max_read_bytes, + int(data['sid'])) else: - p, parent, fd = create_pty_terminal(connection_data) + p, parent, fd = create_pty_terminal(connection_data, int(data['sid'])) non_windows_platform(parent, p, fd, data, max_read_bytes, sid) @@ -567,18 +571,31 @@ def server_disconnect(data): disconnect_socket() +def cleanup_globals(): + del pdata[request.sid] + del cdata[request.sid] + server_id = open_psql_connections[request.sid] + del open_psql_connections[request.sid] + # Check if all the connections of the adhoc server is closed + # then delete the server from the pgadmin database. + from pgadmin.misc.workspaces import check_and_delete_adhoc_server + check_and_delete_adhoc_server(server_id) + + def disconnect_socket(): if _platform == 'win32': if request.sid in app.config['sessions']: process = app.config['sessions'][request.sid] process.terminate() del app.config['sessions'][request.sid] + cleanup_globals() else: os.write(app.config['sessions'][request.sid], r'\q\n'.encode()) sio.sleep(1) os.close(app.config['sessions'][request.sid]) os.close(cdata[request.sid]) del app.config['sessions'][request.sid] + cleanup_globals() def get_connection_status(conn): @@ -607,3 +624,11 @@ def _get_database_role(sid, did): return {'db_name': db_name, 'role': role} except Exception: return None + + +def get_open_psql_connections(): + """ + This function returns open connections + """ + + return open_psql_connections diff --git a/web/pgadmin/tools/sqleditor/__init__.py b/web/pgadmin/tools/sqleditor/__init__.py index ddac4941d4a..cf1595e48a1 100644 --- a/web/pgadmin/tools/sqleditor/__init__.py +++ b/web/pgadmin/tools/sqleditor/__init__.py @@ -64,6 +64,7 @@ get_explain_query_length from pgadmin.browser.server_groups.servers.utils import \ convert_connection_parameter +from pgadmin.misc.workspaces import check_and_delete_adhoc_server MODULE_NAME = 'sqleditor' TRANSACTION_STATUS_CHECK_FAILED = gettext("Transaction status check failed.") @@ -703,6 +704,9 @@ def close_sqleditor_session(trans_id): 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) + # Check if all the connections of the adhoc server is + # closed then delete the server from the pgadmin database. + check_and_delete_adhoc_server(cmd_obj.sid) # Close the auto complete connection if hasattr(cmd_obj, 'conn_id_ac') and cmd_obj.conn_id_ac is not None: