feat(func): impl quote_literal & quote_nullable (#16807)

risingwavelabs · May 23, 2024 · 657a1d5 · 657a1d5
1 parent da162ce
commit 657a1d5
Show file tree

Hide file tree

Showing 7 changed files with 100 additions and 7 deletions.
diff --git a/proto/expr.proto b/proto/expr.proto
@@ -194,6 +194,8 @@ message ExprNode {
     ENCRYPT = 325;
     INET_ATON = 328;
     INET_NTOA = 329;
+    QUOTE_LITERAL = 330;
+    QUOTE_NULLABLE = 331;
 
     // Unary operators
     NEG = 401;

diff --git a/src/common/src/types/to_text.rs b/src/common/src/types/to_text.rs
@@ -47,15 +47,18 @@ pub trait ToText {
     /// - `ScalarRefImpl::Float32` -> `DataType::Float32`
     /// - `ScalarRefImpl::Float64` -> `DataType::Float64`
     /// - `ScalarRefImpl::Decimal` -> `DataType::Decimal`
-    /// - `ScalarRefImpl::Boolean` -> `DataType::Boolean`
+    /// - `ScalarRefImpl::Bool` -> `DataType::Boolean`
     /// - `ScalarRefImpl::Utf8` -> `DataType::Varchar`
     /// - `ScalarRefImpl::Bytea` -> `DataType::Bytea`
     /// - `ScalarRefImpl::Date` -> `DataType::Date`
     /// - `ScalarRefImpl::Time` -> `DataType::Time`
     /// - `ScalarRefImpl::Timestamp` -> `DataType::Timestamp`
+    /// - `ScalarRefImpl::Timestamptz` -> `DataType::Timestamptz`
     /// - `ScalarRefImpl::Interval` -> `DataType::Interval`
+    /// - `ScalarRefImpl::Jsonb` -> `DataType::Jsonb`
     /// - `ScalarRefImpl::List` -> `DataType::List`
     /// - `ScalarRefImpl::Struct` -> `DataType::Struct`
+    /// - `ScalarRefImpl::Serial` -> `DataType::Serial`
     fn to_text(&self) -> String {
         let mut s = String::new();
         self.write(&mut s).unwrap();

diff --git a/src/expr/impl/src/scalar/string.rs b/src/expr/impl/src/scalar/string.rs
@@ -472,6 +472,74 @@ pub fn right(s: &str, n: i32, writer: &mut impl Write) {
         .for_each(|c| writer.write_char(c).unwrap());
 }
 
+/// `quote_literal(string text)`
+/// `quote_literal(value anyelement)`
+///
+/// Return the given string suitably quoted to be used as a string literal in an SQL statement
+/// string. Embedded single-quotes and backslashes are properly doubled.
+/// Note that `quote_literal` returns null on null input; if the argument might be null,
+/// `quote_nullable` is often more suitable.
+///
+/// # Example
+///
+/// Note that the quotes are part of the output string.
+///
+/// ```slt
+/// query T
+/// select quote_literal(E'O\'Reilly')
+/// ----
+/// 'O''Reilly'
+///
+/// query T
+/// select quote_literal(E'C:\\Windows\\')
+/// ----
+/// E'C:\\Windows\\'
+///
+/// query T
+/// select quote_literal(42.5)
+/// ----
+/// '42.5'
+///
+/// query T
+/// select quote_literal('hello'::bytea);
+/// ----
+/// E'\\x68656c6c6f'
+///
+/// query T
+/// select quote_literal('{"hello":"world","foo":233}'::jsonb);
+/// ----
+/// '{"foo": 233, "hello": "world"}'
+/// ```
+#[function("quote_literal(varchar) -> varchar")]
+pub fn quote_literal(s: &str, writer: &mut impl Write) {
+    if s.contains('\\') {
+        // use escape format: E'...'
+        write!(writer, "E").unwrap();
+    }
+    write!(writer, "'").unwrap();
+    for c in s.chars() {
+        match c {
+            '\'' => write!(writer, "''").unwrap(),
+            '\\' => write!(writer, "\\\\").unwrap(),
+            _ => write!(writer, "{}", c).unwrap(),
+        }
+    }
+    write!(writer, "'").unwrap();
+}
+
+/// `quote_nullable(string text)`
+///
+/// Return the given string suitably quoted to be used as a string literal in an SQL statement
+/// string; or, if the argument is null, return NULL.
+/// Embedded single-quotes and backslashes are properly doubled.
+#[function("quote_nullable(varchar) -> varchar")]
+pub fn quote_nullable(s: Option<&str>, writer: &mut impl Write) {
+    match s {
+        Some(s) => quote_literal(s, writer),
+        None => write!(writer, "NULL").unwrap(),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

diff --git a/src/frontend/src/binder/expr/function.rs b/src/frontend/src/binder/expr/function.rs
@@ -40,8 +40,8 @@ use crate::binder::bind_context::Clause;
 use crate::binder::{Binder, UdfContext};
 use crate::error::{ErrorCode, Result, RwError};
 use crate::expr::{
-    AggCall, Expr, ExprImpl, ExprType, FunctionCall, FunctionCallWithLambda, Literal, Now, OrderBy,
-    TableFunction, TableFunctionType, UserDefinedFunction, WindowFunction,
+    AggCall, CastContext, Expr, ExprImpl, ExprType, FunctionCall, FunctionCallWithLambda, Literal,
+    Now, OrderBy, TableFunction, TableFunctionType, UserDefinedFunction, WindowFunction,
 };
 use crate::utils::Condition;
 
@@ -1012,6 +1012,22 @@ impl Binder {
                 ("to_ascii", raw_call(ExprType::ToAscii)),
                 ("to_hex", raw_call(ExprType::ToHex)),
                 ("quote_ident", raw_call(ExprType::QuoteIdent)),
+                ("quote_literal", guard_by_len(1, raw(|_binder, mut inputs| {
+                    if inputs[0].return_type() != DataType::Varchar {
+                        // Support `quote_literal(any)` by converting it to `quote_literal(any::text)`
+                        // Ref. https://github.com/postgres/postgres/blob/REL_16_1/src/include/catalog/pg_proc.dat#L4641
+                        FunctionCall::cast_mut(&mut inputs[0], DataType::Varchar, CastContext::Explicit)?;
+                    }
+                    Ok(FunctionCall::new_unchecked(ExprType::QuoteLiteral, inputs, DataType::Varchar).into())
+                }))),
+                ("quote_nullable", guard_by_len(1, raw(|_binder, mut inputs| {
+                    if inputs[0].return_type() != DataType::Varchar {
+                        // Support `quote_nullable(any)` by converting it to `quote_nullable(any::text)`
+                        // Ref. https://github.com/postgres/postgres/blob/REL_16_1/src/include/catalog/pg_proc.dat#L4650
+                        FunctionCall::cast_mut(&mut inputs[0], DataType::Varchar, CastContext::Explicit)?;
+                    }
+                    Ok(FunctionCall::new_unchecked(ExprType::QuoteNullable, inputs, DataType::Varchar).into())
+                }))),
                 ("string_to_array", raw_call(ExprType::StringToArray)),
                 ("encode", raw_call(ExprType::Encode)),
                 ("decode", raw_call(ExprType::Decode)),

diff --git a/src/frontend/src/expr/pure.rs b/src/frontend/src/expr/pure.rs
@@ -245,7 +245,9 @@ impl ExprVisitor for ImpureAnalyzer {
             | Type::ConvertTo
             | Type::IcebergTransform
             | Type::InetNtoa
-            | Type::InetAton =>
+            | Type::InetAton
+            | Type::QuoteLiteral
+            | Type::QuoteNullable =>
             // expression output is deterministic(same result for the same input)
             {
                 func_call

diff --git a/src/frontend/src/optimizer/plan_expr_visitor/strong.rs b/src/frontend/src/optimizer/plan_expr_visitor/strong.rs
@@ -81,6 +81,7 @@ impl Strong {
             | ExprType::IsDistinctFrom
             | ExprType::IsNotDistinctFrom
             | ExprType::IsTrue
+            | ExprType::QuoteNullable
             | ExprType::IsNotTrue
             | ExprType::IsFalse
             | ExprType::IsNotFalse => false,
@@ -180,6 +181,7 @@ impl Strong {
             | ExprType::ToAscii
             | ExprType::ToHex
             | ExprType::QuoteIdent
+            | ExprType::QuoteLiteral
             | ExprType::Sin
             | ExprType::Cos
             | ExprType::Tan

diff --git a/src/tests/regress/data/sql/text.sql b/src/tests/regress/data/sql/text.sql
@@ -43,9 +43,9 @@ select concat_ws('',10,20,null,30);
 select concat_ws(NULL,10,20,null,30) is null;
 select reverse('abcde');
 select i, left('ahoj', i), right('ahoj', i) from generate_series(-5, 5) t(i) order by i;
---@ select quote_literal('');
---@ select quote_literal('abc''');
---@ select quote_literal(e'\\');
+select quote_literal('');
+select quote_literal('abc''');
+select quote_literal(e'\\');
 -- check variadic labeled argument
 --@ select concat(variadic array[1,2,3]);
 --@ select concat_ws(',', variadic array[1,2,3]);