use async_trait::async_trait;
use nodedb_types::document::Document;
use nodedb_types::error::NodeDbResult;
use nodedb_types::filter::{EdgeFilter, MetadataFilter};
use nodedb_types::id::{EdgeId, NodeId};
use nodedb_types::result::{QueryResult, SearchResult, SubGraph};
use nodedb_types::value::Value;
#[async_trait]
pub trait NodeDb: Send + Sync {
async fn vector_search(
&self,
collection: &str,
query: &[f32],
k: usize,
filter: Option<&MetadataFilter>,
) -> NodeDbResult<Vec<SearchResult>>;
async fn vector_insert(
&self,
collection: &str,
id: &str,
embedding: &[f32],
metadata: Option<Document>,
) -> NodeDbResult<()>;
async fn vector_delete(&self, collection: &str, id: &str) -> NodeDbResult<()>;
async fn graph_traverse(
&self,
start: &NodeId,
depth: u8,
edge_filter: Option<&EdgeFilter>,
) -> NodeDbResult<SubGraph>;
async fn graph_insert_edge(
&self,
from: &NodeId,
to: &NodeId,
edge_type: &str,
properties: Option<Document>,
) -> NodeDbResult<EdgeId>;
async fn graph_delete_edge(&self, edge_id: &EdgeId) -> NodeDbResult<()>;
async fn document_get(&self, collection: &str, id: &str) -> NodeDbResult<Option<Document>>;
async fn document_put(&self, collection: &str, doc: Document) -> NodeDbResult<()>;
async fn document_delete(&self, collection: &str, id: &str) -> NodeDbResult<()>;
async fn vector_insert_field(
&self,
collection: &str,
field_name: &str,
id: &str,
embedding: &[f32],
metadata: Option<Document>,
) -> NodeDbResult<()> {
let _ = field_name;
self.vector_insert(collection, id, embedding, metadata)
.await
}
async fn vector_search_field(
&self,
collection: &str,
field_name: &str,
query: &[f32],
k: usize,
filter: Option<&MetadataFilter>,
) -> NodeDbResult<Vec<SearchResult>> {
let _ = field_name;
self.vector_search(collection, query, k, filter).await
}
async fn graph_shortest_path(
&self,
from: &NodeId,
to: &NodeId,
max_depth: u8,
edge_filter: Option<&EdgeFilter>,
) -> NodeDbResult<Option<Vec<NodeId>>> {
let _ = (from, to, max_depth, edge_filter);
Ok(None)
}
async fn text_search(
&self,
collection: &str,
query: &str,
top_k: usize,
) -> NodeDbResult<Vec<SearchResult>> {
let _ = (collection, query, top_k);
Ok(Vec::new())
}
async fn batch_vector_insert(
&self,
collection: &str,
vectors: &[(&str, &[f32])],
) -> NodeDbResult<()> {
for &(id, embedding) in vectors {
self.vector_insert(collection, id, embedding, None).await?;
}
Ok(())
}
async fn batch_graph_insert_edges(&self, edges: &[(&str, &str, &str)]) -> NodeDbResult<()> {
for &(from, to, label) in edges {
self.graph_insert_edge(&NodeId::new(from), &NodeId::new(to), label, None)
.await?;
}
Ok(())
}
async fn execute_sql(&self, query: &str, params: &[Value]) -> NodeDbResult<QueryResult>;
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashMap;
struct MockDb;
#[async_trait]
impl NodeDb for MockDb {
async fn vector_search(
&self,
_collection: &str,
_query: &[f32],
_k: usize,
_filter: Option<&MetadataFilter>,
) -> NodeDbResult<Vec<SearchResult>> {
Ok(vec![SearchResult {
id: "vec-1".into(),
node_id: None,
distance: 0.1,
metadata: HashMap::new(),
}])
}
async fn vector_insert(
&self,
_collection: &str,
_id: &str,
_embedding: &[f32],
_metadata: Option<Document>,
) -> NodeDbResult<()> {
Ok(())
}
async fn vector_delete(&self, _collection: &str, _id: &str) -> NodeDbResult<()> {
Ok(())
}
async fn graph_traverse(
&self,
_start: &NodeId,
_depth: u8,
_edge_filter: Option<&EdgeFilter>,
) -> NodeDbResult<SubGraph> {
Ok(SubGraph::empty())
}
async fn graph_insert_edge(
&self,
from: &NodeId,
to: &NodeId,
edge_type: &str,
_properties: Option<Document>,
) -> NodeDbResult<EdgeId> {
Ok(EdgeId::from_components(
from.as_str(),
to.as_str(),
edge_type,
))
}
async fn graph_delete_edge(&self, _edge_id: &EdgeId) -> NodeDbResult<()> {
Ok(())
}
async fn document_get(
&self,
_collection: &str,
id: &str,
) -> NodeDbResult<Option<Document>> {
let mut doc = Document::new(id);
doc.set("title", Value::String("test".into()));
Ok(Some(doc))
}
async fn document_put(&self, _collection: &str, _doc: Document) -> NodeDbResult<()> {
Ok(())
}
async fn document_delete(&self, _collection: &str, _id: &str) -> NodeDbResult<()> {
Ok(())
}
async fn execute_sql(&self, _query: &str, _params: &[Value]) -> NodeDbResult<QueryResult> {
Ok(QueryResult::empty())
}
}
#[test]
fn trait_is_object_safe() {
fn _accepts_dyn(_db: &dyn NodeDb) {}
let db = MockDb;
_accepts_dyn(&db);
}
#[test]
fn trait_works_with_arc() {
use std::sync::Arc;
let db: Arc<dyn NodeDb> = Arc::new(MockDb);
let _ = db;
}
#[tokio::test]
async fn mock_vector_search() {
let db = MockDb;
let results = db
.vector_search("embeddings", &[0.1, 0.2, 0.3], 5, None)
.await
.unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].id, "vec-1");
assert!(results[0].distance < 1.0);
}
#[tokio::test]
async fn mock_vector_insert_and_delete() {
let db = MockDb;
db.vector_insert("coll", "v1", &[1.0, 2.0], None)
.await
.unwrap();
db.vector_delete("coll", "v1").await.unwrap();
}
#[tokio::test]
async fn mock_graph_operations() {
let db = MockDb;
let start = NodeId::new("alice");
let subgraph = db.graph_traverse(&start, 2, None).await.unwrap();
assert_eq!(subgraph.node_count(), 0);
let from = NodeId::new("alice");
let to = NodeId::new("bob");
let edge_id = db
.graph_insert_edge(&from, &to, "KNOWS", None)
.await
.unwrap();
assert_eq!(edge_id.as_str(), "alice--KNOWS-->bob");
db.graph_delete_edge(&edge_id).await.unwrap();
}
#[tokio::test]
async fn mock_document_operations() {
let db = MockDb;
let doc = db.document_get("notes", "n1").await.unwrap().unwrap();
assert_eq!(doc.id, "n1");
assert_eq!(doc.get_str("title"), Some("test"));
let mut new_doc = Document::new("n2");
new_doc.set("body", Value::String("hello".into()));
db.document_put("notes", new_doc).await.unwrap();
db.document_delete("notes", "n1").await.unwrap();
}
#[tokio::test]
async fn mock_execute_sql() {
let db = MockDb;
let result = db.execute_sql("SELECT 1", &[]).await.unwrap();
assert_eq!(result.row_count(), 0);
}
#[tokio::test]
async fn unified_api_pattern() {
use std::sync::Arc;
let db: Arc<dyn NodeDb> = Arc::new(MockDb);
let results = db
.vector_search("knowledge_base", &[0.1, 0.2], 5, None)
.await
.unwrap();
assert!(!results.is_empty());
let start = NodeId::new(results[0].id.clone());
let _subgraph = db.graph_traverse(&start, 2, None).await.unwrap();
let doc = Document::new("note-1");
db.document_put("notes", doc).await.unwrap();
}
}