risedev/config/
provide_expander.rs

Help
1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16
17use anyhow::{Result, anyhow};
18use itertools::Itertools;
19use yaml_rust::Yaml;
20
21/// Expands `provide-xxx: ["a", "b", "c"]` to their corresponding IDs.
22pub struct ProvideExpander {
23    all_items: HashMap<String, Yaml>,
24}
25
26impl ProvideExpander {
27    pub fn new(y: &Yaml) -> Result<Self> {
28        let y = y.as_vec().ok_or_else(|| anyhow!("expect an array"))?;
29        let mut all_items = HashMap::new();
30        for v in y {
31            let v = v.as_hash().ok_or_else(|| anyhow!("expect a hashmap"))?;
32            let id = v
33                .get(&Yaml::String("id".into()))
34                .ok_or_else(|| anyhow!("missing id field"))?;
35            let id = id
36                .as_str()
37                .ok_or_else(|| anyhow!("expect id to be a string"))?;
38            all_items.insert(id.to_owned(), Yaml::Hash(v.clone()));
39        }
40        Ok(Self {
41            all_items: Self::remove_provide(all_items)?,
42        })
43    }
44
45    fn remove_provide(all_items: HashMap<String, Yaml>) -> Result<HashMap<String, Yaml>> {
46        let all_items = all_items.into_iter().map(|(k, v)| {
47            let v = v.into_hash().ok_or_else(|| anyhow!("expect a hashmap"))?;
48            let v = v
49                .into_iter()
50                .filter(|(k, _)| {
51                    if let Some(k) = k.as_str() {
52                        !k.starts_with("provide-")
53                    } else {
54                        true
55                    }
56                })
57                .collect();
58            Ok::<_, anyhow::Error>((k, Yaml::Hash(v)))
59        });
60        all_items.try_collect()
61    }
62
63    pub fn visit(&mut self, yaml: Yaml) -> Result<Yaml> {
64        let yaml = yaml
65            .into_vec()
66            .ok_or_else(|| anyhow!("expect an array"))?
67            .into_iter()
68            .map(|yaml| {
69                let map = yaml
70                    .into_hash()
71                    .ok_or_else(|| anyhow!("expect a hashmap"))?;
72                let map = map.into_iter().map(|(k, v)| {
73                    if let Some(k) = k.as_str()
74                        && k.starts_with("provide-")
75                    {
76                        let array = v
77                            .as_vec()
78                            .ok_or_else(|| anyhow!("expect an array of provide-"))?;
79                        let array = array.iter().map(|item| {
80                            let item = item
81                                .as_str()
82                                .ok_or_else(|| anyhow!("expect a string from provide"))?;
83                            Ok::<_, anyhow::Error>(
84                                self.all_items
85                                    .get(item)
86                                    .ok_or_else(|| anyhow!("{} not found", item))?
87                                    .clone(),
88                            )
89                        });
90                        return Ok::<_, anyhow::Error>((
91                            Yaml::String(k.to_owned()),
92                            Yaml::Array(array.try_collect()?),
93                        ));
94                    }
95                    Ok::<_, anyhow::Error>((k, v))
96                });
97                Ok::<_, anyhow::Error>(Yaml::Hash(map.try_collect()?))
98            });
99        Ok(Yaml::Array(yaml.try_collect()?))
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use yaml_rust::YamlLoader;
106
107    use super::*;
108
109    #[test]
110    fn test_expand_provide() {
111        let source = YamlLoader::load_from_str(
112            "
113- provide-b: [\"b\"]
114  test_field: a
115  id: a
116- provide-a: [\"a\"]
117  test_field: a
118  id: b
119      ",
120        )
121        .unwrap()
122        .remove(0);
123
124        let expected_result = YamlLoader::load_from_str(
125            "
126- provide-b:
127    - test_field: a
128      id: b
129  test_field: a
130  id: a
131- provide-a:
132    - test_field: a
133      id: a
134  test_field: a
135  id: b
136  ",
137        )
138        .unwrap()
139        .remove(0);
140
141        let mut visitor = ProvideExpander::new(&source).unwrap();
142
143        assert_eq!(visitor.visit(source).unwrap(), expected_result);
144    }
145}
risedev/config/provide_expander.rs

risedev/config/
provide_expander.rs
