from datetime import datetime from typing import List, Optional, Dict, Any, Tuple import random import json import os import shutil import asyncio from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.core.logging import get_logger from app.core.config import settings from app.db.models.base_entity import LineageNode, LineageEdge from app.db.models.ratio_task import RatioInstance, RatioRelation from app.db.models import Dataset, DatasetFiles from app.db.session import AsyncSessionLocal from app.module.dataset.schema.dataset_file import DatasetFileTag from app.module.shared.common.lineage import LineageService from app.module.shared.schema import TaskStatus, NodeType, EdgeType from app.module.ratio.schema.ratio_task import FilterCondition from app.module.shared.util.resource_utils import get_concurrent_for_ratio_copy, get_system_resource_info logger = get_logger(__name__) class RatioTaskService: """Service for Ratio Task DB operations.""" def __init__(self, db: AsyncSession): self.db = db async def create_task( self, *, name: str, description: Optional[str], totals: int, config: List[Dict[str, Any]], target_dataset_id: Optional[str] = None, ) -> RatioInstance: """Create a ratio task instance and its relations. config item format: {"dataset_id": str, "counts": int, "filter_conditions": str} """ logger.info(f"Creating ratio task: name={name}, totals={totals}, items={len(config or [])}") instance = RatioInstance( name=name, description=description, totals=totals, target_dataset_id=target_dataset_id, status="PENDING", ) self.db.add(instance) await self.db.flush() # populate instance.id for item in config or []: relation = RatioRelation( ratio_instance_id=instance.id, source_dataset_id=item.get("dataset_id"), counts=int(item.get("counts", 0)), filter_conditions=json.dumps({ 'date_range': item.get("filter_conditions").date_range, 'label': { "label":item.get("filter_conditions").label.label, "value":item.get("filter_conditions").label.value, } if item.get("filter_conditions").label else None, }) ) logger.info(f"Relation created: {relation.id}, {relation}, {item}, {config}") self.db.add(relation) await self.db.commit() await self.db.refresh(instance) logger.info(f"Ratio task created: {instance.id}") return instance # ========================= Execution (Background) ========================= # @staticmethod async def execute_dataset_ratio_task(instance_id: str) -> None: async with AsyncSessionLocal() as session: try: inst_res = await session.execute(select(RatioInstance).where(RatioInstance.id == instance_id)) instance: Optional[RatioInstance] = inst_res.scalar_one_or_none() if not instance: logger.error(f"Ratio instance not found: {instance_id}") return logger.info(f"Starting ratio task {instance_id}") logger.info(f"System resources: {get_system_resource_info()}") rel_res = await session.execute( select(RatioRelation).where(RatioRelation.ratio_instance_id == instance_id) ) relations: List[RatioRelation] = list(rel_res.scalars().all()) instance.status = TaskStatus.RUNNING.name ds_res = await session.execute(select(Dataset).where(Dataset.id == instance.target_dataset_id)) target_ds: Optional[Dataset] = ds_res.scalar_one_or_none() if not target_ds: logger.error(f"Target dataset not found for instance {instance_id}") instance.status = TaskStatus.FAILED.name return max_concurrent = get_concurrent_for_ratio_copy(settings) logger.info(f"Using {max_concurrent} concurrent workers for ratio task {instance_id}") added_count, added_size = await RatioTaskService.handle_ratio_relations_parallel( relations, session, target_ds, max_concurrent ) target_ds.file_count = (target_ds.file_count or 0) + added_count target_ds.size_bytes = (target_ds.size_bytes or 0) + added_size if (target_ds.file_count or 0) > 0: target_ds.status = "ACTIVE" instance.status = TaskStatus.COMPLETED.name logger.info(f"Ratio task completed: {instance_id}, files={added_count}, size={added_size}") await RatioTaskService._add_task_to_graph( session=session, src_relations=relations, task=instance, dst_dataset=target_ds, ) except Exception as e: logger.exception(f"Ratio task failed for {instance_id}: {e}") try: inst_res = await session.execute(select(RatioInstance).where(RatioInstance.id == instance_id)) instance = inst_res.scalar_one_or_none() if instance: instance.status = TaskStatus.FAILED.name finally: pass finally: await session.commit() @staticmethod async def handle_ratio_relations_parallel( relations: list[RatioRelation], session: AsyncSession, target_ds: Dataset, max_concurrent: int = 10 ) -> Tuple[int, int]: existing_path_rows = await session.execute( select(DatasetFiles.file_path).where(DatasetFiles.dataset_id == target_ds.id) ) existing_paths = set(p for p in existing_path_rows.scalars().all() if p) source_paths = set() all_copy_tasks: List[Tuple[DatasetFiles, str, str]] = [] for rel in relations: if not rel.source_dataset_id or not rel.counts or rel.counts <= 0: continue files = await RatioTaskService.get_files(rel, session) if not files: continue pick_n = min(rel.counts or 0, len(files)) chosen = random.sample(files, pick_n) if pick_n < len(files) else files for file in chosen: if file.file_path in source_paths: continue dst_prefix = f"/dataset/{target_ds.id}/" file_name = RatioTaskService.get_new_file_name(dst_prefix, existing_paths, file) new_path = dst_prefix + file_name file_record = DatasetFiles( dataset_id=target_ds.id, file_name=file_name, file_path=new_path, file_type=file.file_type, file_size=file.file_size, check_sum=file.check_sum, tags=file.tags, tags_updated_at=datetime.now(), dataset_filemetadata=file.dataset_filemetadata, status="ACTIVE", ) all_copy_tasks.append((file_record, file.file_path, new_path)) existing_paths.add(new_path) source_paths.add(file.file_path) if not all_copy_tasks: return 0, 0 dst_dir = f"/dataset/{target_ds.id}/" await asyncio.to_thread(os.makedirs, dst_dir, exist_ok=True) semaphore = asyncio.Semaphore(max_concurrent) successful_records: List[DatasetFiles] = [] added_count = 0 added_size = 0 async def copy_with_semaphore( file_record: DatasetFiles, src_path: str, dst_path: str ) -> Tuple[bool, DatasetFiles]: async with semaphore: try: file_dst_dir = os.path.dirname(dst_path) if file_dst_dir != dst_dir: await asyncio.to_thread(os.makedirs, file_dst_dir, exist_ok=True) try: await asyncio.to_thread(os.link, src_path, dst_path) except OSError: try: await asyncio.to_thread(os.symlink, src_path, dst_path) except OSError: await asyncio.to_thread(shutil.copy2, src_path, dst_path) return True, file_record except Exception as e: logger.error(f"Copy failed: {src_path} -> {dst_path}: {e}") return False, file_record tasks = [copy_with_semaphore(rec, src, dst) for rec, src, dst in all_copy_tasks] results = await asyncio.gather(*tasks, return_exceptions=True) for i, result in enumerate(results): if isinstance(result, Exception): logger.error(f"Copy task {i} raised exception: {result}") continue success, file_record = result if success: added_count += 1 added_size += int(file_record.file_size or 0) successful_records.append(file_record) if successful_records: session.add_all(successful_records) await session.flush() logger.info( f"Parallel copy completed: {added_count}/{len(all_copy_tasks)} files, " f"{added_size} bytes, {len(results) - added_count} failures" ) return added_count, added_size @staticmethod async def handle_ratio_relations(relations: list[RatioRelation], session, target_ds: Dataset) -> tuple[int, int]: # Preload existing target file paths for deduplication existing_path_rows = await session.execute( select(DatasetFiles.file_path).where(DatasetFiles.dataset_id == target_ds.id) ) existing_paths = set(p for p in existing_path_rows.scalars().all() if p) source_paths = set() added_count = 0 added_size = 0 for rel in relations: if not rel.source_dataset_id or not rel.counts or rel.counts <= 0: continue files = await RatioTaskService.get_files(rel, session) if not files: continue pick_n = min(rel.counts or 0, len(files)) chosen = random.sample(files, pick_n) if pick_n < len(files) else files # Copy into target dataset with de-dup by target path for file in chosen: if file.file_path in source_paths: continue await RatioTaskService.handle_selected_file(existing_paths, file, session, target_ds) source_paths.add(file.file_path) added_count += 1 added_size += int(file.file_size or 0) # Periodically flush to avoid huge transactions await session.flush() return added_count, added_size @staticmethod async def handle_selected_file(existing_paths: set[Any], f, session, target_ds: Dataset): src_path = f.file_path dst_prefix = f"/dataset/{target_ds.id}/" file_name = RatioTaskService.get_new_file_name(dst_prefix, existing_paths, f) new_path = dst_prefix + file_name dst_dir = os.path.dirname(new_path) await asyncio.to_thread(os.makedirs, dst_dir, exist_ok=True) await asyncio.to_thread(shutil.copy2, src_path, new_path) file_data = { "dataset_id": target_ds.id, # type: ignore "file_name": file_name, "file_path": new_path, "file_type": f.file_type, "file_size": f.file_size, "check_sum": f.check_sum, "tags": f.tags, "tags_updated_at": datetime.now(), "dataset_filemetadata": f.dataset_filemetadata, "status": "ACTIVE", } file_record = {k: v for k, v in file_data.items() if v is not None} session.add(DatasetFiles(**file_record)) existing_paths.add(new_path) @staticmethod def get_new_file_name(dst_prefix: str, existing_paths: set[Any], f) -> str: file_name = f.file_name new_path = dst_prefix + file_name # Handle file path conflicts by appending a number to the filename if new_path in existing_paths: file_name_base, file_ext = os.path.splitext(file_name) counter = 1 original_file_name = file_name while new_path in existing_paths: file_name = f"{file_name_base}_{counter}{file_ext}" new_path = f"{dst_prefix}{file_name}" counter += 1 if counter > 1000: # Safety check to prevent infinite loops logger.error(f"Could not find unique filename for {original_file_name} after 1000 attempts") break return file_name @staticmethod async def get_files(rel: RatioRelation, session) -> list[Any]: # Fetch all files for the source dataset (ACTIVE only) files_res = await session.execute( select(DatasetFiles).where( DatasetFiles.dataset_id == rel.source_dataset_id, DatasetFiles.status == "ACTIVE", ) ) files = list(files_res.scalars().all()) # TAG mode: filter by tags according to relation.filter_conditions conditions = RatioTaskService._parse_conditions(rel.filter_conditions) if conditions: files = [f for f in files if RatioTaskService._filter_file(f, conditions)] return files # ------------------------- helpers for TAG filtering ------------------------- # @staticmethod def _parse_conditions(conditions: Optional[str]) -> Optional[FilterCondition]: """Parse filter_conditions JSON string into a FilterCondition object. Args: conditions: JSON string containing filter conditions Returns: FilterCondition object if conditions is not None/empty, otherwise None """ if not conditions: return None try: data = json.loads(conditions) return FilterCondition(**data) except json.JSONDecodeError as e: logger.error(f"Failed to parse filter conditions: {e}") return None except Exception as e: logger.error(f"Error creating FilterCondition: {e}") return None @staticmethod def _filter_file(file: DatasetFiles, conditions: FilterCondition) -> bool: if not conditions: return True logger.info(f"start filter file: {file}, conditions: {conditions}") # Check data range condition if provided if conditions.date_range and len(conditions.date_range) == 2: try: from datetime import datetime, timedelta start_at = datetime.fromisoformat(conditions.date_range[0]) end_at = datetime.fromisoformat(conditions.date_range[1]) if file.tags_updated_at and (file.tags_updated_at < start_at or file.tags_updated_at > end_at): return False except (ValueError, TypeError) as e: logger.warning(f"Invalid data_range value: {conditions.date_range}", e) return False # Check label condition if provided if conditions.label: tags = file.tags if not tags: return False try: # tags could be a list of strings or list of objects with 'name' all_tags = RatioTaskService.get_all_tags(tags) for tag in all_tags: if conditions.label.label and tag.get("label") != conditions.label.label: continue if conditions.label.value is None or len(conditions.label.value) == 0: return True if tag.get("value") == conditions.label.value: return True return False except Exception as e: logger.exception(f"Failed to get tags for {file}", e) return False return True @staticmethod def get_all_tags(tags) -> list[dict]: """获取所有处理后的标签字符串列表""" all_tags = list() if not tags: return all_tags file_tags = [] for tag_data in tags: # 处理可能的命名风格转换（下划线转驼峰） processed_data = {} for key, value in tag_data.items(): # 将驼峰转为下划线以匹配 Pydantic 模型字段 processed_data[key] = value # 创建 DatasetFileTag 对象 file_tag = DatasetFileTag(**processed_data) file_tags.append(file_tag) for file_tag in file_tags: for tag_data in file_tag.get_tags(): all_tags.append(tag_data) return all_tags @staticmethod async def _add_task_to_graph( session: AsyncSession, src_relations: List[RatioRelation], task: RatioInstance, dst_dataset: Dataset, ) -> None: """ 在比例抽取任务完成后，将数据集加入血缘图。 ratio_task(DATASOURCE) -> dataset(DATASET) via DATA_RATIO edge """ try: if not src_relations: logger.warning("Source ratio relations is empty when building lineage graph") return lineage_service = LineageService(db=session) dst_node = LineageNode( id=dst_dataset.id, node_type=NodeType.DATASET.value, name=dst_dataset.name, description=dst_dataset.description, ) for rel in src_relations: ds: Optional[Dataset] = await session.get(Dataset, rel.source_dataset_id) src_node = LineageNode( id=rel.source_dataset_id, node_type=NodeType.DATASET.value, name=ds.name, description=ds.description, ) ratio_edge = LineageEdge( process_id=task.id, name=task.name, edge_type=EdgeType.DATA_RATIO.value, description=task.description, from_node_id=rel.source_dataset_id, to_node_id=dst_node.id, ) await lineage_service.generate_graph(src_node, ratio_edge, dst_node) logger.info("Add dataset lineage graph: %s -> %s -> %s", src_node.id, ratio_edge.id, dst_node.id) await session.commit() logger.info("Add dataset lineage graph success") except Exception as exc: logger.error("Failed to add dataset lineage graph: %s", exc) await session.rollback()