feat: add Parameterized Queries

cli/Cargo.toml

@@ -25,7 +25,7 @@ clap = { version = "4.4", features = ["derive", "env"] }
 comfy-table = "7.1"
 csv = "1.3"
 ctrlc = { version = "3.2.3", features = ["termination"] }
-databend-common-ast = "0.1.3"
+databend-common-ast = "0.2.1"
 duckdb = { version = "0.10.2", features = ["bundled"] }
 fern = { version = "0.6", features = ["colored"] }
 indicatif = "0.17"

driver/Cargo.toml

@@ -34,7 +34,8 @@ databend-driver-core = { workspace = true }
 databend-driver-macros = { workspace = true }
 tokio-stream = { workspace = true }
 tonic = { workspace = true, optional = true }
-databend-common-ast = "0.1.3"
+databend-common-ast = "0.2.1"
+derive-visitor = { version = "0.4.0", features = ["std-types-drive"] }
 
 async-trait = "0.1"
 chrono = { version = "0.4.35", default-features = false, features = ["clock"] }

driver/src/lib.rs

@@ -17,6 +17,7 @@ pub mod conn;
 #[cfg(feature = "flight-sql")]
 mod flight_sql;
 mod params;
+mod place_holder;
 pub mod rest_api;
 
 pub use client::Client;

driver/src/params.rs

@@ -79,28 +79,12 @@ impl Params {
 
     pub fn replace(&self, sql: &str) -> String {
         if !self.is_empty() {
-            if let Ok((stmt, _)) = databend_common_ast::parser::parse_sql(sql, Dialect::PostgreSQL)
+            let tokens = databend_common_ast::parser::tokenize_sql(sql).unwrap();
+            if let Ok((stmt, _)) =
+                databend_common_ast::parser::parse_sql(&tokens, Dialect::PostgreSQL)
             {
-                let mut sql = sql.to_string();
-                let mut positions = Vec::new();
-
-                for token in tokens {
-                    match token.kind {
-                        databend_common_ast::parser::token::TokenKind::Placeholder => {
-                            positions.push(token.span);
-                        }
-                        _ => {}
-                    }
-                }
-                let size = positions.len();
-                for (index, r) in positions.iter().rev().enumerate() {
-                    if let Some(param) = self.get_by_index(size - index) {
-                        let start = r.start as usize;
-                        let end = r.end as usize;
-                        sql.replace_range(start..end, param);
-                    }
-                }
-                return sql;
+                let mut v = super::place_holder::PlaceholderVisitor::new();
+                return v.replace_sql(self, &stmt, sql);
             }
         }
         return sql.to_string();
@@ -248,15 +232,8 @@ impl_from_tuple_for_params! { T1, T2, T3, T4, T5, T6, T7, T8, T9, T10 }
 impl From<Option<serde_json::Value>> for Params {
     fn from(value: Option<serde_json::Value>) -> Self {
         match value {
-            Some(serde_json::Value::Array(obj)) => {
-                let mut array = Vec::new();
-                for v in obj {
-                    array.push(v.as_sql_string());
-                }
-                Params::QuestionParams(array)
-            }
+            Some(v) => v.into(),
             None => Params::default(),
-            _ => unimplemented!("json value must be an array as params"),
         }
     }
 }
@@ -377,5 +354,19 @@ mod tests {
             "SELECT * FROM table WHERE a = ? AND '?' = cj AND b = ? AND c = ? AND d = ? AND e = ? AND f = ?";
         let replaced_sql = params.replace(sql);
         assert_eq!(replaced_sql, "SELECT * FROM table WHERE a = 1 AND '?' = cj AND b = '44' AND c = 2 AND d = 3 AND e = '55' AND f = '66'");
+
+        let params = params! {a => 1, b => "44", c => 2, d => 3, e => "55", f => "66"};
+
+        {
+            let sql = "SELECT * FROM table WHERE a = :a AND '?' = cj AND b = :b AND c = :c AND d = :d AND e = :e AND f = :f";
+            let replaced_sql = params.replace(sql);
+            assert_eq!(replaced_sql, "SELECT * FROM table WHERE a = 1 AND '?' = cj AND b = '44' AND c = 2 AND d = 3 AND e = '55' AND f = '66'");
+        }
+
+        {
+            let sql = "SELECT b = :b, a = :a FROM table WHERE a = :a AND '?' = cj AND b = :b AND c = :c AND d = :d AND e = :e AND f = :f";
+            let replaced_sql = params.replace(sql);
+            assert_eq!(replaced_sql, "SELECT b = '44', a = 1 FROM table WHERE a = 1 AND '?' = cj AND b = '44' AND c = 2 AND d = 3 AND e = '55' AND f = '66'");
+        }
     }
 }

driver/src/place_holder.rs

@@ -0,0 +1,96 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::vec;
+
+use databend_common_ast::ast::Expr;
+use databend_common_ast::ast::Identifier;
+use databend_common_ast::ast::IdentifierType;
+use databend_common_ast::ast::Statement;
+use databend_common_ast::Range;
+use derive_visitor::Drive;
+use derive_visitor::Visitor;
+
+use crate::Params;
+
+#[derive(Visitor)]
+#[visitor(Expr(enter), Identifier(enter))]
+pub(crate) struct PlaceholderVisitor {
+    place_holders: Vec<Range>,
+    names: Vec<(String, Range)>,
+}
+
+impl PlaceholderVisitor {
+    pub fn new() -> Self {
+        PlaceholderVisitor {
+            place_holders: vec![],
+            names: Vec::new(),
+        }
+    }
+
+    fn enter_expr(&mut self, expr: &Expr) {
+        match expr {
+            Expr::Hole {
+                name,
+                span: Some(range),
+            } => {
+                self.names.push((name.clone(), range.clone()));
+            }
+            Expr::Placeholder { span: Some(range) } => {
+                self.place_holders.push(range.clone());
+            }
+            _ => {}
+        }
+    }
+
+    fn enter_identifier(&mut self, ident: &Identifier) {
+        match (ident.ident_type, ident.span) {
+            (IdentifierType::Hole, Some(range)) => {
+                self.names.push((ident.name.clone(), range));
+            }
+            _ => {}
+        }
+    }
+
+    pub fn replace_sql(&mut self, params: &Params, stmt: &Statement, sql: &str) -> String {
+        stmt.drive(self);
+        self.place_holders.sort_by(|l, r| l.start.cmp(&r.start));
+
+        let mut results = vec![];
+
+        for (index, range) in self.place_holders.iter().enumerate() {
+            if let Some(v) = params.get_by_index(index + 1) {
+                results.push((v.to_string(), range.clone()));
+            }
+        }
+
+        for (name, range) in self.names.iter() {
+            if let Some(v) = params.get_by_name(name) {
+                results.push((v.to_string(), range.clone()));
+            }
+        }
+
+        if !results.is_empty() {
+            let mut sql = sql.to_string();
+            results.sort_by(|a, b| a.1.start.cmp(&b.1.start));
+            for (value, r) in results.iter().rev() {
+                let start = r.start as usize;
+                let end = r.end as usize;
+                sql.replace_range(start..end, value);
+            }
+            return sql;
+        }
+        sql.to_string()
+    }
+}