sz_orm_graph/
connection.rs1use crate::error::{sanitize_dsn, GraphError};
6use std::sync::Arc;
7use std::time::Duration;
8use tokio::sync::Mutex;
9
10#[derive(Debug, Clone)]
12pub struct GraphConfig {
13 pub dsn: String,
15 pub connect_timeout_secs: u64,
17 pub query_timeout_secs: u64,
19 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 pub fn sanitized_dsn(&self) -> String {
50 sanitize_dsn(&self.dsn)
51 }
52}
53
54pub 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
95pub 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}