1use 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
16pub const ENGINE_ID: &str = KnownEngines::INDRADB;
18
19pub const PRIMARY: DatabaseFromEngine = Database::from_engine("primary", ENGINE_ID);
21
22type IndraDb = indradb::Database<MemoryDatastore>;
23
24pub 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 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)] fn id_err(e: indradb::ValidationError) -> Error {
89 Error::Validation(format!("invalid indradb identifier: {e:?}"))
90 }
91
92 #[allow(clippy::needless_pass_by_value)] 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 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 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 fn write_fields(&self, table: &str, id: &str, record: serde_json::Value) -> Result<()> {
167 let vertex = self.ensure_vertex(table, id)?;
168 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 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}