140 lines
4.0 KiB
Python
140 lines
4.0 KiB
Python
"""
|
||
数据库管理脚本
|
||
用于初始化数据库、运行迁移、插入测试数据等
|
||
"""
|
||
import os
|
||
import sys
|
||
import asyncio
|
||
from sqlalchemy import create_engine, text
|
||
from app.core.config import settings
|
||
from app.database.models import Base
|
||
|
||
|
||
def create_database():
|
||
"""创建数据库（如果不存在）"""
|
||
# 连接到MySQL服务器（不指定数据库）
|
||
server_url = f"mysql+pymysql://{settings.DB_USER}:{settings.DB_PASSWORD}@{settings.DB_HOST}:{settings.DB_PORT}/"
|
||
engine = create_engine(server_url)
|
||
|
||
with engine.connect() as conn:
|
||
# 创建数据库
|
||
conn.execute(text(f"CREATE DATABASE IF NOT EXISTS `{settings.DB_NAME}` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"))
|
||
print(f"✅ 数据库 '{settings.DB_NAME}' 创建成功")
|
||
|
||
|
||
def create_tables():
|
||
"""创建数据库表"""
|
||
engine = create_engine(settings.DATABASE_URL)
|
||
Base.metadata.create_all(bind=engine)
|
||
print("✅ 数据库表创建成功")
|
||
|
||
|
||
def run_sql_file(file_path):
|
||
"""运行SQL文件"""
|
||
if not os.path.exists(file_path):
|
||
print(f"❌ SQL文件不存在: {file_path}")
|
||
return
|
||
|
||
engine = create_engine(settings.DATABASE_URL)
|
||
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
sql_content = f.read()
|
||
|
||
with engine.connect() as conn:
|
||
# 分割SQL语句并执行
|
||
statements = [stmt.strip() for stmt in sql_content.split(';') if stmt.strip()]
|
||
for statement in statements:
|
||
if statement:
|
||
conn.execute(text(statement))
|
||
conn.commit()
|
||
|
||
print(f"✅ SQL文件执行成功: {file_path}")
|
||
|
||
|
||
def init_database():
|
||
"""初始化数据库"""
|
||
print("🚀 开始初始化数据库...")
|
||
|
||
# 1. 创建数据库
|
||
create_database()
|
||
|
||
# 2. 创建表
|
||
create_tables()
|
||
|
||
# 3. 插入示例数据
|
||
sample_data_file = "sql/seeds/sample_data.sql"
|
||
if os.path.exists(sample_data_file):
|
||
run_sql_file(sample_data_file)
|
||
|
||
print("🎉 数据库初始化完成！")
|
||
|
||
|
||
def reset_database():
|
||
"""重置数据库"""
|
||
print("⚠️ 重置数据库...")
|
||
|
||
# 删除并重新创建数据库
|
||
server_url = f"mysql+pymysql://{settings.DB_USER}:{settings.DB_PASSWORD}@{settings.DB_HOST}:{settings.DB_PORT}/"
|
||
engine = create_engine(server_url)
|
||
|
||
with engine.connect() as conn:
|
||
conn.execute(text(f"DROP DATABASE IF EXISTS `{settings.DB_NAME}`"))
|
||
print(f"🗑️ 数据库 '{settings.DB_NAME}' 已删除")
|
||
|
||
# 重新初始化
|
||
init_database()
|
||
|
||
|
||
def show_help():
|
||
"""显示帮助信息"""
|
||
print("""
|
||
数据库管理脚本使用说明:
|
||
|
||
python database_manager.py init - 初始化数据库（创建数据库和表）
|
||
python database_manager.py reset - 重置数据库（删除并重新创建）
|
||
python database_manager.py create-db - 仅创建数据库
|
||
python database_manager.py create-tables - 仅创建表
|
||
python database_manager.py seed - 仅插入示例数据
|
||
python database_manager.py help - 显示此帮助信息
|
||
|
||
环境变量配置:
|
||
- DB_HOST: 数据库主机 (默认: localhost)
|
||
- DB_PORT: 数据库端口 (默认: 3306)
|
||
- DB_USER: 数据库用户名 (默认: root)
|
||
- DB_PASSWORD: 数据库密码 (默认: 空)
|
||
- DB_NAME: 数据库名称 (默认: fastapi_demo)
|
||
""")
|
||
|
||
|
||
def main():
|
||
"""主函数"""
|
||
if len(sys.argv) < 2:
|
||
show_help()
|
||
return
|
||
|
||
command = sys.argv[1].lower()
|
||
|
||
try:
|
||
if command == "init":
|
||
init_database()
|
||
elif command == "reset":
|
||
reset_database()
|
||
elif command == "create-db":
|
||
create_database()
|
||
elif command == "create-tables":
|
||
create_tables()
|
||
elif command == "seed":
|
||
run_sql_file("sql/seeds/sample_data.sql")
|
||
elif command == "help":
|
||
show_help()
|
||
else:
|
||
print(f"❌ 未知命令: {command}")
|
||
show_help()
|
||
except Exception as e:
|
||
print(f"❌ 执行失败: {str(e)}")
|
||
sys.exit(1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|