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
from typing import List, Optional
import time
 
from nebula3.common.ttypes import ErrorCode
from nebula3.Exception import (
    AuthFailedException,
    NoValidSessionException,
    InValidHostname,
)
 
from nebula3.gclient.net.Session import Session
from nebula3.gclient.net.Connection import Connection
from nebula3.gclient.net.base import BaseExecutor
from nebula3.logger import logger
from nebula3.Config import SessionPoolConfig, SSL_config
 
 
class SessionPool(BaseExecutor, 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: List[Session] = list()
        # sessions that are currently available
        self._idle_sessions: List[Session] = 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: Optional[SessionPoolConfig] = None,
        ssl_configs: Optional[SSL_config] = None,
    ):
        """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.
        """
        if configs is not None:
            assert isinstance(
                configs, SessionPoolConfig
            ), "wrong type of SessionPoolConfig, try this: `from nebula3.Config import SessionPoolConfig`"
            self._configs_configs = configs
        else:
            self._configs_configs = SessionPoolConfig()
        self._ssl_configs_ssl_configs = ssl_configs
        # 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.")
 
        # 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,
                    self._configs_configs.use_http2,
                    self._configs_configs.http_headers,
                )
            else:
                conn.open_SSL(
                    address[0],
                    address[1],
                    1000,
                    self._ssl_configs_ssl_configs,
                    self._configs_configs.use_http2,
                    self._configs_configs.http_headers,
                )
            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 super().execute(stmt)
 
    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)
 
            # Check for session validity based on error code
            if resp.error_code() in [
                ErrorCode.E_SESSION_INVALID,
                ErrorCode.E_SESSION_TIMEOUT,
            ]:
                self._active_sessions.remove(session)
                session = self._get_idle_session_get_idle_session()
                if session is None:
                    logger.warning(
                        "Session invalid or timeout, removed from the pool, but failed to get a new session."
                    )
                    return resp
                logger.warning("Session invalid or timeout, session has been recycled")
                self._add_session_to_idle_add_session_to_idle(session)
 
            else:
                # 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.remove(session)
            raise e
 
    def execute_json(self, stmt):
        """execute statement and return the result as a JSON bytes
            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 bytes
        """
        return super().execute_json(stmt)
 
    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)
            json_obj = json.loads(resp)
            # Check for session validity based on error code
            if json_obj.get("errors", [{}])[0].get("code") in [
                ErrorCode.E_SESSION_INVALID,
                ErrorCode.E_SESSION_TIMEOUT,
            ]:
                self._active_sessions.remove(session)
                session = self._get_idle_session_get_idle_session()
                if session is None:
                    logger.warning(
                        "Session invalid or timeout, removed from the pool, but failed to get a new session."
                    )
                    return resp
                self._add_session_to_idle_add_session_to_idle(session)
                logger.warning("Session invalid or timeout, session has been recycled")
 
            else:
                # reset the space name to the pool config
                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.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:
                session._sign_out()
                session._connection.close()
            for session in self._active_sessions:
                session._sign_out()
                session._connection.close()
            self._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:
                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) > 0:
                return self._idle_sessions.pop(0)
            elif len(self._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
        """
        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:
                if self._ssl_configs_ssl_configs is None:
                    connection.open(
                        addr[0],
                        addr[1],
                        self._configs_configs.timeout,
                        self._configs_configs.use_http2,
                        self._configs_configs.http_headers,
                    )
                else:
                    connection.open_SSL(
                        addr[0],
                        addr[1],
                        self._configs_configs.timeout,
                        self._ssl_configs_ssl_configs,
                        self._configs_configs.use_http2,
                        self._configs_configs.http_headers,
                    )
                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
                try:
                    resp = session.execute("USE {}".format(self._space_name_space_name))
                except Exception:
                    session.release()
                    connection.close()
                    raise RuntimeError(
                        "Failed to get session, execute `use {}` failed.".format(
                            self._space_name_space_name
                        )
                    )
                if not resp.is_succeeded():
                    session.release()
                    connection.close()
                    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()
                else:
                    connection.close()
                raise e
            except Exception:
                connection.close()
                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.remove(session)
            self._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.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.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.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) + len(self._active_sessions)
            if total_sessions <= self._configs_configs.min_size:
                return
            for session in self._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.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")