From 52d5122e07a2f75ba30c43382199e4b99f1ca8de Mon Sep 17 00:00:00 2001
From: xxchan <xxchan22f@gmail.com>
Date: Thu, 9 Nov 2023 12:53:44 +0800
Subject: [PATCH] refactor: extract multiline parsing to a function (#201)

Signed-off-by: xxchan <xxchan22f@gmail.com>
---
 sqllogictest/src/parser.rs | 74 ++++++++++++++++++--------------------
 1 file changed, 34 insertions(+), 40 deletions(-)

diff --git a/sqllogictest/src/parser.rs b/sqllogictest/src/parser.rs
index 940afdc..37b0b11 100644
--- a/sqllogictest/src/parser.rs
+++ b/sqllogictest/src/parser.rs
@@ -487,12 +487,12 @@ fn parse_inner<T: ColumnType>(loc: &Location, script: &str) -> Result<Vec<Record
         if let Some(text) = line.strip_prefix('#') {
             comments.push(text.to_string());
             if lines.peek().is_none() {
+                // Special handling for the case where the last line is a comment.
                 records.push(Record::Comment(comments));
-                comments = vec![];
+                break;
             }
             continue;
         }
-
         if !comments.is_empty() {
             records.push(Record::Comment(comments));
             comments = vec![];
@@ -567,17 +567,7 @@ fn parse_inner<T: ColumnType>(loc: &Location, script: &str) -> Result<Vec<Record
                     }
                     _ => return Err(ParseErrorKind::InvalidLine(line.into()).at(loc)),
                 };
-                let mut sql = match lines.next() {
-                    Some((_, line)) => line.into(),
-                    None => return Err(ParseErrorKind::UnexpectedEOF.at(loc.next_line())),
-                };
-                for (_, line) in &mut lines {
-                    if line.is_empty() {
-                        break;
-                    }
-                    sql += "\n";
-                    sql += line;
-                }
+                let (sql, _) = parse_lines(&mut lines, &loc, None)?;
                 records.push(Record::Statement {
                     loc,
                     conditions: std::mem::take(&mut conditions),
@@ -619,22 +609,7 @@ fn parse_inner<T: ColumnType>(loc: &Location, script: &str) -> Result<Vec<Record
 
                 // The SQL for the query is found on second an subsequent lines of the record
                 // up to first line of the form "----" or until the end of the record.
-                let mut sql = match lines.next() {
-                    Some((_, line)) => line.into(),
-                    None => return Err(ParseErrorKind::UnexpectedEOF.at(loc.next_line())),
-                };
-                let mut has_result = false;
-                for (_, line) in &mut lines {
-                    if line.is_empty() {
-                        break;
-                    }
-                    if line == "----" {
-                        has_result = true;
-                        break;
-                    }
-                    sql += "\n";
-                    sql += line;
-                }
+                let (sql, has_result) = parse_lines(&mut lines, &loc, Some("----"))?;
                 // Lines following the "----" are expected results of the query, one value per line.
                 let mut expected_results = vec![];
                 if has_result {
@@ -659,17 +634,7 @@ fn parse_inner<T: ColumnType>(loc: &Location, script: &str) -> Result<Vec<Record
             }
             ["system", "ok"] => {
                 // TODO: we don't support asserting error message for system command
-                let mut command = match lines.next() {
-                    Some((_, line)) => line.into(),
-                    None => return Err(ParseErrorKind::UnexpectedEOF.at(loc.next_line())),
-                };
-                for (_, line) in &mut lines {
-                    if line.is_empty() {
-                        break;
-                    }
-                    command += "\n";
-                    command += line;
-                }
+                let (command, _) = parse_lines(&mut lines, &loc, None)?;
                 records.push(Record::System {
                     loc,
                     conditions: std::mem::take(&mut conditions),
@@ -742,6 +707,35 @@ fn parse_file_inner<T: ColumnType>(loc: Location) -> Result<Vec<Record<T>>, Pars
     Ok(records)
 }
 
+/// Parse one or more lines until empty line or a delimiter.
+fn parse_lines<'a>(
+    lines: &mut impl Iterator<Item = (usize, &'a str)>,
+    loc: &Location,
+    delimiter: Option<&str>,
+) -> Result<(String, bool), ParseError> {
+    let mut found_delimiter = false;
+    let mut out = match lines.next() {
+        Some((_, line)) => Ok(line.into()),
+        None => Err(ParseErrorKind::UnexpectedEOF.at(loc.clone().next_line())),
+    }?;
+
+    for (_, line) in lines {
+        if line.is_empty() {
+            break;
+        }
+        if let Some(delimiter) = delimiter {
+            if line == delimiter {
+                found_delimiter = true;
+                break;
+            }
+        }
+        out += "\n";
+        out += line;
+    }
+
+    Ok((out, found_delimiter))
+}
+
 #[cfg(test)]
 mod tests {
     use std::io::Write;