feat: knowledge admin role (#5965)

Co-authored-by: JzoNg <jzongcode@gmail.com>
Co-authored-by: jyong <718720800@qq.com>

api/configs/feature/__init__.py

@@ -391,6 +391,11 @@ class DataSetConfig(BaseModel):
         default=30,
     )
 
+    DATASET_OPERATOR_ENABLED: bool = Field(
+        description='whether to enable dataset operator',
+        default=False,
+    )
+
 
 class WorkspaceConfig(BaseModel):
     """

api/controllers/console/datasets/datasets.py

@@ -25,7 +25,7 @@
 from libs.login import login_required
 from models.dataset import Dataset, Document, DocumentSegment
 from models.model import ApiToken, UploadFile
-from services.dataset_service import DatasetService, DocumentService
+from services.dataset_service import DatasetPermissionService, DatasetService, DocumentService
 
 
 def _validate_name(name):
@@ -85,6 +85,12 @@ def get(self):
             else:
                 item['embedding_available'] = True
 
+            if item.get('permission') == 'partial_members':
+                part_users_list = DatasetPermissionService.get_dataset_partial_member_list(item['id'])
+                item.update({'partial_member_list': part_users_list})
+            else:
+                item.update({'partial_member_list': []})
+
         response = {
             'data': data,
             'has_more': len(datasets) == limit,
@@ -108,7 +114,7 @@ def post(self):
                             help='Invalid indexing technique.')
         args = parser.parse_args()
 
-        # The role of the current user in the ta table must be admin, owner, or editor
+        # The role of the current user in the ta table must be admin, owner, or editor, or dataset_operator
         if not current_user.is_editor:
             raise Forbidden()
 
@@ -140,6 +146,10 @@ def get(self, dataset_id):
         except services.errors.account.NoPermissionError as e:
             raise Forbidden(str(e))
         data = marshal(dataset, dataset_detail_fields)
+        if data.get('permission') == 'partial_members':
+            part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str)
+            data.update({'partial_member_list': part_users_list})
+
         # check embedding setting
         provider_manager = ProviderManager()
         configurations = provider_manager.get_configurations(
@@ -163,6 +173,11 @@ def get(self, dataset_id):
                 data['embedding_available'] = False
         else:
             data['embedding_available'] = True
+
+        if data.get('permission') == 'partial_members':
+            part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str)
+            data.update({'partial_member_list': part_users_list})
+
         return data, 200
 
     @setup_required
@@ -188,25 +203,39 @@ def patch(self, dataset_id):
                             nullable=True,
                             help='Invalid indexing technique.')
         parser.add_argument('permission', type=str, location='json', choices=(
-            'only_me', 'all_team_members'), help='Invalid permission.')
+            'only_me', 'all_team_members', 'partial_members'), help='Invalid permission.'
+                            )
         parser.add_argument('embedding_model', type=str,
                             location='json', help='Invalid embedding model.')
         parser.add_argument('embedding_model_provider', type=str,
                             location='json', help='Invalid embedding model provider.')
         parser.add_argument('retrieval_model', type=dict, location='json', help='Invalid retrieval model.')
+        parser.add_argument('partial_member_list', type=list, location='json', help='Invalid parent user list.')
         args = parser.parse_args()
+        data = request.get_json()
 
-        # The role of the current user in the ta table must be admin, owner, or editor
-        if not current_user.is_editor:
-            raise Forbidden()
+        # The role of the current user in the ta table must be admin, owner, editor, or dataset_operator
+        DatasetPermissionService.check_permission(
+            current_user, dataset, data.get('permission'), data.get('partial_member_list')
+        )
 
         dataset = DatasetService.update_dataset(
             dataset_id_str, args, current_user)
 
         if dataset is None:
             raise NotFound("Dataset not found.")
 
-        return marshal(dataset, dataset_detail_fields), 200
+        result_data = marshal(dataset, dataset_detail_fields)
+
+        if data.get('partial_member_list') and data.get('permission') == 'partial_members':
+            DatasetPermissionService.update_partial_member_list(dataset_id_str, data.get('partial_member_list'))
+        else:
+            DatasetPermissionService.clear_partial_member_list(dataset_id_str)
+
+        partial_member_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str)
+        result_data.update({'partial_member_list': partial_member_list})
+
+        return result_data, 200
 
     @setup_required
     @login_required
@@ -215,7 +244,7 @@ def delete(self, dataset_id):
         dataset_id_str = str(dataset_id)
 
         # The role of the current user in the ta table must be admin, owner, or editor
-        if not current_user.is_editor:
+        if not current_user.is_editor or current_user.is_dataset_operator:
             raise Forbidden()
 
         try:
@@ -569,6 +598,27 @@ def get(self, dataset_id):
         }, 200
 
 
+class DatasetPermissionUserListApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, dataset_id):
+        dataset_id_str = str(dataset_id)
+        dataset = DatasetService.get_dataset(dataset_id_str)
+        if dataset is None:
+            raise NotFound("Dataset not found.")
+        try:
+            DatasetService.check_dataset_permission(dataset, current_user)
+        except services.errors.account.NoPermissionError as e:
+            raise Forbidden(str(e))
+
+        partial_members_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str)
+
+        return {
+            'data': partial_members_list,
+        }, 200
+
+
 api.add_resource(DatasetListApi, '/datasets')
 api.add_resource(DatasetApi, '/datasets/<uuid:dataset_id>')
 api.add_resource(DatasetUseCheckApi, '/datasets/<uuid:dataset_id>/use-check')
@@ -582,3 +632,4 @@ def get(self, dataset_id):
 api.add_resource(DatasetApiBaseUrlApi, '/datasets/api-base-info')
 api.add_resource(DatasetRetrievalSettingApi, '/datasets/retrieval-setting')
 api.add_resource(DatasetRetrievalSettingMockApi, '/datasets/retrieval-setting/<string:vector_type>')
+api.add_resource(DatasetPermissionUserListApi, '/datasets/<uuid:dataset_id>/permission-part-users')

api/controllers/console/datasets/datasets_document.py

@@ -228,7 +228,7 @@ def post(self, dataset_id):
             raise NotFound('Dataset not found.')
 
         # The role of the current user in the ta table must be admin, owner, or editor
-        if not current_user.is_editor:
+        if not current_user.is_dataset_editor:
             raise Forbidden()
 
         try:
@@ -294,6 +294,11 @@ def post(self):
         parser.add_argument('retrieval_model', type=dict, required=False, nullable=False,
                             location='json')
         args = parser.parse_args()
+
+        # The role of the current user in the ta table must be admin, owner, or editor, or dataset_operator
+        if not current_user.is_dataset_editor:
+            raise Forbidden()
+
         if args['indexing_technique'] == 'high_quality':
             try:
                 model_manager = ModelManager()
@@ -757,14 +762,18 @@ def patch(self, dataset_id, document_id, action):
         dataset = DatasetService.get_dataset(dataset_id)
         if dataset is None:
             raise NotFound("Dataset not found.")
+
+        # The role of the current user in the ta table must be admin, owner, or editor
+        if not current_user.is_dataset_editor:
+            raise Forbidden()
+
         # check user's model setting
         DatasetService.check_dataset_model_setting(dataset)
 
-        document = self.get_document(dataset_id, document_id)
+        # check user's permission
+        DatasetService.check_dataset_permission(dataset, current_user)
 
-        # The role of the current user in the ta table must be admin, owner, or editor
-        if not current_user.is_editor:
-            raise Forbidden()
+        document = self.get_document(dataset_id, document_id)
 
         indexing_cache_key = 'document_{}_indexing'.format(document.id)
         cache_result = redis_client.get(indexing_cache_key)
@@ -955,10 +964,11 @@ class DocumentRenameApi(DocumentResource):
     @account_initialization_required
     @marshal_with(document_fields)
     def post(self, dataset_id, document_id):
-        # The role of the current user in the ta table must be admin or owner
-        if not current_user.is_admin_or_owner:
+        # The role of the current user in the ta table must be admin, owner, editor, or dataset_operator
+        if not current_user.is_dataset_editor:
             raise Forbidden()
-
+        dataset = DatasetService.get_dataset(dataset_id)
+        DatasetService.check_dataset_operator_permission(current_user, dataset)
         parser = reqparse.RequestParser()
         parser.add_argument('name', type=str, required=True, nullable=False, location='json')
         args = parser.parse_args()

api/controllers/console/tag/tags.py

@@ -36,7 +36,7 @@ def get(self):
     @account_initialization_required
     def post(self):
         # The role of the current user in the ta table must be admin, owner, or editor
-        if not current_user.is_editor:
+        if not (current_user.is_editor or current_user.is_dataset_editor):
             raise Forbidden()
 
         parser = reqparse.RequestParser()
@@ -68,7 +68,7 @@ class TagUpdateDeleteApi(Resource):
     def patch(self, tag_id):
         tag_id = str(tag_id)
         # The role of the current user in the ta table must be admin, owner, or editor
-        if not current_user.is_editor:
+        if not (current_user.is_editor or current_user.is_dataset_editor):
             raise Forbidden()
 
         parser = reqparse.RequestParser()
@@ -109,8 +109,8 @@ class TagBindingCreateApi(Resource):
     @login_required
     @account_initialization_required
     def post(self):
-        # The role of the current user in the ta table must be admin, owner, or editor
-        if not current_user.is_editor:
+        # The role of the current user in the ta table must be admin, owner, editor, or dataset_operator
+        if not (current_user.is_editor or current_user.is_dataset_editor):
             raise Forbidden()
 
         parser = reqparse.RequestParser()
@@ -134,8 +134,8 @@ class TagBindingDeleteApi(Resource):
     @login_required
     @account_initialization_required
     def post(self):
-        # The role of the current user in the ta table must be admin, owner, or editor
-        if not current_user.is_editor:
+        # The role of the current user in the ta table must be admin, owner, editor, or dataset_operator
+        if not (current_user.is_editor or current_user.is_dataset_editor):
             raise Forbidden()
 
         parser = reqparse.RequestParser()

api/controllers/console/workspace/members.py

@@ -131,7 +131,20 @@ def put(self, member_id):
         return {'result': 'success'}
 
 
+class DatasetOperatorMemberListApi(Resource):
+    """List all members of current tenant."""
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @marshal_with(account_with_role_list_fields)
+    def get(self):
+        members = TenantService.get_dataset_operator_members(current_user.current_tenant)
+        return {'result': 'success', 'accounts': members}, 200
+
+
 api.add_resource(MemberListApi, '/workspaces/current/members')
 api.add_resource(MemberInviteEmailApi, '/workspaces/current/members/invite-email')
 api.add_resource(MemberCancelInviteApi, '/workspaces/current/members/<uuid:member_id>')
 api.add_resource(MemberUpdateRoleApi, '/workspaces/current/members/<uuid:member_id>/update-role')
+api.add_resource(DatasetOperatorMemberListApi, '/workspaces/current/dataset-operators')

api/migrations/versions/7e6a8693e07a_add_table_dataset_permissions.py

@@ -0,0 +1,42 @@
+"""add table dataset_permissions
+
+Revision ID: 7e6a8693e07a
+Revises: 4ff534e1eb11
+Create Date: 2024-06-25 03:20:46.012193
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+import models as models
+
+# revision identifiers, used by Alembic.
+revision = '7e6a8693e07a'
+down_revision = 'b2602e131636'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('dataset_permissions',
+    sa.Column('id', models.StringUUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False),
+    sa.Column('dataset_id', models.StringUUID(), nullable=False),
+    sa.Column('account_id', models.StringUUID(), nullable=False),
+    sa.Column('has_permission', sa.Boolean(), server_default=sa.text('true'), nullable=False),
+    sa.Column('created_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP(0)'), nullable=False),
+    sa.PrimaryKeyConstraint('id', name='dataset_permission_pkey')
+    )
+    with op.batch_alter_table('dataset_permissions', schema=None) as batch_op:
+        batch_op.create_index('idx_dataset_permissions_account_id', ['account_id'], unique=False)
+        batch_op.create_index('idx_dataset_permissions_dataset_id', ['dataset_id'], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('dataset_permissions', schema=None) as batch_op:
+        batch_op.drop_index('idx_dataset_permissions_dataset_id')
+        batch_op.drop_index('idx_dataset_permissions_account_id')
+    op.drop_table('dataset_permissions')
+    # ### end Alembic commands ###

api/models/account.py

@@ -80,6 +80,10 @@ def current_tenant_id(self, value):
 
         self._current_tenant = tenant
 
+    @property
+    def current_role(self):
+        return self._current_tenant.current_role
+
     def get_status(self) -> AccountStatus:
         status_str = self.status
         return AccountStatus(status_str)
@@ -110,6 +114,14 @@ def is_admin_or_owner(self):
     def is_editor(self):
         return TenantAccountRole.is_editing_role(self._current_tenant.current_role)
 
+    @property
+    def is_dataset_editor(self):
+        return TenantAccountRole.is_dataset_edit_role(self._current_tenant.current_role)
+
+    @property
+    def is_dataset_operator(self):
+        return self._current_tenant.current_role == TenantAccountRole.DATASET_OPERATOR
+
 class TenantStatus(str, enum.Enum):
     NORMAL = 'normal'
     ARCHIVE = 'archive'
@@ -120,23 +132,30 @@ class TenantAccountRole(str, enum.Enum):
     ADMIN = 'admin'
     EDITOR = 'editor'
     NORMAL = 'normal'
+    DATASET_OPERATOR = 'dataset_operator'
 
     @staticmethod
     def is_valid_role(role: str) -> bool:
-        return role and role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR, TenantAccountRole.NORMAL}
+        return role and role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR,
+                                 TenantAccountRole.NORMAL, TenantAccountRole.DATASET_OPERATOR}
 
     @staticmethod
     def is_privileged_role(role: str) -> bool:
         return role and role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN}
 
     @staticmethod
     def is_non_owner_role(role: str) -> bool:
-        return role and role in {TenantAccountRole.ADMIN, TenantAccountRole.EDITOR, TenantAccountRole.NORMAL}
+        return role and role in {TenantAccountRole.ADMIN, TenantAccountRole.EDITOR, TenantAccountRole.NORMAL,
+                                 TenantAccountRole.DATASET_OPERATOR}
 
     @staticmethod
     def is_editing_role(role: str) -> bool:
         return role and role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR}
 
+    @staticmethod
+    def is_dataset_edit_role(role: str) -> bool:
+        return role and role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR,
+                                 TenantAccountRole.DATASET_OPERATOR}
 
 class Tenant(db.Model):
     __tablename__ = 'tenants'
@@ -172,6 +191,7 @@ class TenantAccountJoinRole(enum.Enum):
     OWNER = 'owner'
     ADMIN = 'admin'
     NORMAL = 'normal'
+    DATASET_OPERATOR = 'dataset_operator'
 
 
 class TenantAccountJoin(db.Model):