Skip to main content

surreal_client/
connection.rs

1//! Connection builder for SurrealDB with authentication and engine creation
2
3use crate::{DebugEngine, Engine, Result, SurrealClient, SurrealError, WsCborEngine};
4
5use serde_json::Value;
6use url::Url;
7
8/// Connection builder for SurrealDB
9#[derive(Default, Debug, Clone)]
10pub struct SurrealConnection {
11    /// URL to connect to
12    pub url: Option<String>,
13
14    /// Namespace to use
15    namespace: Option<String>,
16
17    /// Database to use
18    database: Option<String>,
19
20    /// Authentication credentials
21    auth: Option<AuthParams>,
22
23    /// Whether to check SurrealDB version compatibility
24    version_check: bool,
25
26    /// Whether to enable debug mode for query logging
27    debug: bool,
28}
29
30/// Authentication parameters
31#[derive(Debug, Clone)]
32pub enum AuthParams {
33    /// Root authentication
34    Root { username: String, password: String },
35    /// Namespace authentication
36    Namespace { username: String, password: String },
37    /// Database authentication
38    Database { username: String, password: String },
39    /// Scope authentication
40    Scope {
41        namespace: String,
42        database: String,
43        scope: String,
44        params: Value,
45    },
46    /// JWT token authentication
47    Token(String),
48}
49
50impl SurrealConnection {
51    /// Create a new connection builder
52    pub fn new() -> Self {
53        Self {
54            version_check: true,
55            debug: false,
56            ..Default::default()
57        }
58    }
59
60    /// Parse connection from DSN string
61    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        // Ensure URL has a proper host
66        if url.host().is_none() {
67            return Err(SurrealError::Connection(
68                "URL must have a valid host".to_string(),
69            ));
70        }
71
72        // Store the URL without user credentials and path/query
73        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        // Extract user credentials for root auth
79        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        // Extract namespace and database from path segments
86        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        // Parse query parameters
96        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                // Signin level for the DSN credentials. Without it, DSN
104                // credentials sign in as root — a namespace- or
105                // database-defined user needs `?auth=namespace` /
106                // `?auth=database` (cloud instances have no root users).
107                "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    /// Set the URL to connect to
134    pub fn url(mut self, url: impl Into<String>) -> Self {
135        self.url = Some(url.into());
136        self
137    }
138
139    /// Set the namespace
140    pub fn namespace(mut self, namespace: impl Into<String>) -> Self {
141        self.namespace = Some(namespace.into());
142        self
143    }
144
145    /// Set the database
146    pub fn database(mut self, database: impl Into<String>) -> Self {
147        self.database = Some(database.into());
148        self
149    }
150
151    /// Set root authentication
152    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    /// Set namespace authentication
161    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    /// Set database authentication
174    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    /// Set scope authentication
187    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    /// Set JWT token authentication
204    pub fn auth_token(mut self, token: impl Into<String>) -> Self {
205        self.auth = Some(AuthParams::Token(token.into()));
206        self
207    }
208
209    /// Set version check flag
210    pub fn version_check(mut self, check: bool) -> Self {
211        self.version_check = check;
212        self
213    }
214
215    /// Enable debug mode for query logging
216    pub fn with_debug(mut self, enabled: bool) -> Self {
217        self.debug = enabled;
218        self
219    }
220
221    // /// Configure connection pool with custom settings
222    // pub fn with_pool_config(mut self, config: PoolConfig) -> Self {
223    //     self.pool_config = Some(config);
224    //     self
225    // }
226
227    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                // JWT auth (used for SurrealDB Cloud instances, whose per-instance
292                // access token is brokered out of band). `authenticate` takes the
293                // raw token; the `use` step below selects ns/db when configured.
294                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        // Credential logins (root / NS / DB users) select their working scope
307        // with an explicit `use`. A JWT (`authenticate`) is already scoped to the
308        // namespace/database carried in its claims, and re-selecting the namespace
309        // is a privileged action a database-scoped access actor isn't permitted to
310        // perform — SurrealDB answers `use` with an IAM `NotAllowed`. So skip `use`
311        // for token auth and rely on the token's own scope.
312        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    /// Connect to SurrealDB and return an immutable client
325    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        // Example of the new flow: Connection -> authenticate -> creates engine -> returns immutable client
425
426        // This would be the typical usage:
427        // let client = Connection::new()
428        //     .url("ws://localhost:8000")
429        //     .namespace("bakery")
430        //     .database("inventory")
431        //     .auth_root("root", "root")
432        //     .connect()
433        //     .await
434        //     .unwrap();
435
436        // For testing, we just verify the builder pattern works
437        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        // The client would be immutable once created:
451        // - client.query() - no mut needed
452        // - client.select() - no mut needed
453        // - client.let_var() - changes session but client stays immutable
454        // - Multiple clients can be cloned, each with unique session
455    }
456}