tjdify/api/controllers/console/datasets/datasets_permissions.py

from flask import jsonify, request
from flask_login import current_user
from flask_restful import Resource
from werkzeug.exceptions import Forbidden, NotFound

import services
from controllers.console import api
from controllers.console.wraps import (
    account_initialization_required,
    cloud_edition_billing_resource_check,
    setup_required,
)
from libs.login import login_required
from services.dataset_service import DatasetService


class DatasetPermissionListApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    def get(self, dataset_id):
        dataset_id = str(dataset_id)
        dataset = DatasetService.get_dataset(dataset_id)

        if not dataset:
            raise NotFound("Dataset not found.")
        try:
            DatasetService.check_dataset_permission(dataset, current_user)
        except services.errors.account.NoPermissionError as e:
            raise Forbidden(str(e))

        edit_permission_list = DatasetService.get_datasets_edit_permission(dataset_id)
        read_permission_list = DatasetService.get_datasets_read_permission(dataset_id)

        response = {
            "edit_auth": dataset.edit_auth,
            "edit_permission": edit_permission_list,
            "read_permission": read_permission_list,
        }

        return response

    # 修改知识库编辑权限
    @setup_required
    @login_required
    @account_initialization_required
    @cloud_edition_billing_resource_check("vector_space")
    def post(self, dataset_id):
        dataset_id = str(dataset_id)
        dataset = DatasetService.get_dataset(dataset_id)
        if not dataset:
            raise NotFound("Dataset not found.")
        if not current_user.is_dataset_editor:
            raise Forbidden()
        try:
            # 待完善部门领导权限判断
            DatasetService.check_dataset_permission_new(dataset, current_user)
        except services.errors.account.NoPermissionError as e:
            raise Forbidden(str(e))
        # 解析表单数据
        data = request.get_json()
        if not data:
            raise NotFound("Invalid JSON")

        edit_auth = data.get("edit_auth")
        if edit_auth is None:
            return jsonify({"error": "Missing 'edit_auth' field"}), 400

        edit_permission_list = data.get("edit_permission")
        if edit_permission_list is None:
            return jsonify({"error": "Missing 'edit_permission' field"}), 400
        if not isinstance(edit_permission_list, list):
            return jsonify({"error": "'edit_permission' should be a list"}), 400

        read_permission_list = data.get("read_permission")
        if read_permission_list is None:
            return jsonify({"error": "Missing 'read_permission' field"}), 400
        if not isinstance(read_permission_list, list):
            return jsonify({"error": "'read_permission' should be a list"}), 400
        # 更新知识库的编辑权限
        DatasetService.update_dataset_edit_auth(dataset_id, edit_auth, current_user)
        # 更新知识库的编辑、可见授权
        DatasetService.update_dataset_edit_and_read_permission(
            dataset_id, current_user, edit_permission_list, read_permission_list
        )
        return {"result": "success"}, 204


api.add_resource(DatasetPermissionListApi, "/datasets/<uuid:dataset_id>/permission")