feat: run system commands

sqllogictest-bin/src/engines.rs

@@ -1,4 +1,6 @@
 use std::fmt::Display;
+use std::process::ExitStatus;
+use std::time::Duration;
 
 use async_trait::async_trait;
 use clap::ValueEnum;
@@ -97,14 +99,14 @@ impl std::error::Error for EnginesError {
     }
 }
 
-impl Engines {
-    async fn run(&mut self, sql: &str) -> Result<DBOutput<DefaultColumnType>, anyhow::Error> {
-        Ok(match self {
-            Engines::Postgres(e) => e.run(sql).await?,
-            Engines::PostgresExtended(e) => e.run(sql).await?,
-            Engines::External(e) => e.run(sql).await?,
-        })
-    }
+macro_rules! dispatch_engines {
+    ($impl:expr, $inner:ident, $body:tt) => {{
+        match $impl {
+            Engines::Postgres($inner) => $body,
+            Engines::PostgresExtended($inner) => $body,
+            Engines::External($inner) => $body,
+        }
+    }};
 }
 
 #[async_trait]
@@ -113,6 +115,22 @@ impl AsyncDB for Engines {
     type ColumnType = DefaultColumnType;
 
     async fn run(&mut self, sql: &str) -> Result<DBOutput<Self::ColumnType>, Self::Error> {
-        self.run(sql).await.map_err(EnginesError)
+        dispatch_engines!(self, e, {
+            e.run(sql)
+                .await
+                .map_err(|e| EnginesError(anyhow::Error::from(e)))
+        })
+    }
+
+    fn engine_name(&self) -> &str {
+        dispatch_engines!(self, e, { e.engine_name() })
+    }
+
+    async fn sleep(dur: Duration) {
+        tokio::time::sleep(dur).await
+    }
+
+    async fn run_command(command: std::process::Command) -> std::io::Result<ExitStatus> {
+        Command::from(command).status().await
     }
 }

sqllogictest-engines/src/external.rs

@@ -1,6 +1,6 @@
 use std::io;
 use std::marker::PhantomData;
-use std::process::Stdio;
+use std::process::{ExitStatus, Stdio};
 use std::time::Duration;
 
 use async_trait::async_trait;
@@ -120,6 +120,10 @@ impl AsyncDB for ExternalDriver {
     async fn sleep(dur: Duration) {
         tokio::time::sleep(dur).await
     }
+
+    async fn run_command(command: std::process::Command) -> std::io::Result<ExitStatus> {
+        Command::from(command).status().await
+    }
 }
 
 struct JsonDecoder<T>(PhantomData<T>);

sqllogictest-engines/src/postgres/extended.rs

@@ -1,4 +1,5 @@
 use std::fmt::Write;
+use std::process::{Command, ExitStatus};
 use std::time::Duration;
 
 use async_trait::async_trait;
@@ -317,4 +318,8 @@ impl sqllogictest::AsyncDB for Postgres<Extended> {
     async fn sleep(dur: Duration) {
         tokio::time::sleep(dur).await
     }
+
+    async fn run_command(command: Command) -> std::io::Result<ExitStatus> {
+        tokio::process::Command::from(command).status().await
+    }
 }

sqllogictest-engines/src/postgres/simple.rs

@@ -1,3 +1,4 @@
+use std::process::{Command, ExitStatus};
 use std::time::Duration;
 
 use async_trait::async_trait;
@@ -64,4 +65,8 @@ impl sqllogictest::AsyncDB for Postgres<Simple> {
     async fn sleep(dur: Duration) {
         tokio::time::sleep(dur).await
     }
+
+    async fn run_command(command: Command) -> std::io::Result<ExitStatus> {
+        tokio::process::Command::from(command).status().await
+    }
 }

sqllogictest/Cargo.toml

@@ -24,3 +24,6 @@ regex = "1.9.1"
 owo-colors = "3.5.0"
 md-5 = "0.10"
 fs-err = "2.9.0"
+
+[dev-dependencies]
+pretty_assertions = "1"

sqllogictest/src/parser.rs

@@ -109,6 +109,14 @@ pub enum Record<T: ColumnType> {
         /// The expected results.
         expected_results: Vec<String>,
     },
+    /// A system command is an external command that is to be executed by the shell. Currently it
+    /// must succeed and the output is ignored.
+    System {
+        loc: Location,
+        conditions: Vec<Condition>,
+        /// The external command.
+        command: String,
+    },
     /// A sleep period.
     Sleep {
         loc: Location,
@@ -239,6 +247,13 @@ impl<T: ColumnType> std::fmt::Display for Record<T> {
                 // query always ends with a blank line
                 writeln!(f)
             }
+            Record::System {
+                loc: _,
+                conditions: _,
+                command,
+            } => {
+                writeln!(f, "system ok\n{command}")
+            }
             Record::Sleep { loc: _, duration } => {
                 write!(f, "sleep {}", humantime::format_duration(*duration))
             }
@@ -612,6 +627,25 @@ fn parse_inner<T: ColumnType>(loc: &Location, script: &str) -> Result<Vec<Record
                     expected_error,
                 });
             }
+            ["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;
+                }
+                records.push(Record::System {
+                    loc,
+                    conditions: std::mem::take(&mut conditions),
+                    command,
+                });
+            }
             ["control", res @ ..] => match res {
                 ["sortmode", sort_mode] => match SortMode::try_from_str(sort_mode) {
                     Ok(sort_mode) => records.push(Record::Control(Control::SortMode(sort_mode))),
@@ -739,6 +773,11 @@ mod tests {
         parse_roundtrip::<CustomColumnType>("../tests/custom_type/custom_type.slt")
     }
 
+    #[test]
+    fn test_system_command() {
+        parse_roundtrip::<DefaultColumnType>("../tests/system_command/system_command.slt")
+    }
+
     #[test]
     fn test_fail_unknown_type() {
         let script = "\
@@ -805,18 +844,7 @@ select * from foo;
         let records = normalize_filename(records);
         let reparsed_records = normalize_filename(reparsed_records);
 
-        assert_eq!(
-            records, reparsed_records,
-            "Mismatch in reparsed records\n\
-                    *********\n\
-                    original:\n\
-                    *********\n\
-                    {records:#?}\n\n\
-                    *********\n\
-                    reparsed:\n\
-                    *********\n\
-                    {reparsed_records:#?}\n\n",
-        );
+        pretty_assertions::assert_eq!(records, reparsed_records, "Mismatch in reparsed records");
     }
 
     /// Replaces the actual filename in all Records with
@@ -829,6 +857,7 @@ select * from foo;
                 match &mut record {
                     Record::Include { loc, .. } => normalize_loc(loc),
                     Record::Statement { loc, .. } => normalize_loc(loc),
+                    Record::System { loc, .. } => normalize_loc(loc),
                     Record::Query { loc, .. } => normalize_loc(loc),
                     Record::Sleep { loc, .. } => normalize_loc(loc),
                     Record::Subtest { loc, .. } => normalize_loc(loc),

sqllogictest/src/runner.rs

@@ -3,6 +3,7 @@
 use std::collections::HashSet;
 use std::fmt::{Debug, Display};
 use std::path::Path;
+use std::process::{Command, ExitStatus};
 use std::sync::Arc;
 use std::time::Duration;
 use std::vec;
@@ -21,6 +22,7 @@ use crate::parser::*;
 use crate::{ColumnType, Connections, MakeConnection};
 
 #[derive(Debug, Clone)]
+#[non_exhaustive]
 pub enum RecordOutput<T: ColumnType> {
     Nothing,
     Query {
@@ -32,6 +34,9 @@ pub enum RecordOutput<T: ColumnType> {
         count: u64,
         error: Option<Arc<dyn std::error::Error + Send + Sync>>,
     },
+    System {
+        error: Option<Arc<dyn std::error::Error + Send + Sync>>,
+    },
 }
 
 #[non_exhaustive]
@@ -72,6 +77,15 @@ pub trait AsyncDB {
     async fn sleep(dur: Duration) {
         std::thread::sleep(dur);
     }
+
+    /// [`Runner`] calls this function to run a system command.
+    ///
+    /// The default implementation is `std::process::Command::status`, which is universial to any
+    /// async runtime but would block the current thread. If you are running in tokio runtime, you
+    /// should override this by `tokio::process::Command::status`.
+    async fn run_command(mut command: Command) -> std::io::Result<ExitStatus> {
+        command.status()
+    }
 }
 
 /// The database to be tested.
@@ -225,6 +239,7 @@ impl std::fmt::Display for RecordKind {
 ///
 /// For colored error message, use `self.display()`.
 #[derive(thiserror::Error, Debug, Clone)]
+#[non_exhaustive]
 pub enum TestErrorKind {
     #[error("parse error: {0}")]
     ParseError(ParseErrorKind),
@@ -236,6 +251,11 @@ pub enum TestErrorKind {
         err: Arc<dyn std::error::Error + Send + Sync>,
         kind: RecordKind,
     },
+    #[error("system command failed: {err}\n[CMD] {command}")]
+    SystemFail {
+        command: String,
+        err: Arc<dyn std::error::Error + Send + Sync>,
+    },
     // Remember to also update [`TestErrorKindDisplay`] if this message is changed.
     #[error("{kind} is expected to fail with error:\n\t{expected_err}\nbut got error:\n\t{err}\n[SQL] {sql}")]
     ErrorMismatch {
@@ -570,6 +590,40 @@ impl<D: AsyncDB, M: MakeConnection<Conn = D>> Runner<D, M> {
                     },
                 }
             }
+            Record::System {
+                conditions,
+                command,
+                loc: _,
+            } => {
+                let command = self.replace_keywords(command);
+
+                if should_skip(&self.labels, "", &conditions) {
+                    return RecordOutput::Nothing;
+                }
+
+                let cmd = if cfg!(target_os = "windows") {
+                    let mut cmd = std::process::Command::new("cmd");
+                    cmd.arg("/C").arg(command);
+                    cmd
+                } else {
+                    let mut cmd = std::process::Command::new("sh");
+                    cmd.arg("-c").arg(command);
+                    cmd
+                };
+                let result = D::run_command(cmd).await;
+
+                #[derive(thiserror::Error, Debug)]
+                #[error("external command exited unsuccessfully: {0}")]
+                struct SystemError(ExitStatus);
+
+                let error: Option<Arc<dyn std::error::Error + Send + Sync>> = match result {
+                    Ok(status) if status.success() => None,
+                    Ok(status) => Some(Arc::new(SystemError(status))),
+                    Err(e) => Some(Arc::new(e)),
+                };
+
+                RecordOutput::System { error }
+            }
             Record::Query {
                 conditions,
                 connection,
@@ -839,6 +893,18 @@ impl<D: AsyncDB, M: MakeConnection<Conn = D>> Runner<D, M> {
                     .at(loc));
                 }
             }
+            (
+                Record::System {
+                    loc,
+                    conditions: _,
+                    command,
+                },
+                RecordOutput::System { error },
+            ) => {
+                if let Some(err) = error {
+                    return Err(TestErrorKind::SystemFail { command, err }.at(loc));
+                }
+            }
             _ => unreachable!(),
         }
 

tests/Cargo.toml

@@ -5,6 +5,7 @@ edition = "2021"
 publish = false
 
 [dependencies]
+regex = "1.9.1"
 sqllogictest = { path = "../sqllogictest" }
 
 [[test]]
@@ -23,3 +24,7 @@ path = "./validator/validator.rs"
 [[test]]
 name = "test_dir_escape"
 path = "./test_dir_escape/test_dir_escape.rs"
+
+[[test]]
+name = "system_command"
+path = "./system_command/system_command.rs"