Merge branch 'main' of https://gitlab.matrixone.io/chentao/pyfund

api/bill.py

@@ -5,11 +5,13 @@ from fastapi import APIRouter, Depends, Query
 from motor.core import AgnosticCollection
 from pymongo.operations import UpdateOne
 from dependencies import get_current_user, get_fund_collect, get_bill_collect
+from exception.api import APIError
 from exception.db import NotFundError
 from model import Response, Page, PageResponse, SortParams, FilterTime
 from model.bill import PCFBill, ExchangeBill, AdjustBill, StakingBill
 from schema.bill import CreatePCFBill, PCFBillType, CreateExchangeBill, CreateAdjustBill, CreateStakingBill, \
     UpdatePCFBill, UpdateExchangeBill, UpdateStakingBill, UpdateAdjustBill, AllBillType
+from schema.fund import FundStatus
 from service.bill import update_bill, create_staking
 from service.fund import query_fund_assets, update_assets
 from tools.jwt_tools import User
@@ -114,12 +116,23 @@ async def create_adjust(
              summary='添加质押账目',
              description='添加质押账目')
 async def create_staking_api(
-        item: CreateStakingBill,
+        create_staking_bill: CreateStakingBill,
         user: User = Depends(get_current_user),
         bill_collect: AgnosticCollection = Depends(get_bill_collect),
         fund_collect: AgnosticCollection = Depends(get_fund_collect)
 ):
-    staking_bill = await create_staking(item, user.id, bill_collect, fund_collect)
+    staking_bill = await create_staking(create_staking_bill, user.id, bill_collect, fund_collect)
+
+    assets, adjust_assets, pending_assets, staking_assets = await query_fund_assets(fund_collect,
+                                                                                    user_id=user.id,
+                                                                                    fund_id=staking_bill.fund_id,
+                                                                                    fund_status=FundStatus.active)
+    assert assets.get(create_staking_bill.currency, 0) >= create_staking_bill.volume, APIError(message='余额不足')
+    assets[create_staking_bill.currency] -= create_staking_bill.volume
+    # 防止增加的币种没有 设置默认值
+    pending_assets.setdefault(create_staking_bill.currency, 0)
+    pending_assets[create_staking_bill.currency] += create_staking_bill.volume
+    await update_assets(fund_collect, create_staking_bill.fund_id, pending_assets=pending_assets, assets=assets)
     response = Response[StakingBill](data=staking_bill.dict())
     return response
 

exception/api.py

+from exception import MyException
+from starlette import status
+
+
+class APIError(MyException):
+    status = status.HTTP_500_INTERNAL_SERVER_ERROR
+    message = '接口错误'

exception/db.py

@@ -12,8 +12,3 @@ class NotFundError(MyException):
 class ExistDataError(MyException):
     status = status.HTTP_500_INTERNAL_SERVER_ERROR
     message = '数据已存在'
-
-
-class ErrorDataError(MyException):
-    status = status.HTTP_500_INTERNAL_SERVER_ERROR
-    message = '数据已存在'

service/bill.py

@@ -4,8 +4,9 @@ from typing import TypeVar, Type
 from motor.core import AgnosticCollection
 from pymongo import ReturnDocument
 
-from exception.db import NotFundError, ErrorDataError
-from model import Response, ErrorResponse
+from exception.api import APIError
+from exception.db import NotFundError
+from model import Response
 from model.bill import StakingBill
 from schema.bill import CreateStakingBill
 from schema.fund import FundStatus
@@ -27,7 +28,7 @@ async def create_staking(
     assert fund, NotFundError()
     assets = list(filter(lambda x: x["currency"] == create_bill.currency, fund["assets"]))
     if not assets or assets[0]["volume"] < create_bill.volume >= 0:
-        raise ErrorDataError(message="余额不足")
+        raise APIError(message="余额不足")
 
     assets[0]['volume'] -= create_bill.volume