Feature: GET /datasets to support ids array (#17606)

api/apps/restful_apis/dataset_api.py

@@ -321,6 +321,14 @@ def list_datasets(tenant_id):
         type: string
         required: false
         description: Dataset ID to filter.
+      - in: query
+        name: ids
+        type: array
+        required: false
+        items:
+          type: string
+        collectionFormat: multi
+        description: Dataset IDs to filter.
       - in: query
         name: name
         type: string

api/apps/services/dataset_api_service.py

@@ -403,6 +403,7 @@ def list_datasets(tenant_id: str, args: dict):
     :return: (success, result) or (success, error_message)
     """
     kb_id = args.get("id")
+    kb_ids = args.get("ids")
     name = args.get("name")
     page = int(args.get("page", 1))
     page_size = int(args.get("page_size", 30))
@@ -419,10 +420,13 @@ def list_datasets(tenant_id: str, args: dict):
         # unknown type, default to True
         desc = True
 
+    if kb_id and kb_ids:
+        return False, f"Should not provide both 'id':{kb_id} and 'ids'{kb_ids}"
     if kb_id:
         kbs = KnowledgebaseService.get_kb_by_id(kb_id, tenant_id)
         if not kbs:
             return False, f"User '{tenant_id}' lacks permission for dataset '{kb_id}'"
+
     if name:
         kbs = KnowledgebaseService.get_kb_by_name(name, tenant_id)
         if not kbs:
@@ -438,7 +442,12 @@ def list_datasets(tenant_id: str, args: dict):
         tenants = TenantService.get_joined_tenants_by_user_id(tenant_id)
         tenant_ids = [m["tenant_id"] for m in tenants]
         query_user_id = tenant_id
-    kbs, total = KnowledgebaseService.get_list(tenant_ids, query_user_id, page, page_size, orderby, desc, kb_id, name, keywords, parser_id)
+    if kb_ids:
+        accessible_ids = KnowledgebaseService.get_accessible_ids([m["tenant_id"] for m in tenants], tenant_id, kb_ids)
+        if len(accessible_ids) != len(kb_ids):
+            denied_ids = [kb_id for kb_id in kb_ids if kb_id not in accessible_ids]
+            return False, f"""User '{tenant_id}' lacks permission for datasets: '{", ".join(denied_ids)}'"""
+    kbs, total = KnowledgebaseService.get_list(tenant_ids, query_user_id, page, page_size, orderby, desc, kb_id, name, keywords, parser_id, kb_ids)
     users = UserService.get_by_ids([m["tenant_id"] for m in kbs])
     user_map = {m.id: m.to_dict() for m in users}
 

api/db/services/knowledgebase_service.py

@@ -463,7 +463,7 @@ class KnowledgebaseService(CommonService):
 
     @classmethod
     @DB.connection_context()
-    def get_list(cls, joined_tenant_ids, user_id, page_number, items_per_page, orderby, desc, id, name, keywords, parser_id=None):
+    def get_list(cls, joined_tenant_ids, user_id, page_number, items_per_page, orderby, desc, id, name, keywords, parser_id=None, ids=None):
         # Get list of knowledge bases with filtering and pagination
         # Args:
         #     joined_tenant_ids: List of tenant IDs
@@ -482,6 +482,8 @@ class KnowledgebaseService(CommonService):
         kbs = cls.model.select()
         if id:
             kbs = kbs.where(cls.model.id == id)
+        if ids:
+            kbs = kbs.where(cls.model.id.in_(ids))
         if name:
             kbs = kbs.where(cls.model.name == name)
         if keywords:
@@ -501,6 +503,12 @@ class KnowledgebaseService(CommonService):
 
         return list(kbs.dicts()), total
 
+    @classmethod
+    @DB.connection_context()
+    def get_accessible_ids(cls, joined_tenant_ids, user_id, ids):
+        kbs = cls.model.select(cls.model.id).where(cls.model.id.in_(ids), cls._visibility_and_status_filter(joined_tenant_ids, user_id))
+        return {kb.id for kb in kbs}
+
     @classmethod
     @DB.connection_context()
     def get_owner_filter(cls, joined_tenant_ids, user_id):

api/utils/validation_utils.py

@@ -19,7 +19,8 @@ import pathlib
 import re
 from collections import Counter
 import string
-from typing import Annotated, Any, Literal
+from types import UnionType
+from typing import Annotated, Any, Literal, Union, get_args, get_origin
 from uuid import UUID
 
 from quart import Request
@@ -33,6 +34,15 @@ from api.utils.pagination_utils import validate_rest_api_page_size
 from common.constants import RetCode
 
 
+def _is_list_annotation(annotation: Any) -> bool:
+    origin = get_origin(annotation)
+    if origin is list:
+        return True
+    if origin in (Union, UnionType):
+        return any(_is_list_annotation(arg) for arg in get_args(annotation))
+    return False
+
+
 async def validate_and_parse_json_request(
     request: Request, validator: type[BaseModel], *, extras: dict[str, Any] | None = None, exclude_unset: bool = False
 ) -> tuple[dict[str, Any] | None, str | None]:
@@ -161,6 +171,10 @@ def validate_and_parse_request_args(request: Request, validator: type[BaseModel]
         - Preserves type conversion from Pydantic validation
     """
     args = request.args.to_dict(flat=True)
+    for field_name, field_info in validator.model_fields.items():
+        query_name = field_info.alias or field_name
+        if query_name in request.args and _is_list_annotation(field_info.annotation):
+            args[query_name] = [value for item in request.args.getlist(query_name) for value in item.split(",") if value]
 
     # Handle ext parameter: parse JSON string to dict if it's a string
     if "ext" in args and isinstance(args["ext"], str):
@@ -1002,9 +1016,26 @@ class BaseListReq(BaseModel):
 class ListDatasetReq(BaseListReq):
     """Request model for listing datasets."""
 
+    ids: Annotated[list[str] | None, Field(default=None)]
     include_parsing_status: Annotated[bool, Field(default=False)]
     ext: Annotated[dict, Field(default={})]
 
+    @field_validator("ids", mode="after")
+    @classmethod
+    def validate_ids(cls, v_list: list[str] | None) -> list[str] | None:
+        if v_list is None:
+            return None
+
+        ids_list = []
+        for v in v_list:
+            ids_list.append(validate_uuid1_hex(v))
+
+        duplicates = [item for item, count in Counter(ids_list).items() if count > 1]
+        if duplicates:
+            raise PydanticCustomError("duplicate_uuids", "Duplicate ids: '{duplicate_ids}'", {"duplicate_ids": ", ".join(duplicates)})
+
+        return ids_list
+
 
 # ---- File Management Request Models ----