1use crate::{DebugEngine, Engine, Result, SurrealClient, SurrealError, WsCborEngine};
4
5use serde_json::Value;
6use url::Url;
7
8#[derive(Default, Debug, Clone)]
10pub struct SurrealConnection {
11 pub url: Option<String>,
13
14 namespace: Option<String>,
16
17 database: Option<String>,
19
20 auth: Option<AuthParams>,
22
23 version_check: bool,
25
26 debug: bool,
28}
29
30#[derive(Debug, Clone)]
32pub enum AuthParams {
33 Root { username: String, password: String },
35 Namespace { username: String, password: String },
37 Database { username: String, password: String },
39 Scope {
41 namespace: String,
42 database: String,
43 scope: String,
44 params: Value,
45 },
46 Token(String),
48}
49
50impl SurrealConnection {
51 pub fn new() -> Self {
53 Self {
54 version_check: true,
55 debug: false,
56 ..Default::default()
57 }
58 }
59
60 pub fn dsn(dsn: impl AsRef<str>) -> Result<Self> {
62 let mut conn = Self::new();
63 let url = Url::parse(dsn.as_ref())?;
64
65 if url.host().is_none() {
67 return Err(SurrealError::Connection(
68 "URL must have a valid host".to_string(),
69 ));
70 }
71
72 let base_url = format!("{}://{}", url.scheme(), url.host_str().unwrap());
74 let port = url.port().map(|p| format!(":{}", p)).unwrap_or_default();
75 let final_url = format!("{}{}", base_url, port);
76 conn.url = Some(final_url);
77
78 if !url.username().is_empty() {
80 let username = url.username().to_string();
81 let password = url.password().unwrap_or("").to_string();
82 conn.auth = Some(AuthParams::Root { username, password });
83 }
84
85 let path_segments: Vec<&str> = url.path_segments().map(|c| c.collect()).unwrap_or_default();
87
88 if let Some(namespace) = path_segments.first().filter(|s| !s.is_empty()) {
89 conn.namespace = Some(namespace.to_string());
90 }
91 if let Some(database) = path_segments.get(1).filter(|s| !s.is_empty()) {
92 conn.database = Some(database.to_string());
93 }
94
95 for (key, value) in url.query_pairs() {
97 match key.as_ref() {
98 "namespace" => conn.namespace = Some(value.into_owned()),
99 "database" => conn.database = Some(value.into_owned()),
100 "version_check" => {
101 conn.version_check = value.parse().unwrap_or(true);
102 }
103 "auth" | "auth_level" => {
108 let (username, password) = match conn.auth.take() {
109 Some(AuthParams::Root { username, password }) => (username, password),
110 other => {
111 conn.auth = other;
112 continue;
113 }
114 };
115 conn.auth = Some(match value.as_ref() {
116 "namespace" | "ns" => AuthParams::Namespace { username, password },
117 "database" | "db" => AuthParams::Database { username, password },
118 "root" => AuthParams::Root { username, password },
119 other => {
120 return Err(SurrealError::Connection(format!(
121 "Unknown auth level `{other}` — use root, namespace, or database"
122 )));
123 }
124 });
125 }
126 _ => {}
127 }
128 }
129
130 Ok(conn)
131 }
132
133 pub fn url(mut self, url: impl Into<String>) -> Self {
135 self.url = Some(url.into());
136 self
137 }
138
139 pub fn namespace(mut self, namespace: impl Into<String>) -> Self {
141 self.namespace = Some(namespace.into());
142 self
143 }
144
145 pub fn database(mut self, database: impl Into<String>) -> Self {
147 self.database = Some(database.into());
148 self
149 }
150
151 pub fn auth_root(mut self, username: impl Into<String>, password: impl Into<String>) -> Self {
153 self.auth = Some(AuthParams::Root {
154 username: username.into(),
155 password: password.into(),
156 });
157 self
158 }
159
160 pub fn auth_namespace(
162 mut self,
163 username: impl Into<String>,
164 password: impl Into<String>,
165 ) -> Self {
166 self.auth = Some(AuthParams::Namespace {
167 username: username.into(),
168 password: password.into(),
169 });
170 self
171 }
172
173 pub fn auth_database(
175 mut self,
176 username: impl Into<String>,
177 password: impl Into<String>,
178 ) -> Self {
179 self.auth = Some(AuthParams::Database {
180 username: username.into(),
181 password: password.into(),
182 });
183 self
184 }
185
186 pub fn auth_scope(
188 mut self,
189 namespace: impl Into<String>,
190 database: impl Into<String>,
191 scope: impl Into<String>,
192 params: Value,
193 ) -> Self {
194 self.auth = Some(AuthParams::Scope {
195 namespace: namespace.into(),
196 database: database.into(),
197 scope: scope.into(),
198 params,
199 });
200 self
201 }
202
203 pub fn auth_token(mut self, token: impl Into<String>) -> Self {
205 self.auth = Some(AuthParams::Token(token.into()));
206 self
207 }
208
209 pub fn version_check(mut self, check: bool) -> Self {
211 self.version_check = check;
212 self
213 }
214
215 pub fn with_debug(mut self, enabled: bool) -> Self {
217 self.debug = enabled;
218 self
219 }
220
221 pub(crate) async fn init_engine(&self, engine: &mut crate::WsCborEngine) -> Result<()> {
228 use ciborium::Value as CborValue;
229
230 match self.auth.as_ref().ok_or(SurrealError::Connection(
231 "Attempted to connect without auth".to_string(),
232 ))? {
233 AuthParams::Root { username, password } => {
234 let auth_params = CborValue::Array(vec![CborValue::Map(vec![
235 (
236 CborValue::Text("user".to_string()),
237 CborValue::Text(username.clone()),
238 ),
239 (
240 CborValue::Text("pass".to_string()),
241 CborValue::Text(password.clone()),
242 ),
243 ])]);
244 engine.send_message_cbor("signin", auth_params).await?;
245 }
246 AuthParams::Namespace { username, password } => {
247 let namespace = self.namespace.clone().ok_or(SurrealError::Connection(
248 "Namespace is required for namespace auth".to_string(),
249 ))?;
250 let auth_params = CborValue::Array(vec![CborValue::Map(vec![
251 (
252 CborValue::Text("user".to_string()),
253 CborValue::Text(username.clone()),
254 ),
255 (
256 CborValue::Text("pass".to_string()),
257 CborValue::Text(password.clone()),
258 ),
259 (
260 CborValue::Text("NS".to_string()),
261 CborValue::Text(namespace),
262 ),
263 ])]);
264 engine.send_message_cbor("signin", auth_params).await?;
265 }
266 AuthParams::Database { username, password } => {
267 let namespace = self.namespace.clone().ok_or(SurrealError::Connection(
268 "Namespace is required for database auth".to_string(),
269 ))?;
270 let database = self.database.clone().ok_or(SurrealError::Connection(
271 "Database is required for database auth".to_string(),
272 ))?;
273 let auth_params = CborValue::Array(vec![CborValue::Map(vec![
274 (
275 CborValue::Text("user".to_string()),
276 CborValue::Text(username.clone()),
277 ),
278 (
279 CborValue::Text("pass".to_string()),
280 CborValue::Text(password.clone()),
281 ),
282 (
283 CborValue::Text("NS".to_string()),
284 CborValue::Text(namespace),
285 ),
286 (CborValue::Text("DB".to_string()), CborValue::Text(database)),
287 ])]);
288 engine.send_message_cbor("signin", auth_params).await?;
289 }
290 AuthParams::Token(token) => {
291 let auth_params = CborValue::Array(vec![CborValue::Text(token.clone())]);
295 engine
296 .send_message_cbor("authenticate", auth_params)
297 .await?;
298 }
299 _ => {
300 return Err(SurrealError::Connection(
301 "Unsupported authentication method".to_string(),
302 ));
303 }
304 }
305
306 let is_token_auth = matches!(self.auth, Some(AuthParams::Token(_)));
313 if !is_token_auth && let Some(namespace) = &self.namespace {
314 let use_params = CborValue::Array(vec![
315 CborValue::Text(namespace.clone()),
316 CborValue::Text(self.database.as_ref().unwrap_or(&String::new()).clone()),
317 ]);
318 engine.send_message_cbor("use", use_params).await?;
319 }
320
321 Ok(())
322 }
323
324 pub async fn connect(self) -> Result<SurrealClient> {
326 let url_str = self
327 .url
328 .as_ref()
329 .ok_or_else(|| SurrealError::Connection("URL is required".to_string()))?;
330 let url = Url::parse(url_str)
331 .map_err(|e| SurrealError::Connection(format!("Invalid URL: {}", e)))?;
332
333 let mut engine: Box<dyn Engine> = match url.scheme() {
334 "ws" | "wss" | "cbor" => Box::new(WsCborEngine::from_connection(&self).await?),
335 _ => {
336 return Err(SurrealError::Protocol(
337 "Unsupported protocol. Use ws://, wss://, or cbor://".to_string(),
338 ));
339 }
340 };
341
342 if self.debug {
343 engine = DebugEngine::wrap(engine);
344 }
345
346 let client = SurrealClient::new(engine, self.namespace, self.database);
347 Ok(client.with_debug(self.debug))
348 }
349}
350
351#[cfg(test)]
352mod tests {
353 use super::*;
354
355 #[test]
356 fn test_connection_builder() {
357 let conn = SurrealConnection::new()
358 .url("ws://localhost:8000")
359 .namespace("test_ns")
360 .database("test_db")
361 .auth_root("root", "root")
362 .version_check(false);
363
364 assert_eq!(conn.url, Some("ws://localhost:8000".to_string()));
365 assert_eq!(conn.namespace, Some("test_ns".to_string()));
366 assert_eq!(conn.database, Some("test_db".to_string()));
367 assert!(!conn.version_check);
368 assert!(matches!(conn.auth, Some(AuthParams::Root { .. })));
369 }
370
371 #[test]
372 fn test_dsn_auth_level() {
373 let conn =
374 SurrealConnection::dsn("cbor://user:pw@cloud.example/ns/db?auth=namespace").unwrap();
375 assert!(matches!(conn.auth, Some(AuthParams::Namespace { .. })));
376
377 let conn = SurrealConnection::dsn("cbor://user:pw@cloud.example/ns/db?auth=db").unwrap();
378 assert!(matches!(conn.auth, Some(AuthParams::Database { .. })));
379
380 assert!(SurrealConnection::dsn("cbor://user:pw@h/n/d?auth=nope").is_err());
381 }
382
383 #[test]
384 fn test_dsn_parsing() {
385 let conn = SurrealConnection::dsn(
386 "ws://root:root@localhost:8000/test_ns/test_db?version_check=false",
387 )
388 .unwrap();
389
390 assert_eq!(conn.url, Some("ws://localhost:8000".to_string()));
391 assert_eq!(conn.namespace, Some("test_ns".to_string()));
392 assert_eq!(conn.database, Some("test_db".to_string()));
393 assert!(!conn.version_check);
394 assert!(matches!(conn.auth, Some(AuthParams::Root { .. })));
395 }
396
397 #[test]
398 fn test_dsn_with_query_params() {
399 let conn =
400 SurrealConnection::dsn("http://localhost:8000?namespace=ns&database=db").unwrap();
401
402 assert_eq!(conn.url, Some("http://localhost:8000".to_string()));
403 assert_eq!(conn.namespace, Some("ns".to_string()));
404 assert_eq!(conn.database, Some("db".to_string()));
405 }
406
407 #[test]
408 fn test_auth_methods() {
409 let conn1 = SurrealConnection::new().auth_root("admin", "pass");
410 assert!(matches!(conn1.auth, Some(AuthParams::Root { .. })));
411
412 let conn2 = SurrealConnection::new().auth_namespace("ns_user", "ns_pass");
413 assert!(matches!(conn2.auth, Some(AuthParams::Namespace { .. })));
414
415 let conn3 = SurrealConnection::new().auth_database("db_user", "db_pass");
416 assert!(matches!(conn3.auth, Some(AuthParams::Database { .. })));
417
418 let conn4 = SurrealConnection::new().auth_token("jwt_token");
419 assert!(matches!(conn4.auth, Some(AuthParams::Token(_))));
420 }
421
422 #[tokio::test]
423 async fn test_connection_to_client_flow() {
424 let connection = SurrealConnection::new()
438 .url("ws://localhost:8000")
439 .namespace("test_namespace")
440 .database("test_database")
441 .auth_root("admin", "password")
442 .version_check(false);
443
444 assert_eq!(connection.url, Some("ws://localhost:8000".to_string()));
445 assert_eq!(connection.namespace, Some("test_namespace".to_string()));
446 assert_eq!(connection.database, Some("test_database".to_string()));
447 assert!(!connection.version_check);
448 assert!(matches!(connection.auth, Some(AuthParams::Root { .. })));
449
450 }
456}