From 53014bb381eae7572ecdb9840e19676931fa5cc7 Mon Sep 17 00:00:00 2001 From: caoxiaozhu Date: Mon, 27 Jul 2026 10:43:42 +0800 Subject: [PATCH] =?UTF-8?q?fix(data-process):=20=E5=BB=B6=E8=BF=9F?= =?UTF-8?q?=E9=87=8D=E6=96=B0=E7=94=9F=E6=88=90=E7=A0=B4=E5=9D=8F=E6=80=A7?= =?UTF-8?q?=E5=8F=98=E6=9B=B4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/app/modules/data_process/store.py | 268 ++++++++++++++-------- backend/tests/test_data_process_api.py | 80 +++++-- backend/tests/test_data_process_store.py | 94 +++++--- 3 files changed, 299 insertions(+), 143 deletions(-) diff --git a/backend/app/modules/data_process/store.py b/backend/app/modules/data_process/store.py index b347dda..fc49f60 100644 --- a/backend/app/modules/data_process/store.py +++ b/backend/app/modules/data_process/store.py @@ -41,6 +41,7 @@ _UNSTRUCTURED_PREVIEW_DEFAULTS: dict[str, Any] = { "preserve_code_blocks": True, "preserve_lists": True, } +_REGENERATION_MARKER_KEY = "_regeneration_prepared" 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: if isinstance(value, (datetime, date)): return value.isoformat().replace("+00:00", "Z") @@ -258,7 +287,7 @@ class DataProcessStore: [*params, page_size, (page - 1) * page_size], ).fetchall() return { - "items": [_decode_row(row) for row in rows], + "items": [_public_task(_decode_row(row)) for row in rows], "total": int(total), "page": page, "page_size": page_size, @@ -284,7 +313,13 @@ class DataProcessStore: payload.get("description") or "", payload["process_type"], 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("project_id"), payload.get("owner_id"), @@ -296,7 +331,7 @@ class DataProcessStore: ).fetchone() except psycopg.errors.UniqueViolation as 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]: lock = " FOR UPDATE" if for_update else "" @@ -348,7 +383,7 @@ class DataProcessStore: ).fetchone() if not row: raise NotFoundError("data process task not found") - return _decode_row(row) or {} + return _public_task(_decode_row(row)) or {} def _task_in_connection( self, @@ -370,7 +405,7 @@ class DataProcessStore: def _ensure_editable(task: dict[str, Any]) -> None: if task["status"] not in EDITABLE_STATUSES: 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") def update_task(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]: @@ -381,14 +416,30 @@ class DataProcessStore: "source_dataset_id", } values: dict[str, Any] = {key: value for key, value in payload.items() if key in allowed} - if payload.get("config") is not None: - values["config"] = json_dumps(payload["config"]) - if not values: + if not values and payload.get("config") is None: return self.get_task(task_id) try: with self.connect() as conn: task = self._task_in_connection(conn, task_id, for_update=True) 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 = ( ("config" in payload and payload.get("config") != task.get("config")) or ( @@ -400,7 +451,7 @@ class DataProcessStore: and payload.get("source_dataset_id") != task.get("source_dataset_id") ) ) - if invalidates_results: + if invalidates_results and not regeneration_prepared: values.update( { "status": "pending", @@ -440,18 +491,17 @@ class DataProcessStore: ).fetchone() except psycopg.errors.UniqueViolation as 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( self, task_id: str, payload: dict[str, Any], ) -> dict[str, Any]: - """在一个事务中将已有任务恢复为可重新生成状态。 + """非破坏性地保存重新生成配置。 - 已发布数据集是可被其他任务使用的独立产物，此处只解除当前任务的 - 输出指针，不删除数据集实体。用户完成新一轮生成后，可再次发布 - 以原子替换三个分割数据集的内容。 + 准备阶段保留任务当前状态、结果、切片及已发布数据集。真正开始 + 生成时，才在同一事务内切换运行状态并清理上一轮结果。 """ try: @@ -470,12 +520,18 @@ class DataProcessStore: current_config = dict(task.get("config") or {}) next_config = dict(payload.get("config") or {}) + next_config.pop(_REGENERATION_MARKER_KEY, None) preview_invalidated = _preview_config_changed( process_type, current_config, next_config, ) now = utcnow() + next_config[_REGENERATION_MARKER_KEY] = { + "prepared": True, + "preview_invalidated": preview_invalidated, + "prepared_at": now, + } # 002 迁移前发布的数据集只有 task_id。先补齐新关联字段，保证 # 解除任务输出指针后，详情和后续重新发布仍能定位原来的三份数据集。 conn.execute( @@ -504,35 +560,10 @@ class DataProcessStore: 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( """ UPDATE data_process_tasks - SET name=%s, description=%s, config=%s, status='pending', progress=%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 + SET name=%s, description=%s, config=%s, updated_at=%s WHERE id=%s RETURNING * """, @@ -540,7 +571,6 @@ class DataProcessStore: payload["name"], payload.get("description") or "", json_dumps(next_config), - 20 if preview_count else 0, now, task_id, ), @@ -548,7 +578,7 @@ class DataProcessStore: except psycopg.errors.UniqueViolation as exc: raise ConflictError("data process task name already exists") from exc return { - "task": _decode_row(row) or {}, + "task": _public_task(_decode_row(row)) or {}, "preview_invalidated": preview_invalidated, "published_outputs_preserved": published_outputs_preserved, } @@ -672,9 +702,6 @@ class DataProcessStore: ), ).fetchone() created.append(_decode_row(row) or {}) - # 前端会只对本次新增的源文件构建预览，因此必须保留旧 - # 文件的切片。任何源文件增加都会使旧生成结果失效。 - conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,)) preview_row = conn.execute( """ SELECT COUNT(*) AS count FROM data_process_preview_items @@ -683,20 +710,40 @@ class DataProcessStore: (task_id,), ).fetchone() preview_count = int((preview_row or {}).get("count") or 0) - 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), - ) + if _is_regeneration_prepared(task): + conn.execute( + """ + UPDATE data_process_tasks + SET 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), + ) + 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: raise ConflictError( "the same source file content is already attached to this task" @@ -777,23 +824,39 @@ class DataProcessStore: ).fetchone() if not row: raise NotFoundError("source file not found") - conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,)) conn.execute( "DELETE FROM data_process_preview_items WHERE source_file_id=%s", (file_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, utcnow(), task_id), - ) + now = utcnow() + if _is_regeneration_prepared(task): + conn.execute( + """ + UPDATE data_process_tasks + SET 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), + ) + 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( self, @@ -823,7 +886,9 @@ class DataProcessStore: 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,)) + 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: conn.execute( "DELETE FROM data_process_preview_items WHERE task_id=%s", (task_id,) @@ -878,15 +943,21 @@ class DataProcessStore: ), ).fetchone() created.append(_decode_row(row) or {}) - conn.execute( - """ - 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), - ) + if regeneration_prepared: + conn.execute( + "UPDATE data_process_tasks SET updated_at=%s WHERE id=%s", + (now, task_id), + ) + else: + conn.execute( + """ + 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 def list_preview_items( @@ -980,7 +1051,7 @@ class DataProcessStore: now, ), ).fetchone() - self._invalidate_results(conn, task_id, now) + self._invalidate_results(conn, task, task_id, now) return _decode_row(row) or {} def update_preview_item( @@ -1027,7 +1098,7 @@ class DataProcessStore: task_id, ), ).fetchone() - self._invalidate_results(conn, task_id, now) + self._invalidate_results(conn, task, task_id, now) return _decode_row(row) or {} def delete_preview_item(self, task_id: str, preview_id: str) -> None: @@ -1040,14 +1111,21 @@ class DataProcessStore: ).fetchone() if not row: raise NotFoundError("preview item not found") - self._invalidate_results(conn, task_id, utcnow()) + self._invalidate_results(conn, task, task_id, utcnow()) def _invalidate_results( self, conn: psycopg.Connection[dict[str, Any]], + task: dict[str, Any], task_id: str, now: str, ) -> 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( """ @@ -1065,7 +1143,8 @@ class DataProcessStore: raise DataProcessStoreError("incremental generation is not supported") with self.connect() as conn: 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") if task["status"] == "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,)) now = utcnow() generation_run_id = new_id("dprun") + next_config = dict(task.get("config") or {}) + next_config.pop(_REGENERATION_MARKER_KEY, None) row = conn.execute( """ UPDATE data_process_tasks - SET status='running', progress=30, failure_reason=NULL, started_at=%s, - completed_at=NULL, output_count=0, filtered_count=0, - duplicate_count=0, error_count=0, generation_run_id=%s, updated_at=%s + SET config=%s, status='running', progress=30, failure_reason=NULL, + started_at=%s, completed_at=NULL, output_dataset_id=NULL, + output_count=0, filtered_count=0, duplicate_count=0, error_count=0, + generation_run_id=%s, updated_at=%s WHERE id=%s RETURNING * """, - (now, generation_run_id, now, task_id), + (json_dumps(next_config), now, generation_run_id, now, task_id), ).fetchone() return _decode_row(row) or {} @@ -1405,6 +1487,10 @@ class DataProcessStore: """按精确配额发布训练、验证、测试三个独立数据集。""" with self.connect() as conn: 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": raise InvalidStateError("only a completed task can be published") rows = conn.execute( diff --git a/backend/tests/test_data_process_api.py b/backend/tests/test_data_process_api.py index c0d9ad7..f1b29c5 100644 --- a/backend/tests/test_data_process_api.py +++ b/backend/tests/test_data_process_api.py @@ -30,6 +30,7 @@ class FakeDataProcessStore: self.previews: dict[str, list[dict[str, Any]]] = {} self.results: dict[str, list[dict[str, Any]]] = {} self.datasets: dict[str, dict[str, Any]] = {} + self.regeneration_prepared: set[str] = set() self.sequence = 0 def _id(self, prefix: str) -> str: @@ -109,27 +110,19 @@ class FakeDataProcessStore: raise InvalidStateError("data process task was modified by another request") if payload["process_type"] != task["process_type"]: raise InvalidStateError("process_type cannot be changed during regeneration") - published_outputs_preserved = bool(task.get("output_dataset_id")) - self.results[task_id] = [] + published_outputs_preserved = bool(task.get("output_dataset_id")) or any( + dataset.get("source_task_id") == task_id + for dataset in self.datasets.values() + ) task.update( { "name": payload["name"], "description": payload["description"], "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", } ) + self.regeneration_prepared.add(task_id) return { "task": deepcopy(task), "preview_invalidated": False, @@ -340,14 +333,19 @@ class FakeDataProcessStore: def start_generation(self, task_id: str, *, replace_existing: bool) -> dict[str, Any]: if not self.previews[task_id]: 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: self.results[task_id] = [] - self.tasks[task_id].update( + task.update( status="running", progress=30, + output_dataset_id=None, output_count=0, generation_run_id=self._id("dprun"), ) + self.regeneration_prepared.discard(task_id) return self.get_task(task_id) 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]: task = self.tasks[task_id] + if task_id in self.regeneration_prepared: + raise InvalidStateError("regeneration must start and complete before publishing") published = [ dataset 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, } ) + 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.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 data = response.json()["data"] - assert data["task"]["status"] == "pending" - assert data["task"]["output_dataset_id"] is None + assert data["task"]["status"] == "completed" + assert data["task"]["output_dataset_id"] == "dataset_train" + assert data["task"]["output_count"] == 2 assert data["preview_invalidated"] is False 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" + 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( tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, ) -> None: client, store, _ = make_client(tmp_path) task_id = client.post( @@ -1017,6 +1032,7 @@ def test_published_split_datasets_remain_in_detail_after_regeneration( "output": "答案", } ] + store.previews[task_id] = [{"id": "preview_1", "edited_content": "原切片"}] published = client.post( 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.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}") assert detail.status_code == 200 detail_data = detail.json()["data"] - assert detail_data["status"] == "pending" - assert detail_data["output_dataset_id"] is None + assert detail_data["status"] == "completed" + assert detail_data["output_dataset_id"] == original_output_dataset_id + assert detail_data["output_count"] == 1 assert len(detail_data["output_datasets"]) == 3 assert {item["id"] for item in detail_data["output_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( diff --git a/backend/tests/test_data_process_store.py b/backend/tests/test_data_process_store.py index a2612ff..d316bef 100644 --- a/backend/tests/test_data_process_store.py +++ b/backend/tests/test_data_process_store.py @@ -225,18 +225,7 @@ class _RegenerationConnection: "name": params[0], "description": params[1], "config": params[2], - "status": "pending", - "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], + "updated_at": params[3], } ) return _Result(row=dict(self.task)) @@ -369,39 +358,50 @@ class _TaskListStore(DataProcessStore): 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 = { **_regeneration_task( - status="pending", - output_dataset_id=None, + status="completed" if published_prepared else "pending", + config=config, + output_dataset_id="dataset_train" if published_prepared else None, output_count=28, ), "generation_run_id": None, } + self.preview_count = preview_count self.results = [{"id": "old-result"}] def execute(self, sql: str, params: Any = None) -> _Result: normalized = " ".join(sql.split()) 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"): self.results.clear() 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 self.task.update( { + "config": params[0], "status": "running", "progress": 30, "output_count": 0, + "output_dataset_id": None, "failure_reason": None, - "started_at": params[0], + "started_at": params[1], "completed_at": None, "filtered_count": 0, "duplicate_count": 0, "error_count": 0, - "generation_run_id": params[1], - "updated_at": params[2], + "generation_run_id": params[2], + "updated_at": params[3], } ) return _Result(row=dict(self.task)) @@ -591,6 +591,31 @@ def test_start_generation_clears_previous_output_count() -> None: 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]]: specs = ( ("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["published_outputs_preserved"] is True - assert result["task"]["output_dataset_id"] is None - assert result["task"]["status"] == "pending" - assert result["task"]["progress"] == 20 - assert result["task"]["output_count"] == 0 - assert result["task"]["started_at"] is None - assert result["task"]["completed_at"] is None - assert conn.results == [] + assert result["task"]["output_dataset_id"] == "dataset_train" + assert result["task"]["status"] == "completed" + assert result["task"]["progress"] == 100 + assert result["task"]["output_count"] == 28 + assert result["task"]["started_at"] == "2026-07-25T18:00:00Z" + assert result["task"]["completed_at"] == "2026-07-25T18:05:00Z" + 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.datasets == original_datasets 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["task"]["output_dataset_id"] is None + assert result["task"]["output_dataset_id"] == "dataset_train" assert len(conn.datasets) == 5 assert all( 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)) 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["published_outputs_preserved"] is True - assert result["task"]["progress"] == 0 - assert conn.previews == [] - assert conn.results == [] + assert result["task"]["status"] == "completed" + assert result["task"]["progress"] == 100 + assert result["task"]["output_count"] == 28 + assert conn.previews == [{"id": "preview_1"}] + assert conn.results == [{"id": "result_1"}] assert len(conn.datasets) == 3