DataOps-platform/app/services/neo4j_driver.py

from neo4j import GraphDatabase
from neo4j.exceptions import ServiceUnavailable
import os

class Neo4jDriver:
    def __init__(self, uri=None, user=None, password=None, encrypted=False):
        self._driver = None
        
        # 优先使用传入的参数
        if uri is not None:
            self.uri = uri
        elif user is not None or password is not None or encrypted is not None:
            # 如果只传了部分参数，使用环境变量或默认值填充缺失的
            self.uri = os.environ.get('NEO4J_URI', "bolt://192.168.3.143:7687")
        else:
            # 无参数调用时，尝试从Flask配置获取，否则使用环境变量或默认值
            self.uri = self._get_config_value('NEO4J_URI', "bolt://192.168.3.143:7687")
        
        if user is not None:
            self.user = user
        else:
            self.user = self._get_config_value('NEO4J_USER', "neo4j")
            
        if password is not None:
            self.password = password
        else:
            self.password = self._get_config_value('NEO4J_PASSWORD', "cituneo4j")
            
        if encrypted is not None:
            self.encrypted = encrypted
        else:
            encrypted_str = self._get_config_value('NEO4J_ENCRYPTED', 'false')
            self.encrypted = encrypted_str.lower() == 'true' if isinstance(encrypted_str, str) else encrypted_str
    
    def _get_config_value(self, key, default_value):
        """获取配置值，优先从Flask配置获取，否则从环境变量获取，最后使用默认值"""
        try:
            # 尝试从Flask配置获取
            from flask import current_app
            if current_app and hasattr(current_app, 'config'):
                return current_app.config.get(key, default_value)
        except (ImportError, RuntimeError):
            # 不在Flask环境中或Flask应用上下文外
            pass
        
        # 从环境变量获取
        return os.environ.get(key, default_value)
        
    def connect(self):
        if not self._driver:
            self._driver = GraphDatabase.driver(
                self.uri,
                auth=(self.user, self.password),
                encrypted=self.encrypted
            )
        return self._driver
    
    def close(self):
        if self._driver:
            self._driver.close()
            self._driver = None
            
    def verify_connectivity(self):
        try:
            self.connect().verify_connectivity()
            return True
        except ServiceUnavailable:
            return False
    
    def get_session(self):
        """获取 Neo4j 会话"""
        return self.connect().session()

# 单例实例
neo4j_driver = Neo4jDriver()