deepdata: core product code (backend, frontend, runtime, deepflow module)

This commit is contained in:
figmar
2026-08-30 19:23:38 +08:00
commit bd3db30be9
1405 changed files with 271727 additions and 0 deletions
+37
View File
@@ -0,0 +1,37 @@
# Dev settings
HOST=0.0.0.0
PORT=18000
DEBUG=true
LOG_LEVEL=DEBUG
LOG_FILE_DIR=./logs
LOG_ROTATION_MAX_SIZE=100MB
LOG_ROTATION_BACKUP_COUNT=30
LOG_PVC_MONITOR_ENABLED=true
LOG_PVC_MONITOR_PATH=./logs
LOG_PVC_MONITOR_THRESHOLD=0.9
LOG_PVC_MONITOR_INTERVAL_SECONDS=300
LOG_PVC_MONITOR_DELETE_BATCH_SIZE=10
LOG_PVC_MONITOR_FILE_SUFFIXES=log,out,err
RAG_STORAGE_DIR=./rag_storage
# DataBase
PGSQL_HOST=localhost
PGSQL_PORT=5432
PGSQL_USER=postgres
PGSQL_PASSWORD=password
PGSQL_DATABASE=deepdata
# Label Studio settings
LABEL_STUDIO_BASE_URL=http://localhost:30001
LABEL_STUDIO_USER_TOKEN="demo_dev_token"
DEEPDATA_JWT_ENABLE=false
# Milvus settings (Vector Database)
# Development: use localhost
MILVUS_URI=http://localhost:19530
MILVUS_TOKEN=
# Production: use service name
# MILVUS_URI=http://milvus:19530
+8
View File
@@ -0,0 +1,8 @@
# Local Development Environment Files
.env
.dev.env
# logs
logs/
doc/
+107
View File
@@ -0,0 +1,107 @@
# DeepData Python Service (DeepData)
这是 DeepData 的 Python 服务,负责DeepData的数据合成、数据标注、数据评估等功能。
## 简要说明
- 框架:FastAPI
- 异步数据库/ORM:SQLAlchemy (async)
- 数据库迁移:Alembic
- 运行器:uvicorn
## 快速开始(开发)
### 前置条件
- Python 3.11+
- poetry 包管理器
1. 克隆仓库
```bash
git clone https://gitea.apescale.com/figmar/deepdata.git
```
2. 进入项目目录
```bash
cd runtime/deepdata-python
```
3. 安装依赖
由于项目使用poetry管理依赖,你可以使用以下命令安装::
git clone https://gitea.apescale.com/figmar/deepdata.git
poetry install
```
或者直接使用pip安装(如果poetry不可用):
git clone https://gitea.apescale.com/figmar/deepdata.git
pip install -e .
```
4. 配置环境变量
复制环境变量示例文件并配置:
git clone https://gitea.apescale.com/figmar/deepdata.git
cp .env.example .env
```
编辑.env文件,设置必要的环境变量,如数据库连接、Label Studio配置等。
5. 数据库迁移(开发环境):
git clone https://gitea.apescale.com/figmar/deepdata.git
alembic upgrade head
```
6. 启动开发服务器(示例与常用参数):
- 本地开发(默认 host/port,自动重载):
git clone https://gitea.apescale.com/figmar/deepdata.git
set -a && source .env && set +a && poetry run uvicorn app.main:app --port 18000 --log-level debug --reload
```
或者
git clone https://gitea.apescale.com/figmar/deepdata.git
poetry run python -m app.main
```
- 指定主机与端口并打开调试日志:
git clone https://gitea.apescale.com/figmar/deepdata.git
uvicorn app.main:app --host 0.0.0.0 --port 8000 --reload --log-level debug
```
- 在生产环境使用多个 worker(不使用 --reload):
git clone https://gitea.apescale.com/figmar/deepdata.git
uvicorn app.main:app --host 0.0.0.0 --port 8000 --workers 4 --log-level info --proxy-headers
```
- 使用环境变量启动(示例):
git clone https://gitea.apescale.com/figmar/deepdata.git
HOST=0.0.0.0 PORT=8000 uvicorn app.main:app --reload
```
注意:
- `--reload` 仅用于开发,会监视文件变化并重启进程;不要在生产中使用。
- `--workers` 提供并发处理能力,但会增加内存占用;生产时通常配合进程管理或容器编排(Kubernetes)使用。
- 若需要完整的生产部署建议使用 ASGI 服务器(如 gunicorn + uvicorn workers / 或直接使用 uvicorn 在容器中配合进程管理)。
访问 API 文档:
- Swagger UI: http://127.0.0.1:8000/docs
- ReDoc: http://127.0.0.1:8000/redoc (推荐使用)
## 开发新功能
- 安装开发依赖:
git clone https://gitea.apescale.com/figmar/deepdata.git
poetry add xxx
```
## 使用(简要)
- 所有 API 路径以 `/api` 前缀注册(见 `app/main.py` 中 `app.include_router(api_router, prefix="/api")`)。
- 根路径 `/` 返回服务信息和文档链接。
更多细节请查看 `doc/usage.md`(接口使用)和 `doc/development.md`(开发说明)。
@@ -0,0 +1,7 @@
# app/core/__init__.py
"""
Core module
"""
+120
View File
@@ -0,0 +1,120 @@
from pydantic_settings import BaseSettings
from pydantic import model_validator
from typing import Optional
class Settings(BaseSettings):
"""应用程序配置"""
class Config:
env_file = ".env"
case_sensitive = False
extra = 'ignore'
# Service
app_name: str = "DeepData Python Backend"
app_version: str = "1.0.0"
app_description: str = "Adapter for integrating Data Management System with Label Studio"
host: str = "0.0.0.0"
port: int = 18000
# CORS
# allowed_origins: List[str] = ["*"]
# allowed_methods: List[str] = ["*"]
# allowed_headers: List[str] = ["*"]
# Log
log_level: str = "INFO"
debug: bool = True
log_file_dir: str = "/var/log/deepdata/backend-python"
log_pvc_monitor_enabled: bool = True
log_pvc_monitor_path: str = "/var/log/deepdata"
log_pvc_monitor_threshold: float = 0.9
log_pvc_monitor_interval_seconds: int = 300
log_pvc_monitor_delete_batch_size: int = 10
log_pvc_monitor_file_suffixes: str = "log,out,err"
log_rotation_max_size: str = "100MB"
log_rotation_backup_count: int = 30
rag_storage_dir: str = "/data/rag_storage"
# Database
pgsql_host: str = "deepdata-database"
pgsql_port: int = 5432
pgsql_user: str = "postgres"
pgsql_password: str = ""
pgsql_database: str = "deepdata"
# Database
mysql_host: str = "deepdata-database"
mysql_port: int = 3306
mysql_user: str = "root"
mysql_password: str = ""
mysql_database: str = "deepdata"
database_url: str = "" # Will be overridden by build_database_url() if not provided
@model_validator(mode='after')
def build_database_url(self):
"""如果没有提供 database_url,则根据 MySQL 配置构建"""
if not self.database_url:
if self.pgsql_host:
if self.pgsql_password and self.pgsql_user:
self.database_url = f"postgresql+asyncpg://{self.pgsql_user}:{self.pgsql_password}@{self.pgsql_host}:{self.pgsql_port}/{self.pgsql_database}"
else:
self.database_url = f"postgresql+asyncpg://{self.pgsql_host}:{self.pgsql_port}/{self.pgsql_database}"
elif self.mysql_password and self.mysql_user:
self.database_url = f"mysql+aiomysql://{self.mysql_user}:{self.mysql_password}@{self.mysql_host}:{self.mysql_port}/{self.mysql_database}"
else:
self.database_url = f"mysql+aiomysql://{self.mysql_host}:{self.mysql_port}/{self.mysql_database}"
return self
# Label Studio
label_studio_base_url: str = "http://label-studio:8000"
label_studio_username: Optional[str] = None
label_studio_password: Optional[str] = None
label_studio_user_token: Optional[str] = None # Legacy Token
label_studio_local_document_root: str = "/label-studio/local" # Label Studio local file storage path
label_studio_file_path_prefix: str = "/data/local-files/?d=" # Label Studio local file serving URL prefix
ls_task_page_size: int = 1000
# DeepData
dm_file_path_prefix: str = "/dataset" # DM存储文件夹前缀
deepdata_jwt_enable: bool = False
# Milvus 配置
milvus_uri: str = "http://milvus:19530"
milvus_token: str = ""
# 文件存储配置(共享文件系统)
file_storage_path: str = "/data/files"
# ==================== 配比任务并行复制配置 ====================
# 动态并发计算参数(全闪存储高性能场景默认值)
# 并发下限(最少并发数)
ratio_copy_min_concurrent: int = 8
# 并发上限(最多并发数,防止资源耗尽)
ratio_copy_max_concurrent: int = 128
# CPU核心系数(每个核心贡献的并发数,全闪存储建议4.0)
ratio_copy_cpu_factor: float = 4.0
# 每并发任务预估内存占用(MB)
ratio_copy_memory_per_task_mb: int = 32
# 内存安全保留比例(保留给其他进程)
ratio_copy_memory_reserve_ratio: float = 0.2
# 是否启用动态计算(False则使用固定值)
ratio_copy_dynamic_concurrent: bool = True
# 固定并发数(当 dynamic_concurrent=False 时使用)
ratio_copy_fixed_concurrent: int = 10
# 全局设置实例
settings = Settings()
@@ -0,0 +1,240 @@
"""
核心异常处理模块
为应用程序提供统一、优雅的异常处理系统,确保所有 API 响应都符合标准化格式。
## 快速开始
### 1. 注册异常处理器(在 main.py 中)
```python
from app.core.exception import register_exception_handlers, ExceptionHandlingMiddleware
app = FastAPI()
# 注册全局异常捕获中间件(最外层)
app.add_middleware(ExceptionHandlingMiddleware)
# 注册异常处理器
register_exception_handlers(app)
```
### 2. 在代码中抛出业务异常
```python
from app.core.exception import ErrorCodes, BusinessError
# 资源不存在
async def get_user(user_id: str):
user = await db.get_user(user_id)
if not user:
raise BusinessError(ErrorCodes.NOT_FOUND)
return user
# 参数验证失败
async def create_user(name: str):
if not name:
raise BusinessError(ErrorCodes.BAD_REQUEST)
# ...
```
### 3. 返回成功响应
```python
from app.core.exception import SuccessResponse
@router.get("/users/{user_id}")
async def get_user(user_id: str):
user = await service.get_user(user_id)
return SuccessResponse(data=user)
# 返回格式: {"code": "0", "message": "success", "data": {...}}
```
### 4. 使用事务管理
```python
from app.core.exception import transaction, BusinessError, ErrorCodes
@router.post("/users")
async def create_user(request: CreateUserRequest, db: AsyncSession = Depends(get_db)):
async with transaction(db):
# 所有数据库操作
user = User(name=request.name)
db.add(user)
await db.flush()
# 如果抛出异常,事务会自动回滚
if check_duplicate(user.name):
raise BusinessError(ErrorCodes.OPERATION_FAILED)
# 事务已自动提交
return SuccessResponse(data=user)
```
## 异常类型说明
### BusinessError(业务异常)
用于预期的业务错误,如资源不存在、权限不足、参数错误等。
特点:
- 不会记录完整的堆栈跟踪(因为这是预期内的错误)
- 返回对应的 HTTP 状态码(400、404 等)
- 客户端会收到标准化的错误响应
使用场景:
- 资源不存在
- 参数验证失败
- 权限不足
- 业务规则违反
### SystemError(系统异常)
用于意外的系统错误,如数据库错误、网络错误、配置错误等。
特点:
- 记录完整的堆栈跟踪
- 返回 HTTP 500
- 不暴露敏感的系统信息给客户端
使用场景:
- 数据库连接失败
- 网络超时
- 配置错误
- 编程错误
## 错误码定义
所有错误码在 `ErrorCodes` 类中集中定义,遵循规范:`{module}.{sequence}`
### 通用错误码
- `SUCCESS` (0): 操作成功
- `BAD_REQUEST` (common.0001): 请求参数错误
- `NOT_FOUND` (common.0002): 资源不存在
- `FORBIDDEN` (common.0003): 权限不足
- `UNAUTHORIZED` (common.0004): 未授权访问
- `VALIDATION_ERROR` (common.0005): 数据验证失败
### 系统级错误码
- `INTERNAL_ERROR` (system.0001): 服务器内部错误
- `DATABASE_ERROR` (system.0002): 数据库错误
- `NETWORK_ERROR` (system.0003): 网络错误
### 模块错误码
- `annotation.*`: 标注模块相关错误
- `collection.*`: 归集模块相关错误
- `evaluation.*`: 评估模块相关错误
- `generation.*`: 生成模块相关错误
- `rag.*`: RAG 模块相关错误
- `ratio.*`: 配比模块相关错误
## Result 类型(可选的函数式错误处理)
如果你不喜欢使用异常,可以使用 Result 类型进行函数式错误处理:
```python
from app.core.exception import Result, Ok, Err, ErrorCodes
def get_user(user_id: str) -> Result[User]:
user = db.find_user(user_id)
if user:
return Ok(user)
return Err(ErrorCodes.NOT_FOUND)
# 使用结果
result = get_user("123")
if result.is_ok():
user = result.unwrap()
print(f"User: {user.name}")
else:
error = result.unwrap_err()
print(f"Error: {error.message}")
# 链式操作
result = get_user("123")
.map(lambda user: user.name)
.and_then(validate_name)
```
## 响应格式
### 成功响应
```json
{
"code": "0",
"message": "success",
"data": {
"id": 123,
"name": "张三"
}
}
```
### 错误响应
```json
{
"code": "common.0002",
"message": "资源不存在",
"data": null
}
```
## 最佳实践
1. **始终使用业务异常**:对于可预见的业务错误,使用 `BusinessError` 而不是 HTTPException
2. **集中定义错误码**:所有错误码在 `ErrorCodes` 中定义,不要硬编码
3. **提供有用的数据**:在抛出异常时,可以通过 `data` 参数传递额外的错误信息
4. **使用事务管理**:涉及多个数据库操作时,使用 `transaction` 上下文管理器
5. **不要捕获 SystemError**:让系统错误由全局处理器统一处理
## 迁移指南
如果你有旧的异常处理代码,迁移步骤:
1. 删除旧的异常类定义
2. 将 `raise OldException(...)` 替换为 `raise BusinessError(ErrorCodes.XXX)`
3. 移除 try-except 中的异常转换逻辑,让全局处理器处理
4. 更新导入语句:`from app.core.exception import ErrorCodes, BusinessError`
## 测试
可以使用测试端点验证异常处理:
```bash
curl http://localhost:8000/test-success
curl http://localhost:8000/test-business-error
curl http://localhost:8000/test-system-error
```
"""
from .base import BaseError, ErrorCode, SystemError, BusinessError
from .codes import ErrorCodes
from .handlers import (
register_exception_handlers,
ErrorResponse,
SuccessResponse
)
from .middleware import ExceptionHandlingMiddleware
from .result import Result, Ok, Err
from .transaction import transaction, ensure_transaction_rollback
__all__ = [
# 基础异常类
'BaseError',
'ErrorCode',
'SystemError',
'BusinessError',
# 错误码
'ErrorCodes',
# 处理器
'register_exception_handlers',
'ErrorResponse',
'SuccessResponse',
# 中间件
'ExceptionHandlingMiddleware',
# Result 类型
'Result',
'Ok',
'Err',
# 事务管理
'transaction',
'ensure_transaction_rollback',
]
@@ -0,0 +1,121 @@
"""
基础异常类和错误码定义
"""
from typing import Any
from dataclasses import dataclass
@dataclass(frozen=True)
class ErrorCode:
"""
不可变的错误码定义
属性:
code: 错误码字符串(如 "annotation.0001")
message: 人类可读的错误消息
http_status: HTTP状态码(业务错误默认为400)
"""
code: str
message: str
http_status: int = 400
def __post_init__(self):
"""验证错误码格式"""
if not isinstance(self.code, str):
raise ValueError(f"错误码必须是字符串,实际类型: {type(self.code)}")
if not self.code:
raise ValueError("错误码不能为空")
class BaseError(Exception):
"""
所有应用异常的基类
所有自定义异常都应该继承此类。
提供自动的错误码和消息处理。
使用示例:
raise BusinessError(ErrorCodes.TASK_NOT_FOUND)
客户端将收到:
{
"code": "annotation.0001",
"message": "任务不存在",
"data": null
}
"""
def __init__(
self,
error_code: ErrorCode,
data: Any = None,
*args: Any
):
"""
使用错误码和可选数据初始化异常
Args:
error_code: ErrorCode 实例
data: 附加错误数据(会包含在响应中)
*args: 额外参数(用于兼容性)
"""
self.error_code = error_code
self.data = data
super().__init__(error_code.message, *args)
@property
def code(self) -> str:
"""获取错误码字符串"""
return self.error_code.code
@property
def message(self) -> str:
"""获取错误消息"""
return self.error_code.message
@property
def http_status(self) -> int:
"""获取HTTP状态码"""
return self.error_code.http_status
def to_dict(self) -> dict:
"""将异常转换为响应字典"""
return {
"code": self.code,
"message": self.message,
"data": self.data
}
class SystemError(BaseError):
"""
系统级异常(意外错误)
用于:
- 数据库错误
- 网络错误
- 配置错误
- 编程错误(bug)
系统错误会记录完整的堆栈跟踪
"""
def __init__(self, error_code: ErrorCode, data: Any = None):
super().__init__(error_code, data)
class BusinessError(BaseError):
"""
业务逻辑异常(预期内的错误)
用于:
- 验证失败
- 资源不存在
- 权限不足
- 业务规则违反
业务错误不记录堆栈跟踪(因为它们是预期的)
"""
def __init__(self, error_code: ErrorCode, data: Any = None):
super().__init__(error_code, data)
@@ -0,0 +1,212 @@
"""
集中式错误码定义
所有错误码都在这里定义,遵循规范:{module}.{sequence}
模块代码:
- common: 通用错误
- system: 系统级错误
- annotation: 标注模块
- collection: 归集模块
- evaluation: 评估模块
- generation: 生成模块
- rag: RAG模块
- ratio: 配比模块
"""
from typing import Final
from .base import ErrorCode
class ErrorCodes:
def __init__(self):
self.message = None
self.code = None
"""
集中式错误码定义
所有错误码在此一次性定义,使用时直接通过类属性访问。
使用示例:
from app.core.exception import ErrorCodes, BusinessError
raise BusinessError(ErrorCodes.TASK_NOT_FOUND)
"""
# ========== 通用错误码 ==========
SUCCESS: Final = ErrorCode("0", "Success", 200)
BAD_REQUEST: Final = ErrorCode("common.0001", "Bad request", 400)
NOT_FOUND: Final = ErrorCode("common.0002", "Resource not found", 404)
FORBIDDEN: Final = ErrorCode("common.0003", "Forbidden", 403)
UNAUTHORIZED: Final = ErrorCode("common.0004", "Unauthorized", 401)
VALIDATION_ERROR: Final = ErrorCode("common.0005", "Validation error", 422)
OPERATION_FAILED: Final = ErrorCode("common.0006", "Operation failed", 500)
# ========== 系统级错误码 ==========
INTERNAL_ERROR: Final = ErrorCode("system.0001", "Internal server error", 500)
DATABASE_ERROR: Final = ErrorCode("system.0002", "Database error", 500)
NETWORK_ERROR: Final = ErrorCode("system.0003", "Network error", 500)
CONFIG_ERROR: Final = ErrorCode("system.0004", "Configuration error", 500)
SERVICE_UNAVAILABLE: Final = ErrorCode("system.0005", "Service unavailable", 503)
# ========== 标注模块 ==========
ANNOTATION_TASK_NOT_FOUND: Final = ErrorCode(
"annotation.0001", "Annotation task not found", 404
)
ANNOTATION_PROJECT_NOT_FOUND: Final = ErrorCode(
"annotation.0002", "Annotation project not found", 404
)
ANNOTATION_TEMPLATE_NOT_FOUND: Final = ErrorCode(
"annotation.0003", "Annotation template not found", 404
)
ANNOTATION_FILE_NOT_FOUND: Final = ErrorCode(
"annotation.0004", "File not found", 404
)
ANNOTATION_TAG_UPDATE_FAILED: Final = ErrorCode(
"annotation.0005", "Failed to update tags", 500
)
# ========== 归集模块 ==========
COLLECTION_TASK_NOT_FOUND: Final = ErrorCode(
"collection.0001", "Collection task not found", 404
)
COLLECTION_TEMPLATE_NOT_FOUND: Final = ErrorCode(
"collection.0002", "Collection template not found", 404
)
COLLECTION_EXECUTION_NOT_FOUND: Final = ErrorCode(
"collection.0003", "Execution record not found", 404
)
COLLECTION_LOG_NOT_FOUND: Final = ErrorCode(
"collection.0004", "Log file not found", 404
)
# ========== 评估模块 ==========
EVALUATION_TASK_NOT_FOUND: Final = ErrorCode(
"evaluation.0001", "Evaluation task not found", 404
)
EVALUATION_TASK_TYPE_ERROR: Final = ErrorCode(
"evaluation.0002", "Invalid task type", 400
)
EVALUATION_MODEL_NOT_FOUND: Final = ErrorCode(
"evaluation.0003", "Evaluation model not found", 404
)
# ========== 生成模块 ==========
GENERATION_TASK_NOT_FOUND: Final = ErrorCode(
"generation.0001", "Generation task not found", 404
)
GENERATION_FILE_NOT_FOUND: Final = ErrorCode(
"generation.0002", "Generation file not found", 404
)
GENERATION_CHUNK_NOT_FOUND: Final = ErrorCode(
"generation.0003", "Data chunk not found", 404
)
GENERATION_DATA_NOT_FOUND: Final = ErrorCode(
"generation.0004", "Generation data not found", 404
)
# ========== RAG 模块 ==========
RAG_CONFIG_ERROR: Final = ErrorCode("rag.0001", "RAG configuration error", 400)
RAG_KNOWLEDGE_BASE_NOT_FOUND: Final = ErrorCode(
"rag.0002", "Knowledge base not found", 404
)
RAG_KNOWLEDGE_BASE_ALREADY_EXISTS: Final = ErrorCode(
"rag.0003", "Knowledge base already exists", 400
)
RAG_KNOWLEDGE_BASE_NAME_INVALID: Final = ErrorCode(
"rag.0004", "Knowledge base name is invalid", 400
)
RAG_FILE_NOT_FOUND: Final = ErrorCode("rag.0005", "RAG file not found", 404)
RAG_FILE_PROCESS_FAILED: Final = ErrorCode(
"rag.0006", "File processing failed", 500
)
RAG_FILE_PARSE_FAILED: Final = ErrorCode("rag.0007", "File parsing failed", 500)
RAG_CHUNK_NOT_FOUND: Final = ErrorCode("rag.0008", "Chunk not found", 404)
RAG_MODEL_NOT_FOUND: Final = ErrorCode("rag.0009", "RAG model not found", 404)
RAG_QUERY_FAILED: Final = ErrorCode("rag.0010", "RAG query failed", 500)
RAG_MILVUS_ERROR: Final = ErrorCode("rag.0011", "Milvus operation failed", 500)
RAG_COLLECTION_NOT_FOUND: Final = ErrorCode(
"rag.0012", "Milvus collection not found", 404
)
RAG_EMBEDDING_FAILED: Final = ErrorCode(
"rag.0013", "Embedding generation failed", 500
)
RAG_UNSUPPORTED_TYPE: Final = ErrorCode(
"rag.0014", "Unsupported RAG type", 400
)
RAG_INVALID_REQUEST: Final = ErrorCode(
"rag.0015", "Invalid request", 400
)
# ========== 配比模块 ==========
RATIO_TASK_NOT_FOUND: Final = ErrorCode("ratio.0001", "Ratio task not found", 404)
RATIO_NAME_REQUIRED: Final = ErrorCode("ratio.0002", "Task name is required", 400)
RATIO_ALREADY_EXISTS: Final = ErrorCode("ratio.0003", "Task already exists", 400)
RATIO_DELETE_FAILED: Final = ErrorCode("ratio.0004", "Failed to delete task", 500)
# ========== 清洗模块 ==========
CLEANING_TASK_NOT_FOUND: Final = ErrorCode(
"cleaning.0001", "Cleaning task not found", 404
)
CLEANING_NAME_DUPLICATED: Final = ErrorCode(
"cleaning.0002", "Cleaning task name is duplicated", 400
)
CLEANING_TEMPLATE_NOT_FOUND: Final = ErrorCode(
"cleaning.0003", "Cleaning template not found", 404
)
CLEANING_TEMPLATE_NAME_DUPLICATED: Final = ErrorCode(
"cleaning.0004", "Cleaning template name is duplicated", 400
)
CLEANING_INVALID_OPERATOR_INPUT: Final = ErrorCode(
"cleaning.0005", "Invalid operator input/output types", 400
)
CLEANING_INVALID_EXECUTOR_TYPE: Final = ErrorCode(
"cleaning.0006", "Invalid executor type", 400
)
CLEANING_DATASET_NOT_FOUND: Final = ErrorCode(
"cleaning.0007", "Dataset not found", 404
)
CLEANING_FILE_SYSTEM_ERROR: Final = ErrorCode(
"cleaning.0008", "File system error", 500
)
CLEANING_SETTINGS_PARSE_ERROR: Final = ErrorCode(
"cleaning.0009", "Settings parse error", 400
)
CLEANING_TASK_ID_REQUIRED: Final = ErrorCode(
"cleaning.0010", "Task ID is required", 400
)
CLEANING_CANNOT_DELETE_PRESET_TEMPLATE: Final = ErrorCode(
"cleaning.0011", "Cannot delete preset template", 400
)
CLEANING_TASK_LOG_NOT_FOUND: Final = ErrorCode(
"cleaning.0012", "Cleaning task log file not found", 404
)
CLEANING_TASK_STATUS_INVALID: Final = ErrorCode(
"cleaning.0013", "Cleaning task status is invalid for this operation", 400
)
# ========== 算子市场模块 ==========
OPERATOR_NOT_FOUND: Final = ErrorCode("operator.0001", "Operator not found", 404)
OPERATOR_IN_INSTANCE: Final = ErrorCode("operator.0002", "Operator is in use", 400)
OPERATOR_CANNOT_DELETE_PREDEFINED: Final = ErrorCode(
"operator.0003", "Cannot delete predefined operator", 400
)
OPERATOR_UNSUPPORTED_FILE_TYPE: Final = ErrorCode(
"operator.0004", "Unsupported file type", 400
)
OPERATOR_PARSE_FAILED: Final = ErrorCode(
"operator.0005", "Failed to parse operator package", 400
)
OPERATOR_FIELD_NOT_FOUND: Final = ErrorCode(
"operator.0006", "Required field is missing", 400
)
# ========== 系统模块 ==========
SYSTEM_MODEL_NOT_FOUND: Final = ErrorCode(
"system.0006", "Model configuration not found", 404
)
SYSTEM_MODEL_HEALTH_CHECK_FAILED: Final = ErrorCode(
"system.0007", "Model health check failed", 500
)
@@ -0,0 +1,203 @@
"""
异常处理器和响应构建器
提供异常的自动转换,将异常转换为标准化的 JSON 响应。
"""
from typing import Any, Optional
from fastapi import Request, Response
from fastapi.responses import JSONResponse
from fastapi.exceptions import RequestValidationError
from starlette.exceptions import HTTPException as StarletteHTTPException
from .base import BaseError, BusinessError, SystemError
from .codes import ErrorCodes
from ..logging import get_logger
logger = get_logger(__name__)
def SuccessResponse(data: Any = None, message: str = "success", code: str = "0") -> dict:
"""
构建成功响应
Args:
data: 响应数据
message: 成功消息
code: 自定义成功码(默认为 "0")
Returns:
响应字典,符合 StandardResponse 格式
使用示例:
return SuccessResponse(data={"id": 123})
# 或
return SuccessResponse({"id": 123})
"""
return {
"code": code,
"message": message,
"data": data
}
def ErrorResponse(
error_code: ErrorCodes,
data: Any = None,
custom_message: Optional[str] = None
) -> dict:
"""
构建错误响应
Args:
error_code: 错误码(来自 ErrorCodes)
data: 附加错误数据
custom_message: 覆盖默认错误消息
Returns:
响应字典,符合 StandardResponse 格式
使用示例:
return ErrorResponse(ErrorCodes.TASK_NOT_FOUND)
"""
return {
"code": error_code.code,
"message": custom_message or error_code.message,
"data": data
}
async def business_error_handler(request: Request, exc: BusinessError) -> JSONResponse:
"""
处理业务逻辑异常
业务异常是预期的错误,不需要记录堆栈跟踪。
返回对应的 HTTP 状态码(400、404 等),错误信息在响应体的 code 字段中。
"""
return JSONResponse(
status_code=exc.http_status,
content=exc.to_dict()
)
async def system_error_handler(request: Request, exc: SystemError) -> JSONResponse:
"""
处理系统异常
系统异常是意外的错误,需要记录完整的堆栈跟踪。
返回 HTTP 500 和经过净化的错误消息。
"""
logger.error(
f"System error occurred at {request.method} {request.url.path}: {exc.message}",
exc_info=True
)
return JSONResponse(
status_code=exc.http_status,
content={
"code": exc.code,
"message": exc.message,
"data": None # 绝不暴露系统错误详情
}
)
async def validation_error_handler(
request: Request,
exc: RequestValidationError
) -> JSONResponse:
"""处理请求验证错误"""
errors = exc.errors()
simplified_errors = [
err.get("msg", "Validation error")
for err in errors
]
return JSONResponse(
status_code=422,
content={
"code": ErrorCodes.VALIDATION_ERROR.code,
"message": ErrorCodes.VALIDATION_ERROR.message,
"data": {
"detail": "请求验证失败",
"errors": simplified_errors
}
}
)
async def http_exception_handler(
request: Request,
exc: StarletteHTTPException
) -> JSONResponse:
"""处理 HTTP 异常(404、500 等)"""
# 根据 HTTP 状态码映射到错误码
error_code = ErrorCodes.INTERNAL_ERROR
if exc.status_code == 404:
error_code = ErrorCodes.NOT_FOUND
elif exc.status_code == 401:
error_code = ErrorCodes.UNAUTHORIZED
elif exc.status_code == 403:
error_code = ErrorCodes.FORBIDDEN
elif exc.status_code == 422:
error_code = ErrorCodes.VALIDATION_ERROR
elif exc.status_code == 400:
error_code = ErrorCodes.BAD_REQUEST
return JSONResponse(
status_code=exc.status_code,
content={
"code": error_code.code,
"message": "error",
"data": {
"detail": str(exc.detail)
}
}
)
async def generic_exception_handler(request: Request, exc: Exception) -> JSONResponse:
"""
通用异常处理器(最后的兜底处理)
这是最后一道防线 - 记录完整的堆栈跟踪
并返回经过净化的错误给客户端。
"""
logger.error(
f"Unhandled exception occurred at {request.method} {request.url.path}: {str(exc)}",
exc_info=True
)
return JSONResponse(
status_code=500,
content={
"code": ErrorCodes.INTERNAL_ERROR.code,
"message": ErrorCodes.INTERNAL_ERROR.message,
"data": {
"detail": "服务器内部错误"
}
}
)
def register_exception_handlers(app) -> None:
"""
注册所有异常处理器到 FastAPI 应用
Args:
app: FastAPI 应用实例
使用示例:
from app.core.exception import register_exception_handlers
app = FastAPI()
register_exception_handlers(app)
"""
from .base import BusinessError, SystemError
# 按照特异性顺序注册异常处理器(最具体的在前)
app.add_exception_handler(BusinessError, business_error_handler)
app.add_exception_handler(SystemError, system_error_handler)
app.add_exception_handler(RequestValidationError, validation_error_handler)
app.add_exception_handler(StarletteHTTPException, http_exception_handler)
app.add_exception_handler(Exception, generic_exception_handler)
@@ -0,0 +1,96 @@
"""
全局异常处理中间件,支持自动响应包装
此中间件提供自动地返回值和异常转换,确保所有响应都符合标准化的 JSON 格式。
关键特性:
- 自动包装 dict/list 响应为 StandardResponse 格式
- 异常处理,并记录适当的日志
- 堆栈跟踪隔离(永远不会暴露给客户端)
"""
from fastapi import Request
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import Response as StarletteResponse
from .base import BaseError, BusinessError, SystemError
from .codes import ErrorCodes
from ..logging import get_logger
logger = get_logger(__name__)
class ExceptionHandlingMiddleware(BaseHTTPMiddleware):
"""
全局异常捕获中间件
这是最外层的中间件,捕获所有未处理的异常,
并将它们转换为标准化的响应格式。
堆栈跟踪信息只记录到日志文件,永远不会发送给客户端。
"""
async def dispatch(self, request: Request, call_next):
"""
处理请求并处理任何异常
异常会被捕获、记录(带堆栈跟踪),并转换为标准化响应。
业务异常会被重新抛出,由专门的处理器处理。
"""
try:
response = await call_next(request)
return response
except BusinessError:
# 让业务异常处理器处理
raise
except SystemError as exc:
# 记录系统错误及其完整堆栈跟踪
logger.error(
f"System error occurred at {request.method} {request.url.path}",
exc_info=True
)
return self._error_response(
code=exc.code,
message=exc.message,
http_status=exc.http_status
)
except BaseError as exc:
# 处理其他自定义错误
logger.warning(f"BaseError occurred at {request.url.path}: {exc.message}")
return self._error_response(
code=exc.code,
message=exc.message,
http_status=exc.http_status
)
except Exception as exc:
# 捕获所有未处理的异常
logger.error(
f"Unhandled exception occurred at {request.method} {request.url.path}",
exc_info=True
)
return self._error_response(
code=ErrorCodes.INTERNAL_ERROR.code,
message=ErrorCodes.INTERNAL_ERROR.message,
http_status=500
)
@staticmethod
def _error_response(
code: str,
message: str,
http_status: int
) -> StarletteResponse:
"""构建错误响应"""
return JSONResponse(
status_code=http_status,
content={
"code": code,
"message": message,
"data": None
}
)
@@ -0,0 +1,170 @@
"""
Result 类型用于优雅的错误处理
提供受 Rust 启发的 Result<T, E> 类型,用于处理可能失败的操作而无需使用异常。
使用示例:
# 成功情况
def get_user(user_id: str) -> Result[User]:
user = db.find_user(user_id)
if user:
return Ok(user)
return Err(ErrorCodes.USER_NOT_FOUND)
# 使用结果
result = get_user("123")
if result.is_ok():
user = result.unwrap()
print(f"User: {user.name}")
else:
error = result.unwrap_err()
print(f"Error: {error.message}")
"""
from typing import Generic, TypeVar, Optional, Any
from .base import ErrorCode
from .codes import ErrorCodes
T = TypeVar('T') # 成功类型
E = TypeVar('E', bound=ErrorCode) # 错误类型
class Result(Generic[T, E]):
"""
表示成功(Ok)或失败(Err)的 Result 类型
此类型允许在不需要异常的情况下进行显式错误处理
"""
def __init__(self, value: Optional[T], error: Optional[E], is_ok: bool):
self._value = value
self._error = error
self._is_ok = is_ok
@staticmethod
def ok(value: T) -> 'Result[T, E]':
"""创建一个包含值的成功结果"""
return Result(value, None, True)
@staticmethod
def err(error: E) -> 'Result[T, E]':
"""创建一个包含错误码的失败结果"""
return Result(None, error, False)
@property
def is_ok(self) -> bool:
"""检查结果是否成功"""
return self._is_ok
@property
def is_err(self) -> bool:
"""检查结果是否失败"""
return not self._is_ok
def unwrap(self) -> T:
"""
获取成功值
Returns:
成功值
Raises:
ValueError: 如果结果是错误
"""
if self._is_ok:
return self._value
raise ValueError(
f"Cannot unwrap error result: {self._error.message}"
)
def unwrap_err(self) -> E:
"""
获取错误码
Returns:
错误码
Raises:
ValueError: 如果结果是成功的
"""
if not self._is_ok:
return self._error
raise ValueError("Cannot unwrap error from successful result")
def unwrap_or(self, default: T) -> T:
"""
获取成功值,如果出错则返回默认值
Args:
default: 出错时返回的默认值
Returns:
成功值或默认值
"""
return self._value if self._is_ok else default
def map(self, func) -> 'Result[Any, E]':
"""
如果存在成功值,则应用函数
Args:
func: 要应用的函数
Returns:
包含映射值的新结果或相同的错误
"""
if self._is_ok:
try:
return Result.ok(func(self._value))
except Exception:
# 如果映射失败,转换为错误
return Result.err(ErrorCodes.INTERNAL_ERROR)
return self
def and_then(self, func) -> 'Result[Any, E]':
"""
链式调用返回 Result 的操作
Args:
func: 接收成功值并返回新 Result 的函数
Returns:
来自函数的新结果或相同的错误
"""
if self._is_ok:
return func(self._value)
return self
def or_else(self, func) -> 'Result[T, Any]':
"""
在出错时提供备用结果
Args:
func: 接收错误并返回新 Result 的函数
Returns:
相同的结果或来自函数的新结果
"""
if not self._is_ok:
return func(self._error)
return self
def Ok(value: T) -> Result[T, ErrorCode]:
"""
创建一个成功的结果
使用示例:
return Ok(user_data)
"""
return Result.ok(value)
def Err(error_code: ErrorCode) -> Result[Any, ErrorCode]:
"""
创建一个失败的结果
使用示例:
return Err(ErrorCodes.USER_NOT_FOUND)
"""
return Result.err(error_code)
@@ -0,0 +1,76 @@
"""
数据库事务管理工具
提供自动事务管理的上下文管理器和依赖注入函数,
确保在异常情况下正确回滚事务。
"""
from contextlib import asynccontextmanager
from typing import AsyncGenerator
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.logging import get_logger
logger = get_logger(__name__)
@asynccontextmanager
async def transaction(db: AsyncSession) -> AsyncGenerator[None, None]:
"""
数据库事务上下文管理器
自动处理事务的提交和回滚:
- 如果代码块正常执行完成,则自动提交
- 如果代码块抛出异常,则自动回滚
使用示例:
@router.post("/tasks")
async def create_task(
request: CreateTaskRequest,
db: AsyncSession = Depends(get_db)
):
async with transaction(db):
# 数据库操作
db.add(task)
await db.flush()
# 事务已自动提交
return SuccessResponse(data=task)
Args:
db: SQLAlchemy 异步会话
Yields:
None
Raises:
Exception: 重新抛出代码块中的任何异常
"""
try:
yield
# 如果没有异常,提交事务
await db.commit()
logger.debug("Transaction committed successfully")
except Exception as e:
# 发生异常,回滚事务
try:
await db.rollback()
logger.warning(f"Transaction rolled back due to exception: {e}")
except Exception as rollback_error:
logger.error(f"Failed to rollback transaction: {rollback_error}")
# 重新抛出异常,让上层异常处理器处理
raise
async def ensure_transaction_rollback(db: AsyncSession) -> None:
"""
确保事务回滚(用于错误处理)
Args:
db: SQLAlchemy 异步会话
"""
try:
await db.rollback()
logger.debug("Transaction rolled back")
except Exception as e:
logger.error(f"Failed to rollback transaction: {e}")
+121
View File
@@ -0,0 +1,121 @@
import logging
import re
import sys
from logging.handlers import RotatingFileHandler
from pathlib import Path
from app.core.config import settings
_SIZE_PATTERN = re.compile(r"^\s*(\d+)\s*([KMG]?B?)?\s*$", re.IGNORECASE)
class CenteredLevelNameFormatter(logging.Formatter):
"""Center the level name in the log output"""
def format(self, record):
# 将 levelname 居中对齐到8个字符
record.levelname = record.levelname.center(8)
return super().format(record)
def _parse_size_to_bytes(value: str) -> int:
match = _SIZE_PATTERN.match(value)
if not match:
return 100 * 1024 * 1024
amount = int(match.group(1))
unit = (match.group(2) or "B").upper()
if unit in {"G", "GB"}:
return amount * 1024 * 1024 * 1024
if unit in {"M", "MB"}:
return amount * 1024 * 1024
if unit in {"K", "KB"}:
return amount * 1024
return amount
def _rotated_log_namer(default_name: str) -> str:
path = Path(default_name)
index_suffix = path.suffix
if not index_suffix.lstrip(".").isdigit():
return default_name
base_path = Path(str(path)[: -len(index_suffix)])
if not base_path.suffix:
return default_name
return str(base_path.with_name(f"{base_path.stem}{index_suffix}{base_path.suffix}"))
def _remove_managed_handlers(logger: logging.Logger) -> None:
for handler in list(logger.handlers):
if getattr(handler, "_deepdata_managed", False):
logger.removeHandler(handler)
handler.close()
def _mark_managed(handler: logging.Handler) -> logging.Handler:
setattr(handler, "_deepdata_managed", True)
return handler
def setup_logging():
log_format = "%(asctime)s [%(levelname)s] - %(name)s - %(message)s"
date_format = "%Y-%m-%d %H:%M:%S"
console_handler = _mark_managed(logging.StreamHandler(sys.stdout))
console_handler.setLevel(getattr(logging, settings.log_level.upper()))
log_dir = Path(settings.log_file_dir)
log_dir.mkdir(exist_ok=True)
file_handler = _mark_managed(RotatingFileHandler(
log_dir / "python-backend.log",
maxBytes=_parse_size_to_bytes(settings.log_rotation_max_size),
backupCount=max(1, settings.log_rotation_backup_count),
encoding="utf-8"
))
file_handler.namer = _rotated_log_namer
file_handler.setLevel(getattr(logging, settings.log_level.upper()))
# Style setting - Centered level names
formatter = CenteredLevelNameFormatter(log_format, date_format)
console_handler.setFormatter(formatter)
file_handler.setFormatter(formatter)
# Root Logger
root_logger = logging.getLogger()
_remove_managed_handlers(root_logger)
root_logger.setLevel(getattr(logging, settings.log_level.upper()))
root_logger.addHandler(console_handler)
root_logger.addHandler(file_handler)
# Uvicorn
uvicorn_logger = logging.getLogger("uvicorn")
uvicorn_logger.handlers.clear()
uvicorn_logger.addHandler(console_handler)
uvicorn_logger.setLevel(logging.INFO)
uvicorn_access = logging.getLogger("uvicorn.access")
uvicorn_access.handlers.clear()
uvicorn_access.addHandler(console_handler)
uvicorn_access.setLevel(logging.DEBUG)
uvicorn_error = logging.getLogger("uvicorn.error")
uvicorn_error.handlers.clear()
uvicorn_error.addHandler(console_handler)
uvicorn_error.setLevel(logging.ERROR)
# SQLAlchemy (ERROR only)
sqlalchemy_logger = logging.getLogger("sqlalchemy.engine")
_remove_managed_handlers(sqlalchemy_logger)
sqlalchemy_logger.setLevel(logging.ERROR)
sqlalchemy_logger.addHandler(console_handler)
sqlalchemy_logger.setLevel(logging.ERROR)
# Minimize noise from HTTPX and HTTPCore
logging.getLogger("httpx").setLevel(logging.ERROR)
logging.getLogger("httpcore").setLevel(logging.ERROR)
def get_logger(name: str) -> logging.Logger:
return logging.getLogger(name)
@@ -0,0 +1,159 @@
"""Security utilities for handling sensitive data"""
import re
from typing import Any, Dict
SENSITIVE_FIELDS = [
"secretKey",
"accessKey",
"password",
"passwd",
"pwd",
"secret",
"token",
"apiKey",
"api_key",
"access_key",
"secret_key",
]
MASK_PATTERN = "**********"
def mask_sensitive_value(value: str) -> str:
"""Mask a single sensitive value"""
return MASK_PATTERN
def is_masked_value(value: str) -> bool:
"""Check if a value is already masked"""
return value == MASK_PATTERN
def preserve_sensitive_values(
new_config: Dict[str, Any],
original_config: Dict[str, Any]
) -> Dict[str, Any]:
"""
Preserve original sensitive values if masked pattern is detected in update request.
When frontend receives masked values and sends them back in update request,
this function detects masked values and replaces them with original values
from database.
Args:
new_config: Config from update request (may contain masked values)
original_config: Original config from database (contains real values)
Returns:
Config with original sensitive values preserved
"""
if not isinstance(new_config, dict) or not isinstance(original_config, dict):
return new_config
preserved_config = {}
for key, new_value in new_config.items():
# Check if key is a sensitive field
key_lower = key.lower()
is_sensitive = any(
field.lower() == key_lower or field.lower() in key_lower
for field in SENSITIVE_FIELDS
)
# If value is masked and field is sensitive, use original value
if is_sensitive and isinstance(new_value, str) and is_masked_value(new_value):
original_value = original_config.get(key)
preserved_config[key] = original_value if original_value else new_value
elif isinstance(new_value, dict):
# Recursively process nested dictionaries
original_nested = original_config.get(key, {})
preserved_config[key] = preserve_sensitive_values(new_value, original_nested)
elif isinstance(new_value, list):
original_list = original_config.get(key, [])
preserved_list = []
for i, new_item in enumerate(new_value):
if i < len(original_list):
orig_item = original_list[i]
if isinstance(new_item, dict) and isinstance(orig_item, dict):
preserved_list.append(preserve_sensitive_values(new_item, orig_item))
else:
preserved_list.append(new_item)
else:
preserved_list.append(new_item)
preserved_config[key] = preserved_list
else:
# Keep non-sensitive/non-masked values
preserved_config[key] = new_value
return preserved_config
def mask_sensitive_dict(data: Dict[str, Any]) -> Dict[str, Any]:
"""
Recursively mask sensitive fields in a dictionary.
Args:
data: Dictionary that may contain sensitive fields
Returns:
Dictionary with sensitive values masked
"""
if not isinstance(data, dict):
return data
masked_data = {}
for key, value in data.items():
# Check if the key is a sensitive field (case-insensitive)
key_lower = key.lower()
is_sensitive = any(
field.lower() == key_lower or field.lower() in key_lower
for field in SENSITIVE_FIELDS
)
if is_sensitive and isinstance(value, str):
# Mask the sensitive value
masked_data[key] = MASK_PATTERN
elif isinstance(value, dict):
# Recursively mask nested dictionaries
masked_data[key] = mask_sensitive_dict(value)
elif isinstance(value, list):
# Process list items
masked_data[key] = [
mask_sensitive_dict(item) if isinstance(item, dict) else item
for item in value
]
else:
# Keep non-sensitive values as-is
masked_data[key] = value
return masked_data
def mask_sensitive_info(text: str) -> str:
"""
Mask sensitive information in text by replacing values with **********
Args:
text: Original text that may contain sensitive information
Returns:
Text with sensitive values masked
"""
masked_text = text
for field in SENSITIVE_FIELDS:
patterns = [
rf'"{field}"\s*:\s*"[^"]*"',
rf'{field}\s*=\s*[^\s,\]]+',
rf"'{field}'\s*:\s*'[^']*'",
]
for pattern in patterns:
if '"' in pattern:
masked_text = re.sub(pattern, f'"{field}": "{MASK_PATTERN}"', masked_text, flags=re.IGNORECASE)
elif "'" in pattern:
masked_text = re.sub(pattern, f"'{field}': '{MASK_PATTERN}'", masked_text, flags=re.IGNORECASE)
else:
masked_text = re.sub(pattern, f'{field}={MASK_PATTERN}', masked_text, flags=re.IGNORECASE)
return masked_text
@@ -0,0 +1,37 @@
from contextvars import ContextVar
from typing import List, Optional
_current_user: ContextVar[Optional[str]] = ContextVar("_current_user", default=None)
SYSTEM_USER = "system"
class DataScopeHandle:
"""
Hold current user info in a ContextVar and provide helpers for SQLAlchemy filters.
"""
@staticmethod
def set_user_info(user: Optional[str]) -> None:
if user is None or user == "":
# set explicit None
_current_user.set(None)
else:
_current_user.set(user)
@staticmethod
def remove_user_info() -> None:
_current_user.set(None)
@staticmethod
def get_user_info() -> Optional[str]:
return _current_user.get()
@staticmethod
def allowed_users() -> List[str]:
"""
Return list of allowed creators: current user + system.
"""
user = DataScopeHandle.get_user_info()
if not user:
return []
return [user, SYSTEM_USER]
@@ -0,0 +1,42 @@
from .dataset_management import (
Dataset,
DatasetTag,
DatasetFiles,
DatasetStatistics,
Tag,
)
from .user_management import User
from .annotation_management import AnnotationTemplate, LabelingProject
from .data_evaluation import EvaluationTask, EvaluationItem
from .operator import Operator, Category, CategoryRelation, OperatorRelease
from .chunk_upload import ChunkUploadPreRequest
from .knowledge_gen import KnowledgeBase, RagFile
from .sys_param import SysParam
__all__ = [
"Dataset",
"DatasetTag",
"DatasetFiles",
"DatasetStatistics",
"Tag",
"User",
"AnnotationTemplate",
"LabelingProject",
"EvaluationTask",
"EvaluationItem",
"Operator",
"Category",
"CategoryRelation",
"OperatorRelease",
"ChunkUploadPreRequest",
"KnowledgeBase",
"RagFile",
"SysParam",
]
@@ -0,0 +1,85 @@
"""Tables of Annotation Management Module"""
import uuid
from sqlalchemy import Column, String, Boolean, TIMESTAMP, Text, Integer, JSON, ForeignKey
from sqlalchemy.sql import func
from app.db.models.base_entity import BaseEntity
class AnnotationTemplate(BaseEntity):
"""标注配置模板模型"""
__tablename__ = "t_dm_annotation_templates"
__ignore_data_scope__ = True
id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()), comment="UUID")
name = Column(String(100), nullable=False, comment="模板名称")
description = Column(String(500), nullable=True, comment="模板描述")
data_type = Column(String(50), nullable=False, comment="数据类型: image/text/audio/video/timeseries")
labeling_type = Column(String(50), nullable=False, comment="标注类型: classification/detection/segmentation/ner/relation/etc")
configuration = Column(JSON, nullable=False, comment="标注配置(包含labels定义等)")
style = Column(String(32), nullable=False, comment="样式配置: horizontal/vertical")
category = Column(String(50), default='custom', comment="模板分类: medical/general/custom/system")
built_in = Column(Boolean, default=False, comment="是否系统内置模板")
version = Column(String(20), default='1.0', comment="模板版本")
deleted_at = Column(TIMESTAMP, nullable=True, comment="删除时间(软删除)")
def __repr__(self):
return f"<AnnotationTemplate(id={self.id}, name={self.name}, data_type={self.data_type})>"
@property
def is_deleted(self) -> bool:
"""检查是否已被软删除"""
return self.deleted_at is not None
class LabelingProject(BaseEntity):
"""标注项目模型"""
__tablename__ = "t_dm_labeling_projects"
id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()), comment="UUID")
dataset_id = Column(String(36), nullable=False, comment="数据集ID")
name = Column(String(100), nullable=False, comment="项目名称")
labeling_project_id = Column(String(8), nullable=False, comment="Label Studio项目ID")
template_id = Column(String(36), ForeignKey('t_dm_annotation_templates.id', ondelete='SET NULL'), nullable=True, comment="使用的模板ID")
configuration = Column(JSON, nullable=True, comment="项目配置(可能包含对模板的自定义修改)")
progress = Column(JSON, nullable=True, comment="项目进度信息")
deleted_at = Column(TIMESTAMP, nullable=True, comment="删除时间(软删除)")
def __repr__(self):
return f"<LabelingProject(id={self.id}, name={self.name}, dataset_id={self.dataset_id})>"
@property
def is_deleted(self) -> bool:
"""检查是否已被软删除"""
return self.deleted_at is not None
class AutoAnnotationTask(BaseEntity):
"""自动标注任务模型,对应表 t_dm_auto_annotation_tasks"""
__tablename__ = "t_dm_auto_annotation_tasks"
id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()), comment="UUID")
name = Column(String(255), nullable=False, comment="任务名称")
dataset_id = Column(String(36), nullable=False, comment="数据集ID")
dataset_name = Column(String(255), nullable=True, comment="数据集名称(冗余字段,方便查询)")
config = Column(JSON, nullable=False, comment="任务配置(模型规模、置信度等)")
file_ids = Column(JSON, nullable=True, comment="要处理的文件ID列表,为空则处理数据集所有图像")
status = Column(String(50), nullable=False, default="pending", comment="任务状态: pending/running/completed/failed")
progress = Column(Integer, default=0, comment="任务进度 0-100")
total_images = Column(Integer, default=0, comment="总图片数")
processed_images = Column(Integer, default=0, comment="已处理图片数")
detected_objects = Column(Integer, default=0, comment="检测到的对象总数")
output_path = Column(String(500), nullable=True, comment="输出路径")
error_message = Column(Text, nullable=True, comment="错误信息")
completed_at = Column(TIMESTAMP, nullable=True, comment="完成时间")
deleted_at = Column(TIMESTAMP, nullable=True, comment="删除时间(软删除)")
def __repr__(self) -> str: # pragma: no cover - repr 简单返回
return f"<AutoAnnotationTask(id={self.id}, name={self.name}, status={self.status})>"
@property
def is_deleted(self) -> bool:
"""检查是否已被软删除"""
return self.deleted_at is not None
@@ -0,0 +1,51 @@
from sqlalchemy import Column, String, TIMESTAMP, Text
from sqlalchemy.orm import declarative_base
from sqlalchemy.sql import func
Base = declarative_base()
class BaseEntity(Base):
"""
Common base entity with audit fields.
Subclasses may set `__ignore_data_scope__ = True` to opt-out of data-scope filtering.
"""
__abstract__ = True
created_at = Column(TIMESTAMP, server_default=func.current_timestamp(), comment="创建时间")
updated_at = Column(TIMESTAMP, server_default=func.current_timestamp(), onupdate=func.current_timestamp(),
comment="更新时间")
created_by = Column(String(255), nullable=True, comment="创建者")
updated_by = Column(String(255), nullable=True, comment="更新者")
# default: do enforce data scope unless subclass sets this to True
__ignore_data_scope__ = False
class LineageNode(Base):
"""数据血缘:节点表(实体对象)"""
__tablename__ = "t_lineage_node"
id = Column(String(36), primary_key=True, comment="节点ID")
graph_id = Column(String(36), nullable=True, comment="图ID")
node_type = Column(String(64), nullable=False, comment="节点类型")
name = Column(String(256), nullable=False, comment="节点名称")
description = Column(Text, nullable=True, comment="节点描述")
node_metadata = Column(Text, nullable=True, comment="节点扩展信息(JSON)")
class LineageEdge(Base):
"""数据血缘:边表(处理流程)"""
__tablename__ = "t_lineage_edge"
id = Column(String(36), primary_key=True, comment="边ID")
graph_id = Column(String(36), nullable=True, comment="图ID")
process_id = Column(String(36), nullable=True, comment="处理流程ID")
edge_type = Column(String(64), nullable=False, comment="边类型")
name = Column(String(256), nullable=True, comment="边名称")
description = Column(Text, nullable=True, comment="边描述")
edge_metadata = Column(Text, nullable=True, comment="边扩展信息(JSON)")
from_node_id = Column(String(36), nullable=False, comment="源节点ID")
to_node_id = Column(String(36), nullable=False, comment="目标节点ID")
@@ -0,0 +1,38 @@
"""
Chunk Upload Database Model
分片上传数据库模型
"""
from sqlalchemy import Column, String, Integer, DateTime
from sqlalchemy.sql import func
from app.db.models.base_entity import Base
class ChunkUploadPreRequest(Base):
"""分片上传预请求"""
__tablename__ = "t_chunk_upload_request"
id = Column(String(36), primary_key=True, comment="请求ID")
total_file_num = Column(Integer, nullable=False, comment="总文件数")
uploaded_file_num = Column(Integer, nullable=True, comment="已上传文件数")
upload_path = Column(String(512), nullable=False, comment="文件路径")
timeout = Column(DateTime, nullable=False, comment="上传请求超时时间")
service_id = Column(String(64), nullable=True, comment="上传请求所属服务ID")
check_info = Column(String(512), nullable=True, comment="业务信息")
def increment_uploaded_file_num(self):
"""增加已上传文件数"""
if self.uploaded_file_num is None:
self.uploaded_file_num = 1
else:
self.uploaded_file_num += 1
def is_upload_complete(self) -> bool:
"""检查是否已完成上传"""
return (self.uploaded_file_num is not None and
self.uploaded_file_num == self.total_file_num)
def is_request_timeout(self) -> bool:
"""检查是否已超时"""
from datetime import datetime
return self.timeout is not None and datetime.utcnow() > self.timeout
@@ -0,0 +1,69 @@
from sqlalchemy import Column, String, BigInteger, Integer, TIMESTAMP, Index
from app.db.models.base_entity import BaseEntity, Base
class CleaningTask(BaseEntity):
"""Data cleaning task entity"""
__tablename__ = "t_clean_task"
id = Column(String(36), primary_key=True, comment="Task ID")
name = Column(String(255), nullable=False, comment="Task name")
description = Column(String(1024), nullable=True, comment="Task description")
status = Column(String(50), nullable=False, default="PENDING", comment="Task status: PENDING, RUNNING, COMPLETED, STOPPED, FAILED")
src_dataset_id = Column(String(36), nullable=False, comment="Source dataset ID")
src_dataset_name = Column(String(255), nullable=False, comment="Source dataset name")
dest_dataset_id = Column(String(36), nullable=True, comment="Destination dataset ID")
dest_dataset_name = Column(String(255), nullable=True, comment="Destination dataset name")
before_size = Column(BigInteger, nullable=True, comment="Data size before cleaning")
after_size = Column(BigInteger, nullable=True, comment="Data size after cleaning")
file_count = Column(Integer, nullable=True, comment="Total file count")
retry_count = Column(Integer, default=0, nullable=False, comment="Retry count")
started_at = Column(TIMESTAMP, nullable=True, comment="Task start time")
finished_at = Column(TIMESTAMP, nullable=True, comment="Task finish time")
__table_args__ = (
Index("ix_clean_task_status", "status"),
Index("ix_clean_task_created_at", "created_at"),
)
class CleaningTemplate(BaseEntity):
"""Data cleaning template entity"""
__tablename__ = "t_clean_template"
id = Column(String(36), primary_key=True, comment="Template ID")
name = Column(String(255), nullable=False, comment="Template name")
description = Column(String(1024), nullable=True, comment="Template description")
class CleaningResult(Base):
"""Data cleaning result entity"""
__tablename__ = "t_clean_result"
instance_id = Column(String(36), primary_key=True, comment="Instance ID (task or template ID)")
src_file_id = Column(String(36), primary_key=True, comment="Source file ID")
dest_file_id = Column(String(36), nullable=True, comment="Destination file ID")
src_name = Column(String(512), nullable=True, comment="Source file name")
dest_name = Column(String(512), nullable=True, comment="Destination file name")
src_type = Column(String(50), nullable=True, comment="Source file type")
dest_type = Column(String(50), nullable=True, comment="Destination file type")
src_size = Column(BigInteger, nullable=True, comment="Source file size")
dest_size = Column(BigInteger, nullable=True, comment="Destination file size")
status = Column(String(50), nullable=True, comment="Cleaning status: COMPLETED, FAILED, etc.")
result = Column(String(1024), nullable=True, comment="Cleaning result message")
__table_args__ = (
Index("ix_clean_result_instance_id", "instance_id"),
Index("ix_clean_result_instance_status", "instance_id", "status"),
)
class OperatorInstance(Base):
"""Operator instance in task or template"""
__tablename__ = "t_operator_instance"
instance_id = Column(String(36), primary_key=True, comment="Instance ID (task or template ID)")
operator_id = Column(String(36), primary_key=True, comment="Operator ID")
op_index = Column(Integer, nullable=False, comment="Operator execution order")
settings_override = Column(String(4096), nullable=True, comment="Operator settings override (JSON)")
@@ -0,0 +1,56 @@
import uuid
from sqlalchemy import Column, String, Text, TIMESTAMP, Integer, BigInteger, Numeric, JSON, Boolean
from sqlalchemy.sql import func
from app.db.models.base_entity import BaseEntity
class CollectionTemplate(BaseEntity):
"""归集模板表(UUID 主键) -> t_dc_collection_templates"""
__tablename__ = "t_dc_collection_templates"
__ignore_data_scope__ = True
id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()), comment="模板ID(UUID)")
name = Column(String(255), nullable=False, comment="模板名称")
description = Column(Text, nullable=True, comment="模板描述")
source_type = Column(String(64), nullable=False, comment="源数据源类型")
source_name = Column(String(64), nullable=False, comment="源数据源名称")
target_type = Column(String(64), nullable=False, comment="目标数据源类型")
target_name = Column(String(64), nullable=False, comment="目标数据源名称")
template_content = Column(JSON, nullable=False, comment="模板内容")
built_in = Column(Boolean, default=False, comment="是否系统内置模板")
class CollectionTask(BaseEntity):
"""归集任务表(UUID 主键) -> t_dc_collection_tasks"""
__tablename__ = "t_dc_collection_tasks"
id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()), comment="UUID")
name = Column(String(255), nullable=False, comment="任务名称")
description = Column(Text, nullable=True, comment="任务描述")
sync_mode = Column(String(20), nullable=False, server_default="ONCE", comment="同步模式:ONCE/SCHEDULED")
template_id = Column(String(36), nullable=False, comment="归集模板ID")
template_name = Column(String(255), nullable=False, comment="归集模板名称")
target_path = Column(String(1000), nullable=True, server_default="", comment="目标存储路径")
config = Column(JSON, nullable=False, comment="归集配置(DataX配置),包含源端和目标端配置信息")
schedule_expression = Column(String(255), nullable=True, comment="Cron调度表达式")
status = Column(String(20), nullable=True, server_default="DRAFT", comment="任务状态:DRAFT/READY/RUNNING/SUCCESS/FAILED/STOPPED")
retry_count = Column(Integer, nullable=True, server_default="3", comment="重试次数")
timeout_seconds = Column(Integer, nullable=True, server_default="3600", comment="超时时间(秒)")
last_execution_id = Column(String(36), nullable=True, comment="最后执行ID(UUID)")
class TaskExecution(BaseEntity):
"""任务执行记录表(UUID 主键) -> t_dc_task_executions"""
__tablename__ = "t_dc_task_executions"
id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()), comment="执行记录ID(UUID)")
task_id = Column(String(36), nullable=False, comment="任务ID")
task_name = Column(String(255), nullable=False, comment="任务名称")
status = Column(String(20), nullable=True, server_default="RUNNING", comment="执行状态:RUNNING/SUCCESS/FAILED/STOPPED")
log_path = Column(String(1000), nullable=True, server_default="", comment="日志文件路径")
started_at = Column(TIMESTAMP, nullable=True, comment="开始时间")
completed_at = Column(TIMESTAMP, nullable=True, comment="完成时间")
duration_seconds = Column(Integer, nullable=True, server_default="0", comment="执行时长(秒)")
error_message = Column(Text, nullable=True, comment="错误信息")
@@ -0,0 +1,73 @@
"""
Tables for Data Evaluation module
Derived from scripts/db/data-evaluation-init.sql
- t_de_eval_task
- t_de_eval_item
"""
import uuid
from sqlalchemy import Column, String, Text, Float, TIMESTAMP, ForeignKey, Integer
from sqlalchemy.sql import func
from app.db.models.base_entity import BaseEntity
class EvaluationTask(BaseEntity):
"""评估任务表(UUID 主键) -> t_de_eval_task
Columns per data-evaluation-init.sql:
id, name, description, task_type, source_type, source_id, source_name,
status, eval_process, eval_promt, eval_config, created_at, updated_at,
created_by, updated_by
"""
__tablename__ = "t_de_eval_task"
id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()), comment="UUID")
name = Column(String(255), nullable=False, comment="评估任务名称")
description = Column(Text, nullable=True, comment="评估任务描述")
task_type = Column(String(50), nullable=False, comment="评估任务类型:QA")
source_type = Column(String(36), nullable=True, comment="待评估对象类型:DATASET/SYNTHESIS")
source_id = Column(String(36), nullable=True, comment="待评估对象ID")
source_name = Column(String(255), nullable=True, comment="待评估对象名称")
status = Column(String(50), server_default="PENDING", nullable=False, comment="状态:PENDING/RUNNING/COMPLETED/STOPPED/FAILED")
eval_method = Column(String(50), server_default="AUTO", nullable=False, comment="评估方式:AUTO/MANUAL")
eval_process = Column(Float, nullable=False, server_default="0", comment="评估进度")
eval_prompt = Column(Text, nullable=True, comment="评估提示词")
eval_config = Column(Text, nullable=True, comment="评估配置")
class EvaluationFile(BaseEntity):
"""评估条目表(UUID 主键) -> t_de_eval_file"""
__tablename__ = "t_de_eval_file"
__ignore_data_scope__ = True
id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()), comment="UUID")
task_id = Column(String(36), ForeignKey('t_de_eval_task.id'), nullable=False, comment="评估任务ID")
file_id = Column(String(36), ForeignKey('t_dm_dataset_files.id'), nullable=True, comment="文件ID")
file_name = Column(String(255), nullable=False, comment="文件名")
error_message = Column(Text, nullable=True, comment="错误信息")
total_count = Column(Integer, nullable=False, default=0, comment="总数")
evaluated_count = Column(Integer, nullable=False, default=0, comment="已评估数")
class EvaluationItem(BaseEntity):
"""评估条目表(UUID 主键) -> t_de_eval_item
Columns per data-evaluation-init.sql:
id, task_id, item_id, eval_score, eval_result, status
"""
__tablename__ = "t_de_eval_item"
__ignore_data_scope__ = True
id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()), comment="UUID")
task_id = Column(String(36), ForeignKey('t_de_eval_task.id'), nullable=False, comment="评估任务ID")
file_id = Column(String(36), ForeignKey('t_dm_dataset_files.id'), nullable=True, comment="文件ID")
item_id = Column(String(36), nullable=False, comment="评估条目ID")
eval_content = Column(Text, nullable=True, comment="评估内容")
eval_score = Column(Float, nullable=False, server_default="0", comment="评估分数")
eval_result = Column(Text, nullable=True, comment="评估结果")
status = Column(String(50), server_default="PENDING", nullable=False, comment="状态:PENDING/EVALUATED")
@@ -0,0 +1,177 @@
import uuid
from sqlalchemy import Column, String, Text, Integer, JSON, TIMESTAMP, func
from app.db.models.base_entity import Base, BaseEntity
from app.module.generation.schema.generation import CreateSynthesisTaskRequest
async def save_synthesis_task(db_session, synthesis_task: CreateSynthesisTaskRequest):
"""保存数据合成任务。
注意:当前 MySQL 表 `t_data_synth_instances` 结构中只包含 synth_type / synth_config 等字段,
没有 model_id、text_split_config、source_file_id、result_data_location 等列,因此这里只保存
与表结构一致的字段,其他信息由上层逻辑或其它表负责管理。
"""
gid = str(uuid.uuid4())
# 兼容旧请求结构:从请求对象中提取必要字段,
# - 合成类型:synthesis_type -> synth_type
# - 合成配置:text_split_config + synthesis_config 合并后写入 synth_config
synth_task_instance = DataSynthInstance(
id=gid,
name=synthesis_task.name,
description=synthesis_task.description,
status="pending",
synth_type=synthesis_task.synthesis_type.value,
progress=0,
synth_config=synthesis_task.synth_config.model_dump(),
total_files=len(synthesis_task.source_file_id or []),
processed_files=0,
total_chunks=0,
processed_chunks=0,
total_synth_data=0,
created_at=func.now(),
updated_at=func.now(),
created_by="system",
updated_by="system",
)
db_session.add(synth_task_instance)
await db_session.commit()
await db_session.refresh(synth_task_instance)
return synth_task_instance
class DataSynthInstance(BaseEntity):
"""数据合成任务表,对应表 t_data_synth_instances
create table if not exists t_data_synth_instances
(
id VARCHAR(36) CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci PRIMARY KEY COMMENT 'UUID',
name VARCHAR(255) NOT NULL COMMENT '任务名称',
description TEXT COMMENT '任务描述',
status VARCHAR(20) COMMENT '任务状态',
synth_type VARCHAR(20) NOT NULL COMMENT '合成类型',
progress INT DEFAULT 0 COMMENT '任务进度(百分比)',
synth_config JSON NOT NULL COMMENT '合成配置',
total_files INT DEFAULT 0 COMMENT '总文件数',
processed_files INT DEFAULT 0 COMMENT '已处理文件数',
total_chunks INT DEFAULT 0 COMMENT '总文本块数',
processed_chunks INT DEFAULT 0 COMMENT '已处理文本块数',
total_synth_data INT DEFAULT 0 COMMENT '总合成数据量',
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
created_by VARCHAR(255) COMMENT '创建者',
updated_by VARCHAR(255) COMMENT '更新者'
) COMMENT='数据合成任务表(UUID 主键)';
"""
__tablename__ = "t_data_synth_instances"
id = Column(String(36), primary_key=True, index=True, comment="UUID")
name = Column(String(255), nullable=False, comment="任务名称")
description = Column(Text, nullable=True, comment="任务描述")
status = Column(String(20), nullable=True, comment="任务状态")
# 与数据库字段保持一致:synth_type / synth_config
synth_type = Column(String(20), nullable=False, comment="合成类型")
progress = Column(Integer, nullable=False, default=0, comment="任务进度(百分比)")
synth_config = Column(JSON, nullable=False, comment="合成配置")
total_files = Column(Integer, nullable=False, default=0, comment="总文件数")
processed_files = Column(Integer, nullable=False, default=0, comment="已处理文件数")
total_chunks = Column(Integer, nullable=False, default=0, comment="总文本块数")
processed_chunks = Column(Integer, nullable=False, default=0, comment="已处理文本块数")
total_synth_data = Column(Integer, nullable=False, default=0, comment="总合成数据量")
class DataSynthesisFileInstance(BaseEntity):
"""数据合成文件任务表,对应表 t_data_synthesis_file_instances
create table if not exists t_data_synthesis_file_instances (
id VARCHAR(36) PRIMARY KEY COMMENT 'UUID',
synthesis_instance_id VARCHAR(36) COMMENT '数据合成任务ID',
file_name VARCHAR(255) NOT NULL COMMENT '文件名',
source_file_id VARCHAR(255) NOT NULL COMMENT '原始文件ID',
target_file_location VARCHAR(1000) NOT NULL COMMENT '目标文件存储位置',
status VARCHAR(20) COMMENT '任务状态',
total_chunks INT DEFAULT 0 COMMENT '总文本块数',
processed_chunks INT DEFAULT 0 COMMENT '已处理文本块数',
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
created_by VARCHAR(255) COMMENT '创建者',
updated_by VARCHAR(255) COMMENT '更新者'
) COMMENT='数据合成文件任务表(UUID 主键)';
"""
__tablename__ = "t_data_synthesis_file_instances"
id = Column(String(36), primary_key=True, index=True, comment="UUID")
synthesis_instance_id = Column(
String(36),
nullable=False,
comment="数据合成任务ID",
index=True,
)
file_name = Column(String(255), nullable=False, comment="文件名")
source_file_id = Column(String(255), nullable=False, comment="原始文件ID")
target_file_location = Column(String(1000), nullable=True, comment="目标文件存储位置")
status = Column(String(20), nullable=True, comment="任务状态")
total_chunks = Column(Integer, nullable=False, default=0, comment="总文本块数")
processed_chunks = Column(Integer, nullable=False, default=0, comment="已处理文本块数")
class DataSynthesisChunkInstance(Base):
"""数据合成分块任务表,对应表 t_data_synthesis_chunk_instances
create table if not exists t_data_synthesis_chunk_instances (
id VARCHAR(36) PRIMARY KEY COMMENT 'UUID',
synthesis_file_instance_id VARCHAR(36) COMMENT '数据合成文件任务ID',
chunk_index INT COMMENT '分块索引',
chunk_content TEXT COMMENT '分块内容',
metadata JSON COMMENT '分块元数据'
) COMMENT='数据合成分块任务表(UUID 主键)';
"""
__tablename__ = "t_data_synthesis_chunk_instances"
id = Column(String(36), primary_key=True, index=True, comment="UUID")
synthesis_file_instance_id = Column(
String(36),
nullable=False,
comment="数据合成文件任务ID",
index=True,
)
chunk_index = Column(Integer, nullable=True, comment="分块索引")
chunk_content = Column(Text, nullable=True, comment="分块内容")
# SQLAlchemy Declarative 保留了属性名 'metadata',这里使用 chunk_metadata 作为属性名,
# 底层列名仍为 'metadata' 以保持与表结构兼容。
chunk_metadata = Column("metadata", JSON, nullable=True, comment="分块元数据")
class SynthesisData(Base):
"""数据合成结果表,对应表 t_synthesis_data
create table if not exists t_synthesis_data (
id VARCHAR(36) PRIMARY KEY COMMENT 'UUID',
data json COMMENT '合成的数据',
synthesis_file_instance_id VARCHAR(36) COMMENT '数据合成文件任务ID',
chunk_instance_id VARCHAR(36) COMMENT '分块任务ID'
) COMMENT='数据合成任务队列表(UUID 主键)';
"""
__tablename__ = "t_data_synthesis_data"
id = Column(String(36), primary_key=True, index=True, comment="UUID")
data = Column(JSON, nullable=True, comment="合成的数据")
synthesis_file_instance_id = Column(
String(36),
nullable=False,
comment="数据合成文件任务ID",
index=True,
)
chunk_instance_id = Column(
String(36),
nullable=False,
comment="分块任务ID",
index=True,
)
@@ -0,0 +1,111 @@
"""
Tables of Dataset Management Module
"""
import uuid
from sqlalchemy import Column, String, BigInteger, Boolean, TIMESTAMP, Text, Integer, JSON, Date
from sqlalchemy.sql import func
from app.db.models.base_entity import Base, BaseEntity
class Dataset(BaseEntity):
"""数据集模型(支持医学影像、文本、问答等多种类型)"""
__tablename__ = "t_dm_datasets"
id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()), comment="UUID")
name = Column(String(255), nullable=False, comment="数据集名称")
description = Column(Text, nullable=True, comment="数据集描述")
dataset_type = Column(String(50), nullable=False, comment="数据集类型:IMAGE/TEXT/QA/MULTIMODAL/OTHER")
category = Column(String(100), nullable=True, comment="数据集分类:医学影像/问答/文献等")
path = Column(String(500), nullable=True, comment="数据存储路径")
format = Column(String(50), nullable=True, comment="数据格式:DCM/JPG/JSON/CSV等")
schema_info = Column(JSON, nullable=True, comment="数据结构信息")
size_bytes = Column(BigInteger, default=0, comment="数据大小(字节)")
file_count = Column(BigInteger, default=0, comment="文件数量")
record_count = Column(BigInteger, default=0, comment="记录数量")
retention_days = Column(Integer, default=0, comment="数据保留天数(0表示长期保留)")
tags = Column(JSON, nullable=True, comment="标签列表")
dataset_metadata = Column("metadata", JSON, nullable=True, comment="元数据信息")
status = Column(String(50), default='DRAFT', comment="状态:DRAFT/ACTIVE/ARCHIVED")
is_public = Column(Boolean, default=False, comment="是否公开")
is_featured = Column(Boolean, default=False, comment="是否推荐")
version = Column(BigInteger, nullable=False, default=0, comment="版本号")
def __repr__(self):
return f"<Dataset(id={self.id}, name={self.name}, type={self.dataset_type})>"
class DatasetTag(BaseEntity):
"""数据集标签关联模型"""
__tablename__ = "t_dm_dataset_tags"
dataset_id = Column(String(36), primary_key=True, comment="数据集ID(UUID)")
tag_id = Column(String(36), primary_key=True, comment="标签ID(UUID)")
def __repr__(self):
return f"<DatasetTag(dataset_id={self.dataset_id}, tag_id={self.tag_id})>"
class DatasetFiles(Base):
"""DM数据集文件模型"""
__tablename__ = "t_dm_dataset_files"
id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()), comment="UUID")
dataset_id = Column(String(36), nullable=False, comment="所属数据集ID(UUID)")
file_name = Column(String(255), nullable=False, comment="文件名")
file_path = Column(String(1000), nullable=False, comment="文件路径")
file_type = Column(String(50), nullable=True, comment="文件格式:JPG/PNG/DCM/TXT等")
file_size = Column(BigInteger, default=0, comment="文件大小(字节)")
check_sum = Column(String(64), nullable=True, comment="文件校验和")
tags = Column(JSON, nullable=True, comment="文件标签信息")
tags_updated_at = Column(TIMESTAMP, nullable=True, comment="标签最后更新时间")
annotation = Column(JSON, nullable=True, comment="完整标注结果(原始JSON)")
dataset_filemetadata = Column("metadata", JSON, nullable=True, comment="文件元数据")
status = Column(String(50), default='ACTIVE', comment="文件状态:ACTIVE/DELETED/PROCESSING")
upload_time = Column(TIMESTAMP, server_default=func.current_timestamp(), comment="上传时间")
last_access_time = Column(TIMESTAMP, nullable=True, comment="最后访问时间")
created_at = Column(TIMESTAMP, server_default=func.current_timestamp(), comment="创建时间")
updated_at = Column(TIMESTAMP, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), comment="更新时间")
def __repr__(self):
return f"<DatasetFiles(id={self.id}, dataset_id={self.dataset_id}, file_name={self.file_name})>"
class DatasetStatistics(Base):
"""数据集统计信息模型"""
__tablename__ = "t_dm_dataset_statistics"
id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()), comment="UUID")
dataset_id = Column(String(36), nullable=False, comment="数据集ID(UUID)")
stat_date = Column(Date, nullable=False, comment="统计日期")
total_files = Column(BigInteger, default=0, comment="总文件数")
total_size = Column(BigInteger, default=0, comment="总大小(字节)")
processed_files = Column(BigInteger, default=0, comment="已处理文件数")
error_files = Column(BigInteger, default=0, comment="错误文件数")
download_count = Column(BigInteger, default=0, comment="下载次数")
view_count = Column(BigInteger, default=0, comment="查看次数")
quality_metrics = Column(JSON, nullable=True, comment="质量指标")
created_at = Column(TIMESTAMP, server_default=func.current_timestamp(), comment="创建时间")
updated_at = Column(TIMESTAMP, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), comment="更新时间")
def __repr__(self):
return f"<DatasetStatistics(id={self.id}, dataset_id={self.dataset_id}, date={self.stat_date})>"
class Tag(Base):
"""标签集合模型"""
__tablename__ = "t_dm_tags"
id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()), comment="UUID")
name = Column(String(100), nullable=False, unique=True, comment="标签名称")
description = Column(Text, nullable=True, comment="标签描述")
category = Column(String(50), nullable=True, comment="标签分类")
color = Column(String(7), nullable=True, comment="标签颜色(十六进制)")
usage_count = Column(BigInteger, default=0, comment="使用次数")
created_at = Column(TIMESTAMP, server_default=func.current_timestamp(), comment="创建时间")
updated_at = Column(TIMESTAMP, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), comment="更新时间")
def __repr__(self):
return f"<Tag(id={self.id}, name={self.name}, category={self.category})>"
@@ -0,0 +1,82 @@
"""
知识库(RAG)相关 ORM 模型
表: t_rag_knowledge_base, t_rag_file
与 Java 实体保持一致。
"""
from enum import Enum
from sqlalchemy import Column, String, Integer, JSON
from app.db.models.base_entity import BaseEntity
class RagType(str, Enum):
"""RAG 类型枚举
对应 Java: com.deepdata.rag.indexer.interfaces.dto.RagType
"""
DOCUMENT = "DOCUMENT" # 文档型 RAG(向量检索)
GRAPH = "GRAPH" # 知识图谱型 RAG(LightRAG)
class FileStatus(str, Enum):
"""文件状态枚举
对应 Java: com.deepdata.rag.indexer.domain.model.FileStatus
"""
UNPROCESSED = "UNPROCESSED" # 未处理
PROCESSING = "PROCESSING" # 处理中
PROCESSED = "PROCESSED" # 已处理
PROCESS_FAILED = "PROCESS_FAILED" # 处理失败
class KnowledgeBase(BaseEntity):
"""知识库实体
对应 Java: com.deepdata.rag.indexer.domain.model.KnowledgeBase
表名: t_rag_knowledge_base
"""
__tablename__ = "t_rag_knowledge_base"
__ignore_data_scope__ = True
id = Column(String(36), primary_key=True, comment="知识库ID")
name = Column(String(255), nullable=False, unique=True, comment="知识库名称")
description = Column(String(512), nullable=True, comment="知识库描述")
type = Column(
String(50),
nullable=False,
default=RagType.DOCUMENT,
comment="RAG类型",
)
embedding_model = Column(String(255), nullable=False, comment="嵌入模型ID")
chat_model = Column(String(255), nullable=True, comment="聊天模型ID")
def __repr__(self):
return f"<KnowledgeBase(id={self.id}, name={self.name}, type={self.type})>"
class RagFile(BaseEntity):
"""RAG 文件实体
对应 Java: com.deepdata.rag.indexer.domain.model.RagFile
表名: t_rag_file
"""
__tablename__ = "t_rag_file"
__ignore_data_scope__ = True
id = Column(String(36), primary_key=True, comment="RAG文件ID")
knowledge_base_id = Column(String(36), nullable=False, index=True, comment="知识库ID")
file_name = Column(String(512), nullable=False, comment="文件名")
file_id = Column(String(36), nullable=False, comment="原始文件ID")
chunk_count = Column(Integer, nullable=True, comment="分块数量")
file_metadata = Column("metadata", JSON, nullable=True, comment="元数据(JSON格式)")
status = Column(
String(50),
nullable=False,
default=FileStatus.UNPROCESSED,
comment="处理状态",
)
err_msg = Column(String(2048), nullable=True, comment="错误信息")
progress = Column(Integer, default=0, nullable=False, comment="处理进度(0-100)")
def __repr__(self):
return f"<RagFile(id={self.id}, file_name={self.file_name}, status={self.status})>"
@@ -0,0 +1,46 @@
from sqlalchemy import Boolean, Column, String, TIMESTAMP
from app.db.models.base_entity import BaseEntity
class Models(BaseEntity):
"""模型配置表,对应表 t_models。模型为系统级配置,RAG/生成等按 ID 引用时不受数据权限过滤。
CREATE TABLE IF NOT EXISTS t_models (
id VARCHAR(36) PRIMARY KEY COMMENT '主键ID',
model_name VARCHAR(100) NOT NULL COMMENT '模型名称(如 qwen2)',
provider VARCHAR(50) NOT NULL COMMENT '模型提供商(如 Ollama、OpenAI、DeepSeek)',
base_url VARCHAR(255) NOT NULL COMMENT 'API 基础地址',
api_key VARCHAR(512) DEFAULT '' COMMENT 'API 密钥(无密钥则为空)',
type VARCHAR(50) NOT NULL COMMENT '模型类型(如 chat、embedding)',
is_enabled TINYINT DEFAULT 1 COMMENT '是否启用:1-启用,0-禁用',
is_default TINYINT DEFAULT 0 COMMENT '是否默认:1-默认,0-非默认',
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
created_by VARCHAR(255) COMMENT '创建者',
updated_by VARCHAR(255) COMMENT '更新者',
UNIQUE KEY uk_model_provider (model_name, provider)
) ENGINE = InnoDB DEFAULT CHARSET = utf8mb4 COMMENT ='模型配置表';
"""
__tablename__ = "t_models"
id = Column(String(36), primary_key=True, index=True, comment="主键ID")
model_name = Column(String(100), nullable=False, comment="模型名称(如 qwen2)")
provider = Column(String(50), nullable=False, comment="模型提供商(如 Ollama、OpenAI、DeepSeek)")
base_url = Column(String(255), nullable=False, comment="API 基础地址")
api_key = Column(String(512), nullable=False, default="", comment="API 密钥(无密钥则为空)")
type = Column(String(50), nullable=False, comment="模型类型(如 chat、embedding)")
is_enabled = Column(Boolean, nullable=False, default=True, comment="是否启用")
is_default = Column(Boolean, nullable=False, default=False, comment="是否默认")
is_deleted = Column(Boolean, nullable=False, default=False, comment="是否删除")
__table_args__ = (
# 与 DDL 中的 uk_model_provider 保持一致
{
"mysql_engine": "InnoDB",
"mysql_charset": "utf8mb4",
"comment": "模型配置表",
},
)
@@ -0,0 +1,86 @@
"""
Operator Market Data Models
算子市场数据模型
"""
from sqlalchemy import (
Column,
String,
Integer,
Boolean,
BigInteger,
Text,
JSON,
TIMESTAMP,
Index,
)
from sqlalchemy.sql import func
from app.db.models.base_entity import Base, BaseEntity
class Operator(BaseEntity):
"""算子实体"""
__tablename__ = "t_operator"
id = Column(String(36), primary_key=True, index=True, comment="算子ID")
name = Column(String(255), nullable=False, comment="算子名称")
description = Column(Text, nullable=True, comment="算子描述")
version = Column(String(50), nullable=False, comment="算子版本")
inputs = Column(Text, nullable=True, comment="输入定义(JSON)")
outputs = Column(Text, nullable=True, comment="输出定义(JSON)")
runtime = Column(Text, nullable=True, comment="运行时配置(JSON)")
settings = Column(Text, nullable=True, comment="算子设置(JSON)")
file_name = Column(String(255), nullable=True, comment="文件名")
file_size = Column(BigInteger, nullable=True, comment="文件大小(字节)")
metrics = Column(Text, nullable=True, comment="算子指标(JSON)")
usage_count = Column(Integer, default=0, nullable=False, comment="使用次数")
is_star = Column(Boolean, default=False, nullable=False, comment="是否收藏")
__table_args__ = (Index("idx_is_star", "is_star"),)
class Category(BaseEntity):
"""算子分类实体"""
__tablename__ = "t_operator_category"
id = Column(String(36), primary_key=True, index=True, comment="分类ID")
name = Column(String(255), nullable=False, comment="分类名称(中文)")
name_en = Column(String(255), nullable=True, comment="分类名称(英文)")
value = Column(String(255), nullable=True, comment="分类值")
type = Column(String(50), nullable=True, comment="分类类型")
parent_id = Column(String(36), nullable=False, default="0", comment="父分类ID")
class CategoryRelation(BaseEntity):
"""算子分类关系实体"""
__tablename__ = "t_operator_category_relation"
category_id = Column(String(36), primary_key=True, comment="分类ID")
operator_id = Column(String(36), primary_key=True, comment="算子ID")
__table_args__ = (
Index("idx_category_id", "category_id"),
Index("idx_operator_id", "operator_id"),
)
class OperatorRelease(BaseEntity):
"""算子发布版本实体"""
__tablename__ = "t_operator_release"
id = Column(String(36), primary_key=True, comment="算子ID")
version = Column(String(50), primary_key=True, comment="版本号")
release_date = Column(
TIMESTAMP, nullable=False, default=func.now(), comment="发布时间"
)
changelog = Column(JSON, nullable=True, comment="更新日志列表")
# Ignore data scope for operator models
for model in [Operator, Category, CategoryRelation, OperatorRelease]:
model.__ignore_data_scope__ = True
@@ -0,0 +1,62 @@
"""
Tables for Ratio (Data Synthesis Ratio) module
Derived from scripts/db/data-ratio-init.sql
- t_st_ratio_instances
- t_st_ratio_relations
"""
import uuid
from sqlalchemy import Column, String, Text, BigInteger, TIMESTAMP, JSON, ForeignKey
from sqlalchemy.orm import relationship
from sqlalchemy.sql import func
from app.db.models.base_entity import BaseEntity
class RatioInstance(BaseEntity):
"""配比实例表(UUID 主键) -> t_st_ratio_instances
Columns per data-ratio-init.sql:
id, name, description, target_dataset_id, ratio_method, ratio_parameters,
merge_method, status, totals, created_at, updated_at, created_by, updated_by
"""
__tablename__ = "t_st_ratio_instances"
id = Column(String(64), primary_key=True, default=lambda: str(uuid.uuid4()), comment="UUID")
name = Column(String(64), nullable=True, comment="名称")
description = Column(Text, nullable=True, comment="描述")
target_dataset_id = Column(String(64), nullable=True, comment="模板数据集ID")
ratio_parameters = Column(JSON, nullable=True, comment="配比参数")
merge_method = Column(String(50), nullable=True, comment="合并方式")
status = Column(String(20), nullable=True, comment="状态")
totals = Column(BigInteger, nullable=True, comment="总数")
def __repr__(self) -> str:
return f"<RatioInstance(id={self.id}, name={self.name}, status={self.status})>"
class RatioRelation(BaseEntity):
"""配比关系表(UUID 主键) -> t_st_ratio_relations
Columns per data-ratio-init.sql:
id, ratio_instance_id, source_dataset_id, ratio_value, counts, filter_conditions,
created_at, updated_at, created_by, updated_by
"""
__tablename__ = "t_st_ratio_relations"
id = Column(String(64), primary_key=True, default=lambda: str(uuid.uuid4()), comment="UUID")
ratio_instance_id = Column(String(64), nullable=False, comment="配比实例ID")
source_dataset_id = Column(String(64), nullable=True, comment="源数据集ID")
ratio_value = Column(String(256), nullable=True)
counts = Column(BigInteger, nullable=True, comment="条数")
filter_conditions = Column(Text, nullable=True, comment="过滤条件")
def __repr__(self) -> str:
return (
f"<RatioRelation(id={self.id}, ratio_instance_id={self.ratio_instance_id}, "
f"source_dataset_id={self.source_dataset_id}, counts={self.counts})>"
)
@@ -0,0 +1,30 @@
"""
System Parameter Model
系统参数模型
"""
from sqlalchemy import Column, String, Boolean, Text
from app.db.models.base_entity import BaseEntity
class SysParam(BaseEntity):
"""系统参数实体"""
__tablename__ = "t_sys_param"
__ignore_data_scope__ = True # 系统参数不进行数据权限过滤
id = Column(String(100), primary_key=True, comment="参数ID")
param_value = Column(Text, nullable=False, comment="参数值")
param_type = Column(
String(50),
nullable=False,
default="string",
comment="参数类型(string、integer、boolean)",
)
option_list = Column(
Text, nullable=True, comment="选项列表(JSON格式,仅对enum类型有效)"
)
description = Column(String(255), nullable=True, default="", comment="参数描述")
is_built_in = Column(Boolean, nullable=False, default=False, comment="是否内置")
can_modify = Column(Boolean, nullable=False, default=True, comment="是否可修改")
is_enabled = Column(Boolean, nullable=False, default=True, comment="是否启用")
@@ -0,0 +1,29 @@
"""
Tables of User Management Module
"""
from sqlalchemy import Column, String, BigInteger, Boolean, TIMESTAMP
from sqlalchemy.sql import func
from app.db.models.base_entity import Base
class User(Base):
"""用户模型"""
__tablename__ = "users"
id = Column(BigInteger, primary_key=True, autoincrement=True, comment="用户ID")
username = Column(String(255), nullable=False, unique=True, comment="用户名")
email = Column(String(255), nullable=False, unique=True, comment="邮箱")
password_hash = Column(String(255), nullable=False, comment="密码哈希")
full_name = Column(String(255), nullable=True, comment="真实姓名")
avatar_url = Column(String(500), nullable=True, comment="头像URL")
role = Column(String(50), nullable=False, default='USER', comment="角色:ADMIN/USER")
organization = Column(String(255), nullable=True, comment="所属机构")
enabled = Column(Boolean, nullable=False, default=True, comment="是否启用")
last_login_at = Column(TIMESTAMP, nullable=True, comment="最后登录时间")
created_at = Column(TIMESTAMP, server_default=func.current_timestamp(), comment="创建时间")
updated_at = Column(TIMESTAMP, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), comment="更新时间")
def __repr__(self):
return f"<User(id={self.id}, username={self.username}, role={self.role})>"
+94
View File
@@ -0,0 +1,94 @@
from datetime import datetime
from sqlalchemy import event, true
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
from sqlalchemy.orm import Session, declarative_base, with_loader_criteria
from sqlalchemy.orm.attributes import flag_modified
from app.core.config import settings
from app.core.logging import get_logger
from typing import AsyncGenerator
from app.db.datascope import DataScopeHandle, SYSTEM_USER
from app.db.models.base_entity import BaseEntity
logger = get_logger(__name__)
# 创建数据库引擎
engine = create_async_engine(
settings.database_url,
echo=False, # 关闭SQL调试日志以减少输出
future=True,
pool_pre_ping=True,
pool_recycle=3600
)
# 创建会话工厂
AsyncSessionLocal = async_sessionmaker(
engine,
class_=AsyncSession,
expire_on_commit=False
)
@event.listens_for(Session, "do_orm_execute")
def _apply_data_scope(orm_execute_state):
# only act on SELECT queries generated by ORM
if not orm_execute_state.is_select:
return
allowed = DataScopeHandle.allowed_users()
if not allowed or len(allowed) == 0:
return
# predicate builder: return None to skip for classes that opt-out
def criteria_fn(cls):
# skip if the mapped class explicitly disables data-scope
if getattr(cls, "__ignore_data_scope__", False):
return true()
# some classes may not have created_by column; guard dynamically
col = getattr(cls, "created_by", None)
if col is None:
return true()
return col.in_(allowed)
# apply loader-level criteria to all subclasses of BaseEntity
orm_execute_state.statement = orm_execute_state.statement.options(
with_loader_criteria(BaseEntity, criteria_fn, include_aliases=True)
)
@event.listens_for(Session, "before_flush")
def _audit_before_flush(session, flush_context, instances):
user = DataScopeHandle.get_user_info()
effective_user = (user or "").strip() or SYSTEM_USER
now = datetime.now()
# new -> set created_* and updated_*
for obj in list(session.new):
if isinstance(obj, BaseEntity):
if getattr(obj, "created_at", None) is None:
obj.created_at = now
if getattr(obj, "created_by", None) is None:
obj.created_by = effective_user
obj.updated_at = now
obj.updated_by = effective_user
# ensure SQLAlchemy sees changes
try:
flag_modified(obj, "created_by")
except Exception:
pass
# dirty -> set updated_*
for obj in list(session.dirty):
if isinstance(obj, BaseEntity):
obj.updated_at = now
obj.updated_by = effective_user
try:
flag_modified(obj, "updated_by")
except Exception:
pass
async def get_db() -> AsyncGenerator[AsyncSession, None]:
"""获取数据库会话"""
async with AsyncSessionLocal() as session:
try:
yield session
finally:
await session.close()
+165
View File
@@ -0,0 +1,165 @@
from contextlib import asynccontextmanager
from typing import Literal
from urllib.parse import urlparse, urlunparse
from fastapi import FastAPI, Depends
from fastapi_mcp import FastApiMCP
from fastapi.responses import RedirectResponse
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.config import settings
from app.core.exception import (
register_exception_handlers,
SuccessResponse,
ExceptionHandlingMiddleware,
ErrorCodes,
BusinessError,
)
from app.core.logging import setup_logging, get_logger
from app.db.session import AsyncSessionLocal, get_db
from app.middleware import UserContextMiddleware
from app.module import router
from app.module.collection.schedule import (
load_scheduled_collection_tasks,
set_collection_scheduler,
)
from app.module.shared.schedule import Scheduler
from app.module.generation.service.task_executor import init_executor, shutdown_executor
from app.module.system.service.log_pvc_monitor import schedule_log_pvc_monitor
setup_logging()
logger = get_logger(__name__)
@asynccontextmanager
async def lifespan(app: FastAPI):
# @startup
logger.info("DeepData Python Backend starting...")
# Database connection validation
try:
async with AsyncSessionLocal() as session:
await session.execute(text("SELECT 1"))
def mask_db_url(url: str) -> Literal[b""] | str:
parsed = urlparse(url)
if parsed.password is not None:
# 重建 netloc,将密码替换为 ****
netloc = f"{parsed.username}:*****@{parsed.hostname}"
if parsed.port:
netloc += f":{parsed.port}"
return urlunparse(parsed._replace(netloc=netloc))
return url
# 使用示例
logger.info(f"Database: {mask_db_url(settings.database_url)}")
except Exception as e:
logger.error(f"Database connection validation failed: {e}")
logger.debug(f"Connection details: {settings.database_url}")
raise
# Label Studio
# TODO Add actual connectivity check if needed
logger.info(f"Label Studio: {settings.label_studio_base_url}")
# Collection scheduler
collection_scheduler = Scheduler(name="collection scheduler")
collection_scheduler.start()
set_collection_scheduler(collection_scheduler)
await load_scheduled_collection_tasks()
schedule_log_pvc_monitor(collection_scheduler)
# Initialize generation task executor
init_executor(max_workers=10, max_concurrent_tasks=5)
logger.info("Generation task executor initialized")
# 恢复未完成任务的轮询 + 启动全局状态轮询协程
from app.module.cleaning.service.cleaning_task_scheduler import get_scheduler
scheduler = get_scheduler()
await scheduler.startup()
logger.info("Cleaning task status polling started")
yield
# @shutdown
collection_scheduler.shutdown()
shutdown_executor()
logger.info("DeepData Python Backend shutting down ...\n\n")
# 创建FastAPI应用
app = FastAPI(
title=settings.app_name,
description=settings.app_description,
version=settings.app_version,
debug=settings.debug,
lifespan=lifespan,
)
# 注册全局异常捕获中间件(最外层,确保捕获所有异常)
# 这样即使 debug=True,也不会泄露堆栈信息给客户端
app.add_middleware(ExceptionHandlingMiddleware)
# 注册用户上下文中间件
app.add_middleware(UserContextMiddleware)
# 注册路由
app.include_router(router)
# 注册全局异常处理器
register_exception_handlers(app)
# 测试端点:验证异常处理
@app.get("/test-success", include_in_schema=False)
async def test_success():
"""测试成功响应"""
return SuccessResponse(data={"message": "Test successful"})
@app.get("/test-business-error", include_in_schema=False)
async def test_business_error():
"""测试业务错误响应"""
raise BusinessError(ErrorCodes.ANNOTATION_TASK_NOT_FOUND)
@app.get("/test-system-error", include_in_schema=False)
async def test_system_error():
"""测试系统错误响应"""
raise SystemError(ErrorCodes.DATABASE_ERROR)
# 根路径
@app.get("/", include_in_schema=False)
async def root():
"""根路径,返回服务信息"""
return SuccessResponse(
data={
"message": f"{settings.app_name} is running",
"version": settings.app_version,
"docs_url": "/redoc",
"label_studio_url": settings.label_studio_base_url
}
)
mcp = FastApiMCP(
app,
name="DeepData MCP",
description="DeepData python mcp server",
include_tags=["mcp"],
)
mcp.mount_http(mount_path="/api/mcp")
if __name__ == "__main__":
import uvicorn
uvicorn.run(
"app.main:app",
host=settings.host,
port=settings.port,
reload=settings.debug,
log_level=settings.log_level.lower(),
)
+53
View File
@@ -0,0 +1,53 @@
"""
Application middleware.
Note: Exception handling has been moved to app.core.exception.ExceptionHandlingMiddleware
This file now only contains the UserContextMiddleware.
"""
import json
from typing import Optional
from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.status import HTTP_401_UNAUTHORIZED
from app.core.config import settings
from app.core.logging import get_logger
from app.db.datascope import DataScopeHandle
logger = get_logger(__name__)
class UserContextMiddleware(BaseHTTPMiddleware):
"""
读取 `User` 请求头并设置 DataScopeHandle 的 FastAPI 中间件。
如果 `jwt_enable` 为 True,缺少请求头将返回 401。
"""
def __init__(self, app):
super().__init__(app)
self.jwt_enable = settings.deepdata_jwt_enable
async def dispatch(self, request: Request, call_next):
user: Optional[str] = request.headers.get("User")
logger.info(f"start filter, current user: {user}, need filter: {self.jwt_enable}")
if self.jwt_enable and (user is None or user.strip() == ""):
payload = {"code": "common.401", "message": "unauthorized", "data": None}
return Response(
content=json.dumps(payload),
status_code=HTTP_401_UNAUTHORIZED,
media_type="application/json"
)
DataScopeHandle.set_user_info(user)
try:
response = await call_next(request)
return response
finally:
DataScopeHandle.remove_user_info()
# Re-export ExceptionHandlingMiddleware for backward compatibility
__all__ = ['UserContextMiddleware']
@@ -0,0 +1,29 @@
from fastapi import APIRouter
from .system.interface import router as system_router
from .annotation.interface import router as annotation_router
from .ratio.interface import router as ratio_router
from .generation.interface import router as generation_router
from .evaluation.interface import router as evaluation_router
from .collection.interface import router as collection_route
from .operator.interface import operator_router
from .operator.interface import category_router
from .cleaning.interface import router as cleaning_router
from .rag.interface.knowledge_base import router as knowledge_base_router
router = APIRouter(
prefix="/api"
)
router.include_router(system_router)
router.include_router(annotation_router)
router.include_router(ratio_router)
router.include_router(generation_router)
router.include_router(evaluation_router)
router.include_router(collection_route)
router.include_router(operator_router)
router.include_router(category_router)
router.include_router(cleaning_router)
router.include_router(knowledge_base_router)
__all__ = ["router"]
@@ -0,0 +1,3 @@
from .labelstudio import LabelStudioClient
__all__ = ["LabelStudioClient"]
@@ -0,0 +1,3 @@
from .client import Client as LabelStudioClient
__all__ = ["LabelStudioClient"]
@@ -0,0 +1,824 @@
import httpx
import re
from typing import Optional, Dict, Any, List
from app.core.config import settings
from app.core.logging import get_logger
from .schema import (
LabelStudioProject,
LabelStudioCreateProjectRequest,
LabelStudioCreateTaskRequest
)
logger = get_logger(__name__)
class Client:
"""Label Studio服务客户端
使用 HTTP REST API 直接与 Label Studio 交互
认证方式:使用 Authorization: Token {token} 头部进行认证
"""
# 默认标注配置模板
DEFAULT_LABEL_CONFIGS = {
"image": """
<View>
<Image name="image" value="$image"/>
<RectangleLabels name="label" toName="image">
<Label value="Object" background="red"/>
</RectangleLabels>
</View>
""",
"text": """
<View>
<Text name="text" value="$text"/>
<Choices name="sentiment" toName="text">
<Choice value="positive"/>
<Choice value="negative"/>
<Choice value="neutral"/>
</Choices>
</View>
""",
"audio": """
<View>
<Audio name="audio" value="$audio"/>
<AudioRegionLabels name="label" toName="audio">
<Label value="Speech" background="red"/>
<Label value="Noise" background="blue"/>
</AudioRegionLabels>
</View>
""",
"video": """
<View>
<Video name="video" value="$video"/>
<VideoRegionLabels name="label" toName="video">
<Label value="Action" background="red"/>
</VideoRegionLabels>
</View>
"""
}
def __init__(
self,
base_url: Optional[str] = None,
token: Optional[str] = None,
timeout: float = 30.0
):
"""初始化 Label Studio 客户端
Args:
base_url: Label Studio 服务地址
token: API Token(使用 Authorization: Token {token} 头部)
timeout: 请求超时时间(秒)
"""
self.base_url = (
base_url or settings.label_studio_base_url).rstrip("/")
self.token = token or settings.label_studio_user_token
self.timeout = timeout
if not self.token:
raise ValueError("Label Studio API token is required")
# 初始化 HTTP 客户端
self.client = httpx.AsyncClient(
base_url=self.base_url,
timeout=self.timeout,
headers={
"Authorization": f"Token {self.token}",
"Content-Type": "application/json"
}
)
logger.debug(f"Label Studio client initialized: {self.base_url}")
def get_label_config_by_type(self, data_type: str) -> str:
"""根据数据类型获取标注配置"""
return self.DEFAULT_LABEL_CONFIGS.get(data_type.lower(), self.DEFAULT_LABEL_CONFIGS["image"])
@staticmethod
def get_csrf_token(html: str) -> str:
m = re.search(r'name="csrfmiddlewaretoken"\s+value="([^"]+)"', html)
if not m:
raise IOError("CSRF Token not found")
return m.group(1)
# xuhuanqing
async def login_label_studio(self):
try:
async with httpx.AsyncClient(
base_url=self.base_url,
timeout=self.timeout,
follow_redirects=False,
) as session:
# 1. GET 登录页
get_resp = await session.get("/user/login/")
get_resp.raise_for_status()
# 2. 拿 CSRF(正则兜底 + cookie 兜底)
m = re.search(r'name="csrfmiddlewaretoken"\s+value="([^"]+)"', get_resp.text) \
or re.search(r'value="([^"]+)"\s+name="csrfmiddlewaretoken"', get_resp.text)
csrf = m.group(1) if m else None
if not csrf:
csrf = session.cookies.get("csrftoken")
if not csrf:
logger.error("CSRF token not found in LS login page")
return None
# 3. POST 登录(username 字段,form-urlencoded)
login_resp = await session.post(
"/user/login/",
data={
"username": settings.label_studio_username,
"password": settings.label_studio_password,
"csrfmiddlewaretoken": csrf,
},
headers={
"Content-Type": "application/x-www-form-urlencoded",
"Referer": f"{self.base_url}/user/login/",
"Cookie": "; ".join(get_resp.headers.get_list("set-cookie")),
},
follow_redirects=False,
)
logger.info(
f"LS login resp: status={login_resp.status_code} "
f"set-cookie={login_resp.headers.get_list('set-cookie')}"
)
set_cookies = login_resp.headers.get_list("set-cookie")
has_session = any("sessionid=" in c for c in set_cookies)
if login_resp.status_code in (302, 200) and has_session:
return login_resp
logger.error("LS session login failed: no sessionid obtained")
return None
except httpx.HTTPStatusError as e:
logger.error(
f"LS login HTTP {e.response.status_code}: {e.response.text[:500]}")
return None
except Exception as e:
logger.error(f"Error while login: {e}", exc_info=True)
return None
'''
async def login_label_studio(self):
try:
response = await self.client.get("/user/login/")
response.raise_for_status()
body = response.text
set_cookie_headers = response.headers.get_list("set-cookie")
cookie_header = "; ".join(set_cookie_headers)
form = {
"email": settings.label_studio_username,
"password": settings.label_studio_password,
"csrfmiddlewaretoken": self.get_csrf_token(body),
}
headers = {
"Content-Type": "application/x-www-form-urlencoded",
"Cookie": cookie_header,
}
login_response = await self.client.post("/user/login/", data=form, headers=headers)
logger.info(f"response is: {login_response}, {login_response.text}")
return login_response
except httpx.HTTPStatusError as e:
logger.error(f"Login failed HTTP {e.response.status_code}: {e.response.text}")
return None
except Exception as e:
logger.error(f"Error while login: {e}", e)
return None
'''
async def create_project(
self,
title: str,
description: str = "",
label_config: Optional[str] = None,
) -> Optional[Dict[str, Any]]:
"""创建Label Studio项目"""
try:
logger.debug(f"Creating Label Studio project: {title}")
logger.debug(f"Label Studio URL: {self.base_url}/api/projects")
project_data = {
"title": title,
"description": description,
"label_config": label_config or "<View></View>"
}
# Log the request body for debugging
logger.debug(f"Request body: {project_data}")
logger.debug(
f"Label config being sent:\n{project_data['label_config']}")
response = await self.client.post("/api/projects", json=project_data)
response.raise_for_status()
project = response.json()
project_id = project.get("id")
if not project_id:
raise Exception(
"Label Studio response does not contain project ID")
logger.debug(f"Project created successfully, ID: {project_id}")
return project
except httpx.HTTPStatusError as e:
logger.error(
f"Create project failed - HTTP {e.response.status_code}\n"
f"URL: {e.request.url}\n"
f"Request Body: {e.request.content.decode() if e.request.content else 'None'}\n"
f"Response Body: {e.response.text[:1000]}" # First 1000 chars
)
return None
except httpx.ConnectError as e:
logger.error(
f"Failed to connect to Label Studio at {self.base_url}\n"
f"Error: {str(e)}\n"
f"Possible causes:\n"
f" - Label Studio service is not running\n"
f" - Incorrect URL configuration\n"
f" - Network connectivity issue"
)
return None
except httpx.TimeoutException as e:
logger.error(
f"Request to Label Studio timed out after {self.timeout}s: {str(e)}")
return None
except Exception as e:
logger.error(
f"Error while creating Label Studio project: {str(e)}", exc_info=True)
return None
async def import_tasks(
self,
project_id: int,
tasks: List[Dict[str, Any]],
commit_to_project: bool = True,
return_task_ids: bool = True
) -> Optional[Dict[str, Any]]:
"""批量导入任务到Label Studio项目"""
try:
logger.debug(
f"Importing {len(tasks)} tasks into project {project_id}")
response = await self.client.post(
f"/api/projects/{project_id}/import",
json=tasks,
params={
"commit_to_project": str(commit_to_project).lower(),
"return_task_ids": str(return_task_ids).lower()
}
)
response.raise_for_status()
result = response.json()
task_count = result.get("task_count", len(tasks))
logger.debug(f"Tasks imported successfully: {task_count}")
return result
except httpx.HTTPStatusError as e:
logger.error(
f"Import tasks failed HTTP {e.response.status_code}: {e.response.text}")
return None
except Exception as e:
logger.error(f"Error while importing tasks: {e}")
return None
async def create_tasks_batch(
self,
project_id: str,
tasks: List[Dict[str, Any]]
) -> Optional[Dict[str, Any]]:
"""批量创建任务的便利方法"""
try:
pid = int(project_id)
return await self.import_tasks(pid, tasks)
except ValueError as e:
logger.error(
f"Invalid project ID format: {project_id}, error: {e}")
return None
except Exception as e:
logger.error(f"Error while creating tasks in batch: {e}")
return None
async def create_task(
self,
project_id: str,
data: Dict[str, Any],
meta: Optional[Dict[str, Any]] = None
) -> Optional[Dict[str, Any]]:
"""创建单个任务"""
try:
task = {"data": data}
if meta:
task["meta"] = meta
return await self.create_tasks_batch(project_id, [task])
except Exception as e:
logger.error(f"Error while creating single task: {e}")
return None
async def get_project_tasks(
self,
project_id: str,
page: Optional[int] = None,
page_size: int = 1000
) -> Optional[Dict[str, Any]]:
"""获取项目任务信息
Args:
project_id: 项目ID
page: 页码(从1开始)。如果为None,则获取所有任务
page_size: 每页大小
Returns:
如果指定了page参数,返回包含分页信息的字典:
{
"count": 总任务数,
"page": 当前页码,
"page_size": 每页大小,
"project_id": 项目ID,
"tasks": 当前页的任务列表
}
如果page为None,返回包含所有任务的字典:
"count": 总任务数,
"project_id": 项目ID,
"tasks": 所有任务列表
}
"""
try:
pid = int(project_id)
# 如果指定了page,直接获取单页任务
if page is not None:
logger.debug(
f"Fetching tasks for project {pid}, page {page} (page_size={page_size})")
response = await self.client.get(
f"/api/tasks",
params={
"project": pid,
"page": page,
"page_size": page_size
}
)
response.raise_for_status()
result = response.json()
# 返回单页结果,包含分页信息
return {
"count": result.get("total", len(result.get("tasks", []))),
"page": page,
"page_size": page_size,
"project_id": pid,
"tasks": result.get("tasks", [])
}
# 如果未指定page,获取所有任务
logger.debug(f"(page) not specified, fetching all tasks.")
all_tasks = []
response = await self.client.get(
f"/api/tasks",
params={
"project": pid
}
)
response.raise_for_status()
result = response.json()
tasks = result.get("tasks", [])
if not tasks:
logger.debug(f"No tasks found for this project.")
all_tasks.extend(tasks)
logger.debug(f"Fetched {len(tasks)} tasks.")
# 返回所有任务,不包含分页信息
return {
"count": len(all_tasks),
"project_id": pid,
"tasks": all_tasks
}
except httpx.HTTPStatusError as e:
logger.error(
f"获取项目任务失败 HTTP {e.response.status_code}: {e.response.text}")
return None
except Exception as e:
logger.error(f"获取项目任务时发生错误: {e}")
return None
async def delete_task(
self,
task_id: int
) -> bool:
"""删除单个任务"""
try:
logger.debug(f"Deleting task: {task_id}")
response = await self.client.delete(f"/api/tasks/{task_id}")
response.raise_for_status()
logger.debug(f"Task deleted: {task_id}")
return True
except httpx.HTTPStatusError as e:
logger.error(
f"Delete task {task_id} failed HTTP {e.response.status_code}: {e.response.text}")
return False
except Exception as e:
logger.error(f"Error while deleting task {task_id}: {e}")
return False
async def delete_tasks_batch(
self,
task_ids: List[int]
) -> Dict[str, int]:
"""批量删除任务"""
try:
logger.debug(f"Deleting {len(task_ids)} tasks in batch")
successful_deletions = 0
failed_deletions = 0
for task_id in task_ids:
if await self.delete_task(task_id):
successful_deletions += 1
else:
failed_deletions += 1
logger.debug(
f"Batch deletion finished: success {successful_deletions}, failed {failed_deletions}")
return {
"successful": successful_deletions,
"failed": failed_deletions,
"total": len(task_ids)
}
except Exception as e:
logger.error(f"Error while deleting tasks in batch: {e}")
return {
"successful": 0,
"failed": len(task_ids),
"total": len(task_ids)
}
async def get_project(self, project_id: int) -> Optional[Dict[str, Any]]:
"""获取项目信息"""
try:
logger.debug(f"Fetching project info: {project_id}")
response = await self.client.get(f"/api/projects/{project_id}")
response.raise_for_status()
return response.json()
except httpx.HTTPStatusError as e:
logger.error(
f"Get project info failed HTTP {e.response.status_code}: {e.response.text}")
return None
except Exception as e:
logger.error(f"Error while getting project info: {e}")
return None
async def delete_project(self, project_id: int) -> bool:
"""删除项目"""
try:
logger.debug(f"Deleting project: {project_id}")
response = await self.client.delete(f"/api/projects/{project_id}")
response.raise_for_status()
logger.debug(f"Project deleted: {project_id}")
return True
except httpx.HTTPStatusError as e:
logger.error(
f"Delete project {project_id} failed HTTP {e.response.status_code}: {e.response.text}")
return False
except Exception as e:
logger.error(f"Error while deleting project {project_id}: {e}")
return False
async def export_project(
self,
project_id: int,
export_type: str = "JSON",
) -> Optional[bytes]:
"""导出 Label Studio 项目数据。
对应 Label Studio 的项目导出接口,支持多种 exportType,
例如 JSON/JSON_MIN/CSV/TSV/COCO/YOLO/YOLOv8 等。
返回导出文件的原始二进制内容,调用方负责将其写入本地文件。
"""
try:
logger.debug(
"Exporting Label Studio project %s with exportType=%s",
project_id,
export_type,
)
response = await self.client.get(
f"/api/projects/{project_id}/export",
params={"exportType": export_type},
)
response.raise_for_status()
content = response.content or b""
logger.debug(
"Exported project %s with %d bytes",
project_id,
len(content),
)
return content
except httpx.HTTPStatusError as e:
logger.error(
"Export project %s failed HTTP %s: %s",
project_id,
e.response.status_code,
e.response.text,
)
return None
except Exception as e:
logger.error("Error while exporting project %s: %s", project_id, e)
return None
async def get_task_annotations(
self,
task_id: int
) -> Optional[List[Dict[str, Any]]]:
"""获取任务的标注结果
Args:
task_id: 任务ID
Returns:
标注结果列表,每个标注包含完整的annotation信息
"""
try:
logger.debug(f"Fetching annotations for task: {task_id}")
response = await self.client.get(f"/api/tasks/{task_id}/annotations")
response.raise_for_status()
annotations = response.json()
logger.debug(
f"Fetched {len(annotations)} annotations for task {task_id}")
return annotations
except httpx.HTTPStatusError as e:
logger.error(
f"Get task annotations failed HTTP {e.response.status_code}: {e.response.text}")
return None
except Exception as e:
logger.error(f"Error while getting task annotations: {e}")
return None
async def create_annotation(
self,
task_id: int,
result: List[Dict[str, Any]],
completed_by: Optional[int] = None
) -> Optional[Dict[str, Any]]:
"""为任务创建新的标注
Args:
task_id: 任务ID
result: 标注结果列表
completed_by: 完成标注的用户ID(可选)
Returns:
创建的标注信息,失败返回None
"""
try:
logger.debug(f"Creating annotation for task: {task_id}")
annotation_data = {
"result": result,
"task": task_id
}
if completed_by:
annotation_data["completed_by"] = completed_by
response = await self.client.post(
f"/api/tasks/{task_id}/annotations",
json=annotation_data
)
response.raise_for_status()
annotation = response.json()
logger.debug(
f"Created annotation {annotation.get('id')} for task {task_id}")
return annotation
except httpx.HTTPStatusError as e:
logger.error(
f"Create annotation failed HTTP {e.response.status_code}: {e.response.text}")
return None
except Exception as e:
logger.error(f"Error while creating annotation: {e}")
return None
async def update_annotation(
self,
annotation_id: int,
result: List[Dict[str, Any]]
) -> Optional[Dict[str, Any]]:
"""更新已存在的标注
Args:
annotation_id: 标注ID
result: 新的标注结果列表
Returns:
更新后的标注信息,失败返回None
"""
try:
logger.debug(f"Updating annotation: {annotation_id}")
annotation_data = {
"result": result
}
response = await self.client.patch(
f"/api/annotations/{annotation_id}",
json=annotation_data
)
response.raise_for_status()
annotation = response.json()
logger.debug(f"Updated annotation {annotation_id}")
return annotation
except httpx.HTTPStatusError as e:
logger.error(
f"Update annotation failed HTTP {e.response.status_code}: {e.response.text}")
return None
except Exception as e:
logger.error(f"Error while updating annotation: {e}")
return None
async def delete_annotation(
self,
annotation_id: int
) -> bool:
"""删除标注
Args:
annotation_id: 标注ID
Returns:
成功返回True,失败返回False
"""
try:
logger.debug(f"Deleting annotation: {annotation_id}")
response = await self.client.delete(f"/api/annotations/{annotation_id}")
response.raise_for_status()
logger.debug(f"Deleted annotation {annotation_id}")
return True
except httpx.HTTPStatusError as e:
logger.error(
f"Delete annotation failed HTTP {e.response.status_code}: {e.response.text}")
return False
except Exception as e:
logger.error(f"Error while deleting annotation: {e}")
return False
async def create_prediction(
self,
task_id: int,
result: List[Dict[str, Any]],
model_version: Optional[str] = None,
score: Optional[float] = None,
) -> Optional[Dict[str, Any]]:
"""为任务创建预测结果(predictions)。
该接口对应 Label Studio 的 `/api/predictions`,用于将模型自动标注结果
以 prediction 的形式写入,从而在前端界面里以“预测框”的方式展示,
不会覆盖人工标注(annotations)。
"""
try:
logger.debug(f"Creating prediction for task: {task_id}")
payload: Dict[str, Any] = {
"task": task_id,
"result": result,
}
if model_version is not None:
payload["model_version"] = model_version
if score is not None:
payload["score"] = score
response = await self.client.post("/api/predictions", json=payload)
response.raise_for_status()
prediction = response.json()
logger.debug(
"Created prediction %s for task %s", prediction.get(
"id"), task_id
)
return prediction
except httpx.HTTPStatusError as e:
logger.error(
"Create prediction failed HTTP %s: %s",
e.response.status_code,
e.response.text,
)
return None
except Exception as e:
logger.error("Error while creating prediction: %s", e)
return None
async def create_local_storage(
self,
project_id: int,
path: str,
title: str,
use_blob_urls: bool = True,
regex_filter: Optional[str] = None,
description: Optional[str] = None
) -> Optional[Dict[str, Any]]:
"""创建本地文件存储配置
Args:
project_id: Label Studio 项目 ID
path: 本地文件路径(在 Label Studio 容器中的路径)
title: 存储配置标题
use_blob_urls: 是否使用 blob URLs(建议 True)
regex_filter: 文件过滤正则表达式(可选)
description: 存储描述(可选)
Returns:
创建的存储配置信息,失败返回 None
"""
try:
logger.debug(
f"Creating local storage for project {project_id}: {path}")
storage_data = {
"project": project_id,
"path": path,
"title": title,
"use_blob_urls": use_blob_urls
}
if regex_filter:
storage_data["regex_filter"] = regex_filter
if description:
storage_data["description"] = description
response = await self.client.post(
"/api/storages/localfiles/",
json=storage_data
)
response.raise_for_status()
storage = response.json()
storage_id = storage.get("id")
logger.debug(
f"Local storage created successfully, ID: {storage_id}")
return storage
except httpx.HTTPStatusError as e:
logger.error(
f"Create local storage failed HTTP {e.response.status_code}: {e.response.text}")
return None
except Exception as e:
logger.error(f"Error while creating local storage: {e}")
return None
async def close(self):
"""关闭客户端连接"""
try:
await self.client.aclose()
logger.debug("Label Studio client closed")
except Exception as e:
logger.error(f"Error while closing Label Studio client: {e}")
@@ -0,0 +1,40 @@
from pydantic import Field
from typing import Dict, Any, Optional
from datetime import datetime
from app.module.shared.schema import BaseResponseModel
class LabelStudioProject(BaseResponseModel):
"""Label Studio项目模型"""
id: int = Field(..., description="项目ID")
title: str = Field(..., description="项目标题")
description: Optional[str] = Field(None, description="项目描述")
label_config: str = Field(..., description="标注配置")
created_at: Optional[datetime] = Field(None, description="创建时间")
updated_at: Optional[datetime] = Field(None, description="更新时间")
class LabelStudioTaskData(BaseResponseModel):
"""Label Studio任务数据模型"""
image: Optional[str] = Field(None, description="图像URL")
text: Optional[str] = Field(None, description="文本内容")
audio: Optional[str] = Field(None, description="音频URL")
video: Optional[str] = Field(None, description="视频URL")
filename: Optional[str] = Field(None, description="文件名")
class LabelStudioTask(BaseResponseModel):
"""Label Studio任务模型"""
data: LabelStudioTaskData = Field(..., description="任务数据")
project: Optional[int] = Field(None, description="项目ID")
meta: Optional[Dict[str, Any]] = Field(None, description="元数据")
class LabelStudioCreateProjectRequest(BaseResponseModel):
"""创建Label Studio项目请求模型"""
title: str = Field(..., description="项目标题")
description: str = Field("", description="项目描述")
label_config: str = Field(..., description="标注配置")
class LabelStudioCreateTaskRequest(BaseResponseModel):
"""创建Label Studio任务请求模型"""
data: Dict[str, Any] = Field(..., description="任务数据")
project: Optional[int] = Field(None, description="项目ID")
@@ -0,0 +1,4 @@
"""Tag configuration package"""
from .tag_config import LabelStudioTagConfig
__all__ = ['LabelStudioTagConfig']
@@ -0,0 +1,467 @@
# Label Studio Tag Configuration
# Defines supported tags, their properties, and child element requirements
# Object tags - represent data to be annotated
objects:
Audio:
description: "Display audio files"
required_attrs: [name, value]
optional_attrs: []
category: media
Bitmask:
description: "Display bitmask images for segmentation"
required_attrs: [name, value]
optional_attrs: []
category: image
PDF:
description: "Display PDF documents"
required_attrs: [name, value]
optional_attrs: []
category: document
Markdown:
description: "Display Markdown content"
required_attrs: [name, value]
optional_attrs: []
category: document
ParagraphLabels:
description: "Display paragraphs with label support"
required_attrs: [name, value]
optional_attrs: []
category: text
Timeseries:
description: "Display timeseries data"
required_attrs: [name, value]
optional_attrs: []
category: data
Vector:
description: "Display vector data for annotation"
required_attrs: [name, value]
optional_attrs: []
category: data
Chat:
description: "Display chat data for annotation"
required_attrs: [name, value]
optional_attrs: []
category: text
HyperText:
description: "Display HTML content"
required_attrs: [name, value]
optional_attrs: []
category: document
Image:
description: "Display images for annotation"
required_attrs: [name, value]
optional_attrs: []
category: image
Text:
description: "Display text for annotation"
required_attrs: [name, value]
optional_attrs: []
category: text
Video:
description: "Display video files"
required_attrs: [name, value]
optional_attrs: []
category: media
AudioPlus:
description: "Advanced audio player"
required_attrs: [name, value]
optional_attrs: []
category: media
Paragraphs:
description: "Display paragraphs of text"
required_attrs: [name, value]
optional_attrs: []
category: text
Table:
description: "Display tabular data"
required_attrs: [name, value]
optional_attrs: []
category: data
# Control tags - tools for annotation
# Categories:
# - labeling: Controls used for annotating/labeling objects (shown in template form)
# - layout: UI/layout elements not used for labeling (hidden from template form by default)
controls:
# Choice-based controls (use <Choice> children)
Choices:
description: "Multiple choice classification"
required_attrs: [name, toName]
optional_attrs:
required:
type: boolean
description: "Whether the choice is required"
choice:
type: string
values: [single, multiple]
default: single
description: "Selection mode: single or multiple"
showInline:
type: boolean
default: true
description: "Show choices inline or as dropdown"
requires_children: true
child_tag: Choice
child_required_attrs: [value]
category: labeling
Taxonomy:
description: "Hierarchical multi-label classification"
required_attrs: [name, toName]
optional_attrs:
required:
type: boolean
maxDepth:
type: number
default: 3
description: "Maximum depth of taxonomy tree"
requires_children: true
child_tag: Path
child_required_attrs: [value]
category: labeling
Ranker:
description: "Rank items in order"
required_attrs: [name, toName]
optional_attrs:
required:
type: boolean
maxChoices:
type: number
default: 5
description: "Maximum number of choices to rank"
requires_children: true
child_tag: Choice
child_required_attrs: [value]
category: layout
List:
description: "List selection control"
required_attrs: [name, toName]
optional_attrs:
required:
type: boolean
mode:
type: string
values: [single, multiple]
default: single
requires_children: true
child_tag: Item
child_required_attrs: [value]
category: layout
Filter:
description: "Filter control for annotation"
required_attrs: [name, toName]
optional_attrs:
required:
type: boolean
requires_children: false
category: layout
Collapse:
description: "Collapsible UI section"
required_attrs: [name]
optional_attrs:
collapsed:
type: boolean
default: false
requires_children: false
category: layout
Header:
description: "Section header for UI grouping"
required_attrs: [name]
optional_attrs:
level:
type: number
default: 1
description: "Header level (1-6)"
requires_children: false
category: layout
Shortcut:
description: "Keyboard shortcut definition"
required_attrs: [name, toName]
optional_attrs:
key:
type: string
description: "Shortcut key"
requires_children: false
category: layout
Style:
description: "Custom style for annotation UI"
required_attrs: [name]
optional_attrs:
value:
type: string
description: "CSS style value"
requires_children: false
category: layout
MagicWand:
description: "Magic wand segmentation tool"
required_attrs: [name, toName]
optional_attrs:
required:
type: boolean
requires_children: false
category: labeling
BitmaskLabels:
description: "Bitmask segmentation with labels"
required_attrs: [name, toName]
optional_attrs: [required]
requires_children: true
child_tag: Label
child_required_attrs: [value]
category: labeling
TimeseriesLabels:
description: "Labels for timeseries data"
required_attrs: [name, toName]
optional_attrs: [required]
requires_children: true
child_tag: Label
child_required_attrs: [value]
category: labeling
VectorLabels:
description: "Labels for vector data"
required_attrs: [name, toName]
optional_attrs: [required]
requires_children: true
child_tag: Label
child_required_attrs: [value]
category: labeling
ParagraphLabels:
description: "Labels for paragraphs"
required_attrs: [name, toName]
optional_attrs: [required]
requires_children: true
child_tag: Label
child_required_attrs: [value]
category: labeling
Relation:
description: "Draw relation between objects"
required_attrs: [name, toName]
optional_attrs:
required:
type: boolean
requires_children: false
category: layout
Relations:
description: "Draw multiple relations between objects"
required_attrs: [name, toName]
optional_attrs:
required:
type: boolean
requires_children: false
category: layout
Pairwise:
description: "Pairwise comparison control"
required_attrs: [name, toName]
optional_attrs:
required:
type: boolean
requires_children: false
category: layout
DateTime:
description: "Date and time input"
required_attrs: [name, toName]
optional_attrs:
required:
type: boolean
format:
type: string
default: "YYYY-MM-DD HH:mm:ss"
requires_children: false
category: labeling
Number:
description: "Numeric input field"
required_attrs: [name, toName]
optional_attrs:
required:
type: boolean
min:
type: number
max:
type: number
step:
type: number
default: 1
requires_children: false
category: labeling
# Label-based controls (use <Label> children)
RectangleLabels:
description: "Rectangle bounding boxes with labels"
required_attrs: [name, toName]
optional_attrs:
required:
type: boolean
description: "Whether annotation is required"
strokeWidth:
type: number
default: 3
description: "Width of the bounding box border"
canRotate:
type: boolean
default: true
description: "Allow rotation of rectangles"
requires_children: true
child_tag: Label
child_required_attrs: [value]
category: labeling
PolygonLabels:
description: "Polygon annotations with labels"
required_attrs: [name, toName]
optional_attrs:
required:
type: boolean
strokeWidth:
type: number
default: 3
pointSize:
type: string
values: [small, medium, large]
default: medium
requires_children: true
child_tag: Label
child_required_attrs: [value]
category: labeling
Labels:
description: "Generic labels for classification"
required_attrs: [name, toName]
optional_attrs: [required]
requires_children: true
child_tag: Label
child_required_attrs: [value]
category: labeling
KeyPointLabels:
description: "Keypoint annotations with labels"
required_attrs: [name, toName]
optional_attrs: [required]
requires_children: true
child_tag: Label
child_required_attrs: [value]
category: labeling
BrushLabels:
description: "Brush/semantic segmentation with labels"
required_attrs: [name, toName]
optional_attrs: [required]
requires_children: true
child_tag: Label
child_required_attrs: [value]
category: labeling
EllipseLabels:
description: "Ellipse annotations with labels"
required_attrs: [name, toName]
optional_attrs: [required]
requires_children: true
child_tag: Label
child_required_attrs: [value]
category: labeling
# Simple controls (no children required)
Rectangle:
description: "Rectangle bounding box without labels"
required_attrs: [name, toName]
optional_attrs: [required]
requires_children: false
category: labeling
Polygon:
description: "Polygon annotation without labels"
required_attrs: [name, toName]
optional_attrs: [required]
requires_children: false
category: labeling
Ellipse:
description: "Ellipse annotation without labels"
required_attrs: [name, toName]
optional_attrs: [required]
requires_children: false
category: labeling
KeyPoint:
description: "Keypoint annotation without labels"
required_attrs: [name, toName]
optional_attrs: [required]
requires_children: false
category: labeling
Brush:
description: "Brush annotation without labels"
required_attrs: [name, toName]
optional_attrs: [required]
requires_children: false
category: labeling
TextArea:
description: "Text input field"
required_attrs: [name, toName]
optional_attrs:
required:
type: boolean
placeholder:
type: string
description: "Placeholder text"
maxSubmissions:
type: number
description: "Maximum number of submissions"
rows:
type: number
default: 3
description: "Number of rows in textarea"
editable:
type: boolean
default: true
requires_children: false
category: labeling
Rating:
description: "Star rating or numeric rating"
required_attrs: [name, toName]
optional_attrs:
required:
type: boolean
maxRating:
type: number
default: 5
description: "Maximum rating value"
defaultValue:
type: number
description: "Default rating value"
size:
type: string
values: [small, medium, large]
default: medium
icon:
type: string
values: [star, heart, fire, thumbs]
default: star
requires_children: false
category: labeling
VideoRectangle:
description: "Rectangle annotations for video"
required_attrs: [name, toName]
optional_attrs: [required]
requires_children: false
category: labeling
@@ -0,0 +1,150 @@
"""
Label Studio Tag Configuration Loader
"""
import yaml
from typing import Dict, Any, Optional, Set, Tuple
from pathlib import Path
class LabelStudioTagConfig:
"""Label Studio标签配置管理器"""
_instance: Optional['LabelStudioTagConfig'] = None
_config: Dict[str, Any] = {}
def __new__(cls):
if cls._instance is None:
cls._instance = super().__new__(cls)
return cls._instance
def __init__(self):
"""初始化时加载配置"""
if not self._config:
self._load_config()
@classmethod
def _load_config(cls):
"""加载YAML配置文件"""
config_path = Path(__file__).parent / "label_studio_tags.yaml"
with open(config_path, 'r', encoding='utf-8') as f:
cls._config = yaml.safe_load(f) or {}
@classmethod
def get_object_types(cls) -> Set[str]:
"""获取所有支持的对象类型"""
return set(cls._config.get('objects', {}).keys())
@classmethod
def get_control_types(cls) -> Set[str]:
"""获取所有支持的控件类型"""
return set(cls._config.get('controls', {}).keys())
@classmethod
def get_control_config(cls, control_type: str) -> Optional[Dict[str, Any]]:
"""获取控件的配置信息"""
return cls._config.get('controls', {}).get(control_type)
@classmethod
def get_object_config(cls, object_type: str) -> Optional[Dict[str, Any]]:
"""获取对象的配置信息"""
return cls._config.get('objects', {}).get(object_type)
@classmethod
def requires_children(cls, control_type: str) -> bool:
"""检查控件是否需要子元素"""
config = cls.get_control_config(control_type)
return config.get('requires_children', False) if config else False
@classmethod
def get_child_tag(cls, control_type: str) -> Optional[str]:
"""获取控件的子元素标签名"""
config = cls.get_control_config(control_type)
return config.get('child_tag') if config else None
@classmethod
def get_controls_with_child_tag(cls, child_tag: str) -> Set[str]:
"""获取使用指定子元素标签的所有控件类型"""
controls = set()
for control_type, config in cls._config.get('controls', {}).items():
if config.get('child_tag') == child_tag:
controls.add(control_type)
return controls
@classmethod
def get_optional_attrs(cls, tag_type: str, is_control: bool = True) -> Dict[str, Any]:
"""
获取标签的可选属性配置
Args:
tag_type: 标签类型
is_control: 是否为控件类型(否则为对象类型)
Returns:
可选属性配置字典
"""
config = cls.get_control_config(tag_type) if is_control else cls.get_object_config(tag_type)
if not config:
return {}
optional_attrs = config.get('optional_attrs', {})
# 如果是简单列表格式(旧格式),转换为字典
if isinstance(optional_attrs, list):
return {attr: {} for attr in optional_attrs}
# 确保返回的是字典
return optional_attrs if isinstance(optional_attrs, dict) else {}
@classmethod
def validate_attr_value(cls, tag_type: str, attr_name: str, attr_value: Any, is_control: bool = True) -> Tuple[bool, Optional[str]]:
"""
验证属性值是否符合配置要求
Args:
tag_type: 标签类型
attr_name: 属性名
attr_value: 属性值
is_control: 是否为控件类型
Returns:
(是否有效, 错误信息)
"""
optional_attrs = cls.get_optional_attrs(tag_type, is_control)
if attr_name not in optional_attrs:
return True, None # 不在配置中的属性,不验证
attr_config = optional_attrs.get(attr_name, {})
# 如果配置不是字典,跳过验证
if not isinstance(attr_config, dict):
return True, None
# 检查类型
expected_type = attr_config.get('type')
if expected_type == 'boolean':
if not isinstance(attr_value, (bool, str)) or (isinstance(attr_value, str) and attr_value.lower() not in ['true', 'false']):
return False, f"Attribute '{attr_name}' must be boolean"
elif expected_type == 'number':
try:
float(attr_value)
except (ValueError, TypeError):
return False, f"Attribute '{attr_name}' must be a number"
# 检查枚举值
allowed_values = attr_config.get('values')
if allowed_values and attr_value not in allowed_values:
return False, f"Attribute '{attr_name}' must be one of {allowed_values}, got '{attr_value}'"
return True, None
@classmethod
def get_attr_default(cls, tag_type: str, attr_name: str, is_control: bool = True) -> Optional[Any]:
"""获取属性的默认值"""
optional_attrs = cls.get_optional_attrs(tag_type, is_control)
attr_config = optional_attrs.get(attr_name, {})
# 确保attr_config是字典后再访问
if isinstance(attr_config, dict):
return attr_config.get('default')
return None
@@ -0,0 +1,18 @@
from fastapi import APIRouter
from .config import router as about_router
from .project import router as project_router
from .task import router as task_router
from .template import router as template_router
from .auto import router as auto_router
router = APIRouter(
prefix="/annotation",
tags = ["annotation"]
)
router.include_router(about_router)
router.include_router(project_router)
router.include_router(task_router)
router.include_router(template_router)
router.include_router(auto_router)
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,47 @@
from fastapi import APIRouter
from app.module.shared.schema import StandardResponse
from app.core.logging import get_logger
from app.core.config import settings
from ..schema import (
ConfigResponse,
TagConfigResponse
)
from ..config.tag_config import LabelStudioTagConfig
router = APIRouter(
prefix="/tags",
tags=["annotation/config"]
)
logger = get_logger(__name__)
@router.get("", response_model=StandardResponse[ConfigResponse])
async def get_config():
"""获取配置信息(已废弃,请使用 /api/annotation/about)"""
return StandardResponse(
code="0",
message="success",
data=ConfigResponse(
label_studio_url=settings.label_studio_base_url,
)
)
@router.get("/config", response_model=StandardResponse[TagConfigResponse], summary="获取标签配置")
async def get_tag_config():
"""
获取所有Label Studio标签类型的配置(对象+控件),用于前端动态渲染。
"""
# Ensure config is loaded by instantiating the class
tag_config = LabelStudioTagConfig()
config = LabelStudioTagConfig._config
if not config:
logger.error("Failed to load tag configuration")
return StandardResponse(
code="common.500",
message="Failed to load tag configuration",
data={"objects": {}, "controls": {}}
)
return StandardResponse(code="0", message="success", data=config)
@@ -0,0 +1,993 @@
from typing import Optional
import math
import uuid
from fastapi import APIRouter, Depends, HTTPException, Query, Path, Response, Request
from fastapi.responses import JSONResponse
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from app.db.session import get_db
from app.db.models import LabelingProject
from app.module.shared.schema import StandardResponse, PaginatedData
from app.module.dataset import DatasetManagementService
from app.core.logging import get_logger
from app.core.config import settings
from ..client import LabelStudioClient
from ..service.mapping import DatasetMappingService
from ..service.sync import SyncService
from ..service.ls_annotation_sync import LSAnnotationSyncService
from ..service.template import AnnotationTemplateService
from ..schema import (
DatasetMappingCreateRequest,
DatasetMappingCreateResponse,
DeleteDatasetResponse,
DatasetMappingResponse,
)
from ..schema.auto import ImportFromLabelStudioRequest, UpdateAutoAnnotationTaskFilesRequest
router = APIRouter(
prefix="/project",
tags=["annotation/project"]
)
logger = get_logger(__name__)
# xuhuanqing
@router.get("/{mapping_id}/login")
async def login_label_studio_project(
mapping_id: str = Path(..., description="项目映射ID"),
request: Request = None,
db: AsyncSession = Depends(get_db)
):
"""
登录 Label Studio 项目(Session Cookie 模式)
- 仅负责给浏览器种 Label Studio session
- 不负责跳转
"""
user = request.headers.get("User") if request else None
if settings.deepdata_jwt_enable and (not user or not user.strip()):
logger.warning(
f"Unauthorized access attempt to annotation login: mapping_id={mapping_id}"
)
raise HTTPException(status_code=401, detail="Authentication required")
stmt = select(LabelingProject).where(
LabelingProject.id == mapping_id,
LabelingProject.deleted_at.is_(None)
)
result = await db.execute(stmt)
project = result.scalar_one_or_none()
if not project:
logger.warning(
f"Project not found for login: mapping_id={mapping_id}, user={user or 'anonymous'}"
)
raise HTTPException(status_code=404, detail="Project not found")
logger.info(
f"User '{user or 'anonymous'}' logging into Label Studio project: "
f"mapping_id={mapping_id}, project_name={project.name}"
)
try:
ls_client = LabelStudioClient(
base_url=settings.label_studio_base_url,
token=settings.label_studio_user_token
)
target_response = await ls_client.login_label_studio()
if target_response is None:
raise HTTPException(
status_code=502,
detail="Label Studio login failed"
)
response = Response(status_code=204)
# ✅ Starlette 正确写法:append,不是 add
for cookie in target_response.headers.get_list("set-cookie"):
response.headers.append("set-cookie", cookie)
return response
except HTTPException:
raise
except Exception as e:
logger.error(
f"Error while logging in LabelStudio: {e}",
exc_info=True
)
raise HTTPException(
status_code=500,
detail="Internal server error"
)
'''
@router.get("/{mapping_id}/login")
async def login_label_studio_project(
mapping_id: str = Path(..., description="项目映射ID"),
request: Request = None,
db: AsyncSession = Depends(get_db)
):
"""登录 Label Studio 项目 — 需要认证且校验项目归属"""
# 检查用户认证(仅在 JWT 启用时强制要求)
user = request.headers.get("User") if request else None
if settings.deepdata_jwt_enable and (not user or not user.strip()):
logger.warning(f"Unauthorized access attempt to annotation login: mapping_id={mapping_id}")
raise HTTPException(status_code=401, detail="Authentication required")
# 校验 mapping_id 对应的项目存在
stmt = select(LabelingProject).where(
LabelingProject.id == mapping_id,
LabelingProject.deleted_at.is_(None)
)
result = await db.execute(stmt)
project = result.scalar_one_or_none()
if not project:
logger.warning(f"Project not found for login: mapping_id={mapping_id}, user={user or 'anonymous'}")
raise HTTPException(status_code=404, detail="Project not found")
logger.info(f"User '{user or 'anonymous'}' logging into Label Studio project: mapping_id={mapping_id}, "
f"project_name={project.name}")
try:
ls_client = LabelStudioClient(base_url=settings.label_studio_base_url,
token=settings.label_studio_user_token)
target_response = await ls_client.login_label_studio()
if target_response is None:
raise HTTPException(status_code=502, detail="Label Studio login failed")
headers = dict(target_response.headers)
set_cookies = target_response.headers.get_list("set-cookie")
# 删除合并的 Set-Cookie
if "set-cookie" in headers:
del headers["set-cookie"]
# 创建新响应,添加多个 Set-Cookie
response = Response(
content=target_response.content,
status_code=target_response.status_code,
headers=headers
)
# 分别添加每个 Set-Cookie
for cookie in set_cookies:
response.headers.append("set-cookie", cookie)
return response
except HTTPException:
raise
except Exception as e:
logger.error(f"Error while logining in LabelStudio: {e}", e)
raise HTTPException(status_code=500, detail="Internal server error")
'''
@router.post("", response_model=StandardResponse[DatasetMappingCreateResponse])
async def create_mapping(
request: DatasetMappingCreateRequest,
db: AsyncSession = Depends(get_db)
):
"""
创建数据集映射
根据指定的DM程序中的数据集,创建Label Studio中的数据集,
在数据库中记录这一关联关系,返回Label Studio数据集的ID
注意:一个数据集可以创建多个标注项目
支持通过 template_id 指定标注模板,如果提供了模板ID,则使用模板的配置
"""
try:
dm_client = DatasetManagementService(db)
ls_client = LabelStudioClient(
base_url=settings.label_studio_base_url,
token=settings.label_studio_user_token,
)
mapping_service = DatasetMappingService(db)
sync_service = SyncService(dm_client, ls_client, mapping_service)
template_service = AnnotationTemplateService()
logger.info(
"Create dataset mapping request: dataset_id=%s, file_ids=%s",
request.dataset_id,
request.file_ids,
)
# 从DM服务获取数据集信息
dataset_info = await dm_client.get_dataset(request.dataset_id)
if not dataset_info:
raise HTTPException(
status_code=404,
detail=f"Dataset not found in DM service: {request.dataset_id}",
)
project_name = (
request.name
or dataset_info.name
or "A new project from DeepData"
)
project_description = (
request.description
or dataset_info.description
or f"Imported from DM dataset {dataset_info.name} ({dataset_info.id})"
)
# 如果提供了模板ID,获取模板配置
label_config = None
if request.template_id:
logger.info("Using template: %s", request.template_id)
template = await template_service.get_template(db, request.template_id)
if not template:
raise HTTPException(
status_code=404,
detail=f"Template not found: {request.template_id}",
)
label_config = template.label_config
logger.debug(
"Template label config loaded for template: %s",
template.name,
)
# 在Label Studio中创建项目
project_data = await ls_client.create_project(
title=project_name,
description=project_description,
label_config=label_config,
)
if not project_data:
raise HTTPException(
status_code=500,
detail="Fail to create Label Studio project.",
)
project_id = project_data["id"]
# 配置主数据集的本地存储:dataset/<id>
local_storage_path = (
f"{settings.label_studio_local_document_root}/{request.dataset_id}"
)
storage_result = await ls_client.create_local_storage(
project_id=project_id,
path=local_storage_path,
title="Dataset_BLOB",
use_blob_urls=True,
description=f"Local storage for dataset {dataset_info.name}",
)
if not storage_result:
# 本地存储配置失败,记录警告但不中断流程
logger.warning(
"Failed to configure local storage for project %s",
project_id,
)
else:
logger.info(
"Local storage configured for project %s: %s",
project_id,
local_storage_path,
)
labeling_project = LabelingProject(
id=str(uuid.uuid4()),
dataset_id=request.dataset_id,
labeling_project_id=str(project_id),
name=project_name,
template_id=request.template_id,
)
# 创建映射关系,包含项目名称(先持久化映射以获得 mapping.id)
mapping = await mapping_service.create_mapping(labeling_project)
# 如果未指定 file_ids,保持原有行为:按映射数据集完整同步
if not request.file_ids:
await sync_service.sync_dataset_files(mapping.id, 100)
else:
# 仿照自动标注逻辑:根据 file_ids 反查所属数据集,可跨多个数据集,
# 最终把这些文件同步到同一个 Label Studio 项目中。
try:
from typing import Set as _Set, Dict as _Dict
from app.db.models.dataset_management import DatasetFiles
file_ids = request.file_ids
stmt = (
select(DatasetFiles.dataset_id, DatasetFiles.id)
.where(DatasetFiles.id.in_(file_ids))
)
result = await db.execute(stmt)
rows = result.fetchall()
grouped: _Dict[str, _Set[str]] = {}
resolved_ids: _Set[str] = set()
for ds_id, fid in rows:
if not ds_id or not fid:
continue
fid_str = str(fid)
grouped.setdefault(str(ds_id), set()).add(fid_str)
resolved_ids.add(fid_str)
# 未能解析到数据集的文件,全部归入主数据集,避免丢失
unresolved_ids = {str(fid) for fid in file_ids} - resolved_ids
if unresolved_ids:
logger.warning(
"Some file_ids could not be resolved to dataset_id when syncing manual project files: %s",
",".join(sorted(unresolved_ids)),
)
grouped.setdefault(str(request.dataset_id), set()).update(
unresolved_ids
)
# 为所有涉及到的额外数据集提前配置本地存储
try:
for extra_ds_id in grouped.keys():
# 主数据集已在上方配置过
if str(extra_ds_id) == str(request.dataset_id):
continue
extra_local_storage_path = (
f"{settings.label_studio_local_document_root}/{extra_ds_id}"
)
extra_storage_result = await ls_client.create_local_storage(
project_id=project_id,
path=extra_local_storage_path,
title=f"Dataset_BLOB_{extra_ds_id}",
use_blob_urls=True,
description=(
f"Local storage for dataset {extra_ds_id} "
"(multi-dataset manual project)"
),
)
if not extra_storage_result:
logger.warning(
"Failed to configure extra local storage for project %s (dataset %s)",
project_id,
extra_ds_id,
)
else:
logger.info(
"Extra local storage configured for project %s: %s",
project_id,
extra_local_storage_path,
)
except Exception as e: # pragma: no cover - 容错
logger.warning(
"Error while configuring extra local storage for project %s: %s",
project_id,
e,
)
if not grouped:
# 极端情况:完全无法解析,退回到仅按主数据集+给定 file_ids 同步
await sync_service.sync_files(
mapping,
100,
allowed_file_ids={str(fid) for fid in file_ids},
delete_orphans=False,
)
else:
# 对每个涉及到的数据集,使用 override_dataset_id 将其文件同步到同一个项目
for ds_id, ds_file_ids in grouped.items():
await sync_service.sync_files(
mapping,
100,
allowed_file_ids=ds_file_ids,
override_dataset_id=ds_id,
delete_orphans=False,
)
except Exception as e: # pragma: no cover - 映射创建成功但首次文件同步失败
# 同步失败不影响项目和映射本身的创建,前端可通过“同步”按钮重试
logger.warning(
"Failed to sync dataset files for manual LS project %s with file_ids filter: %s",
project_id,
e,
)
response_data = DatasetMappingCreateResponse(
id=mapping.id,
labeling_project_id=str(mapping.labeling_project_id),
labeling_project_name=mapping.name or project_name,
)
return StandardResponse(
code="0",
message="success",
data=response_data,
)
except HTTPException:
raise
except Exception as e: # pragma: no cover - 兜底错误
logger.error(f"Error while creating dataset mapping: {e}")
raise HTTPException(status_code=500, detail="Internal server error")
@router.get("/{mapping_id}/files", response_model=StandardResponse[list[dict]])
async def get_manual_mapping_files(
mapping_id: str = Path(..., description="映射ID (mapping UUID)"),
db: AsyncSession = Depends(get_db),
):
"""查询手动标注映射当前在 Label Studio 中已存在的 DM 文件列表。
该接口主要用于前端“编辑任务数据集”弹窗的初始选中状态:
- 通过 Label Studio 任务反查当前已有关联的 DM 文件ID;
- 再到 DM 数据库中查询这些文件的基础信息与所属数据集。
"""
from sqlalchemy import select as _select # 本地导入避免循环依赖
from app.db.models.dataset_management import DatasetFiles, Dataset
mapping_service = DatasetMappingService(db)
mapping = await mapping_service.get_mapping_by_uuid(mapping_id)
if not mapping:
raise HTTPException(status_code=404, detail="Mapping not found")
dm_client = DatasetManagementService(db)
ls_client = LabelStudioClient(
base_url=settings.label_studio_base_url,
token=settings.label_studio_user_token,
)
sync_service = SyncService(dm_client, ls_client, mapping_service)
existing_mapping = await sync_service.get_existing_dm_file_mapping(
mapping.labeling_project_id
)
file_ids = list(existing_mapping.keys())
if not file_ids:
return StandardResponse(code="0", message="success", data=[])
files_result = await db.execute(
_select(DatasetFiles).where(DatasetFiles.id.in_(file_ids))
)
files = list(files_result.scalars().all())
dataset_ids = {
str(getattr(f, "dataset_id"))
for f in files
if getattr(f, "dataset_id", None)
}
dataset_name_map: dict[str, str] = {}
if dataset_ids:
ds_result = await db.execute(
_select(Dataset.id, Dataset.name).where(Dataset.id.in_(dataset_ids))
)
for ds_id, ds_name in ds_result.fetchall():
dataset_name_map[str(ds_id)] = ds_name or ""
data: list[dict] = []
for f in files:
fid = str(getattr(f, "id"))
ds_id = (
str(getattr(f, "dataset_id"))
if getattr(f, "dataset_id", None)
else None
)
item: dict = {
"id": fid,
"datasetId": ds_id,
"datasetName": dataset_name_map.get(ds_id or "", ""),
"fileName": getattr(f, "file_name", ""),
"fileSize": int(getattr(f, "file_size", 0) or 0),
"filePath": getattr(f, "file_path", ""),
"fileType": getattr(f, "file_type", None),
}
data.append(item)
return StandardResponse(code="0", message="success", data=data)
@router.put("/{mapping_id}/files", response_model=StandardResponse[bool])
async def update_manual_mapping_files(
mapping_id: str = Path(..., description="映射ID (mapping UUID)"),
body: UpdateAutoAnnotationTaskFilesRequest = ..., # 复用通用结构:datasetId + fileIds
db: AsyncSession = Depends(get_db),
):
"""更新手动标注映射所关联的 DM 文件集合(仅追加,不删除已有任务)。
语义约定:
- 映射创建时所同步的文件集合视为“基础集合”,后续编辑不支持移除;
- 本接口只会为新增的 fileIds 在 Label Studio 中创建任务;
- 不会删除任何现有任务(delete_orphans=False)。
"""
from sqlalchemy import select as _select # 本地导入
from typing import Set as _Set, Dict as _Dict
from app.db.models.dataset_management import DatasetFiles
mapping_service = DatasetMappingService(db)
mapping = await mapping_service.get_mapping_by_uuid(mapping_id)
if not mapping:
raise HTTPException(status_code=404, detail="Mapping not found")
dm_client = DatasetManagementService(db)
ls_client = LabelStudioClient(
base_url=settings.label_studio_base_url,
token=settings.label_studio_user_token,
)
sync_service = SyncService(dm_client, ls_client, mapping_service)
requested_ids = {str(fid) for fid in (body.file_ids or [])}
if not requested_ids:
# 不做任何变更,但认为成功
return StandardResponse(code="0", message="success", data=True)
existing_mapping = await sync_service.get_existing_dm_file_mapping(
mapping.labeling_project_id
)
existing_ids = set(existing_mapping.keys())
# 仅对新增文件创建任务
new_ids = sorted(requested_ids - existing_ids)
if not new_ids:
return StandardResponse(code="0", message="success", data=True)
stmt = (
_select(DatasetFiles.dataset_id, DatasetFiles.id)
.where(DatasetFiles.id.in_(new_ids))
)
result = await db.execute(stmt)
rows = result.fetchall()
grouped: _Dict[str, _Set[str]] = {}
resolved_ids: _Set[str] = set()
for ds_id, fid in rows:
if not ds_id or not fid:
continue
fid_str = str(fid)
grouped.setdefault(str(ds_id), set()).add(fid_str)
resolved_ids.add(fid_str)
unresolved_ids = {str(fid) for fid in new_ids} - resolved_ids
if unresolved_ids:
logger.warning(
"Some file_ids could not be resolved to dataset_id when updating manual mapping files: %s",
",".join(sorted(unresolved_ids)),
)
grouped.setdefault(str(mapping.dataset_id), set()).update(unresolved_ids)
# 为所有涉及到的额外数据集提前配置本地存储(与创建逻辑保持一致)
try:
for extra_ds_id in grouped.keys():
if str(extra_ds_id) == str(mapping.dataset_id):
continue
extra_local_storage_path = (
f"{settings.label_studio_local_document_root}/{extra_ds_id}"
)
extra_storage_result = await ls_client.create_local_storage(
project_id=int(mapping.labeling_project_id),
path=extra_local_storage_path,
title=f"Dataset_BLOB_{extra_ds_id}",
use_blob_urls=True,
description=(
f"Local storage for dataset {extra_ds_id} "
"(multi-dataset manual project, edit)"
),
)
if not extra_storage_result:
logger.warning(
"Failed to configure extra local storage for project %s (dataset %s) when updating manual mapping files",
mapping.labeling_project_id,
extra_ds_id,
)
else:
logger.info(
"Extra local storage configured for project %s: %s (edit)",
mapping.labeling_project_id,
extra_local_storage_path,
)
except Exception as e: # pragma: no cover
logger.warning(
"Error while configuring extra local storage for project %s during update: %s",
mapping.labeling_project_id,
e,
)
# 将新增文件按数据集分组,同步到 Label Studio;不删除已有任务
for ds_id, ds_file_ids in grouped.items():
await sync_service.sync_files(
mapping,
batch_size=100,
allowed_file_ids=ds_file_ids,
override_dataset_id=ds_id,
delete_orphans=False,
)
return StandardResponse(code="0", message="success", data=True)
@router.post("/{mapping_id}/sync-label-studio-back", response_model=StandardResponse[bool])
async def import_manual_from_label_studio_to_dataset(
mapping_id: str = Path(..., description="映射ID (mapping UUID)"),
body: ImportFromLabelStudioRequest | None = None,
db: AsyncSession = Depends(get_db),
):
"""将手动标注工程在 Label Studio 中的标注结果导回到某个数据集。
行为与自动标注的后向同步保持一致:
- 按 mapping_id 定位 Label Studio 项目;
- 通过 exportType 导出项目完整结果(JSON/COCO/YOLO 等);
- 将导出文件作为一个普通文件保存到指定数据集目录,并注册到 t_dm_dataset_files;
- 不解析每条标注、不修改 tags,仅追加一个“导出工件”文件。
"""
import os
import tempfile
from datetime import datetime
if body is None:
raise HTTPException(status_code=400, detail="Request body is required")
# 1. 获取并校验映射(手动标注任务)
mapping_service = DatasetMappingService(db)
mapping = await mapping_service.get_mapping_by_uuid(mapping_id)
if not mapping:
raise HTTPException(status_code=404, detail="Mapping not found")
# 2. 确定目标数据集:若未显式指定,则使用映射绑定的数据集
dm_service = DatasetManagementService(db)
target_dataset_id = body.target_dataset_id or mapping.dataset_id
dataset = await dm_service.get_dataset(target_dataset_id)
if not dataset:
raise HTTPException(status_code=404, detail="Target dataset not found")
# 3. 调用 Label Studio 导出接口
ls_client = LabelStudioClient(
base_url=settings.label_studio_base_url,
token=settings.label_studio_user_token,
)
export_format = (body.export_format or "JSON").upper()
ext_map = {
"JSON": ".json",
"JSON_MIN": ".json",
"CSV": ".csv",
"TSV": ".tsv",
"COCO": ".json",
"YOLO": ".json",
"YOLOV8": ".json",
}
file_ext = ext_map.get(export_format, ".json")
project_id = mapping.labeling_project_id
content = await ls_client.export_project(int(project_id), export_type=export_format)
if content is None or len(content) == 0:
raise HTTPException(status_code=500, detail="Failed to export project from Label Studio")
# 4. 将导出结果写入临时文件,再通过 DatasetManagementService 导入到数据集
tmp_fd, tmp_path = tempfile.mkstemp(suffix=file_ext)
os.close(tmp_fd)
try:
with open(tmp_path, "wb") as f:
f.write(content)
timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
def _sanitize_base_name(raw: str) -> str:
# 去掉路径分隔符,仅保留最后一段
name = (raw or "").strip().replace("\\", "/").split("/")[-1]
# 去掉用户自带的扩展名,避免与服务器推断的后缀冲突
if "." in name:
name = name.rsplit(".", 1)[0]
return name or f"ls_export_{project_id}_{timestamp}"
if getattr(body, "file_name", None):
base_stem = _sanitize_base_name(body.file_name) # type: ignore[arg-type]
base_name = f"{base_stem}{file_ext}"
else:
base_name = f"ls_export_{project_id}_{timestamp}{file_ext}"
tmp_dir = os.path.dirname(tmp_path)
target_tmp_path = os.path.join(tmp_dir, base_name)
os.replace(tmp_path, target_tmp_path)
# 写入目标(通常为源)数据集下的 "标注数据" 目录
await dm_service.add_files_to_dataset_subdir(target_dataset_id, [target_tmp_path], "标注数据")
finally:
if os.path.exists(tmp_path):
try:
os.remove(tmp_path)
except Exception:
pass
if "target_tmp_path" in locals() and os.path.exists(target_tmp_path):
try:
os.remove(target_tmp_path)
except Exception:
pass
return StandardResponse(code="0", message="success", data=True)
@router.post("/{mapping_id}/sync-db", response_model=StandardResponse[int])
async def sync_manual_annotations_to_database(
mapping_id: str = Path(..., description="映射ID (mapping UUID)"),
db: AsyncSession = Depends(get_db),
):
"""从 Label Studio 项目同步当前手动标注结果到 DM 数据库。
行为:
- 基于 mapping_id 定位 Label Studio 项目;
- 先执行 DM -> Label Studio 的文件差异同步(增量创建新文件任务、删除孤儿任务);
- 再执行双向标签同步(LS <-> DM),按时间戳保留较新标注:
* LS 更新时间更新 -> 覆盖 DM;
* DM tags_updated_at 更新 -> 覆盖 LS。
返回值为同步变更数量(LS->DM + DM->LS)。
"""
mapping_service = DatasetMappingService(db)
mapping = await mapping_service.get_mapping_by_uuid(mapping_id)
if not mapping:
raise HTTPException(status_code=404, detail="Mapping not found")
ls_client = LabelStudioClient(
base_url=settings.label_studio_base_url,
token=settings.label_studio_user_token,
)
dm_client = DatasetManagementService(db)
sync_orchestrator = SyncService(dm_client, ls_client, mapping_service)
file_sync_result = await sync_orchestrator.sync_files(mapping, batch_size=50)
logger.info(
"Manual sync-db pre file-sync done: mapping=%s, created=%s, deleted=%s, total=%s",
mapping_id,
file_sync_result.get("created", 0),
file_sync_result.get("deleted", 0),
file_sync_result.get("total", 0),
)
annotation_sync_result = await sync_orchestrator.sync_annotations_bidirectional(
mapping,
batch_size=50,
overwrite=True,
overwrite_ls=True,
sync_files_first=False,
)
total_synced = annotation_sync_result.synced_to_dm + annotation_sync_result.synced_to_ls
return StandardResponse(code="0", message="success", data=total_synced)
@router.get("", response_model=StandardResponse[PaginatedData[DatasetMappingResponse]])
async def list_mappings(
page: int = Query(1, ge=1, description="页码(从1开始)"),
size: int = Query(20, ge=1, le=100, description="每页记录数"),
include_template: bool = Query(False, description="是否包含模板详情", alias="includeTemplate"),
db: AsyncSession = Depends(get_db)
):
"""
查询所有映射关系(分页)
返回所有有效的数据集映射关系(未被软删除的),支持分页查询。
可选择是否包含完整的标注模板信息(默认不包含,以提高列表查询性能)。
参数:
- page: 页码(从1开始)
- pageSize: 每页记录数
- includeTemplate: 是否包含模板详情(默认false)
"""
try:
service = DatasetMappingService(db)
# 计算 skip
skip = (page - 1) * size
logger.info(f"List mappings: page={page}, size={size}, include_template={include_template}")
# 获取数据和总数
mappings, total = await service.get_all_mappings_with_count(
skip=skip,
limit=size,
include_deleted=False,
include_template=include_template
)
# 计算总页数
total_pages = math.ceil(total / size) if total > 0 else 0
# 构造分页响应
paginated_data = PaginatedData(
page=page,
size=size,
total_elements=total,
total_pages=total_pages,
content=mappings
)
logger.info(f"List mappings: page={page}, returned {len(mappings)}/{total}, templates_included: {include_template}")
return StandardResponse(
code="0",
message="success",
data=paginated_data
)
except Exception as e:
logger.error(f"Error listing mappings: {e}")
raise HTTPException(status_code=500, detail="Internal server error")
@router.get("/{mapping_id}", response_model=StandardResponse[DatasetMappingResponse])
async def get_mapping(
mapping_id: str,
db: AsyncSession = Depends(get_db)
):
"""
根据 UUID 查询单个映射关系(包含关联的标注模板详情)
返回数据集映射关系以及关联的完整标注模板信息,包括:
- 映射基本信息
- 数据集信息
- Label Studio 项目信息
- 完整的标注模板配置(如果存在)
"""
try:
service = DatasetMappingService(db)
logger.info(f"Get mapping with template details: {mapping_id}")
# 获取映射,并包含完整的模板信息
mapping = await service.get_mapping_by_uuid(mapping_id, include_template=True)
if not mapping:
raise HTTPException(
status_code=404,
detail=f"Mapping not found: {mapping_id}"
)
logger.info(f"Found mapping: {mapping.id}, template_included: {mapping.template is not None}")
return StandardResponse(
code="0",
message="success",
data=mapping
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error getting mapping: {e}")
raise HTTPException(status_code=500, detail="Internal server error")
@router.get("/by-source/{dataset_id}", response_model=StandardResponse[PaginatedData[DatasetMappingResponse]])
async def get_mappings_by_source(
dataset_id: str,
page: int = Query(1, ge=1, description="页码(从1开始)"),
size: int = Query(20, ge=1, le=100, description="每页记录数"),
include_template: bool = Query(True, description="是否包含模板详情", alias="includeTemplate"),
db: AsyncSession = Depends(get_db)
):
"""
根据源数据集 ID 查询所有映射关系(分页,包含模板详情)
返回该数据集创建的所有标注项目(不包括已删除的),支持分页查询。
默认包含关联的完整标注模板信息。
参数:
- dataset_id: 数据集ID
- page: 页码(从1开始)
- pageSize: 每页记录数
- includeTemplate: 是否包含模板详情(默认true)
"""
try:
service = DatasetMappingService(db)
# 计算 skip
skip = (page - 1) * size
logger.info(f"Get mappings by source dataset id: {dataset_id}, page={page}, size={size}, include_template={include_template}")
# 获取数据和总数(包含模板信息)
mappings, total = await service.get_mappings_by_source_with_count(
dataset_id=dataset_id,
skip=skip,
limit=size,
include_template=include_template
)
# 计算总页数
total_pages = math.ceil(total / size) if total > 0 else 0
# 构造分页响应
paginated_data = PaginatedData(
page=page,
size=size,
total_elements=total,
total_pages=total_pages,
content=mappings
)
logger.info(f"Found {len(mappings)} mappings on page {page}, total: {total}, templates_included: {include_template}")
return StandardResponse(
code="0",
message="success",
data=paginated_data
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error getting mappings: {e}")
raise HTTPException(status_code=500, detail="Internal server error")
@router.delete("/{project_id}", response_model=StandardResponse[DeleteDatasetResponse])
async def delete_mapping(
project_id: str = Path(..., description="映射UUID(path param)"),
db: AsyncSession = Depends(get_db)
):
"""
删除映射关系和对应的 Label Studio 项目
通过 path 参数 `project_id` 指定要删除的映射(映射的 UUID)。
此操作会:
1. 删除 Label Studio 中的项目
2. 软删除数据库中的映射记录
"""
try:
logger.debug(f"Delete mapping request received: project_id={project_id!r}")
ls_client = LabelStudioClient(base_url=settings.label_studio_base_url,
token=settings.label_studio_user_token)
service = DatasetMappingService(db)
# 使用 mapping UUID 查询映射记录
logger.debug(f"Deleting by mapping UUID: {project_id}")
mapping = await service.get_mapping_by_uuid(project_id)
logger.debug(f"Mapping lookup result: {mapping}")
if not mapping:
raise HTTPException(
status_code=404,
detail=f"Mapping either not found or not specified."
)
id = mapping.id
labeling_project_id = mapping.labeling_project_id
logger.debug(f"Found mapping: {id}, Label Studio project ID: {labeling_project_id}")
# 1. 删除 Label Studio 项目
try:
logger.debug(f"Deleting Label Studio project: {labeling_project_id}")
delete_success = await ls_client.delete_project(int(labeling_project_id))
if delete_success:
logger.debug(f"Successfully deleted Label Studio project: {labeling_project_id}")
else:
logger.warning(f"Failed to delete Label Studio project or project not found: {labeling_project_id}")
except Exception as e:
logger.error(f"Error deleting Label Studio project: {e}")
# 继续执行,即使 Label Studio 项目删除失败也要删除映射记录
# 2. 软删除映射记录
soft_delete_success = await service.soft_delete_mapping(id)
logger.debug(f"Soft delete result for mapping {id}: {soft_delete_success}")
if not soft_delete_success:
raise HTTPException(
status_code=500,
detail="Failed to delete mapping record"
)
logger.info(f"Successfully deleted mapping: {id}, Label Studio project: {labeling_project_id}")
return StandardResponse(
code="0",
message="success",
data=DeleteDatasetResponse(
id=id,
status="success"
)
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error deleting mapping: {e}")
raise HTTPException(status_code=500, detail="Internal server error")
@@ -0,0 +1,404 @@
from fastapi import APIRouter, Depends, Query, Path, HTTPException
from sqlalchemy.ext.asyncio import AsyncSession
from typing import List, Optional, Dict, Any
from datetime import datetime
from pydantic import BaseModel, Field, ConfigDict
from app.core.exception import ErrorCodes, BusinessError, SuccessResponse
from app.db.session import get_db
from app.module.shared.schema import StandardResponse
from app.module.dataset import DatasetManagementService
from app.core.logging import get_logger
from app.core.config import settings
from ..client import LabelStudioClient
from ..service.sync import SyncService
from ..service.mapping import DatasetMappingService
from ..schema import (
SyncDatasetRequest,
SyncDatasetResponse,
SyncAnnotationsRequest,
SyncAnnotationsResponse,
UpdateFileTagsRequest,
UpdateFileTagsResponse,
UpdateFileTagsRequest,
UpdateFileTagsResponse
)
router = APIRouter(
prefix="/task",
tags=["annotation/task"]
)
logger = get_logger(__name__)
@router.post("/sync", response_model=StandardResponse[SyncDatasetResponse])
async def sync_dataset_content(
request: SyncDatasetRequest,
db: AsyncSession = Depends(get_db)
):
"""
Sync Dataset Content (Files and Annotations)
根据指定的mapping ID,同步DM程序数据集中的内容到Label Studio数据集中。
默认同时同步文件和标注数据。
"""
try:
ls_client = LabelStudioClient(base_url=settings.label_studio_base_url,
token=settings.label_studio_user_token)
dm_client = DatasetManagementService(db)
mapping_service = DatasetMappingService(db)
sync_service = SyncService(dm_client, ls_client, mapping_service)
logger.debug(f"Sync dataset content request: mapping_id={request.id}, sync_annotations={request.sync_annotations}")
# request.id validation
mapping = await mapping_service.get_mapping_by_uuid(request.id)
if not mapping:
raise HTTPException(
status_code=404,
detail=f"Mapping not found: {request.id}"
)
# Sync dataset files
result = await sync_service.sync_dataset_files(request.id, request.batch_size)
# Sync annotations if requested
if request.sync_annotations:
logger.info(f"Syncing annotations: direction={request.annotation_direction}")
# 根据方向执行标注同步
if request.annotation_direction == "ls_to_dm":
await sync_service.sync_annotations_from_ls_to_dm(
mapping,
request.batch_size,
request.overwrite,
sync_files_first=False,
)
elif request.annotation_direction == "dm_to_ls":
await sync_service.sync_annotations_from_dm_to_ls(
mapping,
request.batch_size,
request.overwrite_labeling_project,
sync_files_first=False,
)
elif request.annotation_direction == "bidirectional":
await sync_service.sync_annotations_bidirectional(
mapping,
request.batch_size,
request.overwrite,
request.overwrite_labeling_project,
sync_files_first=False,
)
logger.info(f"Sync completed: {result.synced_files}/{result.total_files} files")
return StandardResponse(
code="0",
message="success",
data=result
)
except HTTPException:
raise
except BusinessError as e:
# 业务异常已经由全局异常处理器处理
raise
except Exception as e:
logger.error(f"Error syncing dataset content: {e}")
raise
@router.post("/annotation/sync", response_model=StandardResponse[SyncAnnotationsResponse])
async def sync_annotations(
request: SyncAnnotationsRequest,
db: AsyncSession = Depends(get_db)
):
"""
Sync Annotations Only (Bidirectional Support)
同步指定 mapping 下的标注数据,支持单向或双向同步,基于时间戳自动解决冲突。
请求与响应由 Pydantic 模型 `SyncAnnotationsRequest` / `SyncAnnotationsResponse` 定义。
"""
try:
ls_client = LabelStudioClient(base_url=settings.label_studio_base_url,
token=settings.label_studio_user_token)
dm_client = DatasetManagementService(db)
mapping_service = DatasetMappingService(db)
sync_service = SyncService(dm_client, ls_client, mapping_service)
logger.info(f"Sync annotations request: mapping_id={request.id}, direction={request.direction}, overwrite={request.overwrite}, overwrite_ls={request.overwrite_labeling_project}")
# 验证映射是否存在
mapping = await mapping_service.get_mapping_by_uuid(request.id)
if not mapping:
raise HTTPException(
status_code=404,
detail=f"Mapping not found: {request.id}"
)
# 根据方向执行同步
if request.direction == "ls_to_dm":
result = await sync_service.sync_annotations_from_ls_to_dm(
mapping,
request.batch_size,
request.overwrite,
sync_files_first=True,
)
elif request.direction == "dm_to_ls":
result = await sync_service.sync_annotations_from_dm_to_ls(
mapping,
request.batch_size,
request.overwrite_labeling_project,
sync_files_first=True,
)
elif request.direction == "bidirectional":
result = await sync_service.sync_annotations_bidirectional(
mapping,
request.batch_size,
request.overwrite,
request.overwrite_labeling_project,
sync_files_first=True,
)
else:
raise HTTPException(
status_code=400,
detail=f"Invalid direction: {request.direction}"
)
logger.info(f"Annotation sync completed: synced_to_dm={result.synced_to_dm}, synced_to_ls={result.synced_to_ls}, conflicts_resolved={result.conflicts_resolved}")
return StandardResponse(
code="0",
message="success",
data=result
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error syncing annotations: {e}")
raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
@router.get("/check-ls-connection")
async def check_label_studio_connection():
"""
Check Label Studio Connection Status
诊断 Label Studio 连接并返回简要连接信息(状态、base URL、token 摘要、项目统计)。
"""
try:
ls_client = LabelStudioClient(
base_url=settings.label_studio_base_url,
token=settings.label_studio_user_token
)
# 尝试获取项目列表来测试连接
try:
response = await ls_client.client.get("/api/projects")
response.raise_for_status()
projects = response.json()
token_display = settings.label_studio_user_token[:10] + "..." if settings.label_studio_user_token else "None"
return StandardResponse(
code="0",
message="success",
data={
"status": "connected",
"base_url": settings.label_studio_base_url,
"token": token_display,
"projects_count": len(projects.get("results", [])) if isinstance(projects, dict) else len(projects),
"message": "Successfully connected to Label Studio"
}
)
except Exception as e:
token_display = settings.label_studio_user_token[:10] + "..." if settings.label_studio_user_token else "None"
return StandardResponse(
code="common.500",
message="error",
data={
"status": "disconnected",
"base_url": settings.label_studio_base_url,
"token": token_display,
"error": str(e),
"message": f"Failed to connect to Label Studio: {str(e)}",
"troubleshooting": [
"1. Check if Label Studio is running: docker ps | grep label-studio",
"2. Verify LABEL_STUDIO_BASE_URL in .env file",
"3. Verify LABEL_STUDIO_USER_TOKEN is valid",
"4. Check network connectivity between services"
]
}
)
except Exception as e:
logger.error(f"Error checking Label Studio connection: {e}")
raise HTTPException(status_code=500, detail=str(e))
@router.put(
"/{file_id}",
response_model=StandardResponse[UpdateFileTagsResponse],
)
async def update_file_tags(
request: UpdateFileTagsRequest,
file_id: str = Path(..., description="文件ID"),
db: AsyncSession = Depends(get_db)
):
"""
Update File Tags (Partial Update with Auto Format Conversion)
接收部分标签更新并合并到指定文件(只修改提交的标签,其余保持不变),并更新 `tags_updated_at`。
支持两种标签格式:
1. 简化格式(外部用户提交):
[{"from_name": "label", "to_name": "image", "values": ["cat", "dog"]}]
2. 完整格式(内部存储):
[{"id": "...", "from_name": "label", "to_name": "image", "type": "choices",
"value": {"choices": ["cat", "dog"]}}]
系统会自动根据数据集关联的模板将简化格式转换为完整格式。
请求与响应使用 Pydantic 模型 `UpdateFileTagsRequest` / `UpdateFileTagsResponse`。
"""
service = DatasetManagementService(db)
# 首先获取文件所属的数据集
from sqlalchemy.future import select
from app.db.models import DatasetFiles
result = await db.execute(
select(DatasetFiles).where(DatasetFiles.id == file_id)
)
file_record = result.scalar_one_or_none()
if not file_record:
raise HTTPException(status_code=404, detail=f"File not found: {file_id}")
dataset_id = str(file_record.dataset_id) # type: ignore - Convert Column to str
# 查找数据集关联的模板ID
from ..service.mapping import DatasetMappingService
mapping_service = DatasetMappingService(db)
template_id = await mapping_service.get_template_id_by_dataset_id(dataset_id)
if template_id:
logger.info(f"Found template {template_id} for dataset {dataset_id}, will auto-convert tag format")
else:
logger.warning(f"No template found for dataset {dataset_id}, tags must be in full format")
# 更新标签(如果有模板ID则自动转换格式)
success, error_msg, updated_at = await service.update_file_tags_partial(
file_id=file_id,
new_tags=request.tags,
template_id=template_id # 传递模板ID以启用自动转换
)
if not success:
if "not found" in (error_msg or "").lower():
raise HTTPException(status_code=404, detail=error_msg)
raise HTTPException(status_code=500, detail=error_msg or "更新标签失败")
# 重新获取更新后的文件记录(获取完整标签列表)
result = await db.execute(
select(DatasetFiles).where(DatasetFiles.id == file_id)
)
file_record = result.scalar_one_or_none()
if not file_record:
raise HTTPException(status_code=404, detail=f"File not found: {file_id}")
response_data = UpdateFileTagsResponse(
fileId=file_id,
tags=file_record.tags or [], # type: ignore
tagsUpdatedAt=updated_at or datetime.now()
)
return StandardResponse(
code="0",
message="标签更新成功",
data=response_data
)
@router.put(
"/{file_id}/replace",
response_model=StandardResponse[UpdateFileTagsResponse],
)
async def replace_file_tags(
request: UpdateFileTagsRequest,
file_id: str = Path(..., description="文件ID"),
db: AsyncSession = Depends(get_db)
):
"""
Replace File Tags (Full Replacement with Auto Format Conversion)
完全替换文件的标签列表。与部分更新不同,此方法会:
- 空列表会清空所有标签
- 新标签列表完全替换原有标签(不合并)
支持两种标签格式:
1. 简化格式:
[{"from_name": "label", "to_name": "image", "values": ["cat", "dog"]}]
2. 完整格式:
[{"id": "...", "from_name": "label", "to_name": "image", "type": "choices",
"value": {"choices": ["cat", "dog"]}}]
系统会自动根据数据集关联的模板将简化格式转换为完整格式。
"""
service = DatasetManagementService(db)
from sqlalchemy.future import select
from app.db.models import DatasetFiles
result = await db.execute(
select(DatasetFiles).where(DatasetFiles.id == file_id)
)
file_record = result.scalar_one_or_none()
if not file_record:
raise HTTPException(status_code=404, detail=f"File not found: {file_id}")
dataset_id = str(file_record.dataset_id)
mapping_service = DatasetMappingService(db)
template_id = await mapping_service.get_template_id_by_dataset_id(dataset_id)
if template_id:
logger.info(f"Found template {template_id} for dataset {dataset_id}, will auto-convert tag format")
else:
logger.warning(f"No template found for dataset {dataset_id}, tags must be in full format")
success, error_msg, updated_at = await service.replace_file_tags(
file_id=file_id,
new_tags=request.tags,
template_id=template_id
)
if not success:
if "not found" in (error_msg or "").lower():
raise HTTPException(status_code=404, detail=error_msg)
raise HTTPException(status_code=500, detail=error_msg or "替换标签失败")
result = await db.execute(
select(DatasetFiles).where(DatasetFiles.id == file_id)
)
file_record = result.scalar_one_or_none()
if not file_record:
raise HTTPException(status_code=404, detail=f"File not found: {file_id}")
response_data = UpdateFileTagsResponse(
fileId=file_id,
tags=file_record.tags or [],
tagsUpdatedAt=updated_at or datetime.now()
)
return StandardResponse(
code="0",
message="标签替换成功",
data=response_data
)
@@ -0,0 +1,132 @@
"""
Annotation Template API Endpoints
"""
from typing import Optional
from fastapi import APIRouter, Depends, HTTPException, Query
from sqlalchemy.ext.asyncio import AsyncSession
from app.db.session import get_db
from app.module.shared.schema import StandardResponse
from app.module.annotation.schema import (
CreateAnnotationTemplateRequest,
UpdateAnnotationTemplateRequest,
AnnotationTemplateResponse,
AnnotationTemplateListResponse
)
from app.module.annotation.service.template import AnnotationTemplateService
router = APIRouter(prefix="/template", tags=["annotation/template"])
template_service = AnnotationTemplateService()
@router.post(
"",
response_model=StandardResponse[AnnotationTemplateResponse],
)
async def create_template(
request: CreateAnnotationTemplateRequest,
db: AsyncSession = Depends(get_db)
):
"""
创建新的标注模板
- **name**: 模板名称(必填,最多100字符)
- **description**: 模板描述(可选,最多500字符)
- **dataType**: 数据类型(必填)
- **labelingType**: 标注类型(必填)
- **configuration**: 标注配置(必填,包含labels和objects)
- **style**: 样式配置(默认horizontal)
- **category**: 模板分类(默认custom)
"""
template = await template_service.create_template(db, request)
return StandardResponse(code="0", message="success", data=template)
@router.get(
"/{template_id}",
response_model=StandardResponse[AnnotationTemplateResponse],
)
async def get_template(
template_id: str,
db: AsyncSession = Depends(get_db)
):
"""
根据ID获取模板详情
"""
template = await template_service.get_template(db, template_id)
if not template:
raise HTTPException(status_code=404, detail="Template not found")
return StandardResponse(code="0", message="success", data=template)
@router.get(
"",
response_model=StandardResponse[AnnotationTemplateListResponse],
)
async def list_template(
page: int = Query(1, ge=1, description="页码"),
size: int = Query(10, ge=1, le=100, description="每页大小"),
category: Optional[str] = Query(None, description="分类筛选"),
dataType: Optional[str] = Query(None, alias="dataType", description="数据类型筛选"),
labelingType: Optional[str] = Query(None, alias="labelingType", description="标注类型筛选"),
builtIn: Optional[bool] = Query(None, alias="builtIn", description="是否内置模板"),
db: AsyncSession = Depends(get_db)
):
"""
获取模板列表,支持分页和筛选
- **page**: 页码(从1开始)
- **size**: 每页大小(1-100)
- **category**: 模板分类筛选
- **dataType**: 数据类型筛选
- **labelingType**: 标注类型筛选
- **builtIn**: 是否只显示内置模板
"""
templates = await template_service.list_templates(
db=db,
page=page,
size=size,
category=category,
data_type=dataType,
labeling_type=labelingType,
built_in=builtIn
)
return StandardResponse(code="0", message="success", data=templates)
@router.put(
"/{template_id}",
response_model=StandardResponse[AnnotationTemplateResponse],
)
async def update_template(
template_id: str,
request: UpdateAnnotationTemplateRequest,
db: AsyncSession = Depends(get_db)
):
"""
更新模板信息
所有字段都是可选的,只更新提供的字段
"""
template = await template_service.update_template(db, template_id, request)
if not template:
raise HTTPException(status_code=404, detail="Template not found")
return StandardResponse(code="0", message="success", data=template)
@router.delete(
"/{template_id}",
response_model=StandardResponse[bool],
)
async def delete_template(
template_id: str,
db: AsyncSession = Depends(get_db)
):
"""
删除模板(软删除)
"""
success = await template_service.delete_template(db, template_id)
if not success:
raise HTTPException(status_code=404, detail="Template not found")
return StandardResponse(code="0", message="success", data=True)
@@ -0,0 +1,54 @@
from .config import (
ConfigResponse,
TagConfigResponse
)
from .sync import (
SyncDatasetRequest,
SyncDatasetResponse,
SyncAnnotationsRequest,
SyncAnnotationsResponse,
)
from .tag import (
UpdateFileTagsRequest,
UpdateFileTagsResponse,
)
from .template import (
CreateAnnotationTemplateRequest,
UpdateAnnotationTemplateRequest,
AnnotationTemplateResponse,
AnnotationTemplateListResponse
)
from .mapping import (
DatasetMappingCreateRequest,
DatasetMappingCreateResponse,
DatasetMappingUpdateRequest,
DatasetMappingResponse,
DeleteDatasetResponse,
)
# Rebuild model to resolve forward references
DatasetMappingResponse.model_rebuild()
__all__ = [
"ConfigResponse",
"TagConfigResponse",
"DatasetMappingCreateRequest",
"DatasetMappingCreateResponse",
"DatasetMappingUpdateRequest",
"DatasetMappingResponse",
"SyncDatasetRequest",
"SyncDatasetResponse",
"SyncAnnotationsRequest",
"SyncAnnotationsResponse",
"DeleteDatasetResponse",
"UpdateFileTagsRequest",
"UpdateFileTagsResponse",
"CreateAnnotationTemplateRequest",
"UpdateAnnotationTemplateRequest",
"AnnotationTemplateResponse",
"AnnotationTemplateListResponse",
]
@@ -0,0 +1,129 @@
"""Schemas for Auto Annotation tasks"""
from __future__ import annotations
from typing import List, Optional, Dict, Any
from datetime import datetime
from pydantic import BaseModel, Field, ConfigDict
class AutoAnnotationConfig(BaseModel):
"""自动标注任务配置(与前端 payload 对齐)"""
model_size: str = Field(alias="modelSize", description="模型规模: n/s/m/l/x")
conf_threshold: float = Field(alias="confThreshold", description="置信度阈值 0-1")
target_classes: List[int] = Field(
default_factory=list,
alias="targetClasses",
description="目标类别ID列表,空表示全部类别",
)
output_dataset_name: Optional[str] = Field(
default=None,
alias="outputDatasetName",
description="自动标注结果要写入的新数据集名称(可选)",
)
model_config = ConfigDict(populate_by_name=True)
class CreateAutoAnnotationTaskRequest(BaseModel):
"""创建自动标注任务的请求体,对齐前端 CreateAutoAnnotationDialog 发送的结构"""
name: str = Field(..., min_length=1, max_length=255, description="任务名称")
dataset_id: str = Field(..., alias="datasetId", description="数据集ID")
config: AutoAnnotationConfig = Field(..., description="任务配置")
file_ids: Optional[List[str]] = Field(None, alias="fileIds", description="要处理的文件ID列表,为空则处理数据集中所有图像")
model_config = ConfigDict(populate_by_name=True)
class AutoAnnotationTaskResponse(BaseModel):
"""自动标注任务响应模型(列表/详情均可复用)"""
id: str = Field(..., description="任务ID")
name: str = Field(..., description="任务名称")
dataset_id: str = Field(..., alias="datasetId", description="数据集ID")
dataset_name: Optional[str] = Field(None, alias="datasetName", description="数据集名称")
source_datasets: Optional[List[str]] = Field(
default=None,
alias="sourceDatasets",
description="本任务实际处理涉及到的所有数据集名称列表",
)
config: Dict[str, Any] = Field(..., description="任务配置")
status: str = Field(..., description="任务状态")
progress: int = Field(..., description="任务进度 0-100")
total_images: int = Field(..., alias="totalImages", description="总图片数")
processed_images: int = Field(..., alias="processedImages", description="已处理图片数")
detected_objects: int = Field(..., alias="detectedObjects", description="检测到的对象总数")
output_path: Optional[str] = Field(None, alias="outputPath", description="输出路径")
error_message: Optional[str] = Field(None, alias="errorMessage", description="错误信息")
file_ids: Optional[List[str]] = Field(
default=None,
alias="fileIds",
description="本次任务实际处理的文件ID列表,为空表示处理数据集全部文件",
)
created_at: datetime = Field(..., alias="createdAt", description="创建时间")
updated_at: Optional[datetime] = Field(None, alias="updatedAt", description="更新时间")
completed_at: Optional[datetime] = Field(None, alias="completedAt", description="完成时间")
model_config = ConfigDict(populate_by_name=True, from_attributes=True)
class AutoAnnotationTaskListResponse(BaseModel):
"""自动标注任务列表响应,目前前端直接使用数组,这里预留分页结构"""
content: List[AutoAnnotationTaskResponse] = Field(..., description="任务列表")
total: int = Field(..., description="总数")
model_config = ConfigDict(populate_by_name=True)
class UpdateAutoAnnotationTaskFilesRequest(BaseModel):
"""更新自动标注任务所关联的数据集文件。
- dataset_id: 主数据集 ID;若前端未显式传入,则在接口层回退为现有任务的 dataset_id。
- file_ids: 选中的文件 ID 列表,允许跨多个数据集;
实际执行自动标注时仅会对 tags 为空的文件重新推理。
"""
dataset_id: Optional[str] = Field(
default=None,
alias="datasetId",
description="主数据集ID(可选,不传则沿用原任务的datasetId)",
)
file_ids: List[str] = Field(
default_factory=list,
alias="fileIds",
description="要参与本自动标注任务的数据集文件ID列表,允许跨多个数据集",
)
model_config = ConfigDict(populate_by_name=True)
class ImportFromLabelStudioRequest(BaseModel):
"""从 Label Studio 导入标注结果到 DM 数据集的请求体。
- target_dataset_id: 目标数据集ID(可选);若未提供,则自动使用源任务/映射所关联的数据集;
- export_format: Label Studio 导出格式(如 JSON/JSON_MIN/CSV/TSV/COCO/YOLO 等);
- file_name: 可选,自定义保存到数据集中的文件名称(不含路径);
若包含扩展名或路径分隔符,服务端会自动裁剪,仅保留文件名主体并附加正确后缀。
"""
target_dataset_id: Optional[str] = Field(
default=None,
alias="targetDatasetId",
description="导入目标数据集ID(可选,不传则使用源数据集)",
)
export_format: str = Field(
default="JSON",
alias="exportFormat",
description="Label Studio 导出格式 (JSON/COCO/YOLO 等)",
)
file_name: Optional[str] = Field(
default=None,
alias="fileName",
description="自定义导出文件名(可选,不含路径,扩展名将按导出格式自动附加)",
)
model_config = ConfigDict(populate_by_name=True)
@@ -0,0 +1,36 @@
from typing import Dict, Any, List, Optional
from pydantic import BaseModel, Field, ConfigDict
from app.module.shared.schema import BaseResponseModel
from app.module.shared.schema import StandardResponse
class ConfigResponse(BaseResponseModel):
"""配置信息响应模型"""
label_studio_url: str = Field(..., description="Label Studio基础URL")
class _TagAttributeConfig(BaseModel):
"""标签属性配置"""
type: Optional[str] = Field(None, description="属性类型: boolean/string/number")
values: Optional[List[str]] = Field(None, description="允许的枚举值列表")
default: Optional[Any] = Field(None, description="默认值")
description: Optional[str] = Field(None, description="属性描述")
model_config = ConfigDict(populate_by_name=True)
class _TagDefinition(BaseModel):
"""标签定义"""
description: str = Field(..., description="标签描述")
required_attrs: List[str] = Field(default_factory=list, alias="requiredAttrs", description="必需属性列表")
optional_attrs: Dict[str, _TagAttributeConfig] = Field(default_factory=dict, alias="optionalAttrs", description="可选属性配置")
requires_children: bool = Field(default=False, alias="requiresChildren", description="是否需要子元素")
child_tag: Optional[str] = Field(None, alias="childTag", description="子元素标签名")
child_required_attrs: Optional[List[str]] = Field(None, alias="childRequiredAttrs", description="子元素必需属性")
category: Optional[str] = Field(None, description="标签分类")
class TagConfigResponse(BaseResponseModel):
"""标签配置响应"""
objects: Dict[str, _TagDefinition] = Field(default_factory=dict, description="对象标签配置")
controls: Dict[str, _TagDefinition] = Field(default_factory=dict, description="控件标签配置")
@@ -0,0 +1,62 @@
from pydantic import Field, BaseModel
from typing import Optional, TYPE_CHECKING, List
from datetime import datetime
from app.module.shared.schema import BaseResponseModel
from app.module.shared.schema import StandardResponse
if TYPE_CHECKING:
from .template import AnnotationTemplateResponse
class DatasetMappingCreateRequest(BaseModel):
"""数据集映射 创建 请求模型
Accept both snake_case and camelCase field names from frontend JSON by
declaring explicit aliases. Frontend sends `datasetId`, `name`,
`description`, `templateId`, `fileIds` (camelCase), so provide aliases so
pydantic will map them to the internal attributes used in the service code
(dataset_id, name, description, template_id, file_ids).
"""
dataset_id: str = Field(..., alias="datasetId", description="源数据集ID")
name: Optional[str] = Field(None, alias="name", description="标注项目名称")
description: Optional[str] = Field(None, alias="description", description="标注项目描述")
template_id: Optional[str] = Field(None, alias="templateId", description="标注模板ID")
file_ids: Optional[List[str]] = Field(None, alias="fileIds", description="限制同步的文件ID列表,可跨多个数据集")
class Config:
# allow population by field name when constructing model programmatically
validate_by_name = True
class DatasetMappingCreateResponse(BaseResponseModel):
"""数据集映射 创建 响应模型"""
id: str = Field(..., description="映射UUID")
labeling_project_id: str = Field(..., description="Label Studio项目ID")
labeling_project_name: str = Field(..., description="Label Studio项目名称")
class DatasetMappingUpdateRequest(BaseResponseModel):
"""数据集映射 更新 请求模型"""
dataset_id: Optional[str] = Field(None, description="源数据集ID")
class DatasetMappingResponse(BaseModel):
"""数据集映射 查询 响应模型"""
id: str = Field(..., description="映射UUID")
dataset_id: str = Field(..., alias="datasetId", description="源数据集ID")
dataset_name: Optional[str] = Field(None, alias="datasetName", description="数据集名称")
labeling_project_id: str = Field(..., alias="labelingProjectId", description="标注项目ID")
name: Optional[str] = Field(None, description="标注项目名称")
description: Optional[str] = Field(None, description="标注项目描述")
template_id: Optional[str] = Field(None, alias="templateId", description="关联的模板ID")
template: Optional['AnnotationTemplateResponse'] = Field(None, description="关联的标注模板详情")
created_at: datetime = Field(..., alias="createdAt", description="创建时间")
updated_at: Optional[datetime] = Field(None, alias="updatedAt", description="更新时间")
deleted_at: Optional[datetime] = Field(None, alias="deletedAt", description="删除时间")
class Config:
from_attributes = True
populate_by_name = True
class DeleteDatasetResponse(BaseResponseModel):
"""删除数据集响应模型"""
id: str = Field(..., description="映射UUID")
status: str = Field(..., description="删除状态")
@@ -0,0 +1,92 @@
from typing import Literal, List, Dict, Any, Optional
from datetime import datetime
from pydantic import Field, ConfigDict
from app.module.shared.schema import BaseResponseModel
from app.module.shared.schema import StandardResponse
class SyncDatasetRequest(BaseResponseModel):
"""同步数据集请求模型"""
model_config = ConfigDict(populate_by_name=True)
id: str = Field(..., description="映射ID(mapping UUID)")
batch_size: int = Field(50, ge=1, le=100, description="批处理大小", alias="batchSize")
file_priority: Literal[0, 1] = Field(0, description="0 数据集为主,1 标注平台为主", alias="filePriority")
label_priority: Literal[0, 1] = Field(0, description="0 数据集为主,1 标注平台为主", alias="labelPriority")
sync_annotations: bool = Field(True, description="是否同步标注数据", alias="syncAnnotations")
annotation_direction: Literal["ls_to_dm", "dm_to_ls", "bidirectional"] = Field(
"bidirectional",
description="标注同步方向: ls_to_dm(Label Studio到数据集), dm_to_ls(数据集到Label Studio), bidirectional(双向)",
alias="annotationDirection"
)
overwrite: bool = Field(
True,
description="是否覆盖DeepData中的标注(基于时间戳比较)"
)
overwrite_labeling_project: bool = Field(
True,
description="是否覆盖Label Studio中的标注(基于时间戳比较)",
alias="overwriteLabelingProject"
)
class SyncDatasetResponse(BaseResponseModel):
"""同步数据集响应模型"""
id: str = Field(..., description="映射UUID")
status: str = Field(..., description="同步状态")
synced_files: int = Field(..., description="已同步文件数量")
total_files: int = Field(0, description="总文件数量")
message: str = Field(..., description="响应消息")
class SyncDatasetResponseStd(StandardResponse[SyncDatasetResponse]):
pass
class SyncAnnotationsRequest(BaseResponseModel):
"""同步标注请求模型
使用camelCase作为API接口字段名(通过alias),但Python代码内部使用snake_case。
Pydantic会自动处理两种格式的转换。
"""
model_config = ConfigDict(populate_by_name=True)
id: str = Field(..., description="映射ID(mapping UUID)")
batch_size: int = Field(50, ge=1, le=100, description="批处理大小", alias="batchSize")
direction: Literal["ls_to_dm", "dm_to_ls", "bidirectional"] = Field(
"bidirectional",
description="同步方向: ls_to_dm(Label Studio到数据集), dm_to_ls(数据集到Label Studio), bidirectional(双向)"
)
overwrite: bool = Field(
True,
description="是否覆盖DeepData中的标注(基于时间戳比较)。True时,如果Label Studio的标注更新时间更新,则覆盖DeepData的标注"
)
overwrite_labeling_project: bool = Field(
True,
description="是否覆盖Label Studio中的标注(基于时间戳比较)。True时,如果DeepData的标注更新时间更新,则覆盖Label Studio的标注",
alias="overwriteLabelingProject"
)
class TagInfo(BaseResponseModel):
"""标注信息结构(不包含时间戳,时间戳存储在文件级别的tags_updated_at字段)"""
from_name: str = Field(..., description="标注工具名称")
to_name: str = Field(..., description="目标对象名称")
type: str = Field(..., description="标注类型")
values: Dict[str, Any] = Field(..., description="标注值")
class SyncAnnotationsResponse(BaseResponseModel):
"""同步标注响应模型"""
id: str = Field(..., description="映射UUID")
status: str = Field(..., description="同步状态: success/partial/error")
synced_to_dm: int = Field(0, description="同步到数据集的标注数量")
synced_to_ls: int = Field(0, description="同步到Label Studio的标注数量")
skipped: int = Field(0, description="跳过的标注数量")
failed: int = Field(0, description="失败的标注数量")
conflicts_resolved: int = Field(0, description="解决的冲突数量")
message: str = Field(..., description="响应消息")
class SyncAnnotationsResponseStd(StandardResponse[SyncAnnotationsResponse]):
pass
@@ -0,0 +1,17 @@
from datetime import datetime
from typing import List, Dict, Any
from pydantic import BaseModel, Field
from app.module.shared.schema import BaseResponseModel
class UpdateFileTagsRequest(BaseModel):
"""更新文件标签请求"""
tags: List[Dict[str, Any]] = Field(..., description="要更新的标签列表(部分更新)")
class UpdateFileTagsResponse(BaseResponseModel):
"""更新文件标签响应"""
file_id: str = Field(..., alias="fileId", description="文件ID")
tags: List[Dict[str, Any]] = Field(..., description="更新后的完整标签列表")
tags_updated_at: datetime = Field(..., alias="tagsUpdatedAt", description="标签更新时间")
@@ -0,0 +1,94 @@
"""
Annotation Template Schemas
"""
from typing import List, Dict, Any, Optional
from datetime import datetime
from pydantic import BaseModel, Field, ConfigDict
class LabelDefinition(BaseModel):
"""标签定义"""
from_name: str = Field(alias="fromName", description="控件名称")
to_name: str = Field(alias="toName", description="目标对象名称")
type: str = Field(description="控件类型: choices/rectanglelabels/polygonlabels/textarea/etc")
options: Optional[List[str]] = Field(None, description="选项列表(用于choices类型)")
labels: Optional[List[str]] = Field(None, description="标签列表(用于rectanglelabels等类型)")
required: bool = Field(False, description="是否必填")
choice: Optional[str] = Field(None, description="选项设置")
description: Optional[str] = Field(None, description="标签描述")
model_config = ConfigDict(populate_by_name=True)
class ObjectDefinition(BaseModel):
"""对象定义"""
name: str = Field(description="对象标识符")
type: str = Field(description="对象类型: Image/Text/Audio/Video/etc")
value: str = Field(description="变量名,如$image")
model_config = ConfigDict(populate_by_name=True)
class TemplateConfiguration(BaseModel):
"""模板配置结构"""
labels: List[LabelDefinition] = Field(description="标签定义列表")
objects: List[ObjectDefinition] = Field(description="对象定义列表")
metadata: Optional[Dict[str, Any]] = Field(None, description="额外元数据")
model_config = ConfigDict(populate_by_name=True)
class CreateAnnotationTemplateRequest(BaseModel):
"""创建标注模板请求"""
name: str = Field(..., min_length=1, max_length=100, description="模板名称")
description: Optional[str] = Field(None, max_length=500, description="模板描述")
data_type: str = Field(alias="dataType", description="数据类型")
labeling_type: str = Field(alias="labelingType", description="标注类型")
configuration: TemplateConfiguration = Field(..., description="标注配置")
style: str = Field(default="horizontal", description="样式配置")
category: str = Field(default="custom", description="模板分类")
model_config = ConfigDict(populate_by_name=True)
class UpdateAnnotationTemplateRequest(BaseModel):
"""更新标注模板请求"""
name: Optional[str] = Field(None, min_length=1, max_length=100, description="模板名称")
description: Optional[str] = Field(None, max_length=500, description="模板描述")
data_type: Optional[str] = Field(None, alias="dataType", description="数据类型")
labeling_type: Optional[str] = Field(None, alias="labelingType", description="标注类型")
configuration: Optional[TemplateConfiguration] = Field(None, description="标注配置")
style: Optional[str] = Field(None, description="样式配置")
category: Optional[str] = Field(None, description="模板分类")
model_config = ConfigDict(populate_by_name=True)
class AnnotationTemplateResponse(BaseModel):
"""标注模板响应"""
id: str = Field(..., description="模板ID")
name: str = Field(..., description="模板名称")
description: Optional[str] = Field(None, description="模板描述")
data_type: str = Field(alias="dataType", description="数据类型")
labeling_type: str = Field(alias="labelingType", description="标注类型")
configuration: TemplateConfiguration = Field(..., description="标注配置")
label_config: Optional[str] = Field(None, alias="labelConfig", description="生成的Label Studio XML配置")
style: str = Field(..., description="样式配置")
category: str = Field(..., description="模板分类")
built_in: bool = Field(alias="builtIn", description="是否内置模板")
version: str = Field(..., description="版本号")
created_at: datetime = Field(alias="createdAt", description="创建时间")
updated_at: Optional[datetime] = Field(None, alias="updatedAt", description="更新时间")
model_config = ConfigDict(populate_by_name=True, from_attributes=True)
class AnnotationTemplateListResponse(BaseModel):
"""模板列表响应"""
content: List[AnnotationTemplateResponse] = Field(..., description="模板列表")
totalElements: int = Field(..., description="总数")
page: int = Field(..., description="当前页")
size: int = Field(..., description="每页大小")
total_pages: int = Field(alias="totalPages", description="总页数")
model_config = ConfigDict(populate_by_name=True)
@@ -0,0 +1,203 @@
"""Service layer for Auto Annotation tasks"""
from __future__ import annotations
from typing import List, Optional
from datetime import datetime
from uuid import uuid4
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from app.db.models.annotation_management import AutoAnnotationTask
from app.db.models.dataset_management import Dataset, DatasetFiles
from ..schema.auto import (
CreateAutoAnnotationTaskRequest,
AutoAnnotationTaskResponse,
)
class AutoAnnotationTaskService:
"""自动标注任务服务(仅管理任务元数据,真正执行由 runtime 负责)"""
async def create_task(
self,
db: AsyncSession,
request: CreateAutoAnnotationTaskRequest,
dataset_name: Optional[str] = None,
total_images: int = 0,
) -> AutoAnnotationTaskResponse:
"""创建自动标注任务,初始状态为 pending。
这里仅插入任务记录,不负责真正执行 YOLO 推理,
后续可以由调度器/worker 读取该表并更新进度。
"""
now = datetime.now()
task = AutoAnnotationTask(
id=str(uuid4()),
name=request.name,
dataset_id=request.dataset_id,
dataset_name=dataset_name,
config=request.config.model_dump(by_alias=True),
file_ids=request.file_ids, # 存储用户选择的文件ID列表
status="pending",
progress=0,
total_images=total_images,
processed_images=0,
detected_objects=0,
created_at=now,
updated_at=now,
)
db.add(task)
await db.commit()
await db.refresh(task)
# 创建后附带 sourceDatasets 信息(通常只有一个原始数据集)
resp = AutoAnnotationTaskResponse.model_validate(task)
try:
resp.source_datasets = await self._compute_source_datasets(db, task)
except Exception:
resp.source_datasets = [dataset_name] if dataset_name else [request.dataset_id]
return resp
async def list_tasks(self, db: AsyncSession) -> List[AutoAnnotationTaskResponse]:
"""获取未软删除的自动标注任务列表,按创建时间倒序。"""
result = await db.execute(
select(AutoAnnotationTask)
.where(AutoAnnotationTask.deleted_at.is_(None))
.order_by(AutoAnnotationTask.created_at.desc())
)
tasks: List[AutoAnnotationTask] = list(result.scalars().all())
responses: List[AutoAnnotationTaskResponse] = []
for task in tasks:
resp = AutoAnnotationTaskResponse.model_validate(task)
try:
resp.source_datasets = await self._compute_source_datasets(db, task)
except Exception:
# 出错时降级为单个 datasetName/datasetId
fallback_name = getattr(task, "dataset_name", None)
fallback_id = getattr(task, "dataset_id", "")
resp.source_datasets = [fallback_name] if fallback_name else [fallback_id]
responses.append(resp)
return responses
async def get_task(self, db: AsyncSession, task_id: str) -> Optional[AutoAnnotationTaskResponse]:
result = await db.execute(
select(AutoAnnotationTask).where(
AutoAnnotationTask.id == task_id,
AutoAnnotationTask.deleted_at.is_(None),
)
)
task = result.scalar_one_or_none()
if not task:
return None
resp = AutoAnnotationTaskResponse.model_validate(task)
try:
resp.source_datasets = await self._compute_source_datasets(db, task)
except Exception:
fallback_name = getattr(task, "dataset_name", None)
fallback_id = getattr(task, "dataset_id", "")
resp.source_datasets = [fallback_name] if fallback_name else [fallback_id]
return resp
async def _compute_source_datasets(
self,
db: AsyncSession,
task: AutoAnnotationTask,
) -> List[str]:
"""根据任务的 file_ids 推断实际涉及到的所有数据集名称。
- 如果存在 file_ids,则通过 t_dm_dataset_files 反查 dataset_id,再关联 t_dm_datasets 获取名称;
- 如果没有 file_ids,则退回到任务上冗余的 dataset_name/dataset_id。
"""
file_ids = task.file_ids or []
if file_ids:
stmt = (
select(Dataset.name)
.join(DatasetFiles, Dataset.id == DatasetFiles.dataset_id)
.where(DatasetFiles.id.in_(file_ids))
.distinct()
)
result = await db.execute(stmt)
names = [row[0] for row in result.fetchall() if row[0]]
if names:
return names
# 回退:只显示一个数据集
if task.dataset_name:
return [task.dataset_name]
if task.dataset_id:
return [task.dataset_id]
return []
async def soft_delete_task(self, db: AsyncSession, task_id: str) -> bool:
result = await db.execute(
select(AutoAnnotationTask).where(
AutoAnnotationTask.id == task_id,
AutoAnnotationTask.deleted_at.is_(None),
)
)
task = result.scalar_one_or_none()
if not task:
return False
task.deleted_at = datetime.now()
await db.commit()
return True
async def update_task_files(
self,
db: AsyncSession,
task_id: str,
*,
dataset_id: str,
file_ids: List[str],
) -> Optional[AutoAnnotationTaskResponse]:
"""更新自动标注任务关联的数据集与文件列表。
- 覆盖任务的 dataset_id 与 file_ids;
- 将任务重置为 pending,供 worker 重新调度;
- 保留已有的 output_path 与统计信息,确保同一任务始终复用同一个输出数据集,
worker 内部会根据 output_path 中已有的 images 仅对新增文件执行自动标注。
"""
result = await db.execute(
select(AutoAnnotationTask).where(
AutoAnnotationTask.id == task_id,
AutoAnnotationTask.deleted_at.is_(None),
)
)
task = result.scalar_one_or_none()
if not task:
return None
now = datetime.now()
task.dataset_id = dataset_id
task.file_ids = file_ids or []
# 将任务标记为待处理,让 worker 在同一个输出目录下仅对新增文件执行自动标注
task.status = "pending"
task.progress = 0
task.updated_at = now
db.add(task)
await db.commit()
await db.refresh(task)
resp = AutoAnnotationTaskResponse.model_validate(task)
try:
resp.source_datasets = await self._compute_source_datasets(db, task)
except Exception:
fallback_name = getattr(task, "dataset_name", None)
fallback_id = getattr(task, "dataset_id", "")
resp.source_datasets = [fallback_name] if fallback_name else [fallback_id]
return resp
@@ -0,0 +1,229 @@
from __future__ import annotations
from datetime import datetime
from typing import Any, Dict, List, Optional
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.config import settings
from app.core.logging import get_logger
from app.db.models import DatasetFiles
from ..client import LabelStudioClient
logger = get_logger(__name__)
class LSAnnotationSyncService:
"""将 Label Studio 项目中的标注结果同步回 DM 数据库。
约定:
- Label Studio task.data 中包含由 SyncService 写入的 DM 元信息:
{"file_id": "...", "dataset_id": "...", "file_path": "...", "original_name": "..."}
- 对于每个 task,我们会:
* 读取该 task 的 annotations 列表;
* 结合 task.predictions 一并写入 t_dm_dataset_files.annotation;
* 从 annotations/predictions 的 result 字段中提取标签,写入 t_dm_dataset_files.tags;
* 更新时间戳 tags_updated_at。
"""
def __init__(self, db: AsyncSession, ls_client: LabelStudioClient) -> None:
self.db = db
self.ls_client = ls_client
async def sync_project_annotations_to_dm(self, project_id: str) -> int:
"""从指定 Label Studio 项目中同步所有任务的标注到 DM 文件表。
Args:
project_id: Label Studio 项目 ID(字符串形式)。
Returns:
成功更新的文件数量。
"""
try:
page_size = getattr(settings, "ls_task_page_size", 1000)
result = await self.ls_client.get_project_tasks(
project_id=project_id,
page=None,
page_size=page_size,
)
except Exception as e: # pragma: no cover - 防御性日志
logger.error("Failed to fetch tasks from Label Studio project %s: %s", project_id, e)
return 0
if not result:
logger.warning("No tasks returned for project %s when syncing annotations", project_id)
return 0
tasks: List[Dict[str, Any]] = result.get("tasks", []) # type: ignore[assignment]
if not tasks:
logger.info("Project %s has no tasks to sync", project_id)
return 0
updated_files = 0
for task in tasks:
data: Dict[str, Any] = task.get("data") or {}
raw_file_id = data.get("file_id")
if not raw_file_id:
# 旧任务可能没有 DM file_id,跳过但记录日志
logger.debug(
"Skip LS task %s because data.file_id is missing",
task.get("id"),
)
continue
file_id = str(raw_file_id)
# 获取该任务的 annotations 列表
annotations: Optional[List[Dict[str, Any]]] = None
try:
annotations = await self.ls_client.get_task_annotations(int(task["id"])) # type: ignore[arg-type]
except Exception as e: # pragma: no cover
logger.error(
"Failed to fetch annotations for LS task %s (file_id=%s): %s",
task.get("id"),
file_id,
e,
)
predictions: List[Dict[str, Any]] = task.get("predictions") or []
annotation_payload: Dict[str, Any] = {
"task": {
"id": task.get("id"),
"project": task.get("project"),
"data": data,
},
"annotations": annotations or [],
"predictions": predictions,
}
tags = self._extract_tags_from_results(
annotations or [],
predictions,
)
success = await self._update_dataset_file(file_id, tags, annotation_payload)
if success:
updated_files += 1
logger.info(
"Synced annotations from LS project %s to DM, updated %d files",
project_id,
updated_files,
)
return updated_files
async def _update_dataset_file(
self,
file_id: str,
tags: List[Dict[str, Any]],
annotation: Dict[str, Any],
) -> bool:
"""将提取出的 tags 与 annotation 写回 t_dm_dataset_files。
直接更新 DatasetFiles.tags / tags_updated_at / annotation 字段,
不走模板驱动的转换逻辑(模板转换在用户主动批量更新时使用)。
"""
try:
result = await self.db.execute(
select(DatasetFiles).where(DatasetFiles.id == file_id)
)
record = result.scalar_one_or_none()
if not record:
logger.warning("Dataset file not found when syncing LS annotations: %s", file_id)
return False
update_time = datetime.utcnow()
record.tags = tags # type: ignore[assignment]
record.tags_updated_at = update_time # type: ignore[assignment]
# annotation 可能较大,但为 JSONB,直接整体覆盖
setattr(record, "annotation", annotation)
await self.db.commit()
await self.db.refresh(record)
logger.debug(
"Updated DM file %s with %d tags from LS annotations",
file_id,
len(tags),
)
return True
except Exception as e: # pragma: no cover
logger.error("Failed to update DM file %s from LS annotations: %s", file_id, e)
try:
await self.db.rollback()
except Exception:
pass
return False
@staticmethod
def _extract_tags_from_results(
annotations: List[Dict[str, Any]],
predictions: List[Dict[str, Any]],
) -> List[Dict[str, Any]]:
"""从 annotations 和 predictions 的 result 字段中提取通用标签结构。
输出格式与 DM 前端当前解析逻辑兼容:
- 每个元素包含: id/type/from_name/values;
- 其中 values 是一个字典,键通常与 type 一致,例如 "rectanglelabels";
- 前端通过 tag.values[tag.type] 提取字符串标签集合。
"""
def iter_results() -> List[Dict[str, Any]]:
res: List[Dict[str, Any]] = []
for ann in annotations or []:
for item in ann.get("result", []) or []:
if isinstance(item, dict):
res.append(item)
for pred in predictions or []:
for item in pred.get("result", []) or []:
if isinstance(item, dict):
res.append(item)
return res
results = iter_results()
if not results:
return []
normalized: List[Dict[str, Any]] = []
for r in results:
r_type = r.get("type")
if isinstance(r_type, str):
r_type = r_type.strip().lower()
from_name = r.get("from_name") or r.get("fromName")
to_name = r.get("to_name") or r.get("toName")
value_obj = r.get("value") or {}
if not isinstance(value_obj, dict):
continue
# 将 Label Studio 的 value 映射为 values,方便前端统一解析
values: Dict[str, Any] = {}
for key, v in value_obj.items():
normalized_key = str(key).strip().lower() if key is not None else ""
if normalized_key:
values[normalized_key] = v
if isinstance(r_type, str) and r_type:
if r_type in values:
values = {r_type: values[r_type]}
elif len(values) == 1:
only_value = next(iter(values.values()))
values = {r_type: only_value}
tag = {
"id": r.get("id"),
"type": r_type,
"from_name": from_name,
"to_name": to_name,
"values": values,
}
normalized.append(tag)
return normalized
@@ -0,0 +1,484 @@
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select
from sqlalchemy import update, func
from sqlalchemy.orm import aliased
from typing import Optional, List, Tuple
from datetime import datetime
import uuid
from app.core.logging import get_logger
from app.db.models import LabelingProject, AnnotationTemplate
from app.db.models.dataset_management import Dataset
from app.module.annotation.schema import (
DatasetMappingCreateRequest,
DatasetMappingUpdateRequest,
DatasetMappingResponse,
AnnotationTemplateResponse
)
logger = get_logger(__name__)
class DatasetMappingService:
"""数据集映射服务"""
def __init__(self, db: AsyncSession):
self.db = db
def _build_query_with_dataset_name(self):
"""Build base query with dataset name joined"""
return select(
LabelingProject,
Dataset.name.label('dataset_name')
).outerjoin(
Dataset,
LabelingProject.dataset_id == Dataset.id
)
async def _to_response_from_row(
self,
row,
include_template: bool = False
) -> DatasetMappingResponse:
"""
Convert query row (mapping + dataset_name) to response
Args:
row: Query result row containing (LabelingProject, dataset_name)
include_template: If True, fetch and include full template details
"""
mapping = row[0] # LabelingProject object
dataset_name = row[1] # dataset_name from join
# Get template_id from mapping
template_id = getattr(mapping, 'template_id', None)
# Optionally fetch full template details
template_response = None
if include_template and template_id:
from ..service.template import AnnotationTemplateService
template_service = AnnotationTemplateService()
template_response = await template_service.get_template(self.db, template_id)
logger.debug(f"Included template details for template_id: {template_id}")
response_data = {
"id": mapping.id,
"dataset_id": mapping.dataset_id,
"dataset_name": dataset_name,
"labeling_project_id": mapping.labeling_project_id,
"name": mapping.name,
"description": getattr(mapping, 'description', None),
"template_id": template_id,
"template": template_response,
"created_at": mapping.created_at,
"updated_at": mapping.updated_at,
"deleted_at": mapping.deleted_at,
}
return DatasetMappingResponse(**response_data)
async def _to_response(
self,
mapping: LabelingProject,
include_template: bool = False
) -> DatasetMappingResponse:
"""
Convert ORM model to response with dataset name (for single entity operations)
Args:
mapping: LabelingProject ORM object
include_template: If True, fetch and include full template details
"""
# Fetch dataset name
dataset_name = None
dataset_id = getattr(mapping, 'dataset_id', None)
if dataset_id:
dataset_result = await self.db.execute(
select(Dataset.name).where(Dataset.id == dataset_id)
)
dataset_name = dataset_result.scalar_one_or_none()
# Get template_id from mapping
template_id = getattr(mapping, 'template_id', None)
# Optionally fetch full template details
template_response = None
if include_template and template_id:
from ..service.template import AnnotationTemplateService
template_service = AnnotationTemplateService()
template_response = await template_service.get_template(self.db, template_id)
logger.debug(f"Included template details for template_id: {template_id}")
# Create response dict with all fields
response_data = {
"id": mapping.id,
"dataset_id": dataset_id,
"dataset_name": dataset_name,
"labeling_project_id": mapping.labeling_project_id,
"name": mapping.name,
"description": getattr(mapping, 'description', None),
"template_id": template_id,
"template": template_response,
"created_at": mapping.created_at,
"updated_at": mapping.updated_at,
"deleted_at": mapping.deleted_at,
}
return DatasetMappingResponse(**response_data)
async def create_mapping(
self,
labeling_project: LabelingProject
) -> DatasetMappingResponse:
"""创建数据集映射"""
logger.debug(f"Create dataset mapping: {labeling_project.dataset_id} -> {labeling_project.labeling_project_id}")
# Use the passed object directly
self.db.add(labeling_project)
await self.db.commit()
await self.db.refresh(labeling_project)
logger.debug(f"Mapping created: {labeling_project.id}")
return await self._to_response(labeling_project)
async def get_mapping_by_source_uuid(
self,
dataset_id: str
) -> Optional[DatasetMappingResponse]:
"""根据源数据集ID获取映射(返回第一个未删除的)"""
logger.debug(f"Get mapping by source dataset id: {dataset_id}")
result = await self.db.execute(
select(LabelingProject).where(
LabelingProject.dataset_id == dataset_id,
LabelingProject.deleted_at.is_(None)
)
)
mapping = result.scalar_one_or_none()
if mapping:
logger.debug(f"Found mapping: {mapping.id}")
return await self._to_response(mapping)
logger.debug(f"No mapping found for source dataset id: {dataset_id}")
return None
async def get_mappings_by_dataset_id(
self,
dataset_id: str,
include_deleted: bool = False
) -> List[DatasetMappingResponse]:
"""根据源数据集ID获取所有映射关系"""
logger.debug(f"Get all mappings by source dataset id: {dataset_id}")
query = self._build_query_with_dataset_name().where(
LabelingProject.dataset_id == dataset_id
)
if not include_deleted:
query = query.where(LabelingProject.deleted_at.is_(None))
result = await self.db.execute(
query.order_by(LabelingProject.created_at.desc())
)
rows = result.all()
logger.debug(f"Found {len(rows)} mappings")
# Convert rows to responses (async comprehension)
responses = []
for row in rows:
response = await self._to_response_from_row(row, include_template=False)
responses.append(response)
return responses
async def get_mapping_by_labeling_project_id(
self,
labeling_project_id: str
) -> Optional[DatasetMappingResponse]:
"""根据Label Studio项目ID获取映射"""
logger.debug(f"Get mapping by Label Studio project id: {labeling_project_id}")
result = await self.db.execute(
select(LabelingProject).where(
LabelingProject.labeling_project_id == labeling_project_id,
LabelingProject.deleted_at.is_(None)
)
)
mapping = result.scalar_one_or_none()
if mapping:
logger.debug(f"Found mapping: {mapping.id}")
return await self._to_response(mapping)
logger.debug(f"No mapping found for Label Studio project id: {labeling_project_id}")
return None
async def get_mapping_by_uuid(
self,
mapping_id: str,
include_template: bool = False
) -> Optional[DatasetMappingResponse]:
"""
根据映射UUID获取映射
Args:
mapping_id: 映射UUID
include_template: 是否包含完整的模板信息
"""
logger.debug(f"Get mapping: {mapping_id}, include_template={include_template}")
result = await self.db.execute(
select(LabelingProject).where(
LabelingProject.id == mapping_id,
LabelingProject.deleted_at.is_(None)
)
)
mapping = result.scalar_one_or_none()
if mapping:
logger.debug(f"Found mapping: {mapping.id}")
return await self._to_response(mapping, include_template=include_template)
logger.debug(f"No mapping found for mapping id: {mapping_id}")
return None
async def update_mapping(
self,
mapping_id: str,
update_data: DatasetMappingUpdateRequest
) -> Optional[DatasetMappingResponse]:
"""更新映射信息"""
logger.info(f"Update mapping: {mapping_id}")
mapping = await self.get_mapping_by_uuid(mapping_id)
if not mapping:
return None
update_values = update_data.model_dump(exclude_unset=True)
update_values["last_updated_at"] = datetime.now()
result = await self.db.execute(
update(LabelingProject)
.where(LabelingProject.id == mapping_id)
.values(**update_values)
)
await self.db.commit()
if result.rowcount and result.rowcount > 0: # type: ignore
return await self.get_mapping_by_uuid(mapping_id)
return None
async def soft_delete_mapping(self, mapping_id: str) -> bool:
"""软删除映射"""
logger.info(f"Soft delete mapping: {mapping_id}")
result = await self.db.execute(
update(LabelingProject)
.where(
LabelingProject.id == mapping_id,
LabelingProject.deleted_at.is_(None)
)
.values(deleted_at=datetime.now())
)
await self.db.commit()
success = result.rowcount and result.rowcount > 0 # type: ignore
if success:
logger.info(f"Mapping soft-deleted: {mapping_id}")
else:
logger.warning(f"Mapping not exists or already deleted: {mapping_id}")
return success
async def get_all_mappings(
self,
skip: int = 0,
limit: int = 100
) -> List[DatasetMappingResponse]:
"""获取所有有效映射"""
logger.debug(f"List all mappings, skip: {skip}, limit: {limit}")
query = self._build_query_with_dataset_name().where(
LabelingProject.deleted_at.is_(None)
)
result = await self.db.execute(
query
.offset(skip)
.limit(limit)
.order_by(LabelingProject.created_at.desc())
)
rows = result.all()
logger.debug(f"Found {len(rows)} mappings")
# Convert rows to responses (async comprehension)
responses = []
for row in rows:
response = await self._to_response_from_row(row, include_template=False)
responses.append(response)
return responses
async def count_mappings(self, include_deleted: bool = False) -> int:
"""统计映射总数"""
query = select(func.count()).select_from(LabelingProject)
if not include_deleted:
query = query.where(LabelingProject.deleted_at.is_(None))
result = await self.db.execute(query)
return result.scalar_one()
async def get_all_mappings_with_count(
self,
skip: int = 0,
limit: int = 100,
include_deleted: bool = False,
include_template: bool = False
) -> Tuple[List[DatasetMappingResponse], int]:
"""
获取所有映射及总数(用于分页)
Args:
skip: 跳过记录数
limit: 返回记录数
include_deleted: 是否包含已删除的记录
include_template: 是否包含完整的模板信息
"""
logger.debug(f"List all mappings with count, skip: {skip}, limit: {limit}, include_template={include_template}")
# 构建查询
query = self._build_query_with_dataset_name()
if not include_deleted:
query = query.where(LabelingProject.deleted_at.is_(None))
# 获取总数
count_query = select(func.count()).select_from(LabelingProject)
if not include_deleted:
count_query = count_query.where(LabelingProject.deleted_at.is_(None))
count_result = await self.db.execute(count_query)
total = count_result.scalar_one()
# 获取数据
result = await self.db.execute(
query
.offset(skip)
.limit(limit)
.order_by(LabelingProject.created_at.desc())
)
rows = result.all()
# 过滤掉由自动标注任务自动创建的隐藏工程(configuration.autoTaskId 存在)
filtered_rows = []
for row in rows:
mapping: LabelingProject = row[0]
cfg = getattr(mapping, "configuration", None) or {}
if isinstance(cfg, dict) and cfg.get("autoTaskId"):
continue
filtered_rows.append(row)
logger.debug(
f"Found {len(filtered_rows)} visible mappings (filtered from {len(rows)}), total: {total}"
)
# Convert rows to responses (async comprehension)
responses = []
for row in filtered_rows:
response = await self._to_response_from_row(row, include_template=include_template)
responses.append(response)
# total 仍然返回数据库总数,目前前端只关注当前页内容数量,故不调整 total
return responses, total
async def get_template_id_by_dataset_id(self, dataset_id: str) -> Optional[str]:
"""
Get template ID for a dataset by finding its labeling project
Args:
dataset_id: Dataset UUID
Returns:
Template ID or None if no labeling project found or no template associated
"""
logger.debug(f"Looking up template for dataset: {dataset_id}")
result = await self.db.execute(
select(LabelingProject.template_id)
.where(
LabelingProject.dataset_id == dataset_id,
LabelingProject.deleted_at.is_(None)
)
.limit(1)
)
template_id = result.scalar_one_or_none()
if template_id:
logger.debug(f"Found template {template_id} for dataset {dataset_id}")
else:
logger.warning(f"No template found for dataset {dataset_id}")
return template_id
async def get_mappings_by_source_with_count(
self,
dataset_id: str,
skip: int = 0,
limit: int = 100,
include_deleted: bool = False,
include_template: bool = False
) -> Tuple[List[DatasetMappingResponse], int]:
"""
根据源数据集ID获取映射关系及总数(用于分页)
Args:
dataset_id: 数据集ID
skip: 跳过记录数
limit: 返回记录数
include_deleted: 是否包含已删除的记录
include_template: 是否包含完整的模板信息
"""
logger.debug(f"Get mappings by source dataset id with count: {dataset_id}, include_template={include_template}")
# 构建查询
query = self._build_query_with_dataset_name().where(
LabelingProject.dataset_id == dataset_id
)
if not include_deleted:
query = query.where(LabelingProject.deleted_at.is_(None))
# 获取总数
count_query = select(func.count()).select_from(LabelingProject).where(
LabelingProject.dataset_id == dataset_id
)
if not include_deleted:
count_query = count_query.where(LabelingProject.deleted_at.is_(None))
count_result = await self.db.execute(count_query)
total = count_result.scalar_one()
# 获取数据
result = await self.db.execute(
query
.offset(skip)
.limit(limit)
.order_by(LabelingProject.created_at.desc())
)
rows = result.all()
# 过滤掉由自动标注任务自动创建的隐藏工程(configuration.autoTaskId 存在)
filtered_rows = []
for row in rows:
mapping: LabelingProject = row[0]
cfg = getattr(mapping, "configuration", None) or {}
if isinstance(cfg, dict) and cfg.get("autoTaskId"):
continue
filtered_rows.append(row)
logger.debug(
f"Found {len(filtered_rows)} mappings for dataset {dataset_id} (filtered from {len(rows)}), total: {total}"
)
# Convert rows to responses (async comprehension)
responses = []
for row in filtered_rows:
response = await self._to_response_from_row(row, include_template=include_template)
responses.append(response)
return responses, total
@@ -0,0 +1,208 @@
"""服务:将自动标注(YOLO)结果同步到 Label Studio 作为 predictions。
该服务不负责执行模型推理,只负责:
- 读取本地 annotations JSON(由 runtime/ops/annotation/image_object_detection_bounding_box 产出);
- 根据 Label Studio 项目的任务列表,将每个文件对应到具体 task;
- 将检测结果转换为 Label Studio 的 prediction 结构并写入。
"""
from __future__ import annotations
import json
import os
import uuid
from typing import Dict, Any, List, Optional
from app.core.logging import get_logger
from ..client import LabelStudioClient
logger = get_logger(__name__)
class PredictionSyncService:
"""将 YOLO 检测结果推送到 Label Studio 的预测服务。"""
def __init__(self, ls_client: LabelStudioClient) -> None:
self.ls_client = ls_client
async def _build_task_index(self, project_id: str) -> Dict[str, int]:
"""为指定项目构建 "文件名 -> task_id" 索引。
优先使用 data.original_name,其次退回到从 data.image / data.file_path
中解析出的文件名,保证与 YOLO JSON 中的 annotations["image"] 对得上。
"""
result = await self.ls_client.get_project_tasks(project_id=project_id, page=None)
if not result:
logger.warning("No tasks found for project %s", project_id)
return {}
tasks = result.get("tasks", []) or []
index: Dict[str, int] = {}
for task in tasks:
task_id = task.get("id")
data = task.get("data", {}) or {}
if task_id is None:
continue
# candidate names that may appear in YOLO annotations["image"]
candidates: List[str] = []
original_name = data.get("original_name")
if isinstance(original_name, str) and original_name:
candidates.append(original_name)
image_val = data.get("image") or data.get("file_path")
if isinstance(image_val, str) and image_val:
basename = os.path.basename(image_val)
if basename:
candidates.append(basename)
for name in candidates:
# 后写覆盖前写没关系,同一文件名对应同一个 task 的预期
index[name] = task_id
logger.debug(
"Built task index for project %s with %d entries", project_id, len(index)
)
return index
@staticmethod
def _build_ls_result_from_detections(
annotations: Dict[str, Any],
from_name: str = "label",
to_name: str = "image",
) -> List[Dict[str, Any]]:
"""将 YOLO annotations["detections"] 转为 Label Studio result 列表。"""
width = float(annotations.get("width") or 0) or 1.0
height = float(annotations.get("height") or 0) or 1.0
detections = annotations.get("detections") or []
results: List[Dict[str, Any]] = []
for det in detections:
try:
label = det.get("label")
bbox = det.get("bbox_xywh") or det.get("bbox_xyxy")
if not label or not bbox or len(bbox) < 4:
continue
# 优先使用 xywh,单位像素
if len(bbox) == 4:
x, y, w, h = map(float, bbox)
else:
# 回退:如果给的是 xyxy
x1, y1, x2, y2 = map(float, bbox[:4])
x, y, w, h = x1, y1, x2 - x1, y2 - y1
x_pct = max(0.0, min(100.0, x / width * 100.0))
y_pct = max(0.0, min(100.0, y / height * 100.0))
w_pct = max(0.0, min(100.0, w / width * 100.0))
h_pct = max(0.0, min(100.0, h / height * 100.0))
region_id = uuid.uuid4().hex[:10]
results.append(
{
"id": region_id,
"type": "rectanglelabels",
"value": {
"x": x_pct,
"y": y_pct,
"width": w_pct,
"height": h_pct,
"rotation": 0,
"rectanglelabels": [label],
},
"origin": "prediction",
"to_name": to_name,
"from_name": from_name,
"image_rotation": 0,
"original_width": width,
"original_height": height,
}
)
except Exception as exc: # 防御性:单条失败不影响其他
logger.warning("Failed to convert detection to LS result: %s", exc)
continue
return results
async def sync_predictions_from_dir(
self,
project_id: str,
annotations_dir: str,
*,
model_version: Optional[str] = None,
from_name: str = "label",
to_name: str = "image",
) -> int:
"""将某个目录下的 YOLO JSON 同步为给定项目的 predictions。
返回成功写入 prediction 的任务数。
"""
if not os.path.isdir(annotations_dir):
logger.warning("annotations_dir '%s' is not a directory", annotations_dir)
return 0
task_index = await self._build_task_index(project_id)
if not task_index:
logger.warning(
"No tasks found when syncing predictions for project %s", project_id
)
return 0
created_count = 0
for filename in os.listdir(annotations_dir):
if not filename.lower().endswith(".json"):
continue
json_path = os.path.join(annotations_dir, filename)
try:
with open(json_path, "r", encoding="utf-8") as f:
annotations = json.load(f)
except Exception as exc:
logger.warning("Failed to load annotations json %s: %s", json_path, exc)
continue
image_name = annotations.get("image") or filename.rsplit(".", 1)[0]
task_id = task_index.get(image_name)
if task_id is None:
logger.info(
"No matching task for image '%s' in project %s, skip", image_name, project_id
)
continue
result = self._build_ls_result_from_detections(
annotations, from_name=from_name, to_name=to_name
)
if not result:
logger.info("No detections for %s, skip creating prediction", image_name)
continue
mv = model_version
if mv is None:
model_size = annotations.get("model_size") or ""
mv = f"yolov8{model_size}" if model_size else "yolov8"
prediction = await self.ls_client.create_prediction(
task_id=task_id,
result=result,
model_version=mv,
)
if prediction:
created_count += 1
logger.info(
"Synced YOLO predictions from '%s' to project %s, created %d predictions",
annotations_dir,
project_id,
created_count,
)
return created_count
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,341 @@
"""
Annotation Template Service
"""
from typing import Optional, List
from datetime import datetime
from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession
from uuid import uuid4
from fastapi import HTTPException
from app.db.models.annotation_management import AnnotationTemplate
from app.module.annotation.schema.template import (
CreateAnnotationTemplateRequest,
UpdateAnnotationTemplateRequest,
AnnotationTemplateResponse,
AnnotationTemplateListResponse,
TemplateConfiguration
)
from app.module.annotation.utils.config_validator import LabelStudioConfigValidator
from app.module.annotation.config import LabelStudioTagConfig
class AnnotationTemplateService:
"""标注模板服务"""
@staticmethod
def generate_label_studio_config(config: TemplateConfiguration) -> str:
"""
从配置JSON生成Label Studio XML配置
Args:
config: 模板配置对象
Returns:
Label Studio XML字符串
"""
tag_config = LabelStudioTagConfig()
xml_parts = ['<View>']
# 生成对象定义
for obj in config.objects:
obj_attrs = [
f'name="{obj.name}"',
f'value="{obj.value}"'
]
xml_parts.append(f' <{obj.type} {" ".join(obj_attrs)}/>')
# 生成标签定义
for label in config.labels:
label_attrs = [
f'name="{label.from_name}"',
f'toName="{label.to_name}"'
]
# 添加可选属性
if label.required:
label_attrs.append('required="true"')
if label.choice:
label_attrs.append(f'choice="{label.choice}"')
# xuhuanqing
#tag_type = label.type.capitalize() if label.type else "Choices"
tag_type = label.type if label.type else "Choices"
# 检查是否需要子元素
if label.options or label.labels:
choices = label.options or label.labels or []
xml_parts.append(f' <{tag_type} {" ".join(label_attrs)}>')
# 从配置获取子元素标签名
child_tag = tag_config.get_child_tag(tag_type)
if not child_tag:
# 默认使用 Label
child_tag = "Label"
for choice in choices:
xml_parts.append(f' <{child_tag} value="{choice}"/>')
xml_parts.append(f' </{tag_type}>')
else:
# 处理简单标签类型(不需要子元素)
xml_parts.append(f' <{tag_type} {" ".join(label_attrs)}/>')
xml_parts.append('</View>')
return '\n'.join(xml_parts)
async def create_template(
self,
db: AsyncSession,
request: CreateAnnotationTemplateRequest
) -> AnnotationTemplateResponse:
"""
创建标注模板
Args:
db: 数据库会话
request: 创建请求
Returns:
创建的模板响应
"""
# 验证配置JSON
config_dict = request.configuration.model_dump(mode='json', by_alias=False)
valid, error = LabelStudioConfigValidator.validate_configuration_json(config_dict)
if not valid:
raise HTTPException(status_code=400, detail=f"Invalid configuration: {error}")
# 生成Label Studio XML配置(用于验证,但不存储)
label_config = self.generate_label_studio_config(request.configuration)
# 验证生成的XML
valid, error = LabelStudioConfigValidator.validate_xml(label_config)
if not valid:
raise HTTPException(status_code=400, detail=f"Generated XML is invalid: {error}")
# 创建模板对象(不包含label_config字段)
template = AnnotationTemplate(
id=str(uuid4()),
name=request.name,
description=request.description,
data_type=request.data_type,
labeling_type=request.labeling_type,
configuration=config_dict,
style=request.style,
category=request.category,
built_in=False,
version="1.0.0",
created_at=datetime.now()
)
db.add(template)
await db.commit()
await db.refresh(template)
return self._to_response(template)
async def get_template(
self,
db: AsyncSession,
template_id: str
) -> Optional[AnnotationTemplateResponse]:
"""
获取单个模板
Args:
db: 数据库会话
template_id: 模板ID
Returns:
模板响应或None
"""
result = await db.execute(
select(AnnotationTemplate)
.where(
AnnotationTemplate.id == template_id,
AnnotationTemplate.deleted_at.is_(None)
)
)
template = result.scalar_one_or_none()
if template:
return self._to_response(template)
return None
async def list_templates(
self,
db: AsyncSession,
page: int = 1,
size: int = 10,
category: Optional[str] = None,
data_type: Optional[str] = None,
labeling_type: Optional[str] = None,
built_in: Optional[bool] = None
) -> AnnotationTemplateListResponse:
"""
获取模板列表
Args:
db: 数据库会话
page: 页码
size: 每页大小
category: 分类筛选
data_type: 数据类型筛选
labeling_type: 标注类型筛选
built_in: 是否内置模板筛选
Returns:
模板列表响应
"""
# 构建查询条件
conditions: List = [AnnotationTemplate.deleted_at.is_(None)]
if category:
conditions.append(AnnotationTemplate.category == category) # type: ignore
if data_type:
conditions.append(AnnotationTemplate.data_type == data_type) # type: ignore
if labeling_type:
conditions.append(AnnotationTemplate.labeling_type == labeling_type) # type: ignore
if built_in is not None:
conditions.append(AnnotationTemplate.built_in == built_in) # type: ignore
# 查询总数
count_result = await db.execute(
select(func.count()).select_from(AnnotationTemplate).where(*conditions)
)
total = count_result.scalar() or 0
# 分页查询
result = await db.execute(
select(AnnotationTemplate)
.where(*conditions)
.order_by(AnnotationTemplate.created_at.desc())
.limit(size)
.offset((page - 1) * size)
)
templates = result.scalars().all()
return AnnotationTemplateListResponse(
content=[self._to_response(t) for t in templates],
totalElements=total,
page=page,
size=size,
totalPages=(total + size - 1) // size
)
async def update_template(
self,
db: AsyncSession,
template_id: str,
request: UpdateAnnotationTemplateRequest
) -> Optional[AnnotationTemplateResponse]:
"""
更新模板
Args:
db: 数据库会话
template_id: 模板ID
request: 更新请求
Returns:
更新后的模板响应或None
"""
result = await db.execute(
select(AnnotationTemplate)
.where(
AnnotationTemplate.id == template_id,
AnnotationTemplate.deleted_at.is_(None)
)
)
template = result.scalar_one_or_none()
if not template:
return None
# 更新字段
update_data = request.model_dump(exclude_unset=True, by_alias=False)
for field, value in update_data.items():
if field == 'configuration' and value is not None:
# 验证配置JSON
config_dict = value.model_dump(mode='json', by_alias=False)
valid, error = LabelStudioConfigValidator.validate_configuration_json(config_dict)
if not valid:
raise HTTPException(status_code=400, detail=f"Invalid configuration: {error}")
# 重新生成Label Studio XML配置(用于验证)
label_config = self.generate_label_studio_config(value)
# 验证生成的XML
valid, error = LabelStudioConfigValidator.validate_xml(label_config)
if not valid:
raise HTTPException(status_code=400, detail=f"Generated XML is invalid: {error}")
# 只更新configuration字段,不存储label_config
setattr(template, field, config_dict)
else:
setattr(template, field, value)
template.updated_at = datetime.now() # type: ignore
await db.commit()
await db.refresh(template)
return self._to_response(template)
async def delete_template(
self,
db: AsyncSession,
template_id: str
) -> bool:
"""
删除模板(软删除)
Args:
db: 数据库会话
template_id: 模板ID
Returns:
是否删除成功
"""
result = await db.execute(
select(AnnotationTemplate)
.where(
AnnotationTemplate.id == template_id,
AnnotationTemplate.deleted_at.is_(None)
)
)
template = result.scalar_one_or_none()
if not template:
return False
template.deleted_at = datetime.now() # type: ignore
await db.commit()
return True
def _to_response(self, template: AnnotationTemplate) -> AnnotationTemplateResponse:
"""
转换为响应对象
Args:
template: 数据库模型对象
Returns:
模板响应对象
"""
# 将配置JSON转换为TemplateConfiguration对象
from typing import cast, Dict, Any
config_dict = cast(Dict[str, Any], template.configuration)
config = TemplateConfiguration(**config_dict)
# 动态生成Label Studio XML配置
label_config = self.generate_label_studio_config(config)
# 使用model_validate从ORM对象创建响应对象
response = AnnotationTemplateResponse.model_validate(template)
response.configuration = config
response.label_config = label_config # type: ignore
return response
@@ -0,0 +1,11 @@
"""
Annotation Module Utilities
"""
from .config_validator import LabelStudioConfigValidator
from .tag_converter import TagFormatConverter, create_converter_from_template_config
__all__ = [
'LabelStudioConfigValidator',
'TagFormatConverter',
'create_converter_from_template_config'
]
@@ -0,0 +1,273 @@
"""
Label Studio Configuration Validation Utilities
"""
from typing import Dict, List, Tuple, Optional
import xml.etree.ElementTree as ET
from app.module.annotation.config import LabelStudioTagConfig
class LabelStudioConfigValidator:
"""验证Label Studio配置的工具类"""
@staticmethod
def _get_config() -> LabelStudioTagConfig:
"""获取标签配置实例"""
return LabelStudioTagConfig()
@staticmethod
def validate_xml(xml_string: str) -> Tuple[bool, Optional[str]]:
"""
验证XML格式是否正确
Args:
xml_string: Label Studio XML配置字符串
Returns:
(是否有效, 错误信息)
"""
try:
config = LabelStudioConfigValidator._get_config()
root = ET.fromstring(xml_string)
# 检查根元素
if root.tag != 'View':
return False, "Root element must be <View>"
# 检查是否有对象定义
object_types = config.get_object_types()
objects = [child for child in root if child.tag in object_types]
if not objects:
return False, "No data objects (Image, Text, etc.) found"
# 检查是否有控件定义
control_types = config.get_control_types()
controls = [child for child in root if child.tag in control_types]
if not controls:
return False, "No annotation controls found"
# 验证每个控件
for control in controls:
valid, error = LabelStudioConfigValidator._validate_control(control)
if not valid:
return False, f"Control {control.tag}: {error}"
return True, None
except ET.ParseError as e:
return False, f"XML parse error: {str(e)}"
except Exception as e:
return False, f"Validation error: {str(e)}"
@staticmethod
def _validate_control(control: ET.Element) -> Tuple[bool, Optional[str]]:
"""
验证单个控件元素
Args:
control: 控件XML元素
Returns:
(是否有效, 错误信息)
"""
config = LabelStudioConfigValidator._get_config()
# 检查必需属性
if 'name' not in control.attrib:
return False, "Missing 'name' attribute"
if 'toName' not in control.attrib:
return False, "Missing 'toName' attribute"
# 检查控件是否需要子元素
if config.requires_children(control.tag):
child_tag = config.get_child_tag(control.tag)
if not child_tag:
return False, f"Configuration error: no child_tag defined for {control.tag}"
children = control.findall(child_tag)
if not children:
return False, f"{control.tag} must have at least one <{child_tag}> child"
# 检查每个子元素是否有value
for child in children:
if 'value' not in child.attrib:
return False, f"{child_tag} missing 'value' attribute"
return True, None
@staticmethod
def extract_label_values(xml_string: str) -> Dict[str, List[str]]:
"""
从XML中提取所有标签值
Args:
xml_string: Label Studio XML配置字符串
Returns:
字典,键为控件名称,值为标签值列表
"""
result = {}
config = LabelStudioConfigValidator._get_config()
try:
root = ET.fromstring(xml_string)
control_types = config.get_control_types()
controls = [child for child in root if child.tag in control_types]
for control in controls:
if not config.requires_children(control.tag):
continue
control_name = control.get('name', 'unknown')
child_tag = config.get_child_tag(control.tag)
if child_tag:
children = control.findall(child_tag)
label_values = [child.get('value', '') for child in children]
result[control_name] = label_values
except Exception:
pass
return result
@staticmethod
def validate_configuration_json(config: Dict) -> Tuple[bool, Optional[str]]:
"""
验证配置JSON结构
Args:
config: 配置字典
Returns:
(是否有效, 错误信息)
"""
# 检查必需字段
if 'labels' not in config:
return False, "Missing 'labels' field"
if 'objects' not in config:
return False, "Missing 'objects' field"
if not isinstance(config['labels'], list):
return False, "'labels' must be an array"
if not isinstance(config['objects'], list):
return False, "'objects' must be an array"
if not config['labels']:
return False, "'labels' array cannot be empty"
if not config['objects']:
return False, "'objects' array cannot be empty"
# 验证每个标签定义
for idx, label in enumerate(config['labels']):
valid, error = LabelStudioConfigValidator._validate_label_definition(label)
if not valid:
return False, f"Label {idx}: {error}"
# 验证每个对象定义
for idx, obj in enumerate(config['objects']):
valid, error = LabelStudioConfigValidator._validate_object_definition(obj)
if not valid:
return False, f"Object {idx}: {error}"
# 验证toName引用
object_names = {obj['name'] for obj in config['objects']}
for label in config['labels']:
to_name = label.get('toName') or label.get('to_name')
from_name = label.get('fromName') or label.get('from_name')
if to_name not in object_names:
return False, f"Label '{from_name}' references unknown object '{to_name}'"
return True, None
@staticmethod
def _validate_label_definition(label: Dict) -> Tuple[bool, Optional[str]]:
"""验证标签定义"""
config = LabelStudioConfigValidator._get_config()
control_types = config.get_control_types()
# Support both camelCase and snake_case
from_name = label.get('fromName') or label.get('from_name')
to_name = label.get('toName') or label.get('to_name')
label_type = label.get('type')
if not from_name:
return False, "Missing required field 'fromName'"
if not to_name:
return False, "Missing required field 'toName'"
if not label_type:
return False, "Missing required field 'type'"
# 检查类型是否支持
if label_type not in control_types:
return False, f"Unsupported control type '{label_type}'"
# 检查是否需要子元素(options 或 labels)
if config.requires_children(label_type):
if 'options' not in label and 'labels' not in label:
return False, f"{label_type} must have 'options' or 'labels' field"
return True, None
@staticmethod
def _validate_object_definition(obj: Dict) -> Tuple[bool, Optional[str]]:
"""验证对象定义"""
config = LabelStudioConfigValidator._get_config()
object_types = config.get_object_types()
required_fields = ['name', 'type', 'value']
for field in required_fields:
if field not in obj:
return False, f"Missing required field '{field}'"
# 检查类型是否支持
if obj['type'] not in object_types:
return False, f"Unsupported object type '{obj['type']}'"
# 检查value格式
if not obj['value'].startswith('$'):
return False, "Object value must start with '$' (e.g., '$image')"
return True, None
# 使用示例
if __name__ == "__main__":
# 验证XML
xml = """<View>
<Image name="image" value="$image"/>
<Choices name="choice" toName="image" required="true">
<Label value="Cat"/>
<Label value="Dog"/>
</Choices>
</View>"""
valid, error = LabelStudioConfigValidator.validate_xml(xml)
print(f"XML Valid: {valid}, Error: {error}")
# 验证配置JSON
config = {
"labels": [
{
"fromName": "choice",
"toName": "image",
"type": "Choices",
"options": ["Cat", "Dog"],
"required": True
}
],
"objects": [
{
"name": "image",
"type": "Image",
"value": "$image"
}
]
}
valid, error = LabelStudioConfigValidator.validate_configuration_json(config)
print(f"Config Valid: {valid}, Error: {error}")
@@ -0,0 +1,231 @@
"""
Tag Format Converter
Converts simplified external tag format to internal storage format by looking up
the type from the annotation template configuration.
External format (from users):
[
{
"from_name": "label",
"to_name": "image",
"values": ["cat", "dog"]
}
]
Internal storage format:
[
{
"id": "unique_id",
"from_name": "label",
"to_name": "image",
"type": "choices",
"value": {
"choices": ["cat", "dog"]
}
}
]
"""
import uuid
from typing import List, Dict, Any, Optional
from app.core.logging import get_logger
from ..schema.template import TemplateConfiguration
logger = get_logger(__name__)
class TagFormatConverter:
"""Convert between simplified external tag format and internal storage format"""
def __init__(self, template_config: TemplateConfiguration):
"""
Initialize converter with template configuration
Args:
template_config: The template configuration containing label definitions
"""
self.template_config = template_config
# Build a lookup map: from_name -> type
self._type_map = self._build_type_map()
def _build_type_map(self) -> Dict[str, str]:
"""
Build a mapping from from_name to type from template labels
Returns:
Dictionary mapping from_name to control type
"""
type_map = {}
for label_def in self.template_config.labels:
from_name = label_def.from_name
control_type = label_def.type
type_map[from_name] = control_type
logger.debug(f"Registered control: {from_name} -> {control_type}")
return type_map
def get_type_for_from_name(self, from_name: str) -> Optional[str]:
"""
Get the control type for a given from_name
Args:
from_name: The control name
Returns:
Control type or None if not found
"""
return self._type_map.get(from_name)
def convert_simplified_to_full(
self,
simplified_tags: List[Dict[str, Any]]
) -> List[Dict[str, Any]]:
"""
Convert simplified tag format to full internal storage format
Args:
simplified_tags: List of tags in simplified format with structure:
[
{
"from_name": "label",
"to_name": "image",
"values": ["cat", "dog"] # Can be list or single value
}
]
Returns:
List of tags in full internal format:
[
{
"id": "unique_id",
"from_name": "label",
"to_name": "image",
"type": "choices",
"values": {
"choices": ["cat", "dog"]
}
}
]
"""
full_tags = []
for simplified_tag in simplified_tags:
# Support both camelCase and snake_case from external sources
from_name = simplified_tag.get('from_name') or simplified_tag.get('fromName')
to_name = simplified_tag.get('to_name') or simplified_tag.get('toName')
values = simplified_tag.get('values')
tag_id = simplified_tag.get('id') # Use existing ID if provided
if not from_name or not to_name:
logger.warning(f"Skipping tag with missing from_name or to_name: {simplified_tag}")
continue
# Look up the type from template configuration
control_type = self.get_type_for_from_name(from_name)
if not control_type:
logger.warning(
f"Could not find type for from_name '{from_name}' in template. "
f"Tag will be skipped. Available controls: {list(self._type_map.keys())}"
)
continue
# Generate ID if not provided
if not tag_id:
tag_id = str(uuid.uuid4())
# Convert values to the proper nested structure
# The key in the value dict should match the control type
full_tag = {
"id": tag_id,
"from_name": from_name,
"to_name": to_name,
"type": control_type,
"values": {
control_type: values
},
}
full_tags.append(full_tag)
logger.debug(f"Converted tag: {from_name} ({control_type}) with {len(values) if isinstance(values, list) else 1} values")
return full_tags
def is_simplified_format(self, tag: Dict[str, Any]) -> bool:
"""
Check if a tag is in simplified format (missing type field)
Args:
tag: Tag dictionary to check
Returns:
True if tag appears to be in simplified format
"""
# Simplified format has top-level 'values' and no explicit type/value wrapper
has_values = 'values' in tag
has_type = 'type' in tag
has_value = 'value' in tag
# If it has 'values' but no 'type', it's simplified
# If it has 'type' and nested 'value', it's already full format
return has_values and not has_type and not has_value
def convert_if_needed(
self,
tags: List[Dict[str, Any]]
) -> List[Dict[str, Any]]:
"""
Convert tags to full format if they are in simplified format
This method can handle mixed formats - it will convert simplified tags
and pass through tags that are already in full format.
Args:
tags: List of tags in either format
Returns:
List of tags in full internal format
"""
if not tags:
return []
result = []
for tag in tags:
if self.is_simplified_format(tag):
# Convert simplified format
converted = self.convert_simplified_to_full([tag])
result.extend(converted)
else:
# Already in full format, pass through
result.append(tag)
return result
def create_converter_from_template_config(
template_config_dict: Dict[str, Any]
) -> TagFormatConverter:
"""
Create a TagFormatConverter from a template configuration dictionary
Args:
template_config_dict: Template configuration as dict (from database JSON)
Returns:
TagFormatConverter instance
Raises:
ValueError: If template configuration is invalid
"""
try:
# Parse the configuration using Pydantic model
from ..schema.template import TemplateConfiguration
template_config = TemplateConfiguration(**template_config_dict)
return TagFormatConverter(template_config)
except Exception as e:
logger.error(f"Failed to create tag converter from template config: {e}")
raise ValueError(f"Invalid template configuration: {e}")
@@ -0,0 +1,337 @@
"""
Unit tests for TagFormatConverter
Run with: pytest app/module/annotation/utils/test_tag_converter.py -v
"""
import pytest
from .tag_converter import TagFormatConverter, create_converter_from_template_config
from ..schema.template import TemplateConfiguration, LabelDefinition, ObjectDefinition
@pytest.fixture
def sample_template_config():
"""Create a sample template configuration for testing"""
return TemplateConfiguration(
labels=[
LabelDefinition(
fromName="sentiment",
toName="text",
type="choices",
options=["positive", "negative", "neutral"],
required=True,
labels=None,
description=None
),
LabelDefinition(
fromName="bbox",
toName="image",
type="rectanglelabels",
labels=["cat", "dog", "bird"],
required=False,
options=None,
description=None
),
LabelDefinition(
fromName="comment",
toName="text",
type="textarea",
required=False,
options=None,
labels=None,
description=None
)
],
objects=[
ObjectDefinition(name="text", type="Text", value="$text"),
ObjectDefinition(name="image", type="Image", value="$image")
],
metadata=None
)
@pytest.fixture
def converter(sample_template_config):
"""Create a converter instance"""
return TagFormatConverter(sample_template_config)
class TestTagFormatConverter:
"""Test TagFormatConverter functionality"""
def test_type_map_building(self, converter):
"""Test that type map is built correctly from template"""
assert converter.get_type_for_from_name("sentiment") == "choices"
assert converter.get_type_for_from_name("bbox") == "rectanglelabels"
assert converter.get_type_for_from_name("comment") == "textarea"
assert converter.get_type_for_from_name("nonexistent") is None
def test_convert_simplified_to_full_single_value(self, converter):
"""Test conversion of simplified format with single value"""
simplified = [
{
"from_name": "sentiment",
"to_name": "text",
"values": ["positive"]
}
]
result = converter.convert_simplified_to_full(simplified)
assert len(result) == 1
tag = result[0]
assert tag["from_name"] == "sentiment"
assert tag["to_name"] == "text"
assert tag["type"] == "choices"
assert tag["values"] == {"choices": ["positive"]}
assert "id" in tag
def test_convert_simplified_to_full_multiple_values(self, converter):
"""Test conversion of simplified format with multiple values"""
simplified = [
{
"from_name": "bbox",
"to_name": "image",
"values": ["cat", "dog"]
}
]
result = converter.convert_simplified_to_full(simplified)
assert len(result) == 1
tag = result[0]
assert tag["type"] == "rectanglelabels"
assert tag["values"] == {"rectanglelabels": ["cat", "dog"]}
def test_convert_simplified_camelcase(self, converter):
"""Test that camelCase field names are supported"""
simplified = [
{
"fromName": "sentiment", # camelCase
"toName": "text", # camelCase
"values": ["neutral"]
}
]
result = converter.convert_simplified_to_full(simplified)
assert len(result) == 1
assert result[0]["from_name"] == "sentiment"
assert result[0]["to_name"] == "text"
def test_convert_multiple_tags(self, converter):
"""Test conversion of multiple tags at once"""
simplified = [
{
"from_name": "sentiment",
"to_name": "text",
"values": ["positive"]
},
{
"from_name": "bbox",
"to_name": "image",
"values": ["cat"]
}
]
result = converter.convert_simplified_to_full(simplified)
assert len(result) == 2
assert result[0]["type"] == "choices"
assert result[1]["type"] == "rectanglelabels"
def test_convert_with_existing_id(self, converter):
"""Test that existing IDs are preserved"""
existing_id = "my-custom-id-123"
simplified = [
{
"id": existing_id,
"from_name": "sentiment",
"to_name": "text",
"values": ["positive"]
}
]
result = converter.convert_simplified_to_full(simplified)
assert result[0]["id"] == existing_id
def test_skip_unknown_from_name(self, converter):
"""Test that tags with unknown from_name are skipped"""
simplified = [
{
"from_name": "unknown_control",
"to_name": "text",
"values": ["value"]
}
]
result = converter.convert_simplified_to_full(simplified)
assert len(result) == 0 # Should be skipped
def test_skip_missing_fields(self, converter):
"""Test that tags with missing required fields are skipped"""
simplified = [
{
"from_name": "sentiment",
# Missing to_name
"values": ["positive"]
}
]
result = converter.convert_simplified_to_full(simplified)
assert len(result) == 0 # Should be skipped
def test_is_simplified_format(self, converter):
"""Test detection of simplified format"""
# Simplified format
assert converter.is_simplified_format({
"from_name": "x",
"to_name": "y",
"values": ["a"]
}) is True
# Full format
assert converter.is_simplified_format({
"id": "123",
"from_name": "x",
"to_name": "y",
"type": "choices",
"value": {"choices": ["a"]}
}) is False
# Edge case: has both (should not be considered simplified)
assert converter.is_simplified_format({
"from_name": "x",
"to_name": "y",
"type": "choices",
"values": ["a"]
}) is False
def test_convert_if_needed_mixed_formats(self, converter):
"""Test conversion of mixed format tags"""
mixed = [
# Simplified format
{
"from_name": "sentiment",
"to_name": "text",
"values": ["positive"]
},
# Full format
{
"id": "existing-123",
"from_name": "bbox",
"to_name": "image",
"type": "rectanglelabels",
"value": {"rectanglelabels": ["cat"]}
}
]
result = converter.convert_if_needed(mixed)
assert len(result) == 2
# First should be converted
assert result[0]["type"] == "choices"
assert result[0]["values"] == {"choices": ["positive"]}
# Second should pass through unchanged
assert result[1]["id"] == "existing-123"
assert result[1]["type"] == "rectanglelabels"
class TestCreateConverterFromDict:
"""Test the factory function for creating converter from dict"""
def test_create_from_valid_dict(self):
"""Test creating converter from valid configuration dict"""
config_dict = {
"labels": [
{
"fromName": "label",
"toName": "image",
"type": "choices",
"options": ["a", "b"]
}
],
"objects": [
{
"name": "image",
"type": "Image",
"value": "$image"
}
]
}
converter = create_converter_from_template_config(config_dict)
assert isinstance(converter, TagFormatConverter)
assert converter.get_type_for_from_name("label") == "choices"
def test_create_from_invalid_dict(self):
"""Test that invalid config raises ValueError"""
invalid_config = {
"labels": "not-a-list", # Should be a list
"objects": []
}
with pytest.raises(ValueError, match="Invalid template configuration"):
create_converter_from_template_config(invalid_config)
class TestIntegrationScenarios:
"""Test real-world usage scenarios"""
def test_external_api_submission(self, converter):
"""Simulate external user submitting tags via API"""
# User submits simplified format
user_submission = [
{
"fromName": "sentiment", # User uses camelCase
"toName": "text",
"values": ["positive", "negative"]
}
]
# System converts to internal format
internal_tags = converter.convert_if_needed(user_submission)
# Verify correct storage format
assert len(internal_tags) == 1
assert internal_tags[0]["type"] == "choices"
assert internal_tags[0]["values"] == {"choices": ["positive", "negative"]}
assert "id" in internal_tags[0]
def test_update_existing_tags(self, converter):
"""Simulate updating existing tags with new values"""
# Existing tags in database (full format)
existing_tags = [
{
"id": "tag-001",
"from_name": "sentiment",
"to_name": "text",
"type": "choices",
"value": {"choices": ["positive"]}
}
]
# User updates with simplified format
update_request = [
{
"id": "tag-001", # Same ID to update
"from_name": "sentiment",
"to_name": "text",
"values": ["negative"] # New value
}
]
# Convert update request
converted_update = converter.convert_if_needed(update_request)
# Merge logic would replace tag-001
assert converted_update[0]["id"] == "tag-001"
assert converted_update[0]["values"] == {"choices": ["negative"]}
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -0,0 +1,50 @@
from .schema import (
CleaningTaskStatus,
OperatorInstanceDto,
CleaningProcess,
CleaningTaskDto,
CreateCleaningTaskRequest,
CleaningResultDto,
CleaningTaskLog,
CleaningTemplateDto,
CreateCleaningTemplateRequest,
UpdateCleaningTemplateRequest,
)
from .repository import (
CleaningTaskRepository,
CleaningTemplateRepository,
CleaningResultRepository,
OperatorInstanceRepository,
)
from .service import (
CleaningTaskValidator,
CleaningTaskScheduler,
CleaningTemplateService,
CleaningTaskService,
)
from .runtime_client import RuntimeClient
__all__ = [
"CleaningTaskStatus",
"OperatorInstanceDto",
"CleaningProcess",
"CleaningTaskDto",
"CreateCleaningTaskRequest",
"CleaningResultDto",
"CleaningTaskLog",
"CleaningTemplateDto",
"CreateCleaningTemplateRequest",
"UpdateCleaningTemplateRequest",
"CleaningTaskRepository",
"CleaningTemplateRepository",
"CleaningResultRepository",
"OperatorInstanceRepository",
"CleaningTaskValidator",
"CleaningTaskScheduler",
"CleaningTemplateService",
"CleaningTaskService",
"RuntimeClient",
]
@@ -0,0 +1,57 @@
from typing import Optional
class CleaningException(Exception):
"""Base exception for cleaning module"""
def __init__(self, message: str, details: Optional[dict] = None):
self.message = message
self.details = details
super().__init__(self.message)
class CleaningNameDuplicationError(CleaningException):
"""Exception raised when cleaning task name is duplicated"""
def __init__(self, name: str):
super().__init__(f"Cleaning task name '{name}' is duplicated")
class CleaningTaskNotFoundError(CleaningException):
"""Exception raised when cleaning task is not found"""
def __init__(self, task_id: str):
super().__init__(f"Cleaning task '{task_id}' not found")
class CleaningTemplateNotFoundError(CleaningException):
"""Exception raised when cleaning template is not found"""
def __init__(self, template_id: str):
super().__init__(f"Cleaning template '{template_id}' not found")
class InvalidOperatorInputError(CleaningException):
"""Exception raised when operator input/output types are invalid"""
def __init__(self, message: str = "Invalid operator input/output types"):
super().__init__(message)
class ExecutorTypeError(CleaningException):
"""Exception raised when executor type is invalid"""
def __init__(self, message: str = "Invalid executor type"):
super().__init__(message)
class DatasetNotFoundError(CleaningException):
"""Exception raised when dataset is not found"""
def __init__(self, dataset_id: str):
super().__init__(f"Dataset '{dataset_id}' not found")
class FileSystemError(CleaningException):
"""Exception raised when file system operations fail"""
def __init__(self, message: str, details: Optional[dict] = None):
super().__init__(f"File system error: {message}", details)
class SettingsParseError(CleaningException):
"""Exception raised when operator settings parsing fails"""
def __init__(self, message: str, details: Optional[dict] = None):
super().__init__(f"Settings parse error: {message}", details)
@@ -0,0 +1,8 @@
from fastapi import APIRouter
from .cleaning_task_routes import router as task_router
from .cleaning_template_routes import router as template_router
router = APIRouter()
router.include_router(task_router)
router.include_router(template_router)
@@ -0,0 +1,521 @@
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,
)
@@ -0,0 +1,166 @@
import math
from typing import Optional
from fastapi import APIRouter, Depends, Query
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 (
CleaningTemplateDto,
CreateCleaningTemplateRequest,
UpdateCleaningTemplateRequest,
)
from app.module.cleaning.service import CleaningTemplateService
from app.module.shared.schema import StandardResponse, PaginatedData
logger = get_logger(__name__)
router = APIRouter(prefix="/cleaning/templates", tags=["Cleaning Templates"])
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_template_service(db: AsyncSession) -> CleaningTemplateService:
"""Get cleaning template service instance"""
from app.module.cleaning.service import CleaningTaskValidator
from app.module.cleaning.repository import (
CleaningTemplateRepository,
OperatorInstanceRepository,
)
operator_service = _get_operator_service()
template_repo = CleaningTemplateRepository(None)
return CleaningTemplateService(
template_repo=template_repo,
operator_instance_repo=OperatorInstanceRepository(None),
operator_service=operator_service,
validator=CleaningTaskValidator(task_repo=None, template_repo=template_repo),
)
@router.get(
"",
response_model=StandardResponse[PaginatedData[CleaningTemplateDto]],
summary="查询清洗模板列表",
description="分页查询清洗模板",
)
async def get_cleaning_templates(
page: int = Query(0, description="页码"),
size: int = Query(20, description="每页数量"),
keyword: Optional[str] = Query(None, description="关键词搜索"),
db: AsyncSession = Depends(get_db),
):
"""Query cleaning templates with pagination"""
template_service = _get_template_service(db)
items, total = await template_service.get_templates(db, keyword, page, size)
total_pages = math.ceil(total / size) if total > 0 else 0
return StandardResponse(
code="0",
message="success",
data=PaginatedData(
content=items,
total_elements=total,
total_pages=total_pages,
page=page,
size=size,
),
)
@router.post(
"",
response_model=StandardResponse[CleaningTemplateDto],
summary="创建清洗模板",
description="创建新的清洗模板",
)
async def create_cleaning_template(
request: CreateCleaningTemplateRequest,
db: AsyncSession = Depends(get_db),
):
"""Create cleaning template"""
template_service = _get_template_service(db)
template = await template_service.create_template(db, request)
await db.commit()
return StandardResponse(code="0", message="success", data=template)
@router.get(
"/{template_id}",
response_model=StandardResponse[CleaningTemplateDto],
summary="获取清洗模板详情",
description="根据ID获取清洗模板详细信息",
)
async def get_cleaning_template(
template_id: str,
db: AsyncSession = Depends(get_db),
):
"""Get cleaning template by ID"""
template_service = _get_template_service(db)
template = await template_service.get_template(db, template_id)
return StandardResponse(code="0", message="success", data=template)
@router.put(
"/{template_id}",
response_model=StandardResponse[CleaningTemplateDto],
summary="更新清洗模板",
description="更新清洗模板信息",
)
async def update_cleaning_template(
template_id: str,
request: UpdateCleaningTemplateRequest,
db: AsyncSession = Depends(get_db),
):
"""Update cleaning template"""
template_service = _get_template_service(db)
template = await template_service.update_template(db, template_id, request)
await db.commit()
return StandardResponse(code="0", message="success", data=template)
@router.delete(
"/{template_id}",
response_model=StandardResponse[str],
summary="删除清洗模板",
description="删除指定的清洗模板",
)
async def delete_cleaning_template(
template_id: str,
db: AsyncSession = Depends(get_db),
):
"""Delete cleaning template"""
template_service = _get_template_service(db)
await template_service.delete_template(db, template_id)
await db.commit()
return StandardResponse(code="0", message="success", data=template_id)
@@ -0,0 +1,11 @@
from .cleaning_task_repository import CleaningTaskRepository
from .cleaning_template_repository import CleaningTemplateRepository
from .cleaning_result_repository import CleaningResultRepository
from .operator_instance_repository import OperatorInstanceRepository
__all__ = [
"CleaningTaskRepository",
"CleaningTemplateRepository",
"CleaningResultRepository",
"OperatorInstanceRepository",
]
@@ -0,0 +1,110 @@
from typing import List, Optional
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, delete, func, Integer
from app.db.models.cleaning import CleaningResult
from app.module.cleaning.schema import CleaningResultDto
class CleaningResultRepository:
"""Repository for cleaning result operations"""
def __init__(self, model=None):
self.model = model if model else CleaningResult
async def find_by_instance_id(
self,
db: AsyncSession,
instance_id: str,
status: Optional[str] = None
) -> List[CleaningResultDto]:
"""Query results by instance ID"""
query = select(self.model).where(self.model.instance_id == instance_id)
if status:
query = query.where(self.model.status == status)
result = await db.execute(query)
results = result.scalars().all()
return [
CleaningResultDto(
instance_id=res.instance_id,
src_file_id=res.src_file_id,
dest_file_id=res.dest_file_id,
src_name=res.src_name,
dest_name=res.dest_name,
src_type=res.src_type,
dest_type=res.dest_type,
src_size=res.src_size,
dest_size=res.dest_size,
status=res.status,
result=res.result
)
for res in results
]
async def count_by_instance_id(
self,
db: AsyncSession,
instance_id: str
) -> tuple[int, int]:
"""Count results by instance ID (completed, failed) using single SQL"""
query = select(
func.sum(func.cast(self.model.status == "COMPLETED", Integer)).label("completed"),
func.sum(func.cast(self.model.status == "FAILED", Integer)).label("failed"),
).where(self.model.instance_id == instance_id)
result = await db.execute(query)
row = result.one()
return (row.completed or 0, row.failed or 0)
async def count_total_by_instance_id(
self,
db: AsyncSession,
instance_id: str
) -> int:
"""Count total results by instance ID using efficient SQL COUNT"""
query = select(func.count()).select_from(self.model).where(
self.model.instance_id == instance_id
)
result = await db.scalar(query)
return result or 0
async def batch_count_by_instance_ids(
self,
db: AsyncSession,
instance_ids: list[str]
) -> dict[str, tuple[int, int, int]]:
"""Batch count completed/failed/total for multiple instance IDs using single SQL.
Returns dict: {instance_id: (completed, failed, total)}
"""
if not instance_ids:
return {}
query = select(
self.model.instance_id,
func.sum(func.cast(self.model.status == "COMPLETED", Integer)).label("completed"),
func.sum(func.cast(self.model.status == "FAILED", Integer)).label("failed"),
func.count().label("total"),
).where(self.model.instance_id.in_(instance_ids)).group_by(self.model.instance_id)
result = await db.execute(query)
return {
row.instance_id: (row.completed or 0, row.failed or 0, row.total or 0)
for row in result.all()
}
async def delete_by_instance_id(
self,
db: AsyncSession,
instance_id: str,
status: Optional[str] = None
) -> None:
"""Delete results by instance ID"""
query = delete(self.model).where(self.model.instance_id == instance_id)
if status:
query = query.where(self.model.status == status)
await db.execute(query)
await db.flush()
@@ -0,0 +1,161 @@
from typing import List, Optional
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, delete, func
from app.db.models.cleaning import CleaningTask
from app.module.cleaning.schema import CleaningTaskDto
class CleaningTaskRepository:
"""Repository for cleaning task operations"""
def __init__(self, model=None):
self.model = model if model else CleaningTask
async def find_tasks(
self,
db: AsyncSession,
status: Optional[str] = None,
keyword: Optional[str] = None,
page: Optional[int] = None,
size: Optional[int] = None
) -> List[CleaningTaskDto]:
"""Query cleaning tasks"""
query = select(self.model)
if status:
query = query.where(self.model.status == status)
if keyword:
keyword_pattern = f"%{keyword}%"
query = query.where(
self.model.name.ilike(keyword_pattern) | self.model.description.ilike(keyword_pattern)
)
query = query.order_by(self.model.created_at.desc())
if page is not None and size is not None:
offset = max(page * size, 0)
query = query.offset(offset).limit(size)
result = await db.execute(query)
tasks = result.scalars().all()
return [
CleaningTaskDto(
id=task.id,
name=task.name,
description=task.description,
status=task.status,
src_dataset_id=task.src_dataset_id,
src_dataset_name=task.src_dataset_name,
dest_dataset_id=task.dest_dataset_id,
dest_dataset_name=task.dest_dataset_name,
before_size=task.before_size,
after_size=task.after_size,
file_count=task.file_count,
retry_count=task.retry_count,
started_at=task.started_at,
finished_at=task.finished_at,
created_at=task.created_at
)
for task in tasks
]
async def find_task_by_id(self, db: AsyncSession, task_id: str) -> Optional[CleaningTaskDto]:
"""Query task by ID"""
query = select(self.model).where(self.model.id == task_id)
result = await db.execute(query)
task = result.scalar_one_or_none()
if not task:
return None
return CleaningTaskDto(
id=task.id,
name=task.name,
description=task.description,
status=task.status,
src_dataset_id=task.src_dataset_id,
src_dataset_name=task.src_dataset_name,
dest_dataset_id=task.dest_dataset_id,
dest_dataset_name=task.dest_dataset_name,
before_size=task.before_size,
after_size=task.after_size,
file_count=task.file_count,
retry_count=task.retry_count,
started_at=task.started_at,
finished_at=task.finished_at,
created_at=task.created_at
)
async def insert_task(self, db: AsyncSession, task: CleaningTaskDto) -> None:
"""Insert new task"""
from app.db.models.cleaning import CleaningTask as CleaningTaskModel
db_task = CleaningTaskModel(
id=task.id,
name=task.name,
description=task.description,
status=task.status,
src_dataset_id=task.src_dataset_id,
src_dataset_name=task.src_dataset_name,
dest_dataset_id=task.dest_dataset_id,
dest_dataset_name=task.dest_dataset_name,
before_size=task.before_size,
after_size=task.after_size,
file_count=task.file_count,
retry_count=task.retry_count
)
db.add(db_task)
await db.flush()
async def update_task(self, db: AsyncSession, task: CleaningTaskDto) -> None:
"""Update task"""
query = select(CleaningTask).where(CleaningTask.id == task.id)
result = await db.execute(query)
db_task = result.scalar_one_or_none()
if db_task:
if task.status:
db_task.status = task.status
if task.started_at:
db_task.started_at = task.started_at
if task.finished_at:
db_task.finished_at = task.finished_at
if task.retry_count is not None:
db_task.retry_count = task.retry_count
await db.flush()
async def delete_task_by_id(self, db: AsyncSession, task_id: str) -> None:
"""Delete task by ID"""
query = delete(self.model).where(self.model.id == task_id)
await db.execute(query)
await db.flush()
async def count_tasks(
self,
db: AsyncSession,
status: Optional[str] = None,
keyword: Optional[str] = None
) -> int:
"""Count cleaning tasks using SQL COUNT"""
query = select(func.count()).select_from(self.model)
if status:
query = query.where(self.model.status == status)
if keyword:
keyword_pattern = f"%{keyword}%"
query = query.where(
self.model.name.ilike(keyword_pattern) | self.model.description.ilike(keyword_pattern)
)
result = await db.scalar(query)
return result or 0
async def is_name_exist(self, db: AsyncSession, name: str) -> bool:
"""Check if task name exists"""
query = select(func.count()).select_from(self.model).where(self.model.name == name)
result = await db.execute(query)
return result.scalar_one() > 0 if result else False
@@ -0,0 +1,95 @@
from typing import List, Optional
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, delete, func
from app.db.models.cleaning import CleaningTemplate
class CleaningTemplateRepository:
"""Repository for cleaning template operations"""
def __init__(self, model=None):
self.model = model if model else CleaningTemplate
async def find_all_templates(
self,
db: AsyncSession,
keyword: Optional[str] = None,
page: Optional[int] = None,
size: Optional[int] = None,
) -> List[CleaningTemplate]:
"""Query all templates with optional pagination"""
query = select(self.model)
if keyword:
keyword_pattern = f"%{keyword}%"
query = query.where(
self.model.name.ilike(keyword_pattern)
| self.model.description.ilike(keyword_pattern)
)
query = query.order_by(self.model.created_at.desc())
if page is not None and size is not None:
offset = page * size
query = query.offset(offset).limit(size)
result = await db.execute(query)
return result.scalars().all()
async def count_templates(
self, db: AsyncSession, keyword: Optional[str] = None
) -> int:
"""Count templates"""
query = select(func.count()).select_from(self.model)
if keyword:
keyword_pattern = f"%{keyword}%"
query = query.where(
self.model.name.ilike(keyword_pattern)
| self.model.description.ilike(keyword_pattern)
)
result = await db.execute(query)
return result.scalar_one()
async def find_template_by_id(
self, db: AsyncSession, template_id: str
) -> Optional[CleaningTemplate]:
"""Query template by ID"""
query = select(self.model).where(self.model.id == template_id)
result = await db.execute(query)
return result.scalar_one_or_none()
async def insert_template(
self, db: AsyncSession, template: CleaningTemplate
) -> None:
"""Insert new template"""
db.add(template)
await db.flush()
async def update_template(
self, db: AsyncSession, template: CleaningTemplate
) -> None:
"""Update template"""
query = select(self.model).where(self.model.id == template.id)
result = await db.execute(query)
db_template = result.scalar_one_or_none()
if db_template:
db_template.name = template.name
db_template.description = template.description
await db.flush()
async def delete_template(self, db: AsyncSession, template_id: str) -> None:
"""Delete template"""
query = delete(self.model).where(self.model.id == template_id)
await db.execute(query)
await db.flush()
async def is_name_exist(self, db: AsyncSession, name: str) -> bool:
"""Check if template name exists"""
query = (
select(func.count()).select_from(self.model).where(self.model.name == name)
)
result = await db.execute(query)
return result.scalar_one() > 0 if result else False
@@ -0,0 +1,56 @@
import json
from typing import List
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, delete
from app.db.models.cleaning import OperatorInstance
class OperatorInstanceRepository:
"""Repository for operator instance operations"""
def __init__(self, model=None):
self.model = model if model else OperatorInstance
async def find_operator_by_instance_id(
self,
db: AsyncSession,
instance_id: str
) -> List[OperatorInstance]:
"""Query operator instances by instance ID"""
query = select(self.model).where(self.model.instance_id == instance_id)
query = query.order_by(self.model.op_index.asc())
result = await db.execute(query)
return result.scalars().all()
async def find_instance_by_instance_id(
self,
db: AsyncSession,
instance_id: str
) -> List[OperatorInstance]:
"""Query instances for template (same as find_operator_by_instance_id)"""
return await self.find_operator_by_instance_id(db, instance_id)
async def insert_instance(
self,
db: AsyncSession,
instance_id: str,
instances: List
) -> None:
"""Insert operator instances"""
from app.db.models.cleaning import OperatorInstance as OperatorInstanceModel
for idx, instance in enumerate(instances):
db_instance = OperatorInstanceModel(
instance_id=instance_id,
operator_id=instance.id,
op_index=idx,
settings_override=json.dumps(instance.overrides),
)
db.add(db_instance)
await db.flush()
async def delete_by_instance_id(self, db: AsyncSession, instance_id: str) -> None:
"""Delete instances by instance ID"""
query = delete(self.model).where(self.model.instance_id == instance_id)
await db.execute(query)
await db.flush()
@@ -0,0 +1,61 @@
import httpx
from typing import Optional
from app.core.logging import get_logger
logger = get_logger(__name__)
class RuntimeClient:
"""HTTP client for communicating with runtime service"""
def __init__(self, base_url: str = "http://deepdata-runtime:8081"):
self.base_url = base_url
self.client = httpx.AsyncClient(timeout=60.0)
async def submit_task(self, task_id: str, retry_count: int = 0) -> bool:
"""Submit cleaning task to runtime executor"""
try:
url = f"{self.base_url}/api/task/{task_id}/submit"
response = await self.client.post(url, json={"retry_count": retry_count})
response.raise_for_status()
logger.info(f"Task {task_id} submitted successfully with retry_count={retry_count}")
return True
except httpx.HTTPError as e:
logger.error(f"Failed to submit task {task_id}: {e}")
return False
except Exception as e:
logger.error(f"Unexpected error submitting task {task_id}: {e}")
return False
async def stop_task(self, task_id: str) -> bool:
"""Stop running cleaning task"""
try:
url = f"{self.base_url}/api/task/{task_id}/stop"
response = await self.client.post(url)
response.raise_for_status()
logger.info(f"Task {task_id} stopped successfully")
return True
except httpx.HTTPError as e:
logger.error(f"Failed to stop task {task_id}: {e}")
return False
except Exception as e:
logger.error(f"Unexpected error stopping task {task_id}: {e}")
return False
async def get_task_status(self, task_id: str) -> Optional[dict]:
"""Get task status from runtime"""
try:
url = f"{self.base_url}/api/task/{task_id}/status"
response = await self.client.get(url)
response.raise_for_status()
return response.json()
except httpx.HTTPError as e:
logger.error(f"Failed to get task status {task_id}: {e}")
return None
except Exception as e:
logger.error(f"Unexpected error getting task status {task_id}: {e}")
return None
async def close(self):
"""Close HTTP client"""
await self.client.aclose()
@@ -0,0 +1,25 @@
from .cleaning import (
CleaningTaskStatus,
OperatorInstanceDto,
CleaningProcess,
CleaningTaskDto,
CreateCleaningTaskRequest,
CleaningResultDto,
CleaningTaskLog,
CleaningTemplateDto,
CreateCleaningTemplateRequest,
UpdateCleaningTemplateRequest,
)
__all__ = [
"CleaningTaskStatus",
"OperatorInstanceDto",
"CleaningProcess",
"CleaningTaskDto",
"CreateCleaningTaskRequest",
"CleaningResultDto",
"CleaningTaskLog",
"CleaningTemplateDto",
"CreateCleaningTemplateRequest",
"UpdateCleaningTemplateRequest",
]
@@ -0,0 +1,148 @@
from typing import Optional, List, Dict, Any
from pydantic import BaseModel, Field
from datetime import datetime
from app.module.shared.schema.common import BaseResponseModel
class CleaningTaskStatus:
PENDING = "PENDING"
RUNNING = "RUNNING"
COMPLETED = "COMPLETED"
PARTIAL_SUCCESS = "PARTIAL_SUCCESS"
STOPPED = "STOPPED"
FAILED = "FAILED"
class OperatorInstanceDto(BaseResponseModel):
"""Operator instance DTO for task or template"""
id: str = Field(..., description="Operator ID")
name: Optional[str] = Field(None, description="Operator name")
description: Optional[str] = Field(None, description="Operator description")
inputs: Optional[str] = Field(None, description="Input types: text/image/audio/video/multimodal")
outputs: Optional[str] = Field(None, description="Output types: text/image/audio/video/multimodal")
categories: Optional[List[str]] = Field(None, description="Category IDs")
settings: Optional[str] = Field(None, description="算子设置(JSON)")
overrides: Dict[str, Any] = Field(default_factory=dict, description="Operator parameter overrides")
class CleaningProcess(BaseResponseModel):
"""Task progress information"""
process: float = Field(..., description="Progress percentage (based on successful files)")
successRate: float = Field(..., description="Success rate percentage")
totalFileNum: int = Field(..., description="Total file count")
succeedFileNum: int = Field(..., description="Succeeded file count")
failedFileNum: int = Field(..., description="Failed file count")
finishedFileNum: int = Field(..., description="Successfully processed file count (excludes failed)")
@classmethod
def of(cls, total: int, succeed: int, failed: int) -> 'CleaningProcess':
"""Create progress info
- finishedFileNum: only counts successfully processed files (excludes failed)
- process: progress percentage based on successful files only
- successRate: percentage of finished files that succeeded
"""
# finished_file_num only counts successfully processed files
finished_file_num = succeed
if total == 0:
process = 0.0
else:
# Progress only counts successful files
process = round(succeed * 100.0 / total, 2)
# Calculate total processed (succeed + failed) for success rate
total_processed = succeed + failed
if total_processed == 0:
success_rate = 0.0
else:
success_rate = round(succeed * 100.0 / total_processed, 2)
return cls(
process=process,
successRate=success_rate,
totalFileNum=total,
succeedFileNum=succeed,
failedFileNum=failed,
finishedFileNum=finished_file_num,
)
class CleaningTaskDto(BaseResponseModel):
"""Cleaning task DTO"""
id: Optional[str] = Field(None, description="Task ID")
name: Optional[str] = Field(None, description="Task name")
description: Optional[str] = Field(None, description="Task description")
src_dataset_id: Optional[str] = Field(None, description="Source dataset ID")
src_dataset_name: Optional[str] = Field(None, description="Source dataset name")
dest_dataset_id: Optional[str] = Field(None, description="Destination dataset ID")
dest_dataset_name: Optional[str] = Field(None, description="Destination dataset name")
before_size: Optional[int] = Field(None, description="Data size before cleaning")
after_size: Optional[int] = Field(None, description="Data size after cleaning")
file_count: Optional[int] = Field(None, description="Total file count")
retry_count: Optional[int] = Field(None, description="Retry count")
status: Optional[str] = Field(None, description="Task status")
template_id: Optional[str] = Field(None, description="Template ID if created from template")
instance: Optional[List[OperatorInstanceDto]] = Field(None, description="Operator instances")
progress: Optional[CleaningProcess] = Field(None, description="Task progress")
created_at: Optional[datetime] = Field(None, description="Creation time")
started_at: Optional[datetime] = Field(None, description="Start time")
finished_at: Optional[datetime] = Field(None, description="Finish time")
class CreateCleaningTaskRequest(BaseResponseModel):
"""Request to create cleaning task"""
name: str = Field(..., description="Cleaning task name")
description: str = Field(..., description="Cleaning task description")
src_dataset_id: str = Field(..., description="Source dataset ID")
src_dataset_name: str = Field(..., description="Source dataset name")
dest_dataset_id: Optional[str] = Field(None, description="Destination dataset ID")
dest_dataset_name: str = Field(..., description="Destination dataset name, creates new dataset if destDatasetId is empty")
dest_dataset_type: str = Field(..., description="Destination dataset type: TEXT/IMAGE/VIDEO/AUDIO/OTHER")
template_id: Optional[str] = Field(None, description="Template ID (alternative to instance)")
instance: List[OperatorInstanceDto] = Field(default_factory=list, description="Operator list (alternative to templateId)")
class CleaningResultDto(BaseResponseModel):
"""Cleaning result DTO"""
instance_id: Optional[str] = Field(None, description="Instance ID")
src_file_id: Optional[str] = Field(None, description="Source file ID")
dest_file_id: Optional[str] = Field(None, description="Destination file ID")
src_name: Optional[str] = Field(None, description="Source file name")
dest_name: Optional[str] = Field(None, description="Destination file name")
src_type: Optional[str] = Field(None, description="Source file type")
dest_type: Optional[str] = Field(None, description="Destination file type")
src_size: Optional[int] = Field(None, description="Source file size")
dest_size: Optional[int] = Field(None, description="Destination file size")
status: Optional[str] = Field(None, description="Cleaning status")
result: Optional[str] = Field(None, description="Cleaning result message")
class CleaningTaskLog(BaseResponseModel):
"""Task log entry"""
level: str = Field(..., description="Log level: INFO, WARN, ERROR")
message: str = Field(..., description="Log message")
class CleaningTemplateDto(BaseResponseModel):
"""Cleaning template DTO"""
id: Optional[str] = Field(None, description="Template ID")
name: Optional[str] = Field(None, description="Template name")
description: Optional[str] = Field(None, description="Template description")
instance: List[OperatorInstanceDto] = Field(default_factory=list, description="Operator instances")
created_at: Optional[datetime] = Field(None, description="Creation time")
updated_at: Optional[datetime] = Field(None, description="Update time")
class CreateCleaningTemplateRequest(BaseResponseModel):
"""Request to create cleaning template"""
name: str = Field(..., description="Template name")
description: str = Field(..., description="Template description")
instance: List[OperatorInstanceDto] = Field(..., description="Operator instances")
class UpdateCleaningTemplateRequest(BaseResponseModel):
"""Request to update cleaning template"""
name: str = Field(..., description="Template name")
description: str = Field(..., description="Template description")
instance: List[OperatorInstanceDto] = Field(..., description="Operator instances")
@@ -0,0 +1,11 @@
from .cleaning_task_validator import CleaningTaskValidator
from .cleaning_task_scheduler import CleaningTaskScheduler
from .cleaning_template_service import CleaningTemplateService
from .cleaning_task_service import CleaningTaskService
__all__ = [
"CleaningTaskValidator",
"CleaningTaskScheduler",
"CleaningTemplateService",
"CleaningTaskService",
]
@@ -0,0 +1,164 @@
import asyncio
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.logging import get_logger
from app.module.cleaning.repository import CleaningTaskRepository
from app.module.cleaning.runtime_client import RuntimeClient
logger = get_logger(__name__)
# 模块级单例:进程内所有请求共享同一个调度器
_scheduler: "CleaningTaskScheduler | None" = None
def get_scheduler() -> "CleaningTaskScheduler":
"""获取全局调度器单例,不存在则创建"""
global _scheduler
if _scheduler is None:
_scheduler = CleaningTaskScheduler(
CleaningTaskRepository(None),
RuntimeClient()
)
return _scheduler
class CleaningTaskScheduler:
"""Scheduler for executing cleaning tasks"""
def __init__(self, task_repo: CleaningTaskRepository, runtime_client: RuntimeClient):
self.task_repo = task_repo
self.runtime_client = runtime_client
self._polling_task_ids: set[str] = set() # 待轮询的任务 ID 集合
self._polling_started: bool = False
self._poll_failure_count: dict[str, int] = {} # runtime 不可达连续失败计数
self._MAX_POLL_FAILURES = 1 # 一次超时(60s) → 标记 FAILED
async def execute_task(self, db: AsyncSession, task_id: str, retry_count: int) -> bool:
"""Execute cleaning task"""
from app.module.cleaning.schema import CleaningTaskDto, CleaningTaskStatus
from datetime import datetime
task = CleaningTaskDto()
task.id = task_id
task.status = CleaningTaskStatus.RUNNING
task.started_at = datetime.now()
task.retry_count = retry_count
submitted = await self.runtime_client.submit_task(task_id, retry_count)
if submitted:
await self.task_repo.update_task(db, task)
self._polling_task_ids.add(task_id)
return submitted
task.status = CleaningTaskStatus.FAILED
task.finished_at = datetime.now()
await self.task_repo.update_task(db, task)
return submitted
async def stop_task(self, db: AsyncSession, task_id: str) -> bool:
"""Stop cleaning task"""
from app.module.cleaning.schema import CleaningTaskDto, CleaningTaskStatus
self._polling_task_ids.discard(task_id)
self._poll_failure_count.pop(task_id, None)
await self.runtime_client.stop_task(task_id)
task = CleaningTaskDto()
task.id = task_id
task.status = CleaningTaskStatus.STOPPED
await self.task_repo.update_task(db, task)
return True
async def startup(self):
"""进程启动时调用:恢复未完成任务的轮询 + 启动全局轮询协程"""
from app.module.cleaning.schema import CleaningTaskStatus
from app.db.session import AsyncSessionLocal
# 从数据库恢复所有 RUNNING 状态的任务
try:
async with AsyncSessionLocal() as db:
tasks = await self.task_repo.find_tasks(db, status=CleaningTaskStatus.RUNNING)
for task in tasks:
logger.info(f"[Polling] Recovered RUNNING task from DB: {task.id}")
self._polling_task_ids.add(task.id)
except Exception as e:
logger.error(f"[Polling] Failed to recover tasks from DB: {e}")
# 启动全局轮询协程
if not self._polling_started:
self._polling_started = True
asyncio.create_task(self._poll_all_tasks())
logger.info("[Polling] Global polling loop started")
async def _poll_all_tasks(self):
"""全局轮询协程:每 2 秒轮询所有 RUNNING 任务的 runtime 状态"""
from app.module.cleaning.schema import CleaningTaskDto, CleaningTaskStatus
from app.db.session import AsyncSessionLocal
from datetime import datetime
logger.info("[Polling] Global status polling loop started")
terminal_statuses = {"completed", "failed", "cancelled", "stopped", "not_found"}
while True:
task_ids = list(self._polling_task_ids)
if not task_ids:
# 没有待轮询任务,休眠后继续等
await asyncio.sleep(2)
continue
for task_id in task_ids:
try:
status_data = await self.runtime_client.get_task_status(task_id)
if status_data is None:
# runtime 不可达,累计失败计数
count = self._poll_failure_count.get(task_id, 0) + 1
self._poll_failure_count[task_id] = count
logger.warning(
f"[Polling] Task {task_id} unreachable ({count}/{self._MAX_POLL_FAILURES})"
)
if count >= self._MAX_POLL_FAILURES:
logger.error(
f"[Polling] Task {task_id} marked FAILED: "
f"runtime unreachable after {count} attempts"
)
async with AsyncSessionLocal() as db:
task_dto = CleaningTaskDto()
task_dto.id = task_id
task_dto.status = CleaningTaskStatus.FAILED
task_dto.finished_at = datetime.now()
await self.task_repo.update_task(db, task_dto)
await db.commit()
self._polling_task_ids.discard(task_id)
self._poll_failure_count.pop(task_id, None)
continue
# runtime 可达,重置失败计数
self._poll_failure_count.pop(task_id, None)
current_status = (status_data.get("status", "") or "").lower()
logger.debug(f"[Polling] Task {task_id} status: {current_status}")
if current_status in terminal_statuses:
async with AsyncSessionLocal() as db:
task = CleaningTaskDto()
task.id = task_id
if current_status == "completed":
task.status = CleaningTaskStatus.COMPLETED
else:
task.status = CleaningTaskStatus.FAILED
task.finished_at = datetime.now()
await self.task_repo.update_task(db, task)
await db.commit()
logger.info(
f"[Polling] Task {task_id} finished: {current_status}"
)
self._polling_task_ids.discard(task_id)
self._poll_failure_count.pop(task_id, None)
except Exception as e:
logger.error(f"[Polling] Error polling task {task_id}: {e}")
await asyncio.sleep(2)
@@ -0,0 +1,578 @@
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
@@ -0,0 +1,136 @@
import re
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.exception import BusinessError, ErrorCodes
from app.module.cleaning.schema import OperatorInstanceDto
from app.module.operator.constants import CATEGORY_DATA_JUICER_ID, CATEGORY_DEEPDATA_ID
# UUID pattern for task_id validation (prevents path traversal)
_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}$'
)
class CleaningTaskValidator:
"""Validator for cleaning tasks and templates"""
def __init__(self, task_repo=None, template_repo=None):
self.task_repo = task_repo
self.template_repo = template_repo
async def check_task_name_duplication(self, db: AsyncSession, name: str) -> None:
"""Check if task name is duplicated"""
if not name:
raise BusinessError(ErrorCodes.CLEANING_NAME_DUPLICATED)
if await self.task_repo.is_name_exist(db, name):
raise BusinessError(ErrorCodes.CLEANING_NAME_DUPLICATED)
async def check_template_name_duplication(self, db: AsyncSession, name: str) -> None:
"""Check if template name is duplicated"""
if not name:
raise BusinessError(ErrorCodes.CLEANING_TEMPLATE_NAME_DUPLICATED)
if await self.template_repo.is_name_exist(db, name):
raise BusinessError(ErrorCodes.CLEANING_TEMPLATE_NAME_DUPLICATED)
@staticmethod
def check_input_and_output(instances: list[OperatorInstanceDto]) -> None:
"""Validate that operator input/output types are compatible.
Rules:
- multimodal is compatible with any type (acts as wildcard)
- for other types (text, image, audio, video), upstream outputs
must match downstream inputs exactly
"""
if not instances:
return
for i in range(len(instances) - 1):
current = instances[i]
next_op = instances[i + 1]
if not current.outputs:
raise BusinessError(
ErrorCodes.CLEANING_INVALID_OPERATOR_INPUT,
f"Operator {current.id} has no outputs defined"
)
if not next_op.inputs:
raise BusinessError(
ErrorCodes.CLEANING_INVALID_OPERATOR_INPUT,
f"Operator {next_op.id} has no inputs defined"
)
current_outputs = current.outputs.lower().strip()
next_inputs = next_op.inputs.lower().strip()
# multimodal is compatible with everything
if "multimodal" in current_outputs or "multimodal" in next_inputs:
continue
# non-multimodal: types must match exactly
if current_outputs != next_inputs:
raise BusinessError(
ErrorCodes.CLEANING_INVALID_OPERATOR_INPUT,
f"Operator '{current.id}' outputs '{current.outputs}', "
f"but operator '{next_op.id}' requires '{next_op.inputs}'"
)
@staticmethod
def check_and_get_executor_type(instances: list[OperatorInstanceDto]) -> str:
"""Check operator categories and determine executor type (deepdata/datajuicer)"""
if not instances:
return "deepdata"
executor_types = set()
for instance in instances:
if instance.categories:
for category in instance.categories:
if CATEGORY_DATA_JUICER_ID in category.lower():
executor_types.add("default")
elif CATEGORY_DEEPDATA_ID in category.lower():
executor_types.add("deepdata")
if len(executor_types) > 1:
raise BusinessError(
ErrorCodes.CLEANING_INVALID_EXECUTOR_TYPE,
"Cannot mix DeepData and DataJuicer operators in same task"
)
return executor_types.pop() if executor_types else "deepdata"
@staticmethod
def check_task_id(task_id: str) -> None:
"""Validate task ID — rejects non-UUID and path traversal patterns"""
CleaningTaskValidator.sanitize_task_id(task_id)
@staticmethod
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:
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
@staticmethod
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 BusinessError(
ErrorCodes.CLEANING_TASK_ID_REQUIRED,
f"Invalid retry_count: {retry_count}",
)
return retry_count
@@ -0,0 +1,249 @@
import json
import uuid
from typing import List
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.exception import BusinessError, ErrorCodes
from app.core.logging import get_logger
from app.module.cleaning import UpdateCleaningTemplateRequest
from app.module.cleaning.repository import (
CleaningTemplateRepository,
OperatorInstanceRepository,
)
from app.module.cleaning.schema import (
CleaningTemplateDto,
CreateCleaningTemplateRequest,
OperatorInstanceDto,
)
from app.module.cleaning.service.cleaning_task_validator import CleaningTaskValidator
logger = get_logger(__name__)
PRESET_TEMPLATE_IDS = {
"550e8400-e29b-41d4-a716-446655440001",
"661f9500-f3ac-52e5-b827-557766550002",
"772a0611-a4bd-63f6-c938-668877660003",
"883b1722-b5ce-7407-d049-779988770004",
"994c2833-c6df-8518-e150-880099880005",
}
class CleaningTemplateService:
"""Service for managing cleaning templates"""
def __init__(
self,
template_repo: CleaningTemplateRepository,
operator_instance_repo: OperatorInstanceRepository,
operator_service,
validator: CleaningTaskValidator,
):
self.template_repo = template_repo
self.operator_instance_repo = operator_instance_repo
self.operator_service = operator_service
self.validator = validator
async def get_templates(
self,
db: AsyncSession,
keyword: str | None = None,
page: int | None = None,
size: int | None = None,
) -> tuple[List[CleaningTemplateDto], int]:
"""Get templates with pagination"""
total = await self.template_repo.count_templates(db, keyword)
templates = await self.template_repo.find_all_templates(db, keyword, page, size)
# Collect all operator IDs
template_instances_map = {}
for template in templates:
instances = await self.operator_instance_repo.find_operator_by_instance_id(
db, template.id
)
template_instances_map[template.id] = instances
# Batch query all 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}
# Build result
result = []
for template in templates:
template_dto = CleaningTemplateDto(
id=template.id,
name=template.name,
description=template.description,
instance=[],
created_at=template.created_at,
updated_at=template.updated_at,
)
instances = template_instances_map.get(template.id, [])
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}"
)
template_dto.instance.append(operator_dto)
result.append(template_dto)
return result, total
async def get_template(
self, db: AsyncSession, template_id: str
) -> CleaningTemplateDto:
"""Get template by ID"""
template = await self.template_repo.find_template_by_id(db, template_id)
if not template:
raise BusinessError(ErrorCodes.CLEANING_TEMPLATE_NOT_FOUND, template_id)
template_dto = CleaningTemplateDto(
id=template.id,
name=template.name,
description=template.description,
instance=[],
created_at=template.created_at,
updated_at=template.updated_at,
)
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}
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}"
)
template_dto.instance.append(operator_dto)
return template_dto
async def create_template(
self, db: AsyncSession, request: CreateCleaningTemplateRequest
) -> CleaningTemplateDto:
"""Create new template"""
from app.db.models.cleaning import CleaningTemplate
await self.validator.check_template_name_duplication(db, request.name)
self.validator.check_input_and_output(request.instance)
self.validator.check_and_get_executor_type(request.instance)
template_id = str(uuid.uuid4())
template = CleaningTemplate(
id=template_id,
name=request.name,
description=request.description,
)
await self.template_repo.insert_template(db, template)
await self.operator_instance_repo.insert_instance(
db, template_id, request.instance
)
return await self.get_template(db, template_id)
async def update_template(
self, db: AsyncSession, template_id: str, request: UpdateCleaningTemplateRequest
) -> CleaningTemplateDto:
"""Update template"""
template = await self.template_repo.find_template_by_id(db, template_id)
if not template:
raise BusinessError(ErrorCodes.CLEANING_TEMPLATE_NOT_FOUND, template_id)
template.name = request.name
template.description = request.description
await self.template_repo.update_template(db, template)
await self.operator_instance_repo.delete_by_instance_id(db, template_id)
await self.operator_instance_repo.insert_instance(
db, template_id, request.instance
)
return await self.get_template(db, template_id)
async def delete_template(self, db: AsyncSession, template_id: str) -> None:
"""Delete template"""
if template_id in PRESET_TEMPLATE_IDS:
raise BusinessError(ErrorCodes.CLEANING_CANNOT_DELETE_PRESET_TEMPLATE)
await self.template_repo.delete_template(db, template_id)
await self.operator_instance_repo.delete_by_instance_id(db, template_id)
async def get_instance_by_template_id(
self, db: AsyncSession, template_id: str
) -> List[OperatorInstanceDto]:
"""Get operator instances by template ID"""
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
@@ -0,0 +1,251 @@
import json
import threading
import subprocess
from datetime import datetime
from pathlib import Path
from typing import Dict, Any
from app.core.logging import get_logger
from app.core.security import mask_sensitive_info
from app.db.models.data_collection import CollectionTask, TaskExecution, CollectionTemplate
from app.module.collection.schema.collection import CollectionConfig, SyncMode
from app.module.shared.schema import TaskStatus
logger = get_logger(__name__)
class DataxClient:
def __init__(self, task: CollectionTask, execution: TaskExecution, template: CollectionTemplate):
self.execution = execution
self.task = task
self.template = template
# ✅ 确保路径在共享存储 /flow 下,供 Runtime 和 Ray Worker 访问
self.config_file_path = f"/flow/data-collection/{task.id}/config.json"
self.python_path = "python"
self.datax_main = "/opt/datax/bin/datax.py"
Path(self.config_file_path).parent.mkdir(parents=True, exist_ok=True)
def validate_json_string(self) -> Dict[str, Any]:
"""
验证 JSON 字符串
Returns:
解析后的配置字典
"""
try:
config = json.loads(self.task.config)
# 基本验证
if 'job' not in config:
raise ValueError("JSON 必须包含 'job' 字段")
if 'content' not in config.get('job', {}):
raise ValueError("job 必须包含 'content' 字段")
logger.info("JSON 配置验证通过")
return config
except json.JSONDecodeError as e:
raise ValueError(f"JSON 格式错误: {e}")
except Exception as e:
raise ValueError(f"配置验证失败: {e}")
@staticmethod
def generate_datx_config(task_config: CollectionConfig, template: CollectionTemplate, target_path: str):
# 校验参数
dest_path_target = {"nfswriter", "s3writer", "glusterfswriter"}
reader_parameter = {
**(task_config.parameter if task_config.parameter else {}),
**(task_config.reader if task_config.reader else {})
}
dest_parameter = {}
if template.target_type == "txtfilewriter":
dest_parameter = {
"path": target_path,
"fileName": "collection_result",
"writeMode": "truncate",
"fileFormat": "csv"
}
elif dest_path_target.__contains__(template.target_type):
dest_parameter = {
"destPath": target_path
}
writer_parameter = {
**(task_config.parameter if task_config.parameter else {}),
**(task_config.writer if task_config.writer else {}),
**dest_parameter
}
# 生成任务运行配置
job_config = {
"content": [
{
"reader": {
"name": template.source_type,
"parameter": reader_parameter
},
"writer": {
"name": template.target_type,
"parameter": writer_parameter
}
}
],
"setting": {
"speed": {
"channel": 2
}
}
}
task_config.job = job_config
def create__config_file(self) -> str:
"""
创建配置文件
Returns:
临时文件路径
"""
# 验证 JSON
config = self.validate_json_string()
# 写入临时文件
with open(self.config_file_path, 'w', encoding='utf-8') as f:
json.dump(config, f, indent=2, ensure_ascii=False)
logger.debug(f"创建配置文件: {self.config_file_path}")
return self.config_file_path
def generate_config(self) -> str:
"""
✅ 新增:仅生成 DataX JSON 配置,不执行任务。
用于 Runtime/Ray 模式:Backend 生成配置后提交给集群执行。
Returns:
配置文件的绝对路径
"""
return self.create__config_file()
def run_datax_job(self):
"""
启动 DataX 任务 (Local/Fallback 模式)
Returns:
执行结果字典
"""
if not self.execution.started_at:
self.execution.started_at = datetime.now()
try:
# ✅ 复用 generate_config 生成文件
config_path = self.generate_config()
cmd = [self.python_path, str(self.datax_main), str(config_path)]
cmd_str = ' '.join(cmd)
logger.info(f"执行命令: {cmd_str}")
with open(self.execution.log_path, 'w', encoding='utf-8') as log_f:
self.write_header_log(cmd_str, log_f)
exit_code = self._run_process(cmd, log_f)
self.execution.completed_at = datetime.now()
self.execution.duration_seconds = (self.execution.completed_at - self.execution.started_at).total_seconds()
self.write_tail_log(exit_code, log_f)
if exit_code == 0:
logger.info(f"DataX 任务执行成功: {self.execution.id}")
logger.info(f"执行耗时: {self.execution.duration_seconds:.2f} 秒")
self.execution.status = TaskStatus.COMPLETED.name
self.rename_collection_result()
else:
self.execution.error_message = self.execution.error_message or f"DataX 任务执行失败,退出码: {exit_code}"
self.execution.status = TaskStatus.FAILED.name
logger.error(self.execution.error_message)
except Exception as e:
self.execution.completed_at = datetime.now()
if self.execution.started_at:
self.execution.duration_seconds = (self.execution.completed_at - self.execution.started_at).total_seconds()
self.execution.error_message = f"执行异常: {e}"
self.execution.status = TaskStatus.FAILED.name
logger.error(f"执行异常: {e}", exc_info=True)
with open(self.execution.log_path, 'w', encoding='utf-8') as log_f:
log_f.write(f"任务执行失败: {e}\n")
if self.task.sync_mode == SyncMode.ONCE:
self.task.status = self.execution.status
else:
self.task.status = TaskStatus.PENDING.name
def rename_collection_result(self):
if self.template.target_type != "txtfilewriter":
return
target_path = Path(self.task.target_path)
if not target_path.exists():
logger.warning(f"Target path does not exist: {target_path}")
return
# If it's a directory, find all files without extensions
for file_path in target_path.iterdir():
if file_path.is_file() and not file_path.suffix:
new_path = file_path.with_suffix('.csv')
try:
file_path.rename(new_path)
logger.info(f"Renamed {file_path} to {new_path}")
except Exception as e:
logger.error(f"Failed to rename {file_path} to {new_path}: {str(e)}")
def _run_process(self, cmd: list[str], log_f) -> int:
# 启动进程
process = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
encoding='utf-8',
bufsize=1,
universal_newlines=True
)
# 创建读取线程
stdout_thread = threading.Thread(target=lambda stream=process.stdout: self.read_stream(stream, log_f))
stderr_thread = threading.Thread(target=lambda stream=process.stderr: self.read_stream(stream, log_f))
stdout_thread.start()
stderr_thread.start()
# 等待进程完成
try:
exit_code = process.wait(timeout=self.task.timeout_seconds)
except subprocess.TimeoutExpired:
process.kill()
exit_code = -1
self.execution.error_message = f"任务执行超时({self.task.timeout_seconds}秒)"
logger.error(f"任务执行超时({self.task.timeout_seconds}秒)")
# 等待线程完成
stdout_thread.join(timeout=5)
stderr_thread.join(timeout=5)
return exit_code
def write_tail_log(self, exit_code: int, log_f):
log_f.write("\n" + "=" * 100 + "\n")
log_f.write(f"End Time: {self.execution.completed_at}\n")
log_f.write(f"Execution Time: {self.execution.duration_seconds:.2f} seconds\n")
log_f.write(f"Exit Code: {exit_code}\n")
log_f.write(f"Status: {'SUCCESS' if exit_code == 0 else 'FAILED'}\n")
def write_header_log(self, cmd: str, log_f):
log_f.write(f"DataX Task Execution Log\n")
log_f.write(f"Job ID: {self.execution.id}\n")
log_f.write(f"Start Time: {self.execution.started_at}\n")
log_f.write(f"Config Source: JSON String\n")
log_f.write(f"Command: {cmd}\n")
log_f.write("=" * 100 + "\n\n")
@staticmethod
def read_stream(stream, log_f):
"""读取输出流并屏蔽敏感信息"""
for line in stream:
line = line.rstrip('\n')
if line:
# Mask sensitive information before writing to log
masked_line = mask_sensitive_info(line)
# 写入日志文件
log_f.write(f"{masked_line}\n")
log_f.flush()
@@ -0,0 +1,15 @@
from fastapi import APIRouter
router = APIRouter(
prefix="/data-collection",
tags = ["data-collection"]
)
# Include sub-routers
from .collection import router as collection_router
from .execution import router as execution_router
from .template import router as template_router
router.include_router(collection_router)
router.include_router(execution_router)
router.include_router(template_router)
@@ -0,0 +1,398 @@
import json
import math
import uuid
import shutil
import os
import asyncio
from typing import Optional
from fastapi import APIRouter, Depends, Query
from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.exception import ErrorCodes, BusinessError, SuccessResponse, transaction
from app.core.logging import get_logger
from app.core.security import preserve_sensitive_values
from app.db.models import Dataset, DatasetFiles
from app.db.models.data_collection import CollectionTask, TaskExecution, CollectionTemplate
from app.db.session import get_db
from app.module.collection.client.datax_client import DataxClient
from app.module.collection.schema.collection import CollectionTaskBase, CollectionTaskCreate, CollectionTaskUpdate, CollectionConfig, converter_to_response, \
convert_for_create, SyncMode
from app.module.collection.schedule import schedule_collection_task, remove_collection_task
from app.module.collection.service.collection import CollectionTaskService
from app.module.shared.schema import StandardResponse, PaginatedData, TaskStatus
router = APIRouter(
prefix="/tasks",
tags=["data-collection/tasks"],
)
logger = get_logger(__name__)
async def is_hard_link(file_path: str) -> bool:
"""检查文件是否是硬链接"""
try:
stat_info = await asyncio.to_thread(os.stat, file_path)
# 如果链接数大于1,说明是硬链接
return stat_info.st_nlink > 1
except OSError:
return False
async def convert_hardlink_to_real_file(file_path: str) -> bool:
"""
将硬链接转换为实体文件
通过读取并重新写入文件内容,创建一个独立的副本
"""
try:
# 创建临时文件
temp_path = f"{file_path}.tmp"
# 使用 shutil.copy2 创建副本(保留元数据)
await asyncio.to_thread(shutil.copy2, file_path, temp_path)
# 删除原文件(硬链接)
await asyncio.to_thread(os.unlink, file_path)
# 重命名临时文件为原文件名
await asyncio.to_thread(os.replace, temp_path, file_path)
return True
except OSError as e:
logger.warning(f"Failed to convert hard link to real file {file_path}: {e}")
# 清理临时文件(如果存在)
if os.path.exists(f"{file_path}.tmp"):
try:
await asyncio.to_thread(os.remove, f"{file_path}.tmp")
except OSError:
pass
return False
async def convert_dataset_hardlinks_before_delete(task_id: str, db: AsyncSession) -> int:
"""
删除归集任务前,将数据集中的硬链接文件转换为实体文件
Args:
task_id: 归集任务ID
db: 数据库会话
Returns:
转换成功的文件数量
"""
try:
# 查找所有数据集文件(通过文件路径匹配任务ID)
# 注意:归集任务的源文件路径是 tmp/dataset/{task_id}/
# 我们需要找到数据集中所有以这个路径为源的文件
source_prefix = f"tmp/dataset/{task_id}/"
# 查询所有可能相关的数据集文件
result = await db.execute(
select(DatasetFiles).where(
DatasetFiles.file_path.like(f"%/dataset/%"),
DatasetFiles.status == "ACTIVE"
)
)
dataset_files = result.scalars().all()
converted_count = 0
for dataset_file in dataset_files:
file_path = dataset_file.file_path
if not file_path:
continue
# 检查文件是否是硬链接
if await is_hard_link(file_path):
logger.info(f"Converting hard link to real file: {file_path}")
success = await convert_hardlink_to_real_file(file_path)
if success:
converted_count += 1
else:
logger.warning(f"Failed to convert hard link: {file_path}")
if converted_count > 0:
logger.info(f"Converted {converted_count} hard link(s) to real file(s) for task {task_id}")
return converted_count
except Exception as e:
logger.error(f"Error converting hard links for task {task_id}: {e}", exc_info=True)
return 0
@router.post("", response_model=StandardResponse[CollectionTaskBase], operation_id="create_collect_task", tags=["mcp"])
async def create_task(
request: CollectionTaskCreate,
db: AsyncSession = Depends(get_db)
):
"""创建归集任务"""
template = await db.execute(select(CollectionTemplate).where(CollectionTemplate.id == request.template_id))
template = template.scalar_one_or_none()
if not template:
raise BusinessError(ErrorCodes.COLLECTION_TEMPLATE_NOT_FOUND, data={"template_id": request.template_id})
task_id = str(uuid.uuid4())
DataxClient.generate_datx_config(request.config, template, f"/dataset/local/{task_id}")
task = convert_for_create(request, task_id)
task.template_name = template.name
dataset = None
async with transaction(db):
if request.dataset_name:
target_dataset_id = uuid.uuid4()
dataset = Dataset(
id=str(target_dataset_id),
name=request.dataset_name,
description="",
dataset_type=request.dataset_type.name,
status="DRAFT",
path=f"/dataset/{target_dataset_id}",
)
db.add(dataset)
task_service = CollectionTaskService(db)
task = await task_service.create_task(task, dataset)
task = await db.execute(select(CollectionTask).where(CollectionTask.id == task.id))
task = task.scalar_one_or_none()
# 事务已提交,执行调度
if task and task.sync_mode == SyncMode.SCHEDULED.value and task.schedule_expression:
schedule_collection_task(task.id, task.schedule_expression)
return SuccessResponse(data=converter_to_response(task))
@router.get("", response_model=StandardResponse[PaginatedData[CollectionTaskBase]])
async def list_tasks(
page: int = 1,
size: int = 20,
name: Optional[str] = Query(None, description="Fuzzy search by task name"),
status: Optional[str] = Query(None, description="Filter by task status"),
db: AsyncSession = Depends(get_db)
):
"""分页查询归集任务"""
# 构建查询条件
page = page if page > 0 else 1
size = size if size > 0 else 20
query = select(CollectionTask)
if name:
query = query.where(CollectionTask.name.ilike(f"%{name}%"))
if status:
query = query.where(CollectionTask.status == status)
# 获取总数
count_query = select(func.count()).select_from(query.subquery())
total = (await db.execute(count_query)).scalar_one()
# 分页查询
offset = (page - 1) * size
tasks = (await db.execute(
query.order_by(CollectionTask.created_at.desc())
.offset(offset)
.limit(size)
)).scalars().all()
# 转换为响应模型
items = [converter_to_response(task) for task in tasks]
total_pages = math.ceil(total / size) if total > 0 else 0
return SuccessResponse(
data=PaginatedData(
content=items,
total_elements=total,
total_pages=total_pages,
page=page,
size=size,
)
)
@router.post("/{task_id}/execute", response_model=StandardResponse[str], operation_id="execute_task")
async def execute_task(
task_id: str,
db: AsyncSession = Depends(get_db)
):
"""执行归集任务"""
from app.module.collection.service.collection import CollectionTaskService
# 验证任务是否存在
task = await db.execute(select(CollectionTask).where(CollectionTask.id == task_id))
task = task.scalar_one_or_none()
if not task:
raise BusinessError(ErrorCodes.COLLECTION_TASK_NOT_FOUND, data={"task_id": task_id})
# 调用服务执行任务
await CollectionTaskService.run_async(task_id)
return SuccessResponse(
data=task_id
)
@router.put("/{task_id}", response_model=StandardResponse[CollectionTaskBase], operation_id="update_task")
async def update_task(
task_id: str,
request: CollectionTaskUpdate,
db: AsyncSession = Depends(get_db)
):
"""
更新归集任务
支持修改的字段:
- description: 任务描述
- schedule_expression: 定时表达式(仅 SCHEDULED 模式)
- timeout_seconds: 超时时间
- config: 任务配置
状态限制:
- DRAFT/PENDING: 可以修改所有字段
- RUNNING: 不允许修改
- FAILED: 可以修改,修改后可重新执行
- COMPLETED: ONCE 模式不允许修改,SCHEDULED 模式可以修改
Args:
task_id: 任务ID
request: 更新请求
db: 数据库会话
Returns:
更新后的任务信息
"""
# 查询任务
task = await db.get(CollectionTask, task_id)
if not task:
raise BusinessError(ErrorCodes.COLLECTION_TASK_NOT_FOUND, data={"task_id": task_id})
# 检查任务状态,RUNNING 状态不允许修改
if task.status == TaskStatus.RUNNING.name:
raise BusinessError(
ErrorCodes.VALIDATION_ERROR,
data={"task_id": task_id, "current_status": task.status}
)
# 检查任务同步模式,ONCE 模式已完成的不允许修改
if task.sync_mode == SyncMode.ONCE.name and task.status == TaskStatus.COMPLETED.name:
raise BusinessError(
ErrorCodes.VALIDATION_ERROR,
data={"task_id": task_id, "current_status": task.status}
)
# 验证 schedule_expression 只能由 SCHEDULED 模式修改
if request.schedule_expression is not None and task.sync_mode != SyncMode.SCHEDULED.name:
raise BusinessError(
ErrorCodes.VALIDATION_ERROR,
data={"sync_mode": task.sync_mode}
)
# 应用更新
update_data = request.model_dump(exclude_unset=True)
if 'description' in update_data:
task.description = update_data['description']
if 'timeout_seconds' in update_data:
task.timeout_seconds = update_data['timeout_seconds']
if 'schedule_expression' in update_data:
old_schedule_expression = task.schedule_expression
task.schedule_expression = update_data['schedule_expression']
# 如果调度表达式发生变化,需要重新调度任务
if old_schedule_expression != task.schedule_expression:
from app.module.collection.schedule import reschedule_collection_task
reschedule_collection_task(task_id, task.schedule_expression)
if 'config' in update_data:
# Get original config from database
original_config = json.loads(task.config) if task.config else {}
# Preserve sensitive values if masked pattern detected
preserved_config = preserve_sensitive_values(
request.config.dict(),
original_config
)
# Regenerate task config file with preserved sensitive values
template = await db.execute(select(CollectionTemplate).where(CollectionTemplate.id == task.template_id))
template = template.scalar_one_or_none()
if template:
# Use preserved config to generate DataX config
config_obj = CollectionConfig(**preserved_config)
DataxClient.generate_datx_config(
config_obj,
template,
task.target_path
)
# Save the modified config (with regenerated job) to database
task.config = json.dumps(config_obj.model_dump())
# 如果任务处于 FAILED 状态,修改后重置为 PENDING,允许重新执行
if task.status == TaskStatus.FAILED.name:
task.status = TaskStatus.PENDING.name
await db.commit()
await db.refresh(task)
return SuccessResponse(
data=converter_to_response(task)
)
@router.delete("", response_model=StandardResponse[str], status_code=200)
async def delete_collection_tasks(
ids: list[str] = Query(..., description="List of task IDs to delete"),
db: AsyncSession = Depends(get_db),
):
"""
删除归集任务
Args:
ids: 任务ID
db: 数据库会话
Returns:
StandardResponse[str]: 删除结果
"""
# 检查任务是否存在
task_id = ids[0]
task = await db.get(CollectionTask, task_id)
if not task:
raise BusinessError(ErrorCodes.COLLECTION_TASK_NOT_FOUND, data={"task_id": task_id})
# 删除任务执行记录(在事务内)
async with transaction(db):
await db.execute(
TaskExecution.__table__.delete()
.where(TaskExecution.task_id == task_id)
)
# 删除任务
await db.delete(task)
# 事务提交后,先转换硬链接,再删除文件系统和调度
logger.info(f"Converting hard links before deleting task {task_id}")
await convert_dataset_hardlinks_before_delete(task_id, db)
remove_collection_task(task_id)
target_path = f"/dataset/local/{task_id}"
if os.path.exists(target_path):
shutil.rmtree(target_path)
job_path = f"/flow/data-collection/{task_id}"
if os.path.exists(job_path):
shutil.rmtree(job_path)
return SuccessResponse(
data="success"
)
@router.get("/{task_id}", response_model=StandardResponse[CollectionTaskBase])
async def get_task(
task_id: str,
db: AsyncSession = Depends(get_db)
):
"""获取归集任务详情"""
# 根据ID查询任务
task = await db.get(CollectionTask, task_id)
if not task:
raise BusinessError(ErrorCodes.COLLECTION_TASK_NOT_FOUND, data={"task_id": task_id})
return SuccessResponse(data=converter_to_response(task))
@@ -0,0 +1,124 @@
import math
import os
from pathlib import Path
from datetime import datetime
from typing import Optional
from fastapi import APIRouter, Depends, HTTPException, Query
from fastapi.responses import FileResponse
from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.logging import get_logger
from app.db.models.data_collection import TaskExecution
from app.db.session import get_db
from app.module.collection.schema.collection import TaskExecutionBase, converter_execution_to_response
from app.module.shared.schema import StandardResponse, PaginatedData
router = APIRouter(
prefix="/executions",
tags=["data-collection/executions"],
)
logger = get_logger(__name__)
@router.get("", response_model=StandardResponse[PaginatedData[TaskExecutionBase]])
async def list_executions(
page: int = 1,
size: int = 20,
task_id: Optional[str] = Query(None, description="Task ID"),
task_name: Optional[str] = Query(None, description="Fuzzy search by task name"),
status: Optional[str] = Query(None, description="Filter by execution status"),
start_time: Optional[datetime] = Query(None, description="Start time range from (started_at >= start_time)"),
end_time: Optional[datetime] = Query(None, description="Start time range to (started_at <= end_time)"),
db: AsyncSession = Depends(get_db)
):
"""分页查询归集任务执行记录"""
try:
query = select(TaskExecution)
if task_id:
query = query.where(TaskExecution.task_id == task_id)
if task_name:
query = query.where(TaskExecution.task_name.ilike(f"%{task_name}%"))
if status:
query = query.where(TaskExecution.status == status)
if start_time:
query = query.where(TaskExecution.started_at >= start_time)
if end_time:
query = query.where(TaskExecution.started_at <= end_time)
count_query = select(func.count()).select_from(query.subquery())
total = (await db.execute(count_query)).scalar_one()
offset = (page - 1) * size
executions = (await db.execute(
query.order_by(TaskExecution.created_at.desc())
.offset(offset)
.limit(size)
)).scalars().all()
items = [converter_execution_to_response(exe) for exe in executions]
total_pages = math.ceil(total / size) if total > 0 else 0
return StandardResponse(
code="0",
message="Success",
data=PaginatedData(
content=items,
total_elements=total,
total_pages=total_pages,
page=page,
size=size,
)
)
except Exception as e:
logger.error(f"Failed to list task executions: {str(e)}", e)
raise HTTPException(status_code=500, detail="Internal server error")
@router.get("/{execution_id}/log")
async def get_execution_log(
execution_id: str,
db: AsyncSession = Depends(get_db)
):
"""Get log file content for execution record"""
try:
execution = await db.get(TaskExecution, execution_id)
if not execution:
raise HTTPException(status_code=404, detail="Execution record not found")
log_path = getattr(execution, "log_path", None)
if not log_path:
raise HTTPException(status_code=404, detail="Log path not found")
path = Path(str(log_path))
if not path.is_absolute():
path = Path(os.getcwd()) / path
path = path.resolve()
if not path.exists() or not path.is_file():
raise HTTPException(status_code=404, detail="Log file not found")
filename = path.name
headers = {
"Content-Disposition": f'inline; filename={filename}'
}
return FileResponse(
path=str(path),
media_type="text/plain; charset=utf-8",
filename=filename,
headers=headers,
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Failed to get execution log: {str(e)}", e)
raise HTTPException(status_code=500, detail="Internal server error")
@@ -0,0 +1,67 @@
import math
from typing import Optional
from fastapi import APIRouter, Depends, HTTPException, Query
from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.logging import get_logger
from app.db.models.data_collection import CollectionTemplate
from app.db.session import get_db
from app.module.collection.schema.collection import CollectionTemplateBase, converter_template_to_response
from app.module.shared.schema import StandardResponse, PaginatedData
router = APIRouter(
prefix="/templates",
tags=["data-collection/templates"],
)
logger = get_logger(__name__)
@router.get("", response_model=StandardResponse[PaginatedData[CollectionTemplateBase]])
async def list_templates(
page: int = 1,
size: int = 20,
name: Optional[str] = Query(None, description="Fuzzy search by template name"),
built_in: Optional[bool] = Query(None, description="Filter by system built-in template"),
db: AsyncSession = Depends(get_db)
):
"""分页查询归集任务模板"""
try:
query = select(CollectionTemplate)
if name:
query = query.where(CollectionTemplate.name.ilike(f"%{name}%"))
if built_in is not None:
query = query.where(CollectionTemplate.built_in == built_in)
count_query = select(func.count()).select_from(query.subquery())
total = (await db.execute(count_query)).scalar_one()
offset = (page - 1) * size
templates = (await db.execute(
query.order_by(CollectionTemplate.created_at.desc())
.offset(offset)
.limit(size)
)).scalars().all()
items = [converter_template_to_response(tpl) for tpl in templates]
total_pages = math.ceil(total / size) if total > 0 else 0
return StandardResponse(
code="0",
message="Success",
data=PaginatedData(
content=items,
total_elements=total,
total_pages=total_pages,
page=page,
size=size,
)
)
except Exception as e:
logger.error(f"Failed to list collection templates: {str(e)}", e)
raise HTTPException(status_code=500, detail="Internal server error")
@@ -0,0 +1,79 @@
from __future__ import annotations
from typing import Optional
from sqlalchemy import select
from app.core.logging import get_logger
from app.db.models.data_collection import CollectionTask
from app.db.session import AsyncSessionLocal
from app.module.collection.schema.collection import SyncMode
from app.module.shared.schedule import Scheduler
logger = get_logger(__name__)
_scheduler: Optional[Scheduler] = None
def set_collection_scheduler(scheduler: Scheduler) -> None:
global _scheduler
_scheduler = scheduler
def _get_scheduler() -> Scheduler:
if _scheduler is None:
raise RuntimeError("Collection scheduler not initialized")
return _scheduler
def schedule_collection_task(task_id: str, schedule_expression: str, dataset_id: Optional[str] = None) -> None:
scheduler = _get_scheduler()
from app.module.collection.service.collection import CollectionTaskService
scheduler.add_cron_job(
job_id=f"collection:{task_id}",
cron_expression=schedule_expression,
func=CollectionTaskService.run_async,
args=[task_id, dataset_id],
replace_existing=True,
max_instances=1,
coalesce=True,
misfire_grace_time=60,
)
logger.info(f"Scheduled collection task {task_id} with cron {schedule_expression}")
def remove_collection_task(task_id: str) -> None:
scheduler = _get_scheduler()
job_id = f"collection:{task_id}"
if scheduler.has_job(job_id):
scheduler.remove_job(job_id)
logger.info(f"Removed scheduled collection task {task_id}")
def reschedule_collection_task(task_id: str, schedule_expression: str, dataset_id: Optional[str] = None) -> None:
remove_collection_task(task_id)
schedule_collection_task(task_id, schedule_expression, dataset_id)
def validate_schedule_expression(schedule_expression: str) -> None:
Scheduler.validate_cron_expression(schedule_expression)
async def load_scheduled_collection_tasks() -> None:
async with AsyncSessionLocal() as session:
result = await session.execute(
select(CollectionTask).where(
CollectionTask.sync_mode == SyncMode.SCHEDULED.value,
CollectionTask.schedule_expression.isnot(None),
)
)
tasks = result.scalars().all()
for task in tasks:
if not task.schedule_expression:
continue
try:
schedule_collection_task(task.id, task.schedule_expression)
except Exception as exc:
logger.error(f"Failed to schedule task {task.id}: {exc}")
@@ -0,0 +1,226 @@
import json
import uuid
from datetime import datetime
from enum import Enum
from typing import Optional, Dict, Any
from pydantic import BaseModel, Field, validator, ConfigDict, field_validator, field_serializer
from pydantic.alias_generators import to_camel
from app.core.security import mask_sensitive_dict
from app.db.models.data_collection import CollectionTask, TaskExecution, CollectionTemplate
from app.module.dataset.schema import DatasetTypeResponse
from app.module.dataset.schema.dataset import DatasetType
from app.module.shared.schema import TaskStatus
class SyncMode(str, Enum):
ONCE = "ONCE"
SCHEDULED = "SCHEDULED"
class CollectionConfig(BaseModel):
parameter: Optional[dict] = Field(None, description="模板参数")
reader: Optional[dict] = Field(None, description="reader参数")
writer: Optional[dict] = Field(None, description="writer参数")
job: Optional[dict] = Field(None, description="任务配置")
@field_serializer('parameter', 'reader', 'writer', 'job', when_used='json')
def mask_sensitive_fields(self, value: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
"""Mask sensitive fields when serializing to JSON"""
if value is None:
return value
return mask_sensitive_dict(value)
class CollectionTaskBase(BaseModel):
id: str = Field(..., description="任务id")
name: str = Field(..., description="任务名称")
description: Optional[str] = Field(None, description="任务描述")
target_path: str = Field(..., description="目标存放路径")
config: CollectionConfig = Field(..., description="任务配置")
template_id: str = Field(..., description="模板ID")
template_name: Optional[str] = Field(None, description="模板名称")
status: TaskStatus = Field(..., description="任务状态")
sync_mode: SyncMode = Field(default=SyncMode.ONCE, description="同步方式")
schedule_expression: Optional[str] = Field(None, description="调度表达式(cron)")
retry_count: int = Field(default=3, description="重试次数")
timeout_seconds: int = Field(default=3600, description="超时时间")
last_execution_id: Optional[str] = Field(None, description="最后执行id")
created_at: Optional[datetime] = Field(None, description="创建时间")
updated_at: Optional[datetime] = Field(None, description="更新时间")
created_by: Optional[str] = Field(None, description="创建人")
updated_by: Optional[str] = Field(None, description="更新人")
model_config = ConfigDict(
alias_generator=to_camel,
populate_by_name=True
)
class CollectionTaskCreate(BaseModel):
name: str = Field(..., description="任务名称")
description: Optional[str] = Field(None, description="任务描述")
sync_mode: SyncMode = Field(default=SyncMode.ONCE, description="同步方式")
schedule_expression: Optional[str] = Field(None, description="调度表达式(cron)")
config: CollectionConfig = Field(..., description="任务配置")
template_id: str = Field(..., description="模板ID")
dataset_name: Optional[str] = Field(None, description="数据集名称")
dataset_type: Optional[DatasetType] = Field(DatasetType.TEXT, description="数据集类型")
model_config = ConfigDict(
alias_generator=to_camel,
populate_by_name=True
)
class CollectionTaskUpdate(BaseModel):
"""任务更新模型 - 只包含可更新的字段"""
description: Optional[str] = Field(None, description="任务描述")
schedule_expression: Optional[str] = Field(None, description="调度表达式(cron)- 仅SCHEDULED模式可修改")
timeout_seconds: Optional[int] = Field(None, description="超时时间(秒)")
config: Optional[CollectionConfig] = Field(None, description="任务配置")
@field_validator('schedule_expression')
@classmethod
def validate_schedule_expression(cls, v, info):
"""验证调度表达式"""
if v is not None:
# SCHEDULED 模式必须有调度表达式
if not v or not v.strip():
raise ValueError("schedule_expression is required for SCHEDULED mode")
return v
@field_validator('timeout_seconds')
@classmethod
def validate_timeout(cls, v):
"""验证超时时间"""
if v is not None and v <= 0:
raise ValueError("timeout_seconds must be greater than 0")
return v
model_config = ConfigDict(
alias_generator=to_camel,
populate_by_name=True
)
def converter_to_response(task: CollectionTask) -> CollectionTaskBase:
config_dict = json.loads(task.config)
return CollectionTaskBase(
id=task.id,
name=task.name,
description=task.description,
sync_mode=task.sync_mode,
template_id=task.template_id,
template_name=task.template_name,
target_path=task.target_path,
config=CollectionConfig(**config_dict),
schedule_expression=task.schedule_expression,
status=task.status,
retry_count=task.retry_count,
timeout_seconds=task.timeout_seconds,
last_execution_id=task.last_execution_id,
created_at=task.created_at,
updated_at=task.updated_at,
created_by=task.created_by,
updated_by=task.updated_by,
)
def convert_for_create(task: CollectionTaskCreate, task_id: str) -> CollectionTask:
schedule_expression = task.schedule_expression if task.sync_mode == SyncMode.SCHEDULED else None
return CollectionTask(
id=task_id,
name=task.name,
description=task.description,
sync_mode=task.sync_mode,
template_id=task.template_id,
target_path=f"/dataset/local/{task_id}",
config=json.dumps(task.config.dict()),
schedule_expression=schedule_expression,
status=TaskStatus.PENDING.name
)
def create_execute_record(task: CollectionTask) -> TaskExecution:
execution_id = str(uuid.uuid4())
return TaskExecution(
id=execution_id,
task_id=task.id,
task_name=task.name,
status=TaskStatus.RUNNING.name,
started_at=datetime.now(),
log_path=f"/flow/data-collection/{task.id}/{execution_id}.log"
)
class TaskExecutionBase(BaseModel):
id: str = Field(..., description="执行记录ID")
task_id: str = Field(..., description="任务ID")
task_name: str = Field(..., description="任务名称")
status: Optional[str] = Field(None, description="执行状态")
log_path: Optional[str] = Field(None, description="日志文件路径")
started_at: Optional[datetime] = Field(None, description="开始时间")
completed_at: Optional[datetime] = Field(None, description="完成时间")
duration_seconds: Optional[int] = Field(None, description="执行时长(秒)")
error_message: Optional[str] = Field(None, description="错误信息")
created_at: Optional[datetime] = Field(None, description="创建时间")
updated_at: Optional[datetime] = Field(None, description="更新时间")
created_by: Optional[str] = Field(None, description="创建者")
updated_by: Optional[str] = Field(None, description="更新者")
model_config = ConfigDict(
alias_generator=to_camel,
populate_by_name=True
)
def converter_execution_to_response(execution: TaskExecution) -> TaskExecutionBase:
return TaskExecutionBase(
id=execution.id,
task_id=execution.task_id,
task_name=execution.task_name,
status=execution.status,
log_path=execution.log_path,
started_at=execution.started_at,
completed_at=execution.completed_at,
duration_seconds=execution.duration_seconds,
error_message=execution.error_message,
created_at=execution.created_at,
updated_at=execution.updated_at,
created_by=execution.created_by,
updated_by=execution.updated_by,
)
class CollectionTemplateBase(BaseModel):
id: str = Field(..., description="模板ID")
name: str = Field(..., description="模板名称")
description: Optional[str] = Field(None, description="模板描述")
source_type: str = Field(..., description="源数据源类型")
source_name: str = Field(..., description="源数据源名称")
target_type: str = Field(..., description="目标数据源类型")
target_name: str = Field(..., description="目标数据源名称")
template_content: dict = Field(..., description="模板内容")
built_in: Optional[bool] = Field(None, description="是否系统内置模板")
created_at: Optional[datetime] = Field(None, description="创建时间")
updated_at: Optional[datetime] = Field(None, description="更新时间")
created_by: Optional[str] = Field(None, description="创建者")
updated_by: Optional[str] = Field(None, description="更新者")
model_config = ConfigDict(
alias_generator=to_camel,
populate_by_name=True
)
def converter_template_to_response(template: CollectionTemplate) -> CollectionTemplateBase:
return CollectionTemplateBase(
id=template.id,
name=template.name,
description=template.description,
source_type=template.source_type,
source_name=template.source_name,
target_type=template.target_type,
target_name=template.target_name,
template_content=template.template_content,
built_in=template.built_in,
created_at=template.created_at,
updated_at=template.updated_at,
created_by=template.created_by,
updated_by=template.updated_by,
)
@@ -0,0 +1,281 @@
import os
import yaml
import asyncio
from pathlib import Path
from typing import Optional
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.logging import get_logger
from app.db.models import Dataset
# ✅ 确保引入 TaskExecution
from app.db.models.data_collection import CollectionTask, CollectionTemplate, TaskExecution
from app.db.session import AsyncSessionLocal
from app.module.collection.client.datax_client import DataxClient
from app.module.collection.schema.collection import SyncMode, create_execute_record
from app.module.collection.schedule import validate_schedule_expression
from app.module.dataset.service.service import Service
from app.module.shared.schema import TaskStatus, NodeType, EdgeType
from app.module.shared.common.lineage import LineageService
from app.db.models.base_entity import LineageNode, LineageEdge
# ✅ 确保引入 RuntimeClient
from app.module.shared.runtime_client import RuntimeClient
logger = get_logger(__name__)
# 执行模式开关:'local' (默认) | 'runtime'
COLLECTION_EXECUTION_MODE = os.getenv("COLLECTION_EXECUTION_MODE", "local")
class _TempExecution:
"""用于生成配置时的临时占位符"""
def __init__(self, task_id):
self.id = task_id
class CollectionTaskService:
def __init__(self, db: AsyncSession):
self.db = db
async def create_task(self, task: CollectionTask, dataset: Dataset) -> CollectionTask:
self.db.add(task)
if task.sync_mode == SyncMode.ONCE:
task.status = TaskStatus.RUNNING.name
await self.db.commit()
asyncio.create_task(CollectionTaskService.run_async(task.id, dataset.id if dataset else None))
elif task.sync_mode == SyncMode.SCHEDULED:
if not task.schedule_expression:
raise ValueError("schedule_expression is required for scheduled tasks")
validate_schedule_expression(task.schedule_expression)
return task
@staticmethod
async def run_async(task_id: str, dataset_id: str = None):
"""任务执行入口"""
logger.info(f"Start to execute task {task_id}")
# 1. 获取任务基础信息
async with AsyncSessionLocal() as session:
task_result = await session.execute(select(CollectionTask).where(CollectionTask.id == task_id))
task = task_result.scalar_one_or_none()
if not task:
logger.error(f"task {task_id} not exist")
return
template_result = await session.execute(select(CollectionTemplate).where(CollectionTemplate.id == task.template_id))
template = template_result.scalar_one_or_none()
if not template:
logger.error(f"template {task.template_name} not exist")
return
# 2. 创建执行记录 (此时 task_execution.log_path 已经根据 schema 规则被设置好)
task_execution = create_execute_record(task)
session.add(task_execution)
await session.commit()
await session.refresh(task_execution)
exec_id = task_execution.id
# 3. 根据模式分发
try:
if COLLECTION_EXECUTION_MODE == "runtime":
# ✅ 关键修改:将 task_execution 对象传入,以便提取 execution_id
await CollectionTaskService._execute_via_runtime(
task_id, dataset_id, task, template, task_execution
)
# 提交成功后更新状态为 SUBMITTED
async with AsyncSessionLocal() as update_session:
exec_record = await update_session.get(TaskExecution, exec_id)
if exec_record:
exec_record.status = "SUBMITTED"
await update_session.commit()
logger.info(f"Task {task_id} submitted to Runtime. Starting background poller.")
# 启动后台轮询
asyncio.create_task(
CollectionTaskService._poll_runtime_status(task_id, dataset_id, exec_id)
)
else:
# Local 模式
await CollectionTaskService._execute_locally(task_id, dataset_id, task, template, exec_id)
except Exception as e:
logger.error(f"Task execution failed for {task_id}: {e}")
async with AsyncSessionLocal() as fail_session:
exec_record = await fail_session.get(TaskExecution, exec_id)
if exec_record:
exec_record.status = TaskStatus.FAILED.name
exec_record.error_message = str(e)
await fail_session.commit()
# ==========================================
# 🟢 模式 A:本地执行逻辑
# ==========================================
@staticmethod
async def _execute_locally(task_id, dataset_id, task, template, exec_id):
"""本地直接运行 DataX"""
logger.info(f"[Local Mode] Executing DataX for task {task_id}")
async with AsyncSessionLocal() as session:
exec_record = await session.get(TaskExecution, exec_id)
# DataX 运行是阻塞的,放在线程池中运行
datax_client = DataxClient(task=task, execution=exec_record, template=template)
await asyncio.to_thread(datax_client.run_datax_job)
# 同步 Task 状态
task.status = exec_record.status
session.add(task)
if exec_record.status == TaskStatus.COMPLETED.name:
await CollectionTaskService._finalize_task(session, task_id, dataset_id, task)
await session.commit()
# ==========================================
# 🔵 模式 B:Runtime -> Ray 执行逻辑
# ==========================================
@staticmethod
async def _execute_via_runtime(task_id, dataset_id, task, template, task_execution):
"""
生成配置并提交给 Runtime。
不需要 session,因为 task_execution 对象已包含所有必要信息。
"""
logger.info(f"[Runtime Mode] Submitting task {task_id} to deepdata-runtime")
# ✅ 关键步骤:提取 execution_id,这是日志文件名的关键
execution_id = task_execution.id
log_path = task_execution.log_path
# 1. 生成 DataX JSON 配置
datax_client = DataxClient(task=task, execution=_TempExecution(task_id), template=template)
datax_json_path = datax_client.generate_config()
# 2. 构造 process.yaml 指令文件
flow_dir = Path(f"/flow/{task_id}")
flow_dir.mkdir(parents=True, exist_ok=True)
process_cfg = {
"executor_type": "deepdata",
"instance_id": task_id,
"execution_id": execution_id,
"log_path": log_path,
"dataset_id": dataset_id,
"cli_config": {
"command": "python",
"args": ["/opt/datax/bin/datax.py", datax_json_path],
"timeout": 7200
}
}
# 写入指令文件
with open(flow_dir / "process.yaml", "w", encoding="utf-8") as f:
yaml.safe_dump(process_cfg, f)
# 3. 提交给 Runtime
runtime_client = RuntimeClient()
success = await runtime_client.submit_task(task_id)
if not success:
raise RuntimeError("Runtime submission returned False")
@staticmethod
async def _poll_runtime_status(task_id, dataset_id, exec_id):
"""后台轮询 Runtime 状态"""
runtime_client = RuntimeClient()
for _ in range(2000):
try:
status_info = await runtime_client.get_task_status(task_id)
if status_info:
current_status = (status_info.get("status") or "").upper()
async with AsyncSessionLocal() as session:
exec_record = await session.get(TaskExecution, exec_id)
if not exec_record: return
task_result = await session.execute(select(CollectionTask).where(CollectionTask.id == task_id))
task = task_result.scalar_one_or_none()
if current_status in ["COMPLETED", "SUCCEEDED", "SUCCESS"]:
exec_record.status = TaskStatus.COMPLETED.name
await session.commit()
logger.info(f"Task {task_id} completed in Runtime. Finalizing.")
if task:
await CollectionTaskService._finalize_task(session, task_id, dataset_id, task)
return
elif current_status in ["FAILED", "ERROR"]:
exec_record.status = TaskStatus.FAILED.name
exec_record.error_message = status_info.get("error", "Runtime task failed")
await session.commit()
logger.error(f"Task {task_id} failed in Runtime: {exec_record.error_message}")
return
await asyncio.sleep(5)
except Exception as e:
logger.warning(f"Polling error for {task_id}: {e}")
await asyncio.sleep(5)
# ==========================================
# 🛠️ 公共收尾逻辑
# ==========================================
@staticmethod
async def _finalize_task(session, task_id, dataset_id, task):
"""任务执行成功后的后续处理"""
if not dataset_id: return
# 🐛 修复:对于 S3/远端写入器,本地 target_path 可能不存在
# 应当检查 task.target_path 是否为本地路径,或者跳过本地扫描逻辑
target_path = Path(task.target_path)
if target_path.exists() and target_path.is_dir():
service = Service(db=session)
source_paths = []
for file_path in target_path.rglob('*'):
if file_path.is_file():
source_paths.append(str(file_path.absolute()))
if source_paths:
await service.add_files_to_dataset(dataset_id=dataset_id, source_paths=source_paths)
await CollectionTaskService._add_dataset_to_graph(session, dataset_id, task)
else:
logger.warning(f"[Finalize] Target path {target_path} does not exist locally. "
f"Skipping dataset import. (This is expected for s3writer/nfswriter)")
@staticmethod
async def _add_dataset_to_graph(session, dataset_id, task):
"""将数据集加入血缘图"""
try:
dataset = await session.get(Dataset, dataset_id)
if not dataset:
logger.warning(f"dataset {dataset_id} not found when building lineage graph")
return
dataset_node = LineageNode(
id=dataset.id,
node_type=NodeType.DATASET.value,
name=dataset.name,
description=dataset.description
)
collection_node = LineageNode(
id=task.id,
node_type=NodeType.DATASOURCE.value,
name=task.name,
description=task.description
)
collection_edge = LineageEdge(
process_id=task.id,
name=task.name,
edge_type=EdgeType.DATA_COLLECTION.value,
description=dataset.description,
from_node_id=task.id,
to_node_id=dataset.id
)
lineage_service = LineageService(db=session)
await lineage_service.generate_graph(collection_node, collection_edge, dataset_node)
await session.commit()
except Exception as exc:
logger.error(f"Failed to add dataset lineage graph: {exc}")
await session.rollback()
@@ -0,0 +1,3 @@
from .service import DatasetManagementService
__all__ = ["DatasetManagementService"]
@@ -0,0 +1,26 @@
from .dataset_file import (
DatasetFileResponse,
PagedDatasetFileResponse,
BatchUpdateFileTagsRequest,
BatchUpdateFileTagsResponse,
FileTagUpdateResult,
FileTagUpdate,
)
from .dataset import (
DatasetResponse,
DatasetTypeResponse,
CreateDatasetRequest,
)
__all__ = [
"DatasetResponse",
"DatasetFileResponse",
"PagedDatasetFileResponse",
"DatasetTypeResponse",
"BatchUpdateFileTagsRequest",
"BatchUpdateFileTagsResponse",
"FileTagUpdateResult",
"FileTagUpdate",
"CreateDatasetRequest",
]

Some files were not shown because too many files have changed in this diff Show More