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

runtime/deepdata-python/app/module/shared/common/lineage.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+from typing import List, Optional
+from uuid import uuid4
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.core.logging import get_logger
+from app.db.models.base_entity import LineageNode, LineageEdge
+
+log = get_logger(__name__)
+
+
+class LineageService:
+    """血缘服务（Python 版）"""
+
+    def __init__(self, db: AsyncSession):
+        self.db = db
+
+    async def generate_graph(
+        self,
+        from_node: Optional[LineageNode],
+        edge: Optional[LineageEdge],
+        to_node: Optional[LineageNode]
+    ) -> None:
+        """
+        生成血缘图
+
+        Args:
+            from_node: 源节点
+            edge: 边
+            to_node: 目的节点
+        """
+        # 1. 如果 from_node 为空，直接返回不做任何处理
+        if from_node is None:
+            return
+
+        # 2. 如果 from_node 有值，检查 from_node 是否存在
+        existing_from = await self.get_node_by_id(from_node.id) if from_node.id else None
+
+        if existing_from is None:
+            # from_node 不存在，创建 from_node，记录 from_graph_id
+            from_graph_id = str(uuid4())
+            if not from_node.id:
+                from_node.id = str(uuid4())
+            from_node.graph_id = from_graph_id
+            self.db.add(from_node)
+            await self.db.flush()
+        else:
+            # from_node 存在，记录 from_graph_id
+            from_graph_id = existing_from.graph_id
+
+        # 处理 edge 和 to_node
+        await self._generate_graph(edge, to_node, from_graph_id)
+
+    async def _generate_graph(
+        self,
+        edge: Optional[LineageEdge],
+        to_node: Optional[LineageNode],
+        graph_id: Optional[str]
+    ) -> None:
+        # 无 edge，直接返回
+        if edge is None:
+            return
+
+        # 有 edge，处理 edge，检查是否有 to_node
+        await self._handle_lineage_edge(graph_id, edge)
+        if to_node is None:
+            return
+
+        # 有 to_node，检查 to_node 是否存在
+        existing_to = await self.get_node_by_id(to_node.id) if to_node.id else None
+
+        if existing_to is None:
+            # to_node 不存在，创建 to_node 后结束
+            if not to_node.id:
+                to_node.id = str(uuid4())
+            to_node.graph_id = graph_id
+            self.db.add(to_node)
+            await self.db.flush()
+        else:
+            # to_node 存在，将 from_node 所在的图并入 to_node 所在的图后结束
+            await self._merge_graph(graph_id, existing_to.graph_id)
+
+    async def _handle_lineage_edge(self, graph_id: Optional[str], edge: LineageEdge) -> None:
+        stmt = select(LineageEdge).where(
+            LineageEdge.graph_id == graph_id,
+            LineageEdge.from_node_id == edge.from_node_id,
+            LineageEdge.to_node_id == edge.to_node_id
+        )
+        result = await self.db.execute(stmt)
+        existing = result.scalar_one_or_none()
+
+        if existing is None:
+            if not edge.id:
+                edge.id = str(uuid4())
+            edge.graph_id = graph_id
+            self.db.add(edge)
+            await self.db.flush()
+        else:
+            edge.id = existing.id
+            if edge.graph_id is None:
+                edge.graph_id = existing.graph_id
+            # 覆盖更新现有记录
+            existing.graph_id = edge.graph_id
+            existing.process_id = edge.process_id
+            existing.edge_type = edge.edge_type
+            existing.name = edge.name
+            existing.description = edge.description
+            existing.edge_metadata = edge.edge_metadata
+            existing.from_node_id = edge.from_node_id
+            existing.to_node_id = edge.to_node_id
+            await self.db.flush()
+
+    async def _merge_graph(self, from_graph_id: Optional[str], to_graph_id: Optional[str]) -> None:
+        if not from_graph_id or not to_graph_id or from_graph_id == to_graph_id:
+            return
+
+        from_nodes = await self.get_nodes_by_graph_id(from_graph_id)
+        to_nodes = await self.get_nodes_by_graph_id(to_graph_id)
+
+        # choose smaller graph as source, larger as target
+        source_graph_id = from_graph_id if len(from_nodes) <= len(to_nodes) else to_graph_id
+        target_graph_id = to_graph_id if source_graph_id == from_graph_id else from_graph_id
+        source_nodes = from_nodes if source_graph_id == from_graph_id else to_nodes
+
+        # update nodes' graph_id
+        for node in source_nodes:
+            if node is None:
+                continue
+            node.graph_id = target_graph_id
+            self.db.add(node)
+
+        # update edges' graph_id (edges belonging to the source graph)
+        edges = await self.get_edges_by_graph_id(source_graph_id)
+        for edge in edges:
+            if edge is None:
+                continue
+            edge.graph_id = target_graph_id
+            self.db.add(edge)
+
+        await self.db.flush()
+
+    async def get_nodes_by_graph_id(self, graph_id: str) -> List[LineageNode]:
+        """从图ID获取图的节点列表"""
+        stmt = select(LineageNode).where(LineageNode.graph_id == graph_id)
+        result = await self.db.execute(stmt)
+        return list(result.scalars().all())
+
+    async def get_edges_by_graph_id(self, graph_id: str) -> List[LineageEdge]:
+        """从图ID获取图的边列表"""
+        stmt = select(LineageEdge).where(LineageEdge.graph_id == graph_id)
+        result = await self.db.execute(stmt)
+        return list(result.scalars().all())
+
+    async def get_node_by_id(self, node_id: str) -> Optional[LineageNode]:
+        """从节点ID获取节点"""
+        return await self.db.get(LineageNode, node_id)