codekingpro/portable-devtools

1##########################################################################2#3# pgAdmin 4 - PostgreSQL Tools4#5# Copyright (C) 2013 - 2024, The pgAdmin Development Team6# This software is released under the PostgreSQL Licence7#8#########################################################################9 10"""Server heartbeat manager."""11 12 13import threading14import datetime15import config16from flask import session, current_app17from flask_babel import gettext18 19 20def log_server_heartbeat(data):21    """Log Server Heartbeat."""22    from config import PG_DEFAULT_DRIVER23    from pgadmin.utils.driver import get_driver24    manager = get_driver(PG_DEFAULT_DRIVER).connection_manager(int(data['sid'])25                                                               )26 27    _server_heartbeat = getattr(current_app, '_pgadmin_server_heartbeat', {})28 29    if session.sid not in _server_heartbeat:30        _server_heartbeat[session.sid] = {}31 32    if not manager:33        stop_server_heartbeat(data)34        return False, gettext("Manager not found. Stopped Heartbeat logging.")35    else:36        _server_heartbeat[session.sid][data['sid']] = {37            'timestamp': datetime.datetime.now(),38            'conn': manager.connections39        }40        current_app.logger.debug(41            "Heartbeat logged for the session id##server id: {0}##{1}".format(42                session.sid, data['sid']))43 44        setattr(current_app, '_pgadmin_server_heartbeat', _server_heartbeat)45        return True, gettext("Heartbeat logged successfully.")46 47 48def stop_server_heartbeat(data):49    """Stop logging server heartbeat."""50    _server_heartbeat = getattr(current_app, '_pgadmin_server_heartbeat', {})51 52    if session.sid in _server_heartbeat and \53            data['sid'] in _server_heartbeat[session.sid]:54        _server_heartbeat[session.sid].pop(data['sid'])55 56        current_app.logger.debug(57            "Heartbeat logging stopped for the session"58            " id##server id: {0}##{1}".format(session.sid, data['sid']))59 60        setattr(current_app, '_pgadmin_server_heartbeat', _server_heartbeat)61 62    return True, gettext("Stopped Heartbeat logging.")63 64 65def get_server_heartbeat(server_id):66    _server_heartbeat = getattr(current_app, '_pgadmin_server_heartbeat', {})67 68    if session.sid in _server_heartbeat and server_id in _server_heartbeat[69            session.sid]:70        return _server_heartbeat[session.sid][server_id]71    else:72        return None73 74 75class ServerHeartbeatTimer():76    def __init__(self, sec, _app):77        def func_wrapper():78            self.t = threading.Timer(sec, func_wrapper)79            self.t.start()80            self.release_server_heartbeat()81        self.t = threading.Timer(sec, func_wrapper)82        self.t.daemon = True83        self.t.start()84        self._app = _app85 86    def release_server_heartbeat(self):87        with self._app.app_context():88            _server_heartbeat = getattr(self._app,89                                        '_pgadmin_server_heartbeat', {})90            if len(_server_heartbeat) > 0:91                for sess_id in list(_server_heartbeat):92                    for sid in list(_server_heartbeat[sess_id]):93                        last_heartbeat_time = _server_heartbeat[sess_id][sid][94                            'timestamp']95                        current_time = datetime.datetime.now()96                        diff = current_time - last_heartbeat_time97 98                        # Wait for 4 times then the timeout99                        if diff.total_seconds() > (100                                config.SERVER_HEARTBEAT_TIMEOUT * 4):101                            self._release_connections(102                                _server_heartbeat[sess_id][sid]['conn'],103                                sess_id, sid)104                            _server_heartbeat[sess_id].pop(sid)105                            if len(_server_heartbeat[sess_id]) == 0:106                                _server_heartbeat.pop(sess_id)107                setattr(self._app, '_pgadmin_server_heartbeat',108                        _server_heartbeat)109 110    @staticmethod111    def _release_connections(server_conn, sess_id, sid):112        for d in server_conn:113            try:114                # Release the connection only if it is connected115                if server_conn[d].wasConnected:116                    server_conn[d]._release()117                    # Reconnect on the reload118                    server_conn[d].wasConnected = True119                    current_app.logger.debug(120                        "Heartbeat not received. Released "121                        "connection for the session "122                        "id##server id: {0}##{1}".format(123                            sess_id, sid))124            except Exception as e:125                current_app.logger.exception(e)126 127    def cancel(self):128        self.t.cancel()129 130 131def init_app(app):132    setattr(app, '_pgadmin_server_heartbeat', {})133    ServerHeartbeatTimer(sec=config.SERVER_HEARTBEAT_TIMEOUT,134                         _app=app)135