def __init__(self,
trans,
strictRead=False,
client_types=None,
client_type=None):
"""Create a THeaderProtocol instance
@param transport(TTransport) The underlying transport.
@param strictRead(bool) Turn on strictRead if using TBinaryProtocol
@param client_types([CLIENT_TYPE.HEADER, ...])
List of client types to support. Defaults to
CLIENT_TYPE.HEADER only.
"""
if isinstance(trans, THeaderTransport):
trans._THeaderTransport__supported_client_types = set(
client_types or (CLIENT_TYPE.HEADER, ))
if client_type is not None:
trans._THeaderTransport__client_type = client_type
htrans = trans
else:
htrans = THeaderTransport(trans, client_types, client_type)
TProtocolBase.__init__(self, htrans)
self.strictRead = strictRead
self.reset_protocol()
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 it is 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()
def _get_client_by_transport(self, options, transport, socket=None):
# Create the protocol and client
if options.json:
protocol = TJSONProtocol.TJSONProtocol(transport)
elif options.compact:
protocol = TCompactProtocol.TCompactProtocol(transport)
# No explicit option about protocol is specified. Try to infer.
elif options.framed or options.unframed:
protocol = TBinaryProtocol.TBinaryProtocolAccelerated(transport)
elif socket is not None:
# If json, compact, framed, and unframed are not specified,
# THeaderProtocol is the default. Create a protocol using either
# fuzzy or non-fuzzy transport depending on if options.fuzz is set.
if options.fuzz is not None:
transport = TFuzzyHeaderTransport(
socket, fuzz_fields=options.fuzz, verbose=True)
else:
transport = THeaderTransport(socket)
if options.headers is not None:
try:
parsed_headers = eval(options.headers)
except Exception:
self._exit(
error_message='Request headers (--headers) argument'
' failed eval')
if not isinstance(parsed_headers, dict):
self._exit(
error_message='Request headers (--headers) argument'
' must evaluate to a dict')
for header_name, header_value in parsed_headers.items():
transport.set_header(header_name, header_value)
protocol = THeaderProtocol.THeaderProtocol(transport)
else:
self._exit(error_message=('No valid protocol '
'specified for %s' % (type(self))),
status=os.EX_USAGE)
transport.open()
self._transport = transport
client = self.service_class.Client(protocol)
return client
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.set_max_frame_size(MAX_BIG_FRAME_SIZE)
trans.readFrame(len(response))
callback.call(trans.cstringio_buf.read())
except: # noqa
# Don't let exceptions escape back into C++
traceback.print_exc()
def call_processor(self, input, headers, client_type, protocol_type,
context_data, callback):
try:
# TCppServer threads are not created by Python so they are
# missing settrace() hooks. We need to manually set the
# hook here for things to work (e.g. coverage and pdb).
if sys.gettrace() is None and threading._trace_hook is not None:
sys.settrace(threading._trace_hook)
# 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.
should_sample = self._shouldSample()
timestamps = CallTimestamps()
if self.observer and should_sample:
timestamps.setProcessBeginNow()
write_buf = TMemoryBuffer()
trans = THeaderTransport(write_buf)
trans.set_max_frame_size(MAX_BIG_FRAME_SIZE)
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])
prot.trans.set_max_frame_size(MAX_BIG_FRAME_SIZE)
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)
if self.observer:
if should_sample:
timestamps.setProcessEndNow()
# This only bumps counters if `processBegin != 0` and
# `processEnd != 0` and these will only be non-zero if
# we are sampling this request.
self.observer.callCompleted(timestamps)
# 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: # noqa
# Don't let exceptions escape back into C++
traceback.print_exc()