77 lines
2.1 KiB
Python
77 lines
2.1 KiB
Python
# -*- coding = utf-8 -*-
|
|
# @Time : 2025/2/13 下午12:19
|
|
# @File : engine.py
|
|
# @Software : PyCharm
|
|
# @Author : xingc
|
|
# @Desc :
|
|
import functools
|
|
import logging
|
|
from typing import AsyncGenerator, Callable
|
|
|
|
import redis.asyncio as aioredis
|
|
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from config import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
engine = create_async_engine(
|
|
(
|
|
settings.DB_URL or
|
|
'{scheme}://{user}:{password}@{host}:{port}/{db}'.format(
|
|
scheme=settings.DB_SCHEMA,
|
|
user=settings.DB_USER,
|
|
password=settings.DB_PASSWORD,
|
|
host=settings.DB_HOST,
|
|
port=settings.DB_PORT,
|
|
db=settings.DB_NAME,
|
|
)
|
|
),
|
|
pool_size=settings.DB_POOL_SIZE,
|
|
pool_recycle=1800,
|
|
pool_timeout=30,
|
|
connect_args={'charset': 'utf8mb4', "connect_timeout": 5},
|
|
echo=settings.DEBUG
|
|
)
|
|
session_maker = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
|
|
async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
|
|
async with session_maker() as session:
|
|
# async with session.begin():
|
|
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,
|
|
decode_responses=True
|
|
)
|
|
|
|
|
|
redis_pool = create_redis_pool()
|
|
|
|
|
|
async def get_cache() -> aioredis.Redis:
|
|
return aioredis.Redis(connection_pool=redis_pool)
|