feat(data-process): 完善文件解析与切分存储链路

backend/app/modules/data_process/store.py

@@ -15,7 +15,6 @@ from psycopg.rows import dict_row
 from app.core.config import get_settings
 from app.modules.data_process.algorithms import estimate_token_count, stable_split
 
-
 TASK_STATUSES = {"pending", "running", "completed", "failed", "stopped"}
 EDITABLE_STATUSES = {"pending", "failed", "stopped", "completed"}
 
@@ -69,6 +68,34 @@ def _serialize_value(value: Any) -> Any:
     return value
 
 
+def _source_storage_descriptor(
+    payload: dict[str, Any],
+    task_id: str,
+    file_id: str,
+) -> tuple[str, dict[str, Any]]:
+    storage_object_id = str(
+        payload.get("storage_object_id")
+        or f"db://data-process/{task_id}/{file_id}/v1"
+    )
+    expected_local_prefix = f"local://data-process/{task_id}/{file_id}/v1/"
+    expected_database_reference = f"db://data-process/{task_id}/{file_id}/v1"
+    if storage_object_id.startswith(expected_local_prefix) and len(storage_object_id) > len(
+        expected_local_prefix
+    ):
+        storage_backend = "local"
+    elif storage_object_id == expected_database_reference:
+        storage_backend = "database"
+    elif storage_object_id.startswith(("local://data-process/", "db://data-process/")):
+        raise DataProcessStoreError("source storage object owner mismatch")
+    else:
+        raise DataProcessStoreError("unsupported source storage object reference")
+    metadata = {
+        **(payload.get("metadata") or {}),
+        "storage_backend": storage_backend,
+    }
+    return storage_object_id, metadata
+
+
 def _decode_row(row: dict[str, Any] | None) -> dict[str, Any] | None:
     if row is None:
         return None
@@ -343,6 +370,8 @@ class DataProcessStore:
         record_count: int,
         metadata: dict[str, Any] | None = None,
         created_by: str | None = None,
+        source_file_id: str | None = None,
+        storage_object_id: str | None = None,
     ) -> dict[str, Any]:
         return self.add_source_files(
             task_id,
@@ -356,6 +385,8 @@ class DataProcessStore:
                     "record_count": record_count,
                     "metadata": metadata or {},
                     "created_by": created_by,
+                    "id": source_file_id,
+                    "storage_object_id": storage_object_id,
                 }
             ],
         )[0]
@@ -376,12 +407,12 @@ class DataProcessStore:
                 task = self._task_in_connection(conn, task_id, for_update=True)
                 self._ensure_editable(task)
                 for payload in files:
-                    file_id = new_id("dpsf")
-                    storage_object_id = f"db://data-process/{task_id}/{file_id}/v1"
-                    metadata_payload = {
-                        "storage_backend": "database",
-                        **(payload.get("metadata") or {}),
-                    }
+                    file_id = str(payload.get("id") or new_id("dpsf"))
+                    storage_object_id, metadata_payload = _source_storage_descriptor(
+                        payload,
+                        task_id,
+                        file_id,
+                    )
                     row = conn.execute(
                         """
                         INSERT INTO data_process_source_files
@@ -530,14 +561,59 @@ class DataProcessStore:
             )
 
     def replace_preview_items(
-        self, task_id: str, items: Sequence[dict[str, Any]]
+        self,
+        task_id: str,
+        items: Sequence[dict[str, Any]],
+        *,
+        source_file_ids: Sequence[str] | None = None,
     ) -> list[dict[str, Any]]:
+        selected_ids = (
+            list(dict.fromkeys(str(file_id) for file_id in source_file_ids))
+            if source_file_ids is not None
+            else None
+        )
+        if selected_ids is not None:
+            if not selected_ids or any(not file_id for file_id in selected_ids):
+                raise ValueError("source_file_ids must contain non-empty ids")
+            selected_set = set(selected_ids)
+            unexpected = {
+                str(item.get("source_file_id") or "")
+                for item in items
+                if str(item.get("source_file_id") or "") not in selected_set
+            }
+            if unexpected:
+                raise ValueError("preview items contain an unselected source file")
+
         now = utcnow()
         with self.connect() as conn:
             task = self._task_in_connection(conn, task_id, for_update=True)
             self._ensure_editable(task)
             conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
-            conn.execute("DELETE FROM data_process_preview_items WHERE task_id=%s", (task_id,))
+            if selected_ids is None:
+                conn.execute(
+                    "DELETE FROM data_process_preview_items WHERE task_id=%s", (task_id,)
+                )
+            else:
+                rows = conn.execute(
+                    """
+                    SELECT id FROM data_process_source_files
+                    WHERE task_id=%s AND deleted_at IS NULL AND id=ANY(%s)
+                    """,
+                    (task_id, selected_ids),
+                ).fetchall()
+                found = {str(row["id"]) for row in rows}
+                missing = set(selected_ids) - found
+                if missing:
+                    raise NotFoundError(
+                        f"source files not found: {', '.join(sorted(missing))}"
+                    )
+                conn.execute(
+                    """
+                    DELETE FROM data_process_preview_items
+                    WHERE task_id=%s AND source_file_id=ANY(%s)
+                    """,
+                    (task_id, selected_ids),
+                )
             created: list[dict[str, Any]] = []
             for item in items:
                 row = conn.execute(