chenpanliang
6e5b59bde6
- 新增light路由模块，提供灯牌脚本和答案的API接口 - 添加MongoDB连接工具类，支持异步和同步操作 - 实现灯牌服务层，包含数据缓存逻辑 - 移除数据库日志功能，简化日志模块 - 更新依赖项，添加starlette、anyio等MongoDB相关包
95 lines
2.9 KiB
Python
95 lines
2.9 KiB
Python
import os
|
||
import asyncio
|
||
import logging
|
||
from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
|
||
from pymongo import MongoClient
|
||
|
||
# 配置日志
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 从环境变量获取MongoDB配置
|
||
MONGODB_URI = os.getenv("MONGODB_URI", "mongodb://localhost:27017")
|
||
MONGODB_DATABASE = os.getenv("MONGODB_DATABASE", "default_db")
|
||
|
||
class MongoDB:
|
||
"""MongoDB连接管理类"""
|
||
_client: AsyncIOMotorClient = None
|
||
_sync_client: MongoClient = None
|
||
_db: AsyncIOMotorDatabase = None
|
||
_initialized: bool = False
|
||
_sync_initialized: bool = False
|
||
|
||
@classmethod
|
||
async def initialize(cls):
|
||
"""初始化MongoDB连接（异步）"""
|
||
if cls._initialized:
|
||
return
|
||
|
||
try:
|
||
cls._client = AsyncIOMotorClient(MONGODB_URI)
|
||
cls._db = cls._client[MONGODB_DATABASE]
|
||
|
||
# 验证连接
|
||
await cls._client.server_info()
|
||
logger.info("MongoDB异步连接成功")
|
||
cls._initialized = True
|
||
except Exception as e:
|
||
logger.error(f"MongoDB异步连接失败: {str(e)}")
|
||
raise
|
||
|
||
@classmethod
|
||
def initialize_sync(cls):
|
||
"""初始化MongoDB连接（同步）"""
|
||
if cls._sync_initialized:
|
||
return
|
||
|
||
try:
|
||
cls._sync_client = MongoClient(MONGODB_URI)
|
||
# 验证连接
|
||
cls._sync_client.server_info()
|
||
logger.info("MongoDB同步连接成功")
|
||
cls._sync_initialized = True
|
||
except Exception as e:
|
||
logger.error(f"MongoDB同步连接失败: {str(e)}")
|
||
raise
|
||
|
||
@classmethod
|
||
def get_db(cls) -> AsyncIOMotorDatabase:
|
||
"""获取异步数据库实例"""
|
||
if not cls._initialized:
|
||
raise RuntimeError("MongoDB尚未初始化，请先调用initialize()")
|
||
return cls._db
|
||
|
||
@classmethod
|
||
def get_sync_db(cls):
|
||
"""获取同步数据库实例"""
|
||
if not cls._sync_initialized:
|
||
cls.initialize_sync()
|
||
return cls._sync_client[MONGODB_DATABASE]
|
||
|
||
@classmethod
|
||
async def close(cls):
|
||
"""关闭MongoDB连接"""
|
||
if cls._client:
|
||
cls._client.close()
|
||
await asyncio.sleep(0.1) # 等待连接关闭
|
||
cls._initialized = False
|
||
logger.info("MongoDB异步连接已关闭")
|
||
|
||
if cls._sync_client:
|
||
cls._sync_client.close()
|
||
cls._sync_initialized = False
|
||
logger.info("MongoDB同步连接已关闭")
|
||
|
||
# 创建全局MongoDB实例
|
||
mongodb = MongoDB()
|
||
|
||
async def get_mongodb_db() -> AsyncIOMotorDatabase:
|
||
"""获取MongoDB数据库实例（用于依赖注入）"""
|
||
if not MongoDB._initialized:
|
||
await MongoDB.initialize()
|
||
return MongoDB.get_db()
|
||
|
||
def get_mongodb_sync_db():
|
||
"""获取MongoDB同步数据库实例（用于依赖注入）"""
|
||
return MongoDB.get_sync_db() |