1use 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
16pub const ENGINE_ID: &str = KnownEngines::INDRADB;
18
19pub const PRIMARY: DatabaseFromEngine = Database::from_engine("primary", ENGINE_ID);
21
22const BODY_PROPERTY: &str = "body";
23
24type IndraDb = indradb::Database<MemoryDatastore>;
25
26pub 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 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}