MemoryBear/api/app/repositories/neo4j/neo4j_connector.py

# -*- coding: utf-8 -*-
"""Neo4j连接器模块

本模块提供Neo4j图数据库的连接和查询功能。
从 app/core/memory/src/database/neo4j_connector.py 迁移而来。

Classes:
    Neo4jConnector: Neo4j数据库连接器，提供异步查询接口
"""

from typing import Any, List, Dict

from neo4j import AsyncGraphDatabase, basic_auth

from app.core.config import settings


class Neo4jConnector:
    """Neo4j数据库连接器
    
    提供与Neo4j图数据库的连接和查询功能。
    使用异步驱动程序以支持高并发操作。
    
    Attributes:
        driver: Neo4j异步驱动程序实例
        
    Methods:
        close: 关闭数据库连接
        execute_query: 执行Cypher查询
        delete_group: 删除指定组的所有数据
    """
    
    def __init__(self):
        """初始化Neo4j连接器
        
        从配置文件和环境变量中读取连接信息。
        
        Raises:
            RuntimeError: 如果NEO4J_PASSWORD环境变量未设置
        """
        # 从全局配置和环境变量获取 Neo4j 配置
        uri = settings.NEO4J_URI
        username = settings.NEO4J_USERNAME
        password = settings.NEO4J_PASSWORD
        
        if not password:
            raise RuntimeError(
                "NEO4J_PASSWORD is not set. Create a .env with NEO4J_PASSWORD or export it before running."
            )
        self.driver = AsyncGraphDatabase.driver(
            uri,
            auth=basic_auth(username, password)
        )

    async def close(self):
        """关闭数据库连接
        
        释放数据库连接资源。应在应用程序关闭时调用。
        """
        await self.driver.close()

    async def execute_query(self, query: str, **kwargs: Any) -> List[Dict[str, Any]]:
        """执行Cypher查询
        
        Args:
            query: Cypher查询语句
            **kwargs: 查询参数，将作为参数传递给Cypher查询
            
        Returns:
            List[Dict[str, Any]]: 查询结果列表，每个元素是一个字典
            
        Example:
            >>> connector = Neo4jConnector()
            >>> results = await connector.execute_query(
            ...     "MATCH (n:Person {name: $name}) RETURN n",
            ...     name="Alice"
            ... )
        """
        result = await self.driver.execute_query(
            query,
            database="neo4j",
            **kwargs
        )
        records, summary, keys = result
        return [record.data() for record in records]
    
    async def execute_write_transaction(self, transaction_func, **kwargs: Any) -> Any:
        """在写事务中执行操作
        
        提供显式事务支持，确保操作的原子性。
        如果事务函数抛出异常，所有更改将自动回滚。
        
        Args:
            transaction_func: 事务函数，接收 tx 参数并执行查询
            **kwargs: 传递给事务函数的额外参数
            
        Returns:
            Any: 事务函数的返回值
            
        Example:
            >>> async def create_node(tx, name):
            ...     result = await tx.run(
            ...         "CREATE (n:Person {name: $name}) RETURN n",
            ...         name=name
            ...     )
            ...     return await result.single()
            >>> 
            >>> connector = Neo4jConnector()
            >>> result = await connector.execute_write_transaction(
            ...     create_node, name="Alice"
            ... )
        """
        async with self.driver.session(database="neo4j") as session:
            return await session.execute_write(transaction_func, **kwargs)
    
    async def execute_read_transaction(self, transaction_func, **kwargs: Any) -> Any:
        """在读事务中执行操作
        
        提供显式事务支持用于读操作。
        
        Args:
            transaction_func: 事务函数，接收 tx 参数并执行查询
            **kwargs: 传递给事务函数的额外参数
            
        Returns:
            Any: 事务函数的返回值
            
        Example:
            >>> async def get_node(tx, name):
            ...     result = await tx.run(
            ...         "MATCH (n:Person {name: $name}) RETURN n",
            ...         name=name
            ...     )
            ...     return await result.single()
            >>> 
            >>> connector = Neo4jConnector()
            >>> result = await connector.execute_read_transaction(
            ...     get_node, name="Alice"
            ... )
        """
        async with self.driver.session(database="neo4j") as session:
            return await session.execute_read(transaction_func, **kwargs)
    
    async def delete_group(self, group_id: str):
        """删除指定组的所有数据
        
        删除所有属于指定group_id的节点和边。
        这是一个危险操作，会永久删除数据。
        
        Args:
            group_id: 要删除的组ID
            
        Example:
            >>> connector = Neo4jConnector()
            >>> await connector.delete_group("group_123")
            Group group_123 deleted.
        """
        # 删除节点（DETACH DELETE会同时删除相关的边）
        await self.driver.execute_query(
            "MATCH (n) WHERE n.group_id = $group_id DETACH DELETE n",
            database="neo4j",
            group_id=group_id
        )
        # 删除独立的边（如果有的话）
        await self.driver.execute_query(
            "MATCH ()-[r]->() WHERE r.group_id = $group_id DELETE r",
            database="neo4j",
            group_id=group_id
        )
        print(f"Group {group_id} deleted.")