sqlite_graphrag/commands/
normalize_entities.rs

1//! Handler for the `normalize-entities` CLI subcommand (GAP-15).
2//!
3//! Scans all existing entity names in the namespace and normalizes them to
4//! kebab-case ASCII using [`crate::parsers::normalize_entity_name`].
5//! When a normalized name already exists (collision), the source entity is
6//! merged into the target using the same logic as `merge-entities`:
7//! relationships are retargeted via `UPDATE OR IGNORE` + `DELETE`, then
8//! the source row is removed. Otherwise the entity name is updated in place.
9
10use crate::errors::AppError;
11use crate::output::{self, OutputFormat};
12use crate::parsers::normalize_entity_name;
13use crate::paths::AppPaths;
14use crate::storage::connection::open_rw;
15use rusqlite::params;
16use serde::Serialize;
17
18#[derive(clap::Args)]
19#[command(after_long_help = "EXAMPLES:\n  \
20    # Preview which entities would be renamed or merged\n  \
21    sqlite-graphrag normalize-entities --dry-run\n\n  \
22    # Apply normalization to all entity names\n  \
23    sqlite-graphrag normalize-entities --yes\n\n  \
24    # Scope to a specific namespace\n  \
25    sqlite-graphrag normalize-entities --yes --namespace my-project\n\n\
26NOTE:\n  \
27    When a normalized name already exists, the source entity is merged into\n  \
28    the existing target via relationship retargeting (UPDATE OR IGNORE + DELETE).\n  \
29    Run `cleanup-orphans` afterwards to remove any newly orphaned entities.")]
30/// Normalize entities args.
31pub struct NormalizeEntitiesArgs {
32    /// Preview changes without persisting them.
33    #[arg(long, conflicts_with = "yes")]
34    pub dry_run: bool,
35    /// Apply normalization without interactive confirmation.
36    #[arg(long, conflicts_with = "dry_run")]
37    pub yes: bool,
38    /// Namespace scope.
39    #[arg(long)]
40    pub namespace: Option<String>,
41    /// Output format.
42    #[arg(long, value_enum, default_value = "json")]
43    pub format: OutputFormat,
44    /// Emit machine-readable JSON on stdout.
45    #[arg(long, hide = true, help = "No-op; JSON is always emitted on stdout")]
46    pub json: bool,
47    /// Path to the SQLite database file.
48    #[arg(long)]
49    pub db: Option<String>,
50}
51
52#[derive(Serialize)]
53struct NormalizeEntitiesResponse {
54    /// "normalized" when changes were applied, "dry_run" when only previewed.
55    action: String,
56    /// Number of entities whose names were updated in place.
57    normalized_count: usize,
58    /// Number of entities that collided with an existing normalized name and
59    /// were merged into the target.
60    merged_count: usize,
61    namespace: String,
62    /// Total execution time in milliseconds from handler start to serialisation.
63    elapsed_ms: u64,
64}
65
66/// Run.
67pub fn run(args: NormalizeEntitiesArgs) -> Result<(), AppError> {
68    let inicio = std::time::Instant::now();
69
70    if !args.dry_run && !args.yes {
71        return Err(AppError::Validation(
72            "pass --dry-run to preview or --yes to apply changes".to_string(),
73        ));
74    }
75
76    let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
77    let paths = AppPaths::resolve(args.db.as_deref())?;
78
79    crate::storage::connection::ensure_db_ready(&paths)?;
80
81    let mut conn = open_rw(&paths.db)?;
82
83    // Collect all entity (id, name) pairs for the namespace.
84    let entities: Vec<(i64, String)> = {
85        let mut stmt =
86            conn.prepare_cached("SELECT id, name FROM entities WHERE namespace = ?1 ORDER BY id")?;
87        let rows = stmt.query_map(params![namespace], |r| {
88            Ok((r.get::<_, i64>(0)?, r.get::<_, String>(1)?))
89        })?;
90        rows.collect::<Result<Vec<_>, _>>()?
91    };
92
93    // Compute which names need changing.
94    let to_change: Vec<(i64, String, String)> = entities
95        .iter()
96        .filter_map(|(id, name)| {
97            let normalized = normalize_entity_name(name);
98            if normalized != *name {
99                Some((*id, name.clone(), normalized))
100            } else {
101                None
102            }
103        })
104        .collect();
105
106    // G10: classify changes into renames (no collision) and merges (collision).
107    // A collision occurs when two distinct names normalize to the same target,
108    // or when the normalized target already exists in the DB as an already-normalized entity.
109    let already_normalized: std::collections::HashSet<String> = entities
110        .iter()
111        .filter(|(_, name)| normalize_entity_name(name) == *name)
112        .map(|(_, name)| name.clone())
113        .collect();
114
115    let mut target_groups: std::collections::HashMap<String, usize> =
116        std::collections::HashMap::with_capacity(to_change.len());
117    for (_, _, normalized) in &to_change {
118        *target_groups.entry(normalized.clone()).or_insert(0) += 1;
119    }
120
121    let mut merge_count_preview: usize = 0;
122    let mut rename_count_preview: usize = 0;
123    for (target, count) in &target_groups {
124        if *count > 1 || already_normalized.contains(target) {
125            // All sources in this group will merge into the existing or first entity
126            let extra = if already_normalized.contains(target) {
127                *count // all merge into existing
128            } else {
129                count - 1 // first one renames, rest merge
130            };
131            merge_count_preview += extra;
132            rename_count_preview += count - extra;
133        } else {
134            rename_count_preview += 1;
135        }
136    }
137
138    if args.dry_run {
139        let response = NormalizeEntitiesResponse {
140            action: "dry_run".to_string(),
141            normalized_count: rename_count_preview,
142            merged_count: merge_count_preview,
143            namespace,
144            elapsed_ms: inicio.elapsed().as_millis() as u64,
145        };
146        match args.format {
147            OutputFormat::Json => output::emit_json(&response)?,
148            OutputFormat::Text | OutputFormat::Markdown => {
149                output::emit_text(&format!(
150                    "dry_run: {} entity names would be normalized",
151                    response.normalized_count
152                ));
153            }
154        }
155        return Ok(());
156    }
157
158    // Apply changes inside a transaction.
159    let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
160
161    let mut normalized_count: usize = 0;
162    let mut merged_count: usize = 0;
163
164    for (src_id, _original_name, normalized) in &to_change {
165        // Check whether a row with the normalized name already exists.
166        let existing_id: Option<i64> = {
167            let mut stmt =
168                tx.prepare_cached("SELECT id FROM entities WHERE namespace = ?1 AND name = ?2")?;
169            match stmt.query_row(params![namespace, normalized], |r| r.get::<_, i64>(0)) {
170                Ok(id) => Some(id),
171                Err(rusqlite::Error::QueryReturnedNoRows) => None,
172                Err(e) => return Err(AppError::Database(e)),
173            }
174        };
175
176        match existing_id {
177            Some(target_id) if target_id != *src_id => {
178                // Collision: merge source into target using UPDATE OR IGNORE + DELETE.
179                // Step 1a: redirect source_id.
180                tx.execute(
181                    "UPDATE OR IGNORE relationships SET source_id = ?1 WHERE source_id = ?2",
182                    params![target_id, src_id],
183                )?;
184                tx.execute(
185                    "DELETE FROM relationships WHERE source_id = ?1",
186                    params![src_id],
187                )?;
188                // Step 1b: redirect target_id.
189                tx.execute(
190                    "UPDATE OR IGNORE relationships SET target_id = ?1 WHERE target_id = ?2",
191                    params![target_id, src_id],
192                )?;
193                tx.execute(
194                    "DELETE FROM relationships WHERE target_id = ?1",
195                    params![src_id],
196                )?;
197                // Remove self-loops.
198                tx.execute("DELETE FROM relationships WHERE source_id = target_id", [])?;
199                // Retarget memory_entities bindings.
200                tx.execute(
201                    "UPDATE OR IGNORE memory_entities SET entity_id = ?1 WHERE entity_id = ?2",
202                    params![target_id, src_id],
203                )?;
204                tx.execute(
205                    "DELETE FROM memory_entities WHERE entity_id = ?1",
206                    params![src_id],
207                )?;
208                // Remove the source entity row.
209                tx.execute("DELETE FROM entities WHERE id = ?1", params![src_id])?;
210                // Recalculate degree for the surviving target.
211                tx.execute(
212                    "UPDATE entities
213                     SET degree = (SELECT COUNT(*) FROM relationships
214                                   WHERE source_id = entities.id OR target_id = entities.id)
215                     WHERE id = ?1",
216                    params![target_id],
217                )?;
218                tracing::info!(target: "normalize_entities",
219                    src_id = src_id,
220                    target_id = target_id,
221                    normalized = normalized,
222                    "entity merged into existing normalized target"
223                );
224                merged_count += 1;
225            }
226            _ => {
227                // No collision: simple rename.
228                tx.execute(
229                    "UPDATE entities SET name = ?1, updated_at = unixepoch() WHERE id = ?2",
230                    params![normalized, src_id],
231                )?;
232                tracing::info!(target: "normalize_entities",
233                    entity_id = src_id,
234                    normalized = normalized,
235                    "entity name normalized"
236                );
237                normalized_count += 1;
238            }
239        }
240    }
241
242    tx.commit()?;
243    conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);")?;
244
245    let response = NormalizeEntitiesResponse {
246        action: "normalized".to_string(),
247        normalized_count,
248        merged_count,
249        namespace,
250        elapsed_ms: inicio.elapsed().as_millis() as u64,
251    };
252
253    match args.format {
254        OutputFormat::Json => output::emit_json(&response)?,
255        OutputFormat::Text | OutputFormat::Markdown => {
256            output::emit_text(&format!(
257                "normalized: {} renamed, {} merged",
258                response.normalized_count, response.merged_count
259            ));
260        }
261    }
262
263    Ok(())
264}
265
266#[cfg(test)]
267mod tests {
268    use super::*;
269    use crate::storage::connection::register_vec_extension;
270    use rusqlite::Connection;
271    use tempfile::TempDir;
272
273    type TestResult = Result<(), Box<dyn std::error::Error>>;
274
275    /// Opens a temp DB with the full schema applied via migrations.
276    fn setup_db() -> Result<(TempDir, Connection), Box<dyn std::error::Error>> {
277        register_vec_extension();
278        let tmp = TempDir::new()?;
279        let db_path = tmp.path().join("test.db");
280        let mut conn = Connection::open(&db_path)?;
281        crate::migrations::runner().run(&mut conn)?;
282        Ok((tmp, conn))
283    }
284
285    /// Inserts an entity bypassing `upsert_entity` normalization, so tests can
286    /// seed deliberately un-normalized names.
287    fn insert_entity(conn: &Connection, name: &str) -> Result<i64, Box<dyn std::error::Error>> {
288        // Bypass upsert_entity normalization to seed raw (un-normalized) names.
289        conn.execute(
290            "INSERT INTO entities (namespace, name, type, description) VALUES ('global', ?1, 'concept', NULL)",
291            params![name],
292        )?;
293        let id: i64 = conn.query_row(
294            "SELECT id FROM entities WHERE namespace = 'global' AND name = ?1",
295            params![name],
296            |r| r.get(0),
297        )?;
298        Ok(id)
299    }
300
301    #[test]
302    fn dry_run_returns_count_without_changes() -> TestResult {
303        let (_tmp, conn) = setup_db()?;
304        insert_entity(&conn, "Hello World")?;
305        insert_entity(&conn, "already-normalized")?;
306
307        // Verify "Hello World" exists.
308        let count: i64 = conn.query_row(
309            "SELECT COUNT(*) FROM entities WHERE name = 'Hello World' AND namespace = 'global'",
310            [],
311            |r| r.get(0),
312        )?;
313        assert_eq!(count, 1, "entity must exist before dry run");
314
315        // dry_run must not modify anything.
316        let count_after: i64 = conn.query_row(
317            "SELECT COUNT(*) FROM entities WHERE name = 'Hello World' AND namespace = 'global'",
318            [],
319            |r| r.get(0),
320        )?;
321        assert_eq!(count_after, 1, "dry run must not rename entities");
322        Ok(())
323    }
324
325    #[test]
326    fn renames_unnormalized_entity_in_place() -> TestResult {
327        let (_tmp, conn) = setup_db()?;
328        let src_id = insert_entity(&conn, "Hello World")?;
329
330        // Apply normalization directly via the internal logic.
331        {
332            let normalized = normalize_entity_name("Hello World");
333            let existing: Option<i64> = {
334                match conn.query_row(
335                    "SELECT id FROM entities WHERE namespace = 'global' AND name = ?1",
336                    params![normalized],
337                    |r| r.get::<_, i64>(0),
338                ) {
339                    Ok(id) => Some(id),
340                    Err(rusqlite::Error::QueryReturnedNoRows) => None,
341                    Err(e) => return Err(e.into()),
342                }
343            };
344            assert!(existing.is_none(), "no collision expected");
345            conn.execute(
346                "UPDATE entities SET name = ?1 WHERE id = ?2",
347                params![normalized, src_id],
348            )?;
349        }
350
351        let name: String = conn.query_row(
352            "SELECT name FROM entities WHERE id = ?1",
353            params![src_id],
354            |r| r.get(0),
355        )?;
356        assert_eq!(name, "hello-world");
357        Ok(())
358    }
359
360    #[test]
361    fn merges_into_existing_on_collision() -> TestResult {
362        let (_tmp, conn) = setup_db()?;
363        // Target already exists with the normalized name.
364        let target_id = insert_entity(&conn, "hello-world")?;
365        // Source has the un-normalized form that normalizes to the same value.
366        let src_id = insert_entity(&conn, "Hello World")?;
367
368        // Insert a relationship attached to src_id.
369        conn.execute(
370            "INSERT INTO relationships (namespace, source_id, target_id, relation, weight)
371             VALUES ('global', ?1, ?1, 'related', 0.5)",
372            params![src_id],
373        )?;
374
375        // Merge: retarget relationships from src → target.
376        conn.execute(
377            "UPDATE OR IGNORE relationships SET source_id = ?1 WHERE source_id = ?2",
378            params![target_id, src_id],
379        )?;
380        conn.execute(
381            "DELETE FROM relationships WHERE source_id = ?1",
382            params![src_id],
383        )?;
384        conn.execute("DELETE FROM entities WHERE id = ?1", params![src_id])?;
385
386        // Source must be gone.
387        let src_exists: i64 = conn.query_row(
388            "SELECT COUNT(*) FROM entities WHERE id = ?1",
389            params![src_id],
390            |r| r.get(0),
391        )?;
392        assert_eq!(src_exists, 0, "source entity must be deleted after merge");
393
394        // Target must still exist.
395        let target_name: String = conn.query_row(
396            "SELECT name FROM entities WHERE id = ?1",
397            params![target_id],
398            |r| r.get(0),
399        )?;
400        assert_eq!(target_name, "hello-world");
401        Ok(())
402    }
403
404    #[test]
405    fn normalize_entities_response_serializes_correctly() {
406        let resp = NormalizeEntitiesResponse {
407            action: "normalized".to_string(),
408            normalized_count: 3,
409            merged_count: 1,
410            namespace: "global".to_string(),
411            elapsed_ms: 42,
412        };
413        let json = serde_json::to_value(&resp).expect("serialization");
414        assert_eq!(json["action"], "normalized");
415        assert_eq!(json["normalized_count"], 3);
416        assert_eq!(json["merged_count"], 1);
417        assert_eq!(json["namespace"], "global");
418        assert!(json["elapsed_ms"].as_u64().is_some());
419    }
420
421    #[test]
422    fn dry_run_response_has_correct_action() {
423        let resp = NormalizeEntitiesResponse {
424            action: "dry_run".to_string(),
425            normalized_count: 5,
426            merged_count: 0,
427            namespace: "test".to_string(),
428            elapsed_ms: 1,
429        };
430        let json = serde_json::to_value(&resp).expect("serialization");
431        assert_eq!(json["action"], "dry_run");
432    }
433}
sqlite_graphrag/commands/normalize_entities.rs

sqlite_graphrag/commands/
normalize_entities.rs
