import argparse import logging import sys from tenant_manager import TenantManager from model_manager import ModelManager from account_manager import AccountManager from provider_manager import ProviderManager from .config import CONFIG_PATHS # 配置日志 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) def setup_tenant_parser(subparsers): """设置租户管理相关的命令行参数""" tenant_parser = subparsers.add_parser('tenant', help='租户管理') tenant_subparsers = tenant_parser.add_subparsers(dest='tenant_command', help='租户操作') # 创建租户 create_tenant_parser = tenant_subparsers.add_parser('create', help='创建租户') create_tenant_parser.add_argument('name', help='租户名称') # 查询租户 get_tenant_parser = tenant_subparsers.add_parser('get', help='查询租户') get_tenant_parser.add_argument('name', help='租户名称') # 列出所有租户 tenant_subparsers.add_parser('list', help='列出所有租户') def setup_model_parser(subparsers): """设置模型管理相关的命令行参数""" model_parser = subparsers.add_parser('model', help='模型管理') model_subparsers = model_parser.add_subparsers(dest='model_command', help='模型操作') # 添加模型 add_model_parser = model_subparsers.add_parser('add', help='添加模型') add_model_parser.add_argument('--tenant', help='租户名称，不指定则为所有租户添加') add_model_parser.add_argument('--config', default=CONFIG_PATHS['model_config'], help='模型配置文件路径') # 添加火山模型 add_volc_model_parser = model_subparsers.add_parser('add-volc', help='添加火山模型') add_volc_model_parser.add_argument('--tenant', help='租户名称，不指定则为所有租户添加') add_volc_model_parser.add_argument('--config', default=CONFIG_PATHS['volc_model_config'], help='火山模型配置文件路径') # 删除模型 delete_model_parser = model_subparsers.add_parser('delete', help='删除模型') delete_model_parser.add_argument('--tenant', help='租户名称，不指定则删除所有租户下的指定模型') delete_model_parser.add_argument('--model', required=True, help='模型名称') # 删除租户下的所有模型 delete_all_models_parser = model_subparsers.add_parser('delete-all', help='删除租户下的所有模型') delete_all_models_parser.add_argument('--tenant', required=True, help='租户名称') def setup_account_parser(subparsers): """设置账户管理相关的命令行参数""" account_parser = subparsers.add_parser('account', help='账户管理') account_subparsers = account_parser.add_subparsers(dest='account_command', help='账户操作') # 创建账户 create_account_parser = account_subparsers.add_parser('create', help='创建账户') create_account_parser.add_argument('username', help='用户名') create_account_parser.add_argument('email', help='邮箱') create_account_parser.add_argument('password', help='密码') # 查询账户 get_account_parser = account_subparsers.add_parser('get', help='查询账户') get_account_group = get_account_parser.add_mutually_exclusive_group(required=True) get_account_group.add_argument('--username', help='用户名') get_account_group.add_argument('--email', help='邮箱') # 修改密码 update_password_parser = account_subparsers.add_parser('update-password', help='修改密码') update_password_parser.add_argument('username', help='用户名') update_password_parser.add_argument('email', help='邮箱') update_password_parser.add_argument('password', help='新密码') # 关联账户与租户 associate_parser = account_subparsers.add_parser('associate', help='关联账户与租户') associate_parser.add_argument('--username', required=True, help='用户名') associate_parser.add_argument('--tenant', required=True, help='租户名称') associate_parser.add_argument('--role', default='normal', choices=['owner', 'admin', 'editor', 'normal'], help='角色') associate_parser.add_argument('--current', action='store_true', help='是否为当前使用') # 获取租户下的所有账户 get_tenant_accounts_parser = account_subparsers.add_parser('list-tenant-accounts', help='获取租户下的所有账户') get_tenant_accounts_parser.add_argument('--tenant', required=True, help='租户名称') def setup_provider_parser(subparsers): """设置提供商管理相关的命令行参数""" provider_parser = subparsers.add_parser('provider', help='提供商管理') provider_subparsers = provider_parser.add_subparsers(dest='provider_command', help='提供商操作') # 添加提供商 add_provider_parser = provider_subparsers.add_parser('add', help='添加提供商') add_provider_parser.add_argument('--tenant', help='租户名称，不指定则为所有租户添加') add_provider_parser.add_argument('--config', default=CONFIG_PATHS['provider_config'], help='提供商配置文件路径') # 删除提供商 delete_provider_parser = provider_subparsers.add_parser('delete', help='删除提供商') delete_provider_parser.add_argument('--tenant', required=True, help='租户名称') delete_provider_parser.add_argument('--provider', required=True, help='提供商名称') # 删除租户下的所有提供商 delete_all_providers_parser = provider_subparsers.add_parser('delete-all', help='删除租户下的所有提供商') delete_all_providers_parser.add_argument('--tenant', required=True, help='租户名称') # 获取租户下的所有提供商 get_providers_parser = provider_subparsers.add_parser('list', help='获取租户下的所有提供商') get_providers_parser.add_argument('--tenant', required=True, help='租户名称') def handle_tenant_commands(args): """处理租户管理相关的命令""" if not hasattr(args, 'tenant_command') or args.tenant_command is None: print("请指定租户操作命令") return False if args.tenant_command == 'create': tenant_id = TenantManager.create_tenant(args.name) print(f"租户创建成功，ID: {tenant_id}") elif args.tenant_command == 'get': tenant = TenantManager.get_tenant_by_name(args.name) if tenant: print(f"租户信息: ID={tenant['id']}") else: print(f"未找到租户: {args.name}") elif args.tenant_command == 'list': tenants = TenantManager.get_all_tenants() print(f"共找到 {len(tenants)} 个租户:") for tenant in tenants: print(f"ID: {tenant['id']}") else: print("请指定租户操作命令") return False return True def handle_model_commands(args): """处理模型管理相关的命令""" if not hasattr(args, 'model_command') or args.model_command is None: print("请指定模型操作命令") return False if args.model_command == 'add': if args.tenant: count = ModelManager.add_models_for_tenant(args.tenant, args.config) print(f"为租户 {args.tenant} 添加模型完成，共添加 {count} 条记录。") else: count = ModelManager.add_models_for_all_tenants(args.config) print(f"为所有租户添加模型完成，共添加 {count} 条记录。") elif args.model_command == 'add-volc': if args.tenant: count = ModelManager.add_volc_models_for_tenant(args.tenant, args.config) print(f"为租户 {args.tenant} 添加火山模型完成，共添加 {count} 条记录。") else: count = ModelManager.add_volc_models_for_all_tenants(args.config) print(f"为所有租户添加火山模型完成，共添加 {count} 条记录。") elif args.model_command == 'delete': if args.tenant: tenant = TenantManager.get_tenant_by_name(args.tenant) if not tenant: print(f"未找到租户: {args.tenant}") return False count = ModelManager.delete_model_for_tenant(tenant['id'], args.model) print(f"删除租户 {args.tenant} 下的模型 {args.model} 完成，共删除 {count} 条记录。") else: count = ModelManager.delete_specific_model_for_all_tenants(args.model) print(f"删除所有租户下的模型 {args.model} 完成，共删除 {count} 条记录。") elif args.model_command == 'delete-all': tenant = TenantManager.get_tenant_by_name(args.tenant) if not tenant: print(f"未找到租户: {args.tenant}") return False count = ModelManager.delete_models_for_tenant(tenant['id']) print(f"删除租户 {args.tenant} 下的所有模型完成，共删除 {count} 条记录。") else: print("请指定模型操作命令") return False return True def handle_account_commands(args): """处理账户管理相关的命令""" if not hasattr(args, 'account_command') or args.account_command is None: print("请指定账户操作命令") return False if args.account_command == 'create': user_id = AccountManager.create_account(args.username, args.email, args.password) print(f"账户创建成功，ID: {user_id}") elif args.account_command == 'get': if args.username: user = AccountManager.get_user_by_username(args.username) if user and len(user) >= 3: print(f"用户信息: ID={user[0]}, 用户名={user[1]}, 邮箱={user[2]}") else: print(f"未找到用户: {args.username}") elif args.email: user = AccountManager.get_user_by_email(args.email) if user and len(user) >= 3: print(f"用户信息: ID={user[0]}, 用户名={user[1]}, 邮箱={user[2]}") else: print(f"未找到用户: {args.email}") elif args.account_command == 'update-password': success = AccountManager.update_password(args.username, args.email, args.password) if success: print(f"密码更新成功") else: print(f"密码更新失败，未找到用户") elif args.account_command == 'associate': user = AccountManager.get_user_by_username(args.username) if not user or len(user) < 1: print(f"未找到用户: {args.username}") return False tenant = TenantManager.get_tenant_by_name(args.tenant) if not tenant: print(f"未找到租户: {args.tenant}") return False success = AccountManager.associate_with_tenant(user[0], tenant['id'], args.role, None, args.current) if success: print(f"关联成功") else: print(f"关联失败") elif args.account_command == 'list-tenant-accounts': tenant = TenantManager.get_tenant_by_name(args.tenant) if not tenant: print(f"未找到租户: {args.tenant}") return False accounts = AccountManager.get_tenant_accounts(tenant['id']) print(f"租户 {args.tenant} 下共有 {len(accounts)} 个账户:") for account in accounts: print(f"ID: {account[0]}, 用户名: {account[1]}, 邮箱: {account[2]}, 角色: {account[3]}, 当前使用: {account[4]}") else: print("请指定账户操作命令") return False return True def handle_provider_commands(args): """处理提供商管理相关的命令""" if not hasattr(args, 'provider_command') or args.provider_command is None: print("请指定提供商操作命令") return False if args.provider_command == 'add': if args.tenant: count = ProviderManager.add_providers_for_tenant(args.tenant, args.config) print(f"为租户 {args.tenant} 添加提供商完成，共添加 {count} 条记录。") else: count = ProviderManager.add_providers_for_all_tenants(args.config) print(f"为所有租户添加提供商完成，共添加 {count} 条记录。") elif args.provider_command == 'delete': tenant = TenantManager.get_tenant_by_name(args.tenant) if not tenant: print(f"未找到租户: {args.tenant}") return False count = ProviderManager.delete_provider_for_tenant(tenant['id'], args.provider) print(f"删除租户 {args.tenant} 下的提供商 {args.provider} 完成，共删除 {count} 条记录。") elif args.provider_command == 'delete-all': tenant = TenantManager.get_tenant_by_name(args.tenant) if not tenant: print(f"未找到租户: {args.tenant}") return False count = ProviderManager.delete_providers_for_tenant(tenant['id']) print(f"删除租户 {args.tenant} 下的所有提供商完成，共删除 {count} 条记录。") elif args.provider_command == 'list': tenant = TenantManager.get_tenant_by_name(args.tenant) if not tenant: print(f"未找到租户: {args.tenant}") return False providers = ProviderManager.get_providers_for_tenant(tenant['id']) print(f"租户 {args.tenant} 下共有 {len(providers)} 个提供商:") for provider in providers: print(f"ID: {provider[0]}, 名称: {provider[1]}, 类型: {provider[2]}, 有效: {provider[3]}") else: print("请指定提供商操作命令") return False return True def main(): """主函数""" # 创建命令行解析器 parser = argparse.ArgumentParser(description='Dify数据库管理工具') subparsers = parser.add_subparsers(dest='command', help='子命令') # 设置各个子命令的解析器 setup_tenant_parser(subparsers) setup_model_parser(subparsers) setup_account_parser(subparsers) setup_provider_parser(subparsers) # 解析命令行参数 args = parser.parse_args() # 根据命令执行相应的操作 success = False if args.command == 'tenant': success = handle_tenant_commands(args) elif args.command == 'model': success = handle_model_commands(args) elif args.command == 'account': success = handle_account_commands(args) elif args.command == 'provider': success = handle_provider_commands(args) else: parser.print_help() # 根据操作结果设置退出码 sys.exit(0 if success else 1) if __name__ == "__main__": main()