Skip to main content

nexus_common/db/connectors/
neo4j.rs

1use neo4rs::{query, Graph};
2use once_cell::sync::OnceCell;
3use std::fmt;
4use std::sync::Arc;
5use tokio::sync::Mutex;
6use tracing::{debug, info};
7
8use crate::db::setup::setup_graph;
9use crate::db::Neo4JConfig;
10use crate::types::DynError;
11
12pub struct Neo4jConnector {
13    pub graph: OnceCell<Arc<Mutex<Graph>>>,
14}
15
16impl Default for Neo4jConnector {
17    fn default() -> Self {
18        Self {
19            graph: OnceCell::new(),
20        }
21    }
22}
23
24impl Neo4jConnector {
25    /// Initialize and register the global Neo4j connector and verify connectivity
26    pub async fn init(neo4j_config: &Neo4JConfig) -> Result<(), DynError> {
27        let neo4j_connector = Neo4jConnector::new_connection(
28            &neo4j_config.uri,
29            &neo4j_config.user,
30            &neo4j_config.password,
31        )
32        .await?;
33
34        neo4j_connector.ping(&neo4j_config.uri).await?;
35
36        match NEO4J_CONNECTOR.set(neo4j_connector) {
37            Err(e) => debug!("Neo4jConnector was already set: {:?}", e),
38            Ok(()) => info!("Neo4jConnector successfully set up on {}", neo4j_config.uri),
39        }
40
41        // Set Neo4J graph data constraints
42        setup_graph().await?;
43        Ok(())
44    }
45
46    /// Create and return a new connector after defining a database connection
47    async fn new_connection(uri: &str, user: &str, password: &str) -> Result<Self, DynError> {
48        let neo4j_connector = Neo4jConnector::default();
49        match neo4j_connector.connect(uri, user, password).await {
50            Ok(_) => info!("Created Neo4j connector"),
51            Err(e) => return Err(format!("Could not create Neo4J connector: {e}").into()),
52        }
53        Ok(neo4j_connector)
54    }
55
56    /// Dewfine a connection to the Neo4j database and store the graph instance
57    async fn connect(
58        &self,
59        uri: &str,
60        user: &str,
61        password: &str,
62    ) -> Result<(), Box<dyn std::error::Error>> {
63        let graph = Graph::new(uri, user, password).await?;
64        self.graph
65            .set(Arc::new(Mutex::new(graph)))
66            .map_err(|_| "Failed to set graph instance")?;
67        Ok(())
68    }
69
70    /// Perform a health-check PING over the Bolt protocol to the Neo4j server
71    async fn ping(&self, neo4j_uri: &str) -> Result<(), DynError> {
72        let graph = self.graph.get().ok_or("Neo4jConnector not initialized")?;
73        let graph = graph.lock().await;
74        match graph.execute(query("RETURN 1")).await {
75            Ok(_) => info!(
76                "Bolt protocol health-check PING to Neo4j succeeded; server is responsive at {}",
77                neo4j_uri
78            ),
79            Err(neo4j_err) => {
80                return Err(format!("Failed to PING to Neo4j at {neo4j_uri}, {neo4j_err}").into())
81            }
82        };
83        Ok(())
84    }
85}
86
87impl fmt::Debug for Neo4jConnector {
88    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
89        f.debug_struct("Neo4jConnector")
90            .field("graph", &"Graph instance")
91            .finish()
92    }
93}
94
95/// Helper to retrieve a Neo4j graph connection.
96pub fn get_neo4j_graph() -> Result<Arc<Mutex<Graph>>, &'static str> {
97    let neo4j_connector = NEO4J_CONNECTOR
98        .get()
99        .ok_or("Neo4jConnector not initialized")?;
100    let graph = neo4j_connector
101        .graph
102        .get()
103        .ok_or("Not connected to Neo4j")?;
104    Ok(graph.clone())
105}
106
107pub static NEO4J_CONNECTOR: OnceCell<Neo4jConnector> = OnceCell::new();