Source code for luna_tasks.app.handlers.results_proxy_handler

from luna3.common.exceptions import LunaApiException
from sanic.response import HTTPResponse

from app.handlers.base_handler import BaseRequestHandler
from crutches_on_wheels.cow.enums.tasks import TaskStatus
from crutches_on_wheels.cow.errors.errors import Error
from crutches_on_wheels.cow.errors.exception import VLException
from crutches_on_wheels.cow.web.query_getters import uriParamIntGetter, uuidGetter


class TaskResultHandler(BaseRequestHandler):
    """
    Task result handler.
    """

    async def get(self, taskId: str) -> HTTPResponse:
        """
        Get error, see `spec`_

        .. _spec:
            _static/api.html#operation/getTaskResult

        Returns:
            response with result and "Content-Disposition" header
        """
        taskId = uriParamIntGetter(taskId, Error.PageNotFoundError)
        accountId = self.getQueryParam("account_id", uuidGetter)
        task = await self.dbContext.getTask(taskId=taskId, accountId=accountId)
        resultId = task["result_id"]
        if task["task_status"] in (TaskStatus.cancelled.value, TaskStatus.failed.value):
            error = Error.ImpossibleGetTaskResult.format(taskId, task["task_status"])
            raise VLException(error, 400, isCriticalError=False)
        if resultId is None:
            error = Error.TaskDoesNotHaveResultYet.format(taskId)
            raise VLException(error, 404, isCriticalError=False)
        try:
            resultResponse = await self.luna3Client.lunaTasksResultsStore.getObject(
                self.config.taskResultStorage.bucket, objectId=resultId, raiseError=True
            )
        except LunaApiException as e:
            if e.statusCode == 404:
                error = Error.TaskResultNotFound.format(taskId)
                self.logger.warning(f"result of task {taskId} with result id {resultId} not found")
                return self.error(404, error)
            raise

        extension = resultResponse.contentType.split("/")[-1]
        filename = f"task_{taskId}.{extension}"
        return self.success(
            200,
            resultResponse.body,
            contentType=resultResponse.contentType,
            extraHeaders={"Content-Disposition": f"attachment; filename={filename}"},
        )