nodedb_query/msgpack_scan/
kv_body.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Shape-preserving decode / encode of a KV row body for read-modify-write.
4//!
5//! A KV body is stored in one of two shapes (see [`super::kv_row_msgpack`]):
6//! a msgpack map for typed columns, or raw scalar bytes for the single-`value`
7//! SQL form and RESP `SET`. A merge decodes the body into the same row object
8//! reads present, mutates it, and re-encodes into the shape it came from.
9//! Decoding raw bytes as msgpack instead either fails (multi-byte value) or
10//! reads the first byte as a fixint and discards the rest.
11
12use std::collections::HashMap;
13
14use nodedb_types::{MsgpackError, NotScalar, Value, scalar_to_raw_bytes};
15
16use crate::msgpack_scan::{map_header, write_map_header, write_str};
17
18/// The on-disk shape of a KV row body.
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum KvBodyShape {
21    /// Scalar bytes, not msgpack: the single-`value` SQL form and RESP `SET`.
22    Raw,
23    /// A msgpack map of typed columns.
24    Map,
25}
26
27/// A merged row cannot be encoded back into its body shape.
28#[derive(Debug, thiserror::Error)]
29pub enum KvBodyError {
30    /// The stored map body is not well-formed msgpack.
31    #[error("KV map body: {0}")]
32    Decode(#[from] MsgpackError),
33    /// The merged row is not an object.
34    #[error("KV row must be an object, got {kind}")]
35    RowNotObject {
36        /// `Value::type_name()` of the merged row.
37        kind: &'static str,
38    },
39    /// A raw body holds only `value`; the merged row lost it.
40    #[error("KV raw body has no `value` field after merge")]
41    RawMissingValue,
42    /// A raw body holds only `value`; the merged row carries other keys.
43    #[error("KV row holds a bare `value`, not typed columns; cannot set {keys}")]
44    RawExtraKeys {
45        /// The offending keys, sorted, comma-separated.
46        keys: String,
47    },
48    /// A raw body holds a scalar; the merged `value` is not one.
49    #[error("KV raw body `value`: {0}")]
50    RawNotScalar(#[from] NotScalar),
51    /// The merged map row does not encode.
52    #[error("KV map body encode: {0}")]
53    Encode(zerompk::Error),
54}
55
56/// The shape of a stored body: a msgpack map header marks a map, anything
57/// else (an empty body included) is raw scalar bytes.
58pub fn kv_body_shape(body: &[u8]) -> KvBodyShape {
59    if map_header(body, 0).is_some() {
60        KvBodyShape::Map
61    } else {
62        KvBodyShape::Raw
63    }
64}
65
66/// Decode a stored body into the row object a merge operates on.
67///
68/// A map body decodes as is. A raw body becomes `{"value": <string>}`, the
69/// exact row `kv_row_msgpack` presents to reads (non-UTF-8 bytes take the
70/// lossy view there too). An empty body is a raw empty string. Returns the
71/// shape so the writer re-encodes in the same one.
72pub fn kv_body_to_row(body: &[u8]) -> Result<(Value, KvBodyShape), MsgpackError> {
73    if kv_body_shape(body) == KvBodyShape::Map {
74        return Ok((nodedb_types::value_from_msgpack(body)?, KvBodyShape::Map));
75    }
76    let mut row = std::collections::HashMap::with_capacity(1);
77    row.insert(
78        "value".to_string(),
79        Value::String(String::from_utf8_lossy(body).into_owned()),
80    );
81    Ok((Value::Object(row), KvBodyShape::Raw))
82}
83
84/// Encode a merged row back into `shape`.
85///
86/// `Map` writes the fields in key order, so the same logical row encodes to
87/// the same bytes on every node and on WAL replay. `Raw` accepts only an
88/// object whose single key is `value` holding a scalar; the scalar encodes
89/// through [`scalar_to_raw_bytes`], the same rule the SQL lowering uses for
90/// a fresh insert. Any other object is an error naming the extra keys: a raw
91/// row never silently turns into a map.
92pub fn row_to_kv_body(row: &Value, shape: KvBodyShape) -> Result<Vec<u8>, KvBodyError> {
93    let Value::Object(map) = row else {
94        return Err(KvBodyError::RowNotObject {
95            kind: row.type_name(),
96        });
97    };
98    match shape {
99        KvBodyShape::Map => {
100            let mut fields: Vec<(&String, &Value)> = map.iter().collect();
101            fields.sort_unstable_by(|a, b| a.0.cmp(b.0));
102            let mut buf = Vec::with_capacity(map.len() * 16);
103            write_map_header(&mut buf, fields.len());
104            for (key, value) in fields {
105                write_str(&mut buf, key);
106                let encoded = nodedb_types::value_to_msgpack(value).map_err(KvBodyError::Encode)?;
107                buf.extend_from_slice(&encoded);
108            }
109            Ok(buf)
110        }
111        KvBodyShape::Raw => {
112            let mut extra: Vec<&str> = map
113                .keys()
114                .map(String::as_str)
115                .filter(|k| *k != "value")
116                .collect();
117            if !extra.is_empty() {
118                extra.sort_unstable();
119                return Err(KvBodyError::RawExtraKeys {
120                    keys: extra.join(", "),
121                });
122            }
123            let value = map.get("value").ok_or(KvBodyError::RawMissingValue)?;
124            Ok(scalar_to_raw_bytes(value)?)
125        }
126    }
127}
128
129/// The fields a KV body stores for a msgpack-encoded row, and the body shape.
130///
131/// This is the inverse of [`super::kv_row_msgpack`]. `key` is the entry's key.
132/// A `key` field equal to it is the injected primary key and is dropped. A
133/// `key` field that differs is a stored column and is kept. The remaining
134/// fields pick the shape the SQL lowering picks for a fresh insert:
135/// - `value` alone: [`KvBodyShape::Raw`];
136/// - any other field set: [`KvBodyShape::Map`].
137///
138/// The caller encodes the fields in its own body format.
139/// [`row_to_kv_body`] is the Origin format.
140pub fn kv_row_to_body_fields(
141    key: &str,
142    row: &[u8],
143) -> Result<(HashMap<String, Value>, KvBodyShape), KvBodyError> {
144    let row = nodedb_types::value_from_msgpack(row)?;
145    let Value::Object(mut fields) = row else {
146        return Err(KvBodyError::RowNotObject {
147            kind: row.type_name(),
148        });
149    };
150    if matches!(fields.get("key"), Some(Value::String(stored)) if stored == key) {
151        fields.remove("key");
152    }
153    let shape = if fields.len() == 1 && fields.contains_key("value") {
154        KvBodyShape::Raw
155    } else {
156        KvBodyShape::Map
157    };
158    Ok((fields, shape))
159}
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164
165    fn body_from_row(key: &str, row: &[u8]) -> Vec<u8> {
166        let (fields, shape) = kv_row_to_body_fields(key, row).unwrap();
167        row_to_kv_body(&Value::Object(fields), shape).unwrap()
168    }
169
170    #[test]
171    fn a_raw_body_round_trips_through_its_row() {
172        for body in [b"v1".to_vec(), b"1".to_vec(), Vec::new()] {
173            let row = crate::msgpack_scan::kv_row_msgpack("k1", &body);
174            assert_eq!(body_from_row("k1", &row), body);
175        }
176    }
177
178    #[test]
179    fn a_map_body_round_trips_through_its_row() {
180        let mut fields = HashMap::new();
181        fields.insert("n".to_string(), Value::Integer(7));
182        fields.insert("s".to_string(), Value::String("x".into()));
183        let body = row_to_kv_body(&Value::Object(fields), KvBodyShape::Map).unwrap();
184        let row = crate::msgpack_scan::kv_row_msgpack("k2", &body);
185        assert_eq!(body_from_row("k2", &row), body);
186    }
187
188    #[test]
189    fn a_key_column_that_differs_from_the_entry_key_is_kept() {
190        let mut fields = HashMap::new();
191        fields.insert("key".to_string(), Value::String("stored".into()));
192        fields.insert("n".to_string(), Value::Integer(1));
193        let body = row_to_kv_body(&Value::Object(fields), KvBodyShape::Map).unwrap();
194        let row = crate::msgpack_scan::kv_row_msgpack("slot", &body);
195        assert_eq!(body_from_row("slot", &row), body);
196    }
197
198    #[test]
199    fn a_row_that_is_not_a_map_is_an_error() {
200        let not_a_map = nodedb_types::value_to_msgpack(&Value::Integer(118)).unwrap();
201        assert!(matches!(
202            kv_row_to_body_fields("k", &not_a_map),
203            Err(KvBodyError::RowNotObject { kind: "int" })
204        ));
205    }
206
207    #[test]
208    fn bytes_that_are_not_msgpack_are_a_decode_error() {
209        assert!(matches!(
210            kv_row_to_body_fields("k", &[0x81]),
211            Err(KvBodyError::Decode(_))
212        ));
213    }
214
215    fn value_of(row: &Value) -> &Value {
216        row.get("value").expect("row carries `value`")
217    }
218
219    #[test]
220    fn raw_string_round_trips_in_raw_shape() {
221        let (row, shape) = kv_body_to_row(b"second-longer-value").unwrap();
222        assert_eq!(shape, KvBodyShape::Raw);
223        assert_eq!(value_of(&row), &Value::String("second-longer-value".into()));
224        assert_eq!(
225            row_to_kv_body(&row, shape).unwrap(),
226            b"second-longer-value".to_vec()
227        );
228    }
229
230    #[test]
231    fn raw_single_byte_value_is_a_string_not_a_fixint() {
232        // `b"1"` is 0x31, a valid msgpack fixint. It must still read as the
233        // string "1" and write back as the byte 0x31.
234        let (row, shape) = kv_body_to_row(b"1").unwrap();
235        assert_eq!(shape, KvBodyShape::Raw);
236        assert_eq!(value_of(&row), &Value::String("1".into()));
237        assert_eq!(row_to_kv_body(&row, shape).unwrap(), b"1".to_vec());
238    }
239
240    #[test]
241    fn raw_integer_value_encodes_as_decimal_text() {
242        let mut row = HashMap::new();
243        row.insert("value".to_string(), Value::Integer(42));
244        assert_eq!(
245            row_to_kv_body(&Value::Object(row), KvBodyShape::Raw).unwrap(),
246            b"42".to_vec()
247        );
248    }
249
250    #[test]
251    fn map_body_round_trips_in_map_shape() {
252        let mut fields = HashMap::new();
253        fields.insert("n".to_string(), Value::Integer(7));
254        let body = nodedb_types::value_to_msgpack(&Value::Object(fields.clone())).unwrap();
255
256        let (row, shape) = kv_body_to_row(&body).unwrap();
257        assert_eq!(shape, KvBodyShape::Map);
258        assert_eq!(row, Value::Object(fields.clone()));
259
260        let out = row_to_kv_body(&row, shape).unwrap();
261        assert_eq!(
262            nodedb_types::value_from_msgpack(&out).unwrap(),
263            Value::Object(fields)
264        );
265    }
266
267    #[test]
268    fn map_body_encodes_fields_in_key_order() {
269        let mut fields = HashMap::new();
270        fields.insert("b".to_string(), Value::Integer(2));
271        fields.insert("a".to_string(), Value::Integer(1));
272        let out = row_to_kv_body(&Value::Object(fields), KvBodyShape::Map).unwrap();
273        assert_eq!(out[0], 0x82, "fixmap of two entries");
274        let key_at = |name: u8| out.windows(2).position(|w| w == [0xa1, name]).unwrap();
275        assert!(
276            key_at(b'a') < key_at(b'b'),
277            "keys must be written in sorted order"
278        );
279    }
280
281    #[test]
282    fn raw_shape_with_extra_keys_is_an_error_naming_them() {
283        let mut row = HashMap::new();
284        row.insert("value".to_string(), Value::String("x".into()));
285        row.insert("n".to_string(), Value::Integer(1));
286        row.insert("a".to_string(), Value::Integer(2));
287        let err = row_to_kv_body(&Value::Object(row), KvBodyShape::Raw).unwrap_err();
288        match err {
289            KvBodyError::RawExtraKeys { keys } => assert_eq!(keys, "a, n"),
290            other => panic!("expected RawExtraKeys, got {other:?}"),
291        }
292    }
293
294    #[test]
295    fn raw_shape_with_non_scalar_value_is_an_error() {
296        let mut row = HashMap::new();
297        row.insert("value".to_string(), Value::Array(vec![Value::Integer(1)]));
298        let err = row_to_kv_body(&Value::Object(row), KvBodyShape::Raw).unwrap_err();
299        assert!(matches!(err, KvBodyError::RawNotScalar(_)), "{err:?}");
300    }
301
302    #[test]
303    fn non_object_row_is_an_error_in_either_shape() {
304        for shape in [KvBodyShape::Raw, KvBodyShape::Map] {
305            let err = row_to_kv_body(&Value::Integer(1), shape).unwrap_err();
306            assert!(matches!(err, KvBodyError::RowNotObject { kind: "int" }));
307        }
308    }
309
310    #[test]
311    fn empty_body_is_a_raw_empty_string() {
312        let (row, shape) = kv_body_to_row(b"").unwrap();
313        assert_eq!(shape, KvBodyShape::Raw);
314        assert_eq!(value_of(&row), &Value::String(String::new()));
315        assert!(row_to_kv_body(&row, shape).unwrap().is_empty());
316    }
317
318    #[test]
319    fn corrupt_map_body_is_a_decode_error() {
320        // fixmap header claiming one entry, then nothing.
321        assert!(kv_body_to_row(&[0x81]).is_err());
322    }
323}
nodedb_query/msgpack_scan/kv_body.rs

nodedb_query/msgpack_scan/
kv_body.rs
