Skip to main content

valence_backend_indradb/
backend.rs

1//! IndraDB embedded graph backend using in-memory datastore only.
2
3use std::collections::{HashMap, HashSet};
4
5use indradb::{
6    Edge, Identifier, Json, MemoryDatastore, PipePropertyQuery, QueryExt, QueryOutputValue,
7    RangeVertexQuery, SpecificVertexQuery, Vertex,
8};
9use tokio::sync::RwLock;
10use uuid::Uuid;
11use valence_core::{
12    BackendCapabilities, CompiledQuery, Database, DatabaseBackend, DatabaseFromEngine, Error,
13    KnownEngines, RecordId, Result,
14};
15
16/// Stable engine slug for router keys (`indradb:logical_name`).
17pub const ENGINE_ID: &str = KnownEngines::INDRADB;
18
19/// Schema evaluator const for `database:` routing.
20pub const PRIMARY: DatabaseFromEngine = Database::from_engine("primary", ENGINE_ID);
21
22type IndraDb = indradb::Database<MemoryDatastore>;
23
24/// Embedded IndraDB [`DatabaseBackend`] mapping Valence tables to vertex types.
25///
26/// # Examples
27///
28/// ```ignore
29/// use std::sync::Arc;
30/// use valence::{
31///     valence_schema, Database, DatabaseFromEngine, FieldType, IndradbBackend, Valence,
32///     INDRADB_ENGINE_ID,
33/// };
34///
35/// const COUNTER_DB: DatabaseFromEngine =
36///     Database::from_engine("default", INDRADB_ENGINE_ID);
37///
38/// valence_schema! {
39///     Counter {
40///         table: "counter",
41///         version: "0.1.0",
42///         database: COUNTER_DB,
43///         fields: [
44///             id: { r#type: FieldType::String, primary_key: true, required: true },
45///             value: { r#type: FieldType::Integer, required: true },
46///         ],
47///     }
48/// }
49///
50/// let valence = Valence::builder()
51///     .add_backend("default", Arc::new(IndradbBackend::new()))
52///     .build()?;
53/// assert_eq!(
54///     valence.backend_for_table("counter")?.engine_id(),
55///     INDRADB_ENGINE_ID
56/// );
57/// # Ok::<(), valence::Error>(())
58/// ```
59pub struct IndradbBackend {
60    db: IndraDb,
61    unique_indexes: RwLock<HashMap<(String, String), HashSet<String>>>,
62}
63
64impl std::fmt::Debug for IndradbBackend {
65    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
66        f.debug_struct("IndradbBackend")
67            .field("unique_indexes", &self.unique_indexes)
68            .finish_non_exhaustive()
69    }
70}
71
72impl Default for IndradbBackend {
73    fn default() -> Self {
74        Self::new()
75    }
76}
77
78impl IndradbBackend {
79    /// Create an empty in-memory IndraDB backend.
80    pub fn new() -> Self {
81        Self {
82            db: MemoryDatastore::new_db(),
83            unique_indexes: RwLock::new(HashMap::new()),
84        }
85    }
86
87    #[allow(clippy::needless_pass_by_value)] // map_err adapter; value only Display'd
88    fn id_err(e: indradb::ValidationError) -> Error {
89        Error::Validation(format!("invalid indradb identifier: {e:?}"))
90    }
91
92    #[allow(clippy::needless_pass_by_value)] // map_err adapter; value only Display'd
93    fn db_err(e: indradb::Error) -> Error {
94        Error::database(e.to_string())
95    }
96
97    fn table_identifier(table: &str) -> Result<Identifier> {
98        Identifier::new(table).map_err(Self::id_err)
99    }
100
101    fn edge_identifier(edge_table: &str) -> Result<Identifier> {
102        Identifier::new(edge_table).map_err(Self::id_err)
103    }
104
105    fn field_property(name: &str) -> Result<Identifier> {
106        Identifier::new(name).map_err(Self::id_err)
107    }
108
109    fn vertex_uuid(table: &str, id: &str) -> Uuid {
110        Uuid::new_v5(&Uuid::NAMESPACE_URL, format!("{table}:{id}").as_bytes())
111    }
112
113    fn ensure_vertex(&self, table: &str, id: &str) -> Result<Vertex> {
114        let vertex_type = Self::table_identifier(table)?;
115        let vertex = Vertex::with_id(Self::vertex_uuid(table, id), vertex_type);
116        let _ = self.db.create_vertex(&vertex).map_err(Self::db_err)?;
117        Ok(vertex)
118    }
119
120    /// Read one property per field into a Valence JSON object (includes nested `id`).
121    fn read_fields(
122        &self,
123        table: &str,
124        id: &str,
125        vertex_id: Uuid,
126    ) -> Result<Option<serde_json::Value>> {
127        let query = PipePropertyQuery::new(Box::new(SpecificVertexQuery::single(vertex_id).into()))
128            .map_err(Self::id_err)?;
129        let output = self.db.get(query).map_err(Self::db_err)?;
130        let mut map = serde_json::Map::new();
131        let mut found = false;
132        for item in output {
133            if let QueryOutputValue::VertexProperties(vps) = item {
134                for vp in vps {
135                    for prop in vp.props {
136                        found = true;
137                        map.insert(
138                            prop.name.as_str().to_string(),
139                            prop.value.0.as_ref().clone(),
140                        );
141                    }
142                }
143            }
144        }
145        if !found {
146            return Ok(None);
147        }
148        // Prefer the caller-supplied id; fall back to the stored bare id so table
149        // scans (empty probe id) can rebuild the wire RecordId shape.
150        let stored_bare = map.remove("__valence_id").and_then(|v| match v {
151            serde_json::Value::String(s) => Some(s),
152            _ => None,
153        });
154        let bare = if id.is_empty() {
155            stored_bare.unwrap_or_default()
156        } else {
157            id.to_string()
158        };
159        if !bare.is_empty() {
160            map.insert("id".into(), serde_json::json!({"table": table, "id": bare}));
161        }
162        Ok(Some(serde_json::Value::Object(map)))
163    }
164
165    /// Write each top-level field (except nested `id`) as its own Indra property.
166    fn write_fields(&self, table: &str, id: &str, record: serde_json::Value) -> Result<()> {
167        let vertex = self.ensure_vertex(table, id)?;
168        // Bare id for reverse lookup when scanning vertices.
169        let id_prop = Self::field_property("__valence_id")?;
170        self.db
171            .set_properties(
172                SpecificVertexQuery::single(vertex.id),
173                id_prop,
174                &Json::new(serde_json::Value::String(id.to_string())),
175            )
176            .map_err(Self::db_err)?;
177        let obj = match record {
178            serde_json::Value::Object(m) => m,
179            other => {
180                let mut m = serde_json::Map::new();
181                m.insert("_value".into(), other);
182                m
183            }
184        };
185        for (key, value) in obj {
186            if key == "id" {
187                continue;
188            }
189            valence_core::safe_ident::assert_safe_ident(&key)?;
190            let prop = Self::field_property(&key)?;
191            self.db
192                .set_properties(
193                    SpecificVertexQuery::single(vertex.id),
194                    prop,
195                    &Json::new(value),
196                )
197                .map_err(Self::db_err)?;
198        }
199        Ok(())
200    }
201
202    fn record_id_from_props(map: &serde_json::Map<String, serde_json::Value>) -> Option<String> {
203        map.get("__valence_id")
204            .and_then(|v| v.as_str())
205            .map(str::to_string)
206            .or_else(|| {
207                map.get("id").and_then(|v| match v {
208                    serde_json::Value::String(s) => Some(s.clone()),
209                    serde_json::Value::Object(o) => {
210                        o.get("id").and_then(|x| x.as_str()).map(str::to_string)
211                    }
212                    _ => None,
213                })
214            })
215    }
216
217    async fn check_unique_fields(
218        &self,
219        table: &str,
220        record: &serde_json::Value,
221        exclude_id: Option<&str>,
222    ) -> Result<()> {
223        let indexes = self.unique_indexes.read().await.clone();
224        for ((idx_table, field), values) in &indexes {
225            if idx_table != table {
226                continue;
227            }
228            let Some(value) = record.get(field).and_then(|v| v.as_str()) else {
229                continue;
230            };
231            if exclude_id.is_some_and(|id| {
232                self.read_fields(table, id, Self::vertex_uuid(table, id))
233                    .ok()
234                    .flatten()
235                    .and_then(|row| row.get(field).and_then(|v| v.as_str()).map(str::to_string))
236                    .is_some_and(|existing| existing == value)
237            }) {
238                continue;
239            }
240            if values.contains(value) {
241                return Err(Error::database(format!(
242                    "duplicate unique index value for {table}.{field}"
243                )));
244            }
245        }
246        Ok(())
247    }
248
249    async fn track_unique_fields(&self, table: &str, record: &serde_json::Value) {
250        let mut indexes = self.unique_indexes.write().await;
251        for ((idx_table, field), values) in indexes.iter_mut() {
252            if idx_table != table {
253                continue;
254            }
255            if let Some(value) = record.get(field).and_then(|v| v.as_str()) {
256                values.insert(value.to_string());
257            }
258        }
259        drop(indexes);
260    }
261
262    async fn untrack_unique_fields(&self, table: &str, record: &serde_json::Value) {
263        let mut indexes = self.unique_indexes.write().await;
264        for ((idx_table, field), values) in indexes.iter_mut() {
265            if idx_table != table {
266                continue;
267            }
268            if let Some(value) = record.get(field).and_then(|v| v.as_str()) {
269                values.remove(value);
270            }
271        }
272        drop(indexes);
273    }
274
275    fn rows_for_table(&self, table: &str) -> Result<Vec<serde_json::Value>> {
276        let vertex_type = Self::table_identifier(table)?;
277        let output = self
278            .db
279            .get(RangeVertexQuery::new().t(vertex_type))
280            .map_err(Self::db_err)?;
281        let mut rows = Vec::new();
282        for item in output {
283            if let QueryOutputValue::Vertices(vertices) = item {
284                for vertex in vertices {
285                    // Empty probe id: recover bare id from `__valence_id` inside read_fields.
286                    if let Some(row) = self.read_fields(table, "", vertex.id)? {
287                        if row
288                            .as_object()
289                            .and_then(Self::record_id_from_props)
290                            .is_some_and(|bare| !bare.is_empty())
291                        {
292                            rows.push(row);
293                        }
294                    }
295                }
296            }
297        }
298        Ok(rows)
299    }
300
301    fn execute_indra_descriptor(
302        &self,
303        descriptor: &serde_json::Value,
304    ) -> Result<Vec<serde_json::Value>> {
305        let table = descriptor
306            .get("vertex_type")
307            .and_then(|v| v.as_str())
308            .ok_or_else(|| Error::Internal("missing vertex_type in indradb query".into()))?;
309        let mut rows = self.rows_for_table(table)?;
310        if let Some(limit) = descriptor.get("limit").and_then(|v| v.as_u64()) {
311            rows.truncate(usize::try_from(limit).unwrap_or(usize::MAX));
312        }
313        Ok(rows)
314    }
315
316    fn execute_sql_select(&self, q: &str) -> Result<Vec<serde_json::Value>> {
317        let upper = q.to_uppercase();
318        if !upper.starts_with("SELECT ") {
319            return Ok(vec![]);
320        }
321        let from_idx = upper
322            .find(" FROM ")
323            .ok_or_else(|| Error::Internal("missing FROM in select".into()))?;
324        let table = q[from_idx + 6..]
325            .split_whitespace()
326            .next()
327            .unwrap_or("")
328            .trim();
329        if table.is_empty() {
330            return Ok(vec![]);
331        }
332        let rows = self.rows_for_table(table)?;
333        if upper.contains("SELECT id") && !upper.contains("body") {
334            return Ok(rows
335                .iter()
336                .filter_map(|r| {
337                    r.get("id")
338                        .and_then(|id| id.get("id").and_then(|x| x.as_str()))
339                        .or_else(|| r.get("id").and_then(|id| id.as_str()))
340                        .map(|id| serde_json::Value::String(id.to_string()))
341                })
342                .collect());
343        }
344        let mut out = rows;
345        if let Some(limit_idx) = upper.rfind(" LIMIT ") {
346            if let Ok(limit) = q[limit_idx + 7..].trim().parse::<usize>() {
347                out.truncate(limit);
348            }
349        }
350        Ok(out)
351    }
352}
353
354#[async_trait::async_trait]
355impl DatabaseBackend for IndradbBackend {
356    fn engine_id(&self) -> &'static str {
357        ENGINE_ID
358    }
359
360    fn capabilities(&self) -> BackendCapabilities {
361        BackendCapabilities {
362            supports_merge: true,
363            supports_graph_edges: true,
364            telemetry_label: "indradb",
365        }
366    }
367
368    async fn execute_compiled_query(
369        &self,
370        compiled: &CompiledQuery,
371    ) -> Result<Vec<serde_json::Value>> {
372        let q = compiled.query_string.trim();
373        if let Ok(descriptor) = serde_json::from_str::<serde_json::Value>(q) {
374            if descriptor.get("vertex_type").is_some() {
375                return self.execute_indra_descriptor(&descriptor);
376            }
377        }
378        let mut rows = self.execute_sql_select(q)?;
379        rows = valence_core::query::apply_equality_where(rows, compiled);
380        rows = valence_core::query::apply_order_limit_offset(rows, q);
381        Ok(rows)
382    }
383
384    async fn get_record(&self, table: &str, id: &str) -> Result<Option<serde_json::Value>> {
385        self.read_fields(table, id, Self::vertex_uuid(table, id))
386    }
387
388    async fn create_record(
389        &self,
390        table: &str,
391        content: serde_json::Value,
392    ) -> Result<serde_json::Value> {
393        if let Ok(layout) = valence_core::storage_layout::StorageLayout::from_registry_table(table)
394        {
395            valence_core::storage_layout::validate_write_types(&layout, &content)?;
396        }
397        self.check_unique_fields(table, &content, None).await?;
398        let id = storage_id_from_content(&content).unwrap_or_else(uuid_simple);
399        let mut record = content;
400        if let Some(obj) = record.as_object_mut() {
401            let has_string_id = obj.get("id").and_then(|v| v.as_str()).is_some();
402            if !has_string_id {
403                obj.insert("id".into(), record_id_json(table, &id));
404            }
405        }
406        self.write_fields(table, &id, record.clone())?;
407        self.track_unique_fields(table, &record).await;
408        Ok(record)
409    }
410
411    async fn update_record(
412        &self,
413        table: &str,
414        id: &str,
415        content: serde_json::Value,
416    ) -> Result<serde_json::Value> {
417        if self.get_record(table, id).await?.is_none() {
418            return Err(Error::NotFound(format!("{table}:{id}")));
419        }
420        if let Some(existing) = self.get_record(table, id).await? {
421            self.untrack_unique_fields(table, &existing).await;
422        }
423        self.check_unique_fields(table, &content, Some(id)).await?;
424        self.write_fields(table, id, content.clone())?;
425        self.track_unique_fields(table, &content).await;
426        Ok(content)
427    }
428
429    async fn merge_record(
430        &self,
431        table: &str,
432        id: &str,
433        patch: serde_json::Value,
434    ) -> Result<serde_json::Value> {
435        let mut record = self
436            .get_record(table, id)
437            .await?
438            .unwrap_or_else(|| serde_json::json!({}));
439        if let Some(existing) = self.get_record(table, id).await? {
440            self.untrack_unique_fields(table, &existing).await;
441        }
442        if let (Some(base), Some(patch_obj)) = (record.as_object_mut(), patch.as_object()) {
443            for (k, v) in patch_obj {
444                base.insert(k.clone(), v.clone());
445            }
446        }
447        self.check_unique_fields(table, &record, Some(id)).await?;
448        self.write_fields(table, id, record.clone())?;
449        self.track_unique_fields(table, &record).await;
450        Ok(record)
451    }
452
453    async fn upsert_record(
454        &self,
455        table: &str,
456        id: &str,
457        content: serde_json::Value,
458    ) -> Result<serde_json::Value> {
459        if let Some(existing) = self.get_record(table, id).await? {
460            self.untrack_unique_fields(table, &existing).await;
461        }
462        self.check_unique_fields(table, &content, Some(id)).await?;
463        let mut record = content;
464        if let Some(obj) = record.as_object_mut() {
465            obj.insert("id".into(), record_id_json(table, id));
466        }
467        self.write_fields(table, id, record.clone())?;
468        self.track_unique_fields(table, &record).await;
469        Ok(record)
470    }
471
472    async fn delete_record(&self, table: &str, id: &str) -> Result<()> {
473        if let Some(existing) = self.get_record(table, id).await? {
474            self.untrack_unique_fields(table, &existing).await;
475        }
476        let vertex_id = Self::vertex_uuid(table, id);
477        self.db
478            .delete(SpecificVertexQuery::single(vertex_id))
479            .map_err(Self::db_err)
480    }
481
482    async fn relate_edge(&self, from: &RecordId, edge_table: &str, to: &RecordId) -> Result<()> {
483        self.ensure_vertex(from.table(), from.id())?;
484        self.ensure_vertex(to.table(), to.id())?;
485        let edge_type = Self::edge_identifier(edge_table)?;
486        let edge = Edge::new(
487            Self::vertex_uuid(from.table(), from.id()),
488            edge_type,
489            Self::vertex_uuid(to.table(), to.id()),
490        );
491        let _ = self.db.create_edge(&edge).map_err(Self::db_err)?;
492        Ok(())
493    }
494
495    async fn unrelate_edge(&self, from: &RecordId, edge_table: &str, to: &RecordId) -> Result<()> {
496        let edge_type = Self::edge_identifier(edge_table)?;
497        let edge = Edge::new(
498            Self::vertex_uuid(from.table(), from.id()),
499            edge_type,
500            Self::vertex_uuid(to.table(), to.id()),
501        );
502        self.db
503            .delete(indradb::SpecificEdgeQuery::single(edge))
504            .map_err(Self::db_err)
505    }
506
507    async fn get_edge_targets(&self, from: &RecordId, edge_table: &str) -> Result<Vec<RecordId>> {
508        let from_uuid = Self::vertex_uuid(from.table(), from.id());
509        let edge_type = Self::edge_identifier(edge_table)?;
510        let output = self
511            .db
512            .get(
513                SpecificVertexQuery::single(from_uuid)
514                    .outbound()
515                    .map_err(Self::id_err)?,
516            )
517            .map_err(Self::db_err)?;
518        let mut targets = Vec::new();
519        for item in output {
520            if let QueryOutputValue::Edges(edges) = item {
521                for edge in edges {
522                    if edge.t != edge_type {
523                        continue;
524                    }
525                    let inbound_table = self
526                        .db
527                        .get(SpecificVertexQuery::single(edge.inbound_id))
528                        .map_err(Self::db_err)?
529                        .into_iter()
530                        .find_map(|value| {
531                            if let QueryOutputValue::Vertices(vertices) = value {
532                                vertices.first().map(|v| v.t.as_str().to_string())
533                            } else {
534                                None
535                            }
536                        })
537                        .unwrap_or_else(|| from.table().to_string());
538                    let body = self
539                        .read_fields(&inbound_table, "", edge.inbound_id)?
540                        .unwrap_or_else(|| serde_json::json!({}));
541                    let target_id = body
542                        .as_object()
543                        .and_then(Self::record_id_from_props)
544                        .or_else(|| storage_id_from_content(&body))
545                        .unwrap_or_else(|| edge.inbound_id.to_string());
546                    targets.push(RecordId::new(inbound_table, target_id));
547                }
548            }
549        }
550        Ok(targets)
551    }
552
553    async fn define_unique_index(&self, table: &str, field: &str) -> Result<()> {
554        let needs_populate = {
555            let indexes = self.unique_indexes.read().await;
556            indexes
557                .get(&(table.to_string(), field.to_string()))
558                .is_none_or(|entry| entry.is_empty())
559        };
560        let seeded = if needs_populate {
561            let rows = self.rows_for_table(table)?;
562            let mut values = HashSet::new();
563            for row in rows {
564                if let Some(value) = row.get(field).and_then(|v| v.as_str()) {
565                    values.insert(value.to_string());
566                }
567            }
568            Some(values)
569        } else {
570            None
571        };
572        let mut indexes = self.unique_indexes.write().await;
573        let entry = indexes
574            .entry((table.to_string(), field.to_string()))
575            .or_default();
576        if entry.is_empty() {
577            if let Some(values) = seeded {
578                *entry = values;
579            }
580        }
581        drop(indexes);
582        Ok(())
583    }
584}
585
586fn record_id_json(table: &str, id: &str) -> serde_json::Value {
587    serde_json::json!({
588        "table": table,
589        "id": id,
590    })
591}
592
593fn storage_id_from_content(content: &serde_json::Value) -> Option<String> {
594    let id_val = content.get("id")?;
595    if let Some(id) = id_val.get("id").and_then(|v| v.as_str()) {
596        return Some(id.to_string());
597    }
598    id_val.as_str().map(|s| s.to_string())
599}
600
601fn uuid_simple() -> String {
602    use std::time::{SystemTime, UNIX_EPOCH};
603    let nanos = SystemTime::now()
604        .duration_since(UNIX_EPOCH)
605        .map_or(0, |d| d.as_nanos());
606    format!("indradb-{nanos}")
607}