Skip to main content

sqlite_graphrag/commands/
link.rs

1//! Handler for the `link` CLI subcommand.
2
3use crate::cli::RelationKind;
4use crate::constants::DEFAULT_RELATION_WEIGHT;
5use crate::errors::AppError;
6use crate::i18n::{errors_msg, validation};
7use crate::output::{self, OutputFormat};
8use crate::paths::AppPaths;
9use crate::storage::connection::open_rw;
10use crate::storage::entities;
11use serde::Serialize;
12
13#[derive(clap::Args)]
14pub struct LinkArgs {
15    /// Source entity.
16    #[arg(long)]
17    pub from: String,
18    /// Target entity.
19    #[arg(long)]
20    pub to: String,
21    #[arg(long, value_enum)]
22    pub relation: RelationKind,
23    #[arg(long)]
24    pub weight: Option<f64>,
25    #[arg(long)]
26    pub namespace: Option<String>,
27    #[arg(long, value_enum, default_value = "json")]
28    pub format: OutputFormat,
29    #[arg(long, help = "No-op; JSON is always emitted on stdout")]
30    pub json: bool,
31    #[arg(long, env = "SQLITE_GRAPHRAG_DB_PATH")]
32    pub db: Option<String>,
33}
34
35#[derive(Serialize)]
36struct LinkResponse {
37    action: String,
38    from: String,
39    to: String,
40    relation: String,
41    weight: f64,
42    namespace: String,
43    /// Total execution time in milliseconds from handler start to serialisation.
44    elapsed_ms: u64,
45}
46
47pub fn run(args: LinkArgs) -> Result<(), AppError> {
48    let inicio = std::time::Instant::now();
49    let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
50    let paths = AppPaths::resolve(args.db.as_deref())?;
51
52    if args.from == args.to {
53        return Err(AppError::Validation(validation::self_referential_link()));
54    }
55
56    let weight = args.weight.unwrap_or(DEFAULT_RELATION_WEIGHT);
57    if !(0.0..=1.0).contains(&weight) {
58        return Err(AppError::Validation(validation::invalid_link_weight(
59            weight,
60        )));
61    }
62
63    if !paths.db.exists() {
64        return Err(AppError::NotFound(errors_msg::database_not_found(
65            &paths.db.display().to_string(),
66        )));
67    }
68
69    let relation_str = args.relation.as_str();
70
71    let mut conn = open_rw(&paths.db)?;
72
73    let source_id = entities::find_entity_id(&conn, &namespace, &args.from)?
74        .ok_or_else(|| AppError::NotFound(errors_msg::entity_not_found(&args.from, &namespace)))?;
75    let target_id = entities::find_entity_id(&conn, &namespace, &args.to)?
76        .ok_or_else(|| AppError::NotFound(errors_msg::entity_not_found(&args.to, &namespace)))?;
77
78    let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
79    let (_rel_id, was_created) = entities::create_or_fetch_relationship(
80        &tx,
81        &namespace,
82        source_id,
83        target_id,
84        relation_str,
85        weight,
86        None,
87    )?;
88
89    if was_created {
90        entities::recalculate_degree(&tx, source_id)?;
91        entities::recalculate_degree(&tx, target_id)?;
92    }
93    tx.commit()?;
94
95    let action = if was_created {
96        "created".to_string()
97    } else {
98        "already_exists".to_string()
99    };
100
101    let response = LinkResponse {
102        action: action.clone(),
103        from: args.from.clone(),
104        to: args.to.clone(),
105        relation: relation_str.to_string(),
106        weight,
107        namespace: namespace.clone(),
108        elapsed_ms: inicio.elapsed().as_millis() as u64,
109    };
110
111    match args.format {
112        OutputFormat::Json => output::emit_json(&response)?,
113        OutputFormat::Text | OutputFormat::Markdown => {
114            output::emit_text(&format!(
115                "{}: {} --[{}]--> {} [{}]",
116                action, response.from, response.relation, response.to, response.namespace
117            ));
118        }
119    }
120
121    Ok(())
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127
128    #[test]
129    fn link_response_sem_aliases_redundantes() {
130        // P1-O: campos source/target foram removidos do JSON de resposta.
131        let resp = LinkResponse {
132            action: "created".to_string(),
133            from: "entidade-a".to_string(),
134            to: "entidade-b".to_string(),
135            relation: "uses".to_string(),
136            weight: 1.0,
137            namespace: "default".to_string(),
138            elapsed_ms: 0,
139        };
140        let json = serde_json::to_value(&resp).expect("serialização deve funcionar");
141        assert_eq!(json["from"], "entidade-a");
142        assert_eq!(json["to"], "entidade-b");
143        assert!(
144            json.get("source").is_none(),
145            "campo 'source' foi removido em P1-O"
146        );
147        assert!(
148            json.get("target").is_none(),
149            "campo 'target' foi removido em P1-O"
150        );
151    }
152
153    #[test]
154    fn link_response_serializa_todos_campos() {
155        let resp = LinkResponse {
156            action: "already_exists".to_string(),
157            from: "origem".to_string(),
158            to: "destino".to_string(),
159            relation: "mentions".to_string(),
160            weight: 0.8,
161            namespace: "teste".to_string(),
162            elapsed_ms: 5,
163        };
164        let json = serde_json::to_value(&resp).expect("serialização deve funcionar");
165        assert!(json.get("action").is_some());
166        assert!(json.get("from").is_some());
167        assert!(json.get("to").is_some());
168        assert!(json.get("relation").is_some());
169        assert!(json.get("weight").is_some());
170        assert!(json.get("namespace").is_some());
171        assert!(json.get("elapsed_ms").is_some());
172    }
173}