kg_edge_service.py 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114
  1. from sqlalchemy.orm import Session
  2. from sqlalchemy import or_
  3. from typing import Optional
  4. from model.kg_edges import KGEdge
  5. from db.session import get_db
  6. import logging
  7. from sqlalchemy.exc import IntegrityError
  8. from cachetools import TTLCache
  9. from cachetools.keys import hashkey
  10. logger = logging.getLogger(__name__)
  11. class KGEdgeService:
  12. def __init__(self, db: Session):
  13. self.db = db
  14. _cache = TTLCache(maxsize=100000, ttl=60*60*24*30)
  15. def get_edge(self, edge_id: int):
  16. edge = self.db.query(KGEdge).get(edge_id)
  17. if not edge:
  18. raise ValueError("Edge not found")
  19. return edge
  20. def create_edge(self, edge_data: dict):
  21. try:
  22. existing = self.db.query(KGEdge).filter(
  23. KGEdge.src_id == edge_data['src_id'],
  24. KGEdge.dest_id == edge_data['dest_id'],
  25. KGEdge.name == edge_data['name'],
  26. KGEdge.version == edge_data.get('version')
  27. ).first()
  28. if existing:
  29. raise ValueError("Edge already exists")
  30. new_edge = KGEdge(**edge_data)
  31. self.db.add(new_edge)
  32. self.db.commit()
  33. return new_edge
  34. except IntegrityError as e:
  35. self.db.rollback()
  36. logger.error(f"创建边失败: {str(e)}")
  37. raise ValueError("Database integrity error")
  38. def update_edge(self, edge_id: int, update_data: dict):
  39. edge = self.db.query(KGEdge).get(edge_id)
  40. if not edge:
  41. raise ValueError("Edge not found")
  42. try:
  43. for key, value in update_data.items():
  44. setattr(edge, key, value)
  45. self.db.commit()
  46. return edge
  47. except Exception as e:
  48. self.db.rollback()
  49. logger.error(f"更新边失败: {str(e)}")
  50. raise ValueError("Update failed")
  51. def delete_edge(self, edge_id: int):
  52. edge = self.db.query(KGEdge).get(edge_id)
  53. if not edge:
  54. raise ValueError("Edge not found")
  55. try:
  56. self.db.delete(edge)
  57. self.db.commit()
  58. return None
  59. except Exception as e:
  60. self.db.rollback()
  61. logger.error(f"删除边失败: {str(e)}")
  62. raise ValueError("Delete failed")
  63. def get_edges_by_nodes(self, src_id: Optional[int]= None, dest_id: Optional[int]= None, category: Optional[str] = None):
  64. cache_key = f"get_edges_by_nodes_{src_id}_{dest_id}_{category}"
  65. if cache_key in self._cache:
  66. return self._cache[cache_key]
  67. if src_id is None and dest_id is None:
  68. raise ValueError("至少需要提供一个有效的查询条件")
  69. try:
  70. filters = []
  71. if src_id is not None:
  72. filters.append(KGEdge.src_id == src_id)
  73. if dest_id is not None:
  74. filters.append(KGEdge.dest_id == dest_id)
  75. if category is not None:
  76. filters.append(KGEdge.category == category)
  77. edges = self.db.query(KGEdge).filter(*filters).all()
  78. from service.kg_node_service import KGNodeService
  79. node_service = KGNodeService(self.db)
  80. result = []
  81. for edge in edges:
  82. try:
  83. edge_info = {
  84. 'id': edge.id,
  85. 'src_id': edge.src_id,
  86. 'dest_id': edge.dest_id,
  87. 'name': edge.name,
  88. 'version': edge.version,
  89. #'src_node': node_service.get_node(edge.src_id),
  90. 'dest_node': node_service.get_node(edge.dest_id)
  91. }
  92. result.append(edge_info)
  93. except ValueError as e:
  94. logger.warning(f"跳过边关系 {edge.id}: {str(e)}")
  95. continue
  96. self._cache[cache_key] = result
  97. return result
  98. except Exception as e:
  99. logger.error(f"查询边失败: {str(e)}")
  100. raise e