THRIFT-666. python: Allow THttpServer's responses to be overridden
git-svn-id: https://svn.apache.org/repos/asf/incubator/thrift/trunk@991986 13f79535-47bb-0310-9956-ffa450edef68
diff --git a/lib/py/src/server/THttpServer.py b/lib/py/src/server/THttpServer.py
index 9bf80cb..3047d9c 100644
--- a/lib/py/src/server/THttpServer.py
+++ b/lib/py/src/server/THttpServer.py
@@ -22,6 +22,19 @@
from thrift.server import TServer
from thrift.transport import TTransport
+class ResponseException(Exception):
+ """Allows handlers to override the HTTP response
+
+ Normally, THttpServer always sends a 200 response. If a handler wants
+ to override this behavior (e.g., to simulate a misconfigured or
+ overloaded web server during testing), it can raise a ResponseException.
+ The function passed to the constructor will be called with the
+ RequestHandler as its only argument.
+ """
+ def __init__(self, handler):
+ self.handler = handler
+
+
class THttpServer(TServer.TServer):
"""A simple HTTP-based Thrift server
@@ -47,18 +60,21 @@
class RequestHander(BaseHTTPServer.BaseHTTPRequestHandler):
def do_POST(self):
# Don't care about the request path.
- self.send_response(200)
- self.send_header("content-type", "application/x-thrift")
- self.end_headers()
-
itrans = TTransport.TFileObjectTransport(self.rfile)
otrans = TTransport.TFileObjectTransport(self.wfile)
itrans = TTransport.TBufferedTransport(itrans, int(self.headers['Content-Length']))
- otrans = TTransport.TBufferedTransport(otrans)
+ otrans = TTransport.TMemoryBuffer()
iprot = thttpserver.inputProtocolFactory.getProtocol(itrans)
oprot = thttpserver.outputProtocolFactory.getProtocol(otrans)
- thttpserver.processor.process(iprot, oprot)
- otrans.flush()
+ try:
+ thttpserver.processor.process(iprot, oprot)
+ except ResponseException, exn:
+ exn.handler(self)
+ else:
+ self.send_response(200)
+ self.send_header("content-type", "application/x-thrift")
+ self.end_headers()
+ self.wfile.write(otrans.getvalue())
self.httpd = server_class(server_address, RequestHander)