Skip to main content

surreal_client/
client.rs

1use std::sync::Arc;
2
3use ciborium::Value as CborValue;
4use serde_json::{Value, json};
5
6use crate::live::LiveStream;
7use crate::{Engine, RecordId, RecordRange, Result, SessionState, SurrealError, Table};
8
9pub struct SurrealClient {
10    engine: Arc<tokio::sync::Mutex<Box<dyn Engine>>>,
11    session: SessionState,
12    incremental_id: Arc<std::sync::atomic::AtomicU64>,
13    debug: bool,
14}
15
16impl Clone for SurrealClient {
17    /// Clone the client - creates a new client instance sharing the same engine and session
18    fn clone(&self) -> Self {
19        Self {
20            engine: self.engine.clone(),
21            session: self.session.clone(),
22            incremental_id: self.incremental_id.clone(),
23            debug: self.debug,
24        }
25    }
26}
27
28impl SurrealClient {
29    /// Create a new SurrealDB instance with the given engine and optional namespace/database
30    pub fn new(
31        engine: Box<dyn Engine>,
32        namespace: Option<String>,
33        database: Option<String>,
34    ) -> Self {
35        let mut session = SessionState::new();
36        session.set_target(namespace, database);
37
38        Self {
39            engine: Arc::new(tokio::sync::Mutex::new(engine)),
40            session,
41            incremental_id: Arc::new(std::sync::atomic::AtomicU64::new(0)),
42            debug: false,
43        }
44    }
45
46    /// Enable debug mode to log queries
47    pub fn with_debug(mut self, enabled: bool) -> Self {
48        self.debug = enabled;
49        self
50    }
51
52    /// Check if debug mode is enabled
53    pub fn is_debug(&self) -> bool {
54        self.debug
55    }
56
57    /// Set a parameter for the session
58    pub async fn let_var(&mut self, key: &str, value: Value) -> Result<()> {
59        let mut engine = self.engine.lock().await;
60
61        let params = json!([key, value]);
62
63        engine.send_message("let", params).await?;
64
65        // Store the variable in the session
66        self.session.set_param(key.to_string(), value);
67
68        Ok(())
69    }
70
71    /// Unset a parameter from the session
72    pub async fn unset(&mut self, key: &str) -> Result<()> {
73        let mut engine = self.engine.lock().await;
74
75        let params = json!([key]);
76
77        engine.send_message("unset", params).await?;
78
79        // Remove the variable from the session
80        self.session.unset_param(key);
81        Ok(())
82    }
83
84    /// Create a record in the database
85    pub async fn create(&self, resource: &str, data: Option<Value>) -> Result<Value> {
86        let mut engine = self.engine.lock().await;
87
88        let params = if let Some(data) = data {
89            json!([resource, data])
90        } else {
91            json!([resource])
92        };
93
94        engine.send_message("create", params).await
95    }
96
97    /// Select records from the database
98    pub async fn select(&self, resource: &str) -> Result<Value> {
99        let mut engine = self.engine.lock().await;
100
101        let params = json!([resource]);
102
103        engine.send_message("select", params).await
104    }
105
106    /// Select all records from a table
107    pub async fn select_all(&self, table: Table) -> Result<Value> {
108        self.select(&table.to_string()).await
109    }
110
111    /// Select a specific record by ID
112    pub async fn select_record(&self, record: RecordId) -> Result<Value> {
113        self.select(&record.to_string()).await
114    }
115
116    /// Select a range of records
117    pub async fn select_range(&self, range: RecordRange) -> Result<Value> {
118        self.select(&range.to_string()).await
119    }
120
121    /// Update records in the database
122    pub async fn update(&self, resource: &str, data: Option<Value>) -> Result<Value> {
123        let mut engine = self.engine.lock().await;
124
125        let params = if let Some(data) = data {
126            json!([resource, data])
127        } else {
128            json!([resource])
129        };
130
131        engine.send_message("update", params).await
132    }
133
134    /// Update a specific record by ID
135    pub async fn update_record(&self, record: RecordId, data: Value) -> Result<Value> {
136        self.update(&record.to_string(), Some(data)).await
137    }
138
139    /// Update all records in a table
140    pub async fn update_all(&self, table: Table, data: Value) -> Result<Value> {
141        self.update(&table.to_string(), Some(data)).await
142    }
143
144    /// Upsert (insert or update) records in the database
145    pub async fn upsert(&self, resource: &str, data: Option<Value>) -> Result<Value> {
146        let mut engine = self.engine.lock().await;
147
148        let params = if let Some(data) = data {
149            json!([resource, data])
150        } else {
151            json!([resource])
152        };
153
154        engine.send_message("upsert", params).await
155    }
156
157    /// Upsert a specific record by ID
158    pub async fn upsert_record(&self, record: RecordId, data: Value) -> Result<Value> {
159        self.upsert(&record.to_string(), Some(data)).await
160    }
161
162    /// Merge data into records in the database
163    pub async fn merge(&self, resource: &str, data: Value) -> Result<Value> {
164        let mut engine = self.engine.lock().await;
165
166        let params = json!([resource, data]);
167
168        engine.send_message("merge", params).await
169    }
170
171    /// Merge data into a specific record by ID
172    pub async fn merge_record(&self, record: RecordId, data: Value) -> Result<Value> {
173        self.merge(&record.to_string(), data).await
174    }
175
176    /// Merge data into all records in a table
177    pub async fn merge_all(&self, table: Table, data: Value) -> Result<Value> {
178        self.merge(&table.to_string(), data).await
179    }
180
181    /// Apply JSON patches to records
182    /// Apply patches to records in the database
183    pub async fn patch(&self, resource: &str, patches: Vec<Value>) -> Result<Value> {
184        let mut engine = self.engine.lock().await;
185
186        let params = json!([resource, patches]);
187
188        engine.send_message("patch", params).await
189    }
190
191    /// Delete records from the database
192    pub async fn delete(&self, resource: &str) -> Result<Value> {
193        let mut engine = self.engine.lock().await;
194
195        let params = json!([resource]);
196
197        engine.send_message("delete", params).await
198    }
199
200    /// Delete a specific record by ID
201    pub async fn delete_record(&self, record: RecordId) -> Result<Value> {
202        self.delete(&record.to_string()).await
203    }
204
205    /// Delete all records from a table
206    pub async fn delete_all(&self, table: Table) -> Result<Value> {
207        self.delete(&table.to_string()).await
208    }
209
210    /// Insert records into the database
211    /// Insert data into a table
212    pub async fn insert(&self, table: &str, data: Value) -> Result<Value> {
213        let mut engine = self.engine.lock().await;
214
215        let params = json!([table, data]);
216
217        engine.send_message("insert", params).await
218    }
219
220    /// Insert multiple records
221    pub async fn insert_many(&self, table: Table, data: Vec<Value>) -> Result<Value> {
222        // TODO: add single test
223        self.insert(&table.to_string(), Value::Array(data)).await
224    }
225
226    /// Create a relation between records
227    pub async fn relate(
228        &self,
229        from: &str,
230        relation: &str,
231        to: &str,
232        data: Option<Value>,
233    ) -> Result<Value> {
234        let mut engine = self.engine.lock().await;
235
236        let params = if let Some(data) = data {
237            json!([from, relation, to, data])
238        } else {
239            json!([from, relation, to])
240        };
241
242        engine.send_message("relate", params).await
243    }
244
245    /// Create a relation between specific records
246    pub async fn relate_records(
247        &self,
248        from: RecordId,
249        relation: Table,
250        to: RecordId,
251        data: Option<Value>,
252    ) -> Result<Value> {
253        self.relate(
254            &from.to_string(),
255            &relation.to_string(),
256            &to.to_string(),
257            data,
258        )
259        .await
260    }
261
262    /// Run a stored function
263    pub async fn run(&self, func: &str, args: Option<Value>) -> Result<Value> {
264        let mut engine = self.engine.lock().await;
265
266        let params = if let Some(args) = args {
267            json!([func, args])
268        } else {
269            json!([func])
270        };
271
272        engine.send_message("run", params).await
273    }
274
275    /// Execute a custom SurrealQL query
276    pub async fn query(&self, sql: &str, variables: Option<Value>) -> Result<Value> {
277        if self.debug {
278            if let Some(ref vars) = variables {
279                println!("🔍 SQL: {}", sql);
280                println!(
281                    "📊 Params: {}",
282                    serde_json::to_string_pretty(vars).unwrap_or_default()
283                );
284            } else {
285                println!("🔍 SQL: {}", sql);
286            }
287        }
288
289        let mut engine = self.engine.lock().await;
290
291        let params = if let Some(vars) = variables {
292            json!([sql, vars])
293        } else {
294            json!([sql])
295        };
296
297        let response = engine.send_message("query", params).await?;
298
299        if self.debug {
300            // Check if response contains status field to determine icon
301            let icon = if let Value::Array(ref results) = response {
302                if results
303                    .iter()
304                    .any(|r| r.get("status").and_then(|s| s.as_str()) == Some("ERR"))
305                {
306                    "❌"
307                } else {
308                    "✅"
309                }
310            } else {
311                "✅"
312            };
313
314            println!(
315                "{} Response: {}",
316                icon,
317                serde_json::to_string_pretty(&response).unwrap_or_default()
318            );
319        }
320
321        // Handle the query response format
322        match response {
323            Value::Array(results) => {
324                // Return the results array directly
325                Ok(Value::Array(results))
326            }
327            other => Ok(other),
328        }
329    }
330
331    /// Get information about the current session
332    pub async fn info(&self) -> Result<Value> {
333        let mut engine = self.engine.lock().await;
334
335        let params = json!([]);
336
337        engine.send_message("info", params).await
338    }
339
340    /// Get the version of the SurrealDB instance
341    pub async fn version(&self) -> Result<String> {
342        let mut engine = self.engine.lock().await;
343
344        let params = json!([]);
345
346        let response = engine.send_message("version", params).await?;
347
348        match response {
349            Value::String(version) => Ok(version),
350            _ => Err(SurrealError::Protocol(
351                "Invalid version response format".to_string(),
352            )),
353        }
354    }
355
356    /// Close the connection
357    pub async fn close(self) -> Result<()> {
358        // Note: engine is moved here since we're taking ownership
359        // The session will be dropped automatically
360        // Engine trait doesn't have close method in minimal implementation
361        Ok(())
362    }
363
364    /// Import database content (HTTP only)
365    pub async fn import(&self, _content: &str, _username: &str, _password: &str) -> Result<Value> {
366        Err(SurrealError::Protocol(
367            "Import is not supported in minimal engine implementation".to_string(),
368        ))
369    }
370
371    /// Export database content (HTTP only)
372    pub async fn export(&self, _username: &str, _password: &str) -> Result<String> {
373        Err(SurrealError::Protocol(
374            "Export is not supported in minimal engine implementation".to_string(),
375        ))
376    }
377
378    /// Import ML model (HTTP only)
379    pub async fn import_ml(
380        &self,
381        _content: &str,
382        _username: Option<&str>,
383        _password: Option<&str>,
384    ) -> Result<Value> {
385        Err(SurrealError::Protocol(
386            "ML import is not supported in minimal engine implementation".to_string(),
387        ))
388    }
389
390    /// Export ML model (HTTP only)
391    pub async fn export_ml(
392        &self,
393        _name: &str,
394        _version: Option<&str>,
395        _username: Option<&str>,
396        _password: Option<&str>,
397    ) -> Result<String> {
398        Err(SurrealError::Protocol(
399            "ML export is not supported in minimal engine implementation".to_string(),
400        ))
401    }
402
403    /// Start a live query on a table (or record) and stream change
404    /// notifications.
405    ///
406    /// Issues the `live` RPC, then registers the returned live-query id with
407    /// the engine so its read loop forwards every matching `CREATE`/`UPDATE`/
408    /// `DELETE` frame onto the returned [`LiveStream`]. Requires a WebSocket
409    /// connection — request/response-only engines return an error from
410    /// [`Engine::register_live`].
411    ///
412    /// The stream stays open until the connection closes; call [`kill`](Self::kill)
413    /// with the stream's [`query_id`](LiveStream::query_id) to release the
414    /// server-side query early.
415    pub async fn live(&self, resource: &str) -> Result<LiveStream> {
416        let mut engine = self.engine.lock().await;
417
418        let params = CborValue::Array(vec![CborValue::Text(resource.to_string())]);
419        let response = engine.send_message_cbor("live", params).await?;
420
421        let query_id = cbor_uuid_string(&response).ok_or_else(|| {
422            SurrealError::Protocol(format!(
423                "live query did not return a usable id: {:?}",
424                response
425            ))
426        })?;
427
428        let rx = engine.register_live(&query_id).await?;
429        Ok(LiveStream { query_id, rx })
430    }
431
432    /// Stop a live query by its id (the [`LiveStream::query_id`]).
433    ///
434    /// Sends the `kill` RPC and drops the local subscriber so no further
435    /// notifications are delivered for that id.
436    pub async fn kill(&self, query_id: &str) -> Result<()> {
437        let mut engine = self.engine.lock().await;
438        let params = CborValue::Array(vec![CborValue::Text(query_id.to_string())]);
439        engine.send_message_cbor("kill", params).await?;
440        engine.unregister_live(query_id).await;
441        Ok(())
442    }
443
444    /// Execute a custom SurrealQL query with CBOR parameters
445    pub async fn query_cbor(&self, sql: &str, variables: CborValue) -> Result<CborValue> {
446        let mut engine = self.engine.lock().await;
447
448        if self.debug {
449            println!("SQL: {}", sql);
450            println!("Params: {:?}", variables);
451        }
452
453        let params = CborValue::Array(vec![CborValue::Text(sql.to_string()), variables]);
454        let response = engine.send_message_cbor("query", params).await?;
455
456        if self.debug {
457            println!("✅ CBOR Response: {:?}", response);
458        }
459
460        Ok(response)
461    }
462}
463
464/// Render the `live` RPC result (a CBOR UUID) into the same string form the
465/// engine derives from notification frame ids, so a live subscription matches
466/// its notifications. SurrealDB encodes UUIDs as tag 37 over a 16-byte string;
467/// a plain text id is accepted too.
468fn cbor_uuid_string(v: &CborValue) -> Option<String> {
469    match v {
470        CborValue::Text(t) => Some(t.clone()),
471        CborValue::Tag(_, inner) => cbor_uuid_string(inner),
472        CborValue::Bytes(b) if b.len() == 16 => Some(
473            uuid::Uuid::from_slice(b)
474                .map(|u| u.to_string())
475                .unwrap_or_else(|_| hex::encode(b)),
476        ),
477        CborValue::Bytes(b) => Some(hex::encode(b)),
478        _ => None,
479    }
480}
481
482#[cfg(test)]
483mod tests {
484    use super::*;
485    use serde_json::json;
486
487    // Mock engine for testing
488    struct MockEngine;
489
490    #[async_trait::async_trait]
491    impl Engine for MockEngine {
492        async fn send_message(&mut self, _method: &str, _params: Value) -> Result<Value> {
493            Ok(Value::String("mock_response".to_string()))
494        }
495
496        async fn send_message_cbor(
497            &mut self,
498            _method: &str,
499            _params: CborValue,
500        ) -> Result<CborValue> {
501            Ok(CborValue::Text("mock_response".to_string()))
502        }
503    }
504
505    #[tokio::test]
506    async fn test_surrealdb_creation() {
507        let engine = Box::new(MockEngine);
508        let _client = SurrealClient::new(engine, None, None);
509    }
510
511    #[tokio::test]
512    async fn test_connect_and_operations() {
513        let engine = Box::new(MockEngine);
514        let client = SurrealClient::new(
515            engine,
516            Some("test_ns".to_string()),
517            Some("test_db".to_string()),
518        );
519
520        // Test basic operations
521        let result = client.select("user").await.unwrap();
522        assert_eq!(result, Value::String("mock_response".to_string()));
523
524        let result = client
525            .create("user", Some(json!({"name": "John"})))
526            .await
527            .unwrap();
528        assert_eq!(result, Value::String("mock_response".to_string()));
529    }
530
531    #[tokio::test]
532    async fn test_crud_operations() {
533        let engine = Box::new(MockEngine);
534        let client = SurrealClient::new(
535            engine,
536            Some("test_ns".to_string()),
537            Some("test_db".to_string()),
538        );
539
540        // Test Create
541        let create_result = client
542            .create("users", Some(json!({"name": "Alice", "age": 30})))
543            .await
544            .unwrap();
545        assert_eq!(create_result, Value::String("mock_response".to_string()));
546
547        // Test Read
548        let read_result = client.select("users").await.unwrap();
549        assert_eq!(read_result, Value::String("mock_response".to_string()));
550
551        // Test Update
552        let update_result = client
553            .update("users:alice", Some(json!({"age": 31})))
554            .await
555            .unwrap();
556        assert_eq!(update_result, Value::String("mock_response".to_string()));
557
558        // Test Delete
559        let delete_result = client.delete("users:alice").await.unwrap();
560        assert_eq!(delete_result, Value::String("mock_response".to_string()));
561
562        // Test Insert
563        let insert_result = client
564            .insert("users", json!({"name": "Bob", "age": 25}))
565            .await
566            .unwrap();
567        assert_eq!(insert_result, Value::String("mock_response".to_string()));
568
569        // Test Merge
570        let merge_result = client
571            .merge("users:bob", json!({"city": "New York"}))
572            .await
573            .unwrap();
574        assert_eq!(merge_result, Value::String("mock_response".to_string()));
575
576        // Test Upsert
577        let upsert_result = client
578            .upsert("users:charlie", Some(json!({"name": "Charlie", "age": 28})))
579            .await
580            .unwrap();
581        assert_eq!(upsert_result, Value::String("mock_response".to_string()));
582    }
583
584    // Removed test_http_engine as HttpEngine is not available in minimal implementation
585
586    #[tokio::test]
587    async fn test_bakery_queries_with_parameters() {
588        let engine = Box::new(MockEngine);
589        let mut client = SurrealClient::new(
590            engine,
591            Some("bakery".to_string()),
592            Some("inventory".to_string()),
593        );
594
595        // Set variables
596        client.let_var("min_stock", json!(10)).await.unwrap();
597
598        client.let_var("category", json!("bread")).await.unwrap();
599
600        // Test parameterized query
601        let query =
602            "SELECT * FROM products WHERE stock_level < $min_stock AND category = $category";
603        let result = client.query(query, None).await.unwrap();
604        assert_eq!(result, Value::String("mock_response".to_string()));
605
606        // Test query with inline parameters
607        let variables = json!({
608            "supplier": "FreshBake Co",
609            "min_price": 5.0
610        });
611
612        let query_with_params = "SELECT * FROM products WHERE supplier = $supplier AND price >= $min_price ORDER BY price DESC";
613        let result = client
614            .query(query_with_params, Some(variables))
615            .await
616            .unwrap();
617        assert_eq!(result, Value::String("mock_response".to_string()));
618
619        // Test complex aggregation query
620        let analytics_query = r#"
621            SELECT
622                category,
623                COUNT() as total_products,
624                SUM(stock_level) as total_stock,
625                AVG(price) as avg_price,
626                MAX(price) as max_price,
627                MIN(price) as min_price
628            FROM products
629            WHERE stock_level > 0
630            GROUP BY category
631            ORDER BY total_stock DESC
632        "#;
633
634        let result = client.query(analytics_query, None).await.unwrap();
635        assert_eq!(result, Value::String("mock_response".to_string()));
636
637        // Test relation query
638        let relation_query = r#"
639            SELECT *,
640                ->supplied_by->suppliers.* as supplier_info
641            FROM products
642            WHERE category = 'pastries'
643        "#;
644
645        let result = client.query(relation_query, None).await.unwrap();
646        assert_eq!(result, Value::String("mock_response".to_string()));
647
648        // Test time-based query
649        let time_query = r#"
650            SELECT *
651            FROM orders
652            WHERE created_at >= time::now() - 7d
653            ORDER BY created_at DESC
654            LIMIT 50
655        "#;
656
657        let result = client.query(time_query, None).await.unwrap();
658        assert_eq!(result, Value::String("mock_response".to_string()));
659
660        // Clean up variables
661        client.unset("min_stock").await.unwrap();
662        client.unset("category").await.unwrap();
663    }
664
665    #[tokio::test]
666    async fn test_complex_analytics_queries() {
667        let engine = Box::new(MockEngine);
668        let client = SurrealClient::new(
669            engine,
670            Some("analytics".to_string()),
671            Some("business".to_string()),
672        );
673
674        // Test revenue analysis
675        let revenue_query = r#"
676            SELECT
677                date::format(created_at, '%Y-%m') as month,
678                SUM(total_amount) as monthly_revenue,
679                COUNT() as order_count,
680                AVG(total_amount) as avg_order_value
681            FROM orders
682            WHERE created_at >= time::now() - 12mo
683            GROUP BY month
684            ORDER BY month DESC
685        "#;
686
687        let result = client.query(revenue_query, None).await.unwrap();
688        assert_eq!(result, Value::String("mock_response".to_string()));
689
690        // Test customer segmentation
691        let segmentation_query = r#"
692            SELECT
693                CASE
694                    WHEN total_spent >= 1000 THEN 'Premium'
695                    WHEN total_spent >= 500 THEN 'Regular'
696                    ELSE 'Basic'
697                END as segment,
698                COUNT() as customer_count,
699                AVG(total_spent) as avg_spent,
700                SUM(total_spent) as segment_revenue
701            FROM (
702                SELECT
703                    customer_id,
704                    SUM(total_amount) as total_spent
705                FROM orders
706                GROUP BY customer_id
707            ) as customer_totals
708            GROUP BY segment
709            ORDER BY avg_spent DESC
710        "#;
711
712        let result = client.query(segmentation_query, None).await.unwrap();
713        assert_eq!(result, Value::String("mock_response".to_string()));
714
715        // Test product performance with inventory correlation
716        let performance_query = r#"
717            SELECT
718                p.id,
719                p.name,
720                p.category,
721                COUNT(oi.id) as times_ordered,
722                SUM(oi.quantity) as total_quantity_sold,
723                SUM(oi.price * oi.quantity) as total_revenue,
724                p.stock_level as current_stock,
725                CASE
726                    WHEN p.stock_level = 0 THEN 'Out of Stock'
727                    WHEN p.stock_level < 10 THEN 'Low Stock'
728                    WHEN p.stock_level < 50 THEN 'Medium Stock'
729                    ELSE 'High Stock'
730                END as stock_status
731            FROM products p
732            LEFT JOIN order_items oi ON oi.product_id = p.id
733            GROUP BY p.id, p.name, p.category, p.stock_level
734            ORDER BY total_revenue DESC
735        "#;
736
737        let result = client.query(performance_query, None).await.unwrap();
738        assert_eq!(result, Value::String("mock_response".to_string()));
739    }
740}