Connection.py

# --coding:utf-8--
#
# Copyright (c) 2021 vesoft inc. All rights reserved.
#
# This source code is licensed under Apache 2.0 License.
 
 
import time
 
from nebula3.fbthrift.transport import (
    TSocket,
    TSSLSocket,
    TTransport,
    THeaderTransport,
)
from nebula3.fbthrift.transport.TTransport import TTransportException
from nebula3.fbthrift.protocol import THeaderProtocol
 
from nebula3.common.ttypes import ErrorCode
from nebula3.graph import GraphService
from nebula3.graph.ttypes import VerifyClientVersionReq
from nebula3.logger import logger
 
from nebula3.Exception import (
    AuthFailedException,
    IOErrorException,
    ClientServerIncompatibleException,
)
 
from nebula3.gclient.net.AuthResult import AuthResult
 
 
class Connection(object):
    is_used = False
 
    def __init__(self):
        self._connection_connection = None
        self.start_use_timestart_use_time = time.time()
        self._ip_ip = None
        self._port_port = None
        self._timeout_timeout = 0
        self._ssl_conf_ssl_conf = None
 
    def open(self, ip, port, timeout):
        """open the connection
 
        :param ip: the server ip
        :param port: the server port
        :param timeout: the timeout for connect and execute
        :return: void
        """
        self.open_SSLopen_SSL(ip, port, timeout, None)
 
    def open_SSL(self, ip, port, timeout, ssl_config=None):
        """open the SSL connection
 
        :param ip: the server ip
        :param port: the server port
        :param timeout: the timeout for connect and execute
        :ssl_config: configs for SSL
        :return: void
        """
        self._ip_ip = ip
        self._port_port = port
        self._timeout_timeout = timeout
        self._ssl_conf_ssl_conf = ssl_config
        try:
            if ssl_config is not None:
                s = TSSLSocket.TSSLSocket(
                    self._ip_ip,
                    self._port_port,
                    ssl_config.unix_socket,
                    ssl_config.ssl_version,
                    ssl_config.cert_reqs,
                    ssl_config.ca_certs,
                    ssl_config.verify_name,
                    ssl_config.keyfile,
                    ssl_config.certfile,
                    ssl_config.allow_weak_ssl_versions,
                )
            else:
                s = TSocket.TSocket(self._ip_ip, self._port_port)
            if timeout > 0:
                s.setTimeout(timeout)
 
            buffered_transport = TTransport.TBufferedTransport(s)
            header_transport = THeaderTransport.THeaderTransport(buffered_transport)
            protocol = THeaderProtocol.THeaderProtocol(header_transport)
            header_transport.open()
 
            self._connection_connection = GraphService.Client(protocol)
            resp = self._connection_connection.verifyClientVersion(VerifyClientVersionReq())
            if resp.error_code != ErrorCode.SUCCEEDED:
                self._connection_connection._iprot.trans.close()
                raise ClientServerIncompatibleException(resp.error_msg)
        except Exception:
            raise
 
    def _reopen(self):
        """reopen the connection
 
        :return:
        """
        self.closeclose()
        if self._ssl_conf_ssl_conf is not None:
            self.open_SSLopen_SSL(self._ip_ip, self._port_port, self._timeout_timeout, self._ssl_conf_ssl_conf)
        else:
            self.openopen(self._ip_ip, self._port_port, self._timeout_timeout)
 
    def authenticate(self, user_name, password):
        """authenticate to graphd
 
        :param user_name: the user name
        :param password: the password
        :return:
        """
        try:
            resp = self._connection_connection.authenticate(user_name, password)
            if resp.error_code != ErrorCode.SUCCEEDED:
                raise AuthFailedException(resp.error_msg)
            return AuthResult(
                resp.session_id, resp.time_zone_offset_seconds, resp.time_zone_name
            )
        except TTransportException as te:
            if te.message.find("timed out"):
                self._reopen_reopen()
            if te.type == TTransportException.END_OF_FILE:
                self.closeclose()
            raise IOErrorException(IOErrorException.E_CONNECT_BROKEN, te.message)
 
    def execute(self, session_id, stmt):
        """execute interface with session_id and ngql
 
        :param session_id: the session id get from result of authenticate interface
        :param stmt: the ngql
        :return: ExecutionResponse
        """
        return self.execute_parameterexecute_parameter(session_id, stmt, None)
 
    def execute_parameter(self, session_id, stmt, params):
        """execute interface with session_id and ngql
        :param session_id: the session id get from result of authenticate interface
        :param stmt: the ngql
        :param params: parameter map
        :return: ExecutionResponse
        """
        try:
            resp = self._connection_connection.executeWithParameter(session_id, stmt, params)
            return resp
        except Exception as te:
            if isinstance(te, TTransportException):
                if te.message.find("timed out") > 0:
                    self._reopen_reopen()
                    raise IOErrorException(IOErrorException.E_TIMEOUT, te.message)
                elif te.type == TTransportException.END_OF_FILE:
                    raise IOErrorException(
                        IOErrorException.E_CONNECT_BROKEN, te.message
                    )
                elif te.type == TTransportException.NOT_OPEN:
                    raise IOErrorException(IOErrorException.E_NOT_OPEN, te.message)
                else:
                    raise IOErrorException(IOErrorException.E_UNKNOWN, te.message)
            raise
 
    def execute_json(self, session_id, stmt):
        """execute_json interface with session_id and ngql
        :param session_id: the session id get from result of authenticate interface
        :param stmt: the ngql
        :return: string json representing the execution result
        """
        return self.execute_json_with_parameterexecute_json_with_parameter(session_id, stmt, None)
 
    def execute_json_with_parameter(self, session_id, stmt, params):
        """execute_json interface with session_id and ngql with parameter
        :param session_id: the session id get from result of authenticate interface
        :param stmt: the ngql
        :param params: parameter map
        :return: string json representing the execution result
        """
        try:
            resp = self._connection_connection.executeJsonWithParameter(session_id, stmt, params)
            return resp
        except Exception as te:
            if isinstance(te, TTransportException):
                if te.message.find("timed out") > 0:
                    self._reopen_reopen()
                    raise IOErrorException(IOErrorException.E_TIMEOUT, te.message)
                elif te.type == TTransportException.END_OF_FILE:
                    raise IOErrorException(
                        IOErrorException.E_CONNECT_BROKEN, te.message
                    )
                elif te.type == TTransportException.NOT_OPEN:
                    raise IOErrorException(IOErrorException.E_NOT_OPEN, te.message)
                else:
                    raise IOErrorException(IOErrorException.E_UNKNOWN, te.message)
            raise
 
    def signout(self, session_id):
        """tells the graphd can release the session info
 
        :param session_id:the session id
        :return: void
        """
        try:
            self._connection_connection.signout(session_id)
        except TTransportException as te:
            if te.type == TTransportException.END_OF_FILE:
                self.closeclose()
 
    def close(self):
        """close the connection
 
        :return: void
        """
        try:
            self._connection_connection._iprot.trans.close()
        except Exception as e:
            logger.error(
                'Close connection to {}:{} failed:{}'.format(self._ip_ip, self._port_port, e)
            )
 
    def ping(self):
        """check the connection if ok
        :return: True or False
        """
        try:
            resp = self._connection_connection.execute(0, 'YIELD 1;')
            return True
        except Exception:
            return False
 
    def reset(self):
        """reset the idletime
 
        :return: void
        """
        self.start_use_timestart_use_time = time.time()
 
    def idle_time(self):
        """get idletime of connection
 
        :return: idletime
        """
        if self.is_usedis_used:
            return 0
        return (time.time() - self.start_use_timestart_use_time) * 1000
 
    def get_address(self):
        """get the address of the connected service
 
        :return: (ip, port)
        """
        return (self._ip_ip, self._port_port)