Skip to main content

sz_orm_graph/
connection.rs

1//! # Connection — Bolt 协议连接与连接池
2//!
3//! GraphConfig + GraphConnection + GraphPool
4
5use crate::error::{sanitize_dsn, GraphError};
6use std::sync::Arc;
7use std::time::Duration;
8use tokio::sync::Mutex;
9
10/// 图数据库连接配置
11#[derive(Debug, Clone)]
12pub struct GraphConfig {
13    /// Bolt DSN,如 `neo4j://neo4j:password@127.0.0.1:7687`
14    pub dsn: String,
15    /// 连接超时(秒)
16    pub connect_timeout_secs: u64,
17    /// 查询超时(秒)
18    pub query_timeout_secs: u64,
19    /// 连接池最大大小
20    pub max_pool_size: usize,
21}
22
23impl GraphConfig {
24    pub fn new(dsn: &str) -> Self {
25        Self {
26            dsn: dsn.to_string(),
27            connect_timeout_secs: 10,
28            query_timeout_secs: 30,
29            max_pool_size: 10,
30        }
31    }
32
33    pub fn with_connect_timeout(mut self, secs: u64) -> Self {
34        self.connect_timeout_secs = secs;
35        self
36    }
37
38    pub fn with_query_timeout(mut self, secs: u64) -> Self {
39        self.query_timeout_secs = secs;
40        self
41    }
42
43    pub fn with_pool_size(mut self, size: usize) -> Self {
44        self.max_pool_size = size;
45        self
46    }
47
48    /// 脱敏的 DSN(不泄露密码)
49    pub fn sanitized_dsn(&self) -> String {
50        sanitize_dsn(&self.dsn)
51    }
52}
53
54/// Bolt 连接句柄
55pub struct GraphConnection {
56    config: GraphConfig,
57    connected: bool,
58}
59
60impl GraphConnection {
61    pub fn new(config: GraphConfig) -> Self {
62        Self {
63            config,
64            connected: false,
65        }
66    }
67
68    pub fn config(&self) -> &GraphConfig {
69        &self.config
70    }
71
72    pub fn is_connected(&self) -> bool {
73        self.connected
74    }
75
76    pub fn connect(&mut self) -> Result<(), GraphError> {
77        if self.config.dsn.is_empty() {
78            return Err(GraphError::ConnectionError("empty DSN".into()));
79        }
80        if !self.config.dsn.starts_with("neo4j://") && !self.config.dsn.starts_with("bolt://") {
81            return Err(GraphError::ConnectionError(format!(
82                "invalid DSN scheme: {}",
83                self.config.sanitized_dsn()
84            )));
85        }
86        self.connected = true;
87        Ok(())
88    }
89
90    pub fn disconnect(&mut self) {
91        self.connected = false;
92    }
93}
94
95/// 图数据库连接池
96pub struct GraphPool {
97    config: GraphConfig,
98    connections: Arc<Mutex<Vec<GraphConnection>>>,
99}
100
101impl GraphPool {
102    pub fn new(config: GraphConfig) -> Self {
103        Self {
104            config,
105            connections: Arc::new(Mutex::new(Vec::new())),
106        }
107    }
108
109    pub fn config(&self) -> &GraphConfig {
110        &self.config
111    }
112
113    pub async fn acquire(&self) -> Result<GraphConnection, GraphError> {
114        let mut conns = self.connections.lock().await;
115        if let Some(conn) = conns.pop() {
116            return Ok(conn);
117        }
118        if conns.len() >= self.config.max_pool_size {
119            return Err(GraphError::ConnectionError(format!(
120                "pool exhausted (max={}), DSN: {}",
121                self.config.max_pool_size,
122                self.config.sanitized_dsn()
123            )));
124        }
125        let mut conn = GraphConnection::new(self.config.clone());
126        conn.connect()?;
127        Ok(conn)
128    }
129
130    pub async fn release(&self, conn: GraphConnection) {
131        let mut conns = self.connections.lock().await;
132        conns.push(conn);
133    }
134
135    pub async fn size(&self) -> usize {
136        self.connections.lock().await.len()
137    }
138}
139
140impl GraphConfig {
141    pub fn connect_timeout(&self) -> Duration {
142        Duration::from_secs(self.connect_timeout_secs)
143    }
144
145    pub fn query_timeout(&self) -> Duration {
146        Duration::from_secs(self.query_timeout_secs)
147    }
148}