ExceptionStatement raises custom SQL errors

Author: gronke

src/backend/mysql/query.rs

@@ -197,6 +197,17 @@ impl QueryBuilder for MysqlQueryBuilder {
 
     fn prepare_returning(&self, _returning: &Option<ReturningClause>, _sql: &mut dyn SqlWriter) {}
 
+    fn prepare_exception_statement(&self, exception: &ExceptionStatement, sql: &mut dyn SqlWriter) {
+        let mut quoted_exception_message = String::new();
+        self.write_string_quoted(&exception.message, &mut quoted_exception_message);
+        write!(
+            sql,
+            "SIGNAL SQLSTATE '45000' SET MESSAGE_TEXT = {}",
+            quoted_exception_message
+        )
+        .unwrap();
+    }
+
     fn random_function(&self) -> &str {
         "RAND"
     }

src/backend/postgres/query.rs

@@ -178,6 +178,12 @@ impl QueryBuilder for PostgresQueryBuilder {
         sql.push_param(value.clone(), self as _);
     }
 
+    fn prepare_exception_statement(&self, exception: &ExceptionStatement, sql: &mut dyn SqlWriter) {
+        let mut quoted_exception_message = String::new();
+        self.write_string_quoted(&exception.message, &mut quoted_exception_message);
+        write!(sql, "RAISE EXCEPTION {}", quoted_exception_message).unwrap();
+    }
+
     fn write_string_quoted(&self, string: &str, buffer: &mut String) {
         let escaped = self.escape_string(string);
         let string = if escaped.find('\\').is_some() {

src/backend/query_builder.rs

@@ -453,6 +453,9 @@ pub trait QueryBuilder:
             SimpleExpr::Constant(val) => {
                 self.prepare_constant(val, sql);
             }
+            SimpleExpr::Exception(val) => {
+                self.prepare_exception_statement(val, sql);
+            }
         }
     }
 
@@ -1051,6 +1054,15 @@ pub trait QueryBuilder:
         }
     }
 
+    // Translate [`Exception`] into SQL statement.
+    fn prepare_exception_statement(
+        &self,
+        _exception: &ExceptionStatement,
+        _sql: &mut dyn SqlWriter,
+    ) {
+        panic!("Exception handling not implemented for this backend");
+    }
+
     /// Convert a SQL value into syntax-specific string
     fn value_to_string(&self, v: &Value) -> String {
         self.value_to_string_common(v)

src/backend/sqlite/query.rs

@@ -84,6 +84,12 @@ impl QueryBuilder for SqliteQueryBuilder {
         "MIN"
     }
 
+    fn prepare_exception_statement(&self, exception: &ExceptionStatement, sql: &mut dyn SqlWriter) {
+        let mut quoted_exception_message = String::new();
+        self.write_string_quoted(&exception.message, &mut quoted_exception_message);
+        write!(sql, "SELECT RAISE(ABORT, {})", quoted_exception_message).unwrap();
+    }
+
     fn char_length_function(&self) -> &str {
         "LENGTH"
     }

src/exception.rs

@@ -0,0 +1,46 @@
+//! Custom SQL exceptions and errors
+use inherent::inherent;
+
+use crate::backend::SchemaBuilder;
+
+/// SQL Exceptions
+#[derive(Debug, Clone, PartialEq)]
+pub struct ExceptionStatement {
+    pub(crate) message: String,
+}
+
+impl ExceptionStatement {
+    pub fn new(message: String) -> Self {
+        Self { message }
+    }
+}
+
+pub trait ExceptionStatementBuilder {
+    /// Build corresponding SQL statement for certain database backend and return SQL string
+    fn build<T: SchemaBuilder>(&self, schema_builder: T) -> String;
+
+    /// Build corresponding SQL statement for certain database backend and return SQL string
+    fn build_any(&self, schema_builder: &dyn SchemaBuilder) -> String;
+
+    /// Build corresponding SQL statement for certain database backend and return SQL string
+    fn to_string<T: SchemaBuilder>(&self, schema_builder: T) -> String {
+        self.build(schema_builder)
+    }
+}
+
+#[inherent]
+impl ExceptionStatementBuilder for ExceptionStatement {
+    pub fn build<T: SchemaBuilder>(&self, schema_builder: T) -> String {
+        let mut sql = String::with_capacity(256);
+        schema_builder.prepare_exception_statement(self, &mut sql);
+        sql
+    }
+
+    pub fn build_any(&self, schema_builder: &dyn SchemaBuilder) -> String {
+        let mut sql = String::with_capacity(256);
+        schema_builder.prepare_exception_statement(self, &mut sql);
+        sql
+    }
+
+    pub fn to_string<T: SchemaBuilder>(&self, schema_builder: T) -> String;
+}

src/expr.rs

@@ -4,7 +4,7 @@
 //!
 //! [`SimpleExpr`] is the expression common among select fields, where clauses and many other places.
 
-use crate::{func::*, query::*, types::*, value::*};
+use crate::{exception::ExceptionStatement, func::*, query::*, types::*, value::*};
 
 /// Helper to build a [`SimpleExpr`].
 #[derive(Debug, Clone)]
@@ -35,6 +35,7 @@ pub enum SimpleExpr {
     AsEnum(DynIden, Box<SimpleExpr>),
     Case(Box<CaseStatement>),
     Constant(Value),
+    Exception(ExceptionStatement),
 }
 
 /// "Operator" methods for building expressions.

src/lib.rs

@@ -809,6 +809,7 @@
 
 pub mod backend;
 pub mod error;
+pub mod exception;
 pub mod expr;
 pub mod extension;
 pub mod foreign_key;
@@ -827,6 +828,7 @@ pub mod value;
 pub mod tests_cfg;
 
 pub use backend::*;
+pub use exception::*;
 pub use expr::*;
 pub use foreign_key::*;
 pub use func::*;

tests/mysql/exception.rs

@@ -0,0 +1,20 @@
+use super::*;
+use pretty_assertions::assert_eq;
+
+#[test]
+fn signal_sqlstate() {
+    let message = "Some error occurred";
+    assert_eq!(
+        ExceptionStatement::new(message.to_string()).to_string(MysqlQueryBuilder),
+        format!("SIGNAL SQLSTATE '45000' SET MESSAGE_TEXT = '{message}'")
+    );
+}
+
+#[test]
+fn escapes_message() {
+    let unescaped_message = "Does this 'break'?";
+    assert_eq!(
+        ExceptionStatement::new(unescaped_message.to_string()).to_string(MysqlQueryBuilder),
+        format!("SIGNAL SQLSTATE '45000' SET MESSAGE_TEXT = 'Does this \\'break\\'?'")
+    );
+}

tests/mysql/mod.rs

@@ -1,5 +1,6 @@
 use sea_query::{extension::mysql::*, tests_cfg::*, *};
 
+mod exception;
 mod foreign_key;
 mod index;
 mod query;

tests/postgres/exception.rs

@@ -0,0 +1,20 @@
+use super::*;
+use pretty_assertions::assert_eq;
+
+#[test]
+fn raise_exception() {
+    let message = "Some error occurred";
+    assert_eq!(
+        ExceptionStatement::new(message.to_string()).to_string(PostgresQueryBuilder),
+        format!("RAISE EXCEPTION '{message}'")
+    );
+}
+
+#[test]
+fn escapes_message() {
+    let unescaped_message = "Does this 'break'?";
+    assert_eq!(
+        ExceptionStatement::new(unescaped_message.to_string()).to_string(PostgresQueryBuilder),
+        format!("RAISE EXCEPTION E'Does this \\'break\\'?'")
+    );
+}

tests/postgres/mod.rs

@@ -1,5 +1,6 @@
 use sea_query::{tests_cfg::*, *};
 
+mod exception;
 mod foreign_key;
 mod index;
 mod query;

tests/sqlite/exception.rs

@@ -0,0 +1,21 @@
+use super::*;
+use pretty_assertions::assert_eq;
+
+#[test]
+fn select_raise_abort() {
+    let message = "Some error occurred here";
+    assert_eq!(
+        ExceptionStatement::new(message.to_string()).to_string(SqliteQueryBuilder),
+        format!("SELECT RAISE(ABORT, '{}')", message)
+    );
+}
+
+#[test]
+fn escapes_message() {
+    let unescaped_message = "Does this 'break'?";
+    let escaped_message = "Does this ''break''?";
+    assert_eq!(
+        ExceptionStatement::new(unescaped_message.to_string()).to_string(SqliteQueryBuilder),
+        format!("SELECT RAISE(ABORT, '{}')", escaped_message)
+    );
+}

tests/sqlite/mod.rs

@@ -1,5 +1,6 @@
 use sea_query::{tests_cfg::*, *};
 
+mod exception;
 mod foreign_key;
 mod index;
 mod query;