helix_graph_algorithms/
loader.rs

1//! Validation boundary for graph-query response bytes.
2//!
3//! SDKs own query construction. This module owns the shared response aliases
4//! and turns the resulting bytes into a graph only after every row and
5//! topology invariant has been validated.
6
7use std::collections::{BTreeMap, BTreeSet};
8
9use serde::Deserialize;
10use serde_json::Value;
11use thiserror::Error;
12
13use crate::{Edge, ExternalId, Graph, GraphError, GraphKind, IdentitySelection, Node};
14
15/// Prefix reserved for graph loader projection aliases.
16pub const PRIVATE_PREFIX: &str = "__helix_graph_";
17/// Projected internal node ID alias.
18pub const NODE_ID: &str = "__helix_graph_node_id";
19/// Projected public external node ID alias.
20pub const EXTERNAL_ID: &str = "__helix_graph_external_id";
21/// Projected node label alias.
22pub const NODE_LABEL: &str = "__helix_graph_node_label";
23/// Projected stable edge ID alias.
24pub const EDGE_ID: &str = "__helix_graph_edge_id";
25/// Projected Graphify multigraph key alias.
26pub const EDGE_KEY: &str = "__helix_graph_edge_key";
27/// Projected internal source-node ID alias.
28pub const EDGE_SOURCE: &str = "__helix_graph_edge_source";
29/// Projected internal target-node ID alias.
30pub const EDGE_TARGET: &str = "__helix_graph_edge_target";
31/// Projected edge label alias.
32pub const EDGE_LABEL: &str = "__helix_graph_edge_label";
33/// Projected algorithm weight alias.
34pub const EDGE_WEIGHT: &str = "__helix_graph_edge_weight";
35
36/// Metadata required to interpret one graph query response.
37#[derive(Debug, Clone, PartialEq, Eq)]
38pub struct GraphLoadSpec {
39    /// Declared topology contract of the constructed graph.
40    pub kind: GraphKind,
41    /// How projected node identities are decoded.
42    pub node_identity: IdentitySelection,
43    /// How an optional projected Graphify edge key is decoded.
44    pub edge_key_identity: Option<IdentitySelection>,
45    /// Maximum complete node result size.
46    pub node_limit: Option<usize>,
47    /// Maximum complete edge result size.
48    pub edge_limit: Option<usize>,
49}
50
51/// Failure to validate raw graph response bytes.
52#[derive(Debug, Error)]
53pub enum GraphLoadError {
54    /// The response is not the expected two-array result object.
55    #[error("invalid graph query response: {0}")]
56    InvalidResponse(String),
57    /// A projected row is missing or contains an invalid field.
58    #[error("invalid {kind} row {index}: {details}")]
59    InvalidRow {
60        /// `node` or `edge`.
61        kind: &'static str,
62        /// Zero-based row index.
63        index: usize,
64        /// Validation details.
65        details: String,
66    },
67    /// More rows existed than the explicit safety limit permits.
68    #[error("graph selection exceeded the {kind} safety limit of {limit}")]
69    IncompleteSelection {
70        /// `node` or `edge`.
71        kind: &'static str,
72        /// Configured maximum complete result size.
73        limit: usize,
74    },
75    /// Two internal nodes map to the same external identity.
76    #[error("duplicate external node identity: {0}")]
77    DuplicateExternalIdentity(ExternalId),
78    /// Graph metadata selection returned more than one row.
79    #[error("graph metadata selection returned {count} rows; expected at most one")]
80    MultipleGraphMetadataRows {
81        /// Number of returned rows.
82        count: usize,
83    },
84    /// Graph construction found a topology invariant violation.
85    #[error(transparent)]
86    Graph(#[from] GraphError),
87}
88
89#[derive(Debug, Deserialize)]
90struct GraphResponse {
91    nodes: Vec<BTreeMap<String, Value>>,
92    edges: Vec<BTreeMap<String, Value>>,
93    #[serde(default)]
94    metadata: Vec<BTreeMap<String, Value>>,
95}
96
97/// Validate raw `/v2/query` response bytes and construct an immutable graph.
98pub fn graph_from_response(spec: GraphLoadSpec, response: &[u8]) -> Result<Graph, GraphLoadError> {
99    let response: GraphResponse = serde_json::from_slice(response)
100        .map_err(|error| GraphLoadError::InvalidResponse(error.to_string()))?;
101    if let Some(limit) = spec.node_limit
102        && response.nodes.len() > limit
103    {
104        return Err(GraphLoadError::IncompleteSelection {
105            kind: "node",
106            limit,
107        });
108    }
109    if let Some(limit) = spec.edge_limit
110        && response.edges.len() > limit
111    {
112        return Err(GraphLoadError::IncompleteSelection {
113            kind: "edge",
114            limit,
115        });
116    }
117    let graph_attributes = match response.metadata.len() {
118        0 => BTreeMap::new(),
119        1 => {
120            let attributes = response
121                .metadata
122                .into_iter()
123                .next()
124                .expect("metadata length is one");
125            if let Some(alias) = attributes
126                .keys()
127                .find(|alias| alias.starts_with(PRIVATE_PREFIX))
128            {
129                return Err(GraphLoadError::InvalidRow {
130                    kind: "metadata",
131                    index: 0,
132                    details: format!("reserved projection alias {alias}"),
133                });
134            }
135            attributes
136        }
137        count => return Err(GraphLoadError::MultipleGraphMetadataRows { count }),
138    };
139
140    let mut internal_to_external = BTreeMap::new();
141    let mut external_ids = BTreeSet::new();
142    let nodes = response
143        .nodes
144        .into_iter()
145        .enumerate()
146        .map(|(index, mut row)| {
147            let internal = take_identity(&mut row, NODE_ID, "node", index)?;
148            let external =
149                take_external_id(&mut row, EXTERNAL_ID, "node", index, &spec.node_identity)?;
150            if internal_to_external
151                .insert(internal.clone(), external.clone())
152                .is_some()
153            {
154                return Err(invalid_row(
155                    "node",
156                    index,
157                    format!("duplicate internal ID {internal}"),
158                ));
159            }
160            if !external_ids.insert(external.clone()) {
161                return Err(GraphLoadError::DuplicateExternalIdentity(external));
162            }
163            let label = take_optional_string(&mut row, NODE_LABEL, "node", index)?;
164            let mut node = Node::new(external).with_attributes(row);
165            if let Some(label) = label {
166                node = node.with_label(label);
167            }
168            Ok(node)
169        })
170        .collect::<Result<Vec<_>, GraphLoadError>>()?;
171
172    let edges = response
173        .edges
174        .into_iter()
175        .enumerate()
176        .map(|(index, mut row)| {
177            let edge_id = take_identity(&mut row, EDGE_ID, "edge", index)?;
178            let source_internal = take_identity(&mut row, EDGE_SOURCE, "edge", index)?;
179            let target_internal = take_identity(&mut row, EDGE_TARGET, "edge", index)?;
180            let source = internal_to_external
181                .get(&source_internal)
182                .cloned()
183                .ok_or_else(|| {
184                    invalid_row(
185                        "edge",
186                        index,
187                        format!("source {source_internal} is outside the node selection"),
188                    )
189                })?;
190            let target = internal_to_external
191                .get(&target_internal)
192                .cloned()
193                .ok_or_else(|| {
194                    invalid_row(
195                        "edge",
196                        index,
197                        format!("target {target_internal} is outside the node selection"),
198                    )
199                })?;
200            let key = match &spec.edge_key_identity {
201                Some(selection) => {
202                    take_optional_external_id(&mut row, EDGE_KEY, "edge", index, selection)?
203                }
204                None => {
205                    row.remove(EDGE_KEY);
206                    None
207                }
208            };
209            let label = take_optional_string(&mut row, EDGE_LABEL, "edge", index)?;
210            let weight = take_optional_weight(&mut row, EDGE_WEIGHT, index)?;
211            let mut edge = Edge::new(edge_id, source, target).with_attributes(row);
212            if let Some(key) = key {
213                edge = edge.with_graphify_key(key);
214            }
215            if let Some(label) = label {
216                edge = edge.with_label(label);
217            }
218            if let Some(weight) = weight {
219                edge = edge.with_weight(weight);
220            }
221            Ok(edge)
222        })
223        .collect::<Result<Vec<_>, GraphLoadError>>()?;
224    Graph::with_attributes(spec.kind, graph_attributes, nodes, edges).map_err(Into::into)
225}
226
227fn take_identity(
228    row: &mut BTreeMap<String, Value>,
229    field: &str,
230    kind: &'static str,
231    index: usize,
232) -> Result<String, GraphLoadError> {
233    let value = row
234        .remove(field)
235        .ok_or_else(|| invalid_row(kind, index, format!("missing {field}")))?;
236    identity(value).ok_or_else(|| invalid_row(kind, index, format!("invalid {field}")))
237}
238
239fn take_external_id(
240    row: &mut BTreeMap<String, Value>,
241    field: &str,
242    kind: &'static str,
243    index: usize,
244    selection: &IdentitySelection,
245) -> Result<ExternalId, GraphLoadError> {
246    let value = row
247        .remove(field)
248        .ok_or_else(|| invalid_row(kind, index, format!("missing {field}")))?;
249    decode_external_id(value, selection)
250        .map_err(|error| invalid_row(kind, index, format!("invalid {field}: {error}")))
251}
252
253fn take_optional_external_id(
254    row: &mut BTreeMap<String, Value>,
255    field: &str,
256    kind: &'static str,
257    index: usize,
258    selection: &IdentitySelection,
259) -> Result<Option<ExternalId>, GraphLoadError> {
260    match row.remove(field) {
261        None => Ok(None),
262        Some(value) => decode_external_id(value, selection)
263            .map(Some)
264            .map_err(|error| invalid_row(kind, index, format!("invalid {field}: {error}"))),
265    }
266}
267
268fn decode_external_id(
269    value: Value,
270    selection: &IdentitySelection,
271) -> Result<ExternalId, GraphError> {
272    match selection {
273        IdentitySelection::InternalId | IdentitySelection::ScalarProperty(_) => {
274            ExternalId::from_scalar(value)
275        }
276        IdentitySelection::TaggedProperty(_) => ExternalId::from_tagged_value(value),
277    }
278}
279
280fn identity(value: Value) -> Option<String> {
281    match value {
282        Value::String(value) if !value.is_empty() => Some(value),
283        Value::Number(value) => Some(value.to_string()),
284        Value::Null | Value::Bool(_) | Value::String(_) | Value::Array(_) | Value::Object(_) => {
285            None
286        }
287    }
288}
289
290fn take_optional_string(
291    row: &mut BTreeMap<String, Value>,
292    field: &str,
293    kind: &'static str,
294    index: usize,
295) -> Result<Option<String>, GraphLoadError> {
296    match row.remove(field) {
297        None | Some(Value::Null) => Ok(None),
298        Some(Value::String(value)) => Ok(Some(value)),
299        Some(_) => Err(invalid_row(kind, index, format!("invalid {field}"))),
300    }
301}
302
303fn take_optional_weight(
304    row: &mut BTreeMap<String, Value>,
305    field: &str,
306    index: usize,
307) -> Result<Option<f64>, GraphLoadError> {
308    match row.remove(field) {
309        None | Some(Value::Null) => Ok(None),
310        Some(Value::Number(value)) => value
311            .as_f64()
312            .filter(|weight| weight.is_finite() && *weight >= 0.0)
313            .map(Some)
314            .ok_or_else(|| invalid_row("edge", index, format!("invalid {field}"))),
315        Some(_) => Err(invalid_row("edge", index, format!("invalid {field}"))),
316    }
317}
318
319fn invalid_row(kind: &'static str, index: usize, details: String) -> GraphLoadError {
320    GraphLoadError::InvalidRow {
321        kind,
322        index,
323        details,
324    }
325}
326
327#[cfg(test)]
328mod tests {
329    use serde_json::json;
330
331    use super::*;
332    use crate::NodeId;
333
334    fn response(nodes: Value, edges: Value) -> Vec<u8> {
335        serde_json::to_vec(&json!({ "nodes": nodes, "edges": edges })).expect("fixture JSON")
336    }
337
338    fn response_with_metadata(nodes: Value, edges: Value, metadata: Value) -> Vec<u8> {
339        serde_json::to_vec(&json!({
340            "nodes": nodes,
341            "edges": edges,
342            "metadata": metadata,
343        }))
344        .expect("fixture JSON")
345    }
346
347    fn spec() -> GraphLoadSpec {
348        GraphLoadSpec {
349            kind: GraphKind::DiGraph,
350            node_identity: IdentitySelection::ScalarProperty(
351                crate::GraphProperty::new("external_id").unwrap(),
352            ),
353            edge_key_identity: Some(IdentitySelection::ScalarProperty(
354                crate::GraphProperty::new("key").unwrap(),
355            )),
356            node_limit: None,
357            edge_limit: None,
358        }
359    }
360
361    #[test]
362    fn loader_preserves_identity_topology_properties_and_weights() {
363        let bytes = response(
364            json!([
365                { (NODE_ID): "n1", (EXTERNAL_ID): "a", (NODE_LABEL): "File", "name": "A" },
366                { (NODE_ID): "n2", (EXTERNAL_ID): "b", (NODE_LABEL): "File", "name": "B" }
367            ]),
368            json!([{
369                (EDGE_ID): "e1", (EDGE_SOURCE): "n1", (EDGE_TARGET): "n2",
370                (EDGE_KEY): "imports", (EDGE_LABEL): "DEPENDS_ON", (EDGE_WEIGHT): 2.5,
371                "relation": "import"
372            }]),
373        );
374        let graph = graph_from_response(spec(), &bytes).expect("valid graph response");
375        assert_eq!(graph.node_count(), 2);
376        let edge = graph.edge("e1").expect("edge");
377        assert_eq!(
378            (&edge.source, &edge.target),
379            (&NodeId::from("a"), &NodeId::from("b"))
380        );
381        assert_eq!(edge.weight, Some(2.5));
382    }
383
384    #[test]
385    fn loader_rejects_duplicate_external_identity_and_missing_endpoint() {
386        let duplicate = response(
387            json!([
388                { (NODE_ID): "n1", (EXTERNAL_ID): "same", (NODE_LABEL): null },
389                { (NODE_ID): "n2", (EXTERNAL_ID): "same", (NODE_LABEL): null }
390            ]),
391            json!([]),
392        );
393        assert!(matches!(
394            graph_from_response(spec(), &duplicate),
395            Err(GraphLoadError::DuplicateExternalIdentity(id)) if id == "same"
396        ));
397
398        let outside = response(
399            json!([{ (NODE_ID): "n1", (EXTERNAL_ID): "a", (NODE_LABEL): null }]),
400            json!([{
401                (EDGE_ID): "e1", (EDGE_SOURCE): "n1", (EDGE_TARGET): "missing",
402                (EDGE_LABEL): null
403            }]),
404        );
405        assert!(matches!(
406            graph_from_response(spec(), &outside),
407            Err(GraphLoadError::InvalidRow { kind: "edge", .. })
408        ));
409    }
410
411    #[test]
412    fn loader_rejects_malformed_rows_invalid_weights_and_truncation() {
413        assert!(matches!(
414            graph_from_response(spec(), b"[]"),
415            Err(GraphLoadError::InvalidResponse(_))
416        ));
417        let bad_weight = response(
418            json!([{ (NODE_ID): "n1", (EXTERNAL_ID): "a", (NODE_LABEL): null }]),
419            json!([{
420                (EDGE_ID): "e1", (EDGE_SOURCE): "n1", (EDGE_TARGET): "n1",
421                (EDGE_LABEL): null, (EDGE_WEIGHT): "heavy"
422            }]),
423        );
424        assert!(matches!(
425            graph_from_response(spec(), &bad_weight),
426            Err(GraphLoadError::InvalidRow { kind: "edge", .. })
427        ));
428
429        let nodes = response(
430            json!([
431                { (NODE_ID): "n1", (EXTERNAL_ID): "a", (NODE_LABEL): null },
432                { (NODE_ID): "n2", (EXTERNAL_ID): "b", (NODE_LABEL): null }
433            ]),
434            json!([]),
435        );
436        assert!(matches!(
437            graph_from_response(
438                GraphLoadSpec {
439                    node_limit: Some(1),
440                    ..spec()
441                },
442                &nodes
443            ),
444            Err(GraphLoadError::IncompleteSelection {
445                kind: "node",
446                limit: 1
447            })
448        ));
449
450        let edges = response(
451            json!([{ (NODE_ID): 1, (EXTERNAL_ID): 10, (NODE_LABEL): null }]),
452            json!([{
453                (EDGE_ID): 2, (EDGE_SOURCE): 1, (EDGE_TARGET): 1,
454                (EDGE_LABEL): null
455            }]),
456        );
457        let graph = graph_from_response(spec(), &edges).expect("numeric identities are scalar");
458        assert!(graph.contains_node(10_i64));
459        assert!(matches!(
460            graph_from_response(
461                GraphLoadSpec {
462                    edge_limit: Some(0),
463                    ..spec()
464                },
465                &edges
466            ),
467            Err(GraphLoadError::IncompleteSelection {
468                kind: "edge",
469                limit: 0
470            })
471        ));
472    }
473
474    #[test]
475    fn loader_rejects_duplicate_internal_ids_missing_sources_and_invalid_optional_fields() {
476        let duplicate = response(
477            json!([
478                { (NODE_ID): "n1", (EXTERNAL_ID): "a", (NODE_LABEL): null },
479                { (NODE_ID): "n1", (EXTERNAL_ID): "b", (NODE_LABEL): null }
480            ]),
481            json!([]),
482        );
483        assert!(matches!(
484            graph_from_response(spec(), &duplicate),
485            Err(GraphLoadError::InvalidRow { kind: "node", .. })
486        ));
487        let missing_source = response(
488            json!([{ (NODE_ID): "n1", (EXTERNAL_ID): "a", (NODE_LABEL): null }]),
489            json!([{
490                (EDGE_ID): "e", (EDGE_SOURCE): "missing", (EDGE_TARGET): "n1",
491                (EDGE_LABEL): null
492            }]),
493        );
494        assert!(matches!(
495            graph_from_response(spec(), &missing_source),
496            Err(GraphLoadError::InvalidRow { kind: "edge", .. })
497        ));
498        let invalid_key = response(
499            json!([{ (NODE_ID): "n1", (EXTERNAL_ID): "a", (NODE_LABEL): null }]),
500            json!([{
501                (EDGE_ID): "e", (EDGE_SOURCE): "n1", (EDGE_TARGET): "n1",
502                (EDGE_KEY): {}, (EDGE_LABEL): 1
503            }]),
504        );
505        assert!(matches!(
506            graph_from_response(spec(), &invalid_key),
507            Err(GraphLoadError::InvalidRow { kind: "edge", .. })
508        ));
509    }
510
511    #[test]
512    fn scalar_and_tagged_identities_preserve_types_without_collisions() {
513        let scalar = response(
514            json!([
515                { (NODE_ID): "n1", (EXTERNAL_ID): 1, (NODE_LABEL): null },
516                { (NODE_ID): "n2", (EXTERNAL_ID): "1", (NODE_LABEL): null }
517            ]),
518            json!([]),
519        );
520        let graph = graph_from_response(spec(), &scalar).unwrap();
521        assert!(graph.contains_node(1_i64));
522        assert!(graph.contains_node("1"));
523
524        let tuple =
525            ExternalId::tuple(vec![ExternalId::from(1_i64), ExternalId::from("1")]).unwrap();
526        let bytes = ExternalId::Bytes(vec![0, 255]);
527        let tagged = response(
528            json!([
529                {
530                    (NODE_ID): "n1",
531                    (EXTERNAL_ID): serde_json::to_value(&tuple).unwrap(),
532                    (NODE_LABEL): null
533                },
534                {
535                    (NODE_ID): "n2",
536                    (EXTERNAL_ID): serde_json::to_value(&bytes).unwrap(),
537                    (NODE_LABEL): null
538                }
539            ]),
540            json!([]),
541        );
542        let graph = graph_from_response(
543            GraphLoadSpec {
544                node_identity: IdentitySelection::TaggedProperty(
545                    crate::GraphProperty::new("external_id").unwrap(),
546                ),
547                edge_key_identity: None,
548                ..spec()
549            },
550            &tagged,
551        )
552        .unwrap();
553        assert!(graph.contains_node(tuple));
554        assert!(graph.contains_node(bytes));
555    }
556
557    #[test]
558    fn internal_identity_and_no_edge_key_selection_are_explicit() {
559        let bytes = response(
560            json!([
561                { (NODE_ID): "n1", (EXTERNAL_ID): "n1", (NODE_LABEL): null },
562                { (NODE_ID): "n2", (EXTERNAL_ID): "n2", (NODE_LABEL): null }
563            ]),
564            json!([{
565                (EDGE_ID): "e1", (EDGE_SOURCE): "n1", (EDGE_TARGET): "n2",
566                (EDGE_KEY): {"ignored": true}, (EDGE_LABEL): null, (EDGE_WEIGHT): null
567            }]),
568        );
569        let graph = graph_from_response(
570            GraphLoadSpec {
571                node_identity: IdentitySelection::InternalId,
572                edge_key_identity: None,
573                ..spec()
574            },
575            &bytes,
576        )
577        .unwrap();
578        assert!(graph.contains_node("n1"));
579        assert_eq!(graph.edge("e1").unwrap().graphify_key, None);
580        assert!(!graph.edge("e1").unwrap().attributes.contains_key(EDGE_KEY));
581
582        let invalid_label = response(
583            json!([{ (NODE_ID): "n1", (EXTERNAL_ID): "n1", (NODE_LABEL): null }]),
584            json!([{
585                (EDGE_ID): "e1", (EDGE_SOURCE): "n1", (EDGE_TARGET): "n1",
586                (EDGE_LABEL): 1
587            }]),
588        );
589        assert!(matches!(
590            graph_from_response(spec(), &invalid_label),
591            Err(GraphLoadError::InvalidRow { kind: "edge", .. })
592        ));
593    }
594
595    #[test]
596    fn declared_kind_controls_parallel_edges_even_for_small_graphs() {
597        let parallel = response(
598            json!([
599                { (NODE_ID): "n1", (EXTERNAL_ID): "a", (NODE_LABEL): null },
600                { (NODE_ID): "n2", (EXTERNAL_ID): "b", (NODE_LABEL): null }
601            ]),
602            json!([
603                {
604                    (EDGE_ID): "e1", (EDGE_SOURCE): "n1", (EDGE_TARGET): "n2",
605                    (EDGE_LABEL): null
606                },
607                {
608                    (EDGE_ID): "e2", (EDGE_SOURCE): "n1", (EDGE_TARGET): "n2",
609                    (EDGE_LABEL): null
610                }
611            ]),
612        );
613        assert!(matches!(
614            graph_from_response(spec(), &parallel),
615            Err(GraphLoadError::Graph(GraphError::ParallelEdge { .. }))
616        ));
617        let multigraph = graph_from_response(
618            GraphLoadSpec {
619                kind: GraphKind::MultiDiGraph,
620                ..spec()
621            },
622            &parallel,
623        )
624        .unwrap();
625        assert!(multigraph.is_multigraph());
626
627        let empty = graph_from_response(
628            GraphLoadSpec {
629                kind: GraphKind::MultiGraph,
630                ..spec()
631            },
632            &response(json!([]), json!([])),
633        )
634        .unwrap();
635        assert!(empty.is_multigraph());
636    }
637
638    #[test]
639    fn metadata_loads_once_and_rejects_multiple_or_reserved_rows() {
640        let graph = graph_from_response(
641            spec(),
642            &response_with_metadata(json!([]), json!([]), json!([{"name": "demo"}])),
643        )
644        .unwrap();
645        assert_eq!(graph.attributes()["name"], json!("demo"));
646
647        assert!(matches!(
648            graph_from_response(
649                spec(),
650                &response_with_metadata(
651                    json!([]),
652                    json!([]),
653                    json!([{"name": "one"}, {"name": "two"}])
654                )
655            ),
656            Err(GraphLoadError::MultipleGraphMetadataRows { count: 2 })
657        ));
658        assert!(matches!(
659            graph_from_response(
660                spec(),
661                &response_with_metadata(json!([]), json!([]), json!([{(NODE_ID): "private"}]))
662            ),
663            Err(GraphLoadError::InvalidRow {
664                kind: "metadata",
665                ..
666            })
667        ));
668    }
669}
helix_graph_algorithms/loader.rs

helix_graph_algorithms/
loader.rs
