valence_backend_mongodb/
backend.rs

1//! MongoDB wire [`DatabaseBackend`] using the official Rust driver.
2
3use std::collections::HashSet;
4use std::sync::Arc;
5
6use futures::TryStreamExt;
7use mongodb::bson::{doc, Document};
8use mongodb::options::IndexOptions;
9use mongodb::{Client, Collection, IndexModel};
10use serde_json::{Map, Value};
11
12use valence_core::ttl::{SchemaTtlPolicy, EXPIRE_AT_FIELD};
13use valence_core::{
14    BackendCapabilities, CompiledQuery, Database, DatabaseBackend, DatabaseFromEngine, Error,
15    KnownEngines, RecordId, Result,
16};
17
18use crate::config::MongoConfig;
19
20/// Stable engine slug for router keys (`mongodb:logical_name`).
21pub const ENGINE_ID: &str = KnownEngines::MONGODB;
22
23/// Schema evaluator const for `database:` routing.
24pub const PRIMARY: DatabaseFromEngine = Database::from_engine("primary", ENGINE_ID);
25
26const EDGES_COLLECTION: &str = "valence_edges";
27
28/// MongoDB-backed [`DatabaseBackend`] storing JSON documents per collection.
29///
30/// # Examples
31///
32/// ```ignore
33/// use std::sync::Arc;
34/// use valence::{
35///     valence_schema, Database, DatabaseFromEngine, FieldType, MongoBackend, Valence,
36///     MONGODB_ENGINE_ID,
37/// };
38///
39/// const COUNTER_DB: DatabaseFromEngine =
40///     Database::from_engine("default", MONGODB_ENGINE_ID);
41///
42/// valence_schema! {
43///     Counter {
44///         table: "counter",
45///         version: "0.1.0",
46///         database: COUNTER_DB,
47///         fields: [
48///             id: { r#type: FieldType::String, primary_key: true, required: true },
49///             value: { r#type: FieldType::Integer, required: true },
50///         ],
51///     }
52/// }
53///
54/// // Reads VALENCE_MONGODB_URI and optional VALENCE_MONGODB_DATABASE.
55/// let backend = MongoBackend::from_env().await?;
56/// let valence = Valence::builder()
57///     .add_backend("default", Arc::new(backend))
58///     .build()?;
59/// assert_eq!(
60///     valence.backend_for_table("counter")?.engine_id(),
61///     MONGODB_ENGINE_ID
62/// );
63/// # Ok::<(), valence::Error>(())
64/// ```
65#[derive(Clone)]
66pub struct MongoBackend {
67    client: Client,
68    database: String,
69    unique_fields: Arc<tokio::sync::RwLock<HashSet<(String, String)>>>,
70}
71
72impl std::fmt::Debug for MongoBackend {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        f.debug_struct("MongoBackend")
75            .field("database", &self.database)
76            .finish_non_exhaustive()
77    }
78}
79
80impl MongoBackend {
81    /// Start a builder for explicit host wiring.
82    pub fn builder() -> crate::config::MongoBackendBuilder {
83        crate::config::MongoBackendBuilder::new()
84    }
85
86    /// Connect using env defaults via builder (shorthand).
87    ///
88    /// # Errors
89    ///
90    /// Returns [`Error::Internal`] when env config is incomplete, or [`Error::Database`] on connect failure.
91    pub async fn from_env() -> Result<Self> {
92        Self::builder().from_env_defaults().build().await
93    }
94
95    /// Connect to MongoDB at `uri` using database `database`.
96    ///
97    /// # Errors
98    ///
99    /// Returns [`Error::Database`] when the MongoDB connection fails.
100    pub async fn connect(uri: &str, database: &str) -> Result<Self> {
101        Self::builder().uri(uri).database(database).build().await
102    }
103
104    /// Connect using explicit config.
105    ///
106    /// # Errors
107    ///
108    /// Returns [`Error::Database`] when the MongoDB connection fails.
109    pub async fn connect_with_config(config: MongoConfig) -> Result<Self> {
110        let client = Client::with_uri_str(&config.uri)
111            .await
112            .map_err(|e| Error::database(e.to_string()))?;
113        let backend = Self {
114            client,
115            database: config.database,
116            unique_fields: Arc::new(tokio::sync::RwLock::new(HashSet::new())),
117        };
118        backend.ensure_edges_collection().await?;
119        Ok(backend)
120    }
121
122    fn assert_safe_table(table: &str) -> Result<()> {
123        if table.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') {
124            Ok(())
125        } else {
126            Err(Error::Validation(format!("unsafe table name: {table}")))
127        }
128    }
129
130    fn collection(&self, table: &str) -> Collection<Document> {
131        self.client.database(&self.database).collection(table)
132    }
133
134    async fn ensure_edges_collection(&self) -> Result<()> {
135        let coll = self.collection(EDGES_COLLECTION);
136        let index = IndexModel::builder()
137            .keys(doc! {
138                "from_table": 1,
139                "from_id": 1,
140                "edge_type": 1,
141                "to_table": 1,
142                "to_id": 1,
143            })
144            .options(IndexOptions::builder().unique(true).build())
145            .build();
146        coll.create_index(index)
147            .await
148            .map_err(|e| Error::database(e.to_string()))?;
149        Ok(())
150    }
151
152    async fn unique_fields_for(&self, table: &str) -> Vec<String> {
153        self.unique_fields
154            .read()
155            .await
156            .iter()
157            .filter(|(t, _)| t == table)
158            .map(|(_, f)| f.clone())
159            .collect()
160    }
161
162    async fn check_unique_fields(
163        &self,
164        table: &str,
165        record: &Value,
166        exclude_id: Option<&str>,
167    ) -> Result<()> {
168        for field in self.unique_fields_for(table).await {
169            let Some(value) = record.get(&field).and_then(|v| v.as_str()) else {
170                continue;
171            };
172            if let Some(exclude) = exclude_id {
173                if let Ok(Some(row)) = self.get_record(table, exclude).await {
174                    if row.get(&field).and_then(|v| v.as_str()) == Some(value) {
175                        continue;
176                    }
177                }
178            }
179            let coll = self.collection(table);
180            let filter = doc! { field.as_str(): value };
181            if let Some(existing) = coll
182                .find_one(filter)
183                .await
184                .map_err(|e| Error::database(e.to_string()))?
185            {
186                let existing_id = existing.get_str("_id").unwrap_or("");
187                if exclude_id != Some(existing_id) {
188                    return Err(Error::database(format!(
189                        "duplicate unique index value for {table}.{field}"
190                    )));
191                }
192            }
193        }
194        Ok(())
195    }
196
197    fn value_to_bson(value: &Value) -> mongodb::bson::Bson {
198        serde_json::from_value(value.clone()).unwrap_or(mongodb::bson::Bson::Null)
199    }
200
201    fn doc_to_row(table: &str, id: &str, doc: Document) -> Value {
202        let mut map = Map::new();
203        for (k, v) in doc {
204            if k == "_id" {
205                continue;
206            }
207            if k == EXPIRE_AT_FIELD {
208                if let mongodb::bson::Bson::DateTime(dt) = v {
209                    let millis = dt.timestamp_millis();
210                    let chrono_dt = chrono::DateTime::from_timestamp_millis(millis)
211                        .unwrap_or(chrono::DateTime::UNIX_EPOCH);
212                    map.insert(
213                        k,
214                        Value::String(
215                            chrono_dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
216                        ),
217                    );
218                    continue;
219                }
220            }
221            map.insert(k, bson_to_json(v));
222        }
223        row_from_body(table, id, Value::Object(map))
224    }
225
226    async fn rows_for_table(&self, table: &str, limit: Option<usize>) -> Result<Vec<Value>> {
227        Self::assert_safe_table(table)?;
228        let coll = self.collection(table);
229        let mut cursor = coll
230            .find(doc! {})
231            .await
232            .map_err(|e| Error::database(e.to_string()))?;
233        let mut rows = Vec::new();
234        while let Some(doc) = cursor
235            .try_next()
236            .await
237            .map_err(|e| Error::database(e.to_string()))?
238        {
239            let id = doc.get_str("_id").unwrap_or("").to_string();
240            rows.push(Self::doc_to_row(table, &id, doc));
241            if limit.is_some_and(|n| rows.len() >= n) {
242                break;
243            }
244        }
245        Ok(rows)
246    }
247}
248
249#[async_trait::async_trait]
250impl DatabaseBackend for MongoBackend {
251    fn engine_id(&self) -> &'static str {
252        ENGINE_ID
253    }
254
255    fn capabilities(&self) -> BackendCapabilities {
256        BackendCapabilities {
257            supports_merge: true,
258            supports_graph_edges: true,
259            telemetry_label: "mongodb",
260        }
261    }
262
263    async fn execute_compiled_query(&self, compiled: &CompiledQuery) -> Result<Vec<Value>> {
264        let q = compiled.query_string.trim();
265        if let Ok(descriptor) = serde_json::from_str::<Value>(q) {
266            if let Some(collection) = descriptor.get("collection").and_then(|v| v.as_str()) {
267                let limit = descriptor
268                    .get("limit")
269                    .and_then(|v| v.as_u64())
270                    .map(|n| usize::try_from(n).unwrap_or(usize::MAX));
271                let mut rows = self.rows_for_table(collection, limit).await?;
272                rows = valence_core::query::apply_equality_where(rows, compiled);
273                rows = valence_core::query::apply_order_limit_offset(rows, &compiled.query_string);
274                return Ok(rows);
275            }
276        }
277
278        let upper = q.to_uppercase();
279        if !upper.starts_with("SELECT ") {
280            return Ok(vec![]);
281        }
282        let from_idx = upper
283            .find(" FROM ")
284            .ok_or_else(|| Error::Internal("missing FROM in select".into()))?;
285        let table = q[from_idx + 6..]
286            .split_whitespace()
287            .next()
288            .unwrap_or("")
289            .trim();
290        if table.is_empty() {
291            return Ok(vec![]);
292        }
293        // Load all candidates; WHERE / ORDER / LIMIT applied in-process (parity with mem).
294        let mut rows = self.rows_for_table(table, None).await?;
295        rows = valence_core::query::apply_equality_where(rows, compiled);
296        rows = valence_core::query::apply_order_limit_offset(rows, &compiled.query_string);
297        if upper.contains("SELECT ID") && !upper.contains("BODY") {
298            // Match mem: IdOnlyRecord deserializes `{ "id": ... }`, not bare strings.
299            return Ok(rows
300                .iter()
301                .filter_map(|r| {
302                    r.get("id")
303                        .and_then(|id| id.get("id").and_then(|x| x.as_str()))
304                        .or_else(|| r.get("id").and_then(|id| id.as_str()))
305                        .map(|id| serde_json::json!({ "id": id }))
306                })
307                .collect());
308        }
309        Ok(rows)
310    }
311
312    async fn ensure_schemaless_table(&self, _table: &str) -> Result<()> {
313        Ok(())
314    }
315
316    async fn get_record(&self, table: &str, id: &str) -> Result<Option<Value>> {
317        Self::assert_safe_table(table)?;
318        let coll = self.collection(table);
319        let doc = coll
320            .find_one(doc! { "_id": id })
321            .await
322            .map_err(|e| Error::database(e.to_string()))?;
323        Ok(doc.map(|d| Self::doc_to_row(table, id, d)))
324    }
325
326    async fn create_record(&self, table: &str, content: Value) -> Result<Value> {
327        Self::assert_safe_table(table)?;
328        if let Ok(layout) = valence_core::storage_layout::StorageLayout::from_registry_table(table)
329        {
330            valence_core::storage_layout::validate_write_types(&layout, &content)?;
331        }
332        let mut content = content;
333        valence_core::ttl::prepare_create_content(table, self, &mut content)?;
334        self.check_unique_fields(table, &content, None).await?;
335        let id = storage_id(&content).unwrap_or_else(uuid_simple);
336        let mut record = content;
337        if let Some(obj) = record.as_object_mut() {
338            let has_string_id = obj.get("id").and_then(|v| v.as_str()).is_some();
339            if !has_string_id {
340                obj.insert("id".into(), record_id_json(table, &id));
341            }
342        }
343        let mut doc = body_document(&record);
344        doc.insert("_id", id.clone());
345        let coll = self.collection(table);
346        coll.insert_one(doc.clone())
347            .await
348            .map_err(map_duplicate_key)?;
349        // Return the persisted row so TTL DateTime round-trips match get/merge.
350        Ok(Self::doc_to_row(table, &id, doc))
351    }
352
353    async fn update_record(&self, table: &str, id: &str, content: Value) -> Result<Value> {
354        if self.get_record(table, id).await?.is_none() {
355            return Err(Error::NotFound(format!("{table}:{id}")));
356        }
357        self.check_unique_fields(table, &content, Some(id)).await?;
358        let mut record = content;
359        if let Some(obj) = record.as_object_mut() {
360            obj.insert("id".into(), record_id_json(table, id));
361        }
362        let doc = body_document(&record);
363        let coll = self.collection(table);
364        coll.replace_one(doc! { "_id": id }, doc)
365            .await
366            .map_err(map_duplicate_key)?;
367        Ok(record)
368    }
369
370    async fn merge_record(&self, table: &str, id: &str, patch: Value) -> Result<Value> {
371        let existing = self
372            .get_record(table, id)
373            .await?
374            .unwrap_or_else(|| row_from_body(table, id, Value::Object(Map::new())));
375        let mut merged = existing;
376        if let (Some(base), Some(patch_obj)) = (merged.as_object_mut(), patch.as_object()) {
377            for (k, v) in patch_obj {
378                base.insert(k.clone(), v.clone());
379            }
380        }
381        self.check_unique_fields(table, &merged, Some(id)).await?;
382        let doc = body_document(&merged);
383        let coll = self.collection(table);
384        coll.replace_one(doc! { "_id": id }, doc)
385            .await
386            .map_err(|e| {
387                if e.to_string().contains("duplicate key") {
388                    Error::database(format!("duplicate unique index value for {table}"))
389                } else {
390                    Error::database(e.to_string())
391                }
392            })?;
393        Ok(merged)
394    }
395
396    async fn upsert_record(&self, table: &str, id: &str, content: Value) -> Result<Value> {
397        if self.get_record(table, id).await?.is_some() {
398            self.update_record(table, id, content).await
399        } else {
400            let mut record = content;
401            if let Some(obj) = record.as_object_mut() {
402                obj.insert("id".into(), record_id_json(table, id));
403            }
404            self.create_record(table, record).await
405        }
406    }
407
408    async fn delete_record(&self, table: &str, id: &str) -> Result<()> {
409        let coll = self.collection(table);
410        coll.delete_one(doc! { "_id": id })
411            .await
412            .map_err(|e| Error::database(e.to_string()))?;
413        Ok(())
414    }
415
416    async fn relate_edge(&self, from: &RecordId, edge_table: &str, to: &RecordId) -> Result<()> {
417        let coll = self.collection(EDGES_COLLECTION);
418        coll.insert_one(doc! {
419            "from_table": from.table(),
420            "from_id": from.id(),
421            "edge_type": edge_table,
422            "to_table": to.table(),
423            "to_id": to.id(),
424        })
425        .await
426        .map_err(|e| Error::database(e.to_string()))?;
427        Ok(())
428    }
429
430    async fn unrelate_edge(&self, from: &RecordId, edge_table: &str, to: &RecordId) -> Result<()> {
431        let coll = self.collection(EDGES_COLLECTION);
432        coll.delete_one(doc! {
433            "from_table": from.table(),
434            "from_id": from.id(),
435            "edge_type": edge_table,
436            "to_table": to.table(),
437            "to_id": to.id(),
438        })
439        .await
440        .map_err(|e| Error::database(e.to_string()))?;
441        Ok(())
442    }
443
444    async fn get_edge_targets(&self, from: &RecordId, edge_table: &str) -> Result<Vec<RecordId>> {
445        let coll = self.collection(EDGES_COLLECTION);
446        let mut cursor = coll
447            .find(doc! {
448                "from_table": from.table(),
449                "from_id": from.id(),
450                "edge_type": edge_table,
451            })
452            .await
453            .map_err(|e| Error::database(e.to_string()))?;
454        let mut out = Vec::new();
455        while let Some(edge) = cursor
456            .try_next()
457            .await
458            .map_err(|e| Error::database(e.to_string()))?
459        {
460            let to_table = edge.get_str("to_table").unwrap_or("").to_string();
461            let to_id = edge.get_str("to_id").unwrap_or("").to_string();
462            out.push(RecordId::new(to_table, to_id));
463        }
464        Ok(out)
465    }
466
467    async fn define_unique_index(&self, table: &str, field: &str) -> Result<()> {
468        Self::assert_safe_table(table)?;
469        self.unique_fields
470            .write()
471            .await
472            .insert((table.to_string(), field.to_string()));
473        let coll = self.collection(table);
474        let index = IndexModel::builder()
475            .keys(doc! { field: 1 })
476            .options(IndexOptions::builder().unique(true).sparse(true).build())
477            .build();
478        coll.create_index(index)
479            .await
480            .map_err(|e| Error::database(e.to_string()))?;
481        Ok(())
482    }
483
484    fn ttl_capability(&self) -> valence_core::ttl::BackendTtlCapability {
485        crate::ttl::ttl_capability()
486    }
487
488    async fn apply_ttl_policy(&self, table: &str, _policy: &SchemaTtlPolicy) -> Result<()> {
489        Self::assert_safe_table(table)?;
490        let coll = self.collection(table);
491        crate::ttl::apply_ttl_policy(&coll).await
492    }
493}
494
495#[allow(clippy::needless_pass_by_value)] // map_err adapter; value only Display'd
496fn map_duplicate_key(e: mongodb::error::Error) -> Error {
497    if e.to_string().contains("duplicate key") {
498        Error::database("duplicate unique index value")
499    } else {
500        Error::database(e.to_string())
501    }
502}
503
504fn body_document(record: &Value) -> Document {
505    let mut doc = Document::new();
506    if let Some(obj) = record.as_object() {
507        for (k, v) in obj {
508            if k == "id" {
509                continue;
510            }
511            if k == EXPIRE_AT_FIELD {
512                if let Some(s) = v.as_str() {
513                    if let Ok(dt) = chrono::DateTime::parse_from_rfc3339(s) {
514                        doc.insert(
515                            k.clone(),
516                            mongodb::bson::DateTime::from_millis(dt.timestamp_millis()),
517                        );
518                        continue;
519                    }
520                }
521            }
522            doc.insert(k.clone(), MongoBackend::value_to_bson(v));
523        }
524    }
525    doc
526}
527
528fn bson_to_json(bson: mongodb::bson::Bson) -> Value {
529    serde_json::to_value(bson).unwrap_or(Value::Null)
530}
531
532fn row_from_body(table: &str, id: &str, body: Value) -> Value {
533    let mut obj = match body {
534        Value::Object(map) => map,
535        _ => Map::new(),
536    };
537    obj.insert("id".into(), record_id_json(table, id));
538    Value::Object(obj)
539}
540
541fn record_id_json(table: &str, id: &str) -> Value {
542    serde_json::json!({
543        "table": table,
544        "id": id,
545    })
546}
547
548fn storage_id(content: &Value) -> Option<String> {
549    content.get("id").and_then(|v| {
550        v.get("id")
551            .and_then(|x| x.as_str())
552            .map(str::to_string)
553            .or_else(|| v.as_str().map(str::to_string))
554    })
555}
556
557fn uuid_simple() -> String {
558    uuid::Uuid::new_v4().to_string()
559}
valence_backend_mongodb/backend.rs

valence_backend_mongodb/
backend.rs
