import json import os import re import shutil import uuid from datetime import datetime from pathlib import Path from typing import List, Dict, Any, Set from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession from app.core.logging import get_logger from app.db.models.base_entity import LineageNode, LineageEdge from app.core.exception import BusinessError, ErrorCodes from app.module.cleaning.repository import ( CleaningTaskRepository, CleaningResultRepository, OperatorInstanceRepository, ) from app.module.cleaning.schema import ( CleaningTaskDto, CreateCleaningTaskRequest, CleaningResultDto, CleaningTaskLog, OperatorInstanceDto, CleaningProcess, CleaningTaskStatus, ) from app.module.cleaning.service.cleaning_task_validator import CleaningTaskValidator from app.module.cleaning.service.cleaning_task_scheduler import CleaningTaskScheduler from app.module.shared.common.lineage import LineageService from app.module.shared.schema.lineage import NodeType, EdgeType logger = get_logger(__name__) # Module-level sanitizers (CodeQL recognizes return values of module-level # functions as sanitized; instance-method calls are not tracked through) _TASK_ID_PATTERN = re.compile( r'^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$' ) def _sanitize_task_id(task_id: str) -> str: """Validate and return sanitized task_id for safe path construction.""" if not task_id: raise BusinessError(ErrorCodes.CLEANING_TASK_ID_REQUIRED) if not _TASK_ID_PATTERN.match(task_id): raise BusinessError( ErrorCodes.CLEANING_TASK_ID_REQUIRED, f"Invalid task_id format: {task_id}", ) return task_id def _sanitize_retry_count(retry_count: int) -> int: """Validate and return sanitized retry_count for safe path construction.""" if retry_count < 0 or retry_count > 1000: raise BusinessError( ErrorCodes.CLEANING_TASK_ID_REQUIRED, f"Invalid retry_count: {retry_count}", ) return retry_count DATASET_PATH = "/dataset" FLOW_PATH = "/flow" class CleaningTaskService: """Service for managing cleaning tasks""" def __init__( self, task_repo: CleaningTaskRepository, result_repo: CleaningResultRepository, operator_instance_repo: OperatorInstanceRepository, operator_service, scheduler: CleaningTaskScheduler, validator: CleaningTaskValidator, dataset_service, lineage_service: LineageService, ): self.task_repo = task_repo self.result_repo = result_repo self.operator_instance_repo = operator_instance_repo self.operator_service = operator_service self.scheduler = scheduler self.validator = validator self.dataset_service = dataset_service self.lineage_service = lineage_service async def get_tasks( self, db: AsyncSession, status: str | None = None, keyword: str | None = None, page: int | None = None, size: int | None = None, ) -> List[CleaningTaskDto]: """Get cleaning tasks""" tasks = await self.task_repo.find_tasks(db, status, keyword, page, size) if not tasks: return tasks # Batch query progress for all tasks in a single SQL (avoids N+1) task_ids = [task.id for task in tasks] progress_map = await self.result_repo.batch_count_by_instance_ids(db, task_ids) for task in tasks: completed, failed, actual_total = progress_map.get(task.id, (0, 0, 0)) total = max(actual_total, task.file_count or 0) task.progress = CleaningProcess.of(total, completed, failed) return tasks async def _set_process(self, db: AsyncSession, task: CleaningTaskDto) -> None: """Set task progress using actual results from database""" completed, failed = await self.result_repo.count_by_instance_id(db, task.id) # Use actual total from database (t_clean_result table), fallback to task.file_count actual_total = await self.result_repo.count_total_by_instance_id(db, task.id) total = max(actual_total, task.file_count or 0) task.progress = CleaningProcess.of(total, completed, failed) async def count_tasks( self, db: AsyncSession, status: str | None = None, keyword: str | None = None, ) -> int: """Count cleaning tasks using SQL COUNT""" return await self.task_repo.count_tasks(db, status, keyword) async def get_task(self, db: AsyncSession, task_id: str) -> CleaningTaskDto: """Get task by ID""" task = await self.task_repo.find_task_by_id(db, task_id) if not task: raise BusinessError(ErrorCodes.CLEANING_TASK_NOT_FOUND, task_id) await self._set_process(db, task) instances = await self.operator_instance_repo.find_operator_by_instance_id( db, task_id ) # Batch query operators all_operators = await self.operator_service.get_operators( db=db, page=0, size=1000, categories=[], keyword=None, is_star=None ) operator_map = {op.id: op for op in all_operators} task.instance = [] for inst in instances: operator = operator_map.get(inst.operator_id) if operator: task.instance.append( OperatorInstanceDto( id=operator.id, name=operator.name, description=operator.description, inputs=operator.inputs, outputs=operator.outputs, settings=operator.settings, categories=operator.categories, ) ) else: task.instance.append(OperatorInstanceDto(id=inst.operator_id)) return task async def create_task( self, db: AsyncSession, request: CreateCleaningTaskRequest ) -> CleaningTaskDto: """Create new cleaning task""" if request.instance and request.template_id: instances = await self.get_instance_by_template_id(db, request.template_id) request.instance = instances await self.validator.check_task_name_duplication(db, request.name) self.validator.check_input_and_output(request.instance) executor_type = self.validator.check_and_get_executor_type(request.instance) task_id = str(uuid.uuid4()) dest_dataset_id = request.dest_dataset_id dest_dataset_name = request.dest_dataset_name if not dest_dataset_id: logger.info( f"Creating new dataset: {dest_dataset_name}, type: {request.dest_dataset_type}" ) dest_dataset_response = await self.dataset_service.create_dataset( name=dest_dataset_name, dataset_type=request.dest_dataset_type, description="", status="ACTIVE", ) dest_dataset_id = dest_dataset_response.id logger.info(f"Successfully created dataset: {dest_dataset_id}") else: logger.info(f"Using existing dataset: {dest_dataset_id}") dest_dataset_response = await self.dataset_service.get_dataset( dest_dataset_id ) src_dataset = await self.dataset_service.get_dataset(request.src_dataset_id) if not src_dataset: raise BusinessError( ErrorCodes.CLEANING_DATASET_NOT_FOUND, request.src_dataset_id ) task_dto = CleaningTaskDto( id=task_id, name=request.name, description=request.description, status=CleaningTaskStatus.PENDING, src_dataset_id=request.src_dataset_id, src_dataset_name=request.src_dataset_name, dest_dataset_id=dest_dataset_id, dest_dataset_name=dest_dataset_name, before_size=src_dataset.totalSize, file_count=src_dataset.fileCount, retry_count=-1, ) await self.task_repo.insert_task(db, task_dto) await self._add_cleaning_to_graph(src_dataset, task_dto, dest_dataset_response) await self.operator_instance_repo.insert_instance(db, task_id, request.instance) # Increment operator usage count operator_ids = [inst.id for inst in request.instance if inst.id] if operator_ids: await self.operator_service.increment_usage_count(operator_ids, db) all_operators = await self.operator_service.get_operators( db=db, page=0, size=1000, categories=[], keyword=None, is_star=None ) operator_map = {op.id: op for op in all_operators} await self.prepare_task( dest_dataset_id, task_id, request.instance, operator_map, executor_type ) return await self.get_task(db, task_id) async def _add_cleaning_to_graph( self, src_dataset, task: CleaningTaskDto, dest_dataset, ) -> None: """ 添加清洗任务到血缘图 """ from_node = LineageNode( id=src_dataset.id, node_type=NodeType.DATASET.value, name=src_dataset.name, description=src_dataset.description or "", ) to_node = LineageNode( id=dest_dataset.id, node_type=NodeType.DATASET.value, name=dest_dataset.name, description=dest_dataset.description or "", ) edge = LineageEdge( process_id=task.id, name=task.name or "", description=task.description or "", edge_type=EdgeType.DATA_CLEANING.value, from_node_id=from_node.id, to_node_id=to_node.id, ) await self.lineage_service.generate_graph(from_node, edge, to_node) async def prepare_task( self, dataset_id: str, task_id: str, instances: List[OperatorInstanceDto], operator_map: dict, executor_type: str, ) -> None: """Prepare task configuration file""" process_config = { "dataset_id": dataset_id, "instance_id": task_id, "dataset_path": f"{FLOW_PATH}/{task_id}/dataset.jsonl", "export_path": f"{DATASET_PATH}/{dataset_id}", "executor_type": executor_type, "process": [], } for instance in instances: operator = operator_map.get(instance.id) if not operator: continue operator_config = self._get_default_values(operator) operator_config.update(instance.overrides) runtime_config = self._get_runtime_config(operator) operator_config.update(runtime_config) process_config["process"].append({instance.id: operator_config}) config_file_path = Path(f"{FLOW_PATH}/{task_id}/process.yaml") config_file_path.parent.mkdir(parents=True, exist_ok=True) import yaml try: with open(config_file_path, "w", encoding="utf-8") as f: yaml.dump( process_config, f, default_flow_style=False, allow_unicode=True ) except Exception as e: logger.error(f"Failed to write process.yaml: {e}") raise BusinessError(ErrorCodes.CLEANING_FILE_SYSTEM_ERROR, str(e)) def _get_default_values(self, operator) -> Dict[str, Any]: """Get default values from operator settings""" if not operator.settings: return {} try: settings = json.loads(operator.settings) defaults = {} for key, value in settings.items(): setting_type = value.get("type") if "defaultVal" in value: defaults[key] = value["defaultVal"] return defaults except json.JSONDecodeError as e: logger.error(f"Failed to parse settings: {e}") return {} def _get_runtime_config(self, operator) -> Dict[str, Any]: """Get runtime configuration from operator""" if not operator.runtime: return {} try: return json.loads(operator.runtime) except json.JSONDecodeError as e: logger.error(f"Failed to parse runtime config: {e}") return {} async def scan_dataset( self, db: AsyncSession, task_id: str, src_dataset_id: str, succeed_files: Set[str] | None = None, ) -> None: """Scan source dataset and create dataset.jsonl""" target_file_path = Path(f"{FLOW_PATH}/{task_id}/dataset.jsonl") target_file_path.parent.mkdir(parents=True, exist_ok=True) query = text(""" SELECT id, file_name, file_path, file_type, file_size, metadata FROM t_dm_dataset_files WHERE dataset_id = :dataset_id ORDER BY created_at """) result = await db.execute(query, {"dataset_id": src_dataset_id}) files = result.fetchall() with open(target_file_path, "w", encoding="utf-8") as f: for file in files: if succeed_files and file.id in succeed_files: continue metadata_dict = {} if file.metadata: try: parsed = json.loads(file.metadata) if isinstance(parsed, dict): metadata_dict = parsed except (json.JSONDecodeError, TypeError): pass file_info = { "fileId": file.id, "fileName": file.file_name, "filePath": file.file_path, "fileType": file.file_type, "fileSize": file.file_size, "metadata": metadata_dict, } f.write(json.dumps(file_info, ensure_ascii=False) + "\n") async def get_task_results( self, db: AsyncSession, task_id: str ) -> List[CleaningResultDto]: """Get task results""" return await self.result_repo.find_by_instance_id(db, task_id) async def get_task_log( self, db: AsyncSession, task_id: str, retry_count: int ) -> List[CleaningTaskLog]: """Get task log""" safe_task_id = _sanitize_task_id(task_id) safe_retry_count = _sanitize_retry_count(retry_count) flow_root = Path(FLOW_PATH).resolve() log_path = flow_root / safe_task_id / "output.log" if safe_retry_count > 0: log_path = flow_root / safe_task_id / f"output.log.{safe_retry_count}" # 防止路径穿越:规范化后校验仍在 FLOW_PATH 下 resolved_log_path = log_path.resolve() if flow_root not in resolved_log_path.parents: logger.warning(f"Path traversal attempt detected: task_id={task_id}") return [] if not resolved_log_path.exists(): return [] logs = [] last_level = "INFO" standard_level_pattern = re.compile( r"\b(DEBUG|Debug|INFO|Info|WARN|Warn|WARNING|Warning|ERROR|Error|FATAL|Fatal)\b" ) exception_suffix_pattern = re.compile(r"\b\w+(Warning|Error|Exception)\b") with open(resolved_log_path, "r", encoding="utf-8") as f: for line in f: last_level = self._get_log_level( line, last_level, standard_level_pattern, exception_suffix_pattern ) logs.append(CleaningTaskLog(level=last_level, message=line.rstrip())) return logs def _get_log_level( self, line: str, default_level: str, std_pattern, ex_pattern ) -> str: """Extract log level from log line""" if not line or not line.strip(): return default_level std_match = std_pattern.search(line) if std_match: return std_match.group(1).upper() ex_match = ex_pattern.search(line) if ex_match: match = ex_match.group(1).upper() if match == "WARNING": return "WARN" if match in ["ERROR", "EXCEPTION"]: return "ERROR" return default_level async def delete_task(self, db: AsyncSession, task_id: str) -> None: """Delete task""" safe_task_id = _sanitize_task_id(task_id) task = await self.task_repo.find_task_by_id(db, safe_task_id) if not task: raise BusinessError(ErrorCodes.CLEANING_TASK_NOT_FOUND, task_id) # 运行中的任务无法删除 if task.status == CleaningTaskStatus.RUNNING: raise BusinessError( ErrorCodes.CLEANING_TASK_STATUS_INVALID, "Task is running, cannot be deleted. Please stop the task first." ) await self.task_repo.delete_task_by_id(db, safe_task_id) await self.operator_instance_repo.delete_by_instance_id(db, safe_task_id) await self.result_repo.delete_by_instance_id(db, safe_task_id) # 删除任务相关文件 flow_root = Path(FLOW_PATH).resolve() task_path = (flow_root / safe_task_id).resolve() # 防止路径穿越:parents 校验目标路径仍在 flow_root 下 if flow_root not in task_path.parents: logger.warning(f"Path traversal attempt in delete_task: task_id={task_id}") raise BusinessError( ErrorCodes.CLEANING_TASK_NOT_FOUND, f"Invalid task_id: {task_id}", ) if task_path.exists(): try: shutil.rmtree(task_path) except Exception as e: logger.warning(f"Failed to delete task path {task_id}: {e}") async def execute_task(self, db: AsyncSession, task_id: str) -> bool: """Execute task""" succeeded = await self.result_repo.find_by_instance_id(db, task_id, "COMPLETED") succeed_set = {res.src_file_id for res in succeeded} task = await self.task_repo.find_task_by_id(db, task_id) if not task: raise BusinessError(ErrorCodes.CLEANING_TASK_NOT_FOUND, task_id) src_dataset = await self.dataset_service.get_dataset(task.src_dataset_id) if src_dataset: task.before_size = src_dataset.totalSize task.file_count = src_dataset.fileCount await self.task_repo.update_task(db, task) await self.scan_dataset(db, task_id, task.src_dataset_id, succeed_set) await self.result_repo.delete_by_instance_id(db, task_id, "FAILED") try: return await self.scheduler.execute_task( db, task_id, (task.retry_count or 0) + 1 ) except Exception as e: logger.error( f"execute_task failed, task_id={task_id}, retry_count={(task.retry_count or 0) + 1}: {e}" ) task = CleaningTaskDto() task.id = task_id task.status = CleaningTaskStatus.FAILED task.finished_at = datetime.now() await self.task_repo.update_task(db, task) return False async def stop_task(self, db: AsyncSession, task_id: str) -> bool: """Stop task""" return await self.scheduler.stop_task(db, task_id) async def get_instance_by_template_id( self, db: AsyncSession, template_id: str ) -> List[OperatorInstanceDto]: """Get instances by template ID (delegated to template service)""" instances = await self.operator_instance_repo.find_operator_by_instance_id( db, template_id ) # Batch query operators all_operators = await self.operator_service.get_operators( db=db, page=0, size=1000, categories=[], keyword=None, is_star=None ) operator_map = {op.id: op for op in all_operators} result = [] for inst in instances: operator = operator_map.get(inst.operator_id) if operator: operator_dto = OperatorInstanceDto( id=operator.id, name=operator.name, description=operator.description, inputs=operator.inputs, outputs=operator.outputs, settings=operator.settings, categories=operator.categories, ) if inst.settings_override: try: operator_dto.overrides = json.loads(inst.settings_override) except json.JSONDecodeError as e: logger.error( f"Failed to parse settings for {inst.operator_id}: {e}" ) result.append(operator_dto) return result