Skip to main content

nodedb_types/
graph.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Shared graph types used by both Origin and Lite CSR engines.
4
5use serde::{Deserialize, Serialize};
6
7/// Aggregate stats for a single graph collection (or the full edge store
8/// when no collection is specified).
9///
10/// Returned by [`NodeDb::graph_stats`]. Wire-safe: serializes to/from
11/// MessagePack so the value can cross the Data-Plane boundary unchanged.
12///
13/// The `node_count` field counts distinct node IDs observed as edge
14/// endpoints (not necessarily all nodes ever inserted). For Origin, this
15/// equals `distinct_node_count` from the persistent stats table; for Lite,
16/// it is derived from the CRDT edge store at query time.
17#[derive(
18    Debug,
19    Clone,
20    PartialEq,
21    Serialize,
22    Deserialize,
23    zerompk::ToMessagePack,
24    zerompk::FromMessagePack,
25)]
26#[msgpack(map)]
27pub struct GraphStats {
28    /// Name of the collection (or `"__edges"` on Lite when no collection is
29    /// supplied).
30    pub collection: String,
31    /// Distinct node IDs that appear as an edge source or destination.
32    pub node_count: u64,
33    /// Total number of edges in the collection.
34    pub edge_count: u64,
35    /// Number of distinct edge label strings.
36    pub distinct_label_count: u64,
37    /// Per-label edge counts, sorted ascending by label name.
38    pub labels: Vec<(String, u64)>,
39}
40
41/// Edge traversal direction.
42#[derive(
43    Debug,
44    Clone,
45    Copy,
46    PartialEq,
47    Eq,
48    Hash,
49    Serialize,
50    Deserialize,
51    zerompk::ToMessagePack,
52    zerompk::FromMessagePack,
53)]
54#[msgpack(c_enum)]
55pub enum Direction {
56    /// Outgoing edges only.
57    Out,
58    /// Incoming edges only.
59    In,
60    /// Both directions.
61    Both,
62}
63
64impl Direction {
65    pub fn as_str(&self) -> &'static str {
66        match self {
67            Self::Out => "out",
68            Self::In => "in",
69            Self::Both => "both",
70        }
71    }
72}
73
74impl std::fmt::Display for Direction {
75    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
76        f.write_str(self.as_str())
77    }
78}
79
80/// A traversal direction string that names no [`Direction`].
81#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
82#[error("unknown direction '{input}': expected 'out', 'in' or 'both'")]
83pub struct ParseDirectionError {
84    pub input: String,
85}
86
87impl std::str::FromStr for Direction {
88    type Err = ParseDirectionError;
89    /// Case-insensitive. Accepts `out`/`outgoing`, `in`/`incoming` and
90    /// `both`/`any`.
91    fn from_str(s: &str) -> Result<Self, Self::Err> {
92        match s.to_ascii_lowercase().as_str() {
93            "out" | "outgoing" => Ok(Self::Out),
94            "in" | "incoming" => Ok(Self::In),
95            "both" | "any" => Ok(Self::Both),
96            _ => Err(ParseDirectionError {
97                input: s.to_string(),
98            }),
99        }
100    }
101}
102
103impl GraphStats {
104    /// Column names of the wire row shape produced by
105    /// `SHOW GRAPH STATS [<'collection'>]`. Single source of truth — the
106    /// parser validates against this, and clients use it to build the
107    /// `columns` slice in tests so the two never drift.
108    pub const EXPECTED_COLUMNS: [&'static str; 5] = [
109        "collection",
110        "node_count",
111        "edge_count",
112        "distinct_label_count",
113        "labels",
114    ];
115
116    pub fn zero(collection: impl Into<String>) -> Self {
117        Self {
118            collection: collection.into(),
119            node_count: 0,
120            edge_count: 0,
121            distinct_label_count: 0,
122            labels: Vec::new(),
123        }
124    }
125
126    /// Parse the wire shape produced by `SHOW GRAPH STATS [<'collection'>]`
127    /// into a vec of `GraphStats`, one entry per row. Used by both the
128    /// native and the pgwire remote clients — keeping a single parser here
129    /// is what makes the `wire_shape:` smoke-probe error messages a
130    /// meaningful diagnostic rather than a coincidence between two copies.
131    ///
132    /// Expected columns: `(collection, node_count, edge_count,
133    /// distinct_label_count, labels)`, where `labels` is a JSON array of
134    /// `{"label": str, "count": u64}` objects.
135    ///
136    /// Both wire paths can deliver count cells as `Value::Integer`
137    /// (extended-query / native typed) or `Value::String` (pgwire simple
138    /// query text protocol); both shapes are accepted.
139    ///
140    /// Column-shape mismatches surface as errors (this is exactly the bug
141    /// class the smoke probe is designed to trap, so no fallback); empty
142    /// rows return an empty vec — callers decide how to interpret no rows.
143    pub fn parse_show_stats_response(
144        columns: &[String],
145        rows: &[Vec<crate::value::Value>],
146    ) -> crate::error::NodeDbResult<Vec<Self>> {
147        use crate::error::NodeDbError;
148
149        if columns.len() != Self::EXPECTED_COLUMNS.len()
150            || columns
151                .iter()
152                .zip(Self::EXPECTED_COLUMNS.iter())
153                .any(|(a, b)| a != b)
154        {
155            if !columns.is_empty() {
156                return Err(NodeDbError::storage(format!(
157                    "wire_shape: SHOW GRAPH STATS returned unexpected columns: {columns:?}"
158                )));
159            }
160            // No columns and no rows: treat as empty result.
161            return Ok(Vec::new());
162        }
163
164        if rows.is_empty() {
165            return Ok(Vec::new());
166        }
167
168        let mut out = Vec::with_capacity(rows.len());
169        for row in rows {
170            out.push(Self::parse_one_row(row)?);
171        }
172        Ok(out)
173    }
174
175    fn parse_one_row(row: &[crate::value::Value]) -> crate::error::NodeDbResult<Self> {
176        use crate::error::NodeDbError;
177
178        let coll_name = row
179            .first()
180            .and_then(|v| v.as_str())
181            .ok_or_else(|| {
182                NodeDbError::storage("wire_shape: SHOW GRAPH STATS: missing collection cell")
183            })?
184            .to_string();
185        let [_, node_column, edge_column, label_column, _] = Self::EXPECTED_COLUMNS;
186        let node_count = parse_u64_cell(row, 1, node_column)?;
187        let edge_count = parse_u64_cell(row, 2, edge_column)?;
188        let distinct_label_count = parse_u64_cell(row, 3, label_column)?;
189        let labels_json = row.get(4).and_then(|v| v.as_str()).unwrap_or("[]");
190        let parsed: Vec<sonic_rs::Value> = sonic_rs::from_str(labels_json)
191            .map_err(|e| NodeDbError::storage(format!("wire_shape: labels JSON parse: {e}")))?;
192        let mut labels: Vec<(String, u64)> = Vec::with_capacity(parsed.len());
193        for entry in &parsed {
194            use sonic_rs::JsonValueTrait;
195            let label = entry
196                .get("label")
197                .and_then(|v| v.as_str())
198                .ok_or_else(|| NodeDbError::storage("wire_shape: labels entry missing 'label'"))?
199                .to_string();
200            let count = entry
201                .get("count")
202                .and_then(|v| v.as_u64())
203                .ok_or_else(|| NodeDbError::storage("wire_shape: labels entry missing 'count'"))?;
204            labels.push((label, count));
205        }
206        Ok(Self {
207            collection: coll_name,
208            node_count,
209            edge_count,
210            distinct_label_count,
211            labels,
212        })
213    }
214}
215
216/// Parse the count cell `column` at `index` of a `SHOW GRAPH STATS` row.
217///
218/// The cell arrives typed or as pgwire text:
219/// - `Value::Integer`: a count up to `i64::MAX`. A negative value is an error.
220/// - `Value::Decimal` with scale 0: a msgpack `uint64` above `i64::MAX`.
221/// - `Value::String`: the base-10 text form.
222///
223/// A missing cell or any other shape is an error that names the column.
224fn parse_u64_cell(
225    row: &[crate::value::Value],
226    index: usize,
227    column: &str,
228) -> crate::error::NodeDbResult<u64> {
229    use crate::error::NodeDbError;
230    use crate::value::Value;
231
232    let cell = row.get(index).ok_or_else(|| {
233        NodeDbError::storage(format!("wire_shape: SHOW GRAPH STATS: missing {column}"))
234    })?;
235    let parsed = match cell {
236        Value::Integer(i) => u64::try_from(*i).ok(),
237        Value::Decimal(d) => Value::decimal_as_wide_u64(d),
238        Value::String(s) => s.parse::<u64>().ok(),
239        _ => None,
240    };
241    parsed.ok_or_else(|| {
242        NodeDbError::storage(format!(
243            "wire_shape: SHOW GRAPH STATS: {column} is not a u64 count: {cell:?}"
244        ))
245    })
246}
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251
252    #[test]
253    fn direction_parses_case_insensitively_and_names_an_unknown_value() {
254        assert_eq!("IN".parse::<Direction>(), Ok(Direction::In));
255        assert_eq!("outgoing".parse::<Direction>(), Ok(Direction::Out));
256        assert_eq!("Both".parse::<Direction>(), Ok(Direction::Both));
257        let err = "sideways".parse::<Direction>().unwrap_err();
258        assert_eq!(err.input, "sideways");
259        assert!(err.to_string().contains("'sideways'"), "{err}");
260    }
261
262    #[test]
263    fn direction_roundtrip() {
264        for dir in [Direction::Out, Direction::In, Direction::Both] {
265            let s = dir.as_str();
266            let parsed: Direction = s.parse().unwrap();
267            assert_eq!(dir, parsed);
268        }
269    }
270
271    #[test]
272    fn direction_display() {
273        assert_eq!(Direction::Out.to_string(), "out");
274    }
275
276    #[test]
277    fn graph_stats_zero() {
278        let s = GraphStats::zero("my_coll");
279        assert_eq!(s.collection, "my_coll");
280        assert_eq!(s.node_count, 0);
281        assert_eq!(s.edge_count, 0);
282        assert_eq!(s.distinct_label_count, 0);
283        assert!(s.labels.is_empty());
284    }
285
286    #[test]
287    fn graph_stats_serde_round_trip() {
288        let s = GraphStats {
289            collection: "coll".into(),
290            node_count: 10,
291            edge_count: 5,
292            distinct_label_count: 2,
293            labels: vec![("KNOWS".into(), 3), ("OWNS".into(), 2)],
294        };
295        let json = sonic_rs::to_string(&s).unwrap();
296        let back: GraphStats = sonic_rs::from_str(&json).unwrap();
297        assert_eq!(back, s);
298    }
299
300    #[test]
301    fn parse_show_stats_multi_row() {
302        use crate::value::Value;
303        let columns: Vec<String> = GraphStats::EXPECTED_COLUMNS
304            .iter()
305            .map(|s| s.to_string())
306            .collect();
307        let labels_json = r#"[{"label":"KNOWS","count":3},{"label":"OWNS","count":2}]"#;
308        let rows = vec![
309            vec![
310                Value::String("social".into()),
311                // pgwire simple-query text protocol arrives as strings;
312                // the native protocol arrives as integers — both shapes are accepted.
313                Value::String("10".into()),
314                Value::Integer(5),
315                Value::String("2".into()),
316                Value::String(labels_json.into()),
317            ],
318            vec![
319                Value::String("comms".into()),
320                Value::Integer(3),
321                Value::Integer(2),
322                Value::Integer(1),
323                Value::String(r#"[{"label":"CALLS","count":2}]"#.into()),
324            ],
325        ];
326        let result = GraphStats::parse_show_stats_response(&columns, &rows).unwrap();
327        assert_eq!(result.len(), 2);
328        let social = &result[0];
329        assert_eq!(social.collection, "social");
330        assert_eq!(social.node_count, 10);
331        assert_eq!(social.edge_count, 5);
332        assert_eq!(social.distinct_label_count, 2);
333        assert_eq!(social.labels, vec![("KNOWS".into(), 3), ("OWNS".into(), 2)]);
334        let comms = &result[1];
335        assert_eq!(comms.collection, "comms");
336        assert_eq!(comms.edge_count, 2);
337        assert_eq!(comms.labels, vec![("CALLS".into(), 2)]);
338    }
339
340    #[test]
341    fn parse_show_stats_empty_rows_returns_empty_vec() {
342        let columns: Vec<String> = GraphStats::EXPECTED_COLUMNS
343            .iter()
344            .map(|s| s.to_string())
345            .collect();
346        let result = GraphStats::parse_show_stats_response(&columns, &[]).unwrap();
347        assert!(result.is_empty());
348    }
349
350    #[test]
351    fn parse_show_stats_wrong_columns_errors() {
352        let columns = vec!["id".to_string(), "count".to_string()];
353        let err = GraphStats::parse_show_stats_response(&columns, &[]).unwrap_err();
354        assert!(
355            err.to_string().contains("unexpected columns"),
356            "error should mention unexpected columns: {err}"
357        );
358    }
359
360    #[test]
361    fn parse_show_stats_no_columns_no_rows_returns_empty_vec() {
362        let result = GraphStats::parse_show_stats_response(&[], &[]).unwrap();
363        assert!(result.is_empty());
364    }
365
366    fn stats_row(count: crate::value::Value) -> Vec<crate::value::Value> {
367        use crate::value::Value;
368        vec![
369            Value::String("c".into()),
370            count,
371            Value::Integer(0),
372            Value::Integer(0),
373            Value::String("[]".into()),
374        ]
375    }
376
377    #[test]
378    fn parse_u64_cell_accepts_a_count_above_i64_max() {
379        use crate::value::Value;
380        let wide = u64::MAX;
381        let row = stats_row(Value::from_u64(wide));
382        assert!(matches!(row[1], Value::Decimal(_)));
383        assert_eq!(parse_u64_cell(&row, 1, "node_count").unwrap(), wide);
384        let columns: Vec<String> = GraphStats::EXPECTED_COLUMNS
385            .iter()
386            .map(|s| s.to_string())
387            .collect();
388        let parsed = GraphStats::parse_show_stats_response(&columns, &[row]).unwrap();
389        assert_eq!(parsed[0].node_count, wide);
390    }
391
392    #[test]
393    fn parse_u64_cell_refuses_a_negative_count() {
394        use crate::value::Value;
395        let err = parse_u64_cell(&stats_row(Value::Integer(-1)), 1, "node_count").unwrap_err();
396        assert!(err.to_string().contains("node_count"), "{err}");
397    }
398
399    #[test]
400    fn parse_u64_cell_refuses_a_scaled_decimal() {
401        use crate::value::Value;
402        let scaled = rust_decimal::Decimal::new(15, 1);
403        assert!(parse_u64_cell(&stats_row(Value::Decimal(scaled)), 1, "node_count").is_err());
404    }
405
406    #[test]
407    fn parse_u64_cell_names_a_missing_cell() {
408        let err = parse_u64_cell(&[], 2, "edge_count").unwrap_err();
409        assert!(err.to_string().contains("missing edge_count"), "{err}");
410    }
411
412    #[test]
413    fn graph_stats_msgpack_round_trip() {
414        let s = GraphStats {
415            collection: "coll".into(),
416            node_count: 7,
417            edge_count: 3,
418            distinct_label_count: 1,
419            labels: vec![("FOLLOWS".into(), 3)],
420        };
421        let bytes = zerompk::to_msgpack_vec(&s).unwrap();
422        let back: GraphStats = zerompk::from_msgpack(&bytes).unwrap();
423        assert_eq!(back, s);
424    }
425}