SessionPool.py

# --coding:utf-8--
#
# Copyright (c) 2022 vesoft inc. All rights reserved.
#
# This source code is licensed under Apache 2.0 License.
 
 
import json
import socket
 
from threading import RLock, Timer
import time
 
from nebula3.Exception import (
    AuthFailedException,
    NoValidSessionException,
    InValidHostname,
)
 
from nebula3.gclient.net.Session import Session
from nebula3.gclient.net.Connection import Connection
from nebula3.logger import logger
from nebula3.Config import SessionPoolConfig
 
 
class SessionPool(object):
    S_OK = 0
    S_BAD = 1
 
    def __init__(self, username, password, space_name, addresses):
        # user name and password of the session
        self._username_username = username
        self._password_password = password
 
        # space name bonded to the session
        self._space_name_space_name = space_name
 
        # all addresses of servers
        self._addresses_addresses = list()
 
        # server's status
        self._addresses_status_addresses_status = dict()
 
        # validate the addresses
        for address in addresses:
            try:
                ip = socket.gethostbyname(address[0])
            except Exception:
                raise InValidHostname(str(address[0]))
            ip_port = (ip, address[1])
            self._addresses_addresses.append(ip_port)
            self._addresses_status_addresses_status[ip_port] = self.S_BADS_BAD
 
        # sessions that are currently in use
        self._active_sessions_active_sessions = list()
        # sessions that are currently available
        self._idle_sessions_idle_sessions = list()
 
        self._configs_configs = SessionPoolConfig()
        self._ssl_configs_ssl_configs = None
        self._lock_lock = RLock()
 
        # the index of the next address to connect
        self._pos_pos = -1
 
        # the flag of whether the pool is closed
        self._close_close = False
 
    def __del__(self):
        self.closeclose()
 
    def init(self, configs):
        """init the session pool
 
        :param username: the username of the session
        :param password: the password of the session
        :param space_name: the space name of the session
        :param addresses: the addresses of the servers
        :param configs: the config of the pool
 
        :return: if all addresses are valid, return True else return False.
        """
        # check configs
        try:
            self._check_configs_check_configs()
        except Exception as e:
            logger.error('Invalid configs: {}'.format(e))
            return False
 
        if self._close_close:
            logger.error('The pool has init or closed.')
            raise RuntimeError('The pool has init or closed.')
        self._configs_configs = configs
 
        # ping all servers
        self.update_servers_statusupdate_servers_status()
 
        # check services status in the background
        self._period_detect_period_detect()
 
        ok_num = self.get_ok_servers_numget_ok_servers_num()
        if ok_num < len(self._addresses_addresses):
            raise RuntimeError(
                'The services status exception: {}'.format(self._get_services_status_get_services_status())
            )
 
        # iterate all addresses and create sessions to fullfil the min_size
        for i in range(self._configs_configs.min_size):
            session = self._new_session_new_session()
            if session is None:
                raise RuntimeError('Get session failed')
            self._add_session_to_idle_add_session_to_idle(session)
 
        return True
 
    def ping(self, address):
        """check the server is ok
 
        :param address: the server address want to connect
        :return: True or False
        """
        try:
            conn = Connection()
            if self._ssl_configs_ssl_configs is None:
                conn.open(address[0], address[1], 1000)
            else:
                conn.open_SSL(address[0], address[1], 1000, self._ssl_configs_ssl_configs)
            conn.close()
            return True
        except Exception as ex:
            logger.warning(
                'Connect {}:{} failed: {}'.format(address[0], address[1], ex)
            )
            return False
 
    def execute(self, stmt):
        """execute the given query
        Notice there are some limitations:
        1. The query should not be a plain space switch statement, e.g. "USE test_space",
        but queries like "use space xxx; match (v) return v" are accepted.
        2. If the query contains statements like "USE <space name>", the space will be set to the
        one in the pool config after the execution of the query.
        3. The query should not change the user password nor drop a user.
 
        :param stmt: the query string
        :return: ResultSet
        """
        return self.execute_parameterexecute_parameter(stmt, None)
 
    def execute_parameter(self, stmt, params):
        """execute statement
 
        :param stmt: the query string
        :param params: parameter map
        :return: ResultSet
        """
        session = self._get_idle_session_get_idle_session()
        if session is None:
            raise RuntimeError('Get session failed')
        self._add_session_to_active_add_session_to_active(session)
 
        try:
            resp = session.execute_parameter(stmt, params)
 
            # reset the space name to the pool config
            if resp.space_name() != self._space_name_space_name:
                self._set_space_to_default_set_space_to_default(session)
 
            # move the session back to the idle list
            self._return_session_return_session(session)
 
            return resp
        except Exception as e:
            logger.error('Execute failed: {}'.format(e))
            # remove the session from the pool if it is invalid
            self._active_sessions_active_sessions.remove(session)
            raise e
 
    def execute_json(self, stmt):
        """execute statement and return the result as a JSON string
            Date and Datetime will be returned in UTC
            JSON struct:
            {
                "results": [
                {
                    "columns": [],
                    "data": [
                    {
                        "row": [
                        "row-data"
                        ],
                        "meta": [
                        "metadata"
                        ]
                    }
                    ],
                    "latencyInUs": 0,
                    "spaceName": "",
                    "planDesc ": {
                    "planNodeDescs": [
                        {
                        "name": "",
                        "id": 0,
                        "outputVar": "",
                        "description": {
                            "key": ""
                        },
                        "profiles": [
                            {
                            "rows": 1,
                            "execDurationInUs": 0,
                            "totalDurationInUs": 0,
                            "otherStats": {}
                            }
                        ],
                        "branchInfo": {
                            "isDoBranch": false,
                            "conditionNodeId": -1
                        },
                        "dependencies": []
                        }
                    ],
                    "nodeIndexMap": {},
                    "format": "",
                    "optimize_time_in_us": 0
                    },
                    "comment ": ""
                }
                ],
                "errors": [
                {
                    "code": 0,
                    "message": ""
                }
                ]
            }
        :param stmt: the ngql
        :return: JSON string
        """
        return self.execute_json_with_parameterexecute_json_with_parameter(stmt, None)
 
    def execute_json_with_parameter(self, stmt, params):
        session = self._get_idle_session_get_idle_session()
        if session is None:
            raise RuntimeError('Get session failed')
        self._add_session_to_active_add_session_to_active(session)
 
        try:
            resp = session.execute_json_with_parameter(stmt, params)
 
            # reset the space name to the pool config
            json_obj = json.loads(resp)
            if json_obj["results"][0]["spaceName"] != self._space_name_space_name:
                self._set_space_to_default_set_space_to_default(session)
 
            # move the session back to the idle list
            self._return_session_return_session(session)
 
            return resp
        except Exception as e:
            logger.error('Execute failed: {}'.format(e))
            # remove the session from the pool if it is invalid
            self._active_sessions_active_sessions.remove(session)
            raise e
 
    def close(self):
        """log out all sessions and close all connections
 
        :return: void
        """
        with self._lock_lock:
            for session in self._idle_sessions_idle_sessions:
                session._sign_out()
                session._connection.close()
            for session in self._active_sessions_active_sessions:
                session._sign_out()
                session._connection.close()
            self._idle_sessions_idle_sessions.clear()
            self._close_close = True
 
    def get_ok_servers_num(self):
        """get the number of the ok servers
 
        :return: int
        """
        count = 0
        for addr in self._addresses_status_addresses_status.keys():
            if self._addresses_status_addresses_status[addr] == self.S_OKS_OK:
                count = count + 1
        return count
 
    def _get_services_status(self):
        msg_list = []
        for addr in self._addresses_status_addresses_status.keys():
            status = 'OK'
            if self._addresses_status_addresses_status[addr] != self.S_OKS_OK:
                status = 'BAD'
            msg_list.append('[services: {}, status: {}]'.format(addr, status))
        return ', '.join(msg_list)
 
    def update_servers_status(self):
        """update the servers' status"""
        for address in self._addresses_addresses:
            if self.pingping(address):
                self._addresses_status_addresses_status[address] = self.S_OKS_OK
            else:
                self._addresses_status_addresses_status[address] = self.S_BADS_BAD
 
    def ping_sessions(self):
        """ping all sessions in the pool"""
        with self._lock_lock:
            for session in self._idle_sessions_idle_sessions:
                session.execute(r'RETURN "SESSION PING"')
 
    def _get_idle_session(self):
        """get a valid session from the pool idle list.
 
        :return: Session
        """
        with self._lock_lock:
            if len(self._idle_sessions_idle_sessions) > 0:
                return self._idle_sessions_idle_sessions.pop(0)
            elif len(self._active_sessions_active_sessions) < self._configs_configs.max_size:
                return self._new_session_new_session()
            else:
                raise NoValidSessionException(
                    'The total number of sessions reaches the pool max size {}'.format(
                        self._configs_configs.max_size
                    )
                )
 
    def _new_session(self):
        """construct a new session with the username and password in the pool.
            also, the session is bound to the space specified in the configs.
 
        :return: Session
        """
        if self._ssl_configs_ssl_configs is not None:
            raise RuntimeError('SSL is not supported yet')
 
        self._pos_pos = (self._pos_pos + 1) % len(self._addresses_addresses)
        next_addr_index = self._pos_pos
 
        # try to connect with a valid service address, the worst case it to iterate all addresses
        retries = len(self._addresses_addresses)
 
        while retries > 0:
            addr = self._addresses_addresses[next_addr_index]
 
            # if the address is bad, skip it
            if self._addresses_status_addresses_status[addr] == self.S_BADS_BAD:
                logger.warning('The graph service {} is not available'.format(addr))
                retries = retries - 1
                next_addr_index = (next_addr_index + 1) % len(self._addresses_addresses)
                continue
 
            # connect to the valid service
            connection = Connection()
            try:
                connection.open(addr[0], addr[1], self._configs_configs.timeout)
                auth_result = connection.authenticate(self._username_username, self._password_password)
                session = Session(connection, auth_result, self, False)
 
                # switch to the space specified in the configs
                resp = session.execute('USE {}'.format(self._space_name_space_name))
                if not resp.is_succeeded():
                    raise RuntimeError(
                        'Failed to get session, cannot set the session space to {} error: {} {}'.format(
                            self._space_name_space_name, resp.error_code(), resp.error_msg()
                        )
                    )
                return session
            except AuthFailedException as e:
                # if auth failed because of credentials, close the pool
                if e.message.find("Invalid password") or e.message.find(
                    "User not exist"
                ):
                    logger.error(
                        'Authentication failed, because of bad credentials, close the pool {}'.format(
                            e
                        )
                    )
                    self.closeclose()
                raise e
            except Exception:
                raise
 
        raise RuntimeError(
            'Failed to get a valid session, no graph service is available'
        )
 
    def _return_session(self, session):
        """return the session to the pool idle list when query finished.
 
        :param session: the session to return
        :return: void
        """
        with self._lock_lock:
            self._active_sessions_active_sessions.remove(session)
            self._idle_sessions_idle_sessions.append(session)
            session.idle_time_start = time.time()
 
    def _add_session_to_idle(self, session):
        """add the session to the pool idle list
 
        :param session: the session to add
        :return: void
        """
        with self._lock_lock:
            self._idle_sessions_idle_sessions.append(session)
            session.idle_time_start = time.time()
 
    def _add_session_to_active(self, session):
        """add the session to the pool active list
 
        :param session: the session to add
        :return: void
        """
        with self._lock_lock:
            self._active_sessions_active_sessions.append(session)
            # reset the idle time start
            session.idle_time_start = 0
 
    def _set_space_to_default(self, session):
        """set the space to the default space in the pool
 
        :param session: the session to set
        :return: void
        """
        try:
            resp = session.execute('USE {}'.format(self._space_name_space_name))
            if not resp.is_succeeded():
                raise RuntimeError(
                    'Failed to set the session space to {}'.format(self._space_name_space_name)
                )
        except Exception:
            logger.warning(
                'Failed to set the session space to {}, the current session has been dropped'.format(
                    self._space_name_space_name
                )
            )
            session._connection.close()
            with self._lock_lock:
                self._active_sessions_active_sessions.remove(session)
 
    def _remove_idle_unusable_session(self):
        if self._configs_configs.idle_time == 0:
            return
        with self._lock_lock:
            total_sessions = len(self._idle_sessions_idle_sessions) + len(self._active_sessions_active_sessions)
            if total_sessions <= self._configs_configs.min_size:
                return
            for session in self._idle_sessions_idle_sessions:
                # calc session idle time
                idle_time = time.time() - session._idle_time_start
 
                # release idle session and remove from the pool
                if idle_time > self._configs_configs.idle_time:
                    conn = session._connection
                    session.release()
                    conn.close()
                    self._idle_sessions_idle_sessions.remove(session)
 
    def _period_detect(self):
        """periodically detect the services status and remove the sessions from the idle list if they expire"""
        if self._close_close or self._configs_configs.interval_check < 0:
            return
        self.update_servers_statusupdate_servers_status()
        self._remove_idle_unusable_session_remove_idle_unusable_session()
        timer = Timer(self._configs_configs.interval_check, self._period_detect_period_detect)
        timer.setDaemon(True)
        timer.start()
 
    def _check_configs(self):
        """validate the configs"""
        if self._configs_configs.min_size < 0:
            raise RuntimeError('The min_size must be greater than 0')
        if self._configs_configs.max_size < 0:
            raise RuntimeError('The max_size must be greater than 0')
        if self._configs_configs.min_size > self._configs_configs.max_size:
            raise RuntimeError(
                'The min_size must be less than or equal to the max_size'
            )
        if self._configs_configs.idle_time < 0:
            raise RuntimeError('The idle_time must be greater or equal to 0')
        if self._configs_configs.timeout < 0:
            raise RuntimeError('The timeout must be greater or equal to 0')
 
        if self._space_name_space_name == "":
            raise RuntimeError('The space_name must be set')
        if self._username_username == "":
            raise RuntimeError('The username must be set')
        if self._password_password == "":
            raise RuntimeError('The password must be set')
        if self._addresses_addresses is None or len(self._addresses_addresses) == 0:
            raise RuntimeError('The addresses must be set')