From e3067e115afed75a14d5566cf4417bf0b4f3fc37 Mon Sep 17 00:00:00 2001
From: caoqianming <caoqianming@foxmail.com>
Date: Thu, 3 Jul 2025 14:54:20 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E8=8E=B7=E5=8F=96=E8=AF=A5=E6=89=B9?=
 =?UTF-8?q?=E6=AC=A1=E7=9A=84DAG=E5=9B=BE=E6=95=B0=E6=8D=AE=20=E6=94=AF?=
 =?UTF-8?q?=E6=8C=81=E4=BB=85=E7=9B=B4=E6=8E=A5=E5=85=B3=E7=B3=BB?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/wpm/services.py | 91 +++++++++++++++++++++++++-------------------
 apps/wpm/views.py    |  4 +-
 2 files changed, 53 insertions(+), 42 deletions(-)

diff --git a/apps/wpm/services.py b/apps/wpm/services.py
index b7f38915..4a82b727 100644
--- a/apps/wpm/services.py
+++ b/apps/wpm/services.py
@@ -1017,27 +1017,29 @@ def mlog_audit_end(ticket: Ticket):
     mlog_submit(ins, ticket.create_by, now)
 
 
-def get_batch_dag(batch_number: str, version=1):
-        try:
-            batch_ins = BatchSt.objects.get(batch=batch_number, version=version)
-        except Exception:
-            raise ParseError("该批次号未构建关系链")
-        # 收集所有相关批次和边
+def get_batch_dag(batch_number: str, method="full"):
+    try:
+        batch_ins = BatchSt.objects.get(batch=batch_number, version=1)
+    except Exception:
+        raise ParseError("该批次号未构建关系链")
+
+    r_dict = {
+        "split": "分",
+        "merge": "合"
+    }
+
+    if method == "full":
+        # 完整DAG模式 - 收集所有相关批次和边（原逻辑）
         nodes_set = {batch_ins.id}
         edges = []
         prev_size = 0
 
-        r_dict = {
-            "split": "分",
-            "merge": "合"
-        }
         while len(nodes_set) > prev_size:
             prev_size = len(nodes_set)
-            # 查询所有与当前批次相关的记录（作为source或target）
             logs = BatchLog.objects.filter(Q(source__id__in=nodes_set) | Q(target__id__in=nodes_set)).select_related(
                 "source", "target"
             ).order_by("update_time")
-            # 处理每条记录，扩展节点和边
+            
             for log in logs:
                 source = log.source.id
                 target = log.target.id
@@ -1047,33 +1049,42 @@ def get_batch_dag(batch_number: str, version=1):
                     'target': target,
                     "handover": log.handover.id if log.handover else None,
                     "mlog": log.mlog.id if log.mlog else None,
-                    'label': r_dict.get(log.relation_type, ""),  # 使用relation_type作为边的标签
+                    'label': r_dict.get(log.relation_type, ""),
                 })
-        
-        # 去重边
-        # unique_edges = {}
-        # for edge in edges:
-        #     key = (edge['source'], edge['target'])
-        #     if key not in unique_edges:
-        #         unique_edges[key] = edge
-        
-        # 将批次号排序
-        nodes_qs = BatchSt.objects.filter(id__in=nodes_set).order_by('update_time')
-        
-        # batch_to_id = {batch: idx for idx, batch in enumerate(nodes_list)}
-        # 构建节点数据，默认使用'rect'形状
-        nodes = [{
-            'id': item.id,
-            'label': item.batch,
-            'shape': 'rect'  # 可根据业务需求调整形状
-        } for item in nodes_qs]
-        
-        # 构建边数据
-        # edges_converted = [{
-        #     'source': edge['source'],
-        #     'target': edge['target'],
-        #     'label': edge['label']
-        # } for edge in unique_edges.values()]
-        
-        return {'nodes': nodes, 'edges': edges}
+
+    elif method == "direct":
+        # 直接关系模式 - 只查询直接相连的批次
+        nodes_set = {batch_ins.id}
+        edges = []
+
+        # 查询作为source或target的直接关系
+        logs = BatchLog.objects.filter(Q(source__id=batch_ins.id) | Q(target__id=batch_ins.id)).select_related(
+            "source", "target"
+        ).order_by("update_time")
+        
+        for log in logs:
+            source = log.source.id
+            target = log.target.id
+            nodes_set.update([source, target])
+            edges.append({
+                'source': source,
+                'target': target,
+                "handover": log.handover.id if log.handover else None,
+                "mlog": log.mlog.id if log.mlog else None,
+                'label': r_dict.get(log.relation_type, ""),
+            })
+
+    else:
+        raise ParseError("不支持的查询方法，请使用'full'或'direct'")
+
+    # 将批次号排序
+    nodes_qs = BatchSt.objects.filter(id__in=nodes_set).order_by('update_time')
+    
+    nodes = [{
+        'id': item.id,
+        'label': item.batch,
+        'shape': 'rect'
+    } for item in nodes_qs]
+
+    return {'nodes': nodes, 'edges': edges}
 
diff --git a/apps/wpm/views.py b/apps/wpm/views.py
index 0b8c4ee1..8be227da 100644
--- a/apps/wpm/views.py
+++ b/apps/wpm/views.py
@@ -927,10 +927,10 @@ class BatchLogViewSet(ListModelMixin, CustomGenericViewSet):
         获取该批次的DAG图数据
         """
         batch = request.data.get("batch", None)
-        version = request.data.get("version", 1)
+        method = request.data.get("method", "full")
         if not batch:
             raise ParseError("缺少batch参数")
-        return Response(get_batch_dag(batch, version))
+        return Response(get_batch_dag(batch, method))
     
     @action(methods=['post'], detail=False, perms_map={'post': '*'}, serializer_class=Serializer)
     def batches_to(self, request, *args, **kwargs):