taskflow/taskflow/storage.py

# -*- coding: utf-8 -*-

# vim: tabstop=4 shiftwidth=4 softtabstop=4

#    Copyright (C) 2013 Yahoo! Inc. All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

from taskflow import exceptions
from taskflow.openstack.common import uuidutils
from taskflow.persistence import flowdetail
from taskflow.persistence import logbook
from taskflow.persistence import taskdetail
from taskflow import states


def temporary_flow_detail():
    """Creates flow detail class for temporary usage

    Creates in-memory logbook and flow detail in it. Should
    be useful for tests and other use cases where persistence
    is not needed
    """
    lb = logbook.LogBook('tmp', backend='memory')
    fd = flowdetail.FlowDetail(
        name='tmp', uuid=uuidutils.generate_uuid(),
        backend='memory')
    lb.add(fd)
    lb.save()
    fd.save()
    return fd


STATES_WITH_RESULTS = (states.SUCCESS, states.REVERTING, states.FAILURE)


class Storage(object):
    """Interface between engines and logbook

    This class provides simple interface to save task details and
    results to persistence layer for use by engines.
    """

    injector_name = '_TaskFlow_INJECTOR'

    def __init__(self, flow_detail=None):
        self._result_mappings = {}
        self._reverse_mapping = {}

        if flow_detail is None:
            # TODO(imelnikov): this is useful mainly for tests;
            #  maybe we should make flow_detail required parameter?
            self._flowdetail = temporary_flow_detail()
        else:
            self._flowdetail = flow_detail

    def add_task(self, uuid, task_name):
        """Add the task to storage

        Task becomes known to storage by that name and uuid.
        Task state is set to PENDING.
        """
        # TODO(imelnikov): check that task with same uuid or
        #   task name does not exist
        td = taskdetail.TaskDetail(name=task_name, uuid=uuid)
        td.state = states.PENDING
        self._flowdetail.add(td)
        self._flowdetail.save()
        td.save()

    def get_uuid_by_name(self, task_name):
        """Get uuid of task with given name"""
        td = self._flowdetail.find_by_name(task_name)
        if td is not None:
            return td.uuid
        else:
            raise exceptions.NotFound("Unknown task name: %r" % task_name)

    def _taskdetail_by_uuid(self, uuid):
        td = self._flowdetail.find(uuid)
        if td is None:
            raise exceptions.NotFound("Unknown task: %r" % uuid)
        return td

    def set_task_state(self, uuid, state):
        """Set task state"""
        td = self._taskdetail_by_uuid(uuid)
        td.state = state
        td.save()

    def get_task_state(self, uuid):
        """Get state of task with given uuid"""
        return self._taskdetail_by_uuid(uuid).state

    def save(self, uuid, data, state=states.SUCCESS):
        """Put result for task with id 'uuid' to storage"""
        td = self._taskdetail_by_uuid(uuid)
        td.state = state
        td.results = data
        td.save()

    def get(self, uuid):
        """Get result for task with id 'uuid' to storage"""
        td = self._taskdetail_by_uuid(uuid)
        if td.state not in STATES_WITH_RESULTS:
            raise exceptions.NotFound("Result for task %r is not known" % uuid)
        return td.results

    def reset(self, uuid, state=states.PENDING):
        """Remove result for task with id 'uuid' from storage"""
        td = self._taskdetail_by_uuid(uuid)
        td.results = None
        td.state = state
        td.save()

    def inject(self, pairs):
        """Add values into storage

        This method should be used by job in order to put flow parameters
        into storage and put it to action.
        """
        pairs = dict(pairs)
        injector_uuid = uuidutils.generate_uuid()
        self.add_task(injector_uuid, self.injector_name)
        self.save(injector_uuid, pairs)
        self._reverse_mapping.update((name, (injector_uuid, name))
                                     for name in pairs)

    def set_result_mapping(self, uuid, mapping):
        """Set mapping for naming task results

        The result saved with given uuid would be accessible by names
        defined in mapping. Mapping is a dict name => index. If index
        is None, the whole result will have this name; else, only
        part of it, result[index].
        """
        if not mapping:
            return
        self._result_mappings[uuid] = mapping
        for name, index in mapping.iteritems():
            self._reverse_mapping[name] = (uuid, index)

    def fetch(self, name):
        """Fetch named task result"""
        try:
            uuid, index = self._reverse_mapping[name]
        except KeyError:
            raise exceptions.NotFound("Name %r is not mapped" % name)
        result = self.get(uuid)
        if index is None:
            return result
        else:
            return result[index]

    def fetch_all(self):
        """Fetch all named task results known so far

        Should be used for debugging and testing purposes mostly.
        """
        result = {}
        for name in self._reverse_mapping:
            try:
                result[name] = self.fetch(name)
            except exceptions.NotFound:
                pass
        return result

    def fetch_mapped_args(self, args_mapping):
        """Fetch arguments for the task using arguments mapping"""
        return dict((key, self.fetch(name))
                    for key, name in args_mapping.iteritems())

    def set_flow_state(self, state):
        """Set flowdetails state and save it"""
        self._flowdetail.state = state
        self._flowdetail.save()

    def get_flow_state(self):
        """Set state from flowdetails"""
        return self._flowdetail.state