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