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

sqlite_graphrag/commands/
prune_relations.rs
