1pub fn resolve_vars(
14 input: &str,
15 params: Option<&std::collections::HashMap<String, String>>,
16) -> crate::error::Result<String> {
17 let mut result = input.to_string();
18 let mut search_from = 0;
19 while let Some(rel_start) = result[search_from..].find("${") {
20 let start = search_from + rel_start;
21 let Some(rel_end) = result[start..].find('}') else {
22 break;
23 };
24 let end = start + rel_end;
25 let var_name = &result[start + 2..end];
26
27 let value = if var_name.is_empty() {
28 String::new()
31 } else if let Some(v) = params.and_then(|p| p.get(var_name)) {
32 v.clone()
33 } else {
34 match std::env::var(var_name) {
35 Ok(v) => v,
36 Err(_) => anyhow::bail!(
37 "environment variable '{}' referenced in config is not set \
38 (a missing secret silently becomes an empty string — refusing)",
39 var_name
40 ),
41 }
42 };
43
44 result = format!("{}{}{}", &result[..start], value, &result[end + 1..]);
45 search_from = start + value.len();
46 }
47 Ok(result)
48}
49
50pub fn resolve_env_vars(input: &str) -> crate::error::Result<String> {
52 resolve_vars(input, None)
53}
54
55pub fn find_unused_params(
61 haystack: &str,
62 params: Option<&std::collections::HashMap<String, String>>,
63) -> Vec<String> {
64 let Some(p) = params else {
65 return Vec::new();
66 };
67 let mut unused: Vec<String> = p
68 .keys()
69 .filter(|k| !haystack.contains(&format!("${{{k}}}")))
70 .cloned()
71 .collect();
72 unused.sort();
73 unused
74}
75
76pub fn warn_unused_params(
87 haystack: &str,
88 params: Option<&std::collections::HashMap<String, String>>,
89) {
90 for key in find_unused_params(haystack, params) {
91 log::warn!(
92 "--param '{}' was not referenced by any `${{{}}}` placeholder in the config — \
93 check the parameter name (case-sensitive) or remove the unused --param",
94 key,
95 key
96 );
97 }
98}
99
100pub fn parse_file_size(s: &str) -> crate::error::Result<u64> {
102 let s = s.trim().to_uppercase();
103 let (num, multiplier) = if let Some(n) = s.strip_suffix("GB") {
104 (n.trim(), 1024u64 * 1024 * 1024)
105 } else if let Some(n) = s.strip_suffix("MB") {
106 (n.trim(), 1024u64 * 1024)
107 } else if let Some(n) = s.strip_suffix("KB") {
108 (n.trim(), 1024u64)
109 } else if let Some(n) = s.strip_suffix('B') {
110 (n.trim(), 1u64)
111 } else {
112 (s.as_str(), 1u64)
113 };
114 let value: f64 = num
115 .parse()
116 .map_err(|_| anyhow::anyhow!("invalid file size: '{}'", s))?;
117 Ok((value * multiplier as f64) as u64)
118}
119
120#[cfg(test)]
121mod tests {
122 use super::*;
123 use std::collections::HashMap;
124
125 #[test]
128 fn no_placeholders_returned_verbatim() {
129 assert_eq!(resolve_vars("SELECT 1", None).unwrap(), "SELECT 1");
130 }
131
132 #[test]
133 fn empty_string_returned_verbatim() {
134 assert_eq!(resolve_vars("", None).unwrap(), "");
135 }
136
137 #[test]
140 fn param_substitutes_placeholder() {
141 let mut p = HashMap::new();
142 p.insert("TABLE".into(), "orders".into());
143 let result = resolve_vars("SELECT * FROM ${TABLE}", Some(&p)).unwrap();
144 assert_eq!(result, "SELECT * FROM orders");
145 }
146
147 #[test]
148 fn param_takes_precedence_over_env() {
149 unsafe { std::env::set_var("RIVET_TEST_OVERRIDE_VAR", "from_env") };
151 let mut p = HashMap::new();
152 p.insert("RIVET_TEST_OVERRIDE_VAR".into(), "from_param".into());
153 let result = resolve_vars("${RIVET_TEST_OVERRIDE_VAR}", Some(&p)).unwrap();
154 unsafe { std::env::remove_var("RIVET_TEST_OVERRIDE_VAR") };
155 assert_eq!(result, "from_param");
156 }
157
158 #[test]
159 fn multiple_placeholders_all_substituted() {
160 let mut p = HashMap::new();
161 p.insert("A".into(), "hello".into());
162 p.insert("B".into(), "world".into());
163 let result = resolve_vars("${A} ${B}", Some(&p)).unwrap();
164 assert_eq!(result, "hello world");
165 }
166
167 #[test]
170 fn env_var_substituted_when_set() {
171 unsafe { std::env::set_var("RIVET_TEST_RESOLVE_VAR", "secret123") };
172 let result = resolve_vars("pass=${RIVET_TEST_RESOLVE_VAR}", None).unwrap();
173 unsafe { std::env::remove_var("RIVET_TEST_RESOLVE_VAR") };
174 assert_eq!(result, "pass=secret123");
175 }
176
177 #[test]
178 fn missing_env_var_returns_error() {
179 unsafe { std::env::remove_var("RIVET_DEFINITELY_NOT_SET_VAR_XYZ") };
180 let err = resolve_vars("${RIVET_DEFINITELY_NOT_SET_VAR_XYZ}", None).unwrap_err();
181 let msg = err.to_string();
182 assert!(
183 msg.contains("RIVET_DEFINITELY_NOT_SET_VAR_XYZ"),
184 "got: {msg}"
185 );
186 }
187
188 #[test]
191 fn empty_placeholder_expands_to_empty_string() {
192 let result = resolve_vars("pre${}post", None).unwrap();
193 assert_eq!(result, "prepost");
194 }
195
196 #[test]
199 fn unclosed_placeholder_left_as_is() {
200 let result = resolve_vars("${UNCLOSED", None).unwrap();
201 assert_eq!(result, "${UNCLOSED");
202 }
203
204 #[test]
215 fn find_unused_params_returns_empty_when_no_params() {
216 assert!(find_unused_params("SELECT 1", None).is_empty());
217 }
218
219 #[test]
220 fn find_unused_params_used_key_not_flagged() {
221 let mut p = HashMap::new();
222 p.insert("max_id".into(), "20".into());
223 let unused = find_unused_params("SELECT * FROM t WHERE id <= ${max_id}", Some(&p));
224 assert!(unused.is_empty(), "got: {unused:?}");
225 }
226
227 #[test]
228 fn find_unused_params_unused_key_flagged_once() {
229 let mut p = HashMap::new();
230 p.insert("typo_id".into(), "20".into());
231 let unused = find_unused_params("SELECT * FROM t WHERE id <= ${max_id}", Some(&p));
232 assert_eq!(unused, vec!["typo_id".to_string()]);
233 }
234
235 #[test]
236 fn find_unused_params_mixed_used_and_unused() {
237 let mut p = HashMap::new();
238 p.insert("col".into(), "id".into());
239 p.insert("typo".into(), "x".into());
240 let unused = find_unused_params("SELECT ${col} FROM t", Some(&p));
241 assert_eq!(unused, vec!["typo".to_string()]);
242 }
243
244 #[test]
245 fn find_unused_params_partial_match_does_not_count() {
246 let mut p = HashMap::new();
249 p.insert("max".into(), "20".into());
250 let unused = find_unused_params("SELECT * FROM t WHERE id <= ${max_id}", Some(&p));
251 assert_eq!(unused, vec!["max".to_string()]);
252 }
253
254 #[test]
257 fn resolve_env_vars_reads_env() {
258 unsafe { std::env::set_var("RIVET_TEST_ENV_WRAPPER", "wrapped") };
259 let result = resolve_env_vars("v=${RIVET_TEST_ENV_WRAPPER}").unwrap();
260 unsafe { std::env::remove_var("RIVET_TEST_ENV_WRAPPER") };
261 assert_eq!(result, "v=wrapped");
262 }
263
264 #[test]
267 fn parse_1gb() {
268 assert_eq!(parse_file_size("1GB").unwrap(), 1024 * 1024 * 1024);
269 }
270
271 #[test]
272 fn parse_512mb() {
273 assert_eq!(parse_file_size("512MB").unwrap(), 512 * 1024 * 1024);
274 }
275
276 #[test]
277 fn parse_100kb() {
278 assert_eq!(parse_file_size("100KB").unwrap(), 100 * 1024);
279 }
280
281 #[test]
282 fn parse_bytes_suffix() {
283 assert_eq!(parse_file_size("2048B").unwrap(), 2048);
284 }
285
286 #[test]
287 fn parse_no_suffix_treated_as_bytes() {
288 assert_eq!(parse_file_size("4096").unwrap(), 4096);
289 }
290
291 #[test]
292 fn parse_whitespace_trimmed() {
293 assert_eq!(parse_file_size(" 256MB ").unwrap(), 256 * 1024 * 1024);
294 }
295
296 #[test]
297 fn parse_lowercase_accepted() {
298 assert_eq!(parse_file_size("1gb").unwrap(), 1024 * 1024 * 1024);
299 }
300
301 #[test]
302 fn parse_invalid_returns_error() {
303 assert!(parse_file_size("notanumber").is_err());
304 }
305}