pear-admin-flask/applications/common/script/admin.py

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)