mz_sqllogictest/

parser.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

//! A parser for sqllogictest.

use std::borrow::ToOwned;
use std::sync::LazyLock;

use anyhow::{anyhow, bail};
use mz_repr::ColumnName;
use regex::Regex;

use crate::ast::{Location, Mode, Output, QueryOutput, Record, Sort, Type};

static QUERY_OUTPUT_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\r?\n----").unwrap());
static DOUBLE_LINE_REGEX: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"(\n|\r\n|$)(\n|\r\n|$)").unwrap());
static EOF_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"(\n|\r\n)EOF(\n|\r\n)").unwrap());
// Separates the regex from the replacement in a `replace` directive. Two or
// more spaces, so that the regex itself may contain single spaces.
static REPLACE_SEP_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r" {2,}").unwrap());

#[derive(Debug, Clone)]
pub struct Parser<'a> {
    contents: &'a str,
    fname: String,
    curline: usize,
    mode: Mode,
}

impl<'a> Parser<'a> {
    pub fn new(fname: &str, contents: &'a str) -> Self {
        Parser {
            contents,
            fname: fname.to_string(),
            curline: 1,
            mode: Mode::Standard,
        }
    }

    pub fn is_done(&self) -> bool {
        self.contents.is_empty()
    }

    pub fn location(&self) -> Location {
        Location {
            file: self.fname.clone(),
            line: self.curline,
        }
    }

    fn consume(&mut self, upto: usize) {
        for ch in self.contents[..upto].chars() {
            if ch == '\n' {
                self.curline += 1;
            }
        }
        self.contents = &self.contents[upto..];
    }

    pub fn split_at(&mut self, sep: &Regex) -> Result<&'a str, anyhow::Error> {
        match sep.find(self.contents) {
            Some(found) => {
                let result = &self.contents[..found.start()];
                self.consume(found.end());
                Ok(result)
            }
            None => bail!("Couldn't split {:?} at {:?}", self.contents, sep),
        }
    }

    pub fn parse_record(&mut self) -> Result<Record<'a>, anyhow::Error> {
        if self.is_done() {
            return Ok(Record::Halt);
        }

        let line_number = self.curline;

        static COMMENT_AND_LINE_REGEX: LazyLock<Regex> =
            LazyLock::new(|| Regex::new("(#[^\n]*)?\r?(\n|$)").unwrap());
        let first_line = self.split_at(&COMMENT_AND_LINE_REGEX)?.trim();

        if first_line.is_empty() {
            // query starts on the next line
            return self.parse_record();
        }

        let mut words = first_line.split(' ').peekable();
        match words.next().unwrap() {
            "statement" => self.parse_statement(words, first_line),

            "query" => self.parse_query(words, first_line),

            "simple" => self.parse_simple(words),

            "hash-threshold" => {
                let threshold = words
                    .next()
                    .ok_or_else(|| anyhow!("missing threshold in: {}", first_line))?
                    .parse::<u64>()
                    .map_err(|err| anyhow!("invalid threshold ({}) in: {}", err, first_line))?;
                Ok(Record::HashThreshold { threshold })
            }

            // we'll follow the postgresql version of all these tests
            "skipif" => {
                match words.next().unwrap() {
                    "postgresql" => {
                        // discard next record
                        self.parse_record()?;
                        self.parse_record()
                    }
                    _ => self.parse_record(),
                }
            }
            "onlyif" => {
                match words.next().unwrap() {
                    "postgresql" => self.parse_record(),
                    _ => {
                        // discard next record
                        self.parse_record()?;
                        self.parse_record()
                    }
                }
            }

            "halt" => Ok(Record::Halt),

            // this is some cockroach-specific thing, we don't care
            "subtest" | "kv-batch-size" | "skip_on_retry" => self.parse_record(),

            // CockroachDB's `user` directive switches the session user for
            // subsequent records.
            "user" => Ok(Record::User {
                location: self.location(),
                user: words
                    .next()
                    .ok_or_else(|| anyhow!("user directive missing name"))?,
            }),

            // CockroachDB's `let $var` binds the result of the following query
            // to a variable that later records reference. We don't support the
            // binding, so skip the directive and its query block. Records that
            // reference the variable fail at execution time instead.
            "let" => {
                self.split_at(&DOUBLE_LINE_REGEX)?;
                self.parse_record()
            }

            "mode" => {
                self.mode = match words.next() {
                    Some("cockroach") => Mode::Cockroach,
                    Some("standard") | Some("sqlite") => Mode::Standard,
                    other => bail!("unknown parse mode: {:?}", other),
                };
                self.parse_record()
            }

            "copy" => Ok(Record::Copy {
                table_name: words
                    .next()
                    .ok_or_else(|| anyhow!("load directive missing table name"))?,
                tsv_path: words
                    .next()
                    .ok_or_else(|| anyhow!("load directive missing TSV path"))?,
            }),

            "reset-server" => Ok(Record::ResetServer),

            // `replace <regex>  <replacement>`: register a substitution applied
            // to the actual output of subsequent queries before comparison. The
            // regex and replacement are separated by two-or-more spaces, so the
            // regex may contain single spaces. See `Record::Replace`.
            "replace" => {
                let args = first_line
                    .strip_prefix("replace")
                    .expect("dispatched on \"replace\"")
                    .trim_start();
                let mut parts = REPLACE_SEP_REGEX.splitn(args, 2);
                let pattern = parts
                    .next()
                    .filter(|s| !s.is_empty())
                    .ok_or_else(|| anyhow!("replace directive missing regex in: {}", first_line))?;
                let replacement = parts.next().ok_or_else(|| {
                    anyhow!(
                        "replace directive missing replacement (separate the regex \
                         and replacement with two or more spaces) in: {}",
                        first_line
                    )
                })?;
                // Validate the regex now, so an error is located at parse time.
                Regex::new(pattern).map_err(|e| {
                    anyhow!("invalid regex {:?} in replace directive: {}", pattern, e)
                })?;
                Ok(Record::Replace {
                    pattern: pattern.to_owned(),
                    replacement: replacement.to_owned(),
                })
            }

            other => bail!(
                "Unexpected start of record on line {}: {}",
                line_number,
                other
            ),
        }
    }

    pub fn parse_records(&mut self) -> Result<Vec<Record<'a>>, anyhow::Error> {
        let mut records = vec![];
        loop {
            match self.parse_record()? {
                Record::Halt => break,
                record => records.push(record),
            }
        }
        Ok(records)
    }

    fn parse_statement(
        &mut self,
        mut words: impl Iterator<Item = &'a str>,
        first_line: &'a str,
    ) -> Result<Record<'a>, anyhow::Error> {
        let location = self.location();
        let mut expected_error = None;
        let mut rows_affected = None;
        match words.next() {
            Some("count") => {
                rows_affected = Some(
                    words
                        .next()
                        .ok_or_else(|| anyhow!("missing count of rows affected"))?
                        .parse::<u64>()
                        .map_err(|err| anyhow!("parsing count of rows affected: {}", err))?,
                );
            }
            Some("error") => expected_error = Some(parse_expected_error(first_line)),
            // CockroachDB's `statement notice <regex>` expects the statement
            // to succeed and additionally emit a matching notice. We only
            // check for success.
            Some("notice") => (),
            // An `ok` prefix accepts the typos present in files imported from
            // CockroachDB (`oK`, `ok;`, `oko`), which CockroachDB's own
            // lenient runner treats as plain `ok`.
            Some(disposition) if disposition.to_lowercase().starts_with("ok") => (),
            // A bare `statement` with no disposition expects success.
            None => (),
            _ => bail!("invalid statement disposition: {}", first_line),
        };
        let sql = self.split_at(&DOUBLE_LINE_REGEX)?;
        Ok(Record::Statement {
            expected_error,
            rows_affected,
            sql,
            location,
        })
    }

    fn parse_query(
        &mut self,
        mut words: std::iter::Peekable<impl Iterator<Item = &'a str>>,
        first_line: &'a str,
    ) -> Result<Record<'a>, anyhow::Error> {
        let location = self.location();
        if words.peek() == Some(&"error") {
            let error = parse_expected_error(first_line);
            let sql = self.split_at(&DOUBLE_LINE_REGEX)?;
            return Ok(Record::Query {
                sql,
                output: Err(error),
                location,
            });
        }

        let types = words.next().map_or(Ok(vec![]), parse_types)?;
        let mut sort = Sort::No;
        let mut check_column_names = false;
        let mut multiline = false;
        let mut noticetrace = false;
        if let Some(options) = words.next() {
            for option in options.split(',') {
                match option {
                    "nosort" => sort = Sort::No,
                    "rowsort" => sort = Sort::Row,
                    "valuesort" => sort = Sort::Value,
                    "colnames" => check_column_names = true,
                    "multiline" => multiline = true,
                    // CockroachDB re-runs `retry` queries until the output
                    // converges. We run them once, like any other query.
                    "retry" => (),
                    // CockroachDB `noticetrace` queries assert the emitted
                    // notices rather than rows. We can't observe notices, so
                    // the whole record is skipped below.
                    "noticetrace" => noticetrace = true,
                    other => {
                        if other.starts_with("partialsort") {
                            // TODO(jamii) https://github.com/cockroachdb/cockroach/blob/d2f7fbf5dd1fc1a099bbad790a2e1f7c60a66cc3/pkg/sql/logictest/logic.go#L153
                            // partialsort has comma-separated arguments so our parsing is totally broken
                            // luckily it always comes last in the existing tests, so we can just bail out for now
                            sort = Sort::Row;
                            break;
                        } else {
                            bail!("Unrecognized option {:?} in {:?}", other, options);
                        }
                    }
                };
            }
        }
        if multiline && (check_column_names || sort.yes()) {
            bail!("multiline option is incompatible with all other options");
        }
        let label = words.next();
        static LINE_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new("\r?(\n|$)").unwrap());
        static HASH_REGEX: LazyLock<Regex> =
            LazyLock::new(|| Regex::new(r"(\S+) values hashing to (\S+)").unwrap());
        // CockroachDB queries may omit the `----` separator entirely, in
        // which case the query must succeed and return no rows. Detect this
        // by checking whether the record ends before the next `----`, which
        // otherwise belongs to a later record. A blank line only ends the
        // record if what follows it starts a new record (directive, comment,
        // or end of file): hand-written files contain blank lines inside a
        // query's SQL, which CockroachDB's own files never do.
        static RECORD_START_REGEX: LazyLock<Regex> = LazyLock::new(|| {
            Regex::new(
                r"^(#|$|statement( |$)|query( |$)|simple( |$)|halt( |$)|mode |copy |user |subtest( |$)|let |skipif |onlyif |hash-threshold |reset-server( |$)|replace |kv-batch-size |skip_on_retry( |$))",
            )
            .unwrap()
        });
        let sep = QUERY_OUTPUT_REGEX.find(self.contents);
        let end = DOUBLE_LINE_REGEX.find(self.contents);
        let no_separator = match (&sep, &end) {
            (Some(sep), Some(end)) if sep.start() < end.start() => false,
            (_, Some(end)) => {
                let mut rest = &self.contents[end.end()..];
                loop {
                    let line_end = rest.find('\n').map_or(rest.len(), |i| i + 1);
                    let line = rest[..line_end].trim();
                    if line.is_empty() && line_end < rest.len() {
                        rest = &rest[line_end..];
                        continue;
                    }
                    break RECORD_START_REGEX.is_match(line);
                }
            }
            (_, None) => false,
        };
        if no_separator {
            let sql = self.split_at(&DOUBLE_LINE_REGEX)?;
            if noticetrace {
                return self.parse_record();
            }
            return Ok(Record::Query {
                sql,
                output: Ok(QueryOutput {
                    types,
                    sort,
                    multiline,
                    label,
                    column_names: None,
                    mode: self.mode,
                    output: Output::Values(vec![]),
                    // An empty slice at the end of the SQL, so rewriting has
                    // an in-bounds position to work with.
                    output_str: &sql[sql.len()..],
                }),
                location,
            });
        }

        let sql = self.split_at(&QUERY_OUTPUT_REGEX)?;
        let mut output_str = self.split_at(if multiline {
            &EOF_REGEX
        } else {
            &DOUBLE_LINE_REGEX
        })?;

        if noticetrace {
            return self.parse_record();
        }

        // The `split_at(&QUERY_OUTPUT_REGEX)` stopped at the end of `----`, so `output_str` usually
        // starts with a newline, which is not actually part of the expected output. Strip off this
        // newline.
        output_str = if let Some(output_str_stripped) = regexp_strip_prefix(output_str, &LINE_REGEX)
        {
            output_str_stripped
        } else {
            // There should always be a newline after `----`, because we have a lint that there is
            // always a newline at the end of a file. However, we can still get here, when
            // the expected output is empty, in which case the EOF_REGEX or DOUBLE_LINE_REGEX eats
            // the newline at the end of the `----`.
            assert!(output_str.is_empty());
            output_str
        };

        // We don't want to advance the expected output past the column names so rewriting works,
        // but need to be able to parse past them, so remember the position before possible column
        // names.
        let query_output_str = output_str;
        let column_names = if check_column_names {
            Some(
                split_at(&mut output_str, &LINE_REGEX)?
                    .split(' ')
                    .filter(|s| !s.is_empty())
                    .map(|s| ColumnName::from(s.replace('␠', " ")))
                    .collect(),
            )
        } else {
            None
        };
        let output = match HASH_REGEX.captures(output_str) {
            Some(captures) => Output::Hashed {
                num_values: captures.get(1).unwrap().as_str().parse::<usize>()?,
                md5: captures.get(2).unwrap().as_str().to_owned(),
            },
            None => {
                if multiline {
                    Output::Values(vec![output_str.to_owned()])
                } else if output_str.starts_with('\r') || output_str.starts_with('\n') {
                    Output::Values(vec![])
                } else {
                    let mut vals: Vec<String> = output_str.lines().map(|s| s.to_owned()).collect();
                    match self.mode {
                        Mode::Standard => {
                            if !multiline {
                                vals = vals.into_iter().map(|val| val.replace('⏎', "\n")).collect();
                            }

                            if sort == Sort::Value {
                                vals.sort();
                            }
                        }
                        Mode::Cockroach => {
                            let mut rows: Vec<Vec<String>> = vec![];
                            for line in vals {
                                let cols = split_cols(&line, types.len());
                                if sort != Sort::No && cols.len() != types.len() {
                                    // We can't check this condition for
                                    // Sort::No, because some tests use strings
                                    // with whitespace that look like extra
                                    // columns. (Note that these tests never
                                    // use any of the sorting options.)
                                    bail!(
                                        "col len ({}) did not match declared col len ({})",
                                        cols.len(),
                                        types.len()
                                    );
                                }
                                rows.push(
                                    cols.into_iter()
                                        .map(|col| {
                                            let mut col = col.replace('␠', " ");
                                            if !multiline {
                                                col = col.replace('⏎', "\n");
                                            }
                                            col
                                        })
                                        .collect(),
                                );
                            }
                            if sort == Sort::Row {
                                rows.sort();
                            }
                            vals = rows.into_iter().flatten().collect();
                            if sort == Sort::Value {
                                vals.sort();
                            }
                        }
                    }
                    Output::Values(vals)
                }
            }
        };
        Ok(Record::Query {
            sql,
            output: Ok(QueryOutput {
                types,
                sort,
                multiline,
                label,
                column_names,
                mode: self.mode,
                output,
                output_str: query_output_str,
            }),
            location,
        })
    }

    fn parse_simple(
        &mut self,
        mut words: std::iter::Peekable<impl Iterator<Item = &'a str>>,
    ) -> Result<Record<'a>, anyhow::Error> {
        let location = self.location();
        let mut conn = None;
        let mut user = None;
        let mut password = None;
        let mut multiline = false;
        let mut sort = Sort::No;
        if let Some(options) = words.next() {
            for option in options.split(',') {
                if let Some(value) = option.strip_prefix("conn=") {
                    conn = Some(value);
                } else if let Some(value) = option.strip_prefix("user=") {
                    user = Some(value);
                } else if let Some(value) = option.strip_prefix("password=") {
                    password = Some(value);
                } else if option == "rowsort" {
                    sort = Sort::Row;
                } else if option == "multiline" {
                    multiline = true;
                } else {
                    bail!("Unrecognized option {:?} in {:?}", option, options);
                }
            }
        }
        if user.is_some() && conn.is_none() {
            bail!("cannot set user without also setting conn");
        }
        if password.is_some() && user.is_none() {
            bail!("cannot set password without also setting user");
        }
        let sql = self.split_at(&QUERY_OUTPUT_REGEX)?;
        let output_str = self
            .split_at(if multiline {
                &EOF_REGEX
            } else {
                &DOUBLE_LINE_REGEX
            })?
            .trim_start();
        let output = if multiline {
            Output::Values({
                let mut v = vec![output_str.to_owned()];
                // for simple queries we still have to pass the COMPLETE string after the EOF
                let complete_str = self.split_at(&DOUBLE_LINE_REGEX)?.trim_start();
                v.extend(complete_str.lines().map(String::from));
                v
            })
        } else {
            // We only apply rowsort in mode cockroach, for "query" statements,
            // so mirror that here.
            let mut output_lines: Vec<String> = output_str.lines().map(String::from).collect();

            if self.mode == Mode::Cockroach && sort == Sort::Row {
                output_lines.sort();
            }

            Output::Values(output_lines)
        };
        Ok(Record::Simple {
            location,
            conn,
            user,
            password,
            sql,
            sort,
            output,
            output_str,
        })
    }
}

fn split_at<'a>(input: &mut &'a str, sep: &Regex) -> Result<&'a str, anyhow::Error> {
    match sep.find(input) {
        Some(found) => {
            let result = &input[..found.start()];
            *input = &input[found.end()..];
            Ok(result)
        }
        None => bail!("Couldn't split {:?} at {:?}", input, sep),
    }
}

/// Parse a query result type string into a vec of expected types
fn parse_types(input: &str) -> Result<Vec<Type>, anyhow::Error> {
    input
        .chars()
        .map(|char| {
            Ok(match char {
                'T' => Type::Text,
                'I' => Type::Integer,
                'R' => Type::Real,
                // CockroachDB uses `F` for floats and `R` for decimals. We
                // don't distinguish the two.
                'F' => Type::Real,
                'B' => Type::Bool,
                'O' => Type::Oid,
                _ => bail!("Unexpected type char {} in: {}", char, input),
            })
        })
        .collect()
}

fn parse_expected_error(line: &str) -> &str {
    static PGCODE_RE: LazyLock<Regex> =
        LazyLock::new(|| Regex::new("(statement|query) error( pgcode [a-zA-Z0-9]{5})? ?").unwrap());
    // TODO(benesch): one day this should record the expected pgcode, if
    // specified.
    let pos = PGCODE_RE.find(line).unwrap().end();
    &line[pos..]
}

/// Split on whitespace to normalize multiple spaces to one space. This happens
/// unconditionally in Cockroach mode, regardless of the sort option.
///
/// TODO: this doesn't have the whitespace-collapsing behavior for
/// single-column values that cockroach relies on
pub(crate) fn split_cols(line: &str, expected_columns: usize) -> Vec<&str> {
    if expected_columns == 1 {
        vec![line.trim()]
    } else {
        line.split_whitespace().collect()
    }
}

pub fn regexp_strip_prefix<'a>(text: &'a str, regexp: &Regex) -> Option<&'a str> {
    match regexp.find(text) {
        Some(found) => {
            if found.start() == 0 {
                Some(&text[found.end()..])
            } else {
                None
            }
        }
        None => None,
    }
}

#[mz_ore::test]
fn test_parse_query_blank_lines_and_missing_separator() {
    // A blank line inside a query's SQL does not end the record.
    let file = "query I\nSELECT 1\n\nUNION ALL SELECT 2\n----\n1\n2\n\n";
    let records = Parser::new("f", file).parse_records().unwrap();
    assert_eq!(records.len(), 1);
    match &records[0] {
        Record::Query { sql, .. } => {
            assert!(sql.contains("UNION ALL"), "sql: {sql}")
        }
        other => panic!("unexpected record: {other:?}"),
    }

    // A query with no ---- separator expects zero rows, and the blank line
    // ends the record when a new record follows.
    let file = "query I\nSELECT 3\n\nstatement ok\nSELECT 4\n\n# comment\n\nquery I\nSELECT 5\n";
    let records = Parser::new("f", file).parse_records().unwrap();
    assert_eq!(records.len(), 3);
    match &records[0] {
        Record::Query { sql, output, .. } => {
            assert_eq!(*sql, "SELECT 3");
            assert_eq!(output.as_ref().unwrap().output, Output::Values(vec![]));
        }
        other => panic!("unexpected record: {other:?}"),
    }

    // No separator at end of file.
    let file = "query I\nSELECT 6\n";
    let records = Parser::new("f", file).parse_records().unwrap();
    assert_eq!(records.len(), 1);
    match &records[0] {
        Record::Query { output, .. } => {
            assert_eq!(output.as_ref().unwrap().output, Output::Values(vec![]));
        }
        other => panic!("unexpected record: {other:?}"),
    }
}