# -*- coding: utf-8 -*-
|
"""
|
@File : tree_service_impl.py
|
@Time : 2023/1/14 16:10
|
@Author : geekbing
|
@LastEditTime : 2023/7/27 21:10
|
@LastEditors : geekbing
|
@Description : 树形结构操作
|
"""
|
from ast import literal_eval
|
import traceback
|
from typing import Dict, List, Optional
|
|
from loguru import logger
|
|
from crud.base_crud import GenericCURD
|
from lunarlink.dto.tree_dto import TreeOut, TreeUniqueIn, TreeUpdateIn
|
from lunarlink.models import API, Case, Relation
|
from lunarlink.utils.response import (
|
TREE_GET_SUCCESS,
|
TREE_UPDATE_SUCCESS,
|
StandResponse,
|
)
|
from lunarlink.utils.tree import get_tree_max_id
|
from lunarlink.utils.enums.TreeTypeEnum import TreeType
|
|
|
class TreeService:
|
def __init__(self):
|
self.model = Relation
|
self.curd = GenericCURD(self.model)
|
|
def get_or_create(self, query: TreeUniqueIn) -> StandResponse[TreeOut]:
|
default_tree = [
|
{
|
"id": 1,
|
"label": "默认目录",
|
"children": [],
|
}
|
]
|
tree_obj, is_created = self.curd.get_or_create(
|
filter_kwargs=query.dict(),
|
defaults={"tree": default_tree, "project_id": query.project_id},
|
)
|
if is_created:
|
logger.info(f"tree created {query=}")
|
body: List[Dict] = tree_obj.tree
|
else:
|
logger.info(f"tree exist {query}")
|
body: List[Dict] = literal_eval(tree_obj.tree)
|
|
node_api_case_counts = {}
|
for node in body:
|
if query.type == TreeType.API:
|
TreeService.add_api_count_to_node(
|
tree_obj.project_id,
|
node,
|
node_api_case_counts,
|
)
|
else:
|
TreeService.add_case_count_to_node(
|
tree_obj.project_id,
|
node,
|
node_api_case_counts,
|
)
|
|
for root_node in body:
|
TreeService.add_api_case_count_to_tree(root_node, node_api_case_counts)
|
|
tree = {
|
"tree": body,
|
"id": tree_obj.id,
|
"success": True,
|
"max": get_tree_max_id(body),
|
}
|
return StandResponse[TreeOut](**TREE_GET_SUCCESS, data=TreeOut(**tree))
|
|
@staticmethod
|
def add_case_count_to_node(project_id: int, node: Dict, node_case_count: Dict):
|
"""
|
递归获取节点的用例数量
|
:param project_id:
|
:param node:
|
:param node_case_count:
|
:return:
|
"""
|
node_id = node["id"]
|
case_count = Case.objects.filter(
|
project_id=project_id, relation=node_id
|
).count()
|
node_case_count[node_id] = case_count
|
|
for child in node.get("children", []):
|
TreeService.add_case_count_to_node(project_id, child, node_case_count)
|
|
@staticmethod
|
def add_api_count_to_node(project_id: int, node: Dict, node_api_count: Dict):
|
"""
|
递归获取节点的接口数量
|
:param project_id:
|
:param node:
|
:param node_api_count:
|
:return:
|
"""
|
node_id = node["id"]
|
api_count = API.objects.filter(project_id=project_id, relation=node_id).count()
|
node_api_count[node_id] = api_count
|
|
for child in node.get("children", []):
|
TreeService.add_api_count_to_node(project_id, child, node_api_count)
|
|
@staticmethod
|
def add_api_case_count_to_tree(node: Dict, node_api_case_counts: Dict):
|
"""
|
将节点的接口或用例数量嵌入原树形结构中
|
:param node:
|
:param node_api_case_counts:
|
:return:
|
"""
|
data_count = node_api_case_counts.get(node["id"], 0)
|
if "children" in node:
|
for child in node["children"]:
|
data_count += TreeService.add_api_case_count_to_tree(
|
child, node_api_case_counts
|
)
|
node["data_count"] = data_count
|
return data_count
|
|
@staticmethod
|
def check_related_api(node_id: int, project_id: int) -> bool:
|
# 检查是否存在关联的接口
|
return API.objects.filter(relation=node_id, project_id=project_id).exists()
|
|
@staticmethod
|
def check_related_case(node_id: int, project_id: int) -> bool:
|
# 检查是否存在关联的用例
|
return Case.objects.filter(relation=node_id, project_id=project_id).exists()
|
|
def get_all_ids_from_tree(self, tree):
|
"""
|
获取树的所有节点id
|
:param tree:
|
:return:
|
"""
|
ids = [node["id"] for node in tree]
|
for node in tree:
|
if "children" in node:
|
ids.extend(self.get_all_ids_from_tree(node["children"]))
|
return ids
|
|
def _check_related_api_recursive(self, current_tree, payload_tree, project_id):
|
"""
|
提取payload.tree所有节点id，递归检查current_tree中是否存在这些id，如果不存在，检查是否有关联的API
|
"""
|
payload_tree_id_list = self.get_all_ids_from_tree(payload_tree)
|
|
for node in current_tree:
|
if node["id"] not in payload_tree_id_list:
|
# 这个节点正在被删除，检查是否有关联的API
|
if self.check_related_api(node["id"], project_id):
|
return True
|
|
# 如果该节点有子节点，继续递归检查
|
if "children" in node and node["children"]:
|
if self._check_related_api_recursive(
|
node["children"], payload_tree, project_id
|
):
|
return True
|
return False
|
|
def _check_related_case_recursive(self, current_tree, payload_tree, project_id):
|
"""
|
提取payload.tree所有节点id，递归检查current_tree中是否存在这些id，如果不存在，检查是否有关联的用例
|
"""
|
payload_tree_id_list = self.get_all_ids_from_tree(payload_tree)
|
|
for node in current_tree:
|
if node["id"] not in payload_tree_id_list:
|
# 这个节点正在被删除，检查是否有关联的API
|
if self.check_related_case(node["id"], project_id):
|
return True
|
|
# 如果该节点有子节点，继续递归检查
|
if "children" in node and node["children"]:
|
if self._check_related_case_recursive(
|
node["children"], payload_tree, project_id
|
):
|
return True
|
return False
|
|
@staticmethod
|
def get_current_tree(relation_obj) -> List:
|
"""
|
获取当前的树形结构, 转成python对象
|
:param relation_obj:
|
:return:
|
"""
|
if not relation_obj:
|
return []
|
return literal_eval(relation_obj.tree)
|
|
def patch(self, tree_id: int, payload: TreeUpdateIn):
|
"""
|
更新树形结构
|
|
:param tree_id:
|
:param payload:
|
:return:
|
"""
|
# 获取当前的树结构
|
if not payload.tree:
|
return StandResponse[Optional[TreeOut]](
|
code="9999",
|
success=False,
|
msg="删除失败, 至少保留一个根目录",
|
data=None,
|
)
|
relation_obj = self.curd.get_obj_by_pk(pk=tree_id)
|
current_tree = self.get_current_tree(relation_obj)
|
# 检查每个节点，如果在当前树中找到了这个节点，但在新树中找不到，就意味着这个节点被删除了
|
if payload.type == TreeType.API:
|
if self._check_related_api_recursive(
|
current_tree, payload.tree, relation_obj.project_id
|
):
|
# 如果存在关联的接口，返回提示信息
|
return StandResponse[Optional[TreeOut]](
|
code="9999",
|
success=False,
|
msg="目录有关联接口，不能删除",
|
data=None,
|
)
|
elif payload.type == TreeType.CASE:
|
# 如果存在关联的用例，返回提示信息
|
if self._check_related_case_recursive(
|
current_tree, payload.tree, relation_obj.project_id
|
):
|
# 如果存在关联的接口，返回提示信息
|
return StandResponse[Optional[TreeOut]](
|
code="9999",
|
success=False,
|
msg="目录有关联用例，不能删除",
|
data=None,
|
)
|
|
# 如果没有问题，尝试更新
|
try:
|
tree_obj = self.curd.update_obj_by_pk(
|
pk=tree_id, updater="", payload=payload.dict()
|
)
|
except Exception as e:
|
return self._handle_exception(e)
|
|
tree: List[Dict] = tree_obj.tree
|
node_api_case_counts = {}
|
for node in tree:
|
if payload.type == TreeType.API:
|
TreeService.add_api_count_to_node(
|
tree_obj.project_id,
|
node,
|
node_api_case_counts,
|
)
|
else:
|
TreeService.add_case_count_to_node(
|
tree_obj.project_id,
|
node,
|
node_api_case_counts,
|
)
|
|
for root_node in tree:
|
TreeService.add_api_case_count_to_tree(root_node, node_api_case_counts)
|
|
return StandResponse[TreeOut](
|
**TREE_UPDATE_SUCCESS,
|
data=TreeOut(tree=tree, id=tree_obj.id, max=get_tree_max_id(tree)),
|
)
|
|
@staticmethod
|
def _handle_exception(e: Exception) -> StandResponse[Optional[TreeOut]]:
|
err: str = traceback.format_exc()
|
logger.warning(f"Exception {e} occurred with traceback: {err}")
|
return StandResponse[Optional[TreeOut]](
|
code="9999", success=False, msg=err, data=None
|
)
|
|
|
tree_service = TreeService()
|
