IntelligentRailwayCosting/SourceCode/IntelligentRailwayCosting/app/stores/railway_costing_mysql/project_quota.py

from sqlalchemy import and_, or_
from datetime import datetime
from typing import Optional

import tools.db_helper as db_helper
from core.dtos import ProjectQuotaDto
from core.enum import SendStatusEnum
from core.models import ProjectQuotaModel
from core.user_session import UserSession


class ProjectQuotaStore:
    def __init__(self):
        # self._database= None
        self._database = "iwb_railway_costing_v1"
        self._current_user = None

    @property
    def current_user(self):
        if self._current_user is None:
            self._current_user = UserSession.get_current_user()
        return self._current_user

    def get_quotas_paginated(
            self,
            budget_id: int,
            project_id: str,
            item_code: str,
            page: int = 1,
            page_size: int = 10,
            keyword: Optional[str] = None,
            send_status: Optional[int] = None,
    ):
        """分页查询定额列表

        Args:
            page: 页码，从1开始
            page_size: 每页数量
            project_id: 项目编号
            budget_id: 概算序号
            item_code: 条目序号
            keyword: 关键字
            send_status: 发送状态

        Returns:
            Tuple[total_count, quotas]
        """
        with db_helper.mysql_query_session(self._database) as db_session:
            query = db_session.query(ProjectQuotaModel)
            # 构建查询条件
            conditions = [
                ProjectQuotaModel.is_del == 0,
                ProjectQuotaModel.project_id == project_id,
                ProjectQuotaModel.budget_id == budget_id,
                ProjectQuotaModel.item_code.like(f"{item_code}%"),
            ]

            if send_status is not None:
                conditions.append(ProjectQuotaModel.send_status == send_status)
            if keyword:
                conditions.append(
                    or_(
                        ProjectQuotaModel.quota_code.like(f"%{keyword}%"),
                        ProjectQuotaModel.entry_name.like(f"%{keyword}%"),
                    )
                )
            query = query.filter(and_(*conditions))

            # 计算总数
            total_count = query.count()

            # 分页
            query = (
                query.order_by(ProjectQuotaModel.created_at.desc())
                .offset((page - 1) * page_size)
                .limit(page_size)
            )

            quotas = query.all()

            return {"total": total_count, "data": quotas}

    def get_quotas_by_task_paginated(
            self,
            task_id: int,
            budget_id,
            project_id,
            item_code,
            page: int = 1,
            page_size: int = 10,
            keyword: Optional[str] = None,
            send_status: Optional[int] = None,
    ):
        with db_helper.mysql_query_session(self._database) as db_session:
            query = db_session.query(ProjectQuotaModel).filter(
                and_(
                    ProjectQuotaModel.task_id == task_id,
                    ProjectQuotaModel.budget_id == budget_id,
                    ProjectQuotaModel.project_id == project_id,
                    ProjectQuotaModel.item_code.like(f"{item_code}%"),
                    ProjectQuotaModel.is_del == 0,
                )
            )
            if keyword:
                query = query.filter(
                    or_(
                        ProjectQuotaModel.quota_code.like(f"%{keyword}%"),
                        ProjectQuotaModel.entry_name.like(f"%{keyword}%"),
                    )
                )
            if send_status is not None:
                query = query.filter(ProjectQuotaModel.send_status == send_status)
            # 计算总数
            total_count = query.count()

            # 分页
            query = (
                query.order_by(ProjectQuotaModel.id.desc())
                .offset((page - 1) * page_size)
                .limit(page_size)
            )

            quotas = query.all()

            return {"total": total_count, "data": quotas}

    def get_quotas_by_task_id(self, task_id: int, with_quota_code: bool = False):
        with db_helper.mysql_query_session(self._database) as db_session:
            query = db_session.query(ProjectQuotaModel).filter(
                and_(
                    ProjectQuotaModel.task_id == task_id, ProjectQuotaModel.is_del == 0
                )
            )
            if with_quota_code:
                query = query.filter(
                    and_(
                        ProjectQuotaModel.quota_code != None,
                        ProjectQuotaModel.quota_code != "",
                    )
                )
            quotas = query.all()
            return quotas

    def get_quota(self, quota_id: int) -> Optional[ProjectQuotaModel]:
        """根据ID查询定额

        Args:
            quota_id: 定额ID

        Returns:
            Optional[ProjectQuotaDto]
        """
        with db_helper.mysql_query_session(self._database) as db_session:
            quota = (
                db_session.query(ProjectQuotaModel)
                .filter(
                    and_(
                        ProjectQuotaModel.id == quota_id, ProjectQuotaModel.is_del == 0
                    )
                )
                .first()
            )
            return quota

    def get_quota_dto(self, quota_id: int) -> Optional[ProjectQuotaDto]:
        """根据ID查询定额

        Args:
            quota_id: 定额ID

        Returns:
            Optional[ProjectQuotaDto]
        """
        quota = self.get_quota(quota_id)

        return ProjectQuotaDto.from_model(quota) if quota else None

    def get_quota_by_quota_input(
            self, project_id: str, budget_id: int, quota_input_id: int
    ):
        with db_helper.mysql_query_session(self._database) as db_session:
            quota = (
                db_session.query(ProjectQuotaModel)
                .filter(
                    and_(
                        ProjectQuotaModel.project_id == project_id,
                        ProjectQuotaModel.budget_id == budget_id,
                        ProjectQuotaModel.quota_id == quota_input_id,
                        ProjectQuotaModel.is_del == 0,
                    )
                )
                .first()
            )
            return ProjectQuotaDto.from_model(quota) if quota else None

    def create_quota(self, quota_dto: ProjectQuotaDto) -> ProjectQuotaDto:
        """创建定额

        Args:
            quota_dto: 定额DTO

        Returns:
            ProjectQuotaDto
        """
        with db_helper.mysql_session(self._database) as db_session:
            quota = ProjectQuotaModel(
                project_id=quota_dto.project_id,
                budget_id=quota_dto.budget_id,
                task_id=quota_dto.task_id,
                item_id=quota_dto.item_id,
                item_code=quota_dto.item_code,
                quota_id=quota_dto.quota_id,
                quota_code=quota_dto.quota_code,
                entry_name=quota_dto.entry_name,
                units=quota_dto.units,
                amount=quota_dto.amount,
                ex_file=quota_dto.ex_file,
                ex_cell=quota_dto.ex_cell,
                ex_row=quota_dto.ex_row,
                ex_unit=quota_dto.ex_unit,
                ex_amount=quota_dto.ex_amount,
                send_status=SendStatusEnum.NEW.value,
                send_time=quota_dto.send_time,
                send_error=quota_dto.send_error,
                created_by=quota_dto.created_by or self.current_user.username,
                created_at=datetime.now(),
            )

            db_session.add(quota)
            db_session.flush()
            return ProjectQuotaDto.from_model(quota)

    def update_quota(self, quota_dto: ProjectQuotaDto) -> Optional[ProjectQuotaDto]:
        """更新定额

        Args:
            quota_dto: 定额DTO

        Returns:
            Optional[ProjectQuotaDto]
        """
        quota = self.get_quota(quota_dto.id)

        if not quota:
            return None
        with db_helper.mysql_session(self._database) as db_session:
            quota.quota_id = quota_dto.quota_id
            quota.quota_code = quota_dto.quota_code
            quota.entry_name = quota_dto.entry_name
            quota.units = quota_dto.units
            quota.amount = quota_dto.amount
            quota.ex_file = quota_dto.ex_file
            quota.ex_cell = quota_dto.ex_cell
            quota.ex_row = quota_dto.ex_row
            quota.ex_unit = quota_dto.ex_unit
            quota.ex_amount = quota_dto.ex_amount
            quota.send_status = quota_dto.send_status
            quota.send_error = None
            quota.updated_by = self.current_user.username
            quota.updated_at = datetime.now()

            quota = db_session.merge(quota)
            return ProjectQuotaDto.from_model(quota)

    def update_quota_chapter(
            self, quota_id: int, item_id: int, item_code: str
    ) -> Optional[ProjectQuotaDto]:
        quota = self.get_quota(quota_id)
        if not quota:
            raise Exception(f"定额条目[{quota_id}]不存在")
        with db_helper.mysql_session(self._database) as db_session:
            quota.item_id = item_id
            quota.item_code = item_code
            quota.updated_by = self.current_user.username
            quota.updated_at = datetime.now()
            quota = db_session.merge(quota)
            return ProjectQuotaDto.from_model(quota)

    def delete_quota(self, quota_id: int) -> bool:
        """删除定额

        Args:
            quota_id: 定额ID

        Returns:
            bool
        """

        quota = self.get_quota(quota_id)
        if not quota:
            return False

        with db_helper.mysql_session(self._database) as db_session:
            quota.is_del = 1
            quota.deleted_by = self.current_user.username
            quota.deleted_at = datetime.now()
            db_session.merge(quota)
            return True

    def update_send_status(self, quota_id: int, status: int, err: str = None) -> bool:
        """
        更新发送状态
        Args:
            quota_id: int
            status: int
            err: str
        Returns:
            bool
        """
        quota = self.get_quota(quota_id)
        if not quota:
            return False
        with db_helper.mysql_session(self._database) as db_session:
            quota.send_status = status
            quota.send_error = err
            quota.send_time = datetime.now()
            db_session.merge(quota)
            return True

    def update_quota_code(self, quota_dto: ProjectQuotaDto):
        """更新定额编号
        Args:
            quota_dto: 定额DTO
        Returns:
            bool
        """
        quota = self.get_quota(quota_dto.id)
        if not quota:
            return False
        with db_helper.mysql_session(self._database) as db_session:
            quota.quota_code = quota_dto.quota_code
            quota.unit = quota_dto.unit
            quota.unit_price = quota_dto.unit_price
            quota.total_price = quota_dto.total_price
            quota.updated_by = self.current_user.username
            quota.updated_at = datetime.now()
            db_session.merge(quota)
            return True