1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113
|
# -*- coding: utf-8 -*-
from __future__ import absolute_import
import contextlib
import warnings
from thriftpy.protocol import TBinaryProtocolFactory
from thriftpy.server import TThreadedServer
from thriftpy.thrift import TProcessor, TClient
from thriftpy.transport import (
TBufferedTransportFactory,
TServerSocket,
TSSLServerSocket,
TSocket,
TSSLSocket,
)
def make_client(service, host="localhost", port=9090, unix_socket=None,
proto_factory=TBinaryProtocolFactory(),
trans_factory=TBufferedTransportFactory(),
timeout=None,
cafile=None, ssl_context=None, certfile=None, keyfile=None):
if unix_socket:
socket = TSocket(unix_socket=unix_socket)
if certfile:
warnings.warn("SSL only works with host:port, not unix_socket.")
elif host and port:
if cafile or ssl_context:
socket = TSSLSocket(host, port, socket_timeout=timeout,
cafile=cafile,
certfile=certfile, keyfile=keyfile,
ssl_context=ssl_context)
else:
socket = TSocket(host, port, socket_timeout=timeout)
else:
raise ValueError("Either host/port or unix_socket must be provided.")
transport = trans_factory.get_transport(socket)
protocol = proto_factory.get_protocol(transport)
transport.open()
return TClient(service, protocol)
def make_server(service, handler,
host="localhost", port=9090, unix_socket=None,
proto_factory=TBinaryProtocolFactory(),
trans_factory=TBufferedTransportFactory(),
client_timeout=3000, certfile=None):
processor = TProcessor(service, handler)
if unix_socket:
server_socket = TServerSocket(unix_socket=unix_socket)
if certfile:
warnings.warn("SSL only works with host:port, not unix_socket.")
elif host and port:
if certfile:
server_socket = TSSLServerSocket(
host=host, port=port, client_timeout=client_timeout,
certfile=certfile)
else:
server_socket = TServerSocket(
host=host, port=port, client_timeout=client_timeout)
else:
raise ValueError("Either host/port or unix_socket must be provided.")
server = TThreadedServer(processor, server_socket,
iprot_factory=proto_factory,
itrans_factory=trans_factory)
return server
@contextlib.contextmanager
def client_context(service, host="localhost", port=9090, unix_socket=None,
proto_factory=TBinaryProtocolFactory(),
trans_factory=TBufferedTransportFactory(),
timeout=None, socket_timeout=3000, connect_timeout=3000,
cafile=None, ssl_context=None, certfile=None, keyfile=None):
if timeout:
warnings.warn("`timeout` deprecated, use `socket_timeout` and "
"`connect_timeout` instead.")
socket_timeout = connect_timeout = timeout
if unix_socket:
socket = TSocket(unix_socket=unix_socket,
connect_timeout=connect_timeout,
socket_timeout=socket_timeout)
if certfile:
warnings.warn("SSL only works with host:port, not unix_socket.")
elif host and port:
if cafile or ssl_context:
socket = TSSLSocket(host, port,
connect_timeout=connect_timeout,
socket_timeout=socket_timeout,
cafile=cafile,
certfile=certfile, keyfile=keyfile,
ssl_context=ssl_context)
else:
socket = TSocket(host, port,
connect_timeout=connect_timeout,
socket_timeout=socket_timeout)
else:
raise ValueError("Either host/port or unix_socket must be provided.")
try:
transport = trans_factory.get_transport(socket)
protocol = proto_factory.get_protocol(transport)
transport.open()
yield TClient(service, protocol)
finally:
transport.close()
|