mz_testdrive/
parser.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use std::borrow::ToOwned;
11use std::collections::{BTreeMap, btree_map};
12use std::error::Error;
13use std::fmt::Write;
14use std::str::FromStr;
15use std::sync::LazyLock;
16
17use anyhow::{Context, anyhow, bail};
18use regex::Regex;
19
20use crate::error::PosError;
21
22#[derive(Debug, Clone)]
23pub struct PosCommand {
24    pub pos: usize,
25    pub command: Command,
26}
27
28// min and max versions, both inclusive
29#[derive(Debug, Clone)]
30pub struct VersionConstraint {
31    pub min: i32,
32    pub max: i32,
33}
34
35#[derive(Debug, Clone)]
36pub enum Command {
37    Builtin(BuiltinCommand, Option<VersionConstraint>),
38    Sql(SqlCommand, Option<VersionConstraint>),
39    FailSql(FailSqlCommand, Option<VersionConstraint>),
40}
41
42#[derive(Debug, Clone)]
43pub struct BuiltinCommand {
44    pub name: String,
45    pub args: ArgMap,
46    pub input: Vec<String>,
47}
48
49impl BuiltinCommand {
50    pub fn assert_no_input(&self) -> Result<(), anyhow::Error> {
51        if !self.input.is_empty() {
52            bail!("{} action does not take input", self.name);
53        }
54        Ok(())
55    }
56}
57
58#[derive(Debug, Clone)]
59pub enum SqlOutput {
60    Full {
61        column_names: Option<Vec<String>>,
62        expected_rows: Vec<Vec<String>>,
63    },
64    Hashed {
65        num_values: usize,
66        md5: String,
67    },
68}
69#[derive(Debug, Clone)]
70pub struct SqlCommand {
71    pub query: String,
72    pub expected_output: SqlOutput,
73    pub expected_start: usize,
74    pub expected_end: usize,
75}
76
77#[derive(Debug, Clone)]
78pub struct FailSqlCommand {
79    pub query: String,
80    pub expected_error: SqlExpectedError,
81    pub expected_detail: Option<String>,
82    pub expected_hint: Option<String>,
83}
84
85#[derive(Debug, Clone)]
86pub enum SqlExpectedError {
87    Contains(String),
88    Exact(String),
89    Regex(String),
90    Timeout,
91}
92
93pub(crate) fn parse(line_reader: &mut LineReader) -> Result<Vec<PosCommand>, PosError> {
94    let mut out = Vec::new();
95    while let Some((pos, line)) = line_reader.peek() {
96        let pos = *pos;
97        let command = match line.chars().next() {
98            Some('$') => {
99                let version = parse_version_constraint(line_reader)?;
100                Command::Builtin(parse_builtin(line_reader)?, version)
101            }
102            Some('>') => {
103                let version = parse_version_constraint(line_reader)?;
104                Command::Sql(parse_sql(line_reader)?, version)
105            }
106            Some('?') => {
107                let version = parse_version_constraint(line_reader)?;
108                Command::Sql(parse_explain_sql(line_reader)?, version)
109            }
110            Some('!') => {
111                let version = parse_version_constraint(line_reader)?;
112                Command::FailSql(parse_fail_sql(line_reader)?, version)
113            }
114            Some('#') => {
115                // Comment line.
116                line_reader.next();
117                continue;
118            }
119            Some(x) => {
120                return Err(PosError {
121                    source: anyhow!(format!("unexpected input line at beginning of file: {}", x)),
122                    pos: Some(pos),
123                });
124            }
125            None => {
126                return Err(PosError {
127                    source: anyhow!("unexpected input line at beginning of file"),
128                    pos: Some(pos),
129                });
130            }
131        };
132        out.push(PosCommand { command, pos });
133    }
134    Ok(out)
135}
136
137fn parse_builtin(line_reader: &mut LineReader) -> Result<BuiltinCommand, PosError> {
138    let (pos, line) = line_reader.next().unwrap();
139    let mut builtin_reader = BuiltinReader::new(&line, pos);
140    let name = match builtin_reader.next() {
141        Some(Ok((_, s))) => s,
142        Some(Err(e)) => return Err(e),
143        None => {
144            return Err(PosError {
145                source: anyhow!("command line is missing command name"),
146                pos: Some(pos),
147            });
148        }
149    };
150    let mut args = BTreeMap::new();
151    for el in builtin_reader {
152        let (pos, token) = el?;
153        let pieces: Vec<_> = token.splitn(2, '=').collect();
154        let pieces = match pieces.as_slice() {
155            [key, value] => vec![*key, *value],
156            [key] => vec![*key, ""],
157            _ => {
158                return Err(PosError {
159                    source: anyhow!("command argument is not in required key=value format"),
160                    pos: Some(pos),
161                });
162            }
163        };
164        validate_ident(pieces[0]).map_err(|e| PosError::new(e, pos))?;
165
166        if let Some(original) = args.insert(pieces[0].to_owned(), pieces[1].to_owned()) {
167            return Err(PosError {
168                source: anyhow!(
169                    "argument '{}' specified twice: {} & {}",
170                    pieces[0],
171                    original,
172                    pieces[1]
173                ),
174                pos: Some(pos),
175            });
176        };
177    }
178    Ok(BuiltinCommand {
179        name,
180        args: ArgMap(args),
181        input: slurp_all(line_reader),
182    })
183}
184
185/// Validate that the string is an allowed variable name (lowercase letters, numbers and dashes)
186pub fn validate_ident(name: &str) -> Result<(), anyhow::Error> {
187    static VALID_KEY_REGEX: LazyLock<Regex> =
188        LazyLock::new(|| Regex::new("^[a-z0-9\\-]*$").unwrap());
189    if !VALID_KEY_REGEX.is_match(name) {
190        bail!(
191            "invalid builtin argument name '{}': \
192             only lowercase letters, numbers, and hyphens allowed",
193            name
194        );
195    }
196    Ok(())
197}
198
199fn parse_version_constraint(
200    line_reader: &mut LineReader,
201) -> Result<Option<VersionConstraint>, PosError> {
202    let (pos, line) = line_reader.next().unwrap();
203    if line[1..2].to_string() != "[" {
204        line_reader.push(&line);
205        return Ok(None);
206    }
207    let closed_brace_pos = match line.find(']') {
208        Some(x) => x,
209        None => {
210            return Err(PosError {
211                source: anyhow!("version-constraint: found no closing brace"),
212                pos: Some(pos),
213            });
214        }
215    };
216    let mut begin_version_kw = 2;
217    const MIN_VERSION: i32 = 0;
218    let mut min_version = MIN_VERSION;
219    if line.as_bytes()[2].is_ascii_digit() {
220        let Some(op_pos) = line.find('<') else {
221            return Err(PosError {
222                source: anyhow!("version-constraint: initial number but no '<' following"),
223                pos: Some(pos),
224            });
225        };
226        let min_version_str = line[2..op_pos].to_string();
227        match min_version_str.parse::<i32>() {
228            Ok(mv) => min_version = mv,
229            Err(_) => {
230                return Err(PosError {
231                    source: anyhow!(
232                        "version-constraint: invalid version number {}",
233                        min_version_str
234                    ),
235                    pos: Some(pos),
236                });
237            }
238        };
239
240        if line.as_bytes()[op_pos + 1] == b'=' {
241            begin_version_kw = op_pos + 2;
242        } else {
243            begin_version_kw = op_pos + 1;
244            min_version += 1;
245        }
246    };
247
248    let version_start = begin_version_kw + "version".len();
249    if line[begin_version_kw..version_start].to_string() != "version" {
250        return Err(PosError {
251            source: anyhow!(
252                "version-constraint: invalid property {} (found '{}', expected 'version' {begin_version_kw})",
253                line[2..closed_brace_pos].to_string(),
254                line[begin_version_kw..version_start].to_string()
255            ),
256            pos: Some(pos),
257        });
258    }
259    let remainder = line[closed_brace_pos + 1..].to_string();
260    line_reader.push(&remainder);
261    const MAX_VERSION: i32 = 9999999;
262
263    if version_start >= closed_brace_pos && min_version != MIN_VERSION {
264        return Ok(Some(VersionConstraint {
265            min: min_version,
266            max: MAX_VERSION,
267        }));
268    }
269
270    let version_pos = if line.as_bytes()[version_start + 1].is_ascii_digit() {
271        version_start + 1
272    } else {
273        version_start + 2
274    };
275    let version = match line[version_pos..closed_brace_pos].parse::<i32>() {
276        Ok(x) => x,
277        Err(_) => {
278            return Err(PosError {
279                source: anyhow!(
280                    "version-constraint: invalid version number {}",
281                    line[version_pos..closed_brace_pos].to_string()
282                ),
283                pos: Some(pos),
284            });
285        }
286    };
287
288    match &line[version_start..version_pos] {
289        "=" => Ok(Some(VersionConstraint {
290            min: version,
291            max: version,
292        })),
293        "<=" => Ok(Some(VersionConstraint {
294            min: min_version,
295            max: version,
296        })),
297        "<" => Ok(Some(VersionConstraint {
298            min: min_version,
299            max: version - 1,
300        })),
301        ">=" if min_version == MIN_VERSION => Ok(Some(VersionConstraint {
302            min: version,
303            max: MAX_VERSION,
304        })),
305        ">" if min_version == MIN_VERSION => Ok(Some(VersionConstraint {
306            min: version + 1,
307            max: MAX_VERSION,
308        })),
309        ">=" | ">" => Err(PosError {
310            source: anyhow!(
311                "version-constraint: found comparison operator {} with a set minimum version {min_version}",
312                line[version_start..version_pos].to_string()
313            ),
314            pos: Some(pos),
315        }),
316        _ => Err(PosError {
317            source: anyhow!(
318                "version-constraint: unknown comparison operator {}",
319                line[version_start..version_pos].to_string()
320            ),
321            pos: Some(pos),
322        }),
323    }
324}
325
326fn parse_sql(line_reader: &mut LineReader) -> Result<SqlCommand, PosError> {
327    let (_, line1) = line_reader.next().unwrap();
328    let query = line1[1..].trim().to_owned();
329    let expected_start = line_reader.raw_pos;
330    let line2 = slurp_one(line_reader);
331    let line3 = slurp_one(line_reader);
332    let mut column_names = None;
333    let mut expected_rows = Vec::new();
334    static HASH_REGEX: LazyLock<Regex> =
335        LazyLock::new(|| Regex::new(r"^(\S+) values hashing to (\S+)$").unwrap());
336    match (line2, line3) {
337        (Some((pos2, line2)), Some((pos3, line3))) => {
338            if line3.len() >= 3 && line3.chars().all(|c| c == '-') {
339                column_names = Some(split_line(pos2, &line2)?);
340            } else {
341                expected_rows.push(split_line(pos2, &line2)?);
342                expected_rows.push(split_line(pos3, &line3)?);
343            }
344        }
345        (Some((pos2, line2)), None) => match HASH_REGEX.captures(&line2) {
346            Some(captures) => match captures[1].parse::<usize>() {
347                Ok(num_values) => {
348                    return Ok(SqlCommand {
349                        query,
350                        expected_output: SqlOutput::Hashed {
351                            num_values,
352                            md5: captures[2].to_owned(),
353                        },
354                        expected_start: 0,
355                        expected_end: 0,
356                    });
357                }
358                Err(err) => {
359                    return Err(PosError {
360                        source: anyhow!("Error parsing number of expected rows: {}", err),
361                        pos: Some(pos2),
362                    });
363                }
364            },
365            None => expected_rows.push(split_line(pos2, &line2)?),
366        },
367        _ => (),
368    }
369    while let Some((pos, line)) = slurp_one(line_reader) {
370        expected_rows.push(split_line(pos, &line)?)
371    }
372    let expected_end = line_reader.raw_pos;
373    Ok(SqlCommand {
374        query,
375        expected_output: SqlOutput::Full {
376            column_names,
377            expected_rows,
378        },
379        expected_start,
380        expected_end,
381    })
382}
383
384fn parse_explain_sql(line_reader: &mut LineReader) -> Result<SqlCommand, PosError> {
385    let (_, line1) = line_reader.next().unwrap();
386    let expected_start = line_reader.raw_pos;
387    // This is a bit of a hack to extract the next chunk of the file with
388    // blank lines intact. Ideally the `LineReader` would expose the API we
389    // need directly, but that would require a large refactor.
390    let mut expected_output: String = line_reader
391        .inner
392        .lines()
393        .filter(|l| !matches!(l.chars().next(), Some('#')))
394        .take_while(|l| !is_sigil(l.chars().next()))
395        .fold(String::new(), |mut output, l| {
396            let _ = write!(output, "{}\n", l);
397            output
398        });
399    while expected_output.ends_with("\n\n") {
400        expected_output.pop();
401    }
402    // We parsed the multiline expected_output directly using line_reader.inner
403    // above.
404    slurp_all(line_reader);
405    let expected_end = line_reader.raw_pos;
406
407    Ok(SqlCommand {
408        query: line1[1..].trim().to_owned(),
409        expected_output: SqlOutput::Full {
410            column_names: None,
411            expected_rows: vec![vec![expected_output]],
412        },
413        expected_start,
414        expected_end,
415    })
416}
417
418fn parse_fail_sql(line_reader: &mut LineReader) -> Result<FailSqlCommand, PosError> {
419    let (pos, line1) = line_reader.next().unwrap();
420    let line2 = slurp_one(line_reader);
421    let (err_pos, expected_error) = match line2 {
422        Some((err_pos, line2)) => (err_pos, line2),
423        None => {
424            return Err(PosError {
425                pos: Some(pos),
426                source: anyhow!("failing SQL command is missing expected error message"),
427            });
428        }
429    };
430    let query = line1[1..].trim().to_string();
431
432    let expected_error = if let Some(e) = expected_error.strip_prefix("regex:") {
433        SqlExpectedError::Regex(e.trim().into())
434    } else if let Some(e) = expected_error.strip_prefix("contains:") {
435        SqlExpectedError::Contains(e.trim().into())
436    } else if let Some(e) = expected_error.strip_prefix("exact:") {
437        SqlExpectedError::Exact(e.trim().into())
438    } else if expected_error == "timeout" {
439        SqlExpectedError::Timeout
440    } else {
441        return Err(PosError {
442            pos: Some(err_pos),
443            source: anyhow!(
444                "Query error must start with match specifier (`regex:`|`contains:`|`exact:`|`timeout`)"
445            ),
446        });
447    };
448
449    let extra_error = |line_reader: &mut LineReader, prefix| {
450        if let Some((_pos, line)) = line_reader.peek() {
451            if let Some(_) = line.strip_prefix(prefix) {
452                let line = line_reader
453                    .next()
454                    .map(|(_, line)| line)
455                    .unwrap()
456                    .strip_prefix(prefix)
457                    .map(|line| line.to_string())
458                    .unwrap();
459                Some(line.trim().to_string())
460            } else {
461                None
462            }
463        } else {
464            None
465        }
466    };
467    // Expect `hint` to always follow `detail` if they are both present, for now.
468    let expected_detail = extra_error(line_reader, "detail:");
469    let expected_hint = extra_error(line_reader, "hint:");
470
471    Ok(FailSqlCommand {
472        query: query.trim().to_string(),
473        expected_error,
474        expected_detail,
475        expected_hint,
476    })
477}
478
479fn split_line(pos: usize, line: &str) -> Result<Vec<String>, PosError> {
480    let mut out = Vec::new();
481    let mut field = String::new();
482    let mut in_quotes = None;
483    let mut escaping = false;
484    for (i, c) in line.char_indices() {
485        if in_quotes.is_none() && c.is_whitespace() {
486            if !field.is_empty() {
487                out.push(field);
488                field = String::new();
489            }
490        } else if c == '"' && !escaping {
491            if in_quotes.is_none() {
492                in_quotes = Some(i)
493            } else {
494                in_quotes = None;
495                out.push(field);
496                field = String::new();
497            }
498        } else if c == '\\' && !escaping && in_quotes.is_some() {
499            escaping = true;
500        } else if escaping {
501            field.push(match c {
502                'n' => '\n',
503                't' => '\t',
504                'r' => '\r',
505                '0' => '\0',
506                c => c,
507            });
508            escaping = false;
509        } else {
510            field.push(c);
511        }
512    }
513    if let Some(i) = in_quotes {
514        return Err(PosError {
515            source: anyhow!("unterminated quote"),
516            pos: Some(pos + i),
517        });
518    }
519    if !field.is_empty() {
520        out.push(field);
521    }
522    Ok(out)
523}
524
525fn slurp_all(line_reader: &mut LineReader) -> Vec<String> {
526    let mut out = Vec::new();
527    while let Some((_, line)) = slurp_one(line_reader) {
528        out.push(line);
529    }
530    out
531}
532
533fn slurp_one(line_reader: &mut LineReader) -> Option<(usize, String)> {
534    while let Some((_, line)) = line_reader.peek() {
535        match line.chars().next() {
536            Some('#') => {
537                // Comment line. Skip.
538                let _ = line_reader.next();
539            }
540            Some('$') | Some('>') | Some('!') | Some('?') => return None,
541            Some('\\') => {
542                return line_reader.next().map(|(pos, mut line)| {
543                    line.remove(0);
544                    (pos, line)
545                });
546            }
547            _ => return line_reader.next(),
548        }
549    }
550    None
551}
552
553pub struct LineReader<'a> {
554    inner: &'a str,
555    #[allow(clippy::option_option)]
556    next: Option<Option<(usize, String)>>,
557
558    src_line: usize,
559    pos: usize,
560    pos_map: BTreeMap<usize, (usize, usize)>,
561    raw_pos: usize,
562}
563
564impl<'a> LineReader<'a> {
565    pub fn new(inner: &'a str) -> LineReader<'a> {
566        let mut pos_map = BTreeMap::new();
567        pos_map.insert(0, (1, 1));
568        LineReader {
569            inner,
570            src_line: 1,
571            next: None,
572            pos: 0,
573            pos_map,
574            raw_pos: 0,
575        }
576    }
577
578    fn peek(&mut self) -> Option<&(usize, String)> {
579        if self.next.is_none() {
580            self.next = Some(self.next())
581        }
582        self.next.as_ref().unwrap().as_ref()
583    }
584
585    pub fn line_col(&self, pos: usize) -> (usize, usize) {
586        let (base_pos, (line, col)) = self.pos_map.range(..=pos).next_back().unwrap();
587        (*line, col + (pos - base_pos))
588    }
589
590    fn push(&mut self, text: &String) {
591        self.next = Some(Some((0usize, text.to_string())));
592    }
593}
594
595impl<'a> Iterator for LineReader<'a> {
596    type Item = (usize, String);
597
598    fn next(&mut self) -> Option<Self::Item> {
599        if let Some(next) = self.next.take() {
600            return next;
601        }
602        if self.inner.is_empty() {
603            return None;
604        }
605        let mut fold_newlines = is_non_sql_sigil(self.inner.chars().next());
606        let mut handle_newlines = is_sql_sigil(self.inner.chars().next());
607        let mut line = String::new();
608        let mut chars = self.inner.char_indices().fuse().peekable();
609        while let Some((i, c)) = chars.next() {
610            if c == '\n' {
611                self.src_line += 1;
612                if fold_newlines && self.inner.get(i + 1..i + 3) == Some("  ") {
613                    // Chomp the newline and one space. This ensures a SQL query
614                    // that is split over two lines does not become invalid. For $ commands the
615                    // newline should not be removed so that the argument parser can handle the
616                    // arguments correctly.
617                    chars.next();
618                    self.pos_map.insert(self.pos + i, (self.src_line, 2));
619                    continue;
620                } else if handle_newlines && self.inner.get(i + 1..i + 3) == Some("  ") {
621                    // Chomp the two spaces after newline. This ensures a SQL query
622                    // that is split over two lines does not become invalid, and keeping the
623                    // newline ensures that comments don't remove the following lines.
624                    line.push(c);
625                    chars.next();
626                    chars.next();
627                    self.pos_map.insert(self.pos + i + 1, (self.src_line, 2));
628                    continue;
629                } else if line.chars().all(char::is_whitespace) {
630                    line.clear();
631                    fold_newlines = is_non_sql_sigil(chars.peek().map(|c| c.1));
632                    handle_newlines = is_sql_sigil(chars.peek().map(|c| c.1));
633                    self.pos_map.insert(self.pos, (self.src_line, 1));
634                    continue;
635                }
636                let pos = self.pos;
637                self.pos += i;
638                self.raw_pos += i + 1; // Include \n character in count
639                self.pos_map.insert(self.pos, (self.src_line, 1));
640                self.inner = &self.inner[i + 1..];
641                return Some((pos, line));
642            }
643            line.push(c)
644        }
645        self.inner = "";
646        if !line.chars().all(char::is_whitespace) {
647            Some((self.pos, line))
648        } else {
649            None
650        }
651    }
652}
653
654fn is_sigil(c: Option<char>) -> bool {
655    is_sql_sigil(c) || is_non_sql_sigil(c)
656}
657
658fn is_sql_sigil(c: Option<char>) -> bool {
659    matches!(c, Some('>') | Some('!') | Some('?'))
660}
661
662fn is_non_sql_sigil(c: Option<char>) -> bool {
663    matches!(c, Some('$'))
664}
665
666struct BuiltinReader<'a> {
667    inner: &'a str,
668    pos: usize,
669}
670
671impl<'a> BuiltinReader<'a> {
672    fn new(line: &str, pos: usize) -> BuiltinReader {
673        BuiltinReader {
674            inner: &line[1..],
675            pos,
676        }
677    }
678}
679
680impl<'a> Iterator for BuiltinReader<'a> {
681    type Item = Result<(usize, String), PosError>;
682
683    fn next(&mut self) -> Option<Self::Item> {
684        if self.inner.is_empty() {
685            return None;
686        }
687
688        let mut iter = self.inner.char_indices().peekable();
689
690        while let Some((i, c)) = iter.peek() {
691            if c == &' ' {
692                iter.next();
693            } else {
694                self.pos += i;
695                break;
696            }
697        }
698
699        let mut token = String::new();
700        let mut nesting = Vec::new();
701        let mut done = false;
702        let mut quoted = false;
703        for (i, c) in iter {
704            if c == ' ' && nesting.is_empty() && !quoted {
705                done = true;
706                continue;
707            } else if done {
708                if let Some(nested) = nesting.last() {
709                    return Some(Err(PosError {
710                        pos: Some(self.pos + i),
711                        source: anyhow!(
712                            "command argument has unterminated open {}",
713                            if nested == &'{' { "brace" } else { "bracket" }
714                        ),
715                    }));
716                }
717                let pos = self.pos;
718                self.pos += i;
719                self.inner = &self.inner[i..];
720                return Some(Ok((pos, token)));
721            } else if (c == '{' || c == '[') && !quoted {
722                nesting.push(c);
723            } else if (c == '}' || c == ']') && !quoted {
724                if let Some(nested) = nesting.last() {
725                    if (nested == &'{' && c == '}') || (nested == &'[' && c == ']') {
726                        nesting.pop();
727                    } else {
728                        return Some(Err(PosError {
729                            pos: Some(self.pos + i),
730                            source: anyhow!(
731                                "command argument has unterminated open {}",
732                                if nested == &'{' { "brace" } else { "bracket" }
733                            ),
734                        }));
735                    }
736                } else {
737                    return Some(Err(PosError {
738                        pos: Some(self.pos + i),
739                        source: anyhow!(
740                            "command argument has unbalanced close {}",
741                            if c == '}' { "brace" } else { "bracket" }
742                        ),
743                    }));
744                }
745            } else if c == '"' && nesting.is_empty() {
746                // remove the double quote for un-nested commands such as: command="\dt public"
747                // keep the quotes when inside of a nested object such as: schema={ "type" : "array" }
748                quoted = !quoted;
749                continue;
750            }
751            token.push(c);
752        }
753
754        if let Some(nested) = nesting.last() {
755            return Some(Err(PosError {
756                pos: Some(self.pos + self.inner.len() - 1),
757                source: anyhow!(
758                    "command argument has unterminated open {}",
759                    if nested == &'{' { "brace" } else { "bracket" }
760                ),
761            }));
762        }
763
764        if quoted {
765            return Some(Err(PosError {
766                pos: Some(self.pos),
767                source: anyhow!("command argument has unterminated open double quote",),
768            }));
769        }
770
771        self.inner = "";
772        if token.is_empty() {
773            None
774        } else {
775            Some(Ok((self.pos, token)))
776        }
777    }
778}
779
780#[derive(Debug, Clone)]
781pub struct ArgMap(BTreeMap<String, String>);
782
783impl ArgMap {
784    pub fn values_mut(&mut self) -> btree_map::ValuesMut<String, String> {
785        self.0.values_mut()
786    }
787
788    pub fn opt_string(&mut self, name: &str) -> Option<String> {
789        self.0.remove(name)
790    }
791
792    pub fn string(&mut self, name: &str) -> Result<String, anyhow::Error> {
793        self.opt_string(name)
794            .ok_or_else(|| anyhow!("missing {} parameter", name))
795    }
796
797    pub fn opt_parse<T>(&mut self, name: &str) -> Result<Option<T>, anyhow::Error>
798    where
799        T: FromStr,
800        T::Err: Error + Send + Sync + 'static,
801    {
802        match self.opt_string(name) {
803            Some(val) => {
804                let t = val
805                    .parse()
806                    .with_context(|| format!("parsing {} parameter", name))?;
807                Ok(Some(t))
808            }
809            None => Ok(None),
810        }
811    }
812
813    pub fn parse<T>(&mut self, name: &str) -> Result<T, anyhow::Error>
814    where
815        T: FromStr,
816        T::Err: Error + Send + Sync + 'static,
817    {
818        match self.opt_parse(name) {
819            Ok(None) => bail!("missing {} parameter", name),
820            Ok(Some(t)) => Ok(t),
821            Err(err) => Err(err),
822        }
823    }
824
825    pub fn opt_bool(&mut self, name: &str) -> Result<Option<bool>, anyhow::Error> {
826        self.opt_string(name)
827            .map(|val| {
828                if val == "true" {
829                    Ok(true)
830                } else if val == "false" {
831                    Ok(false)
832                } else {
833                    bail!("bad value for boolean parameter {}: {}", name, val);
834                }
835            })
836            .transpose()
837    }
838
839    pub fn done(&self) -> Result<(), anyhow::Error> {
840        if let Some(name) = self.0.keys().next() {
841            bail!("unknown parameter {}", name);
842        }
843        Ok(())
844    }
845}
846
847impl IntoIterator for ArgMap {
848    type Item = (String, String);
849    type IntoIter = btree_map::IntoIter<String, String>;
850
851    fn into_iter(self) -> Self::IntoIter {
852        self.0.into_iter()
853    }
854}
mz_testdrive/parser.rs

mz_testdrive/
parser.rs
