feat(query): Procedure Part2 support arguments (#16453)

* feat(query): Procedure Part2 support arguments

* fix conversation

* refactor var_to_ast with parse_expr

src/meta/app/Cargo.toml

@@ -19,6 +19,7 @@ chrono = { workspace = true }
 chrono-tz = { workspace = true }
 cron = "0.12.0"
 databend-common-ast = { workspace = true }
+databend-common-base = { workspace = true }
 databend-common-exception = { workspace = true }
 databend-common-expression = { workspace = true }
 databend-common-io = { workspace = true }

src/meta/app/src/principal/procedure.rs

@@ -19,6 +19,7 @@ use std::ops::Deref;
 
 use chrono::DateTime;
 use chrono::Utc;
+use databend_common_base::display::display_slice::DisplaySliceExt;
 use databend_common_expression::types::DataType;
 
 use crate::principal::procedure_id_ident::ProcedureIdIdent;
@@ -45,6 +46,7 @@ pub struct ProcedureIdent {
 #[derive(Clone, Debug, PartialEq)]
 pub struct ProcedureMeta {
     pub return_types: Vec<DataType>,
+    pub arg_names: Vec<String>,
     pub created_on: DateTime<Utc>,
     pub updated_on: DateTime<Utc>,
     pub script: String,
@@ -56,6 +58,7 @@ impl Default for ProcedureMeta {
     fn default() -> Self {
         ProcedureMeta {
             return_types: vec![],
+            arg_names: vec![],
             created_on: Utc::now(),
             updated_on: Utc::now(),
             script: "".to_string(),
@@ -69,8 +72,13 @@ impl Display for ProcedureMeta {
     fn fmt(&self, f: &mut Formatter) -> fmt::Result {
         write!(
             f,
-            "Lanuage: {:?}, return_type: {:?}, CreatedOn: {:?}, Script: {:?}, Comment: {:?}",
-            self.procedure_language, self.return_types, self.created_on, self.script, self.comment
+            "Lanuage: {:?}, args {} return_type: {}, CreatedOn: {:?}, Script: {}, Comment: {:?}",
+            self.procedure_language,
+            self.arg_names.display_n::<1000>(),
+            self.return_types.display_n::<1000>(),
+            self.created_on,
+            self.script,
+            self.comment
         )
     }
 }

src/meta/proto-conv/src/procedure_from_to_protobuf_impl.rs

@@ -70,6 +70,7 @@ impl FromToProto for mt::principal::ProcedureMeta {
 
         let v = Self {
             return_types,
+            arg_names: p.arg_names.clone(),
             created_on: DateTime::<Utc>::from_pb(p.created_on)?,
             updated_on: DateTime::<Utc>::from_pb(p.updated_on)?,
             script: p.script,
@@ -94,6 +95,7 @@ impl FromToProto for mt::principal::ProcedureMeta {
             ver: VER,
             min_reader_ver: MIN_READER_VER,
             return_types,
+            arg_names: self.arg_names.clone(),
             created_on: self.created_on.to_pb()?,
             updated_on: self.updated_on.to_pb()?,
             script: self.script.to_string(),

src/meta/proto-conv/src/util.rs

@@ -137,7 +137,8 @@ const META_CHANGE_LOG: &[(u64, &str)] = &[
     (105, "2024-08-05: Add: add Dictionary meta"),
     (106, "2024-08-08: Add: add QueryTokenInfo"),
     (107, "2024-08-09: Add: datatype.proto/DataType Geography type"),
-    (108, "2024-08-29: Add: procedure.proto: ProcedureMeta and ProcedureIdentity")
+    (108, "2024-08-29: Add: procedure.proto: ProcedureMeta and ProcedureIdentity"),
+    (109, "2024-08-29: Refactor: ProcedureMeta add arg_names"),
     // Dear developer:
     //      If you're gonna add a new metadata version, you'll have to add a test for it.
     //      You could just copy an existing test file(e.g., `../tests/it/v024_table_meta.rs`)

src/meta/proto-conv/tests/it/main.rs

@@ -106,3 +106,4 @@ mod v105_dictionary_meta;
 mod v106_query_token;
 mod v107_geography_datatype;
 mod v108_procedure;
+mod v109_procedure_with_args;

src/meta/proto-conv/tests/it/v108_procedure.rs

@@ -22,15 +22,16 @@ use crate::common;
 
 #[test]
 fn v108_procedure_meta() -> anyhow::Result<()> {
-    let procedure_meta_v108: Vec<u8> = vec![
-        34, 9, 146, 2, 0, 160, 6, 108, 168, 6, 24, 82, 23, 50, 48, 49, 52, 45, 49, 49, 45, 50, 56,
+    let procedure_meta_v108 = vec![
+        34, 9, 146, 2, 0, 160, 6, 109, 168, 6, 24, 82, 23, 50, 48, 49, 52, 45, 49, 49, 45, 50, 56,
         32, 49, 50, 58, 48, 48, 58, 48, 57, 32, 85, 84, 67, 90, 23, 50, 48, 49, 52, 45, 49, 49, 45,
         50, 57, 32, 49, 50, 58, 48, 48, 58, 48, 57, 32, 85, 84, 67, 98, 7, 102, 111, 111, 32, 98,
-        97, 114, 114, 3, 83, 81, 76, 160, 6, 108, 168, 6, 24,
+        97, 114, 114, 3, 83, 81, 76, 160, 6, 109, 168, 6, 24,
     ];
 
     let want = || mt::ProcedureMeta {
         return_types: vec![DataType::String],
+        arg_names: vec![],
         created_on: Utc.with_ymd_and_hms(2014, 11, 28, 12, 0, 9).unwrap(),
         updated_on: Utc.with_ymd_and_hms(2014, 11, 29, 12, 0, 9).unwrap(),
         script: "".to_string(),
@@ -39,7 +40,7 @@ fn v108_procedure_meta() -> anyhow::Result<()> {
     };
 
     common::test_pb_from_to(func_name!(), want())?;
-    common::test_load_old(func_name!(), procedure_meta_v108.as_slice(), 108, want())
+    common::test_load_old(func_name!(), procedure_meta_v108.as_slice(), 109, want())
 }
 
 #[test]

src/meta/proto-conv/tests/it/v109_procedure_with_args.rs

@@ -0,0 +1,44 @@
+// Copyright 2023 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 chrono::TimeZone;
+use chrono::Utc;
+use databend_common_expression::types::DataType;
+use databend_common_meta_app::principal as mt;
+use fastrace::func_name;
+
+use crate::common;
+
+#[test]
+fn v109_procedure_meta() -> anyhow::Result<()> {
+    let procedure_meta_v109 = vec![
+        34, 9, 146, 2, 0, 160, 6, 109, 168, 6, 24, 42, 7, 109, 101, 115, 115, 97, 103, 101, 82, 23,
+        50, 48, 49, 52, 45, 49, 49, 45, 50, 56, 32, 49, 50, 58, 48, 48, 58, 48, 57, 32, 85, 84, 67,
+        90, 23, 50, 48, 49, 52, 45, 49, 49, 45, 50, 57, 32, 49, 50, 58, 48, 48, 58, 48, 57, 32, 85,
+        84, 67, 98, 7, 102, 111, 111, 32, 98, 97, 114, 114, 3, 83, 81, 76, 160, 6, 109, 168, 6, 24,
+    ];
+
+    let want = || mt::ProcedureMeta {
+        return_types: vec![DataType::String],
+        arg_names: vec!["message".to_string()],
+        created_on: Utc.with_ymd_and_hms(2014, 11, 28, 12, 0, 9).unwrap(),
+        updated_on: Utc.with_ymd_and_hms(2014, 11, 29, 12, 0, 9).unwrap(),
+        script: "".to_string(),
+        comment: "foo bar".to_string(),
+        procedure_language: "SQL".to_string(),
+    };
+
+    common::test_pb_from_to(func_name!(), want())?;
+    common::test_load_old(func_name!(), procedure_meta_v109.as_slice(), 109, want())
+}

src/meta/protos/proto/procedure.proto

@@ -37,6 +37,7 @@ message ProcedureMeta {
 
   // Procedure return type
   repeated DataType return_types = 4;
+  repeated string arg_names = 5;
 
   // The time database created.
   string created_on = 10;

src/query/ast/src/ast/statements/procedure.rs

@@ -19,8 +19,8 @@ use derive_visitor::Drive;
 use derive_visitor::DriveMut;
 
 use crate::ast::write_comma_separated_list;
-use crate::ast::write_comma_separated_string_list;
 use crate::ast::CreateOption;
+use crate::ast::Expr;
 use crate::ast::TypeName;
 
 #[derive(Debug, Clone, PartialEq, Eq, Drive, DriveMut)]
@@ -169,16 +169,17 @@ impl Display for DescProcedureStmt {
     }
 }
 
-#[derive(Debug, Clone, PartialEq, Eq, Drive, DriveMut)]
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
 pub struct CallProcedureStmt {
     pub name: String,
-    pub args: Vec<String>,
+    pub args: Vec<Expr>,
 }
 
 impl Display for CallProcedureStmt {
     fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
-        write!(f, "CALL PROCEDURE {}(", self.name)?;
-        write_comma_separated_string_list(f, self.args.clone())?;
+        let CallProcedureStmt { name, args } = self;
+        write!(f, "CALL PROCEDURE {}(", name)?;
+        write_comma_separated_list(f, args)?;
         write!(f, ")")?;
         Ok(())
     }

src/query/ast/src/parser/statement.rs

@@ -2125,12 +2125,12 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
 
     let call_procedure = map(
         rule! {
-            CALL ~ PROCEDURE ~ #ident ~ "(" ~ ")"
+            CALL ~ PROCEDURE ~ #ident ~ "(" ~ #comma_separated_list0(subexpr(0))? ~ ")"
         },
-        |(_, _, name, _, _)| {
+        |(_, _, name, _, opt_args, _)| {
             Statement::CallProcedure(CallProcedureStmt {
                 name: name.to_string(),
-                args: vec![],
+                args: opt_args.unwrap_or_default(),
             })
         },
     );

src/query/ast/tests/it/parser.rs

@@ -831,6 +831,7 @@ fn test_statement() {
         r#"drop PROCEDURE p1()"#,
         r#"drop PROCEDURE p1(int, string)"#,
         r#"call PROCEDURE p1()"#,
+        r#"call PROCEDURE p1(1, 'x', '2022-02-02'::Date)"#,
         r#"show PROCEDURES like 'p1%'"#,
         r#"create PROCEDURE p1() returns string not null language sql comment = 'test' as $$
             BEGIN

src/query/ast/tests/it/testdata/stmt.txt

@@ -23066,6 +23066,51 @@ CallProcedure(
 )
 
 
+---------- Input ----------
+call PROCEDURE p1(1, 'x', '2022-02-02'::Date)
+---------- Output ---------
+CALL PROCEDURE p1(1, 'x', '2022-02-02'::DATE)
+---------- AST ------------
+CallProcedure(
+    CallProcedureStmt {
+        name: "p1",
+        args: [
+            Literal {
+                span: Some(
+                    18..19,
+                ),
+                value: UInt64(
+                    1,
+                ),
+            },
+            Literal {
+                span: Some(
+                    21..24,
+                ),
+                value: String(
+                    "x",
+                ),
+            },
+            Cast {
+                span: Some(
+                    38..44,
+                ),
+                expr: Literal {
+                    span: Some(
+                        26..38,
+                    ),
+                    value: String(
+                        "2022-02-02",
+                    ),
+                },
+                target_type: Date,
+                pg_style: true,
+            },
+        ],
+    },
+)
+
+
 ---------- Input ----------
 show PROCEDURES like 'p1%'
 ---------- Output ---------

src/query/script/src/executor.rs

@@ -35,7 +35,7 @@ pub trait Client {
     fn var_to_ast(&self, scalar: &Self::Var) -> Result<Expr>;
     fn read_from_set(&self, block: &Self::Set, row: usize, col: &ColumnAccess)
     -> Result<Self::Var>;
-    fn set_len(&self, block: &Self::Set) -> usize;
+    fn num_rows(&self, block: &Self::Set) -> usize;
     fn is_true(&self, scalar: &Self::Var) -> Result<bool>;
 }
 
@@ -132,7 +132,7 @@ impl<C: Client> Executor<C> {
                 let cursor = Cursor {
                     set: set.clone(),
                     row: 0,
-                    len: self.client.set_len(block),
+                    len: self.client.num_rows(block),
                 };
                 self.iters.insert(to_iter.clone(), cursor);
             }

src/query/script/tests/it/main.rs

@@ -689,7 +689,7 @@ impl Client for MockClient {
         Ok(var)
     }
 
-    fn set_len(&self, set: &Self::Set) -> usize {
+    fn num_rows(&self, set: &Self::Set) -> usize {
         set.data.len()
     }
 

src/query/service/src/interpreters/access/privilege_access.rs

@@ -1260,6 +1260,7 @@ impl AccessChecker for PrivilegeAccess {
             Plan::DescDatamaskPolicy(_) => {}
             Plan::Begin => {}
             Plan::ExecuteImmediate(_)
+            | Plan::CallProcedure(_)
             | Plan::CreateProcedure(_)
             | Plan::DropProcedure(_)
             /*| Plan::ShowCreateProcedure(_)