dify/api/controllers/console/billing/compliance.py

from typing import assert_never

from flask import request
from flask_restx import Resource
from pydantic import BaseModel, Field

from controllers.common.schema import (
    query_params_from_model,
    register_response_schema_models,
    register_schema_models,
)
from controllers.console import console_ns
from controllers.console.billing.error import (
    BillingOperationFailedError,
    BillingOperationFailedErrorResponse,
    BillingRequestInProgressError,
    BillingUnavailableError,
    BillingUnavailableErrorResponse,
    BillingUnprocessableEntityErrorResponse,
    ComplianceDocumentNotFoundError,
    ComplianceDownloadForbiddenError,
    ComplianceRateLimitError,
    ComplianceRateLimitErrorResponse,
    to_billing_request_error,
)
from controllers.console.flask_admission import console_account_admission
from controllers.console.wraps import model_validate
from enums import DeploymentEdition
from extensions.ext_application_services import application_services
from fields.base import ResponseModel
from libs.exception import BaseHTTPException
from libs.external_api import error_responses
from libs.helper import dump_response, extract_remote_ip
from machinery.context import RequestContext
from services.billing_client import billing
from services.errors.billing import ComplianceRateLimitExceededError, UnknownErrorReason


class ComplianceDownloadQuery(BaseModel):
    doc_name: str = Field(..., description="Compliance document name")


class ComplianceDownloadResponse(ResponseModel):
    url: str


register_schema_models(console_ns, ComplianceDownloadQuery)
register_response_schema_models(
    console_ns,
    BillingOperationFailedErrorResponse,
    BillingUnavailableErrorResponse,
    BillingUnprocessableEntityErrorResponse,
    ComplianceDownloadResponse,
    ComplianceRateLimitErrorResponse,
)


def _compliance_download_error(error: billing.BillingComplianceUpstreamError) -> BaseHTTPException:
    match error.reason:
        case "COMPLIANCE_DOCUMENT_NOT_FOUND":
            return ComplianceDocumentNotFoundError()
        case "COMPLIANCE_DOWNLOAD_FORBIDDEN":
            return ComplianceDownloadForbiddenError()
        case "REQUEST_IN_PROGRESS":
            return BillingRequestInProgressError()
        case (
            ""
            | "CODEC"
            | "VALIDATOR"
            | "UNKNOWN"
            | "INVALID_SECRET_KEY"
            | "MISSING_TENANT_ID"
            | "INVALID_TENANT_ID"
            | "MISSING_ACCOUNT_ID"
            | "INVALID_ACCOUNT_ID"
            | "MISSING_PREFILLED_EMAIL"
            | "INVALID_COMPLIANCE_DOWNLOAD_REQUEST"
            | "MISSING_DOC_NAME"
            | "MISSING_IP_ADDRESS"
            | "MISSING_DEVICE_INFO"
            | "COMPLIANCE_DOWNLOAD_FAILED"
        ):
            return to_billing_request_error(error)
        case UnknownErrorReason() | None:
            return to_billing_request_error(error)
        case unreachable:
            assert_never(unreachable)


@console_ns.route("/compliance/download")
class ComplianceApi(Resource):
    @error_responses(
        BillingOperationFailedError,
        BillingRequestInProgressError,
        BillingUnavailableError,
        ComplianceDocumentNotFoundError,
        ComplianceDownloadForbiddenError,
        ComplianceRateLimitError,
    )
    @console_ns.doc(params=query_params_from_model(ComplianceDownloadQuery))
    @console_ns.doc("download_compliance_document")
    @console_ns.doc(description="Get compliance document download link")
    @console_ns.response(200, "Success", console_ns.models[ComplianceDownloadResponse.__name__])
    @console_ns.response(
        422,
        "Invalid compliance download query",
        console_ns.models[BillingUnprocessableEntityErrorResponse.__name__],
    )
    @console_ns.response(
        429,
        "Compliance download rate limit exceeded",
        console_ns.models[ComplianceRateLimitErrorResponse.__name__],
    )
    @console_ns.response(
        502,
        "Compliance download failed",
        console_ns.models[BillingOperationFailedErrorResponse.__name__],
    )
    @console_ns.response(
        503,
        "Billing unavailable",
        console_ns.models[BillingUnavailableErrorResponse.__name__],
    )
    @console_account_admission(editions=frozenset({DeploymentEdition.CLOUD}))
    @model_validate(ComplianceDownloadQuery)
    def get(self, req_data: ComplianceDownloadQuery, request_context: RequestContext):
        ip_address = extract_remote_ip(request)
        device_info = request.headers.get("User-Agent", "Unknown device")
        try:
            data = application_services().compliance_downloads.get_link(
                request_context=request_context,
                document_name=req_data.doc_name,
                ip_address=ip_address,
                device_info=device_info,
            )
        except ComplianceRateLimitExceededError as error:
            raise ComplianceRateLimitError() from error
        except billing.BillingComplianceUpstreamError as error:
            raise _compliance_download_error(error) from error
        return dump_response(ComplianceDownloadResponse, data)