Close the connection and delete adhoc server if all it's Query tool and PSQL connections are closed.

This commit is contained in:
Akshay Joshi 2024-12-17 17:57:25 +05:30 committed by GitHub
parent e84f4cd35b
commit 185e0f61e2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 71 additions and 7 deletions

View File

@ -29,4 +29,5 @@ Housekeeping
Bug fixes
*********
| `Issue #8142 <https://github.com/pgadmin-org/pgadmin4/issues/8142>`_ - Correct the documentation for the MFA configuration.
| `Issue #8208 <https://github.com/pgadmin-org/pgadmin4/issues/8208>`_ - Allow deleting the entry while creating/adding new label to enumeration type.

View File

@ -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

View File

@ -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 @@ class WorkspaceModule(PgAdminModule):
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)

View File

@ -45,6 +45,7 @@ else:
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

View File

@ -64,6 +64,7 @@ from pgadmin.tools.sqleditor.utils.apply_explain_plan_wrapper import \
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: