From af8fd10b26c3494377d20a86f235dd185b00f5f5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefan=20Gr=C3=B6nke?= <stefan@gronke.net>
Date: Sun, 13 Oct 2024 15:58:55 +0200
Subject: [PATCH] implement IfElseStatement

---
 src/backend/query_builder.rs | 15 +++++++++++++
 src/expr.rs                  |  3 ++-
 src/if_else.rs               | 35 ++++++++++++++++++++++++++++++
 src/lib.rs                   |  2 ++
 tests/mysql/if_else.rs       | 42 ++++++++++++++++++++++++++++++++++++
 tests/mysql/mod.rs           |  1 +
 6 files changed, 97 insertions(+), 1 deletion(-)
 create mode 100644 src/if_else.rs
 create mode 100644 tests/mysql/if_else.rs

diff --git a/src/backend/query_builder.rs b/src/backend/query_builder.rs
index 7bf29b55..785971b4 100644
--- a/src/backend/query_builder.rs
+++ b/src/backend/query_builder.rs
@@ -387,9 +387,24 @@ pub trait QueryBuilder:
             SimpleExpr::Constant(val) => {
                 self.prepare_constant(val, sql);
             }
+            SimpleExpr::IfElse(val) => {
+                self.prepare_if_else_statement(val, sql);
+            }
         }
     }
 
+    fn prepare_if_else_statement(&self, val: &Box<IfElseStatement>, sql: &mut dyn SqlWriter) {
+        write!(sql, "IF ").unwrap();
+        self.prepare_simple_expr(&val.when, sql);
+        write!(sql, " THEN\n").unwrap();
+        self.prepare_simple_expr(&val.then, sql);
+        if let Some(otherwise) = &val.otherwise {
+            write!(sql, "\nELSE\n").unwrap();
+            self.prepare_simple_expr(otherwise, sql);
+        };
+        write!(sql, "\nEND IF").unwrap();
+    }
+
     /// Translate [`CaseStatement`] into SQL statement.
     fn prepare_case_statement(&self, stmts: &CaseStatement, sql: &mut dyn SqlWriter) {
         write!(sql, "(CASE").unwrap();
diff --git a/src/expr.rs b/src/expr.rs
index 06cf08ed..8a9f1072 100644
--- a/src/expr.rs
+++ b/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::{func::*, query::*, types::*, value::*, if_else::*};
 
 /// Helper to build a [`SimpleExpr`].
 #[derive(Debug, Clone)]
@@ -35,6 +35,7 @@ pub enum SimpleExpr {
     AsEnum(DynIden, Box<SimpleExpr>),
     Case(Box<CaseStatement>),
     Constant(Value),
+    IfElse(Box<IfElseStatement>),
 }
 
 /// "Operator" methods for building complex expressions.
diff --git a/src/if_else.rs b/src/if_else.rs
new file mode 100644
index 00000000..b6f9779d
--- /dev/null
+++ b/src/if_else.rs
@@ -0,0 +1,35 @@
+use crate::{QueryBuilder, SimpleExpr};
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct IfElseStatement {
+    pub when: SimpleExpr,
+    pub then: SimpleExpr,
+    pub otherwise: Option<SimpleExpr>
+}
+
+impl IfElseStatement {
+
+    pub fn new(when: SimpleExpr, then: SimpleExpr, otherwise: Option<SimpleExpr>) -> Self {
+        Self {
+            when,
+            then,
+            otherwise
+        }
+    }
+
+    pub fn to_string<T: QueryBuilder>(&self, query_builder: T) -> String {
+        let mut sql = String::with_capacity(256);
+        query_builder.prepare_if_else_statement(&Box::new(self.clone()), &mut sql);
+        sql
+    }
+
+}
+pub trait IfElseStatementBuilder {
+    /// Build corresponding SQL statement for certain database backend and return SQL string
+    fn build<T: QueryBuilder>(&self, query_builder: T) -> String;
+
+    /// Build corresponding SQL statement for certain database backend and return SQL string
+    fn to_string<T: QueryBuilder>(&self, query_builder: T) -> String {
+        self.build(query_builder)
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index 15e1a189..e4e8c7cf 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -826,6 +826,7 @@ pub mod table;
 pub mod token;
 pub mod types;
 pub mod value;
+pub mod if_else;
 
 #[doc(hidden)]
 #[cfg(feature = "tests-cfg")]
@@ -843,6 +844,7 @@ pub use table::*;
 pub use token::*;
 pub use types::*;
 pub use value::*;
+pub use if_else::*;
 
 #[cfg(feature = "derive")]
 pub use sea_query_derive::{enum_def, Iden, IdenStatic};
diff --git a/tests/mysql/if_else.rs b/tests/mysql/if_else.rs
new file mode 100644
index 00000000..15c9de4b
--- /dev/null
+++ b/tests/mysql/if_else.rs
@@ -0,0 +1,42 @@
+use super::*;
+use pretty_assertions::assert_eq;
+
+#[test]
+fn if_without_else() {
+    let query = Query::select().column(Asterisk).from(Glyph::Table).take();
+    let then = SimpleExpr::SubQuery(None, Box::new(query.into_sub_query_statement()));
+    let if_statement = IfElseStatement::new(
+        Expr::col(Glyph::Id).eq(1),
+        then,
+        None
+    );
+    assert_eq!(
+        if_statement.to_string(MysqlQueryBuilder),
+        [
+            "IF `id` = 1 THEN",
+            "(SELECT * FROM `glyph`)",
+            "END IF"
+        ].join("\n")
+    )
+}
+
+#[test]
+fn if_with_else() {
+    let query = Query::select().column(Asterisk).from(Glyph::Table).take();
+    let then = SimpleExpr::SubQuery(None, Box::new(query.into_sub_query_statement()));
+    let if_statement = IfElseStatement::new(
+        Expr::col(Glyph::Id).eq(1),
+        then,
+        Some(Expr::val("23").into())
+    );
+    assert_eq!(
+        if_statement.to_string(MysqlQueryBuilder),
+        [
+            "IF `id` = 1 THEN",
+            "(SELECT * FROM `glyph`)",
+            "ELSE",
+            "'23'",
+            "END IF"
+        ].join("\n")
+    )
+}
diff --git a/tests/mysql/mod.rs b/tests/mysql/mod.rs
index d717774f..3facb1e5 100644
--- a/tests/mysql/mod.rs
+++ b/tests/mysql/mod.rs
@@ -4,6 +4,7 @@ mod foreign_key;
 mod index;
 mod query;
 mod table;
+mod if_else;
 
 #[path = "../common.rs"]
 mod common;