1 files changed, 23 insertions, 15 deletions

diff --git a/api/utils/thread.py b/api/utils/thread.py
index 2106548f5..5f4ec7e94 100644
--- a/api/utils/thread.py
+++ b/api/utils/thread.py
@@ -1,37 +1,45 @@
 import threading
+import os
 import logging
 
 from oslo_serialization import jsonutils
 
-from api.database.v1.handlers import TasksHandler
 from yardstick.common import constants as consts
 
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.DEBUG)
+LOG = logging.getLogger(__name__)
+LOG.setLevel(logging.DEBUG)
 
 
 class TaskThread(threading.Thread):
 
-    def __init__(self, target, args):
+    def __init__(self, target, args, handler):
         super(TaskThread, self).__init__(target=target, args=args)
         self.target = target
         self.args = args
+        self.handler = handler
 
     def run(self):
-        task_handler = TasksHandler()
-        data = {'task_id': self.args.task_id, 'status': consts.TASK_NOT_DONE}
-        task_handler.insert(data)
+        if self.handler.__class__.__name__.lower().startswith('v2'):
+            self.handler.update_attr(self.args.task_id, {'status': consts.TASK_NOT_DONE})
+        else:
+            update_data = {'task_id': self.args.task_id, 'status': consts.TASK_NOT_DONE}
+            self.handler.insert(update_data)
 
-        logger.info('Starting run task')
+        LOG.info('Starting run task')
         try:
             data = self.target(self.args)
         except Exception as e:
-            logger.exception('Task Failed')
+            LOG.exception('Task Failed')
             update_data = {'status': consts.TASK_FAILED, 'error': str(e)}
-            task_handler.update_attr(self.args.task_id, update_data)
+            self.handler.update_attr(self.args.task_id, update_data)
         else:
-            logger.info('Task Finished')
-            logger.debug('Result: %s', data)
-
-            data['result'] = jsonutils.dumps(data.get('result', {}))
-            task_handler.update_attr(self.args.task_id, data)
+            LOG.info('Task Finished')
+            LOG.debug('Result: %s', data)
+
+            if self.handler.__class__.__name__.lower().startswith('v2'):
+                new_data = {'status': consts.TASK_DONE, 'result': jsonutils.dumps(data['result'])}
+                self.handler.update_attr(self.args.task_id, new_data)
+                os.remove(self.args.inputfile[0])
+            else:
+                data['result'] = jsonutils.dumps(data.get('result', {}))
+                self.handler.update_attr(self.args.task_id, data)