| # |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| # |
| |
| from thrift.protocol.TBinaryProtocol import TBinaryProtocolAccelerated |
| from thrift.protocol.TCompactProtocol import TCompactProtocolAccelerated |
| from thrift.protocol.TProtocol import TProtocolBase, TProtocolException, TProtocolFactory |
| from thrift.Thrift import TApplicationException, TMessageType |
| from thrift.transport.THeaderTransport import THeaderTransport, THeaderSubprotocolID, THeaderClientType |
| |
| |
| PROTOCOLS_BY_ID = { |
| THeaderSubprotocolID.BINARY: TBinaryProtocolAccelerated, |
| THeaderSubprotocolID.COMPACT: TCompactProtocolAccelerated, |
| } |
| |
| |
| class THeaderProtocol(TProtocolBase): |
| """A framed protocol with headers and payload transforms. |
| |
| THeaderProtocol frames other Thrift protocols and adds support for optional |
| out-of-band headers. The currently supported subprotocols are |
| TBinaryProtocol and TCompactProtocol. When used as a client, the |
| subprotocol to frame can be chosen with the `default_protocol` parameter to |
| the constructor. |
| |
| It's also possible to apply transforms to the encoded message payload. The |
| only transform currently supported is to gzip. |
| |
| When used in a server, THeaderProtocol can accept messages from |
| non-THeaderProtocol clients if allowed (see `allowed_client_types`). This |
| includes framed and unframed transports and both TBinaryProtocol and |
| TCompactProtocol. The server will respond in the appropriate dialect for |
| the connected client. HTTP clients are not currently supported. |
| |
| THeaderProtocol does not currently support THTTPServer, TNonblockingServer, |
| or TProcessPoolServer. |
| |
| See doc/specs/HeaderFormat.md for details of the wire format. |
| |
| """ |
| |
| def __init__(self, transport, allowed_client_types, default_protocol=THeaderSubprotocolID.BINARY): |
| # much of the actual work for THeaderProtocol happens down in |
| # THeaderTransport since we need to do low-level shenanigans to detect |
| # if the client is sending us headers or one of the headerless formats |
| # we support. this wraps the real transport with the one that does all |
| # the magic. |
| if not isinstance(transport, THeaderTransport): |
| transport = THeaderTransport(transport, allowed_client_types, default_protocol) |
| super(THeaderProtocol, self).__init__(transport) |
| self._set_protocol() |
| |
| def get_headers(self): |
| return self.trans.get_headers() |
| |
| def set_header(self, key, value): |
| self.trans.set_header(key, value) |
| |
| def clear_headers(self): |
| self.trans.clear_headers() |
| |
| def add_transform(self, transform_id): |
| self.trans.add_transform(transform_id) |
| |
| def writeMessageBegin(self, name, ttype, seqid): |
| self.trans.sequence_id = seqid |
| return self._protocol.writeMessageBegin(name, ttype, seqid) |
| |
| def writeMessageEnd(self): |
| return self._protocol.writeMessageEnd() |
| |
| def writeStructBegin(self, name): |
| return self._protocol.writeStructBegin(name) |
| |
| def writeStructEnd(self): |
| return self._protocol.writeStructEnd() |
| |
| def writeFieldBegin(self, name, ttype, fid): |
| return self._protocol.writeFieldBegin(name, ttype, fid) |
| |
| def writeFieldEnd(self): |
| return self._protocol.writeFieldEnd() |
| |
| def writeFieldStop(self): |
| return self._protocol.writeFieldStop() |
| |
| def writeMapBegin(self, ktype, vtype, size): |
| return self._protocol.writeMapBegin(ktype, vtype, size) |
| |
| def writeMapEnd(self): |
| return self._protocol.writeMapEnd() |
| |
| def writeListBegin(self, etype, size): |
| return self._protocol.writeListBegin(etype, size) |
| |
| def writeListEnd(self): |
| return self._protocol.writeListEnd() |
| |
| def writeSetBegin(self, etype, size): |
| return self._protocol.writeSetBegin(etype, size) |
| |
| def writeSetEnd(self): |
| return self._protocol.writeSetEnd() |
| |
| def writeBool(self, bool_val): |
| return self._protocol.writeBool(bool_val) |
| |
| def writeByte(self, byte): |
| return self._protocol.writeByte(byte) |
| |
| def writeI16(self, i16): |
| return self._protocol.writeI16(i16) |
| |
| def writeI32(self, i32): |
| return self._protocol.writeI32(i32) |
| |
| def writeI64(self, i64): |
| return self._protocol.writeI64(i64) |
| |
| def writeDouble(self, dub): |
| return self._protocol.writeDouble(dub) |
| |
| def writeBinary(self, str_val): |
| return self._protocol.writeBinary(str_val) |
| |
| def _set_protocol(self): |
| try: |
| protocol_cls = PROTOCOLS_BY_ID[self.trans.protocol_id] |
| except KeyError: |
| raise TApplicationException( |
| TProtocolException.INVALID_PROTOCOL, |
| "Unknown protocol requested.", |
| ) |
| |
| self._protocol = protocol_cls(self.trans) |
| self._fast_encode = self._protocol._fast_encode |
| self._fast_decode = self._protocol._fast_decode |
| |
| def readMessageBegin(self): |
| try: |
| self.trans.readFrame(0) |
| self._set_protocol() |
| except TApplicationException as exc: |
| self._protocol.writeMessageBegin(b"", TMessageType.EXCEPTION, 0) |
| exc.write(self._protocol) |
| self._protocol.writeMessageEnd() |
| self.trans.flush() |
| |
| return self._protocol.readMessageBegin() |
| |
| def readMessageEnd(self): |
| return self._protocol.readMessageEnd() |
| |
| def readStructBegin(self): |
| return self._protocol.readStructBegin() |
| |
| def readStructEnd(self): |
| return self._protocol.readStructEnd() |
| |
| def readFieldBegin(self): |
| return self._protocol.readFieldBegin() |
| |
| def readFieldEnd(self): |
| return self._protocol.readFieldEnd() |
| |
| def readMapBegin(self): |
| return self._protocol.readMapBegin() |
| |
| def readMapEnd(self): |
| return self._protocol.readMapEnd() |
| |
| def readListBegin(self): |
| return self._protocol.readListBegin() |
| |
| def readListEnd(self): |
| return self._protocol.readListEnd() |
| |
| def readSetBegin(self): |
| return self._protocol.readSetBegin() |
| |
| def readSetEnd(self): |
| return self._protocol.readSetEnd() |
| |
| def readBool(self): |
| return self._protocol.readBool() |
| |
| def readByte(self): |
| return self._protocol.readByte() |
| |
| def readI16(self): |
| return self._protocol.readI16() |
| |
| def readI32(self): |
| return self._protocol.readI32() |
| |
| def readI64(self): |
| return self._protocol.readI64() |
| |
| def readDouble(self): |
| return self._protocol.readDouble() |
| |
| def readBinary(self): |
| return self._protocol.readBinary() |
| |
| |
| class THeaderProtocolFactory(TProtocolFactory): |
| def __init__( |
| self, |
| allowed_client_types=(THeaderClientType.HEADERS,), |
| default_protocol=THeaderSubprotocolID.BINARY, |
| ): |
| self.allowed_client_types = allowed_client_types |
| self.default_protocol = default_protocol |
| |
| def getProtocol(self, trans): |
| return THeaderProtocol(trans, self.allowed_client_types, self.default_protocol) |