160 lines
5.8 KiB
Python
160 lines
5.8 KiB
Python
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)
|