use crate::schema::{edge_col, node_col};
use arrow::array::{Array, RecordBatch, StringArray};
use std::sync::Arc;
#[derive(Debug, thiserror::Error)]
pub enum RenameError {
#[error("Node not found: {0}")]
NodeNotFound(String),
#[error("Target ID already exists: {0}")]
TargetExists(String),
#[error("Arrow error: {0}")]
Arrow(#[from] arrow::error::ArrowError),
}
pub type Result<T> = std::result::Result<T, RenameError>;
pub fn rename_node(
nodes_batch: &RecordBatch,
edges_batch: &RecordBatch,
old_id: &str,
new_id: &str,
) -> Result<(RecordBatch, RecordBatch)> {
let ids = nodes_batch
.column(node_col::ID)
.as_any()
.downcast_ref::<StringArray>()
.expect("id column");
let found = (0..nodes_batch.num_rows()).any(|i| ids.value(i) == old_id);
if !found {
return Err(RenameError::NodeNotFound(old_id.to_string()));
}
let conflict = (0..nodes_batch.num_rows()).any(|i| ids.value(i) == new_id);
if conflict {
return Err(RenameError::TargetExists(new_id.to_string()));
}
let new_nodes = rename_in_nodes(nodes_batch, old_id, new_id)?;
let new_edges = rename_in_edges(edges_batch, old_id, new_id)?;
Ok((new_nodes, new_edges))
}
fn rename_in_nodes(batch: &RecordBatch, old_id: &str, new_id: &str) -> Result<RecordBatch> {
let mut columns: Vec<arrow::array::ArrayRef> = Vec::with_capacity(batch.num_columns());
for col_idx in 0..batch.num_columns() {
match col_idx {
node_col::ID => {
let old = batch
.column(col_idx)
.as_any()
.downcast_ref::<StringArray>()
.expect("id");
let vals: Vec<String> = (0..batch.num_rows())
.map(|i| {
let v = old.value(i);
if v == old_id {
new_id.to_string()
} else {
v.to_string()
}
})
.collect();
let refs: Vec<&str> = vals.iter().map(|s| s.as_str()).collect();
columns.push(Arc::new(StringArray::from(refs)));
}
node_col::PARENT_ID => {
let old = batch
.column(col_idx)
.as_any()
.downcast_ref::<StringArray>()
.expect("parent_id");
let vals: Vec<Option<String>> = (0..batch.num_rows())
.map(|i| {
if old.is_null(i) {
None
} else {
let v = old.value(i);
if v == old_id {
Some(new_id.to_string())
} else {
Some(v.to_string())
}
}
})
.collect();
let refs: Vec<Option<&str>> = vals.iter().map(|s| s.as_deref()).collect();
columns.push(Arc::new(StringArray::from(refs)));
}
_ => {
columns.push(batch.column(col_idx).clone());
}
}
}
let schema = Arc::new(batch.schema().as_ref().clone());
Ok(RecordBatch::try_new(schema, columns)?)
}
fn rename_in_edges(batch: &RecordBatch, old_id: &str, new_id: &str) -> Result<RecordBatch> {
if batch.num_rows() == 0 {
return Ok(batch.clone());
}
let mut columns: Vec<arrow::array::ArrayRef> = Vec::with_capacity(batch.num_columns());
for col_idx in 0..batch.num_columns() {
match col_idx {
edge_col::SOURCE_ID | edge_col::TARGET_ID => {
let old = batch
.column(col_idx)
.as_any()
.downcast_ref::<StringArray>()
.expect("edge string column");
let vals: Vec<String> = (0..batch.num_rows())
.map(|i| {
let v = old.value(i);
if v == old_id {
new_id.to_string()
} else {
v.to_string()
}
})
.collect();
let refs: Vec<&str> = vals.iter().map(|s| s.as_str()).collect();
columns.push(Arc::new(StringArray::from(refs)));
}
_ => {
columns.push(batch.column(col_idx).clone());
}
}
}
let schema = Arc::new(batch.schema().as_ref().clone());
Ok(RecordBatch::try_new(schema, columns)?)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::mcp_tools::{QueryFilter, codegraph_query_objects};
use crate::schema::{
CodeEdge, CodeEdgePredicate, CodeNode, CodeNodeKind, build_code_edges_batch,
build_code_nodes_batch,
};
fn sample_nodes() -> Vec<CodeNode> {
vec![
CodeNode {
id: "class:brain/store.py::DualStore".into(),
kind: CodeNodeKind::Class,
parent_id: Some("mod:brain/store.py".into()),
name: "DualStore".into(),
signature: None,
docstring: Some("Dual store.".into()),
body_hash: Some("ds_h".into()),
body: None,
loc: Some(100),
cyclomatic_complexity: None,
coverage_pct: None,
last_modified: None,
..Default::default()
},
CodeNode {
id: "func:brain/store.py::DualStore::promote".into(),
kind: CodeNodeKind::Method,
parent_id: Some("class:brain/store.py::DualStore".into()),
name: "promote".into(),
signature: Some("def promote(self)".into()),
docstring: None,
body_hash: Some("prom_h".into()),
body: None,
loc: Some(20),
cyclomatic_complexity: Some(3),
coverage_pct: None,
last_modified: None,
..Default::default()
},
CodeNode {
id: "func:brain/signal.py::fuse".into(),
kind: CodeNodeKind::Function,
parent_id: None,
name: "fuse".into(),
signature: Some("def fuse(signals)".into()),
docstring: None,
body_hash: Some("fuse_h".into()),
body: None,
loc: Some(30),
cyclomatic_complexity: Some(5),
coverage_pct: None,
last_modified: None,
..Default::default()
},
]
}
fn sample_edges() -> Vec<CodeEdge> {
vec![
CodeEdge {
source_id: "func:brain/signal.py::fuse".into(),
target_id: "func:brain/store.py::DualStore::promote".into(),
predicate: CodeEdgePredicate::Calls,
weight: Some(1.0),
commit_id: None,
},
CodeEdge {
source_id: "class:brain/store.py::DualStore".into(),
target_id: "func:brain/store.py::DualStore::promote".into(),
predicate: CodeEdgePredicate::Contains,
weight: None,
commit_id: None,
},
]
}
#[test]
fn test_rename_function() {
let nodes = build_code_nodes_batch(&sample_nodes()).unwrap();
let edges = build_code_edges_batch(&sample_edges()).unwrap();
let (new_nodes, new_edges) = rename_node(
&nodes,
&edges,
"func:brain/signal.py::fuse",
"func:brain/signal.py::fuse_signals",
)
.unwrap();
let old_result = codegraph_query_objects(
&new_nodes,
&QueryFilter {
name_contains: Some("fuse".into()),
..Default::default()
},
)
.unwrap();
assert_eq!(old_result.nodes.len(), 1);
assert_eq!(old_result.nodes[0].id, "func:brain/signal.py::fuse_signals");
let sources = new_edges
.column(edge_col::SOURCE_ID)
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
assert_eq!(sources.value(0), "func:brain/signal.py::fuse_signals");
}
#[test]
fn test_rename_class_updates_child_parent() {
let nodes = build_code_nodes_batch(&sample_nodes()).unwrap();
let edges = build_code_edges_batch(&sample_edges()).unwrap();
let (new_nodes, _) = rename_node(
&nodes,
&edges,
"class:brain/store.py::DualStore",
"class:brain/store.py::GraphStore",
)
.unwrap();
let parent_ids = new_nodes
.column(node_col::PARENT_ID)
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
assert_eq!(parent_ids.value(1), "class:brain/store.py::GraphStore");
}
#[test]
fn test_rename_updates_edge_target() {
let nodes = build_code_nodes_batch(&sample_nodes()).unwrap();
let edges = build_code_edges_batch(&sample_edges()).unwrap();
let (_, new_edges) = rename_node(
&nodes,
&edges,
"func:brain/store.py::DualStore::promote",
"func:brain/store.py::DualStore::elevate",
)
.unwrap();
let targets = new_edges
.column(edge_col::TARGET_ID)
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
assert_eq!(targets.value(0), "func:brain/store.py::DualStore::elevate");
assert_eq!(targets.value(1), "func:brain/store.py::DualStore::elevate");
}
#[test]
fn test_rename_preserves_unaffected_nodes() {
let nodes = build_code_nodes_batch(&sample_nodes()).unwrap();
let edges = build_code_edges_batch(&sample_edges()).unwrap();
let (new_nodes, _) = rename_node(
&nodes,
&edges,
"func:brain/signal.py::fuse",
"func:brain/signal.py::fuse2",
)
.unwrap();
assert_eq!(new_nodes.num_rows(), 3);
let result = codegraph_query_objects(
&new_nodes,
&QueryFilter {
name_contains: Some("DualStore".into()),
..Default::default()
},
)
.unwrap();
assert_eq!(result.nodes[0].docstring, Some("Dual store.".into()));
}
#[test]
fn test_rename_nonexistent_node_errors() {
let nodes = build_code_nodes_batch(&sample_nodes()).unwrap();
let edges = build_code_edges_batch(&sample_edges()).unwrap();
let result = rename_node(&nodes, &edges, "func:nonexistent::foo", "func:new::foo");
assert!(result.is_err());
}
#[test]
fn test_rename_to_existing_id_errors() {
let nodes = build_code_nodes_batch(&sample_nodes()).unwrap();
let edges = build_code_edges_batch(&sample_edges()).unwrap();
let result = rename_node(
&nodes,
&edges,
"func:brain/signal.py::fuse",
"class:brain/store.py::DualStore", );
assert!(result.is_err());
}
}