payment-platform/backend/app/plugin/module_payment/account/crud.py

from app.api.v1.module_system.auth.schema import AuthSchema
from app.core.base_crud import CRUDBase
from app.core.exceptions import CustomException

from .model import AccountModel, TransferModel, DepositModel, WithdrawModel
from .schema import (
    AccountCreateSchema,
    AccountTransferSchema,
    AccountDepositSchema,
    AccountWithdrawSchema,
)


class AccountCRUD(CRUDBase[AccountModel, AccountCreateSchema, AccountCreateSchema]):
    """资金专户 CRUD 操作"""

    def __init__(self, auth: AuthSchema) -> None:
        self.auth = auth
        super().__init__(model=AccountModel, auth=auth)

    async def get_by_enterprise_id(
        self, enterprise_id: str
    ) -> AccountModel | None:
        return await self.get(enterprise_id=enterprise_id)

    async def get_by_account_book_id(
        self, account_book_id: str
    ) -> AccountModel | None:
        return await self.get(account_book_id=account_book_id)

    async def update_by_enterprise_id(
        self, enterprise_id: str, data: dict
    ) -> AccountModel | None:
        obj = await self.get(enterprise_id=enterprise_id, preload=[])
        if not obj:
            raise CustomException(msg="更新失败！对象不存在")

        if self.auth.user and hasattr(obj, "updated_id"):
            setattr(obj, "updated_id", self.auth.user.id)

        for key, value in data.items():
            if hasattr(obj, key):
                setattr(obj, key, value)

        await self.auth.db.flush()
        await self.auth.db.refresh(obj)

        verify_obj = await self.get(enterprise_id=enterprise_id, preload=[])
        if not verify_obj:
            raise CustomException(msg="更新失败！对象不存在或无权限访问")

        return obj


class TransferCRUD(CRUDBase[TransferModel, AccountTransferSchema, AccountTransferSchema]):
    """转账记录 CRUD 操作"""

    def __init__(self, auth: AuthSchema) -> None:
        self.auth = auth
        super().__init__(model=TransferModel, auth=auth)

    async def get_by_out_biz_no(
        self, out_biz_no: str
    ) -> TransferModel | None:
        return await self.get(out_biz_no=out_biz_no)
    
    async def get_by_order_no(
        self, order_no: str
    ) -> TransferModel | None:
        return await self.get(order_no=order_no)

    async def update_by_order_no(
        self, order_no: str, data: dict
    ) -> TransferModel | None:
        obj = await self.get(order_no=order_no, preload=[])
        if not obj:
            raise CustomException(msg="转账记录不存在")

        if self.auth.user and hasattr(obj, "updated_id"):
            setattr(obj, "updated_id", self.auth.user.id)

        for key, value in data.items():
            if hasattr(obj, key):
                setattr(obj, key, value)

        await self.auth.db.flush()
        await self.auth.db.refresh(obj)

        return obj


class DepositCRUD(CRUDBase[DepositModel, AccountDepositSchema, AccountDepositSchema]):
    """充值记录 CRUD 操作"""

    def __init__(self, auth: AuthSchema) -> None:
        self.auth = auth
        super().__init__(model=DepositModel, auth=auth)

    async def get_by_out_biz_no(
        self, out_biz_no: str
    ) -> DepositModel | None:
        return await self.get(out_biz_no=out_biz_no)
    
    async def get_by_enterprise_id(
        self, enterprise_id: str
    ) -> DepositModel | None:
        return await self.get(enterprise_id=enterprise_id)

    async def update_by_out_biz_no(
        self, out_biz_no: str, data: dict
    ) -> DepositModel | None:
        obj = await self.get(out_biz_no=out_biz_no, preload=[])
        if not obj:
            raise CustomException(msg="充值记录不存在")

        if self.auth.user and hasattr(obj, "updated_id"):
            setattr(obj, "updated_id", self.auth.user.id)

        for key, value in data.items():
            if hasattr(obj, key):
                setattr(obj, key, value)

        await self.auth.db.flush()
        await self.auth.db.refresh(obj)

        return obj


class WithdrawCRUD(CRUDBase[WithdrawModel, AccountWithdrawSchema, AccountWithdrawSchema]):
    """提现记录 CRUD 操作"""

    def __init__(self, auth: AuthSchema) -> None:
        self.auth = auth
        super().__init__(model=WithdrawModel, auth=auth)

    async def get_by_out_biz_no(
        self, out_biz_no: str
    ) -> WithdrawModel | None:
        return await self.get(out_biz_no=out_biz_no)

    async def update_by_out_biz_no(
        self, out_biz_no: str, data: dict
    ) -> WithdrawModel | None:
        obj = await self.get(out_biz_no=out_biz_no, preload=[])
        if not obj:
            raise CustomException(msg="提现记录不存在")

        if self.auth.user and hasattr(obj, "updated_id"):
            setattr(obj, "updated_id", self.auth.user.id)

        for key, value in data.items():
            if hasattr(obj, key):
                setattr(obj, key, value)

        await self.auth.db.flush()
        await self.auth.db.refresh(obj)

        return obj