1use crate::errors::AppError;
11use crate::i18n::errors_msg;
12use crate::output::{self, OutputFormat};
13use crate::paths::AppPaths;
14use crate::storage::connection::open_rw;
15use crate::storage::entities;
16use rusqlite::params;
17use serde::Serialize;
18
19#[derive(clap::Args)]
20#[command(after_long_help = "EXAMPLES:\n \
21 # Merge two source entities into a target\n \
22 sqlite-graphrag merge-entities --names auth,authentication --into auth-service\n\n \
23 # Merge three sources into one target across a namespace\n \
24 sqlite-graphrag merge-entities --names svc-a,svc-b,old-svc --into canonical-service --namespace my-project\n\n \
25 # Merge by ID (unambiguous when homonyms exist across namespaces)\n \
26 sqlite-graphrag merge-entities --ids 12,17 --into-id 3\n\n\
27NOTE:\n \
28 --names is a comma-separated list of source entity names.\n \
29 --into is the target entity name and must already exist.\n \
30 --ids / --into-id select entities by ID; IDs are globally unique so they\n \
31 disambiguate homonyms. They conflict with --names / --into respectively\n \
32 and must belong to the resolved namespace.\n \
33 Source entities are deleted after the merge; the target is preserved.\n \
34 Duplicate relationships (same endpoints + relation) are removed automatically.\n \
35 Run `sqlite-graphrag cleanup-orphans` afterwards if sources had no other links.")]
36pub struct MergeEntitiesArgs {
38 #[arg(
40 long,
41 value_delimiter = ',',
42 value_name = "NAMES",
43 required_unless_present = "ids",
44 conflicts_with = "ids"
45 )]
46 pub names: Vec<String>,
47 #[arg(long, value_delimiter = ',', value_name = "IDS")]
51 pub ids: Vec<i64>,
52 #[arg(
54 long,
55 value_name = "TARGET",
56 required_unless_present = "into_id",
57 conflicts_with = "into_id"
58 )]
59 pub into: Option<String>,
60 #[arg(long, value_name = "TARGET_ID")]
62 pub into_id: Option<i64>,
63 #[arg(long)]
65 pub namespace: Option<String>,
66 #[arg(long, value_enum, default_value = "json")]
68 pub format: OutputFormat,
69 #[arg(long, hide = true, help = "No-op; JSON is always emitted on stdout")]
71 pub json: bool,
72 #[arg(long)]
74 pub db: Option<String>,
75 #[arg(long, default_value_t = false, hide = false)]
80 pub cross_namespace: bool,
81}
82
83#[derive(Serialize)]
84struct MergeEntitiesResponse {
85 action: String,
86 sources: Vec<String>,
87 target: String,
88 namespace: String,
89 target_id: i64,
91 relationships_moved: usize,
92 entities_removed: usize,
93 elapsed_ms: u64,
95}
96
97fn find_entity_name_by_id(
104 conn: &rusqlite::Connection,
105 namespace: &str,
106 id: i64,
107 enforce_namespace: bool,
108) -> Result<(String, String), AppError> {
109 let mut stmt = if enforce_namespace {
110 conn.prepare_cached(
111 "SELECT name, namespace FROM entities WHERE id = ?1 AND namespace = ?2",
112 )?
113 } else {
114 conn.prepare_cached("SELECT name, namespace FROM entities WHERE id = ?1")?
115 };
116 let row = if enforce_namespace {
117 stmt.query_row(params![id, namespace], |r| {
118 Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?))
119 })
120 } else {
121 stmt.query_row(params![id], |r| {
122 Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?))
123 })
124 };
125 match row {
126 Ok((name, ns_actual)) => Ok((name, ns_actual)),
127 Err(rusqlite::Error::QueryReturnedNoRows) => Err(AppError::NotFound(format!(
128 "entity id={id} not found in namespace '{namespace}'"
129 ))),
130 Err(e) => Err(AppError::Database(e)),
131 }
132}
133
134pub fn run(args: MergeEntitiesArgs) -> Result<(), AppError> {
136 let inicio = std::time::Instant::now();
137
138 if args.names.is_empty() && args.ids.is_empty() {
139 return Err(AppError::Validation(
140 "--names or --ids must contain at least one source entity".to_string(),
141 ));
142 }
143
144 if let Some(target_id) = args.into_id {
147 if args.ids.contains(&target_id) {
148 return Err(AppError::Validation(
149 crate::i18n::validation::self_merge_id_in_ids(target_id),
150 ));
151 }
152 }
153 if let Some(ref target_name) = args.into {
154 if args.names.iter().any(|n| n == target_name) {
155 return Err(AppError::Validation(
156 crate::i18n::validation::self_merge_name_in_names(target_name),
157 ));
158 }
159 }
160
161 let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
162 let paths = AppPaths::resolve(args.db.as_deref())?;
163
164 crate::storage::connection::ensure_db_ready(&paths)?;
165
166 let mut conn = open_rw(&paths.db)?;
167
168 let (target_id, target_name) = match args.into_id {
171 Some(id) => {
172 let (name, _ns_actual) = find_entity_name_by_id(&conn, &namespace, id, true)?;
175 (id, name)
176 }
177 None => {
178 let Some(name) = args.into.clone() else {
179 return Err(AppError::Validation(
180 "--into or --into-id is required".to_string(),
181 ));
182 };
183 let id = entities::find_entity_id(&conn, &namespace, &name)?.ok_or_else(|| {
184 AppError::NotFound(errors_msg::entity_not_found(&name, &namespace))
185 })?;
186 (id, name)
187 }
188 };
189
190 let mut source_ids: Vec<i64> = Vec::with_capacity(args.names.len() + args.ids.len());
194 let mut source_names: Vec<String> = Vec::with_capacity(source_ids.capacity());
195 if !args.ids.is_empty() {
196 for &id in &args.ids {
197 if id == target_id {
198 return Err(AppError::Validation(
199 crate::i18n::validation::self_merge_id(id, target_id),
200 ));
201 }
202 let (name, ns_actual) =
206 find_entity_name_by_id(&conn, &namespace, id, !args.cross_namespace)?;
207 if args.cross_namespace && ns_actual != namespace {
208 tracing::warn!(
209 target: "merge_entities",
210 from_id = id,
211 from_namespace = %ns_actual,
212 to_namespace = %namespace,
213 "cross-namespace merge"
214 );
215 }
216 if !source_ids.contains(&id) {
217 source_ids.push(id);
218 source_names.push(name);
219 }
220 }
221 } else {
222 for name in &args.names {
223 if name == &target_name {
224 return Err(AppError::Validation(
225 crate::i18n::validation::self_merge_name(name, &target_name),
226 ));
227 }
228 let id = entities::find_entity_id(&conn, &namespace, name)?.ok_or_else(|| {
229 AppError::NotFound(errors_msg::entity_not_found(name, &namespace))
230 })?;
231 if id == target_id {
232 return Err(AppError::Validation(
233 crate::i18n::validation::self_merge_name_resolves_to_target(name, target_id),
234 ));
235 }
236 if !source_ids.contains(&id) {
237 source_ids.push(id);
238 source_names.push(name.clone());
239 }
240 }
241 }
242
243 if source_ids.is_empty() {
244 return Err(AppError::Validation(
245 "no valid source entities to merge (all names equal the target or were duplicates)"
246 .to_string(),
247 ));
248 }
249
250 let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
251
252 let mut relationships_moved: usize = 0;
253
254 for &src_id in &source_ids {
255 let moved_src = tx.execute(
257 "UPDATE OR IGNORE relationships SET source_id = ?1 WHERE source_id = ?2",
258 params![target_id, src_id],
259 )?;
260 tx.execute(
261 "DELETE FROM relationships WHERE source_id = ?1",
262 params![src_id],
263 )?;
264 let moved_tgt = tx.execute(
266 "UPDATE OR IGNORE relationships SET target_id = ?1 WHERE target_id = ?2",
267 params![target_id, src_id],
268 )?;
269 tx.execute(
270 "DELETE FROM relationships WHERE target_id = ?1",
271 params![src_id],
272 )?;
273 relationships_moved += moved_src + moved_tgt;
274 }
275
276 tx.execute("DELETE FROM relationships WHERE source_id = target_id", [])?;
278
279 tx.execute(
282 "DELETE FROM relationships
283 WHERE id NOT IN (
284 SELECT MIN(id)
285 FROM relationships
286 GROUP BY source_id, target_id, relation
287 )",
288 [],
289 )?;
290
291 for &src_id in &source_ids {
296 tx.execute(
297 "UPDATE OR IGNORE memory_entities SET entity_id = ?1 WHERE entity_id = ?2",
298 params![target_id, src_id],
299 )?;
300 tx.execute(
301 "DELETE FROM memory_entities WHERE entity_id = ?1",
302 params![src_id],
303 )?;
304 }
305
306 tx.execute(
308 "DELETE FROM memory_entities
309 WHERE rowid NOT IN (
310 SELECT MIN(rowid)
311 FROM memory_entities
312 GROUP BY memory_id, entity_id
313 )",
314 [],
315 )?;
316
317 let mut entities_removed: usize = 0;
320 for &src_id in &source_ids {
321 let removed = tx.execute("DELETE FROM entities WHERE id = ?1", params![src_id])?;
322 entities_removed += removed;
323 }
324
325 let adjacent_ids: Vec<i64> = {
327 let mut stmt = tx.prepare(
328 "SELECT DISTINCT CASE WHEN source_id = ?1 THEN target_id ELSE source_id END
329 FROM relationships WHERE source_id = ?1 OR target_id = ?1",
330 )?;
331 let ids: Vec<i64> = stmt
332 .query_map(params![target_id], |r| r.get(0))?
333 .collect::<Result<Vec<_>, _>>()?;
334 ids
335 };
336 entities::recalculate_degree(&tx, target_id)?;
337 for &adj_id in &adjacent_ids {
338 entities::recalculate_degree(&tx, adj_id)?;
339 }
340
341 tx.commit()?;
342
343 conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);")?;
344
345 let response = MergeEntitiesResponse {
346 action: "merged".to_string(),
347 sources: source_names,
348 target: target_name,
349 namespace: namespace.clone(),
350 target_id,
351 relationships_moved,
352 entities_removed,
353 elapsed_ms: inicio.elapsed().as_millis() as u64,
354 };
355
356 match args.format {
357 OutputFormat::Json => output::emit_json(&response)?,
358 OutputFormat::Text | OutputFormat::Markdown => {
359 output::emit_text(&format!(
360 "merged: {} sources into '{}' (relationships_moved={}, entities_removed={}) [{}]",
361 response.sources.len(),
362 response.target,
363 response.relationships_moved,
364 response.entities_removed,
365 response.namespace
366 ));
367 }
368 }
369
370 Ok(())
371}
372
373#[cfg(test)]
374mod tests {
375 use super::*;
376
377 #[test]
380 fn find_entity_name_by_id_disambiguates_homonyms_across_namespaces() {
381 let conn = rusqlite::Connection::open_in_memory().unwrap();
382 conn.execute_batch(
383 "CREATE TABLE entities (
384 id INTEGER PRIMARY KEY,
385 namespace TEXT NOT NULL,
386 name TEXT NOT NULL,
387 UNIQUE(namespace, name)
388 );",
389 )
390 .unwrap();
391 conn.execute(
392 "INSERT INTO entities (id, namespace, name)
393 VALUES (1, 'ns-a', 'auth'), (2, 'ns-b', 'auth')",
394 [],
395 )
396 .unwrap();
397
398 assert_eq!(
400 find_entity_name_by_id(&conn, "ns-a", 1, true).unwrap(),
401 ("auth".to_string(), "ns-a".to_string())
402 );
403 assert_eq!(
404 find_entity_name_by_id(&conn, "ns-b", 2, true).unwrap(),
405 ("auth".to_string(), "ns-b".to_string())
406 );
407 let err = find_entity_name_by_id(&conn, "ns-a", 2, true).unwrap_err();
408 assert_eq!(err.exit_code(), 4, "cross-namespace ID must be NotFound");
409 assert!(err.to_string().contains("id=2"), "obtido: {err}");
410
411 let (name, ns_actual) = find_entity_name_by_id(&conn, "ns-a", 2, false).unwrap();
414 assert_eq!(name, "auth");
415 assert_eq!(ns_actual, "ns-b");
416 }
417
418 #[test]
419 fn find_entity_name_by_id_missing_id_is_not_found() {
420 let conn = rusqlite::Connection::open_in_memory().unwrap();
421 conn.execute_batch(
422 "CREATE TABLE entities (
423 id INTEGER PRIMARY KEY,
424 namespace TEXT NOT NULL,
425 name TEXT NOT NULL
426 );",
427 )
428 .unwrap();
429 let err = find_entity_name_by_id(&conn, "global", 99, true).unwrap_err();
430 assert_eq!(err.exit_code(), 4);
431 }
432
433 #[derive(clap::Parser)]
436 struct TestCli {
437 #[command(flatten)]
438 args: MergeEntitiesArgs,
439 }
440
441 #[test]
442 fn clap_rejects_names_combined_with_ids() {
443 use clap::Parser;
444 let err =
445 match TestCli::try_parse_from(["t", "--names", "a,b", "--ids", "1,2", "--into", "tgt"])
446 {
447 Ok(_) => panic!("expected argument conflict"),
448 Err(e) => e,
449 };
450 assert_eq!(err.kind(), clap::error::ErrorKind::ArgumentConflict);
451 }
452
453 #[test]
454 fn clap_rejects_into_combined_with_into_id() {
455 use clap::Parser;
456 let err =
457 match TestCli::try_parse_from(["t", "--names", "a", "--into", "tgt", "--into-id", "3"])
458 {
459 Ok(_) => panic!("expected argument conflict"),
460 Err(e) => e,
461 };
462 assert_eq!(err.kind(), clap::error::ErrorKind::ArgumentConflict);
463 }
464
465 #[test]
469 fn self_referential_ids_rejected_before_db() {
470 let args = MergeEntitiesArgs {
471 names: vec![],
472 ids: vec![35575, 35340],
473 into: None,
474 into_id: Some(35340),
475 namespace: Some("global".into()),
476 format: OutputFormat::Json,
477 json: false,
478 db: Some("/nonexistent/no-db.sqlite".into()),
479 cross_namespace: true,
480 };
481 let err = run(args).expect_err("self-ref must fail");
482 assert_eq!(err.exit_code(), 1, "validation exit");
483 let msg = err.to_string();
484 assert!(
485 msg.contains("self-referential") || msg.contains("35340"),
486 "expected self-ref message, got: {msg}"
487 );
488 }
489
490 #[test]
491 fn clap_parses_ids_with_target_in_list() {
492 use clap::Parser;
493 let ok = TestCli::try_parse_from([
494 "t",
495 "--ids",
496 "35575,35340",
497 "--into-id",
498 "35340",
499 "--cross-namespace",
500 ])
501 .expect("must parse");
502 assert_eq!(ok.args.ids, vec![35575, 35340]);
503 assert_eq!(ok.args.into_id, Some(35340));
504 assert!(ok.args.cross_namespace);
505 }
506
507 #[test]
508 fn clap_requires_a_source_and_a_target_selector() {
509 use clap::Parser;
510 assert!(TestCli::try_parse_from(["t", "--into", "tgt"]).is_err());
511 assert!(TestCli::try_parse_from(["t", "--names", "a"]).is_err());
512 let ok = match TestCli::try_parse_from(["t", "--ids", "1,2", "--into-id", "3"]) {
513 Ok(cli) => cli,
514 Err(e) => panic!("expected successful parse: {e}"),
515 };
516 assert_eq!(ok.args.ids, vec![1, 2]);
517 assert_eq!(ok.args.into_id, Some(3));
518 assert!(ok.args.names.is_empty());
519 assert!(ok.args.into.is_none());
520 }
521
522 #[test]
523 fn merge_entities_response_serializes_all_fields() {
524 let resp = MergeEntitiesResponse {
525 action: "merged".to_string(),
526 sources: vec!["auth".to_string(), "authentication".to_string()],
527 target: "auth-service".to_string(),
528 namespace: "global".to_string(),
529 target_id: 1,
530 relationships_moved: 7,
531 entities_removed: 2,
532 elapsed_ms: 15,
533 };
534 let json = serde_json::to_value(&resp).expect("serialization failed");
535 assert_eq!(json["action"], "merged");
536 assert_eq!(json["target"], "auth-service");
537 assert_eq!(json["namespace"], "global");
538 assert_eq!(json["relationships_moved"], 7);
539 assert_eq!(json["entities_removed"], 2);
540 let sources = json["sources"].as_array().expect("must be array");
541 assert_eq!(sources.len(), 2);
542 assert!(json["elapsed_ms"].is_number());
543 }
544
545 #[test]
546 fn merge_entities_response_action_is_merged() {
547 let resp = MergeEntitiesResponse {
548 action: "merged".to_string(),
549 sources: vec!["src".to_string()],
550 target: "tgt".to_string(),
551 namespace: "ns".to_string(),
552 target_id: 1,
553 relationships_moved: 0,
554 entities_removed: 1,
555 elapsed_ms: 0,
556 };
557 assert_eq!(resp.action, "merged");
558 }
559
560 #[test]
561 fn merge_entities_response_empty_sources_serializes() {
562 let resp = MergeEntitiesResponse {
563 action: "merged".to_string(),
564 sources: vec![],
565 target: "target".to_string(),
566 namespace: "global".to_string(),
567 target_id: 1,
568 relationships_moved: 0,
569 entities_removed: 0,
570 elapsed_ms: 1,
571 };
572 let json = serde_json::to_value(&resp).expect("serialization failed");
573 let sources = json["sources"].as_array().expect("must be array");
574 assert_eq!(sources.len(), 0);
575 }
576
577 #[test]
578 fn merge_entities_response_with_zero_relationships_moved() {
579 let resp = MergeEntitiesResponse {
580 action: "merged".to_string(),
581 sources: vec!["src-a".to_string()],
582 target: "tgt".to_string(),
583 namespace: "global".to_string(),
584 target_id: 1,
585 relationships_moved: 0,
586 entities_removed: 1,
587 elapsed_ms: 5,
588 };
589 let json = serde_json::to_value(&resp).expect("serialization failed");
590 assert_eq!(json["relationships_moved"], 0);
591 assert_eq!(json["entities_removed"], 1);
592 }
593
594 #[test]
595 fn merge_entities_response_multiple_sources() {
596 let resp = MergeEntitiesResponse {
597 action: "merged".to_string(),
598 sources: vec!["a".into(), "b".into(), "c".into()],
599 target: "canonical".to_string(),
600 namespace: "proj".to_string(),
601 target_id: 1,
602 relationships_moved: 12,
603 entities_removed: 3,
604 elapsed_ms: 42,
605 };
606 let json = serde_json::to_value(&resp).expect("serialization failed");
607 assert_eq!(json["entities_removed"], 3);
608 let sources = json["sources"].as_array().unwrap();
609 assert_eq!(sources.len(), 3);
610 }
611
612 fn setup_migrated_db_on_disk() -> (tempfile::TempDir, rusqlite::Connection, std::path::PathBuf)
616 {
617 crate::storage::connection::register_vec_extension();
618 let tmp = tempfile::TempDir::new().expect("tempdir");
619 let db_path = tmp.path().join("test.db");
620 let mut conn = rusqlite::Connection::open(&db_path).expect("open");
621 crate::migrations::runner().run(&mut conn).expect("migrate");
622 (tmp, conn, db_path)
623 }
624
625 #[test]
629 fn cross_namespace_merges_source_from_other_namespace() {
630 let (_tmp, conn, db_path) = setup_migrated_db_on_disk();
631 conn.execute(
633 "INSERT INTO entities (namespace, name, type) VALUES ('global','tgt','concept')",
634 [],
635 )
636 .unwrap();
637 let tgt_id = conn.last_insert_rowid();
638 conn.execute(
640 "INSERT INTO entities (namespace, name, type) VALUES ('ai-sdd','dup-src','concept')",
641 [],
642 )
643 .unwrap();
644 let src_id = conn.last_insert_rowid();
645 conn.execute(
648 "INSERT INTO entities (namespace, name, type) VALUES ('global','third','concept')",
649 [],
650 )
651 .unwrap();
652 let third_id = conn.last_insert_rowid();
653 conn.execute(
654 "INSERT INTO relationships (namespace, source_id, target_id, relation, weight) \
655 VALUES ('ai-sdd', ?1, ?2, 'related', 0.5)",
656 params![src_id, third_id],
657 )
658 .unwrap();
659 drop(conn);
661
662 let args = MergeEntitiesArgs {
663 names: vec![],
664 ids: vec![src_id],
665 into: None,
666 into_id: Some(tgt_id),
667 namespace: Some("global".to_string()),
668 format: OutputFormat::Json,
669 json: false,
670 db: Some(db_path.to_string_lossy().into_owned()),
671 cross_namespace: true,
672 };
673 run(args).expect("cross-namespace merge must succeed");
674
675 let conn = rusqlite::Connection::open(&db_path).unwrap();
677 let src_remaining: i64 = conn
678 .query_row(
679 "SELECT COUNT(*) FROM entities WHERE id = ?1",
680 params![src_id],
681 |r| r.get(0),
682 )
683 .unwrap();
684 assert_eq!(src_remaining, 0, "cross-namespace source must be deleted");
685 let moved: i64 = conn
686 .query_row(
687 "SELECT COUNT(*) FROM relationships WHERE source_id = ?1 AND target_id = ?2",
688 params![tgt_id, third_id],
689 |r| r.get(0),
690 )
691 .unwrap();
692 assert!(
693 moved > 0,
694 "relationship must be retargeted to the target entity"
695 );
696 }
697
698 #[test]
702 fn cross_namespace_default_false_rejects_cross_id() {
703 let (_tmp, conn, db_path) = setup_migrated_db_on_disk();
704 conn.execute(
705 "INSERT INTO entities (namespace, name, type) VALUES ('global','tgt','concept')",
706 [],
707 )
708 .unwrap();
709 let tgt_id = conn.last_insert_rowid();
710 conn.execute(
711 "INSERT INTO entities (namespace, name, type) VALUES ('ai-sdd','dup-src','concept')",
712 [],
713 )
714 .unwrap();
715 let src_id = conn.last_insert_rowid();
716 drop(conn);
717
718 let args = MergeEntitiesArgs {
719 names: vec![],
720 ids: vec![src_id],
721 into: None,
722 into_id: Some(tgt_id),
723 namespace: Some("global".to_string()),
724 format: OutputFormat::Json,
725 json: false,
726 db: Some(db_path.to_string_lossy().into_owned()),
727 cross_namespace: false,
728 };
729 let err = run(args).expect_err("default must reject cross-namespace ID");
730 assert_eq!(err.exit_code(), 4, "cross-namespace ID must be NotFound");
731 }
732
733 #[test]
736 fn cross_namespace_target_must_still_be_in_resolved_namespace() {
737 let (_tmp, conn, db_path) = setup_migrated_db_on_disk();
738 conn.execute(
740 "INSERT INTO entities (namespace, name, type) VALUES ('ai-sdd','tgt','concept')",
741 [],
742 )
743 .unwrap();
744 let tgt_id = conn.last_insert_rowid();
745 conn.execute(
746 "INSERT INTO entities (namespace, name, type) VALUES ('global','src','concept')",
747 [],
748 )
749 .unwrap();
750 let src_id = conn.last_insert_rowid();
751 drop(conn);
752
753 let args = MergeEntitiesArgs {
754 names: vec![],
755 ids: vec![src_id],
756 into: None,
757 into_id: Some(tgt_id),
758 namespace: Some("global".to_string()),
759 format: OutputFormat::Json,
760 json: false,
761 db: Some(db_path.to_string_lossy().into_owned()),
762 cross_namespace: true,
763 };
764 let err = run(args).expect_err("target in wrong namespace must fail");
765 assert_eq!(
766 err.exit_code(),
767 4,
768 "target must still be NotFound in the resolved namespace"
769 );
770 }
771}