229 lines
7.3 KiB
Python
229 lines
7.3 KiB
Python
import click
|
|
import json
|
|
import datetime
|
|
import os
|
|
from flask.cli import AppGroup
|
|
from applications.extensions import db
|
|
from applications.models import User, Role, Dept, Power
|
|
|
|
# 1. 定义命令组
|
|
admin_cli = AppGroup("admin", help="系统数据管理：包括初始化、导入和导出")
|
|
|
|
# 2. 定义数据存储的文件夹名称
|
|
DATA_DIR = "app_data"
|
|
|
|
|
|
# --- 辅助工具类 ---
|
|
class DateEncoder(json.JSONEncoder):
|
|
def default(self, obj):
|
|
if isinstance(obj, datetime.datetime):
|
|
return obj.strftime('%Y-%m-%d %H:%M:%S')
|
|
return json.JSONEncoder.default(self, obj)
|
|
|
|
|
|
def parse_time(item):
|
|
"""辅助函数：尝试将字典中的时间字符串转换为 datetime 对象"""
|
|
for k, v in item.items():
|
|
if k in ['create_at', 'create_time', 'update_at'] and isinstance(v, str):
|
|
try:
|
|
item[k] = datetime.datetime.strptime(v, '%Y-%m-%d %H:%M:%S')
|
|
except ValueError:
|
|
pass
|
|
return item
|
|
|
|
|
|
def model_to_dict(obj):
|
|
"""将 SQLAlchemy 模型对象转换为字典"""
|
|
data = {}
|
|
for c in obj.__table__.columns:
|
|
val = getattr(obj, c.name)
|
|
data[c.name] = val
|
|
return data
|
|
|
|
|
|
def ensure_dir_exists():
|
|
"""
|
|
确保数据文件夹存在。
|
|
修改：路径基于当前脚本文件的位置，而不是执行命令的位置。
|
|
"""
|
|
# 获取当前脚本(admin_cli.py)所在的目录
|
|
current_script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
# 拼接得到 app_data 的绝对路径
|
|
base_path = os.path.join(current_script_dir, DATA_DIR)
|
|
|
|
if not os.path.exists(base_path):
|
|
os.makedirs(base_path)
|
|
return base_path
|
|
|
|
|
|
# --- 核心逻辑：导入数据 ---
|
|
def import_logic(file_path, is_force):
|
|
"""具体的导入逻辑封装"""
|
|
if not os.path.exists(file_path):
|
|
click.echo(f"错误: 找不到文件 {file_path}")
|
|
return False
|
|
|
|
click.echo(f"正在读取文件: {file_path}")
|
|
with open(file_path, 'r', encoding='utf-8') as f:
|
|
raw_data = json.load(f)
|
|
|
|
if is_force:
|
|
click.echo("警告: 正在清空旧数据...")
|
|
try:
|
|
db.session.query(User).delete()
|
|
db.session.query(Role).delete()
|
|
db.session.query(Power).delete()
|
|
db.session.query(Dept).delete()
|
|
db.session.commit()
|
|
except Exception as e:
|
|
db.session.rollback()
|
|
click.echo(f"清空数据失败: {str(e)}")
|
|
return False
|
|
|
|
try:
|
|
# 1. 导入部门
|
|
count_dept = 0
|
|
for item in raw_data.get("depts", []):
|
|
item = parse_time(item)
|
|
if not Dept.query.get(item['id']):
|
|
db.session.add(Dept(**item))
|
|
count_dept += 1
|
|
db.session.commit()
|
|
|
|
# 2. 导入权限
|
|
count_power = 0
|
|
for item in raw_data.get("powers", []):
|
|
item = parse_time(item)
|
|
if not Power.query.get(item['id']):
|
|
db.session.add(Power(**item))
|
|
count_power += 1
|
|
db.session.commit()
|
|
|
|
# 3. 导入角色
|
|
count_role = 0
|
|
for item in raw_data.get("roles", []):
|
|
item = parse_time(item)
|
|
power_ids = item.pop('power_ids', [])
|
|
|
|
role = Role.query.get(item['id'])
|
|
if not role:
|
|
role = Role(**item)
|
|
db.session.add(role)
|
|
count_role += 1
|
|
|
|
if power_ids:
|
|
powers = Power.query.filter(Power.id.in_(power_ids)).all()
|
|
role.power = powers
|
|
db.session.commit()
|
|
|
|
# 4. 导入用户
|
|
count_user = 0
|
|
for item in raw_data.get("users", []):
|
|
item = parse_time(item)
|
|
role_ids = item.pop('role_ids', [])
|
|
|
|
user = User.query.get(item['id'])
|
|
if not user:
|
|
user = User(**item)
|
|
db.session.add(user)
|
|
count_user += 1
|
|
|
|
if role_ids:
|
|
roles = Role.query.filter(Role.id.in_(role_ids)).all()
|
|
user.role = roles
|
|
db.session.commit()
|
|
|
|
click.echo(f"导入成功! 新增: 部门{count_dept}, 权限{count_power}, 角色{count_role}, 用户{count_user}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
db.session.rollback()
|
|
click.echo(f"导入失败，已回滚。错误信息: {str(e)}")
|
|
return False
|
|
|
|
|
|
# --- 命令：init (系统初始化) ---
|
|
@admin_cli.command("init", help="系统初始化：读取 app_data/init.json 并导入数据库")
|
|
@click.option('--force', is_flag=True, help='是否强制清空旧数据再导入')
|
|
def init_db(force):
|
|
"""
|
|
系统初始化命令。
|
|
默认读取当前脚本同级目录下 app_data/init.json 文件。
|
|
"""
|
|
base_path = ensure_dir_exists()
|
|
target_file = os.path.join(base_path, 'init.json')
|
|
|
|
if not os.path.exists(target_file):
|
|
click.echo(f"错误：在 {base_path} 目录下未找到 init.json 文件。")
|
|
click.echo(f"请先将基础数据文件放入该位置，或运行 export 生成。")
|
|
return
|
|
|
|
import_logic(target_file, force)
|
|
|
|
|
|
# --- 命令：export (导出) ---
|
|
@admin_cli.command("export", help="导出数据：将数据库数据保存为 JSON 文件")
|
|
@click.option('--name', default=None, help='指定导出文件名，不带后缀。如果不填则使用时间戳。')
|
|
def export_data(name):
|
|
"""
|
|
导出数据到当前脚本同级目录下的 app_data 文件夹。
|
|
"""
|
|
base_path = ensure_dir_exists()
|
|
|
|
if name:
|
|
filename = f"{name}.json"
|
|
else:
|
|
now_str = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
|
|
filename = f"data_{now_str}.json"
|
|
|
|
file_path = os.path.join(base_path, filename)
|
|
click.echo(f"准备导出数据到: {file_path} ...")
|
|
|
|
data = {
|
|
"depts": [],
|
|
"powers": [],
|
|
"roles": [],
|
|
"users": []
|
|
}
|
|
|
|
# 1. 导出部门
|
|
data["depts"] = [model_to_dict(d) for d in Dept.query.order_by(Dept.id).all()]
|
|
# 2. 导出权限
|
|
data["powers"] = [model_to_dict(p) for p in Power.query.order_by(Power.id).all()]
|
|
# 3. 导出角色
|
|
roles = Role.query.order_by(Role.id).all()
|
|
for r in roles:
|
|
r_dict = model_to_dict(r)
|
|
r_dict['power_ids'] = [p.id for p in r.power]
|
|
data["roles"].append(r_dict)
|
|
# 4. 导出用户
|
|
users = User.query.order_by(User.id).all()
|
|
for u in users:
|
|
u_dict = model_to_dict(u)
|
|
u_dict['role_ids'] = [r.id for r in u.role]
|
|
data["users"].append(u_dict)
|
|
|
|
try:
|
|
with open(file_path, 'w', encoding='utf-8') as f:
|
|
json.dump(data, f, cls=DateEncoder, ensure_ascii=False, indent=4)
|
|
click.echo(f"导出成功！")
|
|
except Exception as e:
|
|
click.echo(f"导出失败: {str(e)}")
|
|
|
|
|
|
# --- 命令：import (指定文件导入) ---
|
|
@admin_cli.command("import", help="导入数据：从指定的 JSON 文件恢复数据")
|
|
@click.option('--file', required=True, help='要导入的文件名 (需包含路径或位于 app_data 下)')
|
|
@click.option('--force', is_flag=True, help='是否强制清空旧数据')
|
|
def import_data(file, force):
|
|
"""
|
|
从指定文件导入数据。
|
|
"""
|
|
# 如果用户只提供了文件名，没有路径，则假设在 app_data 目录下
|
|
if not os.path.dirname(file):
|
|
file_path = os.path.join(ensure_dir_exists(), file)
|
|
else:
|
|
file_path = file
|
|
|
|
import_logic(file_path, force) |