feat: Add Pydantic validation to import_fail_info_bulk endpoint

- Regenerated Pydantic models from updated OpenAPI specification
- Updated import_fail_info_bulk route handler to use ImportFailInfoBulkRequest/Response models
- Replaced manual JSON validation with Pydantic model validation
- Added proper error handling with ValidationError
- Updated data_models/__init__.py to export new models

Following the process outlined in data_models/README.md for type safety and consistency.
This commit is contained in:
Jin Yi 2025-07-31 14:15:21 +09:00
parent 9c4d6a0773
commit 300c6e7406
3 changed files with 53 additions and 22 deletions

View File

@ -30,6 +30,12 @@ from .generated_models import (
InstalledModelInfo,
ComfyUIVersionInfo,
# Import Fail Info Models
ImportFailInfoBulkRequest,
ImportFailInfoBulkResponse,
ImportFailInfoItem,
ImportFailInfoItem1,
# Other models
OperationType,
OperationResult,
@ -88,6 +94,12 @@ __all__ = [
"InstalledModelInfo",
"ComfyUIVersionInfo",
# Import Fail Info Models
"ImportFailInfoBulkRequest",
"ImportFailInfoBulkResponse",
"ImportFailInfoItem",
"ImportFailInfoItem1",
# Other models
"OperationType",
"OperationResult",

View File

@ -1,6 +1,6 @@
# generated by datamodel-codegen:
# filename: openapi.yaml
# timestamp: 2025-06-27T04:01:45+00:00
# timestamp: 2025-07-31T04:52:26+00:00
from __future__ import annotations
@ -454,6 +454,24 @@ class BatchExecutionRecord(BaseModel):
)
class ImportFailInfoBulkRequest(BaseModel):
cnr_ids: Optional[List[str]] = Field(
None, description="A list of CNR IDs to check."
)
urls: Optional[List[str]] = Field(
None, description="A list of repository URLs to check."
)
class ImportFailInfoItem1(BaseModel):
error: Optional[str] = None
traceback: Optional[str] = None
class ImportFailInfoItem(RootModel[Optional[ImportFailInfoItem1]]):
root: Optional[ImportFailInfoItem1]
class QueueTaskItem(BaseModel):
ui_id: str = Field(..., description="Unique identifier for the task")
client_id: str = Field(..., description="Client identifier that initiated the task")
@ -537,3 +555,7 @@ class HistoryResponse(BaseModel):
history: Optional[Dict[str, TaskHistoryItem]] = Field(
None, description="Map of task IDs to their history items"
)
class ImportFailInfoBulkResponse(RootModel[Optional[Dict[str, ImportFailInfoItem]]]):
root: Optional[Dict[str, ImportFailInfoItem]] = None

View File

@ -61,6 +61,8 @@ from ..data_models import (
ManagerMessageName,
BatchExecutionRecord,
ComfyUISystemState,
ImportFailInfoBulkRequest,
ImportFailInfoBulkResponse,
BatchOperation,
InstalledNodeInfo,
ComfyUIVersionInfo,
@ -1660,12 +1662,12 @@ async def import_fail_info(request):
async def import_fail_info_bulk(request):
try:
json_data = await request.json()
# Basic validation - ensure we have either cnr_ids or urls
if not isinstance(json_data, dict):
return web.Response(status=400, text="Request body must be a JSON object")
if "cnr_ids" not in json_data and "urls" not in json_data:
# Validate input using Pydantic model
request_data = ImportFailInfoBulkRequest.model_validate(json_data)
# Ensure we have either cnr_ids or urls
if not request_data.cnr_ids and not request_data.urls:
return web.Response(
status=400, text="Either 'cnr_ids' or 'urls' field is required"
)
@ -1675,13 +1677,8 @@ async def import_fail_info_bulk(request):
results = {}
if "cnr_ids" in json_data:
if not isinstance(json_data["cnr_ids"], list):
return web.Response(status=400, text="'cnr_ids' must be an array")
for cnr_id in json_data["cnr_ids"]:
if not isinstance(cnr_id, str):
results[cnr_id] = {"error": "cnr_id must be a string"}
continue
if request_data.cnr_ids:
for cnr_id in request_data.cnr_ids:
module_name = core.unified_manager.get_module_name(cnr_id)
if module_name is not None:
info = cm_global.error_dict.get(module_name)
@ -1692,13 +1689,8 @@ async def import_fail_info_bulk(request):
else:
results[cnr_id] = None
if "urls" in json_data:
if not isinstance(json_data["urls"], list):
return web.Response(status=400, text="'urls' must be an array")
for url in json_data["urls"]:
if not isinstance(url, str):
results[url] = {"error": "url must be a string"}
continue
if request_data.urls:
for url in request_data.urls:
module_name = core.unified_manager.get_module_name(url)
if module_name is not None:
info = cm_global.error_dict.get(module_name)
@ -1709,7 +1701,12 @@ async def import_fail_info_bulk(request):
else:
results[url] = None
return web.json_response(results)
# Create response using Pydantic model
response_data = ImportFailInfoBulkResponse(root=results)
return web.json_response(response_data.root)
except ValidationError as e:
logging.error(f"[ComfyUI-Manager] Invalid request data: {e}")
return web.Response(status=400, text=f"Invalid request data: {e}")
except Exception as e:
logging.error(f"[ComfyUI-Manager] Error processing bulk import fail info: {e}")
return web.Response(status=500, text="Internal server error")