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
145    output::emit_progress(&i18n::relations_pruned(count, &args.relation, &namespace));
146
147    let response = PruneRelationsResponse {
148        action: "pruned".to_string(),
149        relation: args.relation.clone(),
150        count,
151        entities_affected: entity_ids.len(),
152        namespace: namespace.clone(),
153        elapsed_ms: inicio.elapsed().as_millis() as u64,
154        affected_entity_names: None,
155    };
156
157    match args.format {
158        OutputFormat::Json => output::emit_json(&response)?,
159        OutputFormat::Text | OutputFormat::Markdown => {
160            output::emit_text(&format!(
161                "pruned: {} '{}' relations removed, {} entities affected [{}]",
162                response.count, response.relation, response.entities_affected, response.namespace
163            ));
164        }
165    }
166
167    Ok(())
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    #[test]
175    fn prune_response_serializes_all_fields() {
176        let resp = PruneRelationsResponse {
177            action: "pruned".to_string(),
178            relation: "mentions".to_string(),
179            count: 3451,
180            entities_affected: 200,
181            namespace: "global".to_string(),
182            elapsed_ms: 42,
183            affected_entity_names: None,
184        };
185        let json = serde_json::to_value(&resp).expect("serialization failed");
186        assert_eq!(json["action"], "pruned");
187        assert_eq!(json["relation"], "mentions");
188        assert_eq!(json["count"], 3451);
189        assert_eq!(json["entities_affected"], 200);
190        assert_eq!(json["namespace"], "global");
191        assert!(json["elapsed_ms"].is_number());
192    }
193
194    #[test]
195    fn prune_response_action_dry_run() {
196        let resp = PruneRelationsResponse {
197            action: "dry_run".to_string(),
198            relation: "mentions".to_string(),
199            count: 100,
200            entities_affected: 0,
201            namespace: "test".to_string(),
202            elapsed_ms: 5,
203            affected_entity_names: None,
204        };
205        let json = serde_json::to_value(&resp).expect("serialization failed");
206        assert_eq!(json["action"], "dry_run");
207        assert_eq!(
208            json["entities_affected"], 0,
209            "dry_run must report zero entities_affected"
210        );
211    }
212
213    #[test]
214    fn prune_response_action_pruned() {
215        let resp = PruneRelationsResponse {
216            action: "pruned".to_string(),
217            relation: "uses".to_string(),
218            count: 50,
219            entities_affected: 10,
220            namespace: "my-project".to_string(),
221            elapsed_ms: 120,
222            affected_entity_names: None,
223        };
224        let json = serde_json::to_value(&resp).expect("serialization failed");
225        assert_eq!(json["action"], "pruned");
226        assert!(json["count"].as_u64().unwrap() > 0);
227        assert!(json["entities_affected"].as_u64().unwrap() > 0);
228    }
229
230    #[test]
231    fn prune_response_zero_count_when_nothing_to_prune() {
232        let resp = PruneRelationsResponse {
233            action: "pruned".to_string(),
234            relation: "nonexistent".to_string(),
235            count: 0,
236            entities_affected: 0,
237            namespace: "global".to_string(),
238            elapsed_ms: 1,
239            affected_entity_names: None,
240        };
241        let json = serde_json::to_value(&resp).expect("serialization failed");
242        assert_eq!(json["count"], 0);
243        assert_eq!(json["entities_affected"], 0);
244    }
245
246    #[test]
247    fn prune_response_verbose_includes_entity_names() {
248        let resp = PruneRelationsResponse {
249            action: "dry_run".to_string(),
250            relation: "mentions".to_string(),
251            count: 10,
252            entities_affected: 3,
253            namespace: "global".to_string(),
254            elapsed_ms: 5,
255            affected_entity_names: Some(vec!["alpha".into(), "beta".into(), "gamma".into()]),
256        };
257        let json = serde_json::to_value(&resp).expect("serialization failed");
258        let names = json["affected_entity_names"]
259            .as_array()
260            .expect("must be array");
261        assert_eq!(names.len(), 3);
262    }
263
264    #[test]
265    fn prune_response_no_verbose_omits_entity_names() {
266        let resp = PruneRelationsResponse {
267            action: "dry_run".to_string(),
268            relation: "mentions".to_string(),
269            count: 10,
270            entities_affected: 0,
271            namespace: "global".to_string(),
272            elapsed_ms: 5,
273            affected_entity_names: None,
274        };
275        let json = serde_json::to_value(&resp).expect("serialization failed");
276        assert!(
277            json.get("affected_entity_names").is_none(),
278            "must be omitted when None"
279        );
280    }
281}
sqlite_graphrag/commands/prune_relations.rs

sqlite_graphrag/commands/
prune_relations.rs
