YG-Soft/YG_FT
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(data-process): 延迟重新生成破坏性变更

Browse Source
This commit is contained in:
caoxiaozhu
2026-07-27 10:43:42 +08:00
parent 03bd0b6d03
commit 53014bb381
3 changed files with 299 additions and 143 deletions
Download Patch File Download Diff File Expand all files Collapse all files

268
backend/app/modules/data_process/store.py
View File

@@ -41,6 +41,7 @@ _UNSTRUCTURED_PREVIEW_DEFAULTS: dict[str, Any] = {
"preserve_code_blocks": True, "preserve_code_blocks": True,
"preserve_lists": True, "preserve_lists": True,
} }
_REGENERATION_MARKER_KEY = "_regeneration_prepared"
class DataProcessStoreError(RuntimeError): class DataProcessStoreError(RuntimeError):
@@ -129,6 +130,34 @@ def _preview_config_changed(
) )
def _regeneration_marker(task: dict[str, Any]) -> dict[str, Any] | None:
config = task.get("config")
if not isinstance(config, dict):
return None
marker = config.get(_REGENERATION_MARKER_KEY)
if not isinstance(marker, dict) or marker.get("prepared") is not True:
return None
return marker
def _is_regeneration_prepared(task: dict[str, Any]) -> bool:
return _regeneration_marker(task) is not None
def _public_task(item: dict[str, Any] | None) -> dict[str, Any] | None:
"""从 API 任务快照中移除服务端内部重新生成标记。"""
if item is None:
return None
public = dict(item)
config = public.get("config")
if isinstance(config, dict) and _REGENERATION_MARKER_KEY in config:
public["config"] = {
key: value for key, value in config.items() if key != _REGENERATION_MARKER_KEY
}
return public
def _serialize_value(value: Any) -> Any: def _serialize_value(value: Any) -> Any:
if isinstance(value, (datetime, date)): if isinstance(value, (datetime, date)):
return value.isoformat().replace("+00:00", "Z") return value.isoformat().replace("+00:00", "Z")
@@ -258,7 +287,7 @@ class DataProcessStore:
[*params, page_size, (page - 1) * page_size], [*params, page_size, (page - 1) * page_size],
).fetchall() ).fetchall()
return { return {
"items": [_decode_row(row) for row in rows], "items": [_public_task(_decode_row(row)) for row in rows],
"total": int(total), "total": int(total),
"page": page, "page": page,
"page_size": page_size, "page_size": page_size,
@@ -284,7 +313,13 @@ class DataProcessStore:
payload.get("description") or "", payload.get("description") or "",
payload["process_type"], payload["process_type"],
payload.get("source_dataset_id"), payload.get("source_dataset_id"),
json_dumps(payload.get("config") or {}), json_dumps(
{
key: value
for key, value in (payload.get("config") or {}).items()
if key != _REGENERATION_MARKER_KEY
}
),
payload.get("tenant_id"), payload.get("tenant_id"),
payload.get("project_id"), payload.get("project_id"),
payload.get("owner_id"), payload.get("owner_id"),
@@ -296,7 +331,7 @@ class DataProcessStore:
).fetchone() ).fetchone()
except psycopg.errors.UniqueViolation as exc: except psycopg.errors.UniqueViolation as exc:
raise ConflictError("data process task name already exists") from exc raise ConflictError("data process task name already exists") from exc
return _decode_row(row) or {} return _public_task(_decode_row(row)) or {}
def get_task(self, task_id: str, *, for_update: bool = False) -> dict[str, Any]: def get_task(self, task_id: str, *, for_update: bool = False) -> dict[str, Any]:
lock = " FOR UPDATE" if for_update else "" lock = " FOR UPDATE" if for_update else ""
@@ -348,7 +383,7 @@ class DataProcessStore:
).fetchone() ).fetchone()
if not row: if not row:
raise NotFoundError("data process task not found") raise NotFoundError("data process task not found")
return _decode_row(row) or {} return _public_task(_decode_row(row)) or {}
def _task_in_connection( def _task_in_connection(
self, self,
@@ -370,7 +405,7 @@ class DataProcessStore:
def _ensure_editable(task: dict[str, Any]) -> None: def _ensure_editable(task: dict[str, Any]) -> None:
if task["status"] not in EDITABLE_STATUSES: if task["status"] not in EDITABLE_STATUSES:
raise InvalidStateError(f"task cannot be edited while status is {task['status']}") raise InvalidStateError(f"task cannot be edited while status is {task['status']}")
if task.get("output_dataset_id"): if task.get("output_dataset_id") and not _is_regeneration_prepared(task):
raise InvalidStateError("published task cannot be edited") raise InvalidStateError("published task cannot be edited")
def update_task(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]: def update_task(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]:
@@ -381,14 +416,30 @@ class DataProcessStore:
"source_dataset_id", "source_dataset_id",
} }
values: dict[str, Any] = {key: value for key, value in payload.items() if key in allowed} values: dict[str, Any] = {key: value for key, value in payload.items() if key in allowed}
if payload.get("config") is not None: if not values and payload.get("config") is None:
values["config"] = json_dumps(payload["config"])
if not values:
return self.get_task(task_id) return self.get_task(task_id)
try: try:
with self.connect() as conn: with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True) task = self._task_in_connection(conn, task_id, for_update=True)
self._ensure_editable(task) self._ensure_editable(task)
regeneration_prepared = _is_regeneration_prepared(task)
if regeneration_prepared and any(
key in payload and payload.get(key) != task.get(key)
for key in ("process_type", "source_dataset_id")
):
raise InvalidStateError(
"process type and source dataset cannot change during regeneration"
)
if payload.get("config") is not None:
next_config = {
key: value
for key, value in payload["config"].items()
if key != _REGENERATION_MARKER_KEY
}
current_marker = _regeneration_marker(task)
if current_marker:
next_config[_REGENERATION_MARKER_KEY] = current_marker
values["config"] = json_dumps(next_config)
invalidates_results = ( invalidates_results = (
("config" in payload and payload.get("config") != task.get("config")) ("config" in payload and payload.get("config") != task.get("config"))
or ( or (
@@ -400,7 +451,7 @@ class DataProcessStore:
and payload.get("source_dataset_id") != task.get("source_dataset_id") and payload.get("source_dataset_id") != task.get("source_dataset_id")
) )
) )
if invalidates_results: if invalidates_results and not regeneration_prepared:
values.update( values.update(
{ {
"status": "pending", "status": "pending",
@@ -440,18 +491,17 @@ class DataProcessStore:
).fetchone() ).fetchone()
except psycopg.errors.UniqueViolation as exc: except psycopg.errors.UniqueViolation as exc:
raise ConflictError("data process task name already exists") from exc raise ConflictError("data process task name already exists") from exc
return _decode_row(row) or {} return _public_task(_decode_row(row)) or {}
def prepare_regeneration( def prepare_regeneration(
self, self,
task_id: str, task_id: str,
payload: dict[str, Any], payload: dict[str, Any],
) -> dict[str, Any]: ) -> dict[str, Any]:
"""在一个事务中将已有任务恢复为可重新生成状态。 """非破坏性地保存重新生成配置。
已发布数据集是可被其他任务使用的独立产物，此处只解除当前任务的 准备阶段保留任务当前状态、结果、切片及已发布数据集。真正开始
输出指针，不删除数据集实体。用户完成新一轮生成后，可再次发布 生成时，才在同一事务内切换运行状态并清理上一轮结果。
以原子替换三个分割数据集的内容。
""" """
try: try:
@@ -470,12 +520,18 @@ class DataProcessStore:
current_config = dict(task.get("config") or {}) current_config = dict(task.get("config") or {})
next_config = dict(payload.get("config") or {}) next_config = dict(payload.get("config") or {})
next_config.pop(_REGENERATION_MARKER_KEY, None)
preview_invalidated = _preview_config_changed( preview_invalidated = _preview_config_changed(
process_type, process_type,
current_config, current_config,
next_config, next_config,
) )
now = utcnow() now = utcnow()
next_config[_REGENERATION_MARKER_KEY] = {
"prepared": True,
"preview_invalidated": preview_invalidated,
"prepared_at": now,
}
# 002 迁移前发布的数据集只有 task_id。先补齐新关联字段，保证 # 002 迁移前发布的数据集只有 task_id。先补齐新关联字段，保证
# 解除任务输出指针后，详情和后续重新发布仍能定位原来的三份数据集。 # 解除任务输出指针后，详情和后续重新发布仍能定位原来的三份数据集。
conn.execute( conn.execute(
@@ -504,35 +560,10 @@ class DataProcessStore:
published_row and published_row.get("exists") published_row and published_row.get("exists")
) )
# 删除顺序保证结果不再指向即将被替换的切片。
conn.execute(
"DELETE FROM data_process_results WHERE task_id=%s",
(task_id,),
)
if preview_invalidated:
conn.execute(
"DELETE FROM data_process_preview_items WHERE task_id=%s",
(task_id,),
)
preview_count = 0
else:
preview_row = conn.execute(
"""
SELECT COUNT(*) AS count FROM data_process_preview_items
WHERE task_id=%s
""",
(task_id,),
).fetchone()
preview_count = int((preview_row or {}).get("count") or 0)
row = conn.execute( row = conn.execute(
""" """
UPDATE data_process_tasks UPDATE data_process_tasks
SET name=%s, description=%s, config=%s, status='pending', progress=%s, SET name=%s, description=%s, config=%s, updated_at=%s
output_dataset_id=NULL, output_count=0, filtered_count=0,
duplicate_count=0, error_count=0, failure_reason=NULL,
generation_run_id=NULL, started_at=NULL, completed_at=NULL,
updated_at=%s
WHERE id=%s WHERE id=%s
RETURNING * RETURNING *
""", """,
@@ -540,7 +571,6 @@ class DataProcessStore:
payload["name"], payload["name"],
payload.get("description") or "", payload.get("description") or "",
json_dumps(next_config), json_dumps(next_config),
20 if preview_count else 0,
now, now,
task_id, task_id,
), ),
@@ -548,7 +578,7 @@ class DataProcessStore:
except psycopg.errors.UniqueViolation as exc: except psycopg.errors.UniqueViolation as exc:
raise ConflictError("data process task name already exists") from exc raise ConflictError("data process task name already exists") from exc
return { return {
"task": _decode_row(row) or {}, "task": _public_task(_decode_row(row)) or {},
"preview_invalidated": preview_invalidated, "preview_invalidated": preview_invalidated,
"published_outputs_preserved": published_outputs_preserved, "published_outputs_preserved": published_outputs_preserved,
} }
@@ -672,9 +702,6 @@ class DataProcessStore:
), ),
).fetchone() ).fetchone()
created.append(_decode_row(row) or {}) created.append(_decode_row(row) or {})
# 前端会只对本次新增的源文件构建预览，因此必须保留旧
# 文件的切片。任何源文件增加都会使旧生成结果失效。
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
preview_row = conn.execute( preview_row = conn.execute(
""" """
SELECT COUNT(*) AS count FROM data_process_preview_items SELECT COUNT(*) AS count FROM data_process_preview_items
@@ -683,20 +710,40 @@ class DataProcessStore:
(task_id,), (task_id,),
).fetchone() ).fetchone()
preview_count = int((preview_row or {}).get("count") or 0) preview_count = int((preview_row or {}).get("count") or 0)
conn.execute( if _is_regeneration_prepared(task):
""" conn.execute(
UPDATE data_process_tasks """
SET status='pending', progress=%s, output_count=0, filtered_count=0, UPDATE data_process_tasks
duplicate_count=0, error_count=0, failure_reason=NULL, SET input_count=(
generation_run_id=NULL, started_at=NULL, completed_at=NULL, input_count=( SELECT COALESCE(SUM(record_count), 0)
SELECT COALESCE(SUM(record_count), 0) FROM data_process_source_files
FROM data_process_source_files WHERE task_id=%s AND deleted_at IS NULL
WHERE task_id=%s AND deleted_at IS NULL ), updated_at=%s
), updated_at=%s WHERE id=%s
WHERE id=%s """,
""", (task_id, now, task_id),
(20 if preview_count else 0, task_id, now, task_id), )
) else:
# 前端会只对本次新增的源文件构建预览，因此保留旧文件切片，
# 但普通未发布任务的旧生成结果已经不再有效。
conn.execute(
"DELETE FROM data_process_results WHERE task_id=%s", (task_id,)
)
conn.execute(
"""
UPDATE data_process_tasks
SET status='pending', progress=%s, output_count=0, filtered_count=0,
duplicate_count=0, error_count=0, failure_reason=NULL,
generation_run_id=NULL, started_at=NULL, completed_at=NULL,
input_count=(
SELECT COALESCE(SUM(record_count), 0)
FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL
), updated_at=%s
WHERE id=%s
""",
(20 if preview_count else 0, task_id, now, task_id),
)
except psycopg.errors.UniqueViolation as exc: except psycopg.errors.UniqueViolation as exc:
raise ConflictError( raise ConflictError(
"the same source file content is already attached to this task" "the same source file content is already attached to this task"
@@ -777,23 +824,39 @@ class DataProcessStore:
).fetchone() ).fetchone()
if not row: if not row:
raise NotFoundError("source file not found") raise NotFoundError("source file not found")
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
conn.execute( conn.execute(
"DELETE FROM data_process_preview_items WHERE source_file_id=%s", (file_id,) "DELETE FROM data_process_preview_items WHERE source_file_id=%s", (file_id,)
) )
conn.execute( now = utcnow()
""" if _is_regeneration_prepared(task):
UPDATE data_process_tasks conn.execute(
SET status='pending', progress=0, output_count=0, filtered_count=0, """
duplicate_count=0, error_count=0, failure_reason=NULL, UPDATE data_process_tasks
input_count=(SELECT COALESCE(SUM(record_count), 0) SET input_count=(SELECT COALESCE(SUM(record_count), 0)
FROM data_process_source_files FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL), WHERE task_id=%s AND deleted_at IS NULL),
updated_at=%s updated_at=%s
WHERE id=%s WHERE id=%s
""", """,
(task_id, utcnow(), task_id), (task_id, now, task_id),
) )
else:
conn.execute(
"DELETE FROM data_process_results WHERE task_id=%s", (task_id,)
)
conn.execute(
"""
UPDATE data_process_tasks
SET status='pending', progress=0, output_count=0, filtered_count=0,
duplicate_count=0, error_count=0, failure_reason=NULL,
input_count=(SELECT COALESCE(SUM(record_count), 0)
FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL),
updated_at=%s
WHERE id=%s
""",
(task_id, now, task_id),
)
def replace_preview_items( def replace_preview_items(
self, self,
@@ -823,7 +886,9 @@ class DataProcessStore:
with self.connect() as conn: with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True) task = self._task_in_connection(conn, task_id, for_update=True)
self._ensure_editable(task) self._ensure_editable(task)
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,)) regeneration_prepared = _is_regeneration_prepared(task)
if not regeneration_prepared:
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
if selected_ids is None: if selected_ids is None:
conn.execute( conn.execute(
"DELETE FROM data_process_preview_items WHERE task_id=%s", (task_id,) "DELETE FROM data_process_preview_items WHERE task_id=%s", (task_id,)
@@ -878,15 +943,21 @@ class DataProcessStore:
), ),
).fetchone() ).fetchone()
created.append(_decode_row(row) or {}) created.append(_decode_row(row) or {})
conn.execute( if regeneration_prepared:
""" conn.execute(
UPDATE data_process_tasks "UPDATE data_process_tasks SET updated_at=%s WHERE id=%s",
SET status='pending', progress=20, output_count=0, filtered_count=0, (now, task_id),
duplicate_count=0, error_count=0, failure_reason=NULL, updated_at=%s )
WHERE id=%s else:
""", conn.execute(
(now, task_id), """
) UPDATE data_process_tasks
SET status='pending', progress=20, output_count=0, filtered_count=0,
duplicate_count=0, error_count=0, failure_reason=NULL, updated_at=%s
WHERE id=%s
""",
(now, task_id),
)
return created return created
def list_preview_items( def list_preview_items(
@@ -980,7 +1051,7 @@ class DataProcessStore:
now, now,
), ),
).fetchone() ).fetchone()
self._invalidate_results(conn, task_id, now) self._invalidate_results(conn, task, task_id, now)
return _decode_row(row) or {} return _decode_row(row) or {}
def update_preview_item( def update_preview_item(
@@ -1027,7 +1098,7 @@ class DataProcessStore:
task_id, task_id,
), ),
).fetchone() ).fetchone()
self._invalidate_results(conn, task_id, now) self._invalidate_results(conn, task, task_id, now)
return _decode_row(row) or {} return _decode_row(row) or {}
def delete_preview_item(self, task_id: str, preview_id: str) -> None: def delete_preview_item(self, task_id: str, preview_id: str) -> None:
@@ -1040,14 +1111,21 @@ class DataProcessStore:
).fetchone() ).fetchone()
if not row: if not row:
raise NotFoundError("preview item not found") raise NotFoundError("preview item not found")
self._invalidate_results(conn, task_id, utcnow()) self._invalidate_results(conn, task, task_id, utcnow())
def _invalidate_results( def _invalidate_results(
self, self,
conn: psycopg.Connection[dict[str, Any]], conn: psycopg.Connection[dict[str, Any]],
task: dict[str, Any],
task_id: str, task_id: str,
now: str, now: str,
) -> None: ) -> None:
if _is_regeneration_prepared(task):
conn.execute(
"UPDATE data_process_tasks SET updated_at=%s WHERE id=%s",
(now, task_id),
)
return
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,)) conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
conn.execute( conn.execute(
""" """
@@ -1065,7 +1143,8 @@ class DataProcessStore:
raise DataProcessStoreError("incremental generation is not supported") raise DataProcessStoreError("incremental generation is not supported")
with self.connect() as conn: with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True) task = self._task_in_connection(conn, task_id, for_update=True)
if task.get("output_dataset_id"): regeneration_prepared = _is_regeneration_prepared(task)
if task.get("output_dataset_id") and not regeneration_prepared:
raise InvalidStateError("published task cannot be regenerated") raise InvalidStateError("published task cannot be regenerated")
if task["status"] == "running": if task["status"] == "running":
raise ConflictError("data process task is already running") raise ConflictError("data process task is already running")
@@ -1078,16 +1157,19 @@ class DataProcessStore:
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,)) conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
now = utcnow() now = utcnow()
generation_run_id = new_id("dprun") generation_run_id = new_id("dprun")
next_config = dict(task.get("config") or {})
next_config.pop(_REGENERATION_MARKER_KEY, None)
row = conn.execute( row = conn.execute(
""" """
UPDATE data_process_tasks UPDATE data_process_tasks
SET status='running', progress=30, failure_reason=NULL, started_at=%s, SET config=%s, status='running', progress=30, failure_reason=NULL,
completed_at=NULL, output_count=0, filtered_count=0, started_at=%s, completed_at=NULL, output_dataset_id=NULL,
duplicate_count=0, error_count=0, generation_run_id=%s, updated_at=%s output_count=0, filtered_count=0, duplicate_count=0, error_count=0,
generation_run_id=%s, updated_at=%s
WHERE id=%s WHERE id=%s
RETURNING * RETURNING *
""", """,
(now, generation_run_id, now, task_id), (json_dumps(next_config), now, generation_run_id, now, task_id),
).fetchone() ).fetchone()
return _decode_row(row) or {} return _decode_row(row) or {}
@@ -1405,6 +1487,10 @@ class DataProcessStore:
"""按精确配额发布训练、验证、测试三个独立数据集。""" """按精确配额发布训练、验证、测试三个独立数据集。"""
with self.connect() as conn: with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True) task = self._task_in_connection(conn, task_id, for_update=True)
if _is_regeneration_prepared(task):
raise InvalidStateError(
"regeneration must start and complete before publishing"
)
if task["status"] != "completed": if task["status"] != "completed":
raise InvalidStateError("only a completed task can be published") raise InvalidStateError("only a completed task can be published")
rows = conn.execute( rows = conn.execute(

80
backend/tests/test_data_process_api.py
View File

@@ -30,6 +30,7 @@ class FakeDataProcessStore:
self.previews: dict[str, list[dict[str, Any]]] = {} self.previews: dict[str, list[dict[str, Any]]] = {}
self.results: dict[str, list[dict[str, Any]]] = {} self.results: dict[str, list[dict[str, Any]]] = {}
self.datasets: dict[str, dict[str, Any]] = {} self.datasets: dict[str, dict[str, Any]] = {}
self.regeneration_prepared: set[str] = set()
self.sequence = 0 self.sequence = 0
def _id(self, prefix: str) -> str: def _id(self, prefix: str) -> str:
@@ -109,27 +110,19 @@ class FakeDataProcessStore:
raise InvalidStateError("data process task was modified by another request") raise InvalidStateError("data process task was modified by another request")
if payload["process_type"] != task["process_type"]: if payload["process_type"] != task["process_type"]:
raise InvalidStateError("process_type cannot be changed during regeneration") raise InvalidStateError("process_type cannot be changed during regeneration")
published_outputs_preserved = bool(task.get("output_dataset_id")) published_outputs_preserved = bool(task.get("output_dataset_id")) or any(
self.results[task_id] = [] dataset.get("source_task_id") == task_id
for dataset in self.datasets.values()
)
task.update( task.update(
{ {
"name": payload["name"], "name": payload["name"],
"description": payload["description"], "description": payload["description"],
"config": deepcopy(payload["config"]), "config": deepcopy(payload["config"]),
"status": "pending",
"progress": 20,
"output_dataset_id": None,
"output_count": 0,
"filtered_count": 0,
"duplicate_count": 0,
"error_count": 0,
"failure_reason": None,
"generation_run_id": None,
"started_at": None,
"completed_at": None,
"updated_at": "2026-07-25T20:00:00Z", "updated_at": "2026-07-25T20:00:00Z",
} }
) )
self.regeneration_prepared.add(task_id)
return { return {
"task": deepcopy(task), "task": deepcopy(task),
"preview_invalidated": False, "preview_invalidated": False,
@@ -340,14 +333,19 @@ class FakeDataProcessStore:
def start_generation(self, task_id: str, *, replace_existing: bool) -> dict[str, Any]: def start_generation(self, task_id: str, *, replace_existing: bool) -> dict[str, Any]:
if not self.previews[task_id]: if not self.previews[task_id]:
raise InvalidStateError("preview must be built before generation") raise InvalidStateError("preview must be built before generation")
task = self.tasks[task_id]
if task.get("output_dataset_id") and task_id not in self.regeneration_prepared:
raise InvalidStateError("published task cannot be regenerated")
if replace_existing: if replace_existing:
self.results[task_id] = [] self.results[task_id] = []
self.tasks[task_id].update( task.update(
status="running", status="running",
progress=30, progress=30,
output_dataset_id=None,
output_count=0, output_count=0,
generation_run_id=self._id("dprun"), generation_run_id=self._id("dprun"),
) )
self.regeneration_prepared.discard(task_id)
return self.get_task(task_id) return self.get_task(task_id)
def generation_is_running(self, task_id: str, generation_run_id: str) -> bool: def generation_is_running(self, task_id: str, generation_run_id: str) -> bool:
@@ -487,6 +485,8 @@ class FakeDataProcessStore:
def publish(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]: def publish(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]:
task = self.tasks[task_id] task = self.tasks[task_id]
if task_id in self.regeneration_prepared:
raise InvalidStateError("regeneration must start and complete before publishing")
published = [ published = [
dataset dataset
for dataset in self.datasets.values() for dataset in self.datasets.values()
@@ -969,6 +969,13 @@ def test_regenerate_endpoint_prepares_an_existing_published_task(tmp_path: Path)
"output_count": 2, "output_count": 2,
} }
) )
store.datasets["dataset_train"] = {
"id": "dataset_train",
"name": "原训练集",
"type": "train",
"source_task_id": task_id,
"deleted_at": None,
}
store.previews[task_id] = [{"id": "preview_1", "edited_content": "原切片"}] store.previews[task_id] = [{"id": "preview_1", "edited_content": "原切片"}]
store.results[task_id] = [{"id": "result_1"}] store.results[task_id] = [{"id": "result_1"}]
@@ -985,16 +992,24 @@ def test_regenerate_endpoint_prepares_an_existing_published_task(tmp_path: Path)
assert response.status_code == 200 assert response.status_code == 200
data = response.json()["data"] data = response.json()["data"]
assert data["task"]["status"] == "pending" assert data["task"]["status"] == "completed"
assert data["task"]["output_dataset_id"] is None assert data["task"]["output_dataset_id"] == "dataset_train"
assert data["task"]["output_count"] == 2
assert data["preview_invalidated"] is False assert data["preview_invalidated"] is False
assert data["published_outputs_preserved"] is True assert data["published_outputs_preserved"] is True
assert store.results[task_id] == [] assert store.results[task_id] == [{"id": "result_1"}]
assert store.previews[task_id][0]["id"] == "preview_1" assert store.previews[task_id][0]["id"] == "preview_1"
detail = client.get(f"/modelTF/data-process/{task_id}").json()["data"]
assert detail["status"] == "completed"
assert detail["output_dataset_id"] == "dataset_train"
assert detail["output_count"] == 2
assert [item["id"] for item in detail["output_datasets"]] == ["dataset_train"]
def test_published_split_datasets_remain_in_detail_after_regeneration( def test_published_split_datasets_remain_in_detail_after_regeneration(
tmp_path: Path, tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None: ) -> None:
client, store, _ = make_client(tmp_path) client, store, _ = make_client(tmp_path)
task_id = client.post( task_id = client.post(
@@ -1017,6 +1032,7 @@ def test_published_split_datasets_remain_in_detail_after_regeneration(
"output": "答案", "output": "答案",
} }
] ]
store.previews[task_id] = [{"id": "preview_1", "edited_content": "原切片"}]
published = client.post( published = client.post(
f"/modelTF/data-process/{task_id}/publish", f"/modelTF/data-process/{task_id}/publish",
@@ -1039,16 +1055,40 @@ def test_published_split_datasets_remain_in_detail_after_regeneration(
}, },
) )
assert regenerated.status_code == 200 assert regenerated.status_code == 200
assert regenerated.json()["data"]["task"]["output_dataset_id"] is None original_output_dataset_id = store.tasks[task_id]["output_dataset_id"]
prepared_task = regenerated.json()["data"]["task"]
assert prepared_task["status"] == "completed"
assert prepared_task["output_dataset_id"] == original_output_dataset_id
assert prepared_task["output_count"] == 1
assert store.results[task_id][0]["id"] == "result_1"
detail = client.get(f"/modelTF/data-process/{task_id}") detail = client.get(f"/modelTF/data-process/{task_id}")
assert detail.status_code == 200 assert detail.status_code == 200
detail_data = detail.json()["data"] detail_data = detail.json()["data"]
assert detail_data["status"] == "pending" assert detail_data["status"] == "completed"
assert detail_data["output_dataset_id"] is None assert detail_data["output_dataset_id"] == original_output_dataset_id
assert detail_data["output_count"] == 1
assert len(detail_data["output_datasets"]) == 3 assert len(detail_data["output_datasets"]) == 3
assert {item["id"] for item in detail_data["output_datasets"]} == published_ids assert {item["id"] for item in detail_data["output_datasets"]} == published_ids
assert set(store.datasets) == published_ids assert set(store.datasets) == published_ids
retained_results = client.get(f"/modelTF/data-process/{task_id}/results").json()["data"]
assert retained_results["total"] == 1
assert retained_results["items"][0]["id"] == "result_1"
monkeypatch.setattr(data_process_endpoint, "_run_generation", lambda *args: None)
started = client.post(f"/modelTF/data-process/{task_id}/generate")
assert started.status_code == 200
running = started.json()["data"]
assert running["status"] == "running"
assert running["output_count"] == 0
assert store.results[task_id] == []
assert set(store.datasets) == published_ids
running_detail = client.get(f"/modelTF/data-process/{task_id}").json()["data"]
assert running_detail["status"] == "running"
assert running_detail["output_dataset_id"] is None
assert running_detail["output_count"] == 0
assert {item["id"] for item in running_detail["output_datasets"]} == published_ids
def test_regenerate_endpoint_validates_snapshot_and_locked_process_type( def test_regenerate_endpoint_validates_snapshot_and_locked_process_type(

94
backend/tests/test_data_process_store.py
View File

@@ -225,18 +225,7 @@ class _RegenerationConnection:
"name": params[0], "name": params[0],
"description": params[1], "description": params[1],
"config": params[2], "config": params[2],
"status": "pending", "updated_at": params[3],
"progress": params[3],
"output_dataset_id": None,
"output_count": 0,
"filtered_count": 0,
"duplicate_count": 0,
"error_count": 0,
"failure_reason": None,
"generation_run_id": None,
"started_at": None,
"completed_at": None,
"updated_at": params[4],
} }
) )
return _Result(row=dict(self.task)) return _Result(row=dict(self.task))
@@ -369,39 +358,50 @@ class _TaskListStore(DataProcessStore):
class _StartGenerationConnection: class _StartGenerationConnection:
def __init__(self) -> None: def __init__(self, *, published_prepared: bool = False, preview_count: int = 1) -> None:
config = {"chunk_method": "fixed", "temperature": 0.7}
if published_prepared:
config["_regeneration_prepared"] = {
"prepared": True,
"preview_invalidated": False,
}
self.task = { self.task = {
**_regeneration_task( **_regeneration_task(
status="pending", status="completed" if published_prepared else "pending",
output_dataset_id=None, config=config,
output_dataset_id="dataset_train" if published_prepared else None,
output_count=28, output_count=28,
), ),
"generation_run_id": None, "generation_run_id": None,
} }
self.preview_count = preview_count
self.results = [{"id": "old-result"}] self.results = [{"id": "old-result"}]
def execute(self, sql: str, params: Any = None) -> _Result: def execute(self, sql: str, params: Any = None) -> _Result:
normalized = " ".join(sql.split()) normalized = " ".join(sql.split())
if normalized.startswith("SELECT COUNT(*) AS count FROM data_process_preview_items"): if normalized.startswith("SELECT COUNT(*) AS count FROM data_process_preview_items"):
return _Result(row={"count": 1}) return _Result(row={"count": self.preview_count})
if normalized.startswith("DELETE FROM data_process_results"): if normalized.startswith("DELETE FROM data_process_results"):
self.results.clear() self.results.clear()
return _Result() return _Result()
assert normalized.startswith("UPDATE data_process_tasks SET status='running'") assert normalized.startswith("UPDATE data_process_tasks SET config=%s, status='running'")
assert "output_dataset_id=NULL" in normalized
assert "output_count=0" in normalized assert "output_count=0" in normalized
self.task.update( self.task.update(
{ {
"config": params[0],
"status": "running", "status": "running",
"progress": 30, "progress": 30,
"output_count": 0, "output_count": 0,
"output_dataset_id": None,
"failure_reason": None, "failure_reason": None,
"started_at": params[0], "started_at": params[1],
"completed_at": None, "completed_at": None,
"filtered_count": 0, "filtered_count": 0,
"duplicate_count": 0, "duplicate_count": 0,
"error_count": 0, "error_count": 0,
"generation_run_id": params[1], "generation_run_id": params[2],
"updated_at": params[2], "updated_at": params[3],
} }
) )
return _Result(row=dict(self.task)) return _Result(row=dict(self.task))
@@ -591,6 +591,31 @@ def test_start_generation_clears_previous_output_count() -> None:
assert conn.results == [] assert conn.results == []
def test_prepared_published_task_is_only_cleared_when_generation_starts() -> None:
conn = _StartGenerationConnection(published_prepared=True)
task = _StartGenerationStore(conn).start_generation("task-1")
assert task["status"] == "running"
assert task["output_dataset_id"] is None
assert task["output_count"] == 0
assert "_regeneration_prepared" not in task["config"]
assert conn.results == []
def test_prepared_published_task_survives_generation_preflight_failure() -> None:
conn = _StartGenerationConnection(published_prepared=True, preview_count=0)
with pytest.raises(InvalidStateError, match="preview must be built"):
_StartGenerationStore(conn).start_generation("task-1")
assert conn.task["status"] == "completed"
assert conn.task["output_dataset_id"] == "dataset_train"
assert conn.task["output_count"] == 28
assert "_regeneration_prepared" in conn.task["config"]
assert conn.results == [{"id": "old-result"}]
def _legacy_published_datasets(task_id: str = "task-1") -> list[dict[str, Any]]: def _legacy_published_datasets(task_id: str = "task-1") -> list[dict[str, Any]]:
specs = ( specs = (
("dataset_train", "制度问答-训练集", "train", "train", 22), ("dataset_train", "制度问答-训练集", "train", "train", 22),
@@ -640,13 +665,16 @@ def test_prepare_regeneration_preserves_outputs_sources_and_generation_only_prev
assert result["preview_invalidated"] is False assert result["preview_invalidated"] is False
assert result["published_outputs_preserved"] is True assert result["published_outputs_preserved"] is True
assert result["task"]["output_dataset_id"] is None assert result["task"]["output_dataset_id"] == "dataset_train"
assert result["task"]["status"] == "pending" assert result["task"]["status"] == "completed"
assert result["task"]["progress"] == 20 assert result["task"]["progress"] == 100
assert result["task"]["output_count"] == 0 assert result["task"]["output_count"] == 28
assert result["task"]["started_at"] is None assert result["task"]["started_at"] == "2026-07-25T18:00:00Z"
assert result["task"]["completed_at"] is None assert result["task"]["completed_at"] == "2026-07-25T18:05:00Z"
assert conn.results == [] assert "_regeneration_prepared" not in result["task"]["config"]
stored_config = json.loads(conn.task["config"])
assert stored_config["_regeneration_prepared"]["prepared"] is True
assert conn.results == [{"id": "result_1"}]
assert conn.previews == [{"id": "preview_1"}] assert conn.previews == [{"id": "preview_1"}]
assert conn.datasets == original_datasets assert conn.datasets == original_datasets
assert conn.sources == original_sources assert conn.sources == original_sources
@@ -694,7 +722,7 @@ def test_prepare_regeneration_backfills_and_keeps_legacy_task_datasets_visible()
) )
assert result["published_outputs_preserved"] is True assert result["published_outputs_preserved"] is True
assert result["task"]["output_dataset_id"] is None assert result["task"]["output_dataset_id"] == "dataset_train"
assert len(conn.datasets) == 5 assert len(conn.datasets) == 5
assert all( assert all(
item["source_task_id"] == "task-1" for item in conn.datasets[:3] item["source_task_id"] == "task-1" for item in conn.datasets[:3]
@@ -712,7 +740,7 @@ def test_prepare_regeneration_backfills_and_keeps_legacy_task_datasets_visible()
] ]
def test_prepare_regeneration_deletes_preview_when_chunk_configuration_changes() -> None: def test_prepare_regeneration_defers_preview_deletion_when_chunk_configuration_changes() -> None:
conn = _RegenerationConnection(_regeneration_task(output_dataset_id=None)) conn = _RegenerationConnection(_regeneration_task(output_dataset_id=None))
result = _RegenerationStore(conn).prepare_regeneration( result = _RegenerationStore(conn).prepare_regeneration(
@@ -728,9 +756,11 @@ def test_prepare_regeneration_deletes_preview_when_chunk_configuration_changes()
assert result["preview_invalidated"] is True assert result["preview_invalidated"] is True
assert result["published_outputs_preserved"] is True assert result["published_outputs_preserved"] is True
assert result["task"]["progress"] == 0 assert result["task"]["status"] == "completed"
assert conn.previews == [] assert result["task"]["progress"] == 100
assert conn.results == [] assert result["task"]["output_count"] == 28
assert conn.previews == [{"id": "preview_1"}]
assert conn.results == [{"id": "result_1"}]
assert len(conn.datasets) == 3 assert len(conn.datasets) == 3