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};
4use std::sync::RwLock;
5
6use indradb::{
7    Edge, Identifier, Json, MemoryDatastore, PipePropertyQuery, QueryExt, QueryOutputValue,
8    RangeVertexQuery, SpecificVertexQuery, Vertex,
9};
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
22const BODY_PROPERTY: &str = "body";
23
24type IndraDb = indradb::Database<MemoryDatastore>;
25
26/// Embedded IndraDB [`DatabaseBackend`] mapping Valence tables to vertex types.
27///
28/// # Examples
29///
30/// ```ignore
31/// use std::sync::Arc;
32/// use valence::{
33///     valence_schema, Database, DatabaseFromEngine, FieldType, IndradbBackend, Valence,
34///     INDRADB_ENGINE_ID,
35/// };
36///
37/// const COUNTER_DB: DatabaseFromEngine =
38///     Database::from_engine("default", INDRADB_ENGINE_ID);
39///
40/// valence_schema! {
41///     Counter {
42///         table: "counter",
43///         version: "0.1.0",
44///         database: COUNTER_DB,
45///         fields: [
46///             id: { r#type: FieldType::String, primary_key: true, required: true },
47///             value: { r#type: FieldType::Integer, required: true },
48///         ],
49///     }
50/// }
51///
52/// let valence = Valence::builder()
53///     .add_backend("default", Arc::new(IndradbBackend::new()))
54///     .build()?;
55/// assert_eq!(
56///     valence.backend_for_table("counter")?.engine_id(),
57///     INDRADB_ENGINE_ID
58/// );
59/// # Ok::<(), valence::Error>(())
60/// ```
61pub struct IndradbBackend {
62    db: IndraDb,
63    unique_indexes: RwLock<HashMap<(String, String), HashSet<String>>>,
64}
65
66impl std::fmt::Debug for IndradbBackend {
67    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
68        f.debug_struct("IndradbBackend")
69            .field("unique_indexes", &self.unique_indexes)
70            .finish_non_exhaustive()
71    }
72}
73
74impl Default for IndradbBackend {
75    fn default() -> Self {
76        Self::new()
77    }
78}
79
80impl IndradbBackend {
81    /// Create an empty in-memory IndraDB backend.
82    pub fn new() -> Self {
83        Self {
84            db: MemoryDatastore::new_db(),
85            unique_indexes: RwLock::new(HashMap::new()),
86        }
87    }
88
89    fn id_err(e: indradb::ValidationError) -> Error {
90        Error::Validation(format!("invalid indradb identifier: {e:?}"))
91    }
92
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 body_property() -> Result<Identifier> {
106        Identifier::new(BODY_PROPERTY).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    fn read_body(&self, vertex_id: Uuid) -> Result<Option<serde_json::Value>> {
121        let query = PipePropertyQuery::new(Box::new(SpecificVertexQuery::single(vertex_id).into()))
122            .map_err(Self::id_err)?;
123        let output = self.db.get(query).map_err(Self::db_err)?;
124        for item in output {
125            if let QueryOutputValue::VertexProperties(vps) = item {
126                for vp in vps {
127                    for prop in vp.props {
128                        if prop.name.as_str() == BODY_PROPERTY {
129                            return Ok(Some(prop.value.0.as_ref().clone()));
130                        }
131                    }
132                }
133            }
134        }
135        Ok(None)
136    }
137
138    fn write_body(&self, table: &str, id: &str, body: serde_json::Value) -> Result<()> {
139        let vertex = self.ensure_vertex(table, id)?;
140        let prop = Self::body_property()?;
141        self.db
142            .set_properties(
143                SpecificVertexQuery::single(vertex.id),
144                prop,
145                &Json::new(body),
146            )
147            .map_err(Self::db_err)
148    }
149
150    fn check_unique_fields(
151        &self,
152        table: &str,
153        record: &serde_json::Value,
154        exclude_id: Option<&str>,
155    ) -> Result<()> {
156        let indexes = self
157            .unique_indexes
158            .read()
159            .map_err(|_| Error::Internal("indradb unique index lock poisoned".into()))?;
160        for ((idx_table, field), values) in indexes.iter() {
161            if idx_table != table {
162                continue;
163            }
164            let Some(value) = record.get(field).and_then(|v| v.as_str()) else {
165                continue;
166            };
167            if exclude_id.is_some_and(|id| {
168                self.read_body(Self::vertex_uuid(table, id))
169                    .ok()
170                    .flatten()
171                    .and_then(|row| row.get(field).and_then(|v| v.as_str()).map(str::to_string))
172                    .is_some_and(|existing| existing == value)
173            }) {
174                continue;
175            }
176            if values.contains(value) {
177                return Err(Error::Database(format!(
178                    "duplicate unique index value for {table}.{field}"
179                )));
180            }
181        }
182        Ok(())
183    }
184
185    fn track_unique_fields(&self, table: &str, record: &serde_json::Value) -> Result<()> {
186        let mut indexes = self
187            .unique_indexes
188            .write()
189            .map_err(|_| Error::Internal("indradb unique index lock poisoned".into()))?;
190        for ((idx_table, field), values) in indexes.iter_mut() {
191            if idx_table != table {
192                continue;
193            }
194            if let Some(value) = record.get(field).and_then(|v| v.as_str()) {
195                values.insert(value.to_string());
196            }
197        }
198        Ok(())
199    }
200
201    fn untrack_unique_fields(&self, table: &str, record: &serde_json::Value) -> Result<()> {
202        let mut indexes = self
203            .unique_indexes
204            .write()
205            .map_err(|_| Error::Internal("indradb unique index lock poisoned".into()))?;
206        for ((idx_table, field), values) in indexes.iter_mut() {
207            if idx_table != table {
208                continue;
209            }
210            if let Some(value) = record.get(field).and_then(|v| v.as_str()) {
211                values.remove(value);
212            }
213        }
214        Ok(())
215    }
216
217    fn rows_for_table(&self, table: &str) -> Result<Vec<serde_json::Value>> {
218        let vertex_type = Self::table_identifier(table)?;
219        let output = self
220            .db
221            .get(RangeVertexQuery::new().t(vertex_type))
222            .map_err(Self::db_err)?;
223        let mut rows = Vec::new();
224        for item in output {
225            if let QueryOutputValue::Vertices(vertices) = item {
226                for vertex in vertices {
227                    if let Some(body) = self.read_body(vertex.id)? {
228                        rows.push(body);
229                    }
230                }
231            }
232        }
233        Ok(rows)
234    }
235
236    fn execute_indra_descriptor(
237        &self,
238        descriptor: &serde_json::Value,
239    ) -> Result<Vec<serde_json::Value>> {
240        let table = descriptor
241            .get("vertex_type")
242            .and_then(|v| v.as_str())
243            .ok_or_else(|| Error::Internal("missing vertex_type in indradb query".into()))?;
244        let mut rows = self.rows_for_table(table)?;
245        if let Some(limit) = descriptor.get("limit").and_then(|v| v.as_u64()) {
246            rows.truncate(limit as usize);
247        }
248        Ok(rows)
249    }
250
251    fn execute_sql_select(&self, q: &str) -> Result<Vec<serde_json::Value>> {
252        let upper = q.to_uppercase();
253        if !upper.starts_with("SELECT ") {
254            return Ok(vec![]);
255        }
256        let from_idx = upper
257            .find(" FROM ")
258            .ok_or_else(|| Error::Internal("missing FROM in select".into()))?;
259        let table = q[from_idx + 6..]
260            .split_whitespace()
261            .next()
262            .unwrap_or("")
263            .trim();
264        if table.is_empty() {
265            return Ok(vec![]);
266        }
267        let rows = self.rows_for_table(table)?;
268        if upper.contains("SELECT id") && !upper.contains("body") {
269            return Ok(rows
270                .iter()
271                .filter_map(|r| {
272                    r.get("id")
273                        .and_then(|id| id.get("id").and_then(|x| x.as_str()))
274                        .or_else(|| r.get("id").and_then(|id| id.as_str()))
275                        .map(|id| serde_json::Value::String(id.to_string()))
276                })
277                .collect());
278        }
279        let mut out = rows;
280        if let Some(limit_idx) = upper.rfind(" LIMIT ") {
281            if let Ok(limit) = q[limit_idx + 7..].trim().parse::<usize>() {
282                out.truncate(limit);
283            }
284        }
285        Ok(out)
286    }
287}
288
289#[async_trait::async_trait]
290impl DatabaseBackend for IndradbBackend {
291    fn engine_id(&self) -> &'static str {
292        ENGINE_ID
293    }
294
295    fn capabilities(&self) -> BackendCapabilities {
296        BackendCapabilities {
297            supports_merge: true,
298            supports_graph_edges: true,
299            telemetry_label: "indradb",
300        }
301    }
302
303    async fn execute_compiled_query(
304        &self,
305        compiled: &CompiledQuery,
306    ) -> Result<Vec<serde_json::Value>> {
307        let q = compiled.query_string.trim();
308        if let Ok(descriptor) = serde_json::from_str::<serde_json::Value>(q) {
309            if descriptor.get("vertex_type").is_some() {
310                return self.execute_indra_descriptor(&descriptor);
311            }
312        }
313        let mut rows = self.execute_sql_select(q)?;
314        rows = valence_core::query::apply_equality_where(rows, compiled);
315        rows = valence_core::query::apply_order_limit_offset(rows, q);
316        Ok(rows)
317    }
318
319    async fn get_record(&self, table: &str, id: &str) -> Result<Option<serde_json::Value>> {
320        self.read_body(Self::vertex_uuid(table, id))
321    }
322
323    async fn create_record(
324        &self,
325        table: &str,
326        content: serde_json::Value,
327    ) -> Result<serde_json::Value> {
328        self.check_unique_fields(table, &content, None)?;
329        let id = storage_id_from_content(&content).unwrap_or_else(uuid_simple);
330        let mut record = content;
331        if let Some(obj) = record.as_object_mut() {
332            let has_string_id = obj.get("id").and_then(|v| v.as_str()).is_some();
333            if !has_string_id {
334                obj.insert("id".into(), record_id_json(table, &id));
335            }
336        }
337        self.write_body(table, &id, record.clone())?;
338        self.track_unique_fields(table, &record)?;
339        Ok(record)
340    }
341
342    async fn update_record(
343        &self,
344        table: &str,
345        id: &str,
346        content: serde_json::Value,
347    ) -> Result<serde_json::Value> {
348        if self.get_record(table, id).await?.is_none() {
349            return Err(Error::NotFound(format!("{table}:{id}")));
350        }
351        if let Some(existing) = self.get_record(table, id).await? {
352            self.untrack_unique_fields(table, &existing)?;
353        }
354        self.check_unique_fields(table, &content, Some(id))?;
355        self.write_body(table, id, content.clone())?;
356        self.track_unique_fields(table, &content)?;
357        Ok(content)
358    }
359
360    async fn merge_record(
361        &self,
362        table: &str,
363        id: &str,
364        patch: serde_json::Value,
365    ) -> Result<serde_json::Value> {
366        let mut record = self
367            .get_record(table, id)
368            .await?
369            .unwrap_or_else(|| serde_json::json!({}));
370        if let Some(existing) = self.get_record(table, id).await? {
371            self.untrack_unique_fields(table, &existing)?;
372        }
373        if let (Some(base), Some(patch_obj)) = (record.as_object_mut(), patch.as_object()) {
374            for (k, v) in patch_obj {
375                base.insert(k.clone(), v.clone());
376            }
377        }
378        self.check_unique_fields(table, &record, Some(id))?;
379        self.write_body(table, id, record.clone())?;
380        self.track_unique_fields(table, &record)?;
381        Ok(record)
382    }
383
384    async fn upsert_record(
385        &self,
386        table: &str,
387        id: &str,
388        content: serde_json::Value,
389    ) -> Result<serde_json::Value> {
390        if let Some(existing) = self.get_record(table, id).await? {
391            self.untrack_unique_fields(table, &existing)?;
392        }
393        self.check_unique_fields(table, &content, Some(id))?;
394        let mut record = content;
395        if let Some(obj) = record.as_object_mut() {
396            obj.insert("id".into(), record_id_json(table, id));
397        }
398        self.write_body(table, id, record.clone())?;
399        self.track_unique_fields(table, &record)?;
400        Ok(record)
401    }
402
403    async fn delete_record(&self, table: &str, id: &str) -> Result<()> {
404        if let Some(existing) = self.get_record(table, id).await? {
405            self.untrack_unique_fields(table, &existing)?;
406        }
407        let vertex_id = Self::vertex_uuid(table, id);
408        self.db
409            .delete(SpecificVertexQuery::single(vertex_id))
410            .map_err(Self::db_err)
411    }
412
413    async fn relate_edge(&self, from: &RecordId, edge_table: &str, to: &RecordId) -> Result<()> {
414        self.ensure_vertex(from.table(), from.id())?;
415        self.ensure_vertex(to.table(), to.id())?;
416        let edge_type = Self::edge_identifier(edge_table)?;
417        let edge = Edge::new(
418            Self::vertex_uuid(from.table(), from.id()),
419            edge_type,
420            Self::vertex_uuid(to.table(), to.id()),
421        );
422        let _ = self.db.create_edge(&edge).map_err(Self::db_err)?;
423        Ok(())
424    }
425
426    async fn unrelate_edge(&self, from: &RecordId, edge_table: &str, to: &RecordId) -> Result<()> {
427        let edge_type = Self::edge_identifier(edge_table)?;
428        let edge = Edge::new(
429            Self::vertex_uuid(from.table(), from.id()),
430            edge_type,
431            Self::vertex_uuid(to.table(), to.id()),
432        );
433        self.db
434            .delete(indradb::SpecificEdgeQuery::single(edge))
435            .map_err(Self::db_err)
436    }
437
438    async fn get_edge_targets(&self, from: &RecordId, edge_table: &str) -> Result<Vec<RecordId>> {
439        let from_uuid = Self::vertex_uuid(from.table(), from.id());
440        let edge_type = Self::edge_identifier(edge_table)?;
441        let output = self
442            .db
443            .get(
444                SpecificVertexQuery::single(from_uuid)
445                    .outbound()
446                    .map_err(Self::id_err)?,
447            )
448            .map_err(Self::db_err)?;
449        let mut targets = Vec::new();
450        for item in output {
451            if let QueryOutputValue::Edges(edges) = item {
452                for edge in edges {
453                    if edge.t != edge_type {
454                        continue;
455                    }
456                    let inbound_table = self
457                        .db
458                        .get(SpecificVertexQuery::single(edge.inbound_id))
459                        .map_err(Self::db_err)?
460                        .into_iter()
461                        .find_map(|value| {
462                            if let QueryOutputValue::Vertices(vertices) = value {
463                                vertices.first().map(|v| v.t.as_str().to_string())
464                            } else {
465                                None
466                            }
467                        })
468                        .unwrap_or_else(|| from.table().to_string());
469                    let body = self
470                        .read_body(edge.inbound_id)?
471                        .unwrap_or_else(|| serde_json::json!({}));
472                    let target_id = storage_id_from_content(&body)
473                        .unwrap_or_else(|| edge.inbound_id.to_string());
474                    targets.push(RecordId::new(inbound_table, target_id));
475                }
476            }
477        }
478        Ok(targets)
479    }
480
481    async fn define_unique_index(&self, table: &str, field: &str) -> Result<()> {
482        let mut indexes = self
483            .unique_indexes
484            .write()
485            .map_err(|_| Error::Internal("indradb unique index lock poisoned".into()))?;
486        let entry = indexes
487            .entry((table.to_string(), field.to_string()))
488            .or_default();
489        if entry.is_empty() {
490            let rows = self.rows_for_table(table)?;
491            for row in rows {
492                if let Some(value) = row.get(field).and_then(|v| v.as_str()) {
493                    entry.insert(value.to_string());
494                }
495            }
496        }
497        Ok(())
498    }
499}
500
501fn record_id_json(table: &str, id: &str) -> serde_json::Value {
502    serde_json::json!({
503        "table": table,
504        "id": id,
505    })
506}
507
508fn storage_id_from_content(content: &serde_json::Value) -> Option<String> {
509    let id_val = content.get("id")?;
510    if let Some(id) = id_val.get("id").and_then(|v| v.as_str()) {
511        return Some(id.to_string());
512    }
513    id_val.as_str().map(|s| s.to_string())
514}
515
516fn uuid_simple() -> String {
517    use std::time::{SystemTime, UNIX_EPOCH};
518    let nanos = SystemTime::now()
519        .duration_since(UNIX_EPOCH)
520        .map_or(0, |d| d.as_nanos());
521    format!("indradb-{nanos}")
522}