use serde::{Deserialize, Serialize};
#[derive(
Debug,
Clone,
PartialEq,
Serialize,
Deserialize,
zerompk::ToMessagePack,
zerompk::FromMessagePack,
)]
#[msgpack(map)]
pub struct GraphStats {
pub collection: String,
pub node_count: u64,
pub edge_count: u64,
pub distinct_label_count: u64,
pub labels: Vec<(String, u64)>,
}
#[derive(
Debug,
Clone,
Copy,
PartialEq,
Eq,
Hash,
Serialize,
Deserialize,
zerompk::ToMessagePack,
zerompk::FromMessagePack,
)]
#[msgpack(c_enum)]
pub enum Direction {
Out,
In,
Both,
}
impl Direction {
pub fn as_str(&self) -> &'static str {
match self {
Self::Out => "out",
Self::In => "in",
Self::Both => "both",
}
}
}
impl std::fmt::Display for Direction {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[error("unknown direction '{input}': expected 'out', 'in' or 'both'")]
pub struct ParseDirectionError {
pub input: String,
}
impl std::str::FromStr for Direction {
type Err = ParseDirectionError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.to_ascii_lowercase().as_str() {
"out" | "outgoing" => Ok(Self::Out),
"in" | "incoming" => Ok(Self::In),
"both" | "any" => Ok(Self::Both),
_ => Err(ParseDirectionError {
input: s.to_string(),
}),
}
}
}
impl GraphStats {
pub const EXPECTED_COLUMNS: [&'static str; 5] = [
"collection",
"node_count",
"edge_count",
"distinct_label_count",
"labels",
];
pub fn zero(collection: impl Into<String>) -> Self {
Self {
collection: collection.into(),
node_count: 0,
edge_count: 0,
distinct_label_count: 0,
labels: Vec::new(),
}
}
pub fn parse_show_stats_response(
columns: &[String],
rows: &[Vec<crate::value::Value>],
) -> crate::error::NodeDbResult<Vec<Self>> {
use crate::error::NodeDbError;
if columns.len() != Self::EXPECTED_COLUMNS.len()
|| columns
.iter()
.zip(Self::EXPECTED_COLUMNS.iter())
.any(|(a, b)| a != b)
{
if !columns.is_empty() {
return Err(NodeDbError::storage(format!(
"wire_shape: SHOW GRAPH STATS returned unexpected columns: {columns:?}"
)));
}
return Ok(Vec::new());
}
if rows.is_empty() {
return Ok(Vec::new());
}
let mut out = Vec::with_capacity(rows.len());
for row in rows {
out.push(Self::parse_one_row(row)?);
}
Ok(out)
}
fn parse_one_row(row: &[crate::value::Value]) -> crate::error::NodeDbResult<Self> {
use crate::error::NodeDbError;
let coll_name = row
.first()
.and_then(|v| v.as_str())
.ok_or_else(|| {
NodeDbError::storage("wire_shape: SHOW GRAPH STATS: missing collection cell")
})?
.to_string();
let [_, node_column, edge_column, label_column, _] = Self::EXPECTED_COLUMNS;
let node_count = parse_u64_cell(row, 1, node_column)?;
let edge_count = parse_u64_cell(row, 2, edge_column)?;
let distinct_label_count = parse_u64_cell(row, 3, label_column)?;
let labels_json = row.get(4).and_then(|v| v.as_str()).unwrap_or("[]");
let parsed: Vec<sonic_rs::Value> = sonic_rs::from_str(labels_json)
.map_err(|e| NodeDbError::storage(format!("wire_shape: labels JSON parse: {e}")))?;
let mut labels: Vec<(String, u64)> = Vec::with_capacity(parsed.len());
for entry in &parsed {
use sonic_rs::JsonValueTrait;
let label = entry
.get("label")
.and_then(|v| v.as_str())
.ok_or_else(|| NodeDbError::storage("wire_shape: labels entry missing 'label'"))?
.to_string();
let count = entry
.get("count")
.and_then(|v| v.as_u64())
.ok_or_else(|| NodeDbError::storage("wire_shape: labels entry missing 'count'"))?;
labels.push((label, count));
}
Ok(Self {
collection: coll_name,
node_count,
edge_count,
distinct_label_count,
labels,
})
}
}
fn parse_u64_cell(
row: &[crate::value::Value],
index: usize,
column: &str,
) -> crate::error::NodeDbResult<u64> {
use crate::error::NodeDbError;
use crate::value::Value;
let cell = row.get(index).ok_or_else(|| {
NodeDbError::storage(format!("wire_shape: SHOW GRAPH STATS: missing {column}"))
})?;
let parsed = match cell {
Value::Integer(i) => u64::try_from(*i).ok(),
Value::Decimal(d) => Value::decimal_as_wide_u64(d),
Value::String(s) => s.parse::<u64>().ok(),
_ => None,
};
parsed.ok_or_else(|| {
NodeDbError::storage(format!(
"wire_shape: SHOW GRAPH STATS: {column} is not a u64 count: {cell:?}"
))
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn direction_parses_case_insensitively_and_names_an_unknown_value() {
assert_eq!("IN".parse::<Direction>(), Ok(Direction::In));
assert_eq!("outgoing".parse::<Direction>(), Ok(Direction::Out));
assert_eq!("Both".parse::<Direction>(), Ok(Direction::Both));
let err = "sideways".parse::<Direction>().unwrap_err();
assert_eq!(err.input, "sideways");
assert!(err.to_string().contains("'sideways'"), "{err}");
}
#[test]
fn direction_roundtrip() {
for dir in [Direction::Out, Direction::In, Direction::Both] {
let s = dir.as_str();
let parsed: Direction = s.parse().unwrap();
assert_eq!(dir, parsed);
}
}
#[test]
fn direction_display() {
assert_eq!(Direction::Out.to_string(), "out");
}
#[test]
fn graph_stats_zero() {
let s = GraphStats::zero("my_coll");
assert_eq!(s.collection, "my_coll");
assert_eq!(s.node_count, 0);
assert_eq!(s.edge_count, 0);
assert_eq!(s.distinct_label_count, 0);
assert!(s.labels.is_empty());
}
#[test]
fn graph_stats_serde_round_trip() {
let s = GraphStats {
collection: "coll".into(),
node_count: 10,
edge_count: 5,
distinct_label_count: 2,
labels: vec![("KNOWS".into(), 3), ("OWNS".into(), 2)],
};
let json = sonic_rs::to_string(&s).unwrap();
let back: GraphStats = sonic_rs::from_str(&json).unwrap();
assert_eq!(back, s);
}
#[test]
fn parse_show_stats_multi_row() {
use crate::value::Value;
let columns: Vec<String> = GraphStats::EXPECTED_COLUMNS
.iter()
.map(|s| s.to_string())
.collect();
let labels_json = r#"[{"label":"KNOWS","count":3},{"label":"OWNS","count":2}]"#;
let rows = vec![
vec![
Value::String("social".into()),
Value::String("10".into()),
Value::Integer(5),
Value::String("2".into()),
Value::String(labels_json.into()),
],
vec![
Value::String("comms".into()),
Value::Integer(3),
Value::Integer(2),
Value::Integer(1),
Value::String(r#"[{"label":"CALLS","count":2}]"#.into()),
],
];
let result = GraphStats::parse_show_stats_response(&columns, &rows).unwrap();
assert_eq!(result.len(), 2);
let social = &result[0];
assert_eq!(social.collection, "social");
assert_eq!(social.node_count, 10);
assert_eq!(social.edge_count, 5);
assert_eq!(social.distinct_label_count, 2);
assert_eq!(social.labels, vec![("KNOWS".into(), 3), ("OWNS".into(), 2)]);
let comms = &result[1];
assert_eq!(comms.collection, "comms");
assert_eq!(comms.edge_count, 2);
assert_eq!(comms.labels, vec![("CALLS".into(), 2)]);
}
#[test]
fn parse_show_stats_empty_rows_returns_empty_vec() {
let columns: Vec<String> = GraphStats::EXPECTED_COLUMNS
.iter()
.map(|s| s.to_string())
.collect();
let result = GraphStats::parse_show_stats_response(&columns, &[]).unwrap();
assert!(result.is_empty());
}
#[test]
fn parse_show_stats_wrong_columns_errors() {
let columns = vec!["id".to_string(), "count".to_string()];
let err = GraphStats::parse_show_stats_response(&columns, &[]).unwrap_err();
assert!(
err.to_string().contains("unexpected columns"),
"error should mention unexpected columns: {err}"
);
}
#[test]
fn parse_show_stats_no_columns_no_rows_returns_empty_vec() {
let result = GraphStats::parse_show_stats_response(&[], &[]).unwrap();
assert!(result.is_empty());
}
fn stats_row(count: crate::value::Value) -> Vec<crate::value::Value> {
use crate::value::Value;
vec![
Value::String("c".into()),
count,
Value::Integer(0),
Value::Integer(0),
Value::String("[]".into()),
]
}
#[test]
fn parse_u64_cell_accepts_a_count_above_i64_max() {
use crate::value::Value;
let wide = u64::MAX;
let row = stats_row(Value::from_u64(wide));
assert!(matches!(row[1], Value::Decimal(_)));
assert_eq!(parse_u64_cell(&row, 1, "node_count").unwrap(), wide);
let columns: Vec<String> = GraphStats::EXPECTED_COLUMNS
.iter()
.map(|s| s.to_string())
.collect();
let parsed = GraphStats::parse_show_stats_response(&columns, &[row]).unwrap();
assert_eq!(parsed[0].node_count, wide);
}
#[test]
fn parse_u64_cell_refuses_a_negative_count() {
use crate::value::Value;
let err = parse_u64_cell(&stats_row(Value::Integer(-1)), 1, "node_count").unwrap_err();
assert!(err.to_string().contains("node_count"), "{err}");
}
#[test]
fn parse_u64_cell_refuses_a_scaled_decimal() {
use crate::value::Value;
let scaled = rust_decimal::Decimal::new(15, 1);
assert!(parse_u64_cell(&stats_row(Value::Decimal(scaled)), 1, "node_count").is_err());
}
#[test]
fn parse_u64_cell_names_a_missing_cell() {
let err = parse_u64_cell(&[], 2, "edge_count").unwrap_err();
assert!(err.to_string().contains("missing edge_count"), "{err}");
}
#[test]
fn graph_stats_msgpack_round_trip() {
let s = GraphStats {
collection: "coll".into(),
node_count: 7,
edge_count: 3,
distinct_label_count: 1,
labels: vec![("FOLLOWS".into(), 3)],
};
let bytes = zerompk::to_msgpack_vec(&s).unwrap();
let back: GraphStats = zerompk::from_msgpack(&bytes).unwrap();
assert_eq!(back, s);
}
}