Merge pull request #93 from happy-game/sqlite

支持 SQLite

.env.example

@@ -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

README.md

@@ -143,12 +143,14 @@ app/
 | 配置项                       | 说明                                                     | 默认值                             |
 | :--------------------------- | :------------------------------------------------------- | :---------------------------------------------------- |
 | **数据库配置**               |                                                          |                                                       |
-| `MYSQL_HOST`                 | 必填，MySQL 数据库主机地址                               | `localhost`                                           |
+| `DATABASE_TYPE`              | 可选，数据库类型，支持 `mysql` 或 `sqlite`               | `mysql`                                              |
-| `MYSQL_SOCKET`               | 可选，MySQL 数据库 socket 地址                             | `/var/run/mysqld/mysqld.sock`                          |
+| `SQLITE_DATABASE`            | 可选，当使用 `sqlite` 时必填，SQLite 数据库文件路径 | `default_db`                                         |
-| `MYSQL_PORT`                 | 必填，MySQL 数据库端口                                   | `3306`                                                |
+| `MYSQL_HOST`                 | 当使用 `mysql` 时必填，MySQL 数据库主机地址    | `localhost`                                          |
-| `MYSQL_USER`                 | 必填，MySQL 数据库用户名                                 | `your_db_user`                                        |
+| `MYSQL_SOCKET`               | 可选，MySQL 数据库 socket 地址                          | `/var/run/mysqld/mysqld.sock`                        |
-| `MYSQL_PASSWORD`             | 必填，MySQL 数据库密码                                   | `your_db_password`                                    |
+| `MYSQL_PORT`                 | 当使用 `mysql` 时必填，MySQL 数据库端口        | `3306`                                               |
-| `MYSQL_DATABASE`             | 必填，MySQL 数据库名称                                   | `defaultdb`                                   |
+| `MYSQL_USER`                 | 当使用 `mysql` 时必填，MySQL 数据库用户名      | `your_db_user`                                       |
+| `MYSQL_PASSWORD`             | 当使用 `mysql` 时必填，MySQL 数据库密码        | `your_db_password`                                   |
+| `MYSQL_DATABASE`             | 当使用 `mysql` 时必填，MySQL 数据库名称        | `defaultdb`                                          |
 | **API 相关配置**             |                                                          |                                                       |
 | `API_KEYS`                   | 必填，Gemini API 密钥列表，用于负载均衡                        | `["your-gemini-api-key-1", "your-gemini-api-key-2"]`  |
 | `ALLOWED_TOKENS`             | 必填，允许访问的 Token 列表                                    | `["your-access-token-1", "your-access-token-2"]`      |

app/config/config.py

@@ -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
+    DATABASE_TYPE: str = "mysql"  # sqlite 或 mysql
-    MYSQL_PORT: int
+    SQLITE_DATABASE: str = "default_db"
-    MYSQL_USER: str
+    MYSQL_HOST: str = ""
-    MYSQL_PASSWORD: str
+    MYSQL_PORT: int = 3306
-    MYSQL_DATABASE: str
+    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]

app/database/connection.py

@@ -12,10 +12,15 @@ from app.log.logger import get_database_logger
 logger = get_database_logger()
 
 # 数据库URL
-if settings.MYSQL_SOCKET:
+if settings.DATABASE_TYPE == "sqlite":
-    DATABASE_URL = f"mysql+pymysql://{settings.MYSQL_USER}:{settings.MYSQL_PASSWORD}@/{settings.MYSQL_DATABASE}?unix_socket={settings.MYSQL_SOCKET}"
+    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)}")

requirements.txt

@@ -14,6 +14,7 @@ cryptography # 支持 MySQL 8+ caching_sha2_password 验证
 pymysql
 sqlalchemy
 aiomysql
+aiosqlite
 databases
 python-dotenv
 apscheduler