From 32bc592eac2ca61ad6db1dff8ec5c76ae633f218 Mon Sep 17 00:00:00 2001
From: Power <26699525@qq.com>
Date: Thu, 15 Oct 2026 09:49:46 +0800
Subject: [PATCH] refactor(api): use Pydantic to validate crawl cache key

Replace raw json.loads and manual dict field inspection in
_crawl_cache_key with a validated Pydantic BaseModel.

Refs #33092
---
 .../dataset_storage_cleanup_service.py        | 28 +++++++++++++------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/api/services/dataset_storage_cleanup_service.py b/api/services/dataset_storage_cleanup_service.py
index 6043f73781e471..f06596f79fc53d 100644
--- a/api/services/dataset_storage_cleanup_service.py
+++ b/api/services/dataset_storage_cleanup_service.py
@@ -1,8 +1,9 @@
 """Remove crawl caches whose remaining document references belong to a deleted dataset."""
 
-import json
 import logging
+from typing import Literal
 
+from pydantic import BaseModel, ValidationError, field_validator
 from sqlalchemy import select
 
 from core.db.session_factory import session_factory
@@ -12,17 +13,26 @@
 logger = logging.getLogger(__name__)
 
 
+class FirecrawlDataSourceInfo(BaseModel):
+    provider: Literal["firecrawl"]
+    job_id: str
+
+    @field_validator("job_id")
+    @classmethod
+    def validate_job_id(cls, v: str) -> str:
+        if not v or any(char in v for char in "/\\\x00"):
+            raise ValueError("Invalid job_id")
+        return v
+
+
 def _crawl_cache_key(data_source_info: str | None) -> str | None:
-    try:
-        info = json.loads(data_source_info or "{}")
-    except (ValueError, TypeError):
-        return None
-    if not isinstance(info, dict) or info.get("provider") != "firecrawl":
+    if not data_source_info:
         return None
-    job_id = info.get("job_id")
-    if not isinstance(job_id, str) or not job_id or any(char in job_id for char in "/\\\x00"):
+    try:
+        info = FirecrawlDataSourceInfo.model_validate_json(data_source_info)
+    except (ValueError, TypeError, ValidationError):
         return None
-    return f"website_files/{job_id}.txt"
+    return f"website_files/{info.job_id}.txt"
 
 
 def cleanup_dataset_crawl_cache(tenant_id: str, dataset_id: str) -> None: