nimbus_cli/
value_utils.rs

Help
1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at https://mozilla.org/MPL/2.0/.
4
5use anyhow::Result;
6use serde::{Deserialize, Serialize};
7use serde_json::{Map, Value};
8use std::collections::HashMap;
9use std::path::Path;
10
11use crate::NimbusApp;
12
13pub(crate) trait CliUtils {
14    fn get_str<'a>(&'a self, key: &str) -> Result<&'a str>;
15    fn get_bool(&self, key: &str) -> Result<bool>;
16    fn get_array<'a>(&'a self, key: &str) -> Result<&'a Vec<Value>>;
17    fn get_mut_array<'a>(&'a mut self, key: &str) -> Result<&'a mut Vec<Value>>;
18    fn get_mut_object<'a>(&'a mut self, key: &str) -> Result<&'a mut Value>;
19    fn get_object<'a>(&'a self, key: &str) -> Result<&'a Value>;
20    fn get_u64(&self, key: &str) -> Result<u64>;
21
22    fn has(&self, key: &str) -> bool;
23    fn set<V>(&mut self, key: &str, value: V) -> Result<()>
24    where
25        V: Serialize;
26}
27
28impl CliUtils for Value {
29    fn get_str<'a>(&'a self, key: &str) -> Result<&'a str> {
30        let v = self
31            .get(key)
32            .ok_or_else(|| {
33                anyhow::Error::msg(format!(
34                    "Expected a string with key '{key}' in the JSONObject"
35                ))
36            })?
37            .as_str()
38            .ok_or_else(|| anyhow::Error::msg("value is not a string"))?;
39
40        Ok(v)
41    }
42
43    fn get_bool(&self, key: &str) -> Result<bool> {
44        let v = self
45            .get(key)
46            .ok_or_else(|| {
47                anyhow::Error::msg(format!(
48                    "Expected a string with key '{key}' in the JSONObject"
49                ))
50            })?
51            .as_bool()
52            .ok_or_else(|| anyhow::Error::msg("value is not a string"))?;
53
54        Ok(v)
55    }
56
57    fn get_array<'a>(&'a self, key: &str) -> Result<&'a Vec<Value>> {
58        let v = self
59            .get(key)
60            .ok_or_else(|| {
61                anyhow::Error::msg(format!(
62                    "Expected an array with key '{key}' in the JSONObject"
63                ))
64            })?
65            .as_array()
66            .ok_or_else(|| anyhow::Error::msg("value is not a array"))?;
67        Ok(v)
68    }
69
70    fn get_mut_array<'a>(&'a mut self, key: &str) -> Result<&'a mut Vec<Value>> {
71        let v = self
72            .get_mut(key)
73            .ok_or_else(|| {
74                anyhow::Error::msg(format!(
75                    "Expected an array with key '{key}' in the JSONObject"
76                ))
77            })?
78            .as_array_mut()
79            .ok_or_else(|| anyhow::Error::msg("value is not a array"))?;
80        Ok(v)
81    }
82
83    fn get_object<'a>(&'a self, key: &str) -> Result<&'a Value> {
84        let v = self.get(key).ok_or_else(|| {
85            anyhow::Error::msg(format!(
86                "Expected an object with key '{key}' in the JSONObject"
87            ))
88        })?;
89        Ok(v)
90    }
91
92    fn get_mut_object<'a>(&'a mut self, key: &str) -> Result<&'a mut Value> {
93        let v = self.get_mut(key).ok_or_else(|| {
94            anyhow::Error::msg(format!(
95                "Expected an object with key '{key}' in the JSONObject"
96            ))
97        })?;
98        Ok(v)
99    }
100
101    fn get_u64(&self, key: &str) -> Result<u64> {
102        let v = self
103            .get(key)
104            .ok_or_else(|| {
105                anyhow::Error::msg(format!(
106                    "Expected an array with key '{key}' in the JSONObject"
107                ))
108            })?
109            .as_u64()
110            .ok_or_else(|| anyhow::Error::msg("value is not a array"))?;
111        Ok(v)
112    }
113
114    fn set<V>(&mut self, key: &str, value: V) -> Result<()>
115    where
116        V: Serialize,
117    {
118        let value = serde_json::to_value(value)?;
119        match self.as_object_mut() {
120            Some(m) => m.insert(key.to_string(), value),
121            _ => anyhow::bail!("Can only insert into JSONObjects"),
122        };
123        Ok(())
124    }
125
126    fn has(&self, key: &str) -> bool {
127        self.get(key).is_some()
128    }
129}
130
131pub(crate) fn try_find_experiment(value: &Value, slug: &str) -> Result<Value> {
132    let array = try_extract_data_list(value)?;
133    let exp = array
134        .iter()
135        .find(|exp| {
136            if let Some(Value::String(s)) = exp.get("slug") {
137                slug == s
138            } else {
139                false
140            }
141        })
142        .ok_or_else(|| anyhow::Error::msg(format!("No experiment with slug {}", slug)))?;
143
144    Ok(exp.clone())
145}
146
147pub(crate) fn try_extract_data_list(value: &Value) -> Result<Vec<Value>> {
148    assert!(value.is_object());
149    Ok(value.get_array("data")?.to_vec())
150}
151
152pub(crate) fn try_find_branches_from_experiment(value: &Value) -> Result<Vec<Value>> {
153    Ok(value.get_array("branches")?.to_vec())
154}
155
156pub(crate) fn try_find_features_from_branch(value: &Value) -> Result<Vec<Value>> {
157    let features = value.get_array("features");
158    Ok(if features.is_ok() {
159        features?.to_vec()
160    } else {
161        let feature = value
162            .get("feature")
163            .expect("Expected a feature or features in a branch");
164        vec![feature.clone()]
165    })
166}
167
168pub(crate) fn try_find_mut_features_from_branch<'a>(
169    value: &'a mut Value,
170) -> Result<HashMap<String, &'a mut Value>> {
171    let mut res = HashMap::new();
172    if value.has("features") {
173        let features = value.get_mut_array("features")?;
174        for f in features {
175            res.insert(
176                f.get_str("featureId")?.to_string(),
177                f.get_mut_object("value")?,
178            );
179        }
180    } else {
181        let f: &'a mut Value = value.get_mut_object("feature")?;
182        res.insert(
183            f.get_str("featureId")?.to_string(),
184            f.get_mut_object("value")?,
185        );
186    }
187    Ok(res)
188}
189
190pub(crate) trait Patch {
191    fn patch(&mut self, patch: &Self) -> bool;
192}
193
194impl Patch for Value {
195    fn patch(&mut self, patch: &Self) -> bool {
196        match (self, patch) {
197            (Value::Object(t), Value::Object(p)) => {
198                t.patch(p);
199            }
200            (Value::String(t), Value::String(p)) => t.clone_from(p),
201            (Value::Bool(t), Value::Bool(p)) => *t = *p,
202            (Value::Number(t), Value::Number(p)) => *t = p.clone(),
203            (Value::Array(t), Value::Array(p)) => t.clone_from(p),
204            (Value::Null, Value::Null) => (),
205            _ => return false,
206        };
207        true
208    }
209}
210
211impl Patch for Map<String, Value> {
212    fn patch(&mut self, patch: &Self) -> bool {
213        for (k, v) in patch {
214            match (self.get_mut(k), v) {
215                (Some(_), Value::Null) => {
216                    self.remove(k);
217                }
218                (_, Value::Null) => {
219                    // If the patch is null, then don't add it to this value.
220                }
221                (Some(t), p) => {
222                    if !t.patch(p) {
223                        println!("Warning: the patched key '{k}' has different types: {t} != {p}");
224                        self.insert(k.clone(), v.clone());
225                    }
226                }
227                (None, _) => {
228                    self.insert(k.clone(), v.clone());
229                }
230            }
231        }
232        true
233    }
234}
235
236fn prepare_recipe(
237    recipe: &Value,
238    params: &NimbusApp,
239    preserve_targeting: bool,
240    preserve_bucketing: bool,
241) -> Result<Value> {
242    let mut recipe = recipe.clone();
243    let slug = recipe.get_str("slug")?;
244    let app_name = params
245        .app_name
246        .as_deref()
247        .expect("An app name is expected. This is a bug in nimbus-cli");
248    if app_name != recipe.get_str("appName")? {
249        anyhow::bail!(format!("'{slug}' is not for {app_name} app"));
250    }
251    recipe.set("channel", &params.channel)?;
252    recipe.set("isEnrollmentPaused", false)?;
253    if !preserve_targeting {
254        recipe.set("targeting", "true")?;
255    }
256    if !preserve_bucketing {
257        let bucketing = recipe.get_mut_object("bucketConfig")?;
258        bucketing.set("start", 0)?;
259        bucketing.set("count", 10_000)?;
260    }
261    Ok(recipe)
262}
263
264pub(crate) fn prepare_rollout(
265    recipe: &Value,
266    params: &NimbusApp,
267    preserve_targeting: bool,
268    preserve_bucketing: bool,
269) -> Result<Value> {
270    let rollout = prepare_recipe(recipe, params, preserve_targeting, preserve_bucketing)?;
271    if !rollout.get_bool("isRollout")? {
272        let slug = rollout.get_str("slug")?;
273        anyhow::bail!(format!("Recipe '{}' isn't a rollout", slug));
274    }
275    Ok(rollout)
276}
277
278pub(crate) fn prepare_experiment(
279    recipe: &Value,
280    params: &NimbusApp,
281    branch: &str,
282    preserve_targeting: bool,
283    preserve_bucketing: bool,
284) -> Result<Value> {
285    let mut experiment = prepare_recipe(recipe, params, preserve_targeting, preserve_bucketing)?;
286
287    if !preserve_bucketing {
288        let branches = experiment.get_mut_array("branches")?;
289        let mut found = false;
290        for b in branches {
291            let slug = b.get_str("slug")?;
292            let ratio = if slug == branch {
293                found = true;
294                100
295            } else {
296                0
297            };
298            b.set("ratio", ratio)?;
299        }
300        if !found {
301            let slug = experiment.get_str("slug")?;
302            anyhow::bail!(format!(
303                "No branch called '{}' was found in '{}'",
304                branch, slug
305            ));
306        }
307    }
308    Ok(experiment)
309}
310
311fn is_yaml<P>(file: P) -> bool
312where
313    P: AsRef<Path>,
314{
315    let ext = file.as_ref().extension().unwrap_or_default();
316    ext == "yaml" || ext == "yml"
317}
318
319pub(crate) fn read_from_file<P, T>(file: P) -> Result<T>
320where
321    P: AsRef<Path>,
322    for<'a> T: Deserialize<'a>,
323{
324    let s = std::fs::read_to_string(&file)?;
325    Ok(if is_yaml(&file) {
326        serde_yaml::from_str(&s)?
327    } else {
328        serde_json::from_str(&s)?
329    })
330}
331
332pub(crate) fn write_to_file_or_print<P, T>(file: Option<P>, contents: &T) -> Result<()>
333where
334    P: AsRef<Path>,
335    T: Serialize,
336{
337    match file {
338        Some(file) => {
339            let s = if is_yaml(&file) {
340                serde_yaml::to_string(&contents)?
341            } else {
342                serde_json::to_string_pretty(&contents)?
343            };
344            std::fs::write(file, s)?;
345        }
346        _ => println!("{}", serde_json::to_string_pretty(&contents)?),
347    }
348
349    Ok(())
350}
351
352#[cfg(test)]
353mod tests {
354    use serde_json::json;
355
356    use super::*;
357
358    #[test]
359    fn test_find_experiment() -> Result<()> {
360        let exp = json!({
361            "slug": "a-name",
362        });
363        let source = json!({ "data": [exp] });
364
365        assert_eq!(try_find_experiment(&source, "a-name")?, exp);
366
367        let source = json!({
368            "data": {},
369        });
370        assert!(try_find_experiment(&source, "a-name").is_err());
371
372        let source = json!({
373            "data": [],
374        });
375        assert!(try_find_experiment(&source, "a-name").is_err());
376
377        Ok(())
378    }
379
380    #[test]
381    fn test_prepare_experiment() -> Result<()> {
382        let src = json!({
383            "appName": "an-app",
384            "slug": "a-name",
385            "branches": [
386                {
387                    "slug": "another-branch",
388                },
389                {
390                    "slug": "a-branch",
391                }
392            ],
393            "bucketConfig": {
394            }
395        });
396
397        let params = NimbusApp::new("an-app", "developer");
398
399        assert_eq!(
400            json!({
401                "appName": "an-app",
402                "channel": "developer",
403                "slug": "a-name",
404                "branches": [
405                    {
406                        "slug": "another-branch",
407                        "ratio": 0,
408                    },
409                    {
410                        "slug": "a-branch",
411                        "ratio": 100,
412                    }
413                ],
414                "bucketConfig": {
415                    "start": 0,
416                    "count": 10_000,
417                },
418                "isEnrollmentPaused": false,
419                "targeting": "true"
420            }),
421            prepare_experiment(&src, &params, "a-branch", false, false)?
422        );
423
424        assert_eq!(
425            json!({
426                "appName": "an-app",
427                "channel": "developer",
428                "slug": "a-name",
429                "branches": [
430                    {
431                        "slug": "another-branch",
432                    },
433                    {
434                        "slug": "a-branch",
435                    }
436                ],
437                "bucketConfig": {
438                },
439                "isEnrollmentPaused": false,
440                "targeting": "true"
441            }),
442            prepare_experiment(&src, &params, "a-branch", false, true)?
443        );
444
445        assert_eq!(
446            json!({
447                "appName": "an-app",
448                "channel": "developer",
449                "slug": "a-name",
450                "branches": [
451                    {
452                        "slug": "another-branch",
453                        "ratio": 0,
454                    },
455                    {
456                        "slug": "a-branch",
457                        "ratio": 100,
458                    }
459                ],
460                "bucketConfig": {
461                    "start": 0,
462                    "count": 10_000,
463                },
464                "isEnrollmentPaused": false,
465            }),
466            prepare_experiment(&src, &params, "a-branch", true, false)?
467        );
468
469        assert_eq!(
470            json!({
471                "appName": "an-app",
472                "slug": "a-name",
473                "channel": "developer",
474                "branches": [
475                    {
476                        "slug": "another-branch",
477                    },
478                    {
479                        "slug": "a-branch",
480                    }
481                ],
482                "bucketConfig": {
483                },
484                "isEnrollmentPaused": false,
485            }),
486            prepare_experiment(&src, &params, "a-branch", true, true)?
487        );
488        Ok(())
489    }
490
491    #[test]
492    fn test_patch_value() -> Result<()> {
493        let mut v1 = json!({
494            "string": "string",
495            "obj": {
496                "string": "string",
497                "num": 1,
498                "bool": false,
499            },
500            "num": 1,
501            "bool": false,
502        });
503        let ov1 = json!({
504            "string": "patched",
505            "obj": {
506                "string": "patched",
507            },
508            "num": 2,
509            "bool": true,
510        });
511
512        v1.patch(&ov1);
513
514        let expected = json!({
515            "string": "patched",
516            "obj": {
517                "string": "patched",
518                "num": 1,
519                "bool": false,
520            },
521            "num": 2,
522            "bool": true,
523        });
524
525        assert_eq!(&expected, &v1);
526
527        let mut v1 = json!({
528            "string": "string",
529            "obj": {
530                "string": "string",
531                "num": 1,
532                "bool": false,
533            },
534            "num": 1,
535            "bool": false,
536        });
537        let ov1 = json!({
538            "obj": null,
539            "never": null,
540        });
541        v1.patch(&ov1);
542        let expected = json!({
543            "string": "string",
544            "num": 1,
545            "bool": false,
546        });
547
548        assert_eq!(&expected, &v1);
549
550        Ok(())
551    }
552}
nimbus_cli/value_utils.rs

nimbus_cli/
value_utils.rs
