Skip to main content

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