Python TSocket.TSocket类代码示例

def connect_to_thrift(conf):
  """
  Connect to a thrift endpoint as determined by the 'conf' parameter.
  Note that this does *not* open the transport.

  Returns a tuple of (service, protocol, transport)
  """
  sock = TSocket(conf.host, conf.port)
  if conf.timeout_seconds:
    # Thrift trivia: You can do this after the fact with
    # _grab_transport_from_wrapper(self.wrapped.transport).setTimeout(seconds*1000)
    sock.setTimeout(conf.timeout_seconds*1000.0)
  if conf.use_sasl:
    def sasl_factory():
      saslc = sasl.Client()
      saslc.setAttr("host", conf.host)
      saslc.setAttr("service", conf.kerberos_principal)
      saslc.init()
      return saslc

    transport = TSaslClientTransport(sasl_factory, "GSSAPI", sock)
  else:
    transport = TBufferedTransport(sock)

  protocol = TBinaryProtocol(transport)
  service = conf.klass(protocol)
  return service, protocol, transport

def get_fast_transport(endpoint, timeout=5000):
    """
    采用cython实现的transport
    :param endpoint:
    :param timeout:
    :return:
    """
    global _fast_transport
    if not _fast_transport:
        if endpoint.find(":") != -1:
            hostport = endpoint.split(":")
            host = hostport[0]
            port = int(hostport[1])
            unix_socket = None
        else:
            host = None
            port = None
            unix_socket = endpoint


        socket = TSocket(host=host, port=port, unix_socket=unix_socket)
        socket.setTimeout(timeout)
        socket.open()

        _fast_transport = TCyFramedTransport(socket, maxIdleTime=1200) # 20分钟没有写数据，则重新打开transport

    return _fast_transport

    def __init__(self, host=None, port=10000, authMechanism=None, user=None, password=None, database=None, configuration=None, timeout=None):
        authMechanisms = set(['NOSASL', 'PLAIN', 'KERBEROS', 'LDAP'])
        if authMechanism not in authMechanisms:
            raise NotImplementedError('authMechanism is either not supported or not implemented')
        #Must set a password for thrift, even if it doesn't need one
        #Open issue with python-sasl
        if authMechanism == 'PLAIN' and (password is None or len(password) == 0):
            password = 'password'
        socket = TSocket(host, port)
        socket.setTimeout(timeout)
        if authMechanism == 'NOSASL':
            transport = TBufferedTransport(socket)
        else:
            sasl_mech = 'PLAIN'
            saslc = sasl.Client()
            saslc.setAttr("username", user)
            saslc.setAttr("password", password)
            if authMechanism == 'KERBEROS':
                krb_host,krb_service = self._get_krb_settings(host, configuration)
                sasl_mech = 'GSSAPI'
                saslc.setAttr("host", krb_host)
                saslc.setAttr("service", krb_service)

            saslc.init()
            transport = TSaslClientTransport(saslc, sasl_mech, socket)

        self.client = TCLIService.Client(TBinaryProtocol(transport))
        transport.open()
        res = self.client.OpenSession(TOpenSessionReq(username=user, password=password, configuration=configuration))
        self.session = res.sessionHandle
        if database is not None:
            with self.cursor() as cur:
                query = "USE {0}".format(database)
                cur.execute(query) 

class TLateInitSSLSocket(TTransportBase):
    def __init__(self, handle):
        self.handle = handle
        self.socket = None

    def getSocket(self):
        if self.socket is None:
            self.handle.do_handshake()
            self.socket = TSocket()
            self.socket.setHandle(self.handle)
        return self.socket

    def isOpen(self):
        return self.getSocket().isOpen()

    def setTimeout(self, ms):
        return self.getSocket().setTimeout(ms)

    def read(self, sz):
        return self.getSocket().read(sz)

    def write(self, sz):
        return self.getSocket().write(sz)

    def flush(self):
        return self.getSocket().flush()

    def close(self):
        if self.socket is None:
            self.handle.close()
        else:
            return self.getSocket().close()

class QsfpServiceClient(QsfpService.Client):
    DEFAULT_PORT = 5910
    DEFAULT_TIMEOUT = 10.0

    # we ignore the value of port
    def __init__(self, host, port=None, timeout=None):
        # In a box with all 32 QSFP ports populated, it takes about 7.5s right
        # now to read all 32 QSFP ports. So, put the defaut timeout to 10s.
        self.host = host

        timeout = timeout or self.DEFAULT_TIMEOUT
        self._socket = TSocket(host, self.DEFAULT_PORT)
        # TSocket.setTimeout() takes a value in milliseconds
        self._socket.setTimeout(timeout * 1000)
        self._transport = THeaderTransport(self._socket)
        self._protocol = THeaderProtocol(self._transport)

        self._transport.open()
        QsfpService.Client.__init__(self, self._protocol)

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        self._transport.close()

    def serveClient(self, socket, address):
        """Process input/output from a client for as long as possible"""
        client = TSocket()
        client.setHandle(socket)
        self.peerName = client.getPeerName()

        itrans = self.inputTransportFactory.getTransport(client)
        otrans = self.outputTransportFactory.getTransport(client)
        iprot = self.inputProtocolFactory.getProtocol(itrans)
        if isinstance(self.inputProtocolFactory, THeaderProtocolFactory):
            oprot = iprot
        else:
            oprot = self.outputProtocolFactory.getProtocol(otrans)

        try:
            while True:
                self.processor._handler.peerName = self.peerName
                self.processor.process(iprot, oprot)
        except TTransportException as tx:
            pass
        except Exception as x:
            self.logger.error('[%s]', x, extra={'clientip':self.peerName})

        itrans.close()
        otrans.close()

    def _refresh_thrift_client(self):
        """Refresh the Thrift socket, transport, and client."""
        socket = TSocket(self.host, self.port)
        if self.timeout is not None:
            socket.setTimeout(self.timeout)

        self.transport = self._transport_class(socket)
        protocol = self._protocol_class(self.transport)
        self.client = Client(protocol)

def connect(server='localhost', port=9090, timeout=None):
    socket = TSocket(server, int(port))
    if timeout is not None:
        socket.setTimeout(timeout)
    transport = TBufferedTransport(socket)
    transport.open()
    protocol = TBinaryProtocol.TBinaryProtocolAccelerated(transport)
    client = Hbase.Client(protocol)
    return client

class HS2TestSuite(ImpalaTestSuite):
  def setup(self):
    host, port = IMPALAD_HS2_HOST_PORT.split(":")
    self.socket = TSocket(host, port)
    self.transport = TBufferedTransport(self.socket)
    self.transport.open()
    self.protocol = TBinaryProtocol.TBinaryProtocol(self.transport)
    self.hs2_client = TCLIService.Client(self.protocol)

  def teardown(self):
    if self.socket:
      self.socket.close()

  @staticmethod
  def check_response(response,
                       expected_status_code = TCLIService.TStatusCode.SUCCESS_STATUS,
                       expected_error_prefix = None):
    assert response.status.statusCode == expected_status_code
    if expected_status_code != TCLIService.TStatusCode.SUCCESS_STATUS\
       and expected_error_prefix is not None:
      assert response.status.errorMessage.startswith(expected_error_prefix)

  def close(self, op_handle):
    close_op_req = TCLIService.TCloseOperationReq()
    close_op_req.operationHandle = op_handle
    close_op_resp = self.hs2_client.CloseOperation(close_op_req)
    assert close_op_resp.status.statusCode == TCLIService.TStatusCode.SUCCESS_STATUS

  def fetch(self, handle, orientation, size, expected_num_rows = None):
    """Fetches at most size number of rows from the query identified by the given
    operation handle. Uses the given fetch orientation. Asserts that the fetch returns
    a success status, and that the number of rows returned is equal to size, or
    equal to the given expected_num_rows (it one was given)."""
    fetch_results_req = TCLIService.TFetchResultsReq()
    fetch_results_req.operationHandle = handle
    fetch_results_req.orientation = orientation
    fetch_results_req.maxRows = size
    fetch_results_resp = self.hs2_client.FetchResults(fetch_results_req)
    HS2TestSuite.check_response(fetch_results_resp)
    num_rows = size
    if expected_num_rows is not None:
      num_rows = expected_num_rows
    assert len(fetch_results_resp.results.rows) == num_rows
    return fetch_results_resp

  def fetch_fail(self, handle, orientation, expected_error_prefix):
    """Attempts to fetch rows from the query identified by the given operation handle.
    Asserts that the fetch returns an error with an error message matching the given
    expected_error_prefix."""
    fetch_results_req = TCLIService.TFetchResultsReq()
    fetch_results_req.operationHandle = handle
    fetch_results_req.orientation = orientation
    fetch_results_req.maxRows = 100
    fetch_results_resp = self.hs2_client.FetchResults(fetch_results_req)
    HS2TestSuite.check_response(fetch_results_resp, TCLIService.TStatusCode.ERROR_STATUS,
                                expected_error_prefix)
    return fetch_results_resp

    def _refresh_thrift_client(self):
        """Refresh the Thrift socket, transport, and client."""
        socket = TSocket(self.host, self.port)
        if self.timeout is not None:
            socket.setTimeout(self.timeout)

        self.transport = self._transport_class(socket)
        protocol = TBinaryProtocol.TBinaryProtocolAccelerated(self.transport)
        self.client = Hbase.Client(protocol)

    def __init__(self, host='localhost', port=9090, unix_socket=None):
        """Initialize a TNoDelaySocket.

        Args:
            host: The host to connect to.
            port: The port to connect to.
            unix_socket: The filename of a unix socket to connect to. In this
                case, host and port will be ignored.
        """
        TSocket.__init__(self, host=host, port=port, unix_socket=unix_socket)

    def _create_conn(self):
        self._host_index += 1
        self._host_index %= len(self._host_list)
        host = self._host_list[self._host_index]
        parts = host.split(':')
        host = parts[0]
        port = int(parts[1])

        conn = TSocket(host, port)
        conn.setTimeout(self._time_out)
        conn.open()
        return conn

 def handle_stream(self, sock, address):
     tsock = TSocket()
     tsock.setHandle(sock)
     itrans = self.itrans.getTransport(tsock) 
     otrans = self.otrans.getTransport(tsock)
     iprot = self.iprot.getProtocol(itrans)
     oprot = self.oprot.getProtocol(otrans)
     try:
         while True:
             self.processor.process(iprot, oprot)
     except TTransportException, ex:
         pass

def construct_client(klass, host, port, service_name, timeout_seconds=45):
  """
  Constructs a thrift client, lazily.
  """
  sock = TSocket(host, port)
  if timeout_seconds:
    # Thrift trivia: You can do this after the fact with
    # self.wrapped.transport._TBufferedTransport__trans.setTimeout(seconds*1000)
    sock.setTimeout(timeout_seconds*1000.0)
  transport = TBufferedTransport(sock)
  protocol = TBinaryProtocol(transport)
  service = klass(protocol)
  return SuperClient(service, transport, timeout_seconds=timeout_seconds)

class FbossAgentClient(FbossCtrl.Client):
    DEFAULT_PORT = 5909

    def __init__(self, host, port=None, timeout=5.0):
        self.host = host
        if port is None:
            port = self.DEFAULT_PORT

        self._socket = TSocket(host, port)
        # TSocket.setTimeout() takes a value in milliseconds
        self._socket.setTimeout(timeout * 1000)
        self._transport = THeaderTransport(self._socket)
        self._protocol = THeaderProtocol(self._transport)

        self._transport.open()
        FbossCtrl.Client.__init__(self, self._protocol)

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        self._transport.close()

    #
    # The getPortStats() thrift API was unfortunately renamed to getPortInfo().
    # Here's a hacky workaround that tries to do the right thing regardless of
    # whether the switch we are talking to supports getPortStats() or
    # getPortInfo().
    #

    def getPortStats(self, *args, **kwargs):
        return self.getPortInfo(*args, **kwargs)

    def getAllPortStats(self, *args, **kwargs):
        return self.getAllPortInfo(*args, **kwargs)

    def getPortInfo(self, *args, **kwargs):
        try:
            return FbossCtrl.Client.getPortInfo(self, *args, **kwargs)
        except TApplicationException as ex:
            if 'Method name getPortInfo not found' in str(ex):
                return FbossCtrl.Client.getPortStats(self, *args, **kwargs)
            raise

    def getAllPortInfo(self, *args, **kwargs):
        try:
            return FbossCtrl.Client.getAllPortInfo(self, *args, **kwargs)
        except TApplicationException as ex:
            if 'Method name getAllPortInfo not found' in str(ex):
                return FbossCtrl.Client.getAllPortStats(self, *args, **kwargs)
            raise