from sqlmodel import SQLModel, create_engine, Session
from contextlib import contextmanager

sqlite_file_name = "./db/safe-algo-pro.db"
sqlite_url = f"sqlite:///{sqlite_file_name}"

connect_args = {"check_same_thread": False}
engine = create_engine(sqlite_url, connect_args=connect_args)


# 初始化数据库表
def init_db():
    SQLModel.metadata.create_all(engine)


# 数据库会话管理
def get_db():
    session = Session(engine)
    try:
        yield session
    finally:
        session.close()