From 680c5c2b638816ebd56aa97a8674928a3ae84ba5 Mon Sep 17 00:00:00 2001 From: xingc Date: Mon, 12 May 2025 20:08:50 +0800 Subject: [PATCH] =?UTF-8?q?feat(task):=20=E5=90=8E=E5=8F=B0=E4=BB=BB?= =?UTF-8?q?=E5=8A=A1=E6=94=B9=E4=B8=BAapscheduler=E6=89=A7=E8=A1=8C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- core/app.py | 3 +++ core/base/scheduler.py | 22 ++++++++++++++++++++++ core/db/engine.py | 25 ++++++++++++++++++++++++- core/routers/auth/services.py | 2 +- core/routers/card/actions.py | 9 +++++++-- core/routers/card/routes.py | 9 ++++++--- core/routers/pay/actions.py | 2 ++ core/routers/pay/routes.py | 11 ++++++++--- requirements.txt | 2 +- 9 files changed, 74 insertions(+), 11 deletions(-) create mode 100644 core/base/scheduler.py diff --git a/core/app.py b/core/app.py index 9dad183..4a3a2ee 100644 --- a/core/app.py +++ b/core/app.py @@ -13,6 +13,7 @@ from fastapi_limiter import FastAPILimiter from config import settings from core.base.exceptions import CardBookException +from core.base.scheduler import scheduler from core.db.engine import get_cache from core.routers.route import routers from core.routers.chat.handlers import socket_manager @@ -44,6 +45,7 @@ def create_app() -> FastAPI: @app.on_event('startup') async def startup_event(): redis_conn = await get_cache() + scheduler.start() await FastAPILimiter.init(redis_conn) logger.info(f'Application ``{settings.NAME}`` is started') @@ -51,6 +53,7 @@ def create_app() -> FastAPI: async def shutdown_event(): logger.info(f'Application ``{settings.NAME}`` is closed') await socket_manager.pubsub_client.disconnect() + scheduler.shutdown() @app.exception_handler(HTTPException) async def http_exception_handler(request: Request, exc: HTTPException): diff --git a/core/base/scheduler.py b/core/base/scheduler.py new file mode 100644 index 0000000..4a38d0b --- /dev/null +++ b/core/base/scheduler.py @@ -0,0 +1,22 @@ +# -*- coding: utf-8 -*- +# @File: scheduler.py +# @Date: 2025/5/12 16:17 +# @Software: PyCharm +# @Author: xingc +# @Desc: +from datetime import datetime, timedelta + +from apscheduler.schedulers.asyncio import AsyncIOScheduler +from apscheduler.triggers.date import DateTrigger + +scheduler = AsyncIOScheduler() + + +def add_task(func, *args, task_id: str = None, **kwargs): + scheduler.add_job( + func, + args=args, + kwargs=kwargs, + trigger=DateTrigger(run_date=(datetime.now() + timedelta(seconds=1))), + id=task_id + ) diff --git a/core/db/engine.py b/core/db/engine.py index 22e08ff..75c056e 100644 --- a/core/db/engine.py +++ b/core/db/engine.py @@ -4,7 +4,9 @@ # @Software : PyCharm # @Author : xingc # @Desc : -from typing import AsyncGenerator +import functools +import logging +from typing import AsyncGenerator, Callable import redis.asyncio as aioredis from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker @@ -12,6 +14,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from config import settings +logger = logging.getLogger(__name__) engine = create_async_engine( ( settings.DB_URL or @@ -39,6 +42,26 @@ async def get_async_session() -> AsyncGenerator[AsyncSession, None]: yield session +def async_db_session(): + """ 异步获取db_session装饰器。 """ + + def decorator(func: Callable): + @functools.wraps(func) + async def wrapper(*args, **kwargs): + async with session_maker() as session: + try: + kwargs.setdefault('db_session', session) + return await func(*args, **kwargs) + except Exception as e: + await session.rollback() + logger.exception(e) + return None + + return wrapper + + return decorator + + def create_redis_pool(): return aioredis.ConnectionPool( host=settings.REDIS_HOST, port=settings.REDIS_PORT, password=settings.REDIS_PASSWORD, db=settings.REDIS_DB, diff --git a/core/routers/auth/services.py b/core/routers/auth/services.py index 9a7c057..e094d2b 100644 --- a/core/routers/auth/services.py +++ b/core/routers/auth/services.py @@ -87,7 +87,7 @@ async def get_current_user_form_ws( """在ws headers中提取token，并验证再获取当前用户""" auth_header = websocket.headers.get("authorization") or websocket.headers.get("Authorization") if auth_header and auth_header.lower().startswith("bearer "): - token = auth_header.removeprefix("Bearer ").strip() + token = auth_header.removeprefix("bearer ").strip() else: token = '' user = await get_current_user(token, db_session, redis_conn) diff --git a/core/routers/card/actions.py b/core/routers/card/actions.py index 0dce19d..233c2a2 100644 --- a/core/routers/card/actions.py +++ b/core/routers/card/actions.py @@ -7,8 +7,11 @@ # @Desc : import logging +from sqlalchemy.ext.asyncio import AsyncSession + from core.base.schema import ProductCategory from core.db.crud import CRUD +from core.db.engine import async_db_session from core.db.models import User, Product from core.routers.auth.schema import LoginUser from core.routers.pay.actions import actions_handler @@ -16,11 +19,13 @@ from core.routers.pay.actions import actions_handler logger = logging.getLogger(__name__) +@async_db_session() async def reward_points_after_uploading_pictures( - db_crud: CRUD, gift_product_id: int, - login_user: LoginUser + login_user: LoginUser, + db_session: AsyncSession ): + db_crud = CRUD(db_session, Product) product: Product = await db_crud.get( model=Product, filters={ 'id': gift_product_id, diff --git a/core/routers/card/routes.py b/core/routers/card/routes.py index 486d429..88e52b2 100644 --- a/core/routers/card/routes.py +++ b/core/routers/card/routes.py @@ -8,6 +8,7 @@ import logging from typing import List, Literal import ujson +from asyncer import asyncify from fastapi import APIRouter, Depends, Query, HTTPException, status import redis.asyncio as aioredis from sqlalchemy import select, func @@ -15,6 +16,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from config import settings from core.base import schema as base_schema +from core.base.scheduler import add_task from core.db.crud import CRUD from core.db.engine import get_async_session, get_cache from core.db.models import Card, SearchCardCache, UserCollection, User @@ -243,11 +245,12 @@ async def created_card( ) # 上传图片奖励 if payload.reward and settings.UPLOAD_REWARD_POINTS_GIFT_ID > 0: - await reward_points_after_uploading_pictures( - db_crud, gift_product_id=settings.UPLOAD_REWARD_POINTS_GIFT_ID, login_user=login_user + await asyncify(add_task)( + reward_points_after_uploading_pictures, + gift_product_id=settings.UPLOAD_REWARD_POINTS_GIFT_ID, login_user=login_user, + task_id=f'{login_user.user_id}-{card.card_number}' ) - else: if card is None: card: Card = await db_crud.create( diff --git a/core/routers/pay/actions.py b/core/routers/pay/actions.py index 1f219e4..aae054d 100644 --- a/core/routers/pay/actions.py +++ b/core/routers/pay/actions.py @@ -11,11 +11,13 @@ from sqlalchemy.ext.asyncio import AsyncSession from core.base.schema import ProductCategory, ProductAction, OrderStatusType from core.db.crud import CRUD +from core.db.engine import async_db_session from core.db.models import User, Order, Product, PointsTransaction logger = logging.getLogger(__name__) +@async_db_session() async def actions_after_payment( notify_params: dict, db_session: AsyncSession diff --git a/core/routers/pay/routes.py b/core/routers/pay/routes.py index 64d0489..f8889e9 100644 --- a/core/routers/pay/routes.py +++ b/core/routers/pay/routes.py @@ -8,10 +8,12 @@ import logging from xml.parsers.expat import ExpatError import xmltodict -from fastapi import APIRouter, Request, HTTPException, status, Depends, BackgroundTasks +from asyncer import asyncify +from fastapi import APIRouter, Request, HTTPException, status, Depends from sqlalchemy.ext.asyncio import AsyncSession from core.base import schema as base_schema +from core.base.scheduler import add_task from core.db.crud import CRUD from core.db.engine import get_async_session from core.db.models import Order, Product @@ -94,7 +96,6 @@ async def create_order( @pay_router.post('/wechat/notify', summary='处理支付结果通知', response_model=base_schema.BaseResponse) async def wechat_notify( request: Request, - background_tasks: BackgroundTasks, db_session: AsyncSession = Depends(get_async_session) ): body = await request.body() @@ -129,7 +130,11 @@ async def wechat_notify( 'status': base_schema.OrderStatusType.PAID }) result = {'status': 200, 'message': 'OK'} - background_tasks.add_task(actions_after_payment, notify_params=params, db_session=db_session) + await asyncify(add_task)( + actions_after_payment, + notify_params=params, + task_id=record['order_id'], + ) else: result = {'status': 400, 'message': '支付失败'} record.update({ diff --git a/requirements.txt b/requirements.txt index f7ae2de..6a17e0d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -23,7 +23,7 @@ uvicorn[standard] pycryptodome xmltodict fastapi-limiter - +asyncer