fix(query): improve private task compatibility (#19918)

* fix(query): assign private task ids

* fix(query): support private task scripts

Author: KKould

src/meta/app/src/id_generator.rs

@@ -31,6 +31,7 @@ pub(crate) const ID_GEN_ROW_POLICY: &str = "row_access";
 pub(crate) const ID_GEN_BACKGROUND_JOB: &str = "background_job";
 
 pub(crate) const ID_GEN_PROCEDURE: &str = "procedure_id";
+pub(crate) const ID_GEN_TASK: &str = "task_id";
 
 /// Key for resource id generator
 ///
@@ -109,6 +110,11 @@ impl IdGenerator {
     pub fn procedure_id() -> Self {
         Self::new(ID_GEN_PROCEDURE)
     }
+
+    /// Create a key for generating task id with kvapi::KVApi
+    pub fn task_id() -> Self {
+        Self::new(ID_GEN_TASK)
+    }
 }
 
 impl kvapi::Key for IdGenerator {
@@ -143,6 +149,7 @@ mod t {
         assert_round_trip(IdGenerator::data_mask_id(), "__fd_id_gen/data_mask");
         assert_round_trip(IdGenerator::table_lock_id(), "__fd_id_gen/table_lock_id");
         assert_round_trip(IdGenerator::procedure_id(), "__fd_id_gen/procedure_id");
+        assert_round_trip(IdGenerator::task_id(), "__fd_id_gen/task_id");
     }
 
     #[test]

src/meta/app/src/principal/mod.rs

@@ -106,6 +106,7 @@ pub use task::Status;
 pub use task::Task;
 pub use task::TaskMessage;
 pub use task::TaskRun;
+pub use task::TaskSql;
 pub use task::WarehouseOptions;
 pub use task_ident::TaskIdent;
 pub use task_ident::TaskIdentRaw;

src/meta/app/src/principal/task.rs

@@ -55,11 +55,34 @@ pub struct WarehouseOptions {
     pub using_warehouse_size: Option<String>,
 }
 
+#[derive(Debug, Clone, PartialEq)]
+pub enum TaskSql {
+    Sql(String),
+    Script(Vec<String>),
+}
+
+impl TaskSql {
+    pub fn query_text(&self) -> String {
+        match self {
+            TaskSql::Sql(sql) => sql.clone(),
+            TaskSql::Script(sqls) => {
+                let mut text = String::from("BEGIN\n");
+                for sql in sqls {
+                    text.push_str(sql);
+                    text.push_str(";\n");
+                }
+                text.push_str("END;");
+                text
+            }
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq)]
 pub struct Task {
     pub task_id: u64,
     pub task_name: String,
-    pub query_text: String,
+    pub task_sql: TaskSql,
     pub when_condition: Option<String>,
     pub after: Vec<String>,
     pub comment: Option<String>,

src/meta/proto-conv/src/impls/task.rs

@@ -72,7 +72,11 @@ impl FromToProto for mt::Task {
         Ok(Self {
             task_id: p.task_id,
             task_name: p.task_name,
-            query_text: p.query_text,
+            task_sql: if p.script_sql.is_empty() {
+                mt::TaskSql::Sql(p.query_text)
+            } else {
+                mt::TaskSql::Script(p.script_sql)
+            },
             when_condition: p.when_condition.clone(),
             after: p.after,
             comment: p.comment,
@@ -97,7 +101,7 @@ impl FromToProto for mt::Task {
             min_reader_ver: MIN_READER_VER,
             task_id: self.task_id,
             task_name: self.task_name.clone(),
-            query_text: self.query_text.clone(),
+            query_text: self.task_sql.query_text(),
             comment: self.comment.clone(),
             owner: self.owner.clone(),
             schedule_options: self.schedule_options.as_ref().map(|s| pb::ScheduleOptions {
@@ -125,6 +129,10 @@ impl FromToProto for mt::Task {
             session_parameters: self.session_params.clone(),
             error_integration: self.error_integration.clone(),
             owner_user: self.owner_user.clone(),
+            script_sql: match &self.task_sql {
+                mt::TaskSql::Sql(_) => vec![],
+                mt::TaskSql::Script(sqls) => sqls.clone(),
+            },
         })
     }
 }

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

@@ -204,7 +204,8 @@ const META_CHANGE_LOG: &[(u64, &str)] = &[
     (172, "2026-04-09: Update: user.proto/CsvFileFormatParams and TextFileFormatParams add trim_space"),
     (173, "2026-04-16: Update: file_format.proto/CsvFileFormatParams add quote_style"),
     (174, "2026-04-28: Add: AuthInfo::KeyPair for key-pair authentication"),
-    (175, "2026-05-08: Add: field_stats_truncate_len per-column string stats truncation in TableMeta")
+    (175, "2026-05-08: Add: field_stats_truncate_len per-column string stats truncation in TableMeta"),
+    (176, "2026-05-25: Add: task.proto/Task.script_sql")
     // 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

@@ -167,3 +167,4 @@ mod v172_file_format_trim_space;
 mod v173_csv_quote_style;
 mod v174_user_key_pair;
 mod v175_field_stats_truncate_len;
+mod v176_task_script_sql;

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

@@ -41,7 +41,7 @@ fn test_decode_v136_add_task() -> anyhow::Result<()> {
     let want = || mt::Task {
         task_id: 11,
         task_name: "task_c".to_string(),
-        query_text: "SELECT * FROM t1".to_string(),
+        task_sql: mt::TaskSql::Sql("SELECT * FROM t1".to_string()),
         when_condition: Some("c1 > 1".to_string()),
         after: vec!["task_a".to_string(), "task_b".to_string()],
         comment: Some("comment".to_string()),
@@ -78,7 +78,7 @@ fn test_decode_v136_task_message() -> anyhow::Result<()> {
     let want_task = || mt::Task {
         task_id: 11,
         task_name: "task_c".to_string(),
-        query_text: "SELECT * FROM t1".to_string(),
+        task_sql: mt::TaskSql::Sql("SELECT * FROM t1".to_string()),
         when_condition: Some("c1 > 1".to_string()),
         after: vec!["task_a".to_string(), "task_b".to_string()],
         comment: Some("comment".to_string()),

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

@@ -27,7 +27,7 @@ fn test_decode_v140_task_message() -> anyhow::Result<()> {
     let want_task = || mt::Task {
         task_id: 11,
         task_name: "task_c".to_string(),
-        query_text: "SELECT * FROM t1".to_string(),
+        task_sql: mt::TaskSql::Sql("SELECT * FROM t1".to_string()),
         when_condition: Some("c1 > 1".to_string()),
         after: vec!["task_a".to_string(), "task_b".to_string()],
         comment: Some("comment".to_string()),

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

@@ -0,0 +1,64 @@
+// Copyright 2026 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::DateTime;
+use databend_common_meta_app::principal as mt;
+use fastrace::func_name;
+use maplit::btreemap;
+
+use crate::common;
+
+#[test]
+fn test_decode_v176_task_script_sql() -> anyhow::Result<()> {
+    let task_v176 = vec![
+        8, 11, 18, 6, 116, 97, 115, 107, 95, 99, 34, 60, 66, 69, 71, 73, 78, 10, 73, 78, 83, 69,
+        82, 84, 32, 73, 78, 84, 79, 32, 116, 32, 86, 65, 76, 85, 69, 83, 40, 49, 41, 59, 10, 73,
+        78, 83, 69, 82, 84, 32, 73, 78, 84, 79, 32, 116, 32, 86, 65, 76, 85, 69, 83, 40, 50, 41,
+        59, 10, 69, 78, 68, 59, 50, 6, 112, 117, 98, 108, 105, 99, 114, 23, 49, 57, 55, 48, 45, 48,
+        49, 45, 48, 49, 32, 48, 48, 58, 48, 48, 58, 49, 49, 32, 85, 84, 67, 122, 23, 49, 57, 55,
+        48, 45, 48, 49, 45, 48, 49, 32, 48, 48, 58, 48, 48, 58, 49, 50, 32, 85, 84, 67, 170, 1, 2,
+        109, 101, 178, 1, 23, 73, 78, 83, 69, 82, 84, 32, 73, 78, 84, 79, 32, 116, 32, 86, 65, 76,
+        85, 69, 83, 40, 49, 41, 178, 1, 23, 73, 78, 83, 69, 82, 84, 32, 73, 78, 84, 79, 32, 116,
+        32, 86, 65, 76, 85, 69, 83, 40, 50, 41, 160, 6, 176, 1, 168, 6, 24,
+    ];
+
+    let want = mt::Task {
+        task_id: 11,
+        task_name: "task_c".to_string(),
+        task_sql: mt::TaskSql::Script(vec![
+            "INSERT INTO t VALUES(1)".to_string(),
+            "INSERT INTO t VALUES(2)".to_string(),
+        ]),
+        when_condition: None,
+        after: vec![],
+        comment: None,
+        owner: mt::BUILTIN_ROLE_PUBLIC.to_string(),
+        owner_user: "me".to_string(),
+        schedule_options: None,
+        warehouse_options: None,
+        next_scheduled_at: None,
+        suspend_task_after_num_failures: None,
+        error_integration: None,
+        status: mt::Status::Suspended,
+        created_at: DateTime::from_timestamp(11, 0).unwrap(),
+        updated_at: DateTime::from_timestamp(12, 0).unwrap(),
+        last_suspended_at: None,
+        session_params: btreemap! {},
+    };
+
+    common::test_pb_from_to(func_name!(), want.clone())?;
+    common::test_load_old(func_name!(), task_v176.as_slice(), 176, want)?;
+
+    Ok(())
+}

src/meta/protos/proto/task.proto

@@ -61,6 +61,8 @@ message Task {
   }
   uint64 task_id = 1;
   string task_name = 2;
+  // Kept for compatibility and display. For single-statement tasks this is the
+  // SQL text; for script tasks it is derived from script_sql.
   string query_text = 4;
   optional string comment = 5;
   string owner = 6;
@@ -77,4 +79,7 @@ message Task {
   map<string, string> session_parameters = 19;
   optional string error_integration = 20;
   string owner_user = 21;
+  // Non-empty only for script tasks created with `AS BEGIN ... END`.
+  // Each element is one statement in execution order.
+  repeated string script_sql = 22;
 }