def __init__(self, host: str,
transport: THeaderTransport.THeaderTransport) -> None:
self.host = host # Just for accessibility
self._transport = transport
self._transport.add_transform(THeaderTransport.TRANSFORM.ZSTD)
OpenrCtrl.Client.__init__(
self, THeaderProtocol.THeaderProtocol(self._transport))
def setUp(self):
if options.http_path:
self.transport = THttpClient.THttpClient(
options.host, options.port, options.http_path)
else:
socket = TSocket.TSocket(options.host, options.port)
# frame or buffer depending upon args
if options.framed:
self.transport = TTransport.TFramedTransport(socket)
else:
self.transport = TTransport.TBufferedTransport(socket)
protocol = self.protocol_factory.getProtocol(self.transport)
if options.header:
protocol = THeaderProtocol.THeaderProtocol(socket)
self.transport = protocol.trans
protocol.trans.add_transform(protocol.trans.ZLIB_TRANSFORM)
self.transport.open()
self.client = ThriftTest.Client(protocol)
if options.multiple:
p = TMultiplexedProtocol.TMultiplexedProtocol(protocol,
"ThriftTest")
self.client = ThriftTest.Client(p)
p = TMultiplexedProtocol.TMultiplexedProtocol(protocol,
"SecondService")
self.client2 = SecondService.Client(p)
else:
self.client = ThriftTest.Client(protocol)
self.client2 = None
def test_throw_populates_headers(self):
handler = self.Handler({7: "hello"})
processor = TestService.Processor(handler)
server = TCppServerTestManager.make_server(processor)
with TCppServerTestManager(server) as server:
host, port = server.addr()
with TSocket.TSocket(host=host, port=port) as sock:
transport = THeaderTransport.THeaderTransport(sock)
protocol = THeaderProtocol.THeaderProtocol(transport)
client = TestService.Client(protocol, protocol)
try:
client.throwUserException()
self.fail("Expect to throw UserException2")
except UserException2:
pass
self.assertEquals(b"UserException2",
transport.get_headers()[b"uex"])
self.assertIn(b"Some message",
transport.get_headers()[b"uexw"])
try:
client.throwUncaughtException("a message!")
self.fail("Expect to throw TApplicationException")
except TApplicationException:
pass
self.assertEquals(b"TApplicationException",
transport.get_headers()[b"uex"])
self.assertIn(b"a message!", transport.get_headers()[b"uexw"])
def _perform_rpc(self, server, val):
host, port = server.addr()
with TSocket.TSocket(host=host, port=port) as sock:
transport = THeaderTransport.THeaderTransport(sock)
protocol = THeaderProtocol.THeaderProtocol(transport)
client = SlotsThrowingService.Client(protocol, protocol)
return client.getDataById(val)
def __init__(
self,
endpoint: config.EndpointConfiguration,
size: int = 10,
max_age: int = 120,
timeout: int = 1,
max_connection_attempts: Optional[int] = None,
max_retries: Optional[int] = None,
protocol_factory: TProtocolFactory = THeaderProtocol.THeaderProtocolFactory(),
):
if max_connection_attempts and max_retries:
raise Exception("do not mix max_retries and max_connection_attempts")
if max_retries:
warn_deprecated(
"ThriftConnectionPool's max_retries is now named max_connection_attempts"
)
max_connection_attempts = max_retries
elif not max_connection_attempts:
max_connection_attempts = 3
self.endpoint = endpoint
self.max_age = max_age
self.retry_policy = RetryPolicy.new(attempts=max_connection_attempts)
self.timeout = timeout
self.protocol_factory = protocol_factory
self.size = size
self.pool: ProtocolPool = queue.LifoQueue()
for _ in range(size):
self.pool.put(None)
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)
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
class HeaderAcceleratedCompactTest(AbstractTest):
protocol_factory = THeaderProtocol.THeaderProtocolFactory(
True, [
THeaderTransport.HEADERS_CLIENT_TYPE,
THeaderTransport.FRAMED_DEPRECATED,
THeaderTransport.UNFRAMED_DEPRECATED,
THeaderTransport.HTTP_CLIENT_TYPE
])
def _make_protocol(endpoint):
if endpoint.family == socket.AF_INET:
trans = TSocket.TSocket(*endpoint.address)
elif endpoint.family == socket.AF_UNIX:
trans = TSocket.TSocket(unix_socket=endpoint.address)
else:
raise Exception("unsupported endpoint family %r" % endpoint.family)
return THeaderProtocol.THeaderProtocol(trans)
class HeaderAcceleratedCompactTest(AbstractTest):
protocol_factory = THeaderProtocol.THeaderProtocolFactory(
True,
[CLIENT_TYPE.HEADER,
CLIENT_TYPE.FRAMED_DEPRECATED,
CLIENT_TYPE.UNFRAMED_DEPRECATED,
CLIENT_TYPE.HTTP_SERVER]
)
def _expiring_rpc(self, server, service, method, tm, *args, **kwargs):
host, port = server.addr()
with TSocket.TSocket(host=host, port=port) as sock:
sock.setTimeout(tm)
transport = THeaderTransport.THeaderTransport(sock)
protocol = THeaderProtocol.THeaderProtocol(transport)
client = service.Client(protocol, protocol)
return getattr(client, method)(*args, **kwargs)
def getHeaderClient(addr, sock_cls=TSocket.TSocket):
transport = sock_cls(addr[0], addr[1])
transport = THeaderTransport.THeaderTransport(transport)
transport.set_header("hello", "world")
protocol = THeaderProtocol.THeaderProtocol(transport)
client = SleepService.Client(protocol)
transport.open()
return client
class HeaderBase(AbstractTest):
protocol_factory = THeaderProtocol.THeaderProtocolFactory(
True,
[
CLIENT_TYPE.HEADER,
CLIENT_TYPE.FRAMED_DEPRECATED,
CLIENT_TYPE.UNFRAMED_DEPRECATED,
CLIENT_TYPE.HTTP_SERVER,
],
)
def testHeader(self):
transport = TSocket.TSocket(self.server_addr[0], self.server_addr[1])
transport = THeaderTransport.THeaderTransport(transport)
transport.set_header("hello", "world")
protocol = THeaderProtocol.THeaderProtocol(transport)
client = SleepService.Client(protocol)
transport.open()
self.assertEquals(client.space("hi"), "h i")
self.stopServer()
def __init__(self, endpoint, size=10, max_age=120, timeout=1, max_retries=3,
protocol_factory=THeaderProtocol.THeaderProtocolFactory()):
self.endpoint = endpoint
self.max_age = max_age
self.retry_policy = RetryPolicy.new(attempts=max_retries)
self.timeout = timeout
self.protocol_factory = protocol_factory
self.pool = queue.LifoQueue()
for _ in range(size):
self.pool.put(None)
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 _get_client_by_transport(self, config, transport, socket=None):
# Create the protocol and client
if config.compact:
protocol = TCompactProtocol.TCompactProtocol(transport)
# No explicit option about protocol is specified. Try to infer.
elif config.framed or config.unframed:
protocol = TBinaryProtocol.TBinaryProtocolAccelerated(transport)
elif socket is not None:
protocol = THeaderProtocol.THeaderProtocol(socket)
transport = protocol.trans
else:
raise ValueError("No protocol specified for HTTP Transport")
transport.open()
self._transport = transport
client = self.client_class(protocol)
return client
def get_thrift_client(self, use_ssl):
socket = (
TSSLSocket.TSSLSocket(
host=self.cli_opts.host,
port=self.cli_opts.openr_ctrl_port,
# verify server
cert_reqs=ssl.CERT_REQUIRED,
ca_certs=self.cli_opts.ca_file,
certfile=self.cli_opts.cert_file,
keyfile=self.cli_opts.key_file,
verify_name=self.cli_opts.acceptable_peer_name,
) if use_ssl else TSocket.TSocket(
host=self.cli_opts.host, port=self.cli_opts.openr_ctrl_port))
socket.setTimeout(self.cli_opts.timeout)
transport = THeaderTransport.THeaderTransport(socket)
protocol = THeaderProtocol.THeaderProtocol(transport)
transport.open()
return OpenrCtrl.Client(protocol)
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.
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 _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)
else:
if socket is None:
self._exit(error_message=('No valid protocol '
'specified for %s' % (type(self))),
status=os.EX_USAGE)
else:
protocol = THeaderProtocol.THeaderProtocol(socket)
transport = protocol.trans
transport.open()
self._transport = transport
client = self.service_class.Client(protocol)
return client
for _iteration in range(1, numberOfFlaps + 1):
# Ignore the SIGHUP signal to maintain session when the server port
# is turned down.
old_handler = signal.signal(signal.SIGHUP, signal.SIG_IGN)
self.log.debug("Flap iteration {}".format(_iteration))
self.check_output(command_for_interface_down)
time.sleep(sleepDuration)
self.check_output(command_for_interface_up)
# Change the SIGHUP settings back to original value.
signal.signal(signal.SIGHUP, old_handler)
if __name__ == '__main__':
handler = TestServer()
transport = TSocket.TServerSocket(DEFAULT_PORT)
tfactory = TTransport.TBufferedTransportFactory()
pfactory = THeaderProtocol.THeaderProtocolFactory()
server = TServer.TSimpleServer(handler, transport, tfactory, pfactory)
# You could do one of these for a multithreaded server
# server = TServer.TThreadedServer(handler, transport, tfactory, pfactory)
print('Starting the server...')
server.serve()
print('Done.')
type="int",
dest="port",
default=9090)
parser.add_option("--timeout",
action="store",
type="int",
dest="timeout",
default=60)
options, args = parser.parse_args()
event_handler = TestServerEventHandler()
if options.header:
pfactory = THeaderProtocol.THeaderProtocolFactory(
True, [
CLIENT_TYPE.HEADER, CLIENT_TYPE.FRAMED_DEPRECATED,
CLIENT_TYPE.UNFRAMED_DEPRECATED, CLIENT_TYPE.HTTP_SERVER
])
else:
pfactory = TBinaryProtocol.TBinaryProtocolFactory()
if options.context:
processor = ThriftTest.ContextProcessor(
TestContextHandler(options.port))
else:
processor = ThriftTest.Processor(TestHandler())
if options.multiple:
processor = TMultiplexedProcessor.TMultiplexedProcessor()
if options.context:
processor.registerProcessor(
dest="port",
default=9090)
parser.add_option("--timeout",
action="store",
type="int",
dest="timeout",
default=60)
options, args = parser.parse_args()
event_handler = TestServerEventHandler()
if options.header:
pfactory = THeaderProtocol.THeaderProtocolFactory(
True, [
THeaderTransport.HEADERS_CLIENT_TYPE,
THeaderTransport.FRAMED_DEPRECATED,
THeaderTransport.UNFRAMED_DEPRECATED,
THeaderTransport.HTTP_CLIENT_TYPE
])
else:
pfactory = TBinaryProtocol.TBinaryProtocolFactory()
if options.context:
processor = ThriftTest.ContextProcessor(
TestContextHandler(options.port))
else:
processor = ThriftTest.Processor(TestHandler())
if options.multiple:
processor = TMultiplexedProcessor.TMultiplexedProcessor()
if options.context:
if options.size is not None:
kwargs.setdefault("size", options.size)
if options.max_age is not None:
kwargs.setdefault("max_age", options.max_age.total_seconds())
if options.timeout is not None:
kwargs.setdefault("timeout", options.timeout.total_seconds())
if options.max_connection_attempts is not None:
kwargs.setdefault("max_connection_attempts",
options.max_connection_attempts)
if options.max_retries is not None:
raise Exception("max_retries was renamed to max_connection_attempts")
return ThriftConnectionPool(endpoint=options.endpoint, **kwargs)
_DEFAULT_PROTOCOL_FACTORY = THeaderProtocol.THeaderProtocolFactory(
default_protocol=THeaderProtocol.THeaderSubprotocolID.COMPACT, )
class ThriftConnectionPool:
"""A pool that maintains a queue of open Thrift connections.
:param endpoint: The remote address
of the Thrift service.
:param size: The maximum number of connections that can be open
before new attempts to open block.
:param max_age: The maximum number of seconds a connection should be
kept alive. Connections older than this will be reaped.
:param timeout: The maximum number of seconds a connection attempt or
RPC call can take before a TimeoutError is raised.
:param max_connection_attempts: The maximum number of times the pool will attempt
to open a connection.
def create_client(port):
socket = TSocket.TSocket("localhost", port)
protocol = THeaderProtocol.THeaderProtocol(socket)
protocol.trans.set_max_frame_size(MAX_BIG_FRAME_SIZE)
protocol.trans.open()
return ThriftTest.Client(protocol)
class HeaderBinaryToDefault(HeaderTest, unittest.TestCase):
serialize_factory = HeaderDefaultFactory(
THeaderProtocol.THeaderProtocol.T_BINARY_PROTOCOL)
deserialize_factory = THeaderProtocol.THeaderProtocolFactory()
class HeaderTest(AbstractTest):
protocol_factory = THeaderProtocol.THeaderProtocolFactory(
True, [
THeaderTransport.HEADERS_CLIENT_TYPE,
THeaderTransport.FRAMED_DEPRECATED,
THeaderTransport.UNFRAMED_DEPRECATED,
THeaderTransport.HTTP_CLIENT_TYPE
])
def testZlibCompression(self):
htrans = self.protocol.trans
if isinstance(htrans, THeaderTransport):
htrans.add_transform(THeaderTransport.ZLIB_TRANSFORM)
self.testStruct()
def testSnappyCompression(self):
htrans = self.protocol.trans
if isinstance(htrans, THeaderTransport):
htrans.add_transform(THeaderTransport.SNAPPY_TRANSFORM)
self.testStruct()
def testMultipleCompression(self):
htrans = self.protocol.trans
if isinstance(htrans, THeaderTransport):
htrans.add_transform(THeaderTransport.ZLIB_TRANSFORM)
htrans.add_transform(THeaderTransport.SNAPPY_TRANSFORM)
self.testStruct()
def testKeyValueHeader(self):
if self.server_header and self.server_type == 'TNonblockingServer':
# TNonblockingServer uses different protocol instances for input
# and output so persistent header won't work
return
htrans = self.protocol.trans
if isinstance(htrans, THeaderTransport):
# Try just persistent header
htrans.set_persistent_header(b"permanent", b"true")
self.client.testString('test')
headers = htrans.get_headers()
self.assertTrue(b"permanent" in headers)
self.assertEquals(headers[b"permanent"], b"true")
# Try with two transient headers
htrans.set_header(b"transient1", b"true")
htrans.set_header(b"transient2", b"true")
self.client.testString('test')
headers = htrans.get_headers()
self.assertTrue(b"permanent" in headers)
self.assertEquals(headers[b"permanent"], b"true")
self.assertTrue(b"transient1" in headers)
self.assertEquals(headers[b"transient1"], b"true")
self.assertTrue(b"transient2" in headers)
self.assertEquals(headers[b"transient2"], b"true")
# Add one, update one and delete one transient header
htrans.set_header(b"transient2", b"false")
htrans.set_header(b"transient3", b"true")
self.client.testString('test')
headers = htrans.get_headers()
self.assertTrue(b"permanent" in headers)
self.assertEquals(headers[b"permanent"], b"true")
self.assertTrue(b"transient1" not in headers)
self.assertTrue(b"transient2" in headers)
self.assertEquals(headers[b"transient2"], b"false")
self.assertTrue(b"transient3" in headers)
self.assertEquals(headers[b"transient3"], b"true")
class HeaderCompactToDefault(HeaderTest, unittest.TestCase):
serialize_factory = HeaderDefaultFactory(
THeaderProtocol.THeaderProtocol.T_COMPACT_PROTOCOL)
deserialize_factory = THeaderProtocol.THeaderProtocolFactory()
