rpc.py 5.0 KB

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