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 boolean.
34#[must_use]
35pub fn boolean(value: bool) -> YsonValue {
36    YsonValue {
37        attributes: None,
38        node: YsonNode::Boolean(value),
39    }
40}
41
42/// A YSON list.
43#[must_use]
44pub fn list(items: impl IntoIterator<Item = YsonValue>) -> YsonValue {
45    YsonValue {
46        attributes: None,
47        node: YsonNode::List(items.into_iter().collect()),
48    }
49}
50
51/// A YSON dict.
52#[must_use]
53pub fn map<K: AsRef<[u8]>>(entries: impl IntoIterator<Item = (K, YsonValue)>) -> YsonValue {
54    let mut out = BTreeMap::new();
55    for (key, value) in entries {
56        out.insert(key.as_ref().to_vec(), value);
57    }
58    YsonValue {
59        attributes: None,
60        node: YsonNode::Map(out),
61    }
62}
63
64/// Attaches attributes to a value, as in `<format=binary>yson`.
65#[must_use]
66pub fn with_attributes<K: AsRef<[u8]>>(
67    value: YsonValue,
68    attributes: impl IntoIterator<Item = (K, YsonValue)>,
69) -> YsonValue {
70    let mut attrs = BTreeMap::new();
71    for (key, v) in attributes {
72        attrs.insert(key.as_ref().to_vec(), v);
73    }
74    YsonValue {
75        attributes: if attrs.is_empty() { None } else { Some(attrs) },
76        node: value.node,
77    }
78}
79
80/// `<format=binary>yson` — the format a `ytsaurus-job` worker expects.
81#[must_use]
82pub fn binary_yson_format() -> YsonValue {
83    with_attributes(string("yson"), [("format", string("binary"))])
84}
85
86/// Inserts into a value that is known to be a dict; panics otherwise.
87///
88/// Only used on values this crate built, so the panic is unreachable in
89/// practice and a `Result` here would be noise at every call site.
90pub(crate) fn insert(target: &mut YsonValue, key: impl AsRef<[u8]>, value: YsonValue) {
91    match &mut target.node {
92        YsonNode::Map(m) => {
93            m.insert(key.as_ref().to_vec(), value);
94        }
95        other => panic!("expected a dict, got {other:?}"),
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102    use ytsaurus_yson::{YsonFormat, to_string};
103
104    #[test]
105    fn builds_the_documents_the_api_expects() {
106        let spec = map([
107            ("input_table_paths", list([string("//tmp/in")])),
108            ("output_table_paths", list([string("//tmp/out")])),
109            (
110                "mapper",
111                map([
112                    ("command", string("./worker")),
113                    ("memory_limit", int(536_870_912)),
114                ]),
115            ),
116        ]);
117
118        let encoded = to_string(&spec, YsonFormat::Text).expect("encodes");
119        assert!(
120            encoded.contains("input_table_paths=[\"//tmp/in\"]"),
121            "{encoded}"
122        );
123        assert!(encoded.contains("memory_limit=536870912"), "{encoded}");
124    }
125
126    #[test]
127    fn format_attributes_render_as_yson_expects() {
128        let encoded = to_string(&binary_yson_format(), YsonFormat::Text).expect("encodes");
129        assert_eq!(encoded, "<format=binary>yson");
130    }
131
132    #[test]
133    fn booleans_use_the_yson_spelling() {
134        let encoded = to_string(&map([("enable", boolean(true))]), YsonFormat::Text).unwrap();
135        assert_eq!(encoded, "{enable=%true}");
136    }
137}