342 lines
15 KiB
Python
Executable File
342 lines
15 KiB
Python
Executable File
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()
|