50 lines
1.4 KiB
Python
50 lines
1.4 KiB
Python
import os
|
||
|
||
from sqlalchemy import create_engine
|
||
from sqlalchemy.orm import sessionmaker, declarative_base, Session
|
||
|
||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||
DB_FILE = os.path.join(BASE_DIR, "avatar.db")
|
||
|
||
engine = create_engine(
|
||
f"sqlite:///{DB_FILE}",
|
||
connect_args={"check_same_thread": False},
|
||
)
|
||
SessionLocal = sessionmaker(bind=engine, autoflush=False, expire_on_commit=False)
|
||
Base = declarative_base()
|
||
|
||
|
||
def get_db():
|
||
db = SessionLocal()
|
||
try:
|
||
yield db
|
||
finally:
|
||
db.close()
|
||
|
||
|
||
def init_db():
|
||
import models
|
||
|
||
Base.metadata.create_all(bind=engine)
|
||
|
||
# 轻量迁移：为已存在的表补充新列（SQLite 不支持自动 ALTER，逐列尝试）
|
||
_try_add_columns(
|
||
("qa_pairs", "enabled", "BOOLEAN DEFAULT 1"),
|
||
("knowledge_docs", "vectorized", "BOOLEAN DEFAULT 0"),
|
||
("knowledge_docs", "embedding_model", "VARCHAR DEFAULT ''"),
|
||
("knowledge_docs", "chunk_count", "INTEGER DEFAULT 0"),
|
||
("knowledge_docs", "vectorized_at", "TIMESTAMP"),
|
||
("avatars", "owner_id", "VARCHAR DEFAULT ''"),
|
||
)
|
||
|
||
|
||
def _try_add_columns(*cols):
|
||
with engine.connect() as conn:
|
||
for table, col, ddl in cols:
|
||
try:
|
||
conn.exec_driver_sql(f"ALTER TABLE {table} ADD COLUMN {col} {ddl}")
|
||
conn.commit()
|
||
except Exception:
|
||
# 列已存在（或全新库由 create_all 建好）则忽略
|
||
pass
|