d3i-szct/dock/dcm/dcm_push_upload.py

"""
数字城管文件上传。
"""
import asyncio
import datetime
import io
import json
import logging
from typing import Optional

from tornado.httpclient import HTTPResponse, HTTPRequest

import dock
from dock.dcm import dcm_api
from dock.oa import PushException
from models.dcm_task import DcmTask
from paste.core.logging import echo_log
from paste.util import udict
from paste.web import requests

ApplyPostponeTmp = "apply_postpone_tmp"
ApplyRollbackTmp = "apply_rollback_tmp"
RollbackTmp = "rollback_tmp"
TransitTmp = "transit_tmp"

MediaUsage = [ApplyPostponeTmp, ApplyRollbackTmp, RollbackTmp, TransitTmp]
"""
可用媒体对象。
"""


async def get_upload_request(dcm_task: DcmTask, media_usage: str, files: dict[str, io.IOBase],
                             relation_sub_id: Optional[int] = None,
                             relation_main_id: Optional[int] = 45):
    """
    创建文件上传请求。因数字城管服务器仅接受单文档上传，因此返回的是请求对象列表。

    :param dcm_task: 待办工单
    :param media_usage: 媒体使用
    :param files: 要上传的文件数据
    :param relation_sub_id: 可选参数
    :param relation_main_id: 可选参数
    :return: 上传请求列表
    """
    api_url = f"/media/upload"
    assert media_usage in MediaUsage, f"参数错误，须为：{MediaUsage} 之一."

    request_list = []
    _id = 0
    for key, value in files.items():
        request_body = {
            "relationTypeID": 10,
            "mediaUsage": media_usage,
            "checkExist": True,
            "relationID": dcm_task.rec_id,
            "relationSubID": relation_sub_id if relation_sub_id is not None else dcm_task.act_id,
            "relationMainID": relation_main_id,
            "tempUsage": media_usage,
            "id": f"WU_FILE_{_id}",
            "name": key,
            "type": 'image/png',
            "lastModifiedDate": datetime.datetime.now().strftime("%a %b %d %Y %H:%M:%S %z").replace(
                "+0800", "GMT+0800"
            ),
            "size": len(value.read()),
            key: value,
        }
        _id += 1
        # 构造 API 请求
        request = await dcm_api.new_api_request(api_url, request_body)
        setattr(request, 'dcm_task', dcm_task)
        setattr(request, 'media_usage', media_usage)
        setattr(request, 'file_name', key)
        request_list.append(request)
    # 返回请求列表
    return request_list


async def get_media_request(dcm_task: DcmTask, media_usage: str):
    """
    创建读取 MediaID 请求。

    :param dcm_task: 待办工单
    :param media_usage: 媒体使用
    :return: 上传请求列表
    """
    api_url = f"/media/get"
    assert media_usage in MediaUsage, f"参数错误，须为：{MediaUsage} 之一."
    request_body = {
        "relationTypeID": 10,
        "mediaUsage": media_usage,
        "relationID": dcm_task.rec_id,
        "relationSubID": dcm_task.act_id,
        "relationMainID": 45,
    }
    # 构造 API 请求
    request = await dcm_api.new_api_request(api_url, request_body, method="GET")
    setattr(request, 'dcm_task', dcm_task)
    setattr(request, 'media_usage', media_usage)
    return request


async def on_upload_error(request: HTTPRequest, exc: Exception, retry_queue: asyncio.Queue = None):
    """
    下载附件时的出错处理。

    :param request: 请求对象
    :param exc: 异常对象
    :param retry_queue:
    :return:
    """
    retry = getattr(request, 'retry', 0)
    max_retry = getattr(request, 'max_retry', 0)
    if retry < max_retry - 1:
        # 非最后一次尝试，不处理
        return

    dcm_task: DcmTask = getattr(request, 'dcm_task')
    media_usage = getattr(request, 'media_usage')
    file_name = getattr(request, 'file_name')
    message = f'工单ID：{dcm_task.id}，类型为：{media_usage}，文件名为：{file_name} 的附件上传失败.'
    echo_log(message)
    echo_log(exc, logging.ERROR, is_log_exc=True)

    # 保存异常
    exc_list = getattr(dcm_task, 'exc_list')
    exc_list.append(PushException(message))


async def after_upload_request(response: HTTPResponse, retry_queue: asyncio.Queue[HTTPRequest]):
    """
    提交数字城管后的处理程序。

    :param response: 响应对象
    :param retry_queue: 重试队列
    """
    dcm_task: DcmTask = getattr(response.request, 'dcm_task')
    media_usage = getattr(response.request, 'media_usage')
    file_name = getattr(response.request, 'file_name')
    echo_log(response.body.decode())
    echo_log(f'工单ID：{dcm_task.id}，类型为：{media_usage}，文件名为：{file_name} 的附件上传成功.')


async def push_upload(dcm_task: DcmTask, media_usage: str, files: dict[str, io.IOBase],
                      relation_sub_id: Optional[int] = None,
                      relation_main_id: Optional[int] = 45):
    """
    向数字城管上传文件。

    :param dcm_task: 待办工单
    :param media_usage: 媒体使用
    :param files: 要上传的文件数据
    :param relation_sub_id: 可选参数
    :param relation_main_id: 可选参数
    :return:
    """
    echo_log(f"正在准备上传队列...")
    request_list = await get_upload_request(dcm_task, media_usage, files, relation_sub_id, relation_main_id)
    upload_request_queue = asyncio.Queue()
    for req in request_list:
        await upload_request_queue.put(req)

    # 并发提交上传请求
    await requests.async_concurrency(
        upload_request_queue, con_count=dock.CONCURRENCY_COUNT, retry=dock.MAX_RETRY_COUNT,
        after_request=after_upload_request, on_error=on_upload_error
    )

    # 上传后，读取文件 MediaID
    media_request = await get_media_request(dcm_task, media_usage)
    media_request_queue = asyncio.Queue()
    await media_request_queue.put(media_request)
    media_response_list: list[HTTPResponse] = await requests.async_concurrency(
        media_request_queue, con_count=dock.CONCURRENCY_COUNT, retry=dock.MAX_RETRY_COUNT
    )
    # 从响应列表中取得媒体ID和媒体类型ID
    media_id_list: list = []
    media_type_list: list = []
    media_num = 0
    if media_response_list:
        response = media_response_list[0]
        response_body = response.body.decode()
        response_data = json.loads(response_body)
        media_list = udict.get_by_path(response_data, 'data.mediaList')
        media_num = len(media_list)
        for media in media_list:
            media_id_list.append(f"{media.get('mediaID')}")
            media_type_list.append(media.get('mediaType'))

    # 用英文逗号拼接后返回
    return ",".join(media_id_list), ",".join(media_type_list), media_num


if __name__ == "__main__":
    from paste.core import aio_pool

    _runner = aio_pool.get_aio_runner()