d3i-szct/models/govs_order_attachment.py

# coding: utf-8
import datetime
from typing import Union

import pandas as pd
from sqlalchemy import select, delete
from tornado_swagger.model import register_swagger_model
from wtforms import StringField, IntegerField
from wtforms.validators import Length

from models.common_model import CommonModel
from models.db_models import TD3iGovsOrderAttachment
from paste.core.logging import echo_log
from paste.rbac.rbac_user import RbacUser
from paste.web.form import ModelForm


class GovsOrderAttachmentForm(ModelForm):
    """工单附件表单验证类"""

    id = IntegerField('附件唯一ID')
    master_id = IntegerField('关联工单主表ID')
    order_id = StringField('工单编号', validators=[Length(max=50)])
    file_path = StringField('文件路径（内网地址）', validators=[Length(max=500)])
    out_file_path = StringField('外网文件路径', validators=[Length(max=500)])
    attach_name = StringField('附件名称', validators=[Length(max=200)])
    to_tenant_id = StringField('目标租户ID', validators=[Length(max=50)])
    create_date = StringField('记录创建时间')

    def process(self, formdata=None, obj=None, **kwargs):
        if formdata:
            for name, values in formdata.items():
                if isinstance(values, list) and values:
                    formdata[name] = [v.strip() if isinstance(v, str) else v for v in values]
                elif isinstance(values, str):
                    formdata[name] = values.strip()
        super().process(formdata, obj, **kwargs)


class GovsOrderAttachmentBase(TD3iGovsOrderAttachment, CommonModel):
    """工单附件业务基类"""

    FieldMapping = {
        # ==================== 主键与关联 ====================
        'id': 'id',                         # 附件唯一ID

        # ==================== 文件信息 ====================
        'file_path': 'filePath',            # 文件路径（内网地址）
        'out_file_path': 'outFilePath',     # 外网文件路径
        'attach_name': 'attachName',        # 附件名称
        'to_tenant_id': 'toTenantId',       # 目标租户ID
    }

    @classmethod
    async def find_by_order_id(cls, order_id: str):
        """根据工单编号查找附件"""
        _query = select(cls).where(cls.order_id == order_id)
        return (await cls.orm_execute_scalars(_query)).all()

    @classmethod
    async def find_by_master_id(cls, master_id: Union[str, int]):
        """根据主工单ID查找附件"""
        _query = select(cls).where(cls.master_id == master_id)
        return (await cls.orm_execute_scalars(_query)).all()

    @classmethod
    async def find_by_attach_name(cls, attach_name: str):
        """根据附件名称查找"""
        _query = select(cls).where(cls.attach_name == attach_name)
        return (await cls.orm_execute_scalars(_query)).all()

    @classmethod
    async def exists_order_id(cls, data_df: pd.DataFrame):
        """根据 order_id 判断数据是否存在
        :param data_df: 输入的数据框架，必须包含 id 和 order_id 列
        :return: (exists_df: pd.DataFrame, latest_df: pd.DataFrame)
            - exists_df: 在数据库中存在的记录（附带数据库中的 id）
            - latest_df: 在数据库中不存在的记录
        """
        if data_df.empty:
            return pd.DataFrame(), pd.DataFrame()

        # 获取待查询的 (id, order_id) 组合
        pairs = data_df[[cls.id.key, cls.order_id.key]].drop_duplicates().values.tolist()
        if not pairs:
            return pd.DataFrame(), data_df.copy()

        # 查询数据库中已存在的记录（使用 IN 批量查询）
        _query = select(cls.id, cls.order_id).where(
            (cls.id.in_([p[0] for p in pairs])) &
            (cls.order_id.in_([p[1] for p in pairs]))
        )
        exists_db_df = await cls.query_as_df(_query)

        if exists_db_df.empty:
            return pd.DataFrame(), data_df.copy()

        exists_db_df[cls.id.key] = exists_db_df[cls.id.key].astype(str)
        exists_db_df[cls.order_id.key] = exists_db_df[cls.order_id.key].astype(str)
        # 构建 (id, order_id) -> id 的映射（用于快速查找）
        key_to_id_map = dict(zip(
            zip(exists_db_df[cls.id.key], exists_db_df[cls.order_id.key]),
            exists_db_df[cls.id.key]
        ))

        # 标记 data_df 中哪些行在数据库中存在
        mask_exists = data_df.apply(
            lambda row: (row[cls.id.key], row[cls.order_id.key]) in key_to_id_map,
            axis=1
        )
        # 提取存在的记录，并补充数据库中的 id（虽然输入中已有 id，但为一致性保留）
        exists_df = data_df[mask_exists].copy()
        exists_df[cls.id.key] = exists_df.apply(
            lambda row: key_to_id_map[(row[cls.id.key], row[cls.order_id.key])],
            axis=1
        )
        # 提取不存在的记录
        latest_df = data_df[~mask_exists].copy()

        return exists_df, latest_df


@register_swagger_model
class GovsOrderAttachment(GovsOrderAttachmentBase):
    """工单附件业务类"""

    @classmethod
    async def create(cls, user: RbacUser = None, **kwargs):
        """创建附件记录"""
        for _k, _v in kwargs.items():
            if isinstance(_v, str):
                kwargs[_k] = _v.strip()

        _form = GovsOrderAttachmentForm(formdata=kwargs)
        _form.validate_form()

        # 检查是否已存在（根据 id）
        _existing = await cls.async_find_by_id(_form.id.data)
        if _existing:
            # 更新已有记录
            _existing.copy_from_dict(_form.data, skip_none=True)
            if user:
                _existing.updated_by = user.username
            await _existing.async_save()
            return _existing

        _obj = cls().copy_from_dict(_form.data, skip_none=True)
        if user:
            _obj.created_by = user.username
            _obj.updated_by = user.username
        await _obj.async_save()
        return _obj

    @classmethod
    async def delete(cls, obj_id: Union[str, int]):
        """删除附件记录"""
        _obj: cls = await cls.async_find_by_id(obj_id)
        assert _obj, f"根据 ID {obj_id} 未找到附件记录。"

        _del_query = delete(cls).where(cls.id == _obj.id)
        await cls.raw_execute(_del_query)
        echo_log(f'已删除附件记录（order_id：{_obj.order_id}，ID：{_obj.id}）.')
        return _obj

    @classmethod
    async def modify(cls, obj_id: Union[str, int], user: RbacUser = None, **kwargs):
        """修改附件记录"""
        for _k, _v in kwargs.items():
            if isinstance(_v, str):
                kwargs[_k] = _v.strip()

        _form = GovsOrderAttachmentForm(formdata=kwargs)
        _form.validate_form()

        _obj: cls = await cls.async_find_by_id(obj_id)
        assert _obj, f'查无此附件记录。'

        _obj.copy_from_dict(_form.data, skip_none=True)
        if user:
            _obj.updated_by = user.username
        await _obj.async_save()
        return _obj

    @classmethod
    async def create_batch(cls, data_df: pd.DataFrame, user: RbacUser = None):
        """批量创建附件记录"""
        if data_df.empty:
            return 0

        if user:
            data_df['created_by'] = user.username
            data_df['updated_by'] = user.username

        records = data_df.to_dict('records')
        objs = [cls().copy_from_dict(record, skip_none=True) for record in records]

        session = cls.get_aio_session()
        try:
            session.add_all(objs)
            await session.commit()
        except Exception as e:
            await session.rollback()
            raise e
        finally:
            await session.close()
        echo_log(f"批量创建成功：创建 {len(objs)} 条附件记录。")
        return len(objs)

    @classmethod
    async def modify_batch(cls, data_df: pd.DataFrame, user: RbacUser = None):
        """批量修改附件记录"""
        if data_df.empty:
            return 0

        if 'id' not in data_df.columns:
            echo_log(f"错误：modify_batch 要求输入数据必须包含 '{cls.id.key}' 列")
            return 0

        data_df['updated_at'] = datetime.datetime.now()
        if user:
            data_df['updated_by'] = user.username

        update_data = data_df.to_dict('records')

        session = cls.get_aio_session()
        try:
            await session.run_sync(
                lambda sync_session: sync_session.bulk_update_mappings(cls, update_data)
            )
            await session.commit()
            updated_count = len(update_data)
        except Exception as e:
            await session.rollback()
            raise e
        finally:
            await session.close()

        echo_log(f"批量修改成功：更新 {updated_count} 条附件记录。")
        return updated_count

    @classmethod
    async def save_batch(cls, data_df: pd.DataFrame, user: RbacUser = None):
        """批量保存数据，自动处理新建和更新"""
        _exists_df, _latest_df = await cls.exists_order_id(data_df)
        _created_count = await cls.create_batch(_latest_df, user)
        _updated_count = await cls.modify_batch(_exists_df, user)
        return _created_count, _updated_count

    @classmethod
    async def create_or_update_by_id(cls, user: RbacUser = None, **kwargs):
        """根据 id 创建或更新附件记录"""
        for _k, _v in kwargs.items():
            if isinstance(_v, str):
                kwargs[_k] = _v.strip()

        _form = GovsOrderAttachmentForm(formdata=kwargs)
        _form.validate_form()

        _existing = await cls.async_find_by_id(_form.id.data)
        if _existing:
            _existing.copy_from_dict(_form.data, skip_none=True)
            if user:
                _existing.updated_by = user.username
            await _existing.async_save()
            return _existing

        _obj = cls().copy_from_dict(_form.data, skip_none=True)
        if user:
            _obj.created_by = user.username
            _obj.updated_by = user.username
        await _obj.async_save()
        return _obj

    @classmethod
    async def delete_by_master_id(cls, master_id: Union[str, int]):
        """根据主工单ID删除附件记录"""
        attachments = await cls.find_by_master_id(master_id)
        if attachments:
            for att in attachments:
                await cls.delete(att.id)
        return attachments

    @classmethod
    async def delete_by_order_id(cls, order_id: str):
        """根据工单编号删除附件记录"""
        attachments = await cls.find_by_order_id(order_id)
        if attachments:
            for att in attachments:
                await cls.delete(att.id)
        return attachments