mirror of
https://github.com/snailyp/gemini-balance.git
synced 2026-06-06 16:19:35 +08:00
feat(database): 支持使用 SQLite 数据库
- 在 `.env.example` 文件中添加了 `DATABASE_TYPE` 变量,用于指定数据库类型,默认使用 mysql - 添加了 `DATABASE_TYPE` 和 `SQLITE_DATABASE` 配置项 - 在使用 mysql 时,对其他 MySQL 配置进行验证 - 添加 `aiosqlite` 依赖
This commit is contained in:
@@ -1,4 +1,6 @@
|
||||
# MySQL数据库配置
|
||||
# 数据库配置
|
||||
DATABASE_TYPE=mysql
|
||||
#SQLITE_DATABASE=default_db
|
||||
MYSQL_HOST=gemini-balance-mysql
|
||||
#MYSQL_SOCKET=/run/mysqld/mysqld.sock
|
||||
MYSQL_PORT=3306
|
||||
|
||||
@@ -5,7 +5,7 @@ import datetime
|
||||
import json
|
||||
from typing import List, Any, Dict, Type
|
||||
|
||||
from pydantic import ValidationError
|
||||
from pydantic import ValidationError, validator
|
||||
from pydantic_settings import BaseSettings
|
||||
from sqlalchemy import insert, update, select
|
||||
|
||||
@@ -15,13 +15,23 @@ from app.log.logger import Logger
|
||||
|
||||
class Settings(BaseSettings):
|
||||
# 数据库配置
|
||||
MYSQL_HOST: str
|
||||
MYSQL_PORT: int
|
||||
MYSQL_USER: str
|
||||
MYSQL_PASSWORD: str
|
||||
MYSQL_DATABASE: str
|
||||
DATABASE_TYPE: str = "mysql" # sqlite 或 mysql
|
||||
SQLITE_DATABASE: str = "default_db"
|
||||
MYSQL_HOST: str = ""
|
||||
MYSQL_PORT: int = 3306
|
||||
MYSQL_USER: str = ""
|
||||
MYSQL_PASSWORD: str = ""
|
||||
MYSQL_DATABASE: str = ""
|
||||
MYSQL_SOCKET: str = ""
|
||||
|
||||
# 验证 MySQL 配置
|
||||
@validator('MYSQL_HOST', 'MYSQL_PORT', 'MYSQL_USER', 'MYSQL_PASSWORD', 'MYSQL_DATABASE')
|
||||
def validate_mysql_config(cls, v, values):
|
||||
if values.get('DATABASE_TYPE') == 'mysql':
|
||||
if v is None or v == "":
|
||||
raise ValueError(f"MySQL configuration is required when DATABASE_TYPE is 'mysql'")
|
||||
return v
|
||||
|
||||
# API相关配置
|
||||
API_KEYS: List[str]
|
||||
ALLOWED_TOKENS: List[str]
|
||||
|
||||
@@ -12,10 +12,15 @@ from app.log.logger import get_database_logger
|
||||
logger = get_database_logger()
|
||||
|
||||
# 数据库URL
|
||||
if settings.MYSQL_SOCKET:
|
||||
DATABASE_URL = f"mysql+pymysql://{settings.MYSQL_USER}:{settings.MYSQL_PASSWORD}@/{settings.MYSQL_DATABASE}?unix_socket={settings.MYSQL_SOCKET}"
|
||||
if settings.DATABASE_TYPE == "sqlite":
|
||||
DATABASE_URL = f"sqlite:///{settings.SQLITE_DATABASE}"
|
||||
elif settings.DATABASE_TYPE == "mysql":
|
||||
if settings.MYSQL_SOCKET:
|
||||
DATABASE_URL = f"mysql+pymysql://{settings.MYSQL_USER}:{settings.MYSQL_PASSWORD}@/{settings.MYSQL_DATABASE}?unix_socket={settings.MYSQL_SOCKET}"
|
||||
else:
|
||||
DATABASE_URL = f"mysql+pymysql://{settings.MYSQL_USER}:{settings.MYSQL_PASSWORD}@{settings.MYSQL_HOST}:{settings.MYSQL_PORT}/{settings.MYSQL_DATABASE}"
|
||||
else:
|
||||
DATABASE_URL = f"mysql+pymysql://{settings.MYSQL_USER}:{settings.MYSQL_PASSWORD}@{settings.MYSQL_HOST}:{settings.MYSQL_PORT}/{settings.MYSQL_DATABASE}"
|
||||
raise ValueError("Unsupported database type. Please set DATABASE_TYPE to 'sqlite' or 'mysql'.")
|
||||
|
||||
# 创建数据库引擎
|
||||
# pool_pre_ping=True: 在从连接池获取连接前执行简单的 "ping" 测试,确保连接有效
|
||||
@@ -27,13 +32,16 @@ metadata = MetaData()
|
||||
# 创建基类
|
||||
Base = declarative_base(metadata=metadata)
|
||||
|
||||
# 创建数据库连接池,并配置连接池参数
|
||||
# 创建数据库连接池,并配置连接池参数,在sqlite中不使用连接池
|
||||
# min_size/max_size: 连接池的最小/最大连接数
|
||||
# pool_recycle=3600: 连接在池中允许存在的最大秒数(生命周期)。
|
||||
# 设置为 3600 秒(1小时),确保在 MySQL 默认的 wait_timeout (通常8小时) 或其他网络超时之前回收连接。
|
||||
# 如果遇到连接失效问题,可以尝试调低此值,使其小于实际的 wait_timeout 或网络超时时间。
|
||||
# databases 库会自动处理连接失效后的重连尝试。
|
||||
database = Database(DATABASE_URL, min_size=5, max_size=20, pool_recycle=1800) # Reduced recycle time to 30 mins
|
||||
if settings.DATABASE_TYPE == "sqlite":
|
||||
database = Database(DATABASE_URL)
|
||||
else:
|
||||
database = Database(DATABASE_URL, min_size=5, max_size=20, pool_recycle=1800) # Reduced recycle time to 30 mins
|
||||
|
||||
# 移除了 SessionLocal 和 get_db 函数
|
||||
|
||||
@@ -44,7 +52,7 @@ async def connect_to_db():
|
||||
"""
|
||||
try:
|
||||
await database.connect()
|
||||
logger.info("Connected to database")
|
||||
logger.info(f"Connected to {settings.DATABASE_TYPE}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to connect to database: {str(e)}")
|
||||
raise
|
||||
@@ -56,6 +64,6 @@ async def disconnect_from_db():
|
||||
"""
|
||||
try:
|
||||
await database.disconnect()
|
||||
logger.info("Disconnected from database")
|
||||
logger.info(f"Disconnected from {settings.DATABASE_TYPE}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to disconnect from database: {str(e)}")
|
||||
|
||||
@@ -14,6 +14,7 @@ cryptography # 支持 MySQL 8+ caching_sha2_password 验证
|
||||
pymysql
|
||||
sqlalchemy
|
||||
aiomysql
|
||||
aiosqlite
|
||||
databases
|
||||
python-dotenv
|
||||
apscheduler
|
||||
|
||||
Reference in New Issue
Block a user