59 lines
1.5 KiB
Python
59 lines
1.5 KiB
Python
# init_db.py
|
||
# 数据库初始化脚本。
|
||
#
|
||
# 用法：
|
||
# python init_db.py 建表（已存在的表不动）
|
||
# python init_db.py --seed 建表 + 灌演示数据
|
||
# python init_db.py --reset 删表重建 + 灌演示数据（会清空数据，慎用）
|
||
#
|
||
# 原版把这件事塞在 main.py 的 if __name__ == "__main__" 里，
|
||
# "启动服务"顺手就把表建了、数据灌了 —— 这在真实项目里是事故来源。
|
||
|
||
import sys
|
||
|
||
from database import Base, engine
|
||
from settings import settings
|
||
|
||
# 必须导入 model 包，否则 Base.metadata 里一张表都没有，create_all 会静默什么都不建
|
||
import model # noqa: F401
|
||
from seed_data import seed_all
|
||
|
||
|
||
def create_tables() -> None:
|
||
Base.metadata.create_all(bind=engine)
|
||
print(f"[ok] 建表完成：{settings.DB_NAME}（{len(Base.metadata.tables)} 张表）")
|
||
for name in sorted(Base.metadata.tables):
|
||
print(f" - {name}")
|
||
|
||
|
||
def drop_tables() -> None:
|
||
Base.metadata.drop_all(bind=engine)
|
||
print("[ok] 已删除所有表")
|
||
|
||
|
||
def reset_tables() -> None:
|
||
drop_tables()
|
||
create_tables()
|
||
|
||
|
||
def main() -> None:
|
||
args = set(sys.argv[1:])
|
||
reset = "--reset" in args
|
||
seed = "--seed" in args or reset
|
||
|
||
if reset:
|
||
print("[!] --reset：正在删表重建，历史数据将丢失")
|
||
reset_tables()
|
||
else:
|
||
create_tables()
|
||
|
||
if seed:
|
||
seed_all()
|
||
|
||
print("\n启动服务：python main.py")
|
||
print(f"接口文档：http://{settings.HOST}:{settings.PORT}/docs")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|