test: add test-private-task-warehouse.sh

KKould · KKould · commit 4c32e2307a74 · 2025-07-29T15:41:28.000+08:00
diff --git a/src/meta/proto-conv/tests/it/v138_task_message.rs b/src/meta/proto-conv/tests/it/v138_task_message.rs
@@ -0,0 +1,156 @@
+// 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::DateTime;
+use databend_common_meta_app::principal as mt;
+use databend_common_meta_app::principal::ScheduleOptions;
+use databend_common_meta_app::principal::ScheduleType;
+use databend_common_meta_app::principal::WarehouseOptions;
+use fastrace::func_name;
+use maplit::btreemap;
+
+use crate::common;
+
+#[test]
+fn test_decode_v138_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(),
+        when_condition: Some("c1 > 1".to_string()),
+        after: vec!["task_a".to_string(), "task_b".to_string()],
+        comment: Some("comment".to_string()),
+        owner: "public".to_string(),
+        owner_user: "me".to_string(),
+        schedule_options: Some(ScheduleOptions {
+            interval: Some(11),
+            cron: Some("30 12 * * *".to_string()),
+            time_zone: Some("UTC".to_string()),
+            schedule_type: ScheduleType::IntervalType,
+            milliseconds_interval: Some(11),
+        }),
+        warehouse_options: Some(WarehouseOptions {
+            warehouse: Some("warehouse_a".to_string()),
+            using_warehouse_size: Some("10".to_string()),
+        }),
+        next_scheduled_at: Some(DateTime::from_timestamp(10, 0).unwrap()),
+        suspend_task_after_num_failures: Some(10),
+        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: Some(DateTime::from_timestamp(13, 0).unwrap()),
+        session_params: btreemap! { s("a") => s("b") },
+    };
+
+    {
+        let task_message_execute_v138 = vec![
+            10, 239, 1, 8, 11, 18, 6, 116, 97, 115, 107, 95, 99, 34, 16, 83, 69, 76, 69, 67, 84,
+            32, 42, 32, 70, 82, 79, 77, 32, 116, 49, 42, 7, 99, 111, 109, 109, 101, 110, 116, 50,
+            6, 112, 117, 98, 108, 105, 99, 58, 22, 8, 11, 18, 11, 51, 48, 32, 49, 50, 32, 42, 32,
+            42, 32, 42, 26, 3, 85, 84, 67, 40, 11, 66, 17, 10, 11, 119, 97, 114, 101, 104, 111,
+            117, 115, 101, 95, 97, 18, 2, 49, 48, 74, 23, 49, 57, 55, 48, 45, 48, 49, 45, 48, 49,
+            32, 48, 48, 58, 48, 48, 58, 49, 48, 32, 85, 84, 67, 80, 10, 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, 130, 1, 23, 49, 57, 55, 48, 45, 48, 49, 45, 48, 49, 32, 48, 48, 58, 48, 48, 58, 49,
+            51, 32, 85, 84, 67, 138, 1, 6, 116, 97, 115, 107, 95, 97, 138, 1, 6, 116, 97, 115, 107,
+            95, 98, 146, 1, 6, 99, 49, 32, 62, 32, 49, 154, 1, 6, 10, 1, 97, 18, 1, 98, 170, 1, 2,
+            109, 101, 160, 6, 137, 1, 168, 6, 24, 160, 6, 137, 1, 168, 6, 24,
+        ];
+        let want_execute = || mt::TaskMessage::ExecuteTask(want_task());
+
+        common::test_pb_from_to(func_name!(), want_execute())?;
+        common::test_load_old(
+            func_name!(),
+            task_message_execute_v138.as_slice(),
+            138,
+            want_execute(),
+        )?;
+    }
+    {
+        let task_message_schedule_v138 = vec![
+            18, 239, 1, 8, 11, 18, 6, 116, 97, 115, 107, 95, 99, 34, 16, 83, 69, 76, 69, 67, 84,
+            32, 42, 32, 70, 82, 79, 77, 32, 116, 49, 42, 7, 99, 111, 109, 109, 101, 110, 116, 50,
+            6, 112, 117, 98, 108, 105, 99, 58, 22, 8, 11, 18, 11, 51, 48, 32, 49, 50, 32, 42, 32,
+            42, 32, 42, 26, 3, 85, 84, 67, 40, 11, 66, 17, 10, 11, 119, 97, 114, 101, 104, 111,
+            117, 115, 101, 95, 97, 18, 2, 49, 48, 74, 23, 49, 57, 55, 48, 45, 48, 49, 45, 48, 49,
+            32, 48, 48, 58, 48, 48, 58, 49, 48, 32, 85, 84, 67, 80, 10, 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, 130, 1, 23, 49, 57, 55, 48, 45, 48, 49, 45, 48, 49, 32, 48, 48, 58, 48, 48, 58, 49,
+            51, 32, 85, 84, 67, 138, 1, 6, 116, 97, 115, 107, 95, 97, 138, 1, 6, 116, 97, 115, 107,
+            95, 98, 146, 1, 6, 99, 49, 32, 62, 32, 49, 154, 1, 6, 10, 1, 97, 18, 1, 98, 170, 1, 2,
+            109, 101, 160, 6, 137, 1, 168, 6, 24, 160, 6, 137, 1, 168, 6, 24,
+        ];
+        let want_schedule = || mt::TaskMessage::ScheduleTask(want_task());
+
+        common::test_pb_from_to(func_name!(), want_schedule())?;
+        common::test_load_old(
+            func_name!(),
+            task_message_schedule_v138.as_slice(),
+            138,
+            want_schedule(),
+        )?;
+    }
+    {
+        let task_message_after_v138 = vec![
+            34, 239, 1, 8, 11, 18, 6, 116, 97, 115, 107, 95, 99, 34, 16, 83, 69, 76, 69, 67, 84,
+            32, 42, 32, 70, 82, 79, 77, 32, 116, 49, 42, 7, 99, 111, 109, 109, 101, 110, 116, 50,
+            6, 112, 117, 98, 108, 105, 99, 58, 22, 8, 11, 18, 11, 51, 48, 32, 49, 50, 32, 42, 32,
+            42, 32, 42, 26, 3, 85, 84, 67, 40, 11, 66, 17, 10, 11, 119, 97, 114, 101, 104, 111,
+            117, 115, 101, 95, 97, 18, 2, 49, 48, 74, 23, 49, 57, 55, 48, 45, 48, 49, 45, 48, 49,
+            32, 48, 48, 58, 48, 48, 58, 49, 48, 32, 85, 84, 67, 80, 10, 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, 130, 1, 23, 49, 57, 55, 48, 45, 48, 49, 45, 48, 49, 32, 48, 48, 58, 48, 48, 58, 49,
+            51, 32, 85, 84, 67, 138, 1, 6, 116, 97, 115, 107, 95, 97, 138, 1, 6, 116, 97, 115, 107,
+            95, 98, 146, 1, 6, 99, 49, 32, 62, 32, 49, 154, 1, 6, 10, 1, 97, 18, 1, 98, 170, 1, 2,
+            109, 101, 160, 6, 137, 1, 168, 6, 24, 160, 6, 137, 1, 168, 6, 24,
+        ];
+        let want_after = || mt::TaskMessage::AfterTask(want_task());
+
+        common::test_pb_from_to(func_name!(), want_after())?;
+        common::test_load_old(
+            func_name!(),
+            task_message_after_v138.as_slice(),
+            138,
+            want_after(),
+        )?;
+    }
+    {
+        let task_message_delete_v138 = vec![
+            26, 27, 10, 6, 116, 97, 115, 107, 95, 99, 18, 17, 10, 11, 119, 97, 114, 101, 104, 111,
+            117, 115, 101, 95, 97, 18, 2, 49, 48, 160, 6, 137, 1, 168, 6, 24,
+        ];
+        let want_delete = || {
+            let task = want_task();
+            mt::TaskMessage::DeleteTask(task.task_name, task.warehouse_options)
+        };
+
+        common::test_pb_from_to(func_name!(), want_delete())?;
+        common::test_load_old(
+            func_name!(),
+            task_message_delete_v138.as_slice(),
+            138,
+            want_delete(),
+        )?;
+    }
+
+    Ok(())
+}
+
+fn s(ss: impl ToString) -> String {
+    ss.to_string()
+}
diff --git a/src/query/service/src/task/service.rs b/src/query/service/src/task/service.rs
@@ -30,6 +30,7 @@ use databend_common_ast::ast::AlterTaskOptions;
 use databend_common_base::base::GlobalInstance;
 use databend_common_base::runtime::Runtime;
 use databend_common_base::runtime::TrySpawn;
+use databend_common_catalog::table_context::TableContext;
 use databend_common_config::GlobalConfig;
 use databend_common_config::InnerConfig;
 use databend_common_exception::ErrorCode;
@@ -208,8 +209,13 @@ impl TaskService {
                 ..
             }) = task_message.warehouse_options()
             {
-                // [WarehouseInfo::SelfManaged] uses ClusterId as warehouse
-                if warehouse != &self.cluster_id {
+                if warehouse
+                    != &self
+                        .create_context(None)
+                        .await?
+                        .get_cluster()
+                        .get_warehouse_id()?
+                {
                     continue;
                 }
             }
@@ -282,7 +288,7 @@ impl TaskService {
                                                         let Some(_guard) = fn_lock(&task_service, &task_key, duration.as_millis() as u64).await? else {
                                                             continue;
                                                         };
-                                                        if !Self::check_when(&task, &owner, &task_service).await.unwrap() {
+                                                        if !Self::check_when(&task, &owner, &task_service).await? {
                                                             continue;
                                                         }
                                                         fn_new_task_run(&task_service, &task).await?;
@@ -540,8 +546,7 @@ impl TaskService {
         };
         let result = task_service
             .execute_sql(Some(user.clone()), &format!("SELECT {when_condition}"))
-            .await
-            .unwrap();
+            .await?;
         Ok(result
             .first()
             .and_then(|block| block.get_by_offset(0).index(0))
diff --git a/src/query/sql/src/planner/binder/ddl/task.rs b/src/query/sql/src/planner/binder/ddl/task.rs
@@ -123,9 +123,7 @@ impl Binder {
             sql,
             session_parameters,
         } = stmt;
-        if (schedule_opts.is_none() && after.is_empty())
-            || (schedule_opts.is_some() && !after.is_empty())
-        {
+        if schedule_opts.is_some() && !after.is_empty() {
             return Err(ErrorCode::SyntaxException(
                 "task must be defined with either given time schedule as a root task or run after other task as a DAG".to_string(),
             ));
diff --git a/tests/task/test-private-task-warehouse.sh b/tests/task/test-private-task-warehouse.sh
@@ -0,0 +1,95 @@
+#!/bin/bash
+# Copyright 2022 The Databend Authors.
+# SPDX-License-Identifier: Apache-2.0.
+
+set -e
+
+SCRIPT_PATH="$(cd "$(dirname "$0")" >/dev/null 2>&1 && pwd)"
+cd "$SCRIPT_PATH/../../" || exit
+
+killall -9 databend-query || true
+killall -9 databend-meta || true
+rm -rf .databend
+
+# enable private task
+CONFIG_FILE="./scripts/ci/deploy/config/databend-query-node-system-managed.toml"
+cat ./tests/task/private_task.toml >> "$CONFIG_FILE"
+
+echo "Starting Databend Query cluster enable private task"
+./scripts/ci/deploy/databend-query-system-managed.sh 2
+
+response=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"CREATE WAREHOUSE wh1 WITH WAREHOUSE_SIZE = '1'\"}")
+echo $response
+create_warehouse_1_query_id=$(echo $response | jq -r '.id')
+echo "Create WareHouse 1 Query ID: $create_warehouse_1_query_id"
+
+response=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"CREATE WAREHOUSE wh2 WITH WAREHOUSE_SIZE = '1'\"}")
+echo $response
+create_warehouse_2_query_id=$(echo $response | jq -r '.id')
+echo "Create WareHouse 2 Query ID: $create_warehouse_2_query_id"
+
+response=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"CREATE TABLE t1 (c1 int)\"}")
+echo $response
+create_table_query_id=$(echo $response | jq -r '.id')
+echo "Create Table Query ID: $create_table_query_id"
+
+response=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"CREATE TASK my_task_1 WAREHOUSE = 'wh1' SCHEDULE = 3 SECOND AS insert into t1 values(1)\"}")
+echo $response
+create_task_1_query_id=$(echo $response | jq -r '.id')
+echo "Create Task 1 Query ID: $create_task_1_query_id"
+
+response=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"CREATE TASK my_task_2 WAREHOUSE = 'wh2' SCHEDULE = 3 SECOND AS insert into t1 values(2)\"}")
+echo $response
+create_task_2_query_id=$(echo $response | jq -r '.id')
+echo "Create Task 2 ID: $create_task_2_query_id"
+
+sleep 2
+
+response=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"ALTER TASK my_task_1 RESUME\"}")
+echo $response
+resume_task_1_query_id=$(echo $response | jq -r '.id')
+echo "RESUME Task 1 ID: $resume_task_1_query_id"
+
+response=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"ALTER TASK my_task_2 RESUME\"}")
+echo $response
+resume_task_2_query_id=$(echo $response | jq -r '.id')
+echo "RESUME Task 2 ID: $resume_task_2_query_id"
+
+sleep 5
+
+response=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"SELECT c1 FROM t1 ORDER BY c1\"}")
+echo $response
+
+actual=$(echo "$response" | jq -c '.data')
+expected='[["1"],["2"]]'
+
+if [ "$actual" = "$expected" ]; then
+    echo "✅ Query result matches expected"
+else
+    echo "❌ Mismatch"
+    echo "Expected: $expected"
+    echo "Actual  : $actual"
+    exit 1
+fi
+
+response=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"SUSPEND WAREHOUSE wh2\"}")
+echo $response
+suspend_warehouse_2_query_id=$(echo $response | jq -r '.id')
+echo "Suspend WareHouse 2 Query ID: $suspend_warehouse_2_query_id"
+
+sleep 4
+
+response=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'X-DATABEND-WAREHOUSE: wh1' -H 'Content-Type: application/json' -d "{\"sql\": \"SELECT c1 FROM t1 ORDER BY c1\"}")
+echo $response
+
+actual=$(echo "$response" | jq -c '.data')
+expected='[["1"],["1"],["2"]]'
+
+if [ "$actual" = "$expected" ]; then
+    echo "✅ Query result matches expected"
+else
+    echo "❌ Mismatch"
+    echo "Expected: $expected"
+    echo "Actual  : $actual"
+    exit 1
+fi
diff --git a/tests/task/test-private-task.sh b/tests/task/test-private-task.sh
@@ -61,15 +61,15 @@ response=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-
 create_table_query_id=$(echo $response | jq -r '.id')
 echo "Create Table Query ID: $create_table_query_id"
 
-response1=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"CREATE TASK my_task_1 WAREHOUSE = 'mywh' SCHEDULE = 1 SECOND AS insert into t1 values(0)\"}")
+response1=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"CREATE TASK my_task_1 SCHEDULE = 1 SECOND AS insert into t1 values(0)\"}")
 create_task_1_query_id=$(echo $response1 | jq -r '.id')
 echo "Create Task 1 Query ID: $create_task_1_query_id"
 
-response2=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"CREATE TASK my_task_2 WAREHOUSE = 'mywh' SCHEDULE = 5 SECOND AS insert into t1 values(1)\"}")
+response2=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"CREATE TASK my_task_2 SCHEDULE = 5 SECOND AS insert into t1 values(1)\"}")
 create_task_2_query_id=$(echo $response2 | jq -r '.id')
 echo "Create Task 2 ID: $create_task_2_query_id"
 
-response3=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"CREATE TASK my_task_3 WAREHOUSE = 'mywh' AFTER 'my_task_1', 'my_task_2' AS insert into t1 values(2)\"}")
+response3=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"CREATE TASK my_task_3 AFTER 'my_task_1', 'my_task_2' AS insert into t1 values(2)\"}")
 create_task_3_query_id=$(echo $response3 | jq -r '.id')
 echo "Create Task 3 ID: $create_task_3_query_id"
 
@@ -234,7 +234,7 @@ response15=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Conten
 create_table_query_id_1=$(echo $response15 | jq -r '.id')
 echo "Create Table Query ID: $create_table_query_id_1"
 
-response16=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"CREATE TASK my_task_4 WAREHOUSE = 'mywh' SCHEDULE = USING CRON '*/5 * * * * ?' AS insert into t2 values(0)\"}")
+response16=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"CREATE TASK my_task_4 SCHEDULE = USING CRON '*/5 * * * * ?' AS insert into t2 values(0)\"}")
 create_task_4_query_id=$(echo $response16 | jq -r '.id')
 echo "Create Task 4 Query ID: $create_task_4_query_id"
 
@@ -274,15 +274,15 @@ response20=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Conten
 create_table_query_id_2=$(echo $response20 | jq -r '.id')
 echo "Create Table Query ID: $create_table_query_id_2"
 
-response21=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"CREATE TASK my_task_5 WAREHOUSE = 'mywh' SCHEDULE = 3 SECOND WHEN EXISTS (SELECT 1 FROM t3 WHERE c2 = 1) AS insert into t3 values(1, 0)\"}")
+response21=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"CREATE TASK my_task_5 SCHEDULE = 3 SECOND WHEN EXISTS (SELECT 1 FROM t3 WHERE c2 = 1) AS insert into t3 values(1, 0)\"}")
 create_task_5_query_id=$(echo $response21 | jq -r '.id')
 echo "Create Task 5 Query ID: $create_task_5_query_id"
 
-response22=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"CREATE TASK my_task_6 WAREHOUSE = 'mywh' SCHEDULE = 5 SECOND WHEN EXISTS (SELECT 1 FROM t3 WHERE c2 = 1) AS insert into t3 values(2, 0)\"}")
+response22=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"CREATE TASK my_task_6 SCHEDULE = 5 SECOND WHEN EXISTS (SELECT 1 FROM t3 WHERE c2 = 1) AS insert into t3 values(2, 0)\"}")
 create_task_6_query_id=$(echo $response22 | jq -r '.id')
 echo "Create Task 6 Query ID: $create_task_6_query_id"
 
-response23=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"CREATE TASK my_task_7 WAREHOUSE = 'mywh' AFTER 'my_task_5', 'my_task_6' WHEN EXISTS (SELECT 1 FROM t3 WHERE c2 = 2)  AS insert into t3 values(3, 0)\"}")
+response23=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d "{\"sql\": \"CREATE TASK my_task_7 AFTER 'my_task_5', 'my_task_6' WHEN EXISTS (SELECT 1 FROM t3 WHERE c2 = 2)  AS insert into t3 values(3, 0)\"}")
 create_task_7_query_id=$(echo $response23 | jq -r '.id')
 echo "Create Task 7 Query ID: $create_task_7_query_id"
 
