Implement LPAD and RPAD functions (gluesql#311)

jekalhong · Aug 23, 2021 · 83db495 · 83db495
1 parent ddef417
commit 83db495
Show file tree

Hide file tree

Showing 7 changed files with 276 additions and 0 deletions.
diff --git a/src/ast/function.rs b/src/ast/function.rs
@@ -14,6 +14,18 @@ pub enum Function {
     Left { expr: Expr, size: Expr },
     #[strum(to_string = "RIGHT")]
     Right { expr: Expr, size: Expr },
+    #[strum(to_string = "LPAD")]
+    Lpad {
+        expr: Expr,
+        size: Expr,
+        fill: Option<Expr>,
+    },
+    #[strum(to_string = "RPAD")]
+    Rpad {
+        expr: Expr,
+        size: Expr,
+        fill: Option<Expr>,
+    },
     #[strum(to_string = "CEIL")]
     Ceil(Expr),
     #[strum(to_string = "ROUND")]

diff --git a/src/executor/evaluate/mod.rs b/src/executor/evaluate/mod.rs
@@ -302,6 +302,59 @@ async fn evaluate_function<'a, T: 'static + Debug>(
 
             Ok(Evaluated::from(Value::Str(converted)))
         }
+        .map(Evaluated::from),
+        Function::Lpad { expr, size, fill } | Function::Rpad { expr, size, fill } => {
+            let name = if matches!(func, Function::Lpad { .. }) {
+                "LPAD"
+            } else {
+                "RPAD"
+            };
+
+            let string = match eval_to_str(expr).await? {
+                Nullable::Value(v) => v,
+                Nullable::Null => {
+                    return Ok(Evaluated::from(Value::Null));
+                }
+            };
+
+            let size = match eval(size).await?.try_into()? {
+                Value::I64(number) => usize::try_from(number)
+                    .map_err(|_| EvaluateError::FunctionRequiresUSizeValue(name.to_owned()))?,
+                Value::Null => {
+                    return Ok(Evaluated::from(Value::Null));
+                }
+                _ => {
+                    return Err(EvaluateError::FunctionRequiresIntegerValue(name.to_owned()).into());
+                }
+            };
+
+            let fill = match fill {
+                Some(expr) => match eval_to_str(expr).await? {
+                    Nullable::Value(v) => v,
+                    Nullable::Null => {
+                        return Ok(Evaluated::from(Value::Null));
+                    }
+                },
+                None => " ".to_string(),
+            };
+
+            let result = if size > string.len() {
+                let padding_size = size - string.len();
+                let repeat_count = padding_size / fill.len();
+                let plus_count = padding_size % fill.len();
+                let fill = fill.repeat(repeat_count) + &fill[0..plus_count];
+
+                if name == "LPAD" {
+                    fill + &string
+                } else {
+                    string + &fill
+                }
+            } else {
+                string[0..size].to_string()
+            };
+
+            Ok(Evaluated::from(Value::Str(result)))
+        }
         Function::Ceil(expr) => match eval_to_float(expr).await? {
             Nullable::Value(v) => Ok(Value::F64(v.ceil())),
             Nullable::Null => Ok(Value::Null),

diff --git a/src/tests/function/lpad_rpad.rs b/src/tests/function/lpad_rpad.rs
@@ -0,0 +1,153 @@
+use crate::*;
+
+test_case!(lpad_rpad, async move {
+    use Value::{Null, Str};
+
+    let test_cases = vec![
+        ("CREATE TABLE Item (name TEXT)", Ok(Payload::Create)),
+        (
+            r#"INSERT INTO Item VALUES ("hello")"#,
+            Ok(Payload::Insert(1)),
+        ),
+        (
+            "CREATE TABLE NullName (name TEXT NULL)",
+            Ok(Payload::Create),
+        ),
+        (
+            r#"INSERT INTO NullName VALUES (NULL)"#,
+            Ok(Payload::Insert(1)),
+        ),
+        (
+            "CREATE TABLE NullNumber (number INTEGER NULL)",
+            Ok(Payload::Create),
+        ),
+        (
+            r#"INSERT INTO NullNumber VALUES (NULL)"#,
+            Ok(Payload::Insert(1)),
+        ),
+        (
+            "SELECT LPAD(name, 10), RPAD(name, 10) FROM Item",
+            Ok(select!(
+                "LPAD(name, 10)"         | "RPAD(name, 10)"
+                Str                      | Str;
+                "     hello".to_owned()    "hello     ".to_owned()
+            )),
+        ),
+        (
+            "SELECT LPAD(name, 10, 'ab'), RPAD(name, 10, 'ab') FROM Item",
+            Ok(select!(
+                "LPAD(name, 10, 'ab')"   | "RPAD(name, 10, 'ab')"
+                Str                      | Str;
+                "ababahello".to_owned()    "helloababa".to_owned()
+            )),
+        ),
+        (
+            "SELECT LPAD(name, 3), RPAD(name, 3) FROM Item",
+            Ok(select!(
+                "LPAD(name, 3)"   | "RPAD(name, 3)"
+                Str               | Str;
+                "hel".to_owned()    "hel".to_owned()
+            )),
+        ),
+        (
+            "SELECT LPAD(name, 3, 'ab'), RPAD(name, 3, 'ab') FROM Item",
+            Ok(select!(
+                "LPAD(name, 3, 'ab')"   | "RPAD(name, 3, 'ab')"
+                Str                     | Str;
+                "hel".to_owned()          "hel".to_owned()
+            )),
+        ),
+        (
+            "SELECT LPAD(name, 10, 'ab') AS lpad FROM NullName",
+            Ok(select_with_null!(lpad; Null)),
+        ),
+        (
+            "SELECT RPAD(name, 10, 'ab') AS rpad FROM NullName",
+            Ok(select_with_null!(rpad; Null)),
+        ),
+        (
+            "SELECT LPAD('hello', number, 'ab') AS lpad FROM NullNumber",
+            Ok(select_with_null!(lpad; Null)),
+        ),
+        (
+            "SELECT RPAD('hello', number, 'ab') AS rpad FROM NullNumber",
+            Ok(select_with_null!(rpad; Null)),
+        ),
+        (
+            "SELECT LPAD('hello', 10, name) AS lpad FROM NullName",
+            Ok(select_with_null!(lpad; Null)),
+        ),
+        (
+            "SELECT RPAD('hello', 10, name) AS rpad FROM NullName",
+            Ok(select_with_null!(rpad; Null)),
+        ),
+        (
+            "SELECT LPAD(name) FROM Item",
+            Err(TranslateError::FunctionArgsLengthNotWithinRange {
+                name: "LPAD".to_string(),
+                expected_minimum: 2,
+                expected_maximum: 3,
+                found: 1,
+            }
+            .into()),
+        ),
+        (
+            "SELECT RPAD(name) FROM Item",
+            Err(TranslateError::FunctionArgsLengthNotWithinRange {
+                name: "RPAD".to_string(),
+                expected_minimum: 2,
+                expected_maximum: 3,
+                found: 1,
+            }
+            .into()),
+        ),
+        (
+            "SELECT LPAD(name, 10, 'ab', 'cd') FROM Item",
+            Err(TranslateError::FunctionArgsLengthNotWithinRange {
+                name: "LPAD".to_string(),
+                expected_minimum: 2,
+                expected_maximum: 3,
+                found: 4,
+            }
+            .into()),
+        ),
+        (
+            "SELECT RPAD(name, 10, 'ab', 'cd') FROM Item",
+            Err(TranslateError::FunctionArgsLengthNotWithinRange {
+                name: "RPAD".to_string(),
+                expected_minimum: 2,
+                expected_maximum: 3,
+                found: 4,
+            }
+            .into()),
+        ),
+        (
+            "SELECT LPAD(1, 10, 'ab') FROM Item",
+            Err(EvaluateError::FunctionRequiresStringValue("LPAD".to_string()).into()),
+        ),
+        (
+            "SELECT RPAD(1, 10, 'ab') FROM Item",
+            Err(EvaluateError::FunctionRequiresStringValue("RPAD".to_string()).into()),
+        ),
+        (
+            "SELECT LPAD(name, -10, 'ab') FROM Item",
+            Err(EvaluateError::FunctionRequiresUSizeValue("LPAD".to_string()).into()),
+        ),
+        (
+            "SELECT RPAD(name, -10, 'ab') FROM Item",
+            Err(EvaluateError::FunctionRequiresUSizeValue("RPAD".to_string()).into()),
+        ),
+        (
+            "SELECT LPAD(name, 10.1, 'ab') FROM Item",
+            Err(EvaluateError::FunctionRequiresIntegerValue("LPAD".to_string()).into()),
+        ),
+        (
+            "SELECT RPAD(name, 10.1, 'ab') FROM Item",
+            Err(EvaluateError::FunctionRequiresIntegerValue("RPAD".to_string()).into()),
+        ),
+    ];
+
+    for (sql, expected) in test_cases.into_iter() {
+        test!(expected, sql);
+    }
+});
diff --git a/src/tests/function/mod.rs b/src/tests/function/mod.rs
@@ -4,6 +4,7 @@ pub mod div_mod;
 pub mod floor;
 pub mod gcd_lcm;
 pub mod left_right;
+pub mod lpad_rpad;
 pub mod math_function;
 pub mod round;
 pub mod trim;

diff --git a/src/tests/mod.rs b/src/tests/mod.rs
@@ -58,6 +58,7 @@ macro_rules! generate_tests {
         glue!(function_upper_lower, function::upper_lower::upper_lower);
         glue!(function_gcd_lcm, function::gcd_lcm::gcd_lcm);
         glue!(function_left_right, function::left_right::left_right);
+        glue!(function_lpad_rpad, function::lpad_rpad::lpad_rpad);
         glue!(function_trim, function::trim::trim);
         glue!(function_div_mod, function::div_mod::div_mod);
         glue!(function_cast_literal, function::cast::cast_literal);

diff --git a/src/translate/error.rs b/src/translate/error.rs
@@ -21,6 +21,14 @@ pub enum TranslateError {
         found: usize,
     },
 
+    #[error("function args.length not matching: {name}, expected: {expected_minimum} ~ {expected_maximum}, found: {found}")]
+    FunctionArgsLengthNotWithinRange {
+        name: String,
+        expected_minimum: usize,
+        expected_maximum: usize,
+        found: usize,
+    },
+
     #[error("named function arg is not supported")]
     NamedFunctionArgNotSupported,
 

diff --git a/src/translate/function.rs b/src/translate/function.rs
@@ -37,6 +37,20 @@ pub fn translate_function(sql_function: &SqlFunction) -> Result<Expr> {
         }
     };
 
+    let check_len_range = |name, found, expected_minimum, expected_maximum| -> Result<_> {
+        if found >= expected_minimum && found <= expected_maximum {
+            Ok(())
+        } else {
+            Err(TranslateError::FunctionArgsLengthNotWithinRange {
+                name,
+                expected_minimum,
+                expected_maximum,
+                found,
+            }
+            .into())
+        }
+    };
+
     macro_rules! aggr {
         ($aggregate: expr) => {{
             check_len(name, args.len(), 1)?;
@@ -78,6 +92,40 @@ pub fn translate_function(sql_function: &SqlFunction) -> Result<Expr> {
 
             Ok(Expr::Function(Box::new(Function::Right { expr, size })))
         }
+        "LPAD" => {
+            check_len_range(name, args.len(), 2, 3)?;
+
+            let expr = translate_expr(args[0])?;
+            let size = translate_expr(args[1])?;
+            let fill = if args.len() == 2 {
+                None
+            } else {
+                Some(translate_expr(args[2])?)
+            };
+
+            Ok(Expr::Function(Box::new(Function::Lpad {
+                expr,
+                size,
+                fill,
+            })))
+        }
+        "RPAD" => {
+            check_len_range(name, args.len(), 2, 3)?;
+
+            let expr = translate_expr(args[0])?;
+            let size = translate_expr(args[1])?;
+            let fill = if args.len() == 2 {
+                None
+            } else {
+                Some(translate_expr(args[2])?)
+            };
+
+            Ok(Expr::Function(Box::new(Function::Rpad {
+                expr,
+                size,
+                fill,
+            })))
+        }
         "CEIL" => func_with_one_arg!(Function::Ceil),
         "ROUND" => func_with_one_arg!(Function::Round),
         "FLOOR" => func_with_one_arg!(Function::Floor),