Source code for luna_tasks.app.handlers.results_proxy_handler

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, ErrorInfo
from crutches_on_wheels.cow.errors.exception import VLException
from crutches_on_wheels.cow.web.query_getters import uriParamIntGetter, uuidGetter
from crutches_on_wheels.cow.web.utils import decodeBytes, loadJson


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
        """
        self.request.streamResponse = True
        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)

        resultResponse = await self.luna3Client.lunaTasksResultsStore.getObject(
            self.config.taskResultStorage.bucket, objectId=resultId, accountId=accountId, stream=True, totalTimeout=0
        )

        async with resultResponse.sendRequest() as resultResp:
            if resultResp.rawResponse.status == 404:
                # Need to use separate error for task result, we don't want to proxy error from image store
                return self.error(404, Error.TaskResultNotFound.format(taskId))
            elif resultResp.rawResponse.status != 200:
                errorBody = await resultResp.rawResponse.read()
                errorInfo = loadJson(decodeBytes(body=errorBody))
                return self.error(resultResp.rawResponse.status, ErrorInfo.fromDict(errorInfo))

            extension = resultResp.contentType.split("/")[-1]
            filename = f"task_{taskId}.{extension}"
            resp = await self.request.respond(
                content_type=resultResponse.contentType,
                headers={"Content-Disposition": f"attachment; filename={filename}"},
            )
            async for chunk in resultResp.iter():
                await resp.send(chunk)
            await resp.eof()