feat: support matching sqlstate with [statement|query] error (XX000)

sqllogictest-bin/src/engines.rs

@@ -73,17 +73,17 @@ pub(crate) async fn connect(
         EngineConfig::MySql => Engines::MySql(
             MySql::connect(config.into())
                 .await
-                .map_err(|e| EnginesError(e.into()))?,
+                .map_err(EnginesError::without_state)?,
         ),
         EngineConfig::Postgres => Engines::Postgres(
             PostgresSimple::connect(config.into())
                 .await
-                .map_err(|e| EnginesError(e.into()))?,
+                .map_err(EnginesError::without_state)?,
         ),
         EngineConfig::PostgresExtended => Engines::PostgresExtended(
             PostgresExtended::connect(config.into())
                 .await
-                .map_err(|e| EnginesError(e.into()))?,
+                .map_err(EnginesError::without_state)?,
         ),
         EngineConfig::External(cmd_tmpl) => {
             let (host, port) = config.random_addr();
@@ -98,24 +98,36 @@ pub(crate) async fn connect(
             Engines::External(
                 ExternalDriver::connect(cmd)
                     .await
-                    .map_err(|e| EnginesError(e.into()))?,
+                    .map_err(EnginesError::without_state)?,
             )
         }
     })
 }
 
 #[derive(Debug)]
-pub(crate) struct EnginesError(anyhow::Error);
+pub(crate) struct EnginesError {
+    error: anyhow::Error,
+    sqlstate: Option<String>,
+}
+
+impl EnginesError {
+    fn without_state(error: impl Into<anyhow::Error>) -> Self {
+        Self {
+            error: error.into(),
+            sqlstate: None,
+        }
+    }
+}
 
 impl Display for EnginesError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        self.0.fmt(f)
+        self.error.fmt(f)
     }
 }
 
 impl std::error::Error for EnginesError {
     fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
-        self.0.source()
+        self.error.source()
     }
 }
 
@@ -130,16 +142,21 @@ macro_rules! dispatch_engines {
     }};
 }
 
+fn error_sql_state<E: AsyncDB>(_engine: &E, error: &E::Error) -> Option<String> {
+    E::error_sql_state(error)
+}
+
 #[async_trait]
 impl AsyncDB for Engines {
     type Error = EnginesError;
     type ColumnType = DefaultColumnType;
 
     async fn run(&mut self, sql: &str) -> Result<DBOutput<Self::ColumnType>, Self::Error> {
         dispatch_engines!(self, e, {
-            e.run(sql)
-                .await
-                .map_err(|e| EnginesError(anyhow::Error::from(e)))
+            e.run(sql).await.map_err(|error| EnginesError {
+                sqlstate: error_sql_state(e, &error),
+                error: anyhow::Error::from(error),
+            })
         })
     }
 
@@ -158,4 +175,8 @@ impl AsyncDB for Engines {
     async fn shutdown(&mut self) {
         dispatch_engines!(self, e, { e.shutdown().await })
     }
+
+    fn error_sql_state(err: &Self::Error) -> Option<String> {
+        err.sqlstate.clone()
+    }
 }

sqllogictest-engines/src/mysql.rs

@@ -81,4 +81,12 @@ impl sqllogictest::AsyncDB for MySql {
     async fn run_command(command: Command) -> std::io::Result<std::process::Output> {
         tokio::process::Command::from(command).output().await
     }
+
+    fn error_sql_state(err: &Self::Error) -> Option<String> {
+        if let mysql_async::Error::Server(err) = err {
+            Some(err.state.clone())
+        } else {
+            None
+        }
+    }
 }

sqllogictest-engines/src/postgres/extended.rs

@@ -326,4 +326,8 @@ impl sqllogictest::AsyncDB for Postgres<Extended> {
     async fn run_command(command: Command) -> std::io::Result<std::process::Output> {
         tokio::process::Command::from(command).output().await
     }
+
+    fn error_sql_state(err: &Self::Error) -> Option<String> {
+        err.code().map(|s| s.code().to_owned())
+    }
 }

sqllogictest-engines/src/postgres/simple.rs

@@ -77,4 +77,8 @@ impl sqllogictest::AsyncDB for Postgres<Simple> {
     async fn run_command(command: Command) -> std::io::Result<std::process::Output> {
         tokio::process::Command::from(command).output().await
     }
+
+    fn error_sql_state(err: &Self::Error) -> Option<String> {
+        err.code().map(|s| s.code().to_owned())
+    }
 }

sqllogictest/src/parser.rs

@@ -377,12 +377,28 @@ pub enum ExpectedError {
     /// The actual error message that's exactly the same as the expected one is considered as a
     /// match.
     Multiline(String),
+    /// An expected SQL state code.
+    ///
+    /// The actual SQL state that matches the expected one is considered as a match.
+    SqlState(String),
 }
 
 impl ExpectedError {
     /// Parses an inline regex variant from tokens.
     fn parse_inline_tokens(tokens: &[&str]) -> Result<Self, ParseErrorKind> {
-        Self::new_inline(tokens.join(" "))
+        let joined = tokens.join(" ");
+
+        // Check if this is a sqlstate error pattern: error(sqlstate)
+        if let Some(captures) = regex::Regex::new(r"^\(([^)]+)\)$")
+            .unwrap()
+            .captures(&joined)
+        {
+            if let Some(sqlstate) = captures.get(1) {
+                return Ok(Self::SqlState(sqlstate.as_str().to_string()));
+            }
+        }
+
+        Self::new_inline(joined)
     }
 
     /// Creates an inline expected error message from a regex string.
@@ -406,8 +422,10 @@ impl ExpectedError {
     /// Unparses the expected message after `statement`.
     fn fmt_inline(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "error")?;
-        if let Self::Inline(regex) = self {
-            write!(f, " {regex}")?;
+        match self {
+            Self::Inline(regex) => write!(f, " {regex}")?,
+            Self::SqlState(sqlstate) => write!(f, " ({sqlstate})")?,
+            Self::Empty | Self::Multiline(_) => {}
         }
         Ok(())
     }
@@ -423,11 +441,14 @@ impl ExpectedError {
     }
 
     /// Returns whether the given error message matches the expected one.
-    pub fn is_match(&self, err: &str) -> bool {
+    pub fn is_match(&self, err: &str, sqlstate: Option<&str>) -> bool {
         match self {
             Self::Empty => true,
             Self::Inline(regex) => regex.is_match(err),
             Self::Multiline(results) => results.trim() == err.trim(),
+            Self::SqlState(expected_state) => {
+                sqlstate.map_or(false, |state| state == expected_state)
+            }
         }
     }
 
@@ -460,6 +481,7 @@ impl std::fmt::Display for ExpectedError {
             ExpectedError::Empty => write!(f, "(any)"),
             ExpectedError::Inline(regex) => write!(f, "(regex) {}", regex),
             ExpectedError::Multiline(results) => write!(f, "(multiline) {}", results.trim()),
+            ExpectedError::SqlState(sqlstate) => write!(f, "(sqlstate) {}", sqlstate),
         }
     }
 }
@@ -470,6 +492,7 @@ impl PartialEq for ExpectedError {
             (Self::Empty, Self::Empty) => true,
             (Self::Inline(l0), Self::Inline(r0)) => l0.as_str() == r0.as_str(),
             (Self::Multiline(l0), Self::Multiline(r0)) => l0 == r0,
+            (Self::SqlState(l0), Self::SqlState(r0)) => l0 == r0,
             _ => false,
         }
     }

sqllogictest/src/runner.rs

@@ -98,6 +98,11 @@ pub trait AsyncDB {
     async fn run_command(mut command: Command) -> std::io::Result<std::process::Output> {
         command.output()
     }
+
+    /// Extract the SQL state from the error.
+    fn error_sql_state(_err: &Self::Error) -> Option<String> {
+        None
+    }
 }
 
 /// The database to be tested.
@@ -117,6 +122,11 @@ pub trait DB {
     fn engine_name(&self) -> &str {
         ""
     }
+
+    /// Extract the SQL state from the error.
+    fn error_sql_state(_err: &Self::Error) -> Option<String> {
+        None
+    }
 }
 
 /// Compat-layer for the new AsyncDB and DB trait
@@ -139,6 +149,10 @@ where
     fn engine_name(&self) -> &str {
         D::engine_name(self)
     }
+
+    fn error_sql_state(err: &Self::Error) -> Option<String> {
+        D::error_sql_state(err)
+    }
 }
 
 /// The error type for running sqllogictest.
@@ -282,12 +296,14 @@ pub enum TestErrorKind {
         actual_stdout: String,
     },
     // 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}")]
+    #[error("{kind} is expected to fail with error:\n\t{expected_err}\nbut got error:\n\t{}{err}\n[SQL] {sql}", .actual_sqlstate.as_ref().map(|s| format!("(sqlstate {s}) ")).unwrap_or_default())]
     ErrorMismatch {
         sql: String,
         err: AnyError,
         expected_err: String,
         kind: RecordKind,
+        /// The actual SQL state when the expected error was a SqlState type
+        actual_sqlstate: Option<String>,
     },
     #[error("statement is expected to affect {expected} rows, but actually {actual}\n[SQL] {sql}")]
     StatementResultMismatch {
@@ -355,12 +371,16 @@ impl Display for TestErrorKindDisplay<'_> {
                 err,
                 expected_err,
                 kind,
-            } => write!(
-                f,
-                "{kind} is expected to fail with error:\n\t{}\nbut got error:\n\t{}\n[SQL] {sql}",
-                expected_err.bright_green(),
-                err.bright_red(),
-            ),
+                actual_sqlstate,
+            } => {
+                write!(
+                    f,
+                    "{kind} is expected to fail with error:\n\t{}\nbut got error:\n\t{}{}\n[SQL] {sql}",
+                    expected_err.bright_green(),
+                    actual_sqlstate.as_ref().map(|s| format!("(sqlstate {s}) ")).unwrap_or_default(),
+                    err.bright_red(),
+                )
+            }
             TestErrorKind::QueryResultMismatch {
                 sql,
                 expected,
@@ -1112,12 +1132,16 @@ impl<D: AsyncDB, M: MakeConnection<Conn = D>> Runner<D, M> {
                 }
                 (None, StatementExpect::Ok) => {}
                 (Some(e), StatementExpect::Error(expected_error)) => {
-                    if !expected_error.is_match(&e.to_string()) {
+                    let sqlstate = e
+                        .downcast_ref::<D::Error>()
+                        .and_then(|concrete_err| D::error_sql_state(concrete_err));
+                    if !expected_error.is_match(&e.to_string(), sqlstate.as_deref()) {
                         return Err(TestErrorKind::ErrorMismatch {
                             sql,
                             err: Arc::clone(e),
                             expected_err: expected_error.to_string(),
                             kind: RecordKind::Statement,
+                            actual_sqlstate: sqlstate,
                         }
                         .at(loc));
                     }
@@ -1151,12 +1175,16 @@ impl<D: AsyncDB, M: MakeConnection<Conn = D>> Runner<D, M> {
                         .at(loc));
                     }
                     (Some(e), QueryExpect::Error(expected_error)) => {
-                        if !expected_error.is_match(&e.to_string()) {
+                        let sqlstate = e
+                            .downcast_ref::<D::Error>()
+                            .and_then(|concrete_err| D::error_sql_state(concrete_err));
+                        if !expected_error.is_match(&e.to_string(), sqlstate.as_deref()) {
                             return Err(TestErrorKind::ErrorMismatch {
                                 sql,
                                 err: Arc::clone(e),
                                 expected_err: expected_error.to_string(),
                                 kind: RecordKind::Query,
+                                actual_sqlstate: sqlstate,
                             }
                             .at(loc));
                         }
@@ -1701,7 +1729,7 @@ pub fn update_record_with_output<T: ColumnType>(
             }),
             // Error match
             (Some(e), StatementExpect::Error(expected_error))
-                if expected_error.is_match(&e.to_string()) =>
+                if expected_error.is_match(&e.to_string(), None) =>
             {
                 None
             }
@@ -1738,7 +1766,7 @@ pub fn update_record_with_output<T: ColumnType>(
         ) => match (error, expected) {
             // Error match
             (Some(e), QueryExpect::Error(expected_error))
-                if expected_error.is_match(&e.to_string()) =>
+                if expected_error.is_match(&e.to_string(), None) =>
             {
                 None
             }