from __future__ import annotations import hashlib import json import uuid from collections.abc import Iterator, Sequence from contextlib import contextmanager from datetime import UTC, date, datetime from decimal import Decimal from functools import lru_cache from pathlib import Path from typing import Any import psycopg 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_assignments TASK_STATUSES = {"pending", "running", "completed", "failed", "stopped"} EDITABLE_STATUSES = {"pending", "failed", "stopped", "completed"} class DataProcessStoreError(RuntimeError): pass class NotFoundError(DataProcessStoreError): pass class ConflictError(DataProcessStoreError): pass class InvalidStateError(DataProcessStoreError): pass def utcnow() -> str: return datetime.now(UTC).replace(microsecond=0).isoformat().replace("+00:00", "Z") def new_id(prefix: str) -> str: return f"{prefix}_{uuid.uuid4().hex[:20]}" def json_dumps(value: Any) -> str: return json.dumps(value, ensure_ascii=False, separators=(",", ":")) def _database_url(value: str) -> str: return value.replace("postgresql+psycopg://", "postgresql://") def _json_value(value: Any, default: Any) -> Any: if value is None or value == "": return default if isinstance(value, (dict, list)): return value try: return json.loads(value) except (TypeError, json.JSONDecodeError): return default def _serialize_value(value: Any) -> Any: if isinstance(value, (datetime, date)): return value.isoformat().replace("+00:00", "Z") if isinstance(value, Decimal): return float(value) 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 item = {key: _serialize_value(value) for key, value in row.items()} for key, default in { "config": {}, "metadata": {}, "quality_score": {}, "versions": [], }.items(): if key in item: item[key] = _json_value(item[key], default) return item class DataProcessStore: """数据处理持久层。 构造函数不会连接数据库或执行迁移。部署方必须显式执行 002 SQL， 或在受控的管理命令中调用 :meth:`ensure_schema`，避免应用启动时 修改远程数据库。 """ def __init__(self, database_url: str | None = None) -> None: self.database_url = _database_url(database_url or get_settings().database_url) @contextmanager def connect(self) -> Iterator[psycopg.Connection[dict[str, Any]]]: with psycopg.connect(self.database_url, row_factory=dict_row) as conn: try: yield conn conn.commit() except Exception: conn.rollback() raise def ensure_schema(self) -> None: """显式安装数据处理表；API 路由和应用启动流程不会调用此方法。""" schema_path = Path(__file__).resolve().parents[2] / "db" / "sql" / "002_data_process.sql" sql = schema_path.read_text(encoding="utf-8") with self.connect() as conn, conn.cursor() as cursor: cursor.execute(sql) def list_tasks( self, *, page: int = 1, page_size: int = 20, keyword: str | None = None, status: str | None = None, process_type: str | None = None, tenant_id: str | None = None, project_id: str | None = None, ) -> dict[str, Any]: clauses = ["deleted_at IS NULL"] params: list[Any] = [] if keyword: clauses.append("(name ILIKE %s OR COALESCE(description, '') ILIKE %s)") pattern = f"%{keyword.strip()}%" params.extend([pattern, pattern]) if status: clauses.append("status = %s") params.append(status) if process_type: clauses.append("process_type = %s") params.append(process_type) if tenant_id: clauses.append("tenant_id = %s") params.append(tenant_id) if project_id: clauses.append("project_id = %s") params.append(project_id) where = " AND ".join(clauses) with self.connect() as conn: total = conn.execute( f"SELECT COUNT(*) AS count FROM data_process_tasks WHERE {where}", params ).fetchone()["count"] rows = conn.execute( f""" SELECT * FROM data_process_tasks WHERE {where} ORDER BY created_at DESC, id DESC LIMIT %s OFFSET %s """, [*params, page_size, (page - 1) * page_size], ).fetchall() return { "items": [_decode_row(row) for row in rows], "total": int(total), "page": page, "page_size": page_size, } def create_task(self, payload: dict[str, Any]) -> dict[str, Any]: task_id = new_id("dpt") now = utcnow() try: with self.connect() as conn: row = conn.execute( """ INSERT INTO data_process_tasks (id, name, description, status, process_type, source_dataset_id, config, progress, tenant_id, project_id, owner_id, created_by, updated_by, created_at, updated_at) VALUES (%s, %s, %s, 'pending', %s, %s, %s, 0, %s, %s, %s, %s, %s, %s, %s) RETURNING * """, ( task_id, payload["name"], payload.get("description") or "", payload["process_type"], payload.get("source_dataset_id"), json_dumps(payload.get("config") or {}), payload.get("tenant_id"), payload.get("project_id"), payload.get("owner_id"), payload.get("created_by"), payload.get("created_by"), now, now, ), ).fetchone() except psycopg.errors.UniqueViolation as exc: raise ConflictError("data process task name already exists") from exc return _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 "" with self.connect() as conn: if for_update: row = conn.execute( f"SELECT * FROM data_process_tasks WHERE id=%s AND deleted_at IS NULL{lock}", (task_id,), ).fetchone() else: row = conn.execute( """ SELECT task.*, (SELECT COUNT(*) FROM data_process_source_files source WHERE source.task_id=task.id AND source.deleted_at IS NULL) AS source_file_count, (SELECT COUNT(*) FROM data_process_preview_items preview WHERE preview.task_id=task.id) AS preview_count, CASE WHEN task.started_at IS NOT NULL AND task.completed_at IS NOT NULL THEN EXTRACT(EPOCH FROM (task.completed_at - task.started_at)) ELSE NULL END AS duration_seconds FROM data_process_tasks task WHERE task.id=%s AND task.deleted_at IS NULL """, (task_id,), ).fetchone() if not row: raise NotFoundError("data process task not found") return _decode_row(row) or {} def _task_in_connection( self, conn: psycopg.Connection[dict[str, Any]], task_id: str, *, for_update: bool = False, ) -> dict[str, Any]: lock = " FOR UPDATE" if for_update else "" row = conn.execute( f"SELECT * FROM data_process_tasks WHERE id=%s AND deleted_at IS NULL{lock}", (task_id,), ).fetchone() if not row: raise NotFoundError("data process task not found") return _decode_row(row) or {} @staticmethod 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"): raise InvalidStateError("published task cannot be edited") def update_task(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]: allowed = { "name", "description", "process_type", "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: 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) invalidates_results = ( ("config" in payload and payload.get("config") != task.get("config")) or ( "process_type" in payload and payload.get("process_type") != task.get("process_type") ) or ( "source_dataset_id" in payload and payload.get("source_dataset_id") != task.get("source_dataset_id") ) ) if invalidates_results: values.update( { "status": "pending", "progress": 0, "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, } ) 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 ( "process_type" in payload and payload.get("process_type") != task.get("process_type") ): conn.execute( """ UPDATE data_process_source_files SET deleted_at=%s, updated_at=%s WHERE task_id=%s AND deleted_at IS NULL """, (utcnow(), utcnow(), task_id), ) values["input_count"] = 0 values["updated_at"] = utcnow() assignments = ", ".join(f"{key}=%s" for key in values) row = conn.execute( f"UPDATE data_process_tasks SET {assignments} WHERE id=%s RETURNING *", [*values.values(), task_id], ).fetchone() except psycopg.errors.UniqueViolation as exc: raise ConflictError("data process task name already exists") from exc return _decode_row(row) or {} def delete_task(self, task_id: str, *, deleted_by: str | None = None) -> None: with self.connect() as conn: task = self._task_in_connection(conn, task_id, for_update=True) if task["status"] == "running": raise InvalidStateError("running task must be stopped before deletion") now = utcnow() conn.execute( """ UPDATE data_process_tasks SET deleted_at=%s, deleted_by=%s, updated_at=%s WHERE id=%s """, (now, deleted_by, now, task_id), ) def list_source_files(self, task_id: str) -> list[dict[str, Any]]: self.get_task(task_id) with self.connect() as conn: rows = conn.execute( """ SELECT id, task_id, storage_object_id, name, size_bytes, record_count, file_format, checksum_sha256, version_no, content_preview, metadata, tenant_id, project_id, created_by, created_at, updated_at FROM data_process_source_files WHERE task_id=%s AND deleted_at IS NULL ORDER BY created_at, id """, (task_id,), ).fetchall() return [_decode_row(row) or {} for row in rows] def add_source_file( self, task_id: str, *, name: str, content: str, raw_size: int, checksum_sha256: str, file_format: str, 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, [ { "name": name, "content": content, "raw_size": raw_size, "checksum_sha256": checksum_sha256, "file_format": file_format, "record_count": record_count, "metadata": metadata or {}, "created_by": created_by, "id": source_file_id, "storage_object_id": storage_object_id, } ], )[0] def add_source_files( self, task_id: str, files: Sequence[dict[str, Any]], ) -> list[dict[str, Any]]: """在同一事务中登记一个上传批次，任一文件失败则全部回滚。""" if not files: raise DataProcessStoreError("at least one source file is required") now = utcnow() created: list[dict[str, Any]] = [] try: with self.connect() as conn: task = self._task_in_connection(conn, task_id, for_update=True) self._ensure_editable(task) for payload in files: 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 (id, task_id, storage_object_id, name, size_bytes, record_count, file_format, checksum_sha256, version_no, content, content_preview, metadata, tenant_id, project_id, created_by, created_at, updated_at) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, 1, %s, %s, %s, %s, %s, %s, %s, %s) RETURNING id, task_id, storage_object_id, name, size_bytes, record_count, file_format, checksum_sha256, version_no, content_preview, metadata, tenant_id, project_id, created_by, created_at, updated_at """, ( file_id, task_id, storage_object_id, payload["name"], payload["raw_size"], payload["record_count"], payload["file_format"], payload["checksum_sha256"], payload["content"], str(payload["content"])[:2000], json_dumps(metadata_payload), task.get("tenant_id"), task.get("project_id"), payload.get("created_by") or task.get("created_by"), now, now, ), ).fetchone() created.append(_decode_row(row) or {}) 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,)) 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, 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 """, (task_id, now, task_id), ) except psycopg.errors.UniqueViolation as exc: raise ConflictError( "the same source file content is already attached to this task" ) from exc return created def get_source_file( self, task_id: str, file_id: str, *, include_content: bool = True ) -> dict[str, Any]: # 先验证父任务仍然可见，避免软删除任务后通过已知文件 ID 读取正文。 self.get_task(task_id) content_column = ", content" if include_content else "" with self.connect() as conn: row = conn.execute( f""" SELECT id, task_id, storage_object_id, name, size_bytes, record_count, file_format, checksum_sha256, version_no, content_preview, metadata, tenant_id, project_id, created_by, created_at, updated_at{content_column} FROM data_process_source_files WHERE id=%s AND task_id=%s AND deleted_at IS NULL """, (file_id, task_id), ).fetchone() if not row: raise NotFoundError("source file not found") return _decode_row(row) or {} def source_content_window( self, task_id: str, file_id: str, offset: int, limit: int ) -> dict[str, Any]: source_file = self.get_source_file(task_id, file_id, include_content=True) content = str(source_file.pop("content", "")) window = content[offset : offset + limit] return { "file": source_file, "content": window, "offset": offset, "limit": limit, "total_chars": len(content), "has_more": offset + len(window) < len(content), } def source_content_lines( self, task_id: str, file_id: str, start_line: int, line_count: int, ) -> dict[str, Any]: source_file = self.get_source_file(task_id, file_id, include_content=True) content = str(source_file.pop("content", "")) lines = content.splitlines(keepends=True) start_index = min(len(lines), start_line - 1) selected = lines[start_index : start_index + line_count] end_line = start_index + len(selected) return { "file": source_file, "content": "".join(selected), "start_line": start_line, "end_line": end_line, "line_count": len(selected), "total_lines": len(lines), "has_more": end_line < len(lines), } def delete_source_file(self, task_id: str, file_id: str) -> None: with self.connect() as conn: task = self._task_in_connection(conn, task_id, for_update=True) self._ensure_editable(task) row = conn.execute( """ UPDATE data_process_source_files SET deleted_at=%s, updated_at=%s WHERE id=%s AND task_id=%s AND deleted_at IS NULL RETURNING id """, (utcnow(), utcnow(), file_id, task_id), ).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), ) def replace_preview_items( 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,)) 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( """ INSERT INTO data_process_preview_items (id, task_id, source_file_id, original_content, edited_content, source_start, source_end, source_start_line, source_end_line, token_count, status, quality_score, created_at, updated_at) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) RETURNING * """, ( item.get("id") or new_id("dpp"), task_id, item.get("source_file_id"), item.get("original_content") or "", item.get("edited_content", item.get("original_content") or ""), item.get("source_start"), item.get("source_end"), item.get("source_start_line"), item.get("source_end_line"), max(0, int(item.get("token_count") or 0)), item.get("status") or "original", json_dumps(item.get("quality_score") or {}), now, now, ), ).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), ) return created def list_preview_items( self, task_id: str, *, source_file_id: str | None = None, page: int = 1, page_size: int = 200, keyword: str | None = None, ) -> dict[str, Any]: self.get_task(task_id) clauses = ["task_id=%s"] params: list[Any] = [task_id] if source_file_id: clauses.append("source_file_id=%s") params.append(source_file_id) if keyword: clauses.append("(original_content ILIKE %s OR edited_content ILIKE %s)") pattern = f"%{keyword.strip()}%" params.extend([pattern, pattern]) where = " AND ".join(clauses) with self.connect() as conn: total = conn.execute( f"SELECT COUNT(*) AS count FROM data_process_preview_items WHERE {where}", params ).fetchone()["count"] rows = conn.execute( f""" SELECT * FROM data_process_preview_items WHERE {where} ORDER BY source_file_id NULLS LAST, source_start NULLS LAST, created_at, id LIMIT %s OFFSET %s """, [*params, page_size, (page - 1) * page_size], ).fetchall() return { "items": [_decode_row(row) for row in rows], "total": int(total), "page": page, "page_size": page_size, } def get_preview_item(self, task_id: str, preview_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute( "SELECT * FROM data_process_preview_items WHERE id=%s AND task_id=%s", (preview_id, task_id), ).fetchone() if not row: raise NotFoundError("preview item not found") return _decode_row(row) or {} def create_preview_item(self, task_id: str, item: dict[str, Any]) -> dict[str, Any]: now = utcnow() with self.connect() as conn: task = self._task_in_connection(conn, task_id, for_update=True) self._ensure_editable(task) if item.get("source_file_id"): source = conn.execute( """ SELECT id FROM data_process_source_files WHERE id=%s AND task_id=%s AND deleted_at IS NULL """, (item["source_file_id"], task_id), ).fetchone() if not source: raise NotFoundError("source file not found") row = conn.execute( """ INSERT INTO data_process_preview_items (id, task_id, source_file_id, original_content, edited_content, source_start, source_end, source_start_line, source_end_line, token_count, status, quality_score, created_at, updated_at) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) RETURNING * """, ( new_id("dpp"), task_id, item.get("source_file_id"), item.get("original_content") or "", item.get("edited_content") or "", item.get("source_start"), item.get("source_end"), item.get("source_start_line"), item.get("source_end_line"), max(0, int(item.get("token_count") or 0)), item.get("status") or "manual", json_dumps(item.get("quality_score") or {}), now, now, ), ).fetchone() self._invalidate_results(conn, task_id, now) return _decode_row(row) or {} def update_preview_item( self, task_id: str, preview_id: str, payload: dict[str, Any] ) -> dict[str, Any]: with self.connect() as conn: task = self._task_in_connection(conn, task_id, for_update=True) self._ensure_editable(task) existing = conn.execute( "SELECT * FROM data_process_preview_items WHERE id=%s AND task_id=%s", (preview_id, task_id), ).fetchone() if not existing: raise NotFoundError("preview item not found") expected_updated_at = payload.get("expected_updated_at") current_updated_at = _serialize_value(existing.get("updated_at")) if expected_updated_at and expected_updated_at != current_updated_at: raise ConflictError("preview item was modified by another request") edited = payload["edited_content"] status = payload.get("status") if not status: if not edited.strip(): status = "invalid" elif edited == existing["original_content"]: status = "original" else: status = "modified" now = utcnow() row = conn.execute( """ UPDATE data_process_preview_items SET edited_content=%s, token_count=%s, status=%s, quality_score=%s, updated_at=%s WHERE id=%s AND task_id=%s RETURNING * """, ( edited, estimate_token_count(edited), status, json_dumps(payload.get("quality_score") or {}), now, preview_id, task_id, ), ).fetchone() self._invalidate_results(conn, task_id, now) return _decode_row(row) or {} def delete_preview_item(self, task_id: str, preview_id: str) -> None: with self.connect() as conn: task = self._task_in_connection(conn, task_id, for_update=True) self._ensure_editable(task) row = conn.execute( "DELETE FROM data_process_preview_items WHERE id=%s AND task_id=%s RETURNING id", (preview_id, task_id), ).fetchone() if not row: raise NotFoundError("preview item not found") self._invalidate_results(conn, task_id, utcnow()) def _invalidate_results( self, conn: psycopg.Connection[dict[str, Any]], task_id: str, now: str, ) -> None: conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,)) 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, generation_run_id=NULL, updated_at=%s WHERE id=%s """, (now, task_id), ) def start_generation(self, task_id: str, *, replace_existing: bool = True) -> dict[str, Any]: if not replace_existing: 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"): raise InvalidStateError("published task cannot be regenerated") if task["status"] == "running": raise ConflictError("data process task is already running") preview_count = conn.execute( "SELECT COUNT(*) AS count FROM data_process_preview_items WHERE task_id=%s", (task_id,), ).fetchone()["count"] if not preview_count: raise InvalidStateError("preview must be built before generation") conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,)) now = utcnow() generation_run_id = new_id("dprun") row = conn.execute( """ UPDATE data_process_tasks SET status='running', progress=30, failure_reason=NULL, started_at=%s, completed_at=NULL, 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), ).fetchone() return _decode_row(row) or {} def stop_task(self, task_id: str) -> dict[str, Any]: with self.connect() as conn: task = self._task_in_connection(conn, task_id, for_update=True) if task["status"] != "running": raise InvalidStateError("only a running task can be stopped") now = utcnow() row = conn.execute( """ UPDATE data_process_tasks SET status='stopped', failure_reason=NULL, generation_run_id=NULL, updated_at=%s WHERE id=%s RETURNING * """, (now, task_id), ).fetchone() return _decode_row(row) or {} def generation_is_running(self, task_id: str, generation_run_id: str) -> bool: task = self.get_task(task_id) return ( task["status"] == "running" and task.get("generation_run_id") == generation_run_id ) def update_generation_progress( self, task_id: str, generation_run_id: str, processed_count: int, total_count: int, ) -> bool: ratio = processed_count / max(1, total_count) progress = min(95.0, 30.0 + ratio * 65.0) with self.connect() as conn: row = conn.execute( """ UPDATE data_process_tasks SET progress=%s, updated_at=%s WHERE id=%s AND status='running' AND generation_run_id=%s RETURNING id """, (progress, utcnow(), task_id, generation_run_id), ).fetchone() return row is not None def complete_generation( self, task_id: str, results: Sequence[dict[str, Any]], *, generation_run_id: str, filtered_count: int = 0, duplicate_count: int = 0, error_count: int = 0, ) -> dict[str, Any]: now = utcnow() with self.connect() as conn: task = self._task_in_connection(conn, task_id, for_update=True) if ( task["status"] != "running" or task.get("generation_run_id") != generation_run_id ): return task conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,)) for result in results: conn.execute( """ INSERT INTO data_process_results (id, task_id, preview_item_id, instruction, input, output, original_instruction, original_input, original_output, status, error, split, quality_score, created_at, updated_at) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """, ( result.get("id") or new_id("dpr"), task_id, result.get("preview_item_id"), result.get("instruction") or "", result.get("input") or "", result.get("output") or "", result.get("original_instruction", result.get("instruction") or ""), result.get("original_input", result.get("input") or ""), result.get("original_output", result.get("output") or ""), result.get("status") or "valid", result.get("error"), result.get("split"), json_dumps(result.get("quality_score") or {}), now, now, ), ) row = conn.execute( """ UPDATE data_process_tasks SET status='completed', progress=100, output_count=%s, filtered_count=%s, duplicate_count=%s, error_count=%s, failure_reason=NULL, completed_at=%s, generation_run_id=NULL, updated_at=%s WHERE id=%s AND generation_run_id=%s RETURNING * """, ( len(results), filtered_count, duplicate_count, error_count, now, now, task_id, generation_run_id, ), ).fetchone() return _decode_row(row) or {} def mark_failed( self, task_id: str, reason: str, *, generation_run_id: str ) -> dict[str, Any]: with self.connect() as conn: task = self._task_in_connection(conn, task_id, for_update=True) if ( task["status"] != "running" or task.get("generation_run_id") != generation_run_id ): return task now = utcnow() row = conn.execute( """ UPDATE data_process_tasks SET status='failed', failure_reason=%s, completed_at=%s, generation_run_id=NULL, updated_at=%s WHERE id=%s AND generation_run_id=%s RETURNING * """, (reason[:4000], now, now, task_id, generation_run_id), ).fetchone() return _decode_row(row) or {} def progress(self, task_id: str) -> dict[str, Any]: task = self.get_task(task_id) return { "task_id": task["id"], "status": task["status"], "progress": float(task.get("progress") or 0), "input_count": int(task.get("input_count") or 0), "output_count": int(task.get("output_count") or 0), "filtered_count": int(task.get("filtered_count") or 0), "duplicate_count": int(task.get("duplicate_count") or 0), "error_count": int(task.get("error_count") or 0), "failure_reason": task.get("failure_reason"), "started_at": task.get("started_at"), "completed_at": task.get("completed_at"), } def list_results( self, task_id: str, *, page: int = 1, page_size: int = 100, status: str | None = None, split: str | None = None, keyword: str | None = None, ) -> dict[str, Any]: self.get_task(task_id) clauses = ["task_id=%s"] params: list[Any] = [task_id] if status: clauses.append("status=%s") params.append(status) if split: clauses.append("split=%s") params.append(split) if keyword: clauses.append("(instruction ILIKE %s OR input ILIKE %s OR output ILIKE %s)") pattern = f"%{keyword.strip()}%" params.extend([pattern, pattern, pattern]) where = " AND ".join(clauses) with self.connect() as conn: total = conn.execute( f"SELECT COUNT(*) AS count FROM data_process_results WHERE {where}", params ).fetchone()["count"] rows = conn.execute( f""" SELECT * FROM data_process_results WHERE {where} ORDER BY created_at, id LIMIT %s OFFSET %s """, [*params, page_size, (page - 1) * page_size], ).fetchall() return { "items": [_decode_row(row) for row in rows], "total": int(total), "page": page, "page_size": page_size, } def get_result(self, task_id: str, result_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute( "SELECT * FROM data_process_results WHERE id=%s AND task_id=%s", (result_id, task_id), ).fetchone() if not row: raise NotFoundError("data process result not found") return _decode_row(row) or {} def update_result( self, task_id: str, result_id: str, payload: dict[str, Any] ) -> dict[str, Any]: allowed = {"instruction", "input", "output", "quality_score"} values = {key: value for key, value in payload.items() if key in allowed} if "quality_score" in values: values["quality_score"] = json_dumps(values["quality_score"]) if not values: raise DataProcessStoreError("no result fields supplied") with self.connect() as conn: task = self._task_in_connection(conn, task_id, for_update=True) if task["status"] == "running": raise InvalidStateError("results cannot be edited while generation is running") if task.get("output_dataset_id"): raise InvalidStateError("published results cannot be edited") current = conn.execute( "SELECT * FROM data_process_results WHERE id=%s AND task_id=%s", (result_id, task_id), ).fetchone() if not current: raise NotFoundError("data process result not found") expected_updated_at = payload.get("expected_updated_at") current_updated_at = _serialize_value(current.get("updated_at")) if expected_updated_at and expected_updated_at != current_updated_at: raise ConflictError("data process result was modified by another request") merged = {**current, **values} quality = payload.get("quality_score") or {} hard_valid = bool( str(merged.get("instruction") or "").strip() and str(merged.get("output") or "").strip() ) quality_valid = bool(quality.get("is_valid", hard_valid)) changed = any( str(merged.get(field) or "") != str(merged.get(f"original_{field}") or "") for field in ("instruction", "input", "output") ) status = "invalid" if not hard_valid or not quality_valid else ( "modified" if changed else "valid" ) values["status"] = status flags = quality.get("flags") if isinstance(quality, dict) else None values["error"] = ", ".join(str(flag) for flag in flags or []) or ( "quality validation failed" if status == "invalid" else None ) values["updated_at"] = utcnow() assignments = ", ".join(f"{key}=%s" for key in values) row = conn.execute( f"""UPDATE data_process_results SET {assignments} WHERE id=%s AND task_id=%s RETURNING *""", [*values.values(), result_id, task_id], ).fetchone() conn.execute( """ UPDATE data_process_tasks SET error_count=( SELECT COUNT(*) FROM data_process_results WHERE task_id=%s AND status='invalid' ), updated_at=%s WHERE id=%s """, (task_id, utcnow(), task_id), ) return _decode_row(row) or {} def get_generation_model(self, model_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute( """ SELECT id, name, type, purpose, model_source, description, path, api_url, api_key, online_model_name, create_time FROM models WHERE id=%s """, (model_id,), ).fetchone() if not row: raise NotFoundError("generation model not found") return _decode_row(row) or {} def save_generation_model_snapshot( self, task_id: str, model_snapshot: dict[str, Any], *, generation_run_id: str, ) -> dict[str, Any]: # API 密钥仅用于本次调用，绝不能进入任务配置、详情响应或审计快照。 safe_snapshot = { key: value for key, value in model_snapshot.items() if key != "api_key" } with self.connect() as conn: task = self._task_in_connection(conn, task_id, for_update=True) if ( task["status"] != "running" or task.get("generation_run_id") != generation_run_id ): raise InvalidStateError("generation run is no longer active") config = dict(task.get("config") or {}) config["generation_model_snapshot"] = safe_snapshot row = conn.execute( """ UPDATE data_process_tasks SET config=%s, updated_at=%s WHERE id=%s AND generation_run_id=%s RETURNING * """, (json_dumps(config), utcnow(), task_id, generation_run_id), ).fetchone() return _decode_row(row) or {} def publish(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]: """按精确配额发布三个切分文件；重复发布会同步修复既有发布物。""" with self.connect() as conn: task = self._task_in_connection(conn, task_id, for_update=True) existing_dataset = None if task.get("output_dataset_id"): existing_dataset = conn.execute( "SELECT * FROM datasets WHERE id=%s", (task["output_dataset_id"],) ).fetchone() if not existing_dataset: # 数据集被外部流程清理后，解除断链并重新发布。 conn.execute( "UPDATE data_process_tasks SET output_dataset_id=NULL WHERE id=%s", (task_id,), ) if task["status"] != "completed": raise InvalidStateError("only a completed task can be published") rows = conn.execute( """ SELECT * FROM data_process_results WHERE task_id=%s ORDER BY created_at, id """, (task_id,), ).fetchall() if not rows: raise InvalidStateError("task has no results to publish") invalid_count = sum( 1 for row in rows if row["status"] == "invalid" or not str(row.get("instruction") or "").strip() or not str(row.get("output") or "").strip() ) if invalid_count: raise InvalidStateError(f"task contains {invalid_count} invalid results") dataset_id = ( str(existing_dataset["id"]) if existing_dataset else new_id("dataset") ) now = utcnow() requested_split = payload.get("split") or { "train": 80, "validation": 10, "test": 10, } assignments = stable_split_assignments( [str(row["id"]) for row in rows], requested_split, seed=task_id, ) records = [ { "instruction": row["instruction"], "input": row["input"], "output": row["output"], "split": assignment, } for row, assignment in zip(rows, assignments, strict=True) ] split_order = ("train", "validation", "test") split_counts = { split_name: assignments.count(split_name) for split_name in split_order } file_specs: list[dict[str, Any]] = [] for split_name in split_order: split_records = [ (source_row, record) for source_row, record in zip(rows, records, strict=True) if record["split"] == split_name ] if not split_records: continue file_id = new_id("dfile") version_id = new_id("dfv") content = "".join( json_dumps(record) + "\n" for _, record in split_records ) raw = content.encode("utf-8") file_specs.append( { "split": split_name, "records": split_records, "file_id": file_id, "version_id": version_id, "content": content, "raw": raw, "checksum": hashlib.sha256(raw).hexdigest(), "storage_object_id": ( f"db://data-process/{task_id}/{file_id}/v1" ), } ) total_size = sum(len(spec["raw"]) for spec in file_specs) source_result_ids = [row["id"] for row in rows] metadata = { "source": "data_process", "storage_backend": "database", "source_task_id": task_id, "source_file_ids": [item["id"] for item in self._source_ids(conn, task_id)], "source_result_ids": source_result_ids, "format": payload.get("format") or "alpaca_jsonl", "split": requested_split, "split_counts": split_counts, } try: if existing_dataset: # 用户显式重新发布时，在同一数据集 ID 下修复旧的混合文件。 conn.execute( "DELETE FROM dataset_records WHERE dataset_id=%s", (dataset_id,) ) conn.execute( """DELETE FROM dataset_file_versions WHERE dataset_file_id IN (SELECT id FROM dataset_files WHERE dataset_id=%s)""", (dataset_id,), ) conn.execute("DELETE FROM dataset_files WHERE dataset_id=%s", (dataset_id,)) dataset = conn.execute( """ UPDATE datasets SET size=%s, size_bytes=%s, count=%s, record_count=%s, description=%s, metadata=%s, updated_at=%s WHERE id=%s RETURNING * """, ( f"{total_size} B", total_size, len(records), len(records), payload.get("description") or task.get("description") or "", json_dumps(metadata), now, dataset_id, ), ).fetchone() else: dataset = conn.execute( """ INSERT INTO datasets (id, name, type, storage_type, source, task_id, source_task_id, size, size_bytes, count, record_count, description, metadata, tenant_id, project_id, owner_id, created_by, create_time, created_at, updated_at) VALUES (%s, %s, %s, %s, 'task', %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) RETURNING * """, ( dataset_id, payload["dataset_name"], payload.get("dataset_type") or "train", payload.get("storage_type") or "local", task_id, task_id, f"{total_size} B", total_size, len(records), len(records), payload.get("description") or task.get("description") or "", json_dumps(metadata), task.get("tenant_id"), task.get("project_id"), task.get("owner_id"), payload.get("created_by") or task.get("created_by"), now, now, now, ), ).fetchone() for spec in file_specs: file_metadata = {**metadata, "file_split": spec["split"]} version = { "id": spec["version_id"], "version_no": 1, "version": 1, "description": f"data process {spec['split']} publish", "checksum_sha256": spec["checksum"], "size_bytes": len(spec["raw"]), "record_count": len(spec["records"]), "created_at": now, "create_time": now, "source_task_id": task_id, "storage_object_id": spec["storage_object_id"], } conn.execute( """ INSERT INTO dataset_files (id, dataset_id, name, storage_object_id, size, content, active_version_id, versions, create_time, current_version_id, size_bytes, record_count, file_format, checksum_sha256, version_no, source_task_id, tenant_id, project_id, created_by, metadata, created_at, updated_at) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, 1, %s, %s, %s, %s, %s, %s, %s) """, ( spec["file_id"], dataset_id, f"{payload['dataset_name']}.{spec['split']}.jsonl", spec["storage_object_id"], f"{len(spec['raw'])} B", spec["content"], spec["version_id"], json_dumps([version]), now, spec["version_id"], len(spec["raw"]), len(spec["records"]), "jsonl", spec["checksum"], task_id, task.get("tenant_id"), task.get("project_id"), payload.get("created_by") or task.get("created_by"), json_dumps(file_metadata), now, now, ), ) conn.execute( """ INSERT INTO dataset_file_versions (id, dataset_file_id, version_no, storage_object_id, content_preview, description, size_bytes, record_count, checksum_sha256, source_task_id, metadata, created_by, created_at) VALUES (%s, %s, 1, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """, ( spec["version_id"], spec["file_id"], spec["storage_object_id"], spec["content"][:2000], f"data process {spec['split']} publish", len(spec["raw"]), len(spec["records"]), spec["checksum"], task_id, json_dumps(file_metadata), payload.get("created_by") or task.get("created_by"), now, ), ) for line_number, (source_row, record) in enumerate( spec["records"], start=1 ): conn.execute( """ INSERT INTO dataset_records (id, dataset_id, dataset_file_id, version_id, line_no, split, instruction, input, output, raw, status, source_task_id, source_result_id, preview_item_id, created_at) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """, ( new_id("drec"), dataset_id, spec["file_id"], spec["version_id"], line_number, record["split"], record["instruction"], record["input"], record["output"], json_dumps( { **record, "source_task_id": task_id, "source_result_id": source_row["id"], "preview_item_id": source_row.get("preview_item_id"), } ), source_row["status"], task_id, source_row["id"], source_row.get("preview_item_id"), now, ), ) except psycopg.errors.UniqueViolation as exc: raise ConflictError("dataset name already exists") from exc conn.execute( """ UPDATE data_process_tasks SET output_dataset_id=%s, updated_at=%s, updated_by=%s WHERE id=%s """, (dataset_id, now, payload.get("created_by"), task_id), ) return { "dataset": _decode_row(dataset), "created": existing_dataset is None, "split_counts": split_counts, } @staticmethod def _source_ids( conn: psycopg.Connection[dict[str, Any]], task_id: str ) -> list[dict[str, Any]]: return conn.execute( """ SELECT id FROM data_process_source_files WHERE task_id=%s AND deleted_at IS NULL ORDER BY created_at, id """, (task_id,), ).fetchall() @lru_cache def get_data_process_store() -> DataProcessStore: return DataProcessStore()