ZhAgent/apps/system_manage/migrations/0003_alter_workspaceuserresourcepermission_target.py

# Generated by Django 5.2.6 on 2025-10-11 02:54
from concurrent.futures import ThreadPoolExecutor
from functools import reduce

from django.db import migrations, models
from django.db.models import QuerySet

from common.constants.permission_constants import WorkspaceUserRoleMapping
from common.utils.common import group_by


def workspace_user_role_mapping_model_exists(workspace_user_role_mapping_model):
    try:
        QuerySet(workspace_user_role_mapping_model).first()
    except Exception as e:
        return False
    return False

def delete_auth(apps,folder_model):
    workspace_user_resource_permission_model = apps.get_model('system_manage', 'WorkspaceUserResourcePermission')
    QuerySet(workspace_user_resource_permission_model).filter(target__in=QuerySet(folder_model).values_list('id')).delete()


def get_workspace_user_resource_permission_list(apps, auth_target_type, workspace_user_role_mapping_model_workspace_dict,
                                                folder_model):
    workspace_user_resource_permission_model = apps.get_model('system_manage', 'WorkspaceUserResourcePermission')
    return reduce(lambda x, y: [*x, *y], [
        [workspace_user_resource_permission_model(target=f.id, workspace_id=f.workspace_id, user_id=wurm.user_id,
                                         auth_target_type=auth_target_type, auth_type="RESOURCE_PERMISSION_GROUP",
                                         permission_list=['VIEW','MANAGE'] if wurm.user_id == f.user_id else ['VIEW']) for wurm in
         workspace_user_role_mapping_model_workspace_dict.get(f.workspace_id, [])] for f in
        QuerySet(folder_model).all()], [])


def auth_folder(apps, schema_editor):
    from common.database_model_manage.database_model_manage import DatabaseModelManage
    DatabaseModelManage.init()

    user_model = apps.get_model('users', 'User')
    application_folder_model = apps.get_model('application', 'ApplicationFolder')
    knowledge_folder_model = apps.get_model('knowledge', 'KnowledgeFolder')
    tool_folder_model = apps.get_model('tools', 'ToolFolder')
    workspace_user_resource_permission_model = apps.get_model('system_manage', 'WorkspaceUserResourcePermission')

    workspace_user_role_mapping_model = DatabaseModelManage.get_model("workspace_user_role_mapping")

    with ThreadPoolExecutor(max_workers=3) as executor:
        future = executor.submit(workspace_user_role_mapping_model_exists, workspace_user_role_mapping_model)
        exists = future.result()
        if not exists:
            workspace_user_role_mapping_model = None

    if workspace_user_role_mapping_model is None:
        workspace_user_role_mapping_model_workspace_dict = {
            'default': [WorkspaceUserRoleMapping('default', '', u.id) for u in QuerySet(user_model).all()]}
    else:
        workspace_user_role_mapping_model_workspace_dict = group_by(
            [v for v in {str(wurm.user_id) + str(wurm.workspace_id): wurm for wurm in
                         QuerySet(workspace_user_role_mapping_model)}.values()],
            lambda item: item.workspace_id)

    workspace_user_resource_permission_list = get_workspace_user_resource_permission_list(apps,"APPLICATION",
                                                                                          workspace_user_role_mapping_model_workspace_dict,
                                                                                          application_folder_model)

    workspace_user_resource_permission_list += get_workspace_user_resource_permission_list(apps,"TOOL",
                                                                                           workspace_user_role_mapping_model_workspace_dict,
                                                                                           tool_folder_model)

    workspace_user_resource_permission_list += get_workspace_user_resource_permission_list(apps,"KNOWLEDGE",
                                                                                           workspace_user_role_mapping_model_workspace_dict,
                                                                                           knowledge_folder_model)
    delete_auth(apps,application_folder_model)
    delete_auth(apps,knowledge_folder_model)
    delete_auth(apps,tool_folder_model)
    QuerySet(workspace_user_resource_permission_model).bulk_create(workspace_user_resource_permission_list)


class Migration(migrations.Migration):
    dependencies = [
        ('system_manage', '0002_refresh_collation_reindex'),
        ('tools', '0001_initial'),
        ('application', '0001_initial'),
        ('knowledge', '0001_initial'),
    ]

    operations = [
        migrations.AlterField(
            model_name='workspaceuserresourcepermission',
            name='target',
            field=models.CharField(db_index=True, max_length=128, verbose_name='知识库/应用id'),
        ),
        migrations.RunPython(auth_folder, atomic=False)
    ]