Python THeaderTransport.THeaderTransport类代码示例

    def call_processor(self, input, client_type, protocol_type,
                       client_principal):
        try:
            # The input string has already had the header removed, but
            # the python processor will expect it to be there.  In
            # order to reconstitute the message with headers, we use
            # the THeaderProtocol object to write into a memory
            # buffer, then pass that buffer to the python processor.

            write_buf = TMemoryBuffer()
            trans = THeaderTransport(write_buf, client_types=[client_type])
            trans.set_protocol_id(protocol_type)
            trans.write(input)
            trans.flush()

            prot_buf = TMemoryBuffer(write_buf.getvalue())
            prot = THeaderProtocol(prot_buf)

            ctx = TCppConnectionContext(client_principal)

            self.processor.process(prot, prot, ctx)

            # And on the way out, we need to strip off the header,
            # because the C++ code will expect to add it.

            read_buf = TMemoryBuffer(prot_buf.getvalue())
            trans = THeaderTransport(read_buf, client_types=[client_type])
            trans.readFrame(0)

            return trans.cstringio_buf.read()
        except:
            # Don't let exceptions escape back into C++
            traceback.print_exc()

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()

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

    def done(prot_buf, client_type, callback):
        try:
            response = prot_buf.getvalue()

            if len(response) == 0:
                callback.call(response)
            else:
                # And on the way out, we need to strip off the header,
                # because the C++ code will expect to add it.

                read_buf = TMemoryBuffer(response)
                trans = THeaderTransport(read_buf, client_types=[client_type])
                trans.readFrame(len(response))
                callback.call(trans.cstringio_buf.read())
        except:
            traceback.print_exc()

    def setUp(self):
        self.trans = TTransport.TMemoryBuffer()
        self.trans = THeaderTransport(self.trans)
        self.prot = TBinaryProtocol.TBinaryProtocol(self.trans)

        self.x = Xtruct()
        self.x.string_thing = "Zero"
        self.x.byte_thing = 0

    def schedule_timeout(self, fname, seqid):
        timeout = self.timeouts[fname]
        if not timeout:
            return

        tmo = TMemoryBuffer()
        thp = THeaderTransport(tmo)
        oprot = THeaderProtocol(thp)
        exc = TApplicationException(
            TApplicationException.TIMEOUT, "Call to {} timed out".format(fname)
        )
        oprot.writeMessageBegin(fname, TMessageType.EXCEPTION, seqid)
        exc.write(oprot)
        oprot.writeMessageEnd()
        thp.flush()
        timeout_task = self.loop.create_task(
            self.message_received(tmo.getvalue(), delay=timeout),
        )
        self.update_pending_tasks(seqid, timeout_task)

    def testOversizeFrameRecv(self):
        """Test that an oversize frame on recv gets a TTransportException"""

        self.trans.set_max_frame_size(200)

        self.x.write(self.prot)
        self.trans.flush()

        inTrans = TTransport.TMemoryBuffer(self.trans.getTransport().getvalue())
        inTrans = THeaderTransport(inTrans)
        inProt = TBinaryProtocol.TBinaryProtocol(inTrans)

        inTrans.set_max_frame_size(2)

        try:
            self.x.read(inProt)
        except TTransport.TTransportException:
            return

        self.fail("Should have gotten TTransportException")

    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)
        self._socket.setTimeout(timeout * 1000)
        self._transport = THeaderTransport(self._socket)
        self._protocol = THeaderProtocol(self._transport)

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

class NetlinkManagerClient(NetlinkManagerService.Client):
    DEFAULT_PORT = 5912

    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)
        self._socket.setTimeout(timeout * 1000)
        self._transport = THeaderTransport(self._socket)
        self._protocol = THeaderProtocol(self._transport)

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

    def __enter__(self):
        return self

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

    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()
        PcapPushSubscriber.Client.__init__(self, self._protocol)

class PcapPushSubClient(PcapPushSubscriber.Client):
    DEFAULT_PORT = 5911

    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()
        PcapPushSubscriber.Client.__init__(self, self._protocol)

    def __enter__(self):
        return self

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

class TestClient(TestService.Client):
    DEFAULT_PORT = fboss.system_tests.test.constants.DEFAULT_PORT

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

        self._socket = TSocket(host, port)
        self._socket.setTimeout(timeout * 1000)
        self._transport = THeaderTransport(self._socket)
        self._protocol = THeaderProtocol(self._transport)

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

    def __enter__(self):
        return self

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

    def serialize_texception(cls, fname, seqid, exception):
        """This saves us a bit of processing time for timeout handling by
        reusing the Thrift structs involved in exception serialization.

        NOTE: this is not thread-safe nor is it meant to be.
        """
        # the serializer is a singleton
        if cls._exception_serializer is None:
            buffer = TWriteOnlyBuffer()
            transport = THeaderTransport(buffer)
            cls._exception_serializer = THeaderProtocol(transport)
        else:
            transport = cls._exception_serializer.trans
            buffer = transport.getTransport()
            buffer.reset()

        serializer = cls._exception_serializer
        serializer.writeMessageBegin(fname, TMessageType.EXCEPTION, seqid)
        exception.write(serializer)
        serializer.writeMessageEnd()
        serializer.trans.flush()
        return buffer.getvalue()

class TestEof(unittest.TestCase):

    def setUp(self):
        self.trans = TTransport.TMemoryBuffer()
        self.trans = THeaderTransport(self.trans)
        self.prot = TBinaryProtocol.TBinaryProtocol(self.trans)

        self.x = Xtruct()
        self.x.string_thing = "Zero"
        self.x.byte_thing = 0

    def testOversizeFrameRecv(self):
        """Test that an oversize frame on recv gets a TTransportException"""

        self.trans.set_max_frame_size(200)

        self.x.write(self.prot)
        self.trans.flush()

        inTrans = TTransport.TMemoryBuffer(self.trans.getTransport().getvalue())
        inTrans = THeaderTransport(inTrans)
        inProt = TBinaryProtocol.TBinaryProtocol(inTrans)

        inTrans.set_max_frame_size(2)

        try:
            self.x.read(inProt)
        except TTransport.TTransportException:
            return

        self.fail("Should have gotten TTransportException")

    def testOversizeFrameSend(self):
        """Test that an oversize frame on send gets a TTransportException"""
        self.trans.set_max_frame_size(2)

        self.x.write(self.prot)

        try:
            self.trans.flush()
        except TTransport.TTransportException:
            return

        self.fail("Should have gotten TTransportException")

    def call_processor(self, input, headers, client_type, protocol_type,
                       context_data, callback):
        try:
            # The input string has already had the header removed, but
            # the python processor will expect it to be there.  In
            # order to reconstitute the message with headers, we use
            # the THeaderProtocol object to write into a memory
            # buffer, then pass that buffer to the python processor.

            write_buf = TMemoryBuffer()
            trans = THeaderTransport(write_buf)
            trans._THeaderTransport__client_type = client_type
            trans._THeaderTransport__write_headers = headers
            trans.set_protocol_id(protocol_type)
            trans.write(input)
            trans.flush()

            prot_buf = TMemoryBuffer(write_buf.getvalue())
            prot = THeaderProtocol(prot_buf, client_types=[client_type])

            ctx = TCppConnectionContext(context_data)

            ret = self.processor.process(prot, prot, ctx)

            done_callback = partial(_ProcessorAdapter.done,
                                    prot_buf=prot_buf,
                                    client_type=client_type,
                                    callback=callback)
            # This future is created by and returned from the processor's
            # ThreadPoolExecutor, which keeps a reference to it. So it is
            # fine for this future to end its lifecycle here.
            if isinstance(ret, Future):
                ret.add_done_callback(lambda x, d=done_callback: d())
            else:
                done_callback()
        except:
            # Don't let exceptions escape back into C++
            traceback.print_exc()