src/utils/neo4j_helper.py

from neo4j import GraphDatabase
from typing import List, Dict, Any
import logging

class Neo4jConnection:
    def __init__(self, uri: str = "bolt://localhost:7687", 
                 user: str = "neo4j", 
                 password: str = "your_password"):
        """初始化 Neo4j 连接"""
        try:
            self.driver = GraphDatabase.driver(uri, auth=(user, password))
            logging.info("Neo4j 连接成功")
        except Exception as e:
            logging.error(f"Neo4j 连接失败: {str(e)}")
            raise

    def close(self):
        """关闭连接"""
        if self.driver:
            self.driver.close()

    def run_query(self, query: str, parameters: Dict[str, Any] = None) -> List[Dict]:
        """执行 Cypher 查询"""
        try:
            with self.driver.session() as session:
                result = session.run(query, parameters or {})
                return [record.data() for record in result]
        except Exception as e:
            logging.error(f"执行查询失败: {str(e)}")
            raise

    def create_node(self, label: str, properties: Dict[str, Any]) -> Dict:
        """创建节点"""
        query = f"""
        CREATE (n:{label} $properties)
        RETURN n
        """
        return self.run_query(query, {"properties": properties})

    def create_relationship(self, start_node_label: str, start_node_props: Dict,
                          end_node_label: str, end_node_props: Dict,
                          relationship_type: str, relationship_props: Dict = None) -> Dict:
        """创建关系"""
        query = f"""
        MATCH (a:{start_node_label}), (b:{end_node_label})
        WHERE a.id = $start_props.id AND b.id = $end_props.id
        CREATE (a)-[r:{relationship_type} $rel_props]->(b)
        RETURN a, r, b
        """
        params = {
            "start_props": start_node_props,
            "end_props": end_node_props,
            "rel_props": relationship_props or {}
        }
        return self.run_query(query, params)

    def get_node(self, label: str, properties: Dict[str, Any]) -> Dict:
        """获取节点"""
        query = f"""
        MATCH (n:{label})
        WHERE n.id = $properties.id
        RETURN n
        """
        return self.run_query(query, {"properties": properties})

# 使用示例：
if __name__ == "__main__":
    # 创建连接
    neo4j = Neo4jConnection()
    
    try:
        # 创建示例节点
        node = neo4j.create_node("Place", {
            "id": "1",
            "name": "香港迪士尼乐园",
            "type": "景点"
        })
        print("创建节点成功:", node)
        
    finally:
        # 关闭连接
        neo4j.close()