deepdata: core product code (backend, frontend, runtime, deepflow module)
This commit is contained in:
@@ -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
|
||||
@@ -0,0 +1,8 @@
|
||||
# Local Development Environment Files
|
||||
.env
|
||||
.dev.env
|
||||
|
||||
# logs
|
||||
logs/
|
||||
|
||||
doc/
|
||||
@@ -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
|
||||
|
||||
|
||||
"""
|
||||
@@ -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}")
|
||||
@@ -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})>"
|
||||
@@ -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()
|
||||
@@ -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(),
|
||||
)
|
||||
@@ -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
|
||||
+95
@@ -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
|
||||
+56
@@ -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
Reference in New Issue
Block a user