cr-connect-workflow/crc/scripts/complete_template.py

from io import StringIO, BytesIO

from jinja2 import UndefinedError

from crc import session
from crc.api.common import ApiError
from crc.models.file import FileModel, FileDataModel, CONTENT_TYPES
from crc.models.workflow import WorkflowSpecModel
from docxtpl import DocxTemplate
import jinja2

from crc.scripts.script import Script
from crc.services.file_service import FileService
from crc.services.workflow_processor import WorkflowProcessor


class CompleteTemplate(Script):

    def get_description(self):
        return """Takes one argument, which is the name of a MS Word docx file to use as a template.
        All data currently collected up to this Task will be available for use in the template."""

    def do_task(self, task, study_id, *args, **kwargs):
        """Entry point, mostly worried about wiring it all up."""
        if len(args) != 1:
            raise ApiError(code="missing_argument",
                           message="The CompleteTask script requires a single argument with "
                                   "the name of the docx template to use.")
        file_name = args[0]
        workflow_spec_model = self.find_spec_model_in_db(task.workflow)

        if workflow_spec_model is None:
            raise ApiError(code="workflow_model_error",
                           message="Something is wrong.  I can't find the workflow you are using.")

        file_data_model = session.query(FileDataModel) \
            .join(FileModel) \
            .filter(FileModel.name == file_name) \
            .filter(FileModel.workflow_spec_id == workflow_spec_model.id).first()

        if file_data_model is None:
            raise ApiError(code="file_missing",
                           message="Can not find a file called '%s' "
                                   "within workflow specification '%s'") % (args[0], workflow_spec_model.id)

        final_document_stream = self.make_template(BytesIO(file_data_model.data), task.data)
        study_id = task.workflow.data[WorkflowProcessor.STUDY_ID_KEY]
        workflow_id = task.workflow.data[WorkflowProcessor.WORKFLOW_ID_KEY]
        FileService.add_task_file(study_id=study_id, workflow_id=workflow_id, task_id=task.id,
                                  name=file_name,
                                  content_type=CONTENT_TYPES['docx'],
                                  binary_data=final_document_stream.read())

        print("Complete Task was called with %s" % str(args))

    def make_template(self, binary_stream, context):
        doc = DocxTemplate(binary_stream)
        jinja_env = jinja2.Environment(autoescape=True)
        doc.render(context, jinja_env)
        target_stream = BytesIO()
        doc.save(target_stream)
        target_stream.seek(0) # move to the beginning of the stream.
        return target_stream

    def find_spec_model_in_db(self, workflow):
        """ Search for the workflow """
        # When the workflow spec model is created, we record the primary process id,
        # then we can look it up.  As there is the potential for sub-workflows, we
        # may need to travel up to locate the primary process.
        spec = workflow.spec
        workflow_model = session.query(WorkflowSpecModel). \
            filter(WorkflowSpecModel.primary_process_id == spec.name).first()
        if workflow_model is None and workflow != workflow.outer_workflow:
            return self.find_spec_model_in_db(workflow.outer_workflow)

        return workflow_model