svn commit: r991986 - /incubator/thrift/trunk/lib/py/src/server/THttpServer.py

[dreiss]

Author: dreiss
Date: Thu Sep  2 15:36:03 2010
New Revision: 991986

URL: http://svn.apache.org/viewvc?rev=991986&view=rev
Log:
THRIFT-666. python: Allow THttpServer's responses to be overridden

Modified:
    incubator/thrift/trunk/lib/py/src/server/THttpServer.py

Modified: incubator/thrift/trunk/lib/py/src/server/THttpServer.py
URL: 
http://svn.apache.org/viewvc/incubator/thrift/trunk/lib/py/src/server/THttpServer.py?rev=991986&r1=991985&r2=991986&view=diff
==============================================================================
--- incubator/thrift/trunk/lib/py/src/server/THttpServer.py (original)
+++ incubator/thrift/trunk/lib/py/src/server/THttpServer.py Thu Sep  2 15:36:03 
2010
@@ -22,6 +22,19 @@ import BaseHTTPServer
 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 THttpServer(TServer.TServer):
     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)