Skip to main content

ytsaurus_client/
yson_build.rs

1//! Small constructors for the YSON documents the API expects.
2//!
3//! Command parameters and operation specs are YSON dicts. `YsonValue` can model
4//! all of them, but building one by hand is verbose — its map keys are `Vec<u8>`
5//! and every leaf needs wrapping. These helpers keep the call sites readable.
6//!
7//! The client encodes parameters with this project's own codec rather than
8//! reaching for JSON, which keeps the dependency list short and exercises
9//! `ytsaurus-yson` against the real cluster on every request.
10
11use std::collections::BTreeMap;
12
13use ytsaurus_yson::{YsonNode, YsonValue};
14
15/// A YSON string.
16#[must_use]
17pub fn string(value: impl AsRef<[u8]>) -> YsonValue {
18    YsonValue {
19        attributes: None,
20        node: YsonNode::String(value.as_ref().to_vec()),
21    }
22}
23
24/// A YSON int64.
25#[must_use]
26pub fn int(value: i64) -> YsonValue {
27    YsonValue {
28        attributes: None,
29        node: YsonNode::Int64(value),
30    }
31}
32
33/// A YSON uint64.
34///
35/// A different YSON type from [`int`], but on a `uint64` key column the
36/// cluster does not hold that against a value it can read: measured,
37/// `{exact={key=[42]}}` and `{exact={key=[42u]}}` returned the same row, so
38/// [`Key::from`](crate::Key)`(42_i64)` finds it. The difference that bites is
39/// *range*. `Key::from(i64)` tops out at `i64::MAX`, so every key above that
40/// is one it cannot spell, and such a row is reachable only through this
41/// helper — `Key::new([yson_build::uint(u64::MAX)])` returned the row keyed
42/// `18446744073709551615u`.
43#[must_use]
44pub fn uint(value: u64) -> YsonValue {
45    YsonValue {
46        attributes: None,
47        node: YsonNode::Uint64(value),
48    }
49}
50
51/// A YSON double.
52///
53/// The type a scheduler weight has: `update_operation_parameters` takes
54/// `weight=2.5`, and an int64 in its place is a different YSON value.
55#[must_use]
56pub fn double(value: f64) -> YsonValue {
57    YsonValue {
58        attributes: None,
59        node: YsonNode::Double(value),
60    }
61}
62
63/// A YSON boolean.
64#[must_use]
65pub fn boolean(value: bool) -> YsonValue {
66    YsonValue {
67        attributes: None,
68        node: YsonNode::Boolean(value),
69    }
70}
71
72/// A YSON list.
73#[must_use]
74pub fn list(items: impl IntoIterator<Item = YsonValue>) -> YsonValue {
75    YsonValue {
76        attributes: None,
77        node: YsonNode::List(items.into_iter().collect()),
78    }
79}
80
81/// A YSON dict.
82#[must_use]
83pub fn map<K: AsRef<[u8]>>(entries: impl IntoIterator<Item = (K, YsonValue)>) -> YsonValue {
84    let mut out = BTreeMap::new();
85    for (key, value) in entries {
86        out.insert(key.as_ref().to_vec(), value);
87    }
88    YsonValue {
89        attributes: None,
90        node: YsonNode::Map(out),
91    }
92}
93
94/// An empty YSON dict — the parameters of a command that takes none.
95///
96/// `map([])` cannot express this: the key type has nothing to be inferred from,
97/// and `map` takes its entries as an `impl Trait` argument, so a turbofish is
98/// not allowed either. A command like `get_supported_features` still has to
99/// send `{}` in `X-YT-Parameters`, so the shorthand exists.
100#[must_use]
101pub fn empty_map() -> YsonValue {
102    YsonValue {
103        attributes: None,
104        node: YsonNode::Map(BTreeMap::new()),
105    }
106}
107
108/// Attaches attributes to a value, as in `<format=binary>yson`.
109#[must_use]
110pub fn with_attributes<K: AsRef<[u8]>>(
111    value: YsonValue,
112    attributes: impl IntoIterator<Item = (K, YsonValue)>,
113) -> YsonValue {
114    let mut attrs = BTreeMap::new();
115    for (key, v) in attributes {
116        attrs.insert(key.as_ref().to_vec(), v);
117    }
118    YsonValue {
119        attributes: if attrs.is_empty() { None } else { Some(attrs) },
120        node: value.node,
121    }
122}
123
124/// `<format=binary>yson` — the format a `ytsaurus-job` worker expects.
125#[must_use]
126pub fn binary_yson_format() -> YsonValue {
127    with_attributes(string("yson"), [("format", string("binary"))])
128}
129
130/// Inserts into a value that is known to be a dict; panics otherwise.
131///
132/// A `Result` here would be noise at every call site, so the invariant is kept
133/// by the callers instead: a builder reading back a value a caller supplied
134/// through `with_raw` passes it through `map_or_empty` first, and the raw
135/// command doors refuse parameters that are not a dict before anything is
136/// stamped onto them. Reach for one of those rather than widening this.
137pub(crate) fn insert(target: &mut YsonValue, key: impl AsRef<[u8]>, value: YsonValue) {
138    match &mut target.node {
139        YsonNode::Map(m) => {
140            m.insert(key.as_ref().to_vec(), value);
141        }
142        other => panic!("expected a dict, got {other:?}"),
143    }
144}
145
146#[cfg(test)]
147mod tests {
148    use super::*;
149    use ytsaurus_yson::{YsonFormat, to_string};
150
151    #[test]
152    fn no_parameters_still_encodes_as_a_dict() {
153        // `X-YT-Parameters` is a YSON dict on every command, including the ones
154        // that take nothing. An empty *list* or a missing header would be a
155        // different statement to the proxy.
156        assert_eq!(
157            to_string(&empty_map(), YsonFormat::Text).expect("encodes"),
158            "{}"
159        );
160    }
161
162    #[test]
163    fn builds_the_documents_the_api_expects() {
164        let spec = map([
165            ("input_table_paths", list([string("//tmp/in")])),
166            ("output_table_paths", list([string("//tmp/out")])),
167            (
168                "mapper",
169                map([
170                    ("command", string("./worker")),
171                    ("memory_limit", int(536_870_912)),
172                ]),
173            ),
174        ]);
175
176        let encoded = to_string(&spec, YsonFormat::Text).expect("encodes");
177        assert!(
178            encoded.contains("input_table_paths=[\"//tmp/in\"]"),
179            "{encoded}"
180        );
181        assert!(encoded.contains("memory_limit=536870912"), "{encoded}");
182    }
183
184    #[test]
185    fn format_attributes_render_as_yson_expects() {
186        let encoded = to_string(&binary_yson_format(), YsonFormat::Text).expect("encodes");
187        assert_eq!(encoded, "<format=binary>yson");
188    }
189
190    #[test]
191    fn booleans_use_the_yson_spelling() {
192        let encoded = to_string(&map([("enable", boolean(true))]), YsonFormat::Text).unwrap();
193        assert_eq!(encoded, "{enable=%true}");
194    }
195
196    #[test]
197    fn an_unsigned_integer_is_not_the_same_value_as_a_signed_one() {
198        // YSON writes a uint64 with a `u` suffix, and these are two distinct
199        // values in this crate's own model.
200        let encoded = to_string(&map([("n", uint(42))]), YsonFormat::Text).unwrap();
201        assert_eq!(encoded, "{n=42u}");
202        assert_ne!(uint(42), int(42));
203        // The cluster, though, does not make the caller pick: measured on a
204        // `uint64`-keyed table, `{exact={key=[42]}}` and `{exact={key=[42u]}}`
205        // both returned the row. What only this helper can do is name a key
206        // above `i64::MAX` at all — that half of `u64` is unreachable from
207        // `Key::from(i64)`, and the row keyed `18446744073709551615u` came
208        // back for exactly this spelling.
209        let encoded = to_string(&map([("n", uint(u64::MAX))]), YsonFormat::Text).unwrap();
210        assert_eq!(encoded, "{n=18446744073709551615u}");
211        assert!(i64::try_from(u64::MAX).is_err());
212    }
213}