92 lines
2.3 KiB
Python
92 lines
2.3 KiB
Python
"""
|
|
数据库模块
|
|
|
|
提供 SQLite 数据库连接和会话管理功能。
|
|
"""
|
|
|
|
import os
|
|
import sqlite3
|
|
from pathlib import Path
|
|
from contextlib import contextmanager
|
|
from typing import Optional
|
|
|
|
# 默认数据库路径
|
|
DEFAULT_DB_PATH = Path(__file__).parent.parent / "data" / "chat.db"
|
|
|
|
|
|
def init_db(db_path: Optional[str] = None):
|
|
"""
|
|
初始化数据库
|
|
创建必要的表结构
|
|
"""
|
|
if db_path is None:
|
|
db_path = os.getenv("DB_PATH", str(DEFAULT_DB_PATH))
|
|
|
|
# 确保数据目录存在
|
|
Path(db_path).parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
conn = sqlite3.connect(db_path)
|
|
cursor = conn.cursor()
|
|
|
|
# 创建会话表
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS conversations (
|
|
id TEXT PRIMARY KEY,
|
|
title TEXT,
|
|
model TEXT,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
""")
|
|
|
|
# 创建消息表
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS messages (
|
|
id TEXT PRIMARY KEY,
|
|
conversation_id TEXT,
|
|
role TEXT,
|
|
content TEXT,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE
|
|
)
|
|
""")
|
|
|
|
# 创建文件表
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS files (
|
|
id TEXT PRIMARY KEY,
|
|
conversation_id TEXT,
|
|
filename TEXT,
|
|
file_path TEXT,
|
|
file_type TEXT,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE
|
|
)
|
|
""")
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
print(f"[数据库] 初始化完成: {db_path}")
|
|
|
|
|
|
@contextmanager
|
|
def get_db(db_path: Optional[str] = None):
|
|
"""
|
|
获取数据库连接的上下文管理器
|
|
|
|
用法:
|
|
with get_db() as db:
|
|
cursor = db.execute("SELECT * FROM conversations")
|
|
rows = cursor.fetchall()
|
|
"""
|
|
if db_path is None:
|
|
db_path = os.getenv("DB_PATH", str(DEFAULT_DB_PATH))
|
|
|
|
conn = sqlite3.connect(db_path)
|
|
conn.row_factory = sqlite3.Row
|
|
try:
|
|
yield conn
|
|
finally:
|
|
conn.close()
|