use crate::types::{ETypeId, NodeId, PropKeyId, PropValue};
use std::collections::HashMap;
#[derive(Debug, Clone)]
pub struct NodeRef {
pub id: NodeId,
pub key: String,
pub props: HashMap<String, PropValue>,
}
impl NodeRef {
pub fn new(id: NodeId, key: String) -> Self {
Self {
id,
key,
props: HashMap::new(),
}
}
pub fn with_props(id: NodeId, key: String, props: HashMap<String, PropValue>) -> Self {
Self { id, key, props }
}
pub fn get_prop(&self, name: &str) -> Option<&PropValue> {
self.props.get(name)
}
}
pub struct InsertBuilder<'a, F, R>
where
F: FnMut(InsertData) -> R,
{
executor: &'a mut F,
_node_type: String,
}
#[derive(Debug, Clone)]
pub struct InsertData {
pub key: Option<String>,
pub props: HashMap<PropKeyId, PropValue>,
}
impl InsertData {
pub fn new() -> Self {
Self {
key: None,
props: HashMap::new(),
}
}
pub fn with_key(mut self, key: impl Into<String>) -> Self {
self.key = Some(key.into());
self
}
pub fn with_prop(mut self, key_id: PropKeyId, value: PropValue) -> Self {
self.props.insert(key_id, value);
self
}
}
impl Default for InsertData {
fn default() -> Self {
Self::new()
}
}
impl<'a, F, R> InsertBuilder<'a, F, R>
where
F: FnMut(InsertData) -> R,
{
pub fn new(node_type: impl Into<String>, executor: &'a mut F) -> Self {
Self {
executor,
_node_type: node_type.into(),
}
}
pub fn values(self, data: InsertData) -> R {
(self.executor)(data)
}
}
pub struct UpdateBuilder<'a, F, R>
where
F: FnMut(NodeId, HashMap<PropKeyId, Option<PropValue>>) -> R,
{
executor: &'a mut F,
node_id: Option<NodeId>,
updates: HashMap<PropKeyId, Option<PropValue>>,
}
impl<'a, F, R> UpdateBuilder<'a, F, R>
where
F: FnMut(NodeId, HashMap<PropKeyId, Option<PropValue>>) -> R,
{
pub fn new(executor: &'a mut F) -> Self {
Self {
executor,
node_id: None,
updates: HashMap::new(),
}
}
pub fn where_id(mut self, node_id: NodeId) -> Self {
self.node_id = Some(node_id);
self
}
pub fn set(mut self, prop_key_id: PropKeyId, value: PropValue) -> Self {
self.updates.insert(prop_key_id, Some(value));
self
}
pub fn unset(mut self, prop_key_id: PropKeyId) -> Self {
self.updates.insert(prop_key_id, None);
self
}
pub fn execute(self) -> R {
let node_id = self
.node_id
.expect("Update requires a node ID (use where_id())");
(self.executor)(node_id, self.updates)
}
}
pub struct DeleteBuilder<'a, F, R>
where
F: FnMut(NodeId) -> R,
{
executor: &'a mut F,
node_id: Option<NodeId>,
}
impl<'a, F, R> DeleteBuilder<'a, F, R>
where
F: FnMut(NodeId) -> R,
{
pub fn new(executor: &'a mut F) -> Self {
Self {
executor,
node_id: None,
}
}
pub fn where_id(mut self, node_id: NodeId) -> Self {
self.node_id = Some(node_id);
self
}
pub fn execute(self) -> R {
let node_id = self
.node_id
.expect("Delete requires a node ID (use where_id())");
(self.executor)(node_id)
}
}
pub struct LinkBuilder<'a, F, R>
where
F: FnMut(NodeId, ETypeId, NodeId, HashMap<PropKeyId, PropValue>) -> R,
{
executor: &'a mut F,
src: NodeId,
etype: ETypeId,
dst: Option<NodeId>,
props: HashMap<PropKeyId, PropValue>,
}
impl<'a, F, R> LinkBuilder<'a, F, R>
where
F: FnMut(NodeId, ETypeId, NodeId, HashMap<PropKeyId, PropValue>) -> R,
{
pub fn new(src: NodeId, etype: ETypeId, executor: &'a mut F) -> Self {
Self {
executor,
src,
etype,
dst: None,
props: HashMap::new(),
}
}
pub fn to(mut self, dst: NodeId) -> Self {
self.dst = Some(dst);
self
}
pub fn with_prop(mut self, prop_key_id: PropKeyId, value: PropValue) -> Self {
self.props.insert(prop_key_id, value);
self
}
pub fn execute(self) -> R {
let dst = self.dst.expect("Link requires a destination (use to())");
(self.executor)(self.src, self.etype, dst, self.props)
}
}
pub struct UnlinkBuilder<'a, F, R>
where
F: FnMut(NodeId, ETypeId, NodeId) -> R,
{
executor: &'a mut F,
src: NodeId,
etype: ETypeId,
dst: Option<NodeId>,
}
impl<'a, F, R> UnlinkBuilder<'a, F, R>
where
F: FnMut(NodeId, ETypeId, NodeId) -> R,
{
pub fn new(src: NodeId, etype: ETypeId, executor: &'a mut F) -> Self {
Self {
executor,
src,
etype,
dst: None,
}
}
pub fn from_node(mut self, dst: NodeId) -> Self {
self.dst = Some(dst);
self
}
pub fn execute(self) -> R {
let dst = self
.dst
.expect("Unlink requires a destination (use from_node())");
(self.executor)(self.src, self.etype, dst)
}
}
pub struct UpdateEdgeBuilder<'a, F, R>
where
F: FnMut(NodeId, ETypeId, NodeId, HashMap<PropKeyId, Option<PropValue>>) -> R,
{
executor: &'a mut F,
src: NodeId,
etype: ETypeId,
dst: NodeId,
updates: HashMap<PropKeyId, Option<PropValue>>,
}
impl<'a, F, R> UpdateEdgeBuilder<'a, F, R>
where
F: FnMut(NodeId, ETypeId, NodeId, HashMap<PropKeyId, Option<PropValue>>) -> R,
{
pub fn new(src: NodeId, etype: ETypeId, dst: NodeId, executor: &'a mut F) -> Self {
Self {
executor,
src,
etype,
dst,
updates: HashMap::new(),
}
}
pub fn set(mut self, prop_key_id: PropKeyId, value: PropValue) -> Self {
self.updates.insert(prop_key_id, Some(value));
self
}
pub fn unset(mut self, prop_key_id: PropKeyId) -> Self {
self.updates.insert(prop_key_id, None);
self
}
pub fn execute(self) -> R {
(self.executor)(self.src, self.etype, self.dst, self.updates)
}
}
#[derive(Debug, Clone)]
pub enum BatchOp {
Insert(InsertData),
Update {
node_id: NodeId,
updates: HashMap<PropKeyId, Option<PropValue>>,
},
Delete { node_id: NodeId },
Link {
src: NodeId,
etype: ETypeId,
dst: NodeId,
props: HashMap<PropKeyId, PropValue>,
},
Unlink {
src: NodeId,
etype: ETypeId,
dst: NodeId,
},
UpdateEdge {
src: NodeId,
etype: ETypeId,
dst: NodeId,
updates: HashMap<PropKeyId, Option<PropValue>>,
},
}
#[derive(Debug, Default)]
pub struct BatchBuilder {
ops: Vec<BatchOp>,
}
impl BatchBuilder {
pub fn new() -> Self {
Self { ops: Vec::new() }
}
pub fn insert(mut self, data: InsertData) -> Self {
self.ops.push(BatchOp::Insert(data));
self
}
pub fn update(mut self, node_id: NodeId, updates: HashMap<PropKeyId, Option<PropValue>>) -> Self {
self.ops.push(BatchOp::Update { node_id, updates });
self
}
pub fn delete(mut self, node_id: NodeId) -> Self {
self.ops.push(BatchOp::Delete { node_id });
self
}
pub fn link(mut self, src: NodeId, etype: ETypeId, dst: NodeId) -> Self {
self.ops.push(BatchOp::Link {
src,
etype,
dst,
props: HashMap::new(),
});
self
}
pub fn link_with_props(
mut self,
src: NodeId,
etype: ETypeId,
dst: NodeId,
props: HashMap<PropKeyId, PropValue>,
) -> Self {
self.ops.push(BatchOp::Link {
src,
etype,
dst,
props,
});
self
}
pub fn unlink(mut self, src: NodeId, etype: ETypeId, dst: NodeId) -> Self {
self.ops.push(BatchOp::Unlink { src, etype, dst });
self
}
pub fn build(self) -> Vec<BatchOp> {
self.ops
}
pub fn len(&self) -> usize {
self.ops.len()
}
pub fn is_empty(&self) -> bool {
self.ops.is_empty()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_node_ref() {
let node = NodeRef::new(1, "alice".to_string());
assert_eq!(node.id, 1);
assert_eq!(node.key, "alice");
assert!(node.props.is_empty());
}
#[test]
fn test_node_ref_with_props() {
let mut props = HashMap::new();
props.insert("name".to_string(), PropValue::String("Alice".to_string()));
props.insert("age".to_string(), PropValue::I64(30));
let node = NodeRef::with_props(1, "alice".to_string(), props);
assert_eq!(
node.get_prop("name"),
Some(&PropValue::String("Alice".to_string()))
);
assert_eq!(node.get_prop("age"), Some(&PropValue::I64(30)));
assert_eq!(node.get_prop("unknown"), None);
}
#[test]
fn test_insert_data() {
let data = InsertData::new()
.with_key("alice")
.with_prop(1, PropValue::String("Alice".to_string()))
.with_prop(2, PropValue::I64(30));
assert_eq!(data.key, Some("alice".to_string()));
assert_eq!(data.props.len(), 2);
}
#[test]
fn test_insert_builder() {
let mut executed = false;
let mut captured_data: Option<InsertData> = None;
let mut executor = |data: InsertData| {
executed = true;
captured_data = Some(data);
1u64 };
let data = InsertData::new().with_key("test");
let result = InsertBuilder::new("User", &mut executor).values(data);
assert!(executed);
assert_eq!(result, 1);
assert_eq!(captured_data.unwrap().key, Some("test".to_string()));
}
#[test]
fn test_update_builder() {
let mut captured: Option<(NodeId, HashMap<PropKeyId, Option<PropValue>>)> = None;
let mut executor = |node_id: NodeId, updates: HashMap<PropKeyId, Option<PropValue>>| {
captured = Some((node_id, updates));
};
UpdateBuilder::new(&mut executor)
.where_id(42)
.set(1, PropValue::String("Updated".to_string()))
.unset(2)
.execute();
let (node_id, updates) = captured.unwrap();
assert_eq!(node_id, 42);
assert_eq!(updates.len(), 2);
assert!(updates.get(&1).unwrap().is_some());
assert!(updates.get(&2).unwrap().is_none());
}
#[test]
fn test_delete_builder() {
let mut deleted_id: Option<NodeId> = None;
let mut executor = |node_id: NodeId| {
deleted_id = Some(node_id);
true
};
let result = DeleteBuilder::new(&mut executor).where_id(42).execute();
assert!(result);
assert_eq!(deleted_id, Some(42));
}
#[test]
fn test_link_builder() {
let mut captured: Option<(NodeId, ETypeId, NodeId, HashMap<PropKeyId, PropValue>)> = None;
let mut executor = |src, etype, dst, props| {
captured = Some((src, etype, dst, props));
};
LinkBuilder::new(1, 10, &mut executor)
.to(2)
.with_prop(100, PropValue::F64(1.5))
.execute();
let (src, etype, dst, props) = captured.unwrap();
assert_eq!(src, 1);
assert_eq!(etype, 10);
assert_eq!(dst, 2);
assert_eq!(props.get(&100), Some(&PropValue::F64(1.5)));
}
#[test]
fn test_unlink_builder() {
let mut captured: Option<(NodeId, ETypeId, NodeId)> = None;
let mut executor = |src, etype, dst| {
captured = Some((src, etype, dst));
true
};
let result = UnlinkBuilder::new(1, 10, &mut executor)
.from_node(2)
.execute();
assert!(result);
let (src, etype, dst) = captured.unwrap();
assert_eq!(src, 1);
assert_eq!(etype, 10);
assert_eq!(dst, 2);
}
#[test]
fn test_batch_builder() {
let batch = BatchBuilder::new()
.insert(InsertData::new().with_key("alice"))
.insert(InsertData::new().with_key("bob"))
.link(1, 10, 2)
.update(1, HashMap::new())
.delete(3)
.unlink(1, 10, 2)
.build();
assert_eq!(batch.len(), 6);
assert!(matches!(batch[0], BatchOp::Insert(_)));
assert!(matches!(batch[1], BatchOp::Insert(_)));
assert!(matches!(batch[2], BatchOp::Link { .. }));
assert!(matches!(batch[3], BatchOp::Update { .. }));
assert!(matches!(batch[4], BatchOp::Delete { .. }));
assert!(matches!(batch[5], BatchOp::Unlink { .. }));
}
#[test]
fn test_batch_builder_empty() {
let batch = BatchBuilder::new();
assert!(batch.is_empty());
assert_eq!(batch.len(), 0);
}
}