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() # 新增：清空 Nav 表 from applications.models import Nav db.session.query(Nav).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() # 5. 导入公开导航（Nav 表） from applications.models import Nav count_nav = 0 for item in raw_data.get("navs", []): item = parse_time(item) if not Nav.query.get(item.get('id')) if item.get('id') else True: db.session.add(Nav(**item)) count_nav += 1 db.session.commit() click.echo(f"导入成功! 新增: 部门{count_dept}, 权限{count_power}, 角色{count_role}, 用户{count_user}, 导航{count_nav}") 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": [], "navs": [] } # 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) # 5. 导出公开导航 from applications.models import Nav data["navs"] = [model_to_dict(n) for n in Nav.query.order_by(Nav.category, Nav.sort, Nav.id).all()] 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)