diff --git a/ziggy/tornado_utils.py b/ziggy/tornado_utils.py index 11afddc..55b60d0 100644 --- a/ziggy/tornado_utils.py +++ b/ziggy/tornado_utils.py @@ -19,6 +19,7 @@ import tornado.web import tornado.gen +import tornado.simple_httpclient import ziggy @@ -143,3 +144,29 @@ def run(self): finally: if self.ziggy_ctx: self.ziggy_ctx.stop() + +class AsyncHTTPClient(tornado.simple_httpclient.SimpleAsyncHTTPClient): + def __init__(self, *args, **kwargs): + self.ziggy_name = 'tornado.httpclient' + return super(AsyncHTTPClient, self).__init__(*args, **kwargs) + + def fetch(self, request, callback, **kwargs): + ctx = ziggy.Context(self.ziggy_name) + ctx.start() + if isinstance(request, basestring): + ctx.set('request.uri', request) + else: + ctx.set('request.uri', request.url) + ctx.set('request.size', len(request.body) if request.body else 0) + + ctx.stop() + + def wrap_callback(response): + ctx.start() + ctx.set('response.code', response.code) + ctx.set('response.size', len(response.body)) + ctx.done() + callback(response) + + return super(AsyncHTTPClient, self).fetch(request, wrap_callback, **kwargs) +