优化数据库初始化逻辑,确保数据库连接有效并检查表及字段的存在性

This commit is contained in:
墨梓柒
2025-05-16 17:34:43 +08:00
parent b698d17a76
commit e4f7c1fe62
2 changed files with 35 additions and 13 deletions

View File

@@ -383,17 +383,22 @@ class EmojiManager:
def initialize(self): def initialize(self):
"""初始化数据库连接和表情目录""" """初始化数据库连接和表情目录"""
if not self._initialized:
try:
# Ensure Peewee database connection is up and tables are created
if not peewee_db.is_closed():
peewee_db.connect(reuse_if_open=True) peewee_db.connect(reuse_if_open=True)
Emoji.create_table(safe=True) # Ensures table exists if peewee_db.is_closed():
raise RuntimeError("数据库连接失败")
_ensure_emoji_dir() _ensure_emoji_dir()
self._initialized = True Emoji.create_table(safe=True) # Ensures table exists
except Exception as e: # if not self._initialized:
logger.exception(f"初始化表情管理器失败: {e}") # try:
# # Ensure Peewee database connection is up and tables are created
# self._initialized = True
# except Exception as e:
# logger.exception(f"初始化表情管理器失败: {e}")
def _ensure_db(self): def _ensure_db(self):
"""确保数据库已初始化""" """确保数据库已初始化"""

View File

@@ -300,7 +300,10 @@ def create_tables():
def initialize_database(): def initialize_database():
""" """
检查所有定义的表是否存在,如果不存在则创建它们。 检查所有定义的表是否存在,如果不存在则创建它们。
检查所有表的所有字段是否存在,如果缺失则警告用户并退出程序。
""" """
import sys
models = [ models = [
ChatStreams, ChatStreams,
LLMUsage, LLMUsage,
@@ -319,12 +322,26 @@ def initialize_database():
try: try:
with db: # 管理 table_exists 检查的连接 with db: # 管理 table_exists 检查的连接
for model in models: for model in models:
table_name = model._meta.table_name
if not db.table_exists(model): if not db.table_exists(model):
logger.warning(f"'{model._meta.table_name}' 未找到。") logger.warning(f"'{table_name}' 未找到。")
needs_creation = True needs_creation = True
break # 一个表丢失,无需进一步检查。 break # 一个表丢失,无需进一步检查。
if not needs_creation:
# 检查字段
for model in models:
table_name = model._meta.table_name
cursor = db.execute_sql(f"PRAGMA table_info('{table_name}')")
existing_columns = {row[1] for row in cursor.fetchall()}
model_fields = model._meta.fields
for field_name in model_fields:
if field_name not in existing_columns:
logger.error(
f"'{table_name}' 缺失字段 '{field_name}',请手动迁移数据库结构后重启程序。"
)
sys.exit(1)
except Exception as e: except Exception as e:
logger.exception(f"检查表是否存在时出错: {e}") logger.exception(f"检查表或字段是否存在时出错: {e}")
# 如果检查失败(例如数据库不可用),则退出 # 如果检查失败(例如数据库不可用),则退出
return return
@@ -336,7 +353,7 @@ def initialize_database():
except Exception as e: except Exception as e:
logger.exception(f"创建表期间出错: {e}") logger.exception(f"创建表期间出错: {e}")
else: else:
logger.info("所有数据库表均已存在。") logger.info("所有数据库表及字段均已存在。")
# 模块加载时调用初始化函数 # 模块加载时调用初始化函数