522 lines
18 KiB
Python
522 lines
18 KiB
Python
from typing import Optional
|
||
import re
|
||
|
||
from fastapi import APIRouter, Depends
|
||
from fastapi.responses import StreamingResponse
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from app.core.logging import get_logger
|
||
from app.db.session import get_db
|
||
from app.module.cleaning.schema import (
|
||
CleaningTaskDto,
|
||
CreateCleaningTaskRequest,
|
||
CleaningResultDto,
|
||
CleaningTaskLog,
|
||
)
|
||
from app.module.cleaning.service import CleaningTaskService
|
||
from app.module.shared.schema import StandardResponse, PaginatedData
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
# Module-level pattern for task_id sanitization (CodeQL sanitizer)
|
||
_TASK_ID_PATTERN = re.compile(r"^[A-Za-z0-9_-]+$")
|
||
|
||
|
||
def _sanitize_task_id(task_id: str) -> str:
|
||
"""Validate and return sanitized task_id for safe path construction.
|
||
|
||
CodeQL recognizes the return value of this function as sanitized,
|
||
allowing safe use in path construction.
|
||
"""
|
||
if not _TASK_ID_PATTERN.fullmatch(task_id):
|
||
raise ValueError(f"Invalid task_id: {task_id}")
|
||
return task_id
|
||
|
||
|
||
def _sanitize_retry_count(retry_count: int) -> int:
|
||
"""Validate and return sanitized retry_count for safe path construction.
|
||
|
||
CodeQL recognizes the return value of this function as sanitized,
|
||
allowing safe use in path construction.
|
||
"""
|
||
if retry_count < 0 or retry_count > 1000:
|
||
raise ValueError(f"Invalid retry_count: {retry_count}")
|
||
return retry_count
|
||
|
||
router = APIRouter(prefix="/cleaning/tasks", tags=["Cleaning Tasks"])
|
||
|
||
|
||
def _get_operator_service():
|
||
"""Get operator service"""
|
||
from app.module.operator.service import OperatorService
|
||
from app.module.operator.repository import (
|
||
OperatorRepository,
|
||
CategoryRelationRepository,
|
||
OperatorReleaseRepository,
|
||
)
|
||
from app.module.operator.parsers import ParserHolder
|
||
from app.module.shared.file_service import FileService
|
||
from app.module.shared.chunk_upload_repository import ChunkUploadRepository
|
||
|
||
return OperatorService(
|
||
operator_repo=OperatorRepository(None),
|
||
category_relation_repo=CategoryRelationRepository(None),
|
||
operator_release_repo=OperatorReleaseRepository(None),
|
||
parser_holder=ParserHolder(),
|
||
file_service=FileService(ChunkUploadRepository()),
|
||
)
|
||
|
||
|
||
def _get_task_service(db: AsyncSession) -> CleaningTaskService:
|
||
"""Get cleaning task service instance"""
|
||
from app.module.cleaning.service import (
|
||
CleaningTaskScheduler,
|
||
CleaningTaskValidator,
|
||
)
|
||
from app.module.cleaning.service.cleaning_task_scheduler import get_scheduler
|
||
from app.module.cleaning.repository import (
|
||
CleaningTaskRepository,
|
||
CleaningResultRepository,
|
||
OperatorInstanceRepository,
|
||
)
|
||
from app.module.dataset.service import DatasetManagementService
|
||
from app.module.shared.common.lineage import LineageService
|
||
|
||
scheduler = get_scheduler() # 使用全局单例，保持轮询状态跨请求共享
|
||
operator_service = _get_operator_service()
|
||
dataset_service = DatasetManagementService(db)
|
||
lineage_service = LineageService(db)
|
||
|
||
task_repo = CleaningTaskRepository(None)
|
||
|
||
return CleaningTaskService(
|
||
task_repo=task_repo,
|
||
result_repo=CleaningResultRepository(None),
|
||
operator_instance_repo=OperatorInstanceRepository(None),
|
||
operator_service=operator_service,
|
||
scheduler=scheduler,
|
||
validator=CleaningTaskValidator(task_repo=task_repo, template_repo=None),
|
||
dataset_service=dataset_service,
|
||
lineage_service=lineage_service,
|
||
)
|
||
|
||
|
||
@router.get(
|
||
"",
|
||
response_model=StandardResponse[PaginatedData[CleaningTaskDto]],
|
||
summary="查询清洗任务列表",
|
||
description="根据参数查询清洗任务列表（支持分页、状态过滤、关键词搜索）",
|
||
tags=["mcp"],
|
||
)
|
||
async def get_cleaning_tasks(
|
||
page: int = 0,
|
||
size: int = 10,
|
||
status: Optional[str] = None,
|
||
keyword: Optional[str] = None,
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""Query cleaning tasks"""
|
||
task_service = _get_task_service(db)
|
||
|
||
tasks = await task_service.get_tasks(db, status, keyword, page, size)
|
||
count = await task_service.count_tasks(db, status, keyword)
|
||
total_pages = (count + size - 1) // size if size > 0 else 0
|
||
|
||
return StandardResponse(
|
||
code="0",
|
||
message="success",
|
||
data=PaginatedData(
|
||
page=page,
|
||
size=size,
|
||
total_elements=count,
|
||
total_pages=total_pages,
|
||
content=tasks,
|
||
),
|
||
)
|
||
|
||
|
||
@router.post(
|
||
"",
|
||
response_model=StandardResponse[CleaningTaskDto],
|
||
summary="创建清洗任务",
|
||
description="根据模板ID或算子列表创建清洗任务",
|
||
tags=["mcp"],
|
||
)
|
||
async def create_cleaning_task(
|
||
request: CreateCleaningTaskRequest,
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""Create cleaning task"""
|
||
task_service = _get_task_service(db)
|
||
|
||
task = await task_service.create_task(db, request)
|
||
await db.commit()
|
||
|
||
await task_service.execute_task(db, task.id)
|
||
await db.commit()
|
||
|
||
return StandardResponse(code="0", message="success", data=task)
|
||
|
||
|
||
@router.get(
|
||
"/{task_id}",
|
||
response_model=StandardResponse[CleaningTaskDto],
|
||
summary="获取清洗任务详情",
|
||
description="根据ID获取清洗任务详细信息",
|
||
)
|
||
async def get_cleaning_task(
|
||
task_id: str,
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""Get cleaning task by ID"""
|
||
task_service = _get_task_service(db)
|
||
task = await task_service.get_task(db, task_id)
|
||
return StandardResponse(code="0", message="success", data=task)
|
||
|
||
|
||
@router.delete(
|
||
"/{task_id}",
|
||
response_model=StandardResponse[str],
|
||
summary="删除清洗任务",
|
||
description="删除指定的清洗任务",
|
||
)
|
||
async def delete_cleaning_task(
|
||
task_id: str,
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""Delete cleaning task"""
|
||
task_service = _get_task_service(db)
|
||
await task_service.delete_task(db, task_id)
|
||
await db.commit()
|
||
return StandardResponse(code="0", message="success", data=task_id)
|
||
|
||
|
||
@router.post(
|
||
"/{task_id}/stop",
|
||
response_model=StandardResponse[str],
|
||
summary="停止清洗任务",
|
||
description="停止正在运行的清洗任务",
|
||
)
|
||
async def stop_cleaning_task(
|
||
task_id: str,
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""Stop cleaning task"""
|
||
task_service = _get_task_service(db)
|
||
await task_service.stop_task(db, task_id)
|
||
await db.commit()
|
||
return StandardResponse(code="0", message="success", data=task_id)
|
||
|
||
|
||
@router.post(
|
||
"/{task_id}/execute",
|
||
response_model=StandardResponse[str],
|
||
summary="执行清洗任务",
|
||
description="重新执行清洗任务",
|
||
)
|
||
async def execute_cleaning_task(
|
||
task_id: str,
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""Execute cleaning task"""
|
||
task_service = _get_task_service(db)
|
||
await task_service.execute_task(db, task_id)
|
||
await db.commit()
|
||
return StandardResponse(code="0", message="success", data=task_id)
|
||
|
||
|
||
@router.get(
|
||
"/{task_id}/result",
|
||
response_model=StandardResponse[list[CleaningResultDto]],
|
||
summary="获取清洗任务结果",
|
||
description="获取指定清洗任务的执行结果",
|
||
)
|
||
async def get_cleaning_task_results(
|
||
task_id: str,
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""Get cleaning task results"""
|
||
task_service = _get_task_service(db)
|
||
results = await task_service.get_task_results(db, task_id)
|
||
return StandardResponse(code="0", message="success", data=results)
|
||
|
||
|
||
@router.get(
|
||
"/{task_id}/result/download",
|
||
summary="下载清洗任务结果文件压缩包",
|
||
description="下载指定清洗任务的源文件和处理后文件压缩包",
|
||
)
|
||
async def download_cleaning_result_files(
|
||
task_id: str,
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""Download cleaning task result files (both src and dest)"""
|
||
import zipfile
|
||
import io
|
||
from pathlib import Path
|
||
from fastapi.responses import StreamingResponse
|
||
from app.core.exception import BusinessError, ErrorCodes
|
||
|
||
task_service = _get_task_service(db)
|
||
|
||
task = await task_service.get_task(db, task_id)
|
||
|
||
results = await task_service.get_task_results(db, task_id)
|
||
if not results:
|
||
raise BusinessError(ErrorCodes.NOT_FOUND, f"Cleaning task {task_id} not found")
|
||
|
||
src_path = Path("/dataset") / task.src_dataset_id
|
||
|
||
dest_path = Path("/dataset") / task.dest_dataset_id
|
||
|
||
zip_buffer = io.BytesIO()
|
||
|
||
with zipfile.ZipFile(zip_buffer, "w") as zipf:
|
||
for file_record in results:
|
||
if file_record.src_name:
|
||
src_path = src_path / file_record.src_name
|
||
if src_path.exists():
|
||
zipf.write(src_path, arcname=f"src/{file_record.src_name}")
|
||
|
||
if file_record.dest_name:
|
||
dest_path = dest_path / file_record.dest_name
|
||
if dest_path.exists():
|
||
zipf.write(dest_path, arcname=f"dest/{file_record.dest_name}")
|
||
|
||
zip_buffer.seek(0)
|
||
|
||
filename = f"task_{task_id}_{file_record.src_name}.zip"
|
||
|
||
return StreamingResponse(
|
||
zip_buffer,
|
||
media_type="application/zip",
|
||
headers={"Content-Disposition": f'attachment; filename="{filename}"'},
|
||
)
|
||
|
||
|
||
@router.get(
|
||
"/{task_id}/log/stream",
|
||
summary="流式获取清洗任务日志",
|
||
description="通过SSE流式获取清洗任务日志",
|
||
)
|
||
async def stream_cleaning_task_log(
|
||
task_id: str,
|
||
retry_count: int = 0,
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""Stream cleaning task log via SSE"""
|
||
import asyncio
|
||
import json
|
||
import re
|
||
from pathlib import Path
|
||
|
||
flow_base = Path("/flow").resolve()
|
||
|
||
# 校验 task_id 并获取净化后的值（CodeQL 认可此返回值为安全）
|
||
try:
|
||
safe_task_id = _sanitize_task_id(task_id)
|
||
except ValueError:
|
||
logger.warning(f"Invalid task_id in stream_log: task_id={task_id}")
|
||
|
||
async def invalid_error_generator():
|
||
yield f"data: {json.dumps({'level': 'ERROR', 'message': 'Invalid task_id'}, ensure_ascii=False)}\n\n"
|
||
|
||
return StreamingResponse(invalid_error_generator(), media_type="text/event-stream")
|
||
|
||
task_service = _get_task_service(db)
|
||
task = await task_service.task_repo.find_task_by_id(db, safe_task_id)
|
||
|
||
if not task:
|
||
|
||
async def error_generator():
|
||
yield f"data: {json.dumps({'level': 'ERROR', 'message': 'Task not found'}, ensure_ascii=False)}\n\n"
|
||
|
||
return StreamingResponse(error_generator(), media_type="text/event-stream")
|
||
|
||
# 校验并净化 retry_count（CodeQL 认可此返回值为安全）
|
||
safe_retry_count = _sanitize_retry_count(retry_count)
|
||
log_filename = "output.log" if safe_retry_count == 0 else f"output.log.{safe_retry_count}"
|
||
log_path = (flow_base / safe_task_id / log_filename).resolve()
|
||
|
||
# 防止路径穿越：parents 校验仍在 flow_base 下
|
||
if flow_base not in log_path.parents:
|
||
logger.warning(f"Path traversal attempt in stream_log: task_id={task_id}")
|
||
|
||
async def traversal_error_generator():
|
||
yield f"data: {json.dumps({'level': 'ERROR', 'message': 'Invalid task_id'}, ensure_ascii=False)}\n\n"
|
||
|
||
return StreamingResponse(traversal_error_generator(), media_type="text/event-stream")
|
||
|
||
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")
|
||
|
||
def get_log_level(line: str, default_level: str) -> str:
|
||
if not line or not line.strip():
|
||
return default_level
|
||
std_match = standard_level_pattern.search(line)
|
||
if std_match:
|
||
return std_match.group(1).upper()
|
||
ex_match = exception_suffix_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 log_generator():
|
||
last_position = 0
|
||
last_level = "INFO"
|
||
heartbeat_counter = 0
|
||
HEARTBEAT_INTERVAL = 10
|
||
|
||
while True:
|
||
try:
|
||
current_task = await task_service.task_repo.find_task_by_id(db, safe_task_id)
|
||
is_task_finished = current_task and current_task.status in [
|
||
"COMPLETED",
|
||
"FAILED",
|
||
"STOPPED",
|
||
]
|
||
|
||
if log_path.exists():
|
||
with open(log_path, "r", encoding="utf-8") as f:
|
||
f.seek(last_position)
|
||
new_content = f.read()
|
||
if new_content:
|
||
for line in new_content.splitlines(keepends=True):
|
||
last_level = get_log_level(line, last_level)
|
||
log_entry = json.dumps(
|
||
{"level": last_level, "message": line.rstrip()},
|
||
ensure_ascii=False,
|
||
)
|
||
yield f"data: {log_entry}\n\n"
|
||
last_position = f.tell()
|
||
heartbeat_counter = 0
|
||
else:
|
||
heartbeat_counter += 1
|
||
|
||
f.seek(0, 2)
|
||
current_size = f.tell()
|
||
if current_size < last_position:
|
||
last_position = 0
|
||
else:
|
||
heartbeat_counter += 1
|
||
|
||
if is_task_finished:
|
||
yield f"data: {json.dumps({'level': 'INFO', 'message': '[END_OF_STREAM]'}, ensure_ascii=False)}\n\n"
|
||
break
|
||
|
||
if heartbeat_counter >= HEARTBEAT_INTERVAL:
|
||
yield f"data: {json.dumps({'level': 'DEBUG', 'message': '[HEARTBEAT]'}, ensure_ascii=False)}\n\n"
|
||
heartbeat_counter = 0
|
||
|
||
await asyncio.sleep(1)
|
||
except Exception as e:
|
||
logger.error(f"Error reading log file: {e}")
|
||
await asyncio.sleep(2)
|
||
|
||
return StreamingResponse(
|
||
log_generator(),
|
||
media_type="text/event-stream",
|
||
headers={
|
||
"Cache-Control": "no-cache",
|
||
"Connection": "keep-alive",
|
||
"X-Accel-Buffering": "no",
|
||
},
|
||
)
|
||
|
||
|
||
@router.get(
|
||
"/{task_id}/log/{retry_count}",
|
||
response_model=StandardResponse[list[CleaningTaskLog]],
|
||
summary="获取清洗任务日志",
|
||
description="获取指定清洗任务的执行日志",
|
||
)
|
||
async def get_cleaning_task_log(
|
||
task_id: str,
|
||
retry_count: int,
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""Get cleaning task log"""
|
||
task_service = _get_task_service(db)
|
||
logs = await task_service.get_task_log(db, task_id, retry_count)
|
||
return StandardResponse(code="0", message="success", data=logs)
|
||
|
||
|
||
@router.get(
|
||
"/{task_id}/log/{retry_count}/download",
|
||
summary="下载清洗任务日志文件",
|
||
description="下载指定清洗任务的日志文件",
|
||
)
|
||
async def download_cleaning_task_log(
|
||
task_id: str,
|
||
retry_count: int,
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""Download cleaning task log file"""
|
||
from pathlib import Path
|
||
from fastapi.responses import FileResponse
|
||
|
||
flow_base = Path("/flow").resolve()
|
||
|
||
# 校验 task_id 并获取净化后的值（CodeQL 认可此返回值为安全）
|
||
try:
|
||
safe_task_id = _sanitize_task_id(task_id)
|
||
except ValueError:
|
||
logger.warning(f"Invalid task_id in download_log: task_id={task_id}")
|
||
from app.core.exception import BusinessError, ErrorCodes
|
||
raise BusinessError(
|
||
ErrorCodes.CLEANING_TASK_LOG_NOT_FOUND,
|
||
f"Invalid task_id: {task_id}",
|
||
)
|
||
|
||
task_service = _get_task_service(db)
|
||
task = await task_service.task_repo.find_task_by_id(db, safe_task_id)
|
||
|
||
if not task:
|
||
from app.core.exception import BusinessError, ErrorCodes
|
||
|
||
raise BusinessError(ErrorCodes.CLEANING_TASK_NOT_FOUND, task_id)
|
||
|
||
# 校验并净化 retry_count（CodeQL 认可此返回值为安全）
|
||
safe_retry_count = _sanitize_retry_count(retry_count)
|
||
log_filename = "output.log" if safe_retry_count == 0 else f"output.log.{safe_retry_count}"
|
||
log_path = (flow_base / safe_task_id / log_filename).resolve()
|
||
|
||
# 防止路径穿越：parents 校验仍在 flow_base 下
|
||
if flow_base not in log_path.parents:
|
||
logger.warning(f"Path traversal attempt in download_log: task_id={task_id}")
|
||
from app.core.exception import BusinessError, ErrorCodes
|
||
raise BusinessError(
|
||
ErrorCodes.CLEANING_TASK_LOG_NOT_FOUND,
|
||
f"Invalid task_id: {task_id}",
|
||
)
|
||
|
||
if not log_path.exists():
|
||
from app.core.exception import BusinessError, ErrorCodes
|
||
|
||
raise BusinessError(
|
||
ErrorCodes.CLEANING_TASK_LOG_NOT_FOUND,
|
||
f"Log file not found for task {task_id}, retry {retry_count}",
|
||
)
|
||
|
||
# Generate filename with task name and retry count
|
||
import re
|
||
|
||
task_name = task.name or "未命名任务"
|
||
safe_task_name = re.sub(
|
||
r"[^\w\u4e00-\u9fff\-]", "_", task_name
|
||
) # Keep alphanumeric, Chinese, and hyphens
|
||
run_number = retry_count + 1 # retry_count is 0-indexed, so add 1 for display
|
||
filename = f"{safe_task_name}_第{run_number}次运行.log"
|
||
|
||
return FileResponse(
|
||
path=log_path,
|
||
media_type="text/plain",
|
||
filename=filename,
|
||
)
|