/cvat/apps/redis_handler/views.py

from __future__ import annotations

import functools
from collections import namedtuple
from collections.abc import Iterable
from typing import cast

import django_rq
from django.conf import settings
from django.http import HttpResponseBadRequest, HttpResponseNotFound
from django.utils.decorators import method_decorator
from django.views.decorators.cache import never_cache
from django_rq.queues import DjangoRQ
from drf_spectacular.utils import (
    OpenApiParameter,
    OpenApiResponse,
    OpenApiTypes,
    extend_schema,
    extend_schema_view,
)
from redis.exceptions import ConnectionError as RedisConnectionError
from rest_framework import status, viewsets
from rest_framework.decorators import action
from rest_framework.response import Response
from rq.command import send_stop_job_command
from rq.exceptions import InvalidJobOperation
from rq.job import Job as RQJob
from rq.job import JobStatus as RQJobStatus
from rq.worker import Worker as RQWorker

from cvat.apps.engine.filters import (
    NonModelJsonLogicFilter,
    NonModelOrderingFilter,
    NonModelSimpleFilter,
)
from cvat.apps.engine.log import ServerLogManager
from cvat.apps.engine.models import RequestAction, RequestSubresource
from cvat.apps.engine.rq import is_rq_job_owner
from cvat.apps.engine.types import ExtendedRequest
from cvat.apps.redis_handler.apps import SELECTOR_TO_QUEUE
from cvat.apps.redis_handler.permissions import RequestPermission
from cvat.apps.redis_handler.rq import CustomRQJob, RequestId
from cvat.apps.redis_handler.serializers import RequestSerializer, RequestStatus

slogger = ServerLogManager(__name__)

CVAT_CAN_STOP_STARTED_JOBS_KEY = "cvat_can_stop_started_jobs"


@extend_schema(tags=["requests"])
@extend_schema_view(
    list=extend_schema(
        summary="List requests",
        responses={
            "200": RequestSerializer(many=True),
        },
    ),
    retrieve=extend_schema(
        summary="Get request details",
        parameters=[
            OpenApiParameter(
                "id",
                OpenApiTypes.STR,
                OpenApiParameter.PATH,
            )
        ],
        responses={
            "200": RequestSerializer,
        },
    ),
)
class RequestViewSet(viewsets.GenericViewSet):
    serializer_class = RequestSerializer
    iam_supports_organization_params = False
    iam_permission_class = RequestPermission
    filter_backends = [
        NonModelSimpleFilter,
        NonModelJsonLogicFilter,
        NonModelOrderingFilter,
    ]

    ordering_fields = ["created_date", "status", "action"]
    ordering = "-created_date"

    simple_filters = [
        # RQ job fields
        "status",
        # derivatives fields (from meta)
        "project_id",
        "task_id",
        "job_id",
        # derivatives fields (from parsed rq_id)
        "action",
        "target",
        "subresource",
        "format",
        # request context
        "org",
        "org_id",
    ]

    filter_fields = [*simple_filters]

    lookup_fields = {
        "created_date": "created_at",
        "action": "parsed_id.action",
        "target": "parsed_id.target",
        "subresource": "parsed_id.subresource",
        "format": "parsed_id.format",
        "status": "get_status",
        "project_id": "meta.project_id",
        "task_id": "meta.task_id",
        "job_id": "meta.job_id",
        "org": "meta.org_slug",
        "org_id": "meta.org_id",
    }

    SchemaField = namedtuple("SchemaField", ["type", "choices"], defaults=(None,))

    simple_filters_schema = {
        "status": SchemaField("string", RequestStatus.choices),
        "project_id": SchemaField("integer"),
        "task_id": SchemaField("integer"),
        "job_id": SchemaField("integer"),
        "action": SchemaField("string"),
        "target": SchemaField("string"),
        "subresource": SchemaField("string"),
        "format": SchemaField("string"),
        "org": SchemaField("string"),
        "org_id": SchemaField("integer"),
    }

    class _SchemaQuerySet:
        # drf-spectacular expects schema views to expose a queryset-like object
        # with a model attribute. Requests are backed by RQ jobs, not Django models.
        # See https://drf-spectacular.readthedocs.io/en/latest/faq.html#my-get-queryset-depends-on-some-attributes-not-available-at-schema-generation-time
        model = None

    def get_queryset(self):
        if getattr(self, "swagger_fake_view", False):
            return self._SchemaQuerySet()

        return None

    @property
    def queues(self) -> Iterable[DjangoRQ]:
        return (django_rq.get_queue(queue_name) for queue_name in set(SELECTOR_TO_QUEUE.values()))

    def _get_rq_jobs_from_queue(self, queue: DjangoRQ, user_id: int) -> list[RQJob]:
        job_ids = set(
            queue.get_job_ids()
            + queue.started_job_registry.get_job_ids()
            + queue.finished_job_registry.get_job_ids()
            + queue.failed_job_registry.get_job_ids()
            + queue.deferred_job_registry.get_job_ids()
        )
        jobs = []

        for job in queue.job_class.fetch_many(job_ids, queue.connection):
            if job and is_rq_job_owner(job, user_id):
                if job.get_status(refresh=False) in {RQJobStatus.CANCELED, RQJobStatus.STOPPED}:
                    continue

                job = cast(CustomRQJob, job)
                try:
                    parsed_request_id = RequestId.parse_and_validate_queue(
                        job.id, expected_queue=queue.name
                    )
                except Exception:  # nosec B112
                    continue

                job.parsed_id = parsed_request_id
                jobs.append(job)

        return jobs

    def _get_rq_jobs(self, user_id: int) -> list[RQJob]:
        """
        Get all RQ jobs for a specific user and return them as a list of RQJob objects.

        Parameters:
            user_id (int): The ID of the user for whom to retrieve jobs.

        Returns:
            list[RQJob]: A list of RQJob objects representing all jobs for the specified user.
        """
        all_jobs = []
        for queue in self.queues:
            jobs = self._get_rq_jobs_from_queue(queue, user_id)
            all_jobs.extend(jobs)

        return all_jobs

    def _get_rq_job_by_id(self, rq_id: str) -> CustomRQJob | None:
        """
        Get a RQJob by its ID from the queues.

        Args:
            rq_id (str): The ID of the RQJob to retrieve.

        Returns:
            RQJob | None: The retrieved RQJob, or None if not found.
        """
        try:
            parsed_request_id, queue_name = RequestId.parse(rq_id, try_legacy_format=True)
            rq_id = parsed_request_id.render()
        except Exception:
            return None

        queue: DjangoRQ = django_rq.get_queue(queue_name)
        job: CustomRQJob | None = queue.fetch_job(rq_id)

        if job:
            job.parsed_id = parsed_request_id

        return job

    @staticmethod
    def _is_export_request(rq_job: CustomRQJob) -> bool:
        return rq_job.parsed_id.action == RequestAction.EXPORT and rq_job.parsed_id.subresource in {
            RequestSubresource.ANNOTATIONS,
            RequestSubresource.DATASET,
            RequestSubresource.BACKUP,
        }

    @classmethod
    def _is_started_process_cancellable(cls, rq_job: CustomRQJob) -> bool:
        if not cls._is_export_request(rq_job):
            return False

        worker_key = f"{RQWorker.redis_worker_namespace_prefix}{rq_job.worker_name}"
        # The marker is CVAT-specific and can be absent on already running or older production
        # workers. Only SimpleWorker writes "0", so default to "can stop" for backward
        # compatibility.
        can_stop_started_jobs = (
            rq_job.connection.hget(worker_key, CVAT_CAN_STOP_STARTED_JOBS_KEY) or "1"
        )

        if isinstance(can_stop_started_jobs, bytes):
            can_stop_started_jobs = can_stop_started_jobs.decode()

        return can_stop_started_jobs != "0"

    def _handle_redis_exceptions(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            try:
                return func(*args, **kwargs)
            except RedisConnectionError as ex:
                msg = "Redis service is not available"
                slogger.glob.exception(f"{msg}: {str(ex)}")
                return Response(msg, status=status.HTTP_503_SERVICE_UNAVAILABLE)

        return wrapper

    @method_decorator(never_cache)
    @_handle_redis_exceptions
    def retrieve(self, request: ExtendedRequest, pk: str):
        job = self._get_rq_job_by_id(pk)

        if not job:
            return HttpResponseNotFound("There is no request with specified id")

        self.check_object_permissions(request, job)

        if job.get_status(refresh=False) in {RQJobStatus.CANCELED, RQJobStatus.STOPPED}:
            job.delete()
            return HttpResponseNotFound("There is no request with specified id")

        serializer = self.get_serializer(job, context={"request": request})
        return Response(data=serializer.data, status=status.HTTP_200_OK)

    @method_decorator(never_cache)
    @_handle_redis_exceptions
    def list(self, request: ExtendedRequest):
        user_id = request.user.id
        user_jobs = self._get_rq_jobs(user_id)

        filtered_jobs = self.filter_queryset(user_jobs)

        page = self.paginate_queryset(filtered_jobs)
        if page is not None:
            serializer = self.get_serializer(page, many=True, context={"request": request})
            return self.get_paginated_response(serializer.data)

        serializer = self.get_serializer(filtered_jobs, many=True, context={"request": request})
        return Response(data=serializer.data, status=status.HTTP_200_OK)

    @extend_schema(
        summary="Cancel request",
        parameters=[
            OpenApiParameter(
                "id",
                OpenApiTypes.STR,
                OpenApiParameter.PATH,
            )
        ],
        request=None,
        responses={
            "200": OpenApiResponse(description="The request has been cancelled"),
        },
    )
    @method_decorator(never_cache)
    @action(detail=True, methods=["POST"], url_path="cancel")
    @_handle_redis_exceptions
    def cancel(self, request: ExtendedRequest, pk: str):
        rq_job = self._get_rq_job_by_id(pk)

        if not rq_job:
            return HttpResponseNotFound("There is no request with specified id")

        self.check_object_permissions(request, rq_job)
        rq_job_status = rq_job.get_status(refresh=False)

        # Terminal canceled jobs are not useful for users in the requests API.
        if rq_job_status in {RQJobStatus.CANCELED, RQJobStatus.STOPPED}:
            return HttpResponseNotFound("There is no request with specified id")

        # Jobs that have not started yet are safe to cancel directly. RQ will also enqueue
        # dependents when ONE_RUNNING_JOB_IN_QUEUE_PER_USER is enabled.
        if rq_job_status in {RQJobStatus.QUEUED, RQJobStatus.DEFERRED}:
            # FUTURE-TODO: race condition is possible here
            rq_job.cancel(enqueue_dependents=settings.ONE_RUNNING_JOB_IN_QUEUE_PER_USER)
            rq_job.delete()
            return Response(status=status.HTTP_200_OK)

        # Finished, failed, scheduled, and other non-started states cannot be interrupted.
        if rq_job_status != RQJobStatus.STARTED:
            return HttpResponseBadRequest(
                f"Requests with status {rq_job_status!r} cannot be cancelled"
            )

        # Started cancellation is intentionally limited to non-mutating export jobs running in
        # production-style workers.
        if not self._is_started_process_cancellable(rq_job):
            return HttpResponseBadRequest(
                "Cancellation of started requests is supported only for export requests "
                "executed by production workers"
            )

        # RQ will stop the forked work horse and enqueue dependent jobs from its stopped-job path.
        queue = django_rq.get_queue(rq_job.origin)
        try:
            send_stop_job_command(queue.connection, rq_job.id, serializer=queue.serializer)
        except InvalidJobOperation as ex:
            slogger.glob.warning("Failed to stop RQ job %s: %s", rq_job.id, ex)
            return HttpResponseBadRequest("The request cannot be cancelled")

        return Response(status=status.HTTP_200_OK)