Fix how async method are created.

  * Use Tornado coroutines instead of the callback
    approach.

Change-Id: I23545b621185e8ad2a525a615cd9aab0fbd4f8a0

5 files changed, 77 insertions, 90 deletions

diff --git a/app/handlers/base.py b/app/handlers/base.py
index f96f3fd..a0d0f64 100644
--- a/app/handlers/base.py
+++ b/app/handlers/base.py
@@ -19,10 +19,12 @@ except ImportError:
     import json
 
 import bson
-import functools
 import httplib
 import tornado
+import tornado.escape
+import tornado.gen
 import tornado.web
+import types
 
 import handlers.common as hcommon
 import handlers.response as hresponse
@@ -44,6 +46,9 @@ STATUS_MESSAGES = {
 }
 
 
+# pylint: disable=unused-argument
+# pylint: disable=too-many-public-methods
+# pylint: disable=no-self-use
 class BaseHandler(tornado.web.RequestHandler):
     """The base handler."""
 
@@ -63,6 +68,7 @@ class BaseHandler(tornado.web.RequestHandler):
 
     @property
     def content_type(self):
+        """The accepted content-type header."""
         return hcommon.ACCEPTED_CONTENT_TYPE
 
     @property
@@ -99,6 +105,10 @@ class BaseHandler(tornado.web.RequestHandler):
 
     @staticmethod
     def _token_validation_func():
+        """The function that should be used to validate the token.
+
+        :return A function.
+        """
         return hcommon.valid_token_general
 
     def _get_status_message(self, status_code):
@@ -116,33 +126,36 @@ class BaseHandler(tornado.web.RequestHandler):
                 status_code, "Unknown status code returned")
         return message
 
-    def _create_valid_response(self, response):
-        """Create a valid JSON response based on its type.
-
-        :param response: The response we have from a query to the database.
-        :type HandlerResponse
-        """
+    def write(self, future):
+        """Write the response back to the requestor."""
         status_code = 200
         headers = {}
         result = {}
 
-        if isinstance(response, hresponse.HandlerResponse):
-            status_code = response.status_code
-            reason = response.reason or self._get_status_message(status_code)
-            headers = response.headers
-            result = json.dumps(
-                response.to_dict(),
-                default=bson.json_util.default,
-                ensure_ascii=False,
-                separators=(",", ":")
-            )
+        if isinstance(future, hresponse.HandlerResponse):
+            status_code = future.status_code
+            reason = future.reason or self._get_status_message(status_code)
+            headers = future.headers
+            to_dump = future.to_dict()
+        elif isinstance(future, types.DictionaryType):
+            status_code = future.get("code", 200)
+            reason = future.get(
+                "reason", self._get_status_message(status_code))
+            to_dump = future
         else:
             status_code = 506
             reason = self._get_status_message(status_code)
-            result = dict(code=status_code, reason=reason)
+            to_dump = dict(code=status_code, reason=reason)
+
+        result = json.dumps(
+            to_dump,
+            default=bson.json_util.default,
+            ensure_ascii=False,
+            separators=(",", ":")
+        )
 
         self.set_status(status_code=status_code, reason=reason)
-        self.write(result)
+        self._write_buffer.append(tornado.escape.utf8(result))
         self.set_header("Content-Type", hcommon.DEFAULT_RESPONSE_TYPE)
 
         if headers:
@@ -151,6 +164,17 @@ class BaseHandler(tornado.web.RequestHandler):
 
         self.finish()
 
+    def write_error(self, status_code, **kwargs):
+        if kwargs.get("message", None):
+            status_message = kwargs["message"]
+        else:
+            status_message = self._get_status_message(status_code)
+
+        if status_message:
+            self.write(dict(code=status_code, reason=status_message))
+        else:
+            super(BaseHandler, self).write_error(status_code, kwargs)
+
     def _has_valid_content_type(self):
         """Check if the request content type is the one expected.
 
@@ -173,15 +197,10 @@ class BaseHandler(tornado.web.RequestHandler):
 
         return valid_content
 
-    @tornado.web.asynchronous
+    @tornado.gen.coroutine
     def put(self, *args, **kwargs):
-        self.executor.submit(
-            functools.partial(self.execute_put, *args, **kwargs)
-        ).add_done_callback(
-            lambda future: tornado.ioloop.IOLoop.instance().add_callback(
-                functools.partial(self._create_valid_response, future.result())
-            )
-        )
+        future = yield self.executor.submit(self.execute_put, *args, **kwargs)
+        self.write(future)
 
     def execute_put(self, *args, **kwargs):
         """Execute the PUT pre-operations."""
@@ -208,15 +227,10 @@ class BaseHandler(tornado.web.RequestHandler):
         """
         return hresponse.HandlerResponse(501)
 
-    @tornado.web.asynchronous
+    @tornado.gen.coroutine
     def post(self, *args, **kwargs):
-        self.executor.submit(
-            functools.partial(self.execute_post, *args, **kwargs)
-        ).add_done_callback(
-            lambda future: tornado.ioloop.IOLoop.instance().add_callback(
-                functools.partial(self._create_valid_response, future.result())
-            )
-        )
+        future = yield self.executor.submit(self.execute_post, *args, **kwargs)
+        self.write(future)
 
     def execute_post(self, *args, **kwargs):
         """Execute the POST pre-operations.
@@ -302,16 +316,11 @@ class BaseHandler(tornado.web.RequestHandler):
         """
         return hresponse.HandlerResponse(501)
 
-    @tornado.web.asynchronous
+    @tornado.gen.coroutine
     def delete(self, *args, **kwargs):
-        self.executor.submit(
-            functools.partial(self.execute_delete, *args, **kwargs)
-        ).add_done_callback(
-            lambda future:
-            tornado.ioloop.IOLoop.instance().add_callback(
-                functools.partial(self._create_valid_response, future.result())
-            )
-        )
+        future = yield self.executor.submit(
+            self.execute_delete, *args, **kwargs)
+        self.write(future)
 
     def execute_delete(self, *args, **kwargs):
         """Perform DELETE pre-operations.
@@ -347,16 +356,10 @@ class BaseHandler(tornado.web.RequestHandler):
         """
         return hresponse.HandlerResponse(501)
 
-    @tornado.web.asynchronous
+    @tornado.gen.coroutine
     def get(self, *args, **kwargs):
-        self.executor.submit(
-            functools.partial(self.execute_get, *args, **kwargs)
-        ).add_done_callback(
-            lambda future:
-            tornado.ioloop.IOLoop.instance().add_callback(
-                functools.partial(self._create_valid_response, future.result())
-            )
-        )
+        future = yield self.executor.submit(self.execute_get, *args, **kwargs)
+        self.write(future)
 
     def execute_get(self, *args, **kwargs):
         """This is the actual GET operation.
@@ -482,19 +485,6 @@ class BaseHandler(tornado.web.RequestHandler):
 
         return (spec, sort, fields, skip, limit, unique)
 
-    def write_error(self, status_code, **kwargs):
-        if kwargs.get("message", None):
-            status_message = kwargs["message"]
-        else:
-            status_message = self._get_status_message(status_code)
-
-        if status_message:
-            self.set_status(status_code, status_message)
-            self.write(dict(code=status_code, message=status_message))
-            self.finish()
-        else:
-            super(BaseHandler, self).write_error(status_code, kwargs)
-
     # TODO: cache the validated token.
     def validate_req_token(self, method):
         """Validate the request token.
@@ -504,7 +494,7 @@ class BaseHandler(tornado.web.RequestHandler):
         """
         valid_token = False
 
-        req_token = self.get_request_token()
+        req_token = self.request.headers.get(hcommon.API_TOKEN_HEADER, None)
         remote_ip = self.request.remote_ip
         master_key = self.settings.get(hcommon.MASTER_KEY, None)
 
@@ -521,14 +511,17 @@ class BaseHandler(tornado.web.RequestHandler):
 
         return valid_token
 
-    def get_request_token(self):
-        """Retrieve the Authorization token of this request.
-
-        :return The authorization token as string.
-        """
-        return self.request.headers.get(hcommon.API_TOKEN_HEADER, None)
-
     def _token_validation(self, req_token, method, remote_ip, master_key):
+        """Perform the real token validation.
+
+        :param req_token: The token as taken from the request.
+        :type req_token: string
+        :param method: The HTTP verb to validate.
+        :type method: string
+        :param remote_ip: The IP address originating the request.
+        :param master_key: The default master key.
+        :return True or False.
+        """
         valid_token = False
         token_obj = self._find_token(req_token, self.db)
 
diff --git a/app/handlers/bisect.py b/app/handlers/bisect.py
index 2a242f2..23c64a0 100644
--- a/app/handlers/bisect.py
+++ b/app/handlers/bisect.py
@@ -14,8 +14,8 @@
 """The request handler for bisect URLs."""
 
 import bson
-import functools
 import tornado
+import tornado.gen
 import tornado.web
 
 import handlers.base as hbase
@@ -32,16 +32,10 @@ class BisectHandler(hbase.BaseHandler):
     def __init__(self, application, request, **kwargs):
         super(BisectHandler, self).__init__(application, request, **kwargs)
 
-    @tornado.web.asynchronous
+    @tornado.gen.coroutine
     def get(self, *args, **kwargs):
-        self.executor.submit(
-            functools.partial(self.execute_get, *args, **kwargs)
-        ).add_done_callback(
-            lambda future:
-            tornado.ioloop.IOLoop.instance().add_callback(
-                functools.partial(self._create_valid_response, future.result())
-            )
-        )
+        future = yield self.executor.submit(self.execute_get, *args, **kwargs)
+        self.write(future)
 
     @property
     def collection(self):
diff --git a/app/handlers/boot.py b/app/handlers/boot.py
index 9f314aa..a7cbe4c 100644
--- a/app/handlers/boot.py
+++ b/app/handlers/boot.py
@@ -46,7 +46,7 @@ class BootHandler(hbase.BaseHandler):
         return hcommon.valid_token_bh
 
     def _post(self, *args, **kwargs):
-        req_token = self.get_request_token()
+        req_token = self.request.headers.get(hcommon.API_TOKEN_HEADER, None)
         lab_name = kwargs["json_obj"].get(models.LAB_NAME_KEY, None)
 
         if self._is_valid_token(req_token, lab_name):
@@ -177,7 +177,7 @@ class BootHandler(hbase.BaseHandler):
         :return True or False.
         """
         valid_token = True
-        req_token = self.get_request_token()
+        req_token = self.request.headers.get(hcommon.API_TOKEN_HEADER, None)
         token = self._find_token(req_token, self.db)
 
         if token:
diff --git a/app/handlers/count.py b/app/handlers/count.py
index 3092ed7..b24fd40 100644
--- a/app/handlers/count.py
+++ b/app/handlers/count.py
@@ -15,7 +15,7 @@
 
 """Handle the /count URLs used to count objects in the database."""
 
-import tornado.web
+import tornado.gen
 
 import handlers.base as hbase
 import handlers.common as hcommon
@@ -66,12 +66,12 @@ class CountHandler(hbase.BaseHandler):
 
         return response
 
-    @tornado.web.asynchronous
+    @tornado.gen.coroutine
     def post(self, *args, **kwargs):
         """Not implemented."""
         self.write_error(status_code=501)
 
-    @tornado.web.asynchronous
+    @tornado.gen.coroutine
     def delete(self, *args, **kwargs):
         """Not implemented."""
         self.write_error(status_code=501)
diff --git a/app/handlers/tests/test_job_handler.py b/app/handlers/tests/test_job_handler.py
index 4e2350d..8058308 100644
--- a/app/handlers/tests/test_job_handler.py
+++ b/app/handlers/tests/test_job_handler.py
@@ -285,7 +285,7 @@ class TestJobHandler(
         self.assertEqual(
             response.headers["Content-Type"], DEFAULT_CONTENT_TYPE)
 
-    @mock.patch('handlers.job.JobHandler._get_one')
+    @mock.patch('handlers.base.BaseHandler._get_one')
     def test_get_wrong_handler_response(self, mock_get_one):
         mock_get_one.return_value = ""