| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113 |
- # -*- 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()
|