sqlite_graphrag/commands/
prune_relations.rs

1//! Handler for the `prune-relations` CLI subcommand.
2
3use crate::errors::AppError;
4use crate::i18n;
5use crate::output::{self, OutputFormat};
6use crate::paths::AppPaths;
7use crate::storage::connection::open_rw;
8use crate::storage::entities;
9use serde::Serialize;
10
11#[derive(clap::Args)]
12#[command(after_long_help = "EXAMPLES:\n  \
13    # Preview how many 'mentions' relations would be removed\n  \
14    sqlite-graphrag prune-relations --relation mentions --dry-run\n\n  \
15    # Remove all 'mentions' relations without confirmation prompt\n  \
16    sqlite-graphrag prune-relations --relation mentions --yes\n\n\
17NOTE:\n  \
18    This command permanently deletes relationships. Use --dry-run first.\n  \
19    Entity degree counts are automatically recalculated after pruning.")]
20pub struct PruneRelationsArgs {
21    /// Relation type to delete (e.g. mentions, related, uses).
22    /// Accepts canonical and custom kebab-case/snake_case values.
23    #[arg(long, value_parser = crate::parsers::parse_relation, value_name = "RELATION")]
24    pub relation: String,
25    #[arg(long)]
26    pub namespace: Option<String>,
27    /// Preview count without deleting.
28    #[arg(long)]
29    pub dry_run: bool,
30    /// Skip confirmation for destructive operation.
31    #[arg(long)]
32    pub yes: bool,
33    /// Show affected entity names during --dry-run preview.
34    #[arg(long, default_value_t = false)]
35    pub show_entities: bool,
36    #[arg(long, value_enum, default_value = "json")]
37    pub format: OutputFormat,
38    #[arg(long, hide = true, help = "No-op; JSON is always emitted on stdout")]
39    pub json: bool,
40    #[arg(long, env = "SQLITE_GRAPHRAG_DB_PATH")]
41    pub db: Option<String>,
42}
43
44#[derive(Serialize)]
45struct PruneRelationsResponse {
46    action: String,
47    relation: String,
48    count: usize,
49    entities_affected: usize,
50    namespace: String,
51    /// Total execution time in milliseconds from handler start to serialisation.
52    elapsed_ms: u64,
53    #[serde(skip_serializing_if = "Option::is_none")]
54    affected_entity_names: Option<Vec<String>>,
55}
56
57pub fn run(args: PruneRelationsArgs) -> Result<(), AppError> {
58    let inicio = std::time::Instant::now();
59    let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
60    let paths = AppPaths::resolve(args.db.as_deref())?;
61
62    crate::storage::connection::ensure_db_ready(&paths)?;
63
64    crate::parsers::warn_if_non_canonical(&args.relation);
65
66    let mut conn = open_rw(&paths.db)?;
67
68    if args.dry_run {
69        let count = entities::count_relationships_by_relation(&conn, &namespace, &args.relation)?;
70
71        let affected_names = if args.show_entities {
72            Some(entities::list_entity_names_by_relation(
73                &conn,
74                &namespace,
75                &args.relation,
76            )?)
77        } else {
78            None
79        };
80
81        let entities_affected_count = affected_names.as_ref().map_or(0, |v| v.len());
82
83        output::emit_progress(&i18n::prune_dry_run(count, &args.relation));
84
85        let response = PruneRelationsResponse {
86            action: "dry_run".to_string(),
87            relation: args.relation.clone(),
88            count,
89            entities_affected: entities_affected_count,
90            namespace: namespace.clone(),
91            elapsed_ms: inicio.elapsed().as_millis() as u64,
92            affected_entity_names: affected_names,
93        };
94
95        match args.format {
96            OutputFormat::Json => output::emit_json(&response)?,
97            OutputFormat::Text | OutputFormat::Markdown => {
98                output::emit_text(&format!(
99                    "dry_run: {} '{}' relations would be removed [{}]",
100                    response.count, response.relation, response.namespace
101                ));
102            }
103        }
104
105        return Ok(());
106    }
107
108    if !args.yes {
109        output::emit_progress(&i18n::prune_requires_yes());
110
111        let count = entities::count_relationships_by_relation(&conn, &namespace, &args.relation)?;
112
113        let response = PruneRelationsResponse {
114            action: "aborted".to_string(),
115            relation: args.relation.clone(),
116            count,
117            entities_affected: 0,
118            namespace: namespace.clone(),
119            elapsed_ms: inicio.elapsed().as_millis() as u64,
120            affected_entity_names: None,
121        };
122
123        match args.format {
124            OutputFormat::Json => output::emit_json(&response)?,
125            OutputFormat::Text | OutputFormat::Markdown => {
126                output::emit_text(&format!(
127                    "aborted: {} '{}' relations would be removed; pass --yes to confirm [{}]",
128                    response.count, response.relation, response.namespace
129                ));
130            }
131        }
132
133        return Ok(());
134    }
135
136    // Destructive path: delete relationships.
137    let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
138    let (count, entity_ids) =
139        entities::delete_relationships_by_relation(&tx, &namespace, &args.relation)?;
140    tx.commit()?;
141
142    // Run ANALYZE to refresh query planner statistics after bulk deletion.
143    conn.execute_batch("ANALYZE relationships; ANALYZE memory_relationships;")?;
144    conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);")?;
145
146    output::emit_progress(&i18n::relations_pruned(count, &args.relation, &namespace));
147
148    let response = PruneRelationsResponse {
149        action: "pruned".to_string(),
150        relation: args.relation.clone(),
151        count,
152        entities_affected: entity_ids.len(),
153        namespace: namespace.clone(),
154        elapsed_ms: inicio.elapsed().as_millis() as u64,
155        affected_entity_names: None,
156    };
157
158    match args.format {
159        OutputFormat::Json => output::emit_json(&response)?,
160        OutputFormat::Text | OutputFormat::Markdown => {
161            output::emit_text(&format!(
162                "pruned: {} '{}' relations removed, {} entities affected [{}]",
163                response.count, response.relation, response.entities_affected, response.namespace
164            ));
165        }
166    }
167
168    Ok(())
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174
175    #[test]
176    fn prune_response_serializes_all_fields() {
177        let resp = PruneRelationsResponse {
178            action: "pruned".to_string(),
179            relation: "mentions".to_string(),
180            count: 3451,
181            entities_affected: 200,
182            namespace: "global".to_string(),
183            elapsed_ms: 42,
184            affected_entity_names: None,
185        };
186        let json = serde_json::to_value(&resp).expect("serialization failed");
187        assert_eq!(json["action"], "pruned");
188        assert_eq!(json["relation"], "mentions");
189        assert_eq!(json["count"], 3451);
190        assert_eq!(json["entities_affected"], 200);
191        assert_eq!(json["namespace"], "global");
192        assert!(json["elapsed_ms"].is_number());
193    }
194
195    #[test]
196    fn prune_response_action_dry_run() {
197        let resp = PruneRelationsResponse {
198            action: "dry_run".to_string(),
199            relation: "mentions".to_string(),
200            count: 100,
201            entities_affected: 0,
202            namespace: "test".to_string(),
203            elapsed_ms: 5,
204            affected_entity_names: None,
205        };
206        let json = serde_json::to_value(&resp).expect("serialization failed");
207        assert_eq!(json["action"], "dry_run");
208        assert_eq!(
209            json["entities_affected"], 0,
210            "dry_run must report zero entities_affected"
211        );
212    }
213
214    #[test]
215    fn prune_response_action_pruned() {
216        let resp = PruneRelationsResponse {
217            action: "pruned".to_string(),
218            relation: "uses".to_string(),
219            count: 50,
220            entities_affected: 10,
221            namespace: "my-project".to_string(),
222            elapsed_ms: 120,
223            affected_entity_names: None,
224        };
225        let json = serde_json::to_value(&resp).expect("serialization failed");
226        assert_eq!(json["action"], "pruned");
227        assert!(json["count"].as_u64().unwrap() > 0);
228        assert!(json["entities_affected"].as_u64().unwrap() > 0);
229    }
230
231    #[test]
232    fn prune_response_zero_count_when_nothing_to_prune() {
233        let resp = PruneRelationsResponse {
234            action: "pruned".to_string(),
235            relation: "nonexistent".to_string(),
236            count: 0,
237            entities_affected: 0,
238            namespace: "global".to_string(),
239            elapsed_ms: 1,
240            affected_entity_names: None,
241        };
242        let json = serde_json::to_value(&resp).expect("serialization failed");
243        assert_eq!(json["count"], 0);
244        assert_eq!(json["entities_affected"], 0);
245    }
246
247    #[test]
248    fn prune_response_verbose_includes_entity_names() {
249        let resp = PruneRelationsResponse {
250            action: "dry_run".to_string(),
251            relation: "mentions".to_string(),
252            count: 10,
253            entities_affected: 3,
254            namespace: "global".to_string(),
255            elapsed_ms: 5,
256            affected_entity_names: Some(vec!["alpha".into(), "beta".into(), "gamma".into()]),
257        };
258        let json = serde_json::to_value(&resp).expect("serialization failed");
259        let names = json["affected_entity_names"]
260            .as_array()
261            .expect("must be array");
262        assert_eq!(names.len(), 3);
263    }
264
265    #[test]
266    fn prune_response_no_verbose_omits_entity_names() {
267        let resp = PruneRelationsResponse {
268            action: "dry_run".to_string(),
269            relation: "mentions".to_string(),
270            count: 10,
271            entities_affected: 0,
272            namespace: "global".to_string(),
273            elapsed_ms: 5,
274            affected_entity_names: None,
275        };
276        let json = serde_json::to_value(&resp).expect("serialization failed");
277        assert!(
278            json.get("affected_entity_names").is_none(),
279            "must be omitted when None"
280        );
281    }
282
283    #[test]
284    fn prune_response_action_values_are_exhaustive() {
285        for action in &["pruned", "dry_run", "aborted"] {
286            let resp = PruneRelationsResponse {
287                action: action.to_string(),
288                relation: "mentions".to_string(),
289                count: 0,
290                entities_affected: 0,
291                namespace: "global".to_string(),
292                elapsed_ms: 0,
293                affected_entity_names: None,
294            };
295            let json = serde_json::to_value(&resp).expect("serialization");
296            assert_eq!(json["action"], *action);
297        }
298    }
299}
sqlite_graphrag/commands/prune_relations.rs

sqlite_graphrag/commands/
prune_relations.rs
