Python TTransport.TMemoryBuffer类代码示例

    def test_with_headers(self):
        client_memory_trans = TMemoryBuffer()
        client_prot = THeaderProtocol(client_memory_trans)
        client_header_trans = client_prot.trans
        client_header_trans.set_header("Trace", "1234")
        client_header_trans.set_header("Parent", "2345")
        client_header_trans.set_header("Span", "3456")
        client_header_trans.set_header("Sampled", "1")
        client_header_trans.set_header("Flags", "1")
        client = TestService.Client(client_prot)
        try:
            client.example_simple()
        except TTransportException:
            pass  # we don't have a test response for the client
        self.itrans._readBuffer = StringIO(client_memory_trans.getvalue())

        self.processor.process(self.iprot, self.oprot, self.server_context)
        self.assertEqual(self.observer.on_server_span_created.call_count, 1)

        context, server_span = self.observer.on_server_span_created.call_args[0]
        self.assertEqual(server_span.trace_id, 1234)
        self.assertEqual(server_span.parent_id, 2345)
        self.assertEqual(server_span.id, 3456)
        self.assertTrue(server_span.sampled)
        self.assertEqual(server_span.flags, 1)

        self.assertEqual(self.server_observer.on_start.call_count, 1)
        self.assertEqual(self.server_observer.on_finish.call_count, 1)
        self.assertEqual(self.server_observer.on_finish.call_args[0], (None,))

 def inner_processor(message_buffer):
     in_transport = TMemoryBuffer(message_buffer.getvalue())
     out_transport = TMemoryBuffer()
     in_prot = proto_factory.getProtocol(in_transport)
     out_prot = proto_factory.getProtocol(out_transport)
     method = processor.process(in_prot, out_prot)
     return (method, out_transport.getvalue())

  def __auth_headers(self, headers, body, support_account_key):
    auth_headers = dict()
    if self.credential and self.credential.type and self.credential.secretKeyId:
      if self.credential.type in SIGNATURE_SUPPORT:
        auth_headers[HOST] = self.host
        # timestamp
        auth_headers[TIMESTAMP] = str(int(time.time() + self.__clock_offset))
        auth_headers[MI_DATE] = formatdate(usegmt=True)
        # content md5
        auth_headers[CONTENT_MD5] = hashlib.md5(body).hexdigest()

        headers_to_sign = defaultdict(lambda :[])
        for k, v in headers.iteritems():
          headers_to_sign[str(k).lower()].append(v)

        for k, v in auth_headers.iteritems():
          headers_to_sign[str(k).lower()].append(v)

        signature = base64.b64encode(self.sign(self.__form_sign_content("POST", self.uri,
          headers_to_sign))).strip()
        auth_string = "Galaxy-V2 %s:%s" % (self.credential.secretKeyId, signature)

        auth_headers[AUTHORIZATION] = auth_string
      else:
        auth_header = HttpAuthorizationHeader()
        auth_header.secretKeyId = self.credential.secretKeyId
        auth_header.userType = self.credential.type
        auth_header.secretKey = self.credential.secretKey
        auth_header.supportAccountKey = support_account_key
        mb = TMemoryBuffer()
        protocol = TJSONProtocol(mb)
        auth_header.write(protocol)
        auth_headers[AUTHORIZATION] = str(mb.getvalue())
    return auth_headers

    def test_with_headers(self):
        client_memory_trans = TMemoryBuffer()
        client_prot = THeaderProtocol(client_memory_trans)
        client_header_trans = client_prot.trans
        client_header_trans.set_header("Trace", "1234")
        client_header_trans.set_header("Parent", "2345")
        client_header_trans.set_header("Span", "3456")
        client = BaseplateService.Client(client_prot)
        try:
            client.is_healthy()
        except:
            pass  # we don't have a test response for the client
        self.itrans._readBuffer = StringIO(client_memory_trans.getvalue())

        self.processor.process(self.iprot, self.oprot, self.server_context)

        self.assertEqual(self.observer.on_root_span_created.call_count, 1)

        context, root_span = self.observer.on_root_span_created.call_args[0]
        self.assertEqual(root_span.trace_id, "1234")
        self.assertEqual(root_span.parent_id, "2345")
        self.assertEqual(root_span.id, "3456")

        self.assertTrue(self.root_observer.on_start.called)
        self.assertTrue(self.root_observer.on_stop.called)

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

    def handle_request(self, event):

        # t = time.time()
        # 0.2ms
        # 1. 将zeromq的消息转换成为 thrift的 protocols
        trans_input = TMemoryBuffer(event.msg)
        trans_output = TMemoryBuffer()

        proto_input = self.proto_factory_input.getProtocol(trans_input)
        proto_output = self.proto_factory_output.getProtocol(trans_output)

        # 2. 交给processor来处理
        try:
            self.processor.process(proto_input, proto_output)
            # 3. 将thirft的结果转换成为 zeromq 格式的数据
            msg = trans_output.getvalue()
            # print "Return Msg: ", msg, event.id
            if self.profile:
                event.id.extend(["", "%.4f" % time.time()])
                self.events.emit(msg, event.id)
            else:
                self.events.emit(msg, event.id)
        except Exception as e:
            # 如何出现了异常该如何处理呢
            # 程序不能挂
            logging.exception("Exception: %s", e)

    def encode_and_decode(self, obj):
        trans = TMemoryBuffer()
        if self.PROTO == 0:
            proto = TBinaryProtocol.TBinaryProtocol(trans)
        else:
            proto = TCompactProtocol.TCompactProtocol(trans)

        obj.write(proto)

        obj_new = obj.__class__()
        trans = TMemoryBuffer(trans.getvalue())
        proto = proto.__class__(trans)

        obj_new.read(proto)

def http_handler(request):
    if request.method != 'POST':
        return HttpResponseNotAllowed(['GET', 'PUT'])

    server_info.client_ip = request.META['REMOTE_ADDR']
    if server_info.client_ip[0:7] == '::ffff:':
        server_info.client_ip = server_info.client_ip[7:]
    server_info.client_port = None

    itrans = TMemoryBuffer(request.body)
    otrans = TMemoryBuffer()
    iproto = TBinaryProtocol(itrans)
    oproto = TBinaryProtocol(otrans)
    thrift_handler(iproto, oproto)
    return HttpResponse(otrans.getvalue(), content_type="application/x-thrift")

    def message_received(self, frame):
        tmi = TMemoryBuffer(frame)
        tmo = TMemoryBuffer()

        iprot = THeaderProtocol(tmi)
        oprot = THeaderProtocol(tmo)

        try:
            yield from self.processor.process(iprot, oprot, self.server_context)
            msg = tmo.getvalue()
            if len(msg) > 0:
                self.transport.write(msg)
        except Exception:
            logging.exception("Exception while processing request")
            self.transport.close()

    def test_expected_exception_not_passed_to_server_span_finish(self):
        client_memory_trans = TMemoryBuffer()
        client_prot = THeaderProtocol(client_memory_trans)
        client = TestService.Client(client_prot)
        try:
            client.example_throws(crash=False)
        except TTransportException:
            pass  # we don't have a test response for the client
        self.itrans._readBuffer = StringIO(client_memory_trans.getvalue())

        self.processor.process(self.iprot, self.oprot, self.server_context)

        self.assertEqual(self.server_observer.on_start.call_count, 1)
        self.assertEqual(self.server_observer.on_finish.call_count, 1)
        self.assertEqual(self.server_observer.on_finish.call_args[0], (None,))

def write_communication_to_buffer(comm):
    '''
    Serialize communication to buffer (binary string) and return
    buffer.

    Args:
        comm (Communication): communication to serialize

    Returns:
        Communication: Communication read from buffer
    '''
    transport = TMemoryBuffer()
    protocol = factory.createProtocol(transport)
    comm.write(protocol)
    return transport.getvalue()

    def call(self, procedure, args):
        if isinstance(procedure, str):
            try:
                procedure = processor._procedures[procedure]
            except KeyError:
                raise TApplicationException(TApplicationException.UNKNOWN_METHOD,
                    "Unknown method '{0}'".format(name))

        otrans = TMemoryBuffer()
        oproto = TBinaryProtocol(otrans)

#        perf.begin('send')
        oproto.writeMessageBegin(procedure.name, TMessageType.CALL, self._processor.seqid)
        self._processor.seqid = self._processor.seqid + 1
        self._processor.send_struct(Namespace(args), procedure.parameters_struct, oproto)
        oproto.writeMessageEnd()
#        perf.end('send')

        self._http.request('POST', self._url, otrans.getvalue(), {})
        resp = self._http.getresponse()
        data = resp.read()
        iproto = TBinaryProtocol(TMemoryBuffer(data))

#        perf.begin('wait')
        (fname, mtype, rseqid) = iproto.readMessageBegin()
#        perf.end('wait')

#        perf.begin('recv')
        if mtype == TMessageType.EXCEPTION:
            x = TApplicationException()
            x.read(iproto)
            iproto.readMessageEnd()
            x.args = (x.message,)
            raise x
        result = self._processor.recv_struct(procedure.results_struct, iproto)
        iproto.readMessageEnd()
#        perf.end('recv')
#        perf.end('call')

        if result.result is not None:
            return result.result

        for field in procedure.results_struct.fields:
            if getattr(result, field.name) is not None:
                raise getattr(result, field.name)

        return None

    def round_robin(self, compress=None):
        original = b'A' * MAX_FRAME_SIZE
        mb = TMemoryBuffer()
        trans = THeaderTransport(mb, client_type=CLIENT_TYPE.HEADER)
        trans.set_max_frame_size(MAX_FRAME_SIZE + MIN_HEADER_SIZE)
        if compress:
            trans.add_transform(compress)
        trans.write(original)
        trans.flush()
        frame = mb.getvalue()
        # Cleanup the memory buffer
        mb.close()
        del mb

        if compress is None:
            # Partial Decode the frame and see if its correct size wise
            sz = struct.unpack('!I', frame[:4])[0]
            self.assertEqual(sz, BIG_FRAME_MAGIC)
            sz = struct.unpack('!Q', frame[4:12])[0]
            self.assertEqual(len(frame), sz + 12)

        # Read it back
        mb = TMemoryBuffer(frame)
        trans = THeaderTransport(mb, client_type=CLIENT_TYPE.HEADER)
        trans.set_max_frame_size(len(frame))
        trans.readFrame(0)
        result = trans.read(MAX_FRAME_SIZE)
        mb.close()
        del mb
        self.assertEqual(result, original, 'round-robin different from original')

  def DeserializeThriftCall(self, buf):
    """Deserialize a stream and context to a MethodReturnMessage.

    Args:
      buf - The buffer.
      ctx - The context from serialization.

    Returns:
      A MethodCallMessage.
    """

    thrift_buffer = TMemoryBuffer()
    thrift_buffer._buffer = buf
    protocol = self._protocol_factory.getProtocol(thrift_buffer)

    (fn_name, msg_type, seq_id) = protocol.readMessageBegin()
    if msg_type == TMessageType.EXCEPTION:
      x = TApplicationException()
      x.read(protocol)
      protocol.readMessageEnd()
      return MethodReturnMessage(error=x)

    result_cls = self._FindClass('%s_result' % fn_name)
    if result_cls:
      result = result_cls()
      result.read(protocol)
    else:
      result = None
    protocol.readMessageEnd()

    if not result:
      return MethodReturnMessage()
    if getattr(result, 'success', None) is not None:
      return MethodReturnMessage(return_value=result.success)

    result_spec = getattr(result_cls, 'thrift_spec', None)
    if result_spec:
      exceptions = result_spec[1:]
      for e in exceptions:
        attr_val = getattr(result, e[2], None)
        if attr_val is not None:
          return MethodReturnMessage(error=attr_val)

    return MethodReturnMessage(TApplicationException(
      TApplicationException.MISSING_RESULT, "%s failed: unknown result" % fn_name))

    def decode_helper(self, obj, split=1.0):
        trans = TMemoryBuffer()
        if self.PROTO == 0:
            proto = TBinaryProtocol.TBinaryProtocol(trans)
        else:
            proto = TCompactProtocol.TCompactProtocol(trans)

        obj.write(proto)
        index = int(split * len(trans.getvalue()))
        trans = ReadOnlyBufferWithRefill(index, trans.getvalue())
        obj_new = obj.__class__()
        fastproto.decode(obj_new, trans, [obj.__class__, obj.thrift_spec,
            obj.isUnion()], utf8strings=0, protoid=self.PROTO)
        self.assertEqual(obj, obj_new)
        # Verify the entire buffer is read
        self.assertEqual(len(trans._readBuffer.read()), 0)
        if split != 1.0:
            self.assertEqual(1, trans.refill_called)