refactor: database.py 添加 MySQL 连接池管理

This commit is contained in:
2026-01-23 17:44:54 +08:00
parent 185ac8e40f
commit c44de8f52e

View File

@@ -1,40 +1,61 @@
"""SQLite 数据库操作"""
"""数据库操作 - 支持 MySQL"""
import json
import os
import sqlite3
from contextlib import contextmanager
from datetime import date, time, datetime, timedelta
from pathlib import Path
from typing import Optional
import mysql.connector
from mysql.connector import pooling
from .models import Exercise, Meal, Sleep, Weight, UserConfig, User, Reading, Invite
def get_db_path() -> Path:
"""获取数据库路径(支持环境变量配置)"""
# 优先使用环境变量
env_path = os.environ.get("VITALS_DB_PATH")
if env_path:
db_path = Path(env_path)
db_path.parent.mkdir(parents=True, exist_ok=True)
return db_path
# 数据库连接池(全局)
_connection_pool = None
# 默认路径
db_dir = Path.home() / ".vitals"
db_dir.mkdir(exist_ok=True)
return db_dir / "vitals.db"
def get_mysql_config() -> dict:
"""获取 MySQL 配置"""
return {
"host": os.environ.get("MYSQL_HOST", "localhost"),
"port": int(os.environ.get("MYSQL_PORT", "3306")),
"user": os.environ.get("MYSQL_USER", "vitals"),
"password": os.environ.get("MYSQL_PASSWORD", ""),
"database": os.environ.get("MYSQL_DATABASE", "vitals"),
}
def init_connection_pool():
"""初始化数据库连接池"""
global _connection_pool
if _connection_pool is None:
config = get_mysql_config()
_connection_pool = pooling.MySQLConnectionPool(
pool_name="vitals_pool",
pool_size=5,
pool_reset_session=True,
**config
)
return _connection_pool
@contextmanager
def get_connection():
"""获取数据库连接"""
conn = sqlite3.connect(get_db_path())
conn.row_factory = sqlite3.Row
pool = init_connection_pool()
conn = pool.get_connection()
try:
yield conn
cursor = conn.cursor(dictionary=True)
yield conn, cursor
conn.commit()
except Exception:
conn.rollback()
raise
finally:
cursor.close()
conn.close()