feat: imp when_condition for Task & store TaskRun on system table

Author: KKould

src/binaries/query/entry.rs

@@ -44,6 +44,7 @@ use databend_query::servers::MySQLHandler;
 use databend_query::servers::MySQLTlsConfig;
 use databend_query::servers::Server;
 use databend_query::servers::ShutdownHandle;
+use databend_query::task::TaskService;
 use databend_query::GlobalServices;
 use log::info;
 
@@ -302,6 +303,7 @@ pub async fn start_services(conf: &InnerConfig) -> Result<(), MainError> {
         }
         println!("    system history tables: {}", conf.log.history);
     }
+    TaskService::instance().initialized();
 
     println!();
     println!(

src/common/exception/src/exception_code.rs

@@ -403,6 +403,8 @@ build_exceptions! {
     TaskCronInvalid(2614),
     /// Task schedule and after conflict
     TaskScheduleAndAfterConflict(2615),
+    /// Task when condition not met
+    TaskWhenConditionNotMet(2616),
 }
 
 // Search and External Service Errors [1901-1903, 1910]

src/query/management/src/task/task_mgr.rs

@@ -103,13 +103,21 @@ impl TaskMgr {
     ) -> Result<Result<(), TaskError>, TaskApiError> {
         task.created_at = Utc::now();
 
-        self.create_task_inner(task, create_option).await
+        self.create_task_inner(task, create_option, false).await
+    }
+
+    #[async_backtrace::framed]
+    #[fastrace::trace]
+    pub async fn update_task(&self, task: Task) -> Result<Result<(), TaskError>, TaskApiError> {
+        self.create_task_inner(task, &CreateOption::CreateOrReplace, true)
+            .await
     }
 
     async fn create_task_inner(
         &self,
         task: Task,
         create_option: &CreateOption,
+        without_schedule: bool,
     ) -> Result<Result<(), TaskError>, TaskApiError> {
         assert!(task.after.is_empty() || task.schedule_options.is_none());
         // check
@@ -152,9 +160,13 @@ impl TaskMgr {
             }
         }
         if !task.after.is_empty() {
-            let _ = TaskChannel::instance().send(TaskMessage::AfterTask(task));
-        } else if task.schedule_options.is_some() {
-            let _ = TaskChannel::instance().send(TaskMessage::ScheduleTask(task));
+            let _ = TaskChannel::instance()
+                .send(TaskMessage::AfterTask(task))
+                .await;
+        } else if task.schedule_options.is_some() && !without_schedule {
+            let _ = TaskChannel::instance()
+                .send(TaskMessage::ScheduleTask(task))
+                .await;
         }
 
         Ok(Ok(()))
@@ -174,7 +186,9 @@ impl TaskMgr {
                 context: "while execute task".to_string(),
             }));
         };
-        let _ = TaskChannel::instance().send(TaskMessage::ExecuteTask(Task::clone(&task)));
+        let _ = TaskChannel::instance()
+            .send(TaskMessage::ExecuteTask(Task::clone(&task)))
+            .await;
 
         Ok(Ok(()))
     }
@@ -256,7 +270,7 @@ impl TaskMgr {
             }
         }
         if let Err(e) = self
-            .create_task_inner(task, &CreateOption::CreateOrReplace)
+            .create_task_inner(task, &CreateOption::CreateOrReplace, false)
             .await?
         {
             return Ok(Err(TaskError::NotFound {
@@ -322,50 +336,6 @@ impl TaskMgr {
         Ok(tasks)
     }
 
-    #[async_backtrace::framed]
-    #[fastrace::trace]
-    pub async fn update_task_run(
-        &self,
-        task_run: TaskRun,
-    ) -> Result<Result<(), TaskError>, TaskApiError> {
-        let seq = MatchSeq::from(CreateOption::CreateOrReplace);
-        let key = TaskRunIdent::new(&self.tenant, task_run.key());
-        let req = UpsertPB::insert(key, task_run.clone()).with(seq);
-        let _ = self.kv_api.upsert_pb(&req).await?;
-
-        Ok(Ok(()))
-    }
-
-    #[async_backtrace::framed]
-    #[fastrace::trace]
-    pub async fn lasted_task_run(
-        &self,
-        task_name: &str,
-    ) -> Result<Result<Option<TaskRun>, TaskError>, TaskApiError> {
-        let key = DirName::new(TaskRunIdent::new(&self.tenant, task_name));
-        let result = self
-            .kv_api
-            .list_pb_values(&key)
-            .await?
-            .next()
-            .await
-            .transpose()?;
-
-        Ok(Ok(result))
-    }
-
-    #[async_backtrace::framed]
-    #[fastrace::trace]
-    pub async fn show_task_runs_full(
-        &self,
-        task_name: &str,
-    ) -> Result<Result<Vec<TaskRun>, TaskApiError>, TaskApiError> {
-        let key = DirName::new(TaskRunIdent::new(&self.tenant, task_name));
-        let stream = self.kv_api.list_pb_values(&key).await?;
-
-        Ok(Ok(stream.try_collect().await?))
-    }
-
     pub fn make_schedule_options(opt: ScheduleOptions) -> task::ScheduleOptions {
         match opt {
             ScheduleOptions::IntervalSecs(secs, ms) => {

src/query/service/src/global_services.rs

@@ -58,7 +58,7 @@ use crate::servers::http::v1::ClientSessionManager;
 use crate::servers::http::v1::HttpQueryManager;
 use crate::sessions::QueriesQueueManager;
 use crate::sessions::SessionManager;
-use crate::task_service::TaskService;
+use crate::task::task_service::TaskService;
 
 pub struct GlobalServices;
 
@@ -123,8 +123,7 @@ impl GlobalServices {
         SessionManager::init(config)?;
         LockManager::init()?;
         AuthMgr::init(config)?;
-        
-        
+
         let task_rx = TaskChannel::init(GlobalConfig::instance().query.tasks_channel_len)?;
         // Init user manager.
         // Builtin users and udfs are created here.
@@ -146,7 +145,6 @@ impl GlobalServices {
             )
             .await?;
         }
-        TaskService::init(task_rx, config.query.tenant_id.clone())?;
         RoleCacheManager::init()?;
 
         DataOperator::init(&config.storage, config.spill.storage_params.clone()).await?;
@@ -181,6 +179,7 @@ impl GlobalServices {
         if config.log.history.on {
             GlobalHistoryLog::init(config).await?;
         }
+        TaskService::init(task_rx, config)?;
 
         GLOBAL_QUERIES_MANAGER.set_gc_handle(memory_gc_handle);
 

src/query/service/src/interpreters/task/mod.rs

@@ -15,7 +15,6 @@
 use std::sync::Arc;
 
 use databend_common_cloud_control::task_utils;
-use databend_common_config::GlobalConfig;
 use databend_common_exception::Result;
 use databend_common_sql::plans::AlterTaskPlan;
 use databend_common_sql::plans::CreateTaskPlan;
@@ -37,7 +36,7 @@ impl TaskInterpreterFactory {
     pub fn build() -> TaskInterpreterImpl {
         // TODO: for test
         if true {
-        // if GlobalConfig::instance().query.enable_private_task {
+            // if GlobalConfig::instance().query.enable_private_task {
             return TaskInterpreterImpl::Private(PrivateTaskInterpreter);
         }
         TaskInterpreterImpl::Cloud(CloudTaskInterpreter)

src/query/service/src/interpreters/task/private.rs

@@ -163,9 +163,6 @@ impl TaskInterpreter for PrivateTaskInterpreter {
     ) -> Result<Vec<task_utils::Task>> {
         let tasks = UserApiProvider::instance().show_tasks(&plan.tenant).await?;
 
-        tasks
-            .into_iter()
-            .map(Self::task_trans)
-            .try_collect()
+        tasks.into_iter().map(Self::task_trans).try_collect()
     }
 }

src/query/service/src/lib.rs

@@ -61,7 +61,7 @@ pub mod table_functions;
 pub mod test_kits;
 
 mod global_services;
-mod task_service;
+pub mod task;
 
 pub use databend_common_sql as sql;
 pub use databend_common_storages_factory as storages;

src/query/service/src/task/meta.rs

@@ -0,0 +1,133 @@
+// 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::sync::Arc;
+use std::time::Duration;
+
+use databend_common_base::runtime::block_on;
+use databend_common_exception::Result;
+use databend_common_meta_client::ClientHandle;
+use databend_common_meta_kvapi::kvapi::KVApi;
+use databend_common_meta_semaphore::acquirer::Permit;
+use databend_common_meta_semaphore::Semaphore;
+use databend_common_meta_types::MatchSeq;
+use databend_common_meta_types::Operation;
+use databend_common_meta_types::UpsertKV;
+
+/// refer to [crate::history_tables::meta::PermitGuard]
+pub struct PermitGuard {
+    _permit: Permit,
+    meta_handle: Arc<TaskMetaHandle>,
+    meta_key: String,
+}
+
+impl PermitGuard {
+    pub fn new(permit: Permit, meta_handle: Arc<TaskMetaHandle>, meta_key: String) -> Self {
+        Self {
+            _permit: permit,
+            meta_handle,
+            meta_key,
+        }
+    }
+}
+
+impl Drop for PermitGuard {
+    fn drop(&mut self) {
+        let meta_handle = self.meta_handle.clone();
+        let meta_key = self.meta_key.clone();
+
+        block_on(async move {
+            let _ = meta_handle.update_last_execution_timestamp(&meta_key).await;
+        });
+    }
+}
+
+/// refer to [crate::history_tables::meta::HistoryMetaHandle]
+pub struct TaskMetaHandle {
+    meta_client: Arc<ClientHandle>,
+    node_id: String,
+}
+
+impl TaskMetaHandle {
+    pub fn new(meta_client: Arc<ClientHandle>, node_id: String) -> Self {
+        Self {
+            meta_client,
+            node_id,
+        }
+    }
+
+    pub async fn acquire(&self, meta_key: &str, interval: u64) -> Result<Option<Permit>> {
+        let acquired_guard = Semaphore::new_acquired(
+            self.meta_client.clone(),
+            meta_key,
+            1,
+            self.node_id.clone(),
+            Duration::from_secs(3),
+        )
+        .await
+        .map_err(|_e| "acquire semaphore failed from TaskService")?;
+        if interval == 0 {
+            return Ok(Some(acquired_guard));
+        }
+        if match self
+            .meta_client
+            .get_kv(&format!("{}/last_timestamp", meta_key))
+            .await?
+        {
+            Some(v) => {
+                let last: u64 = serde_json::from_slice(&v.data)?;
+                chrono::Utc::now().timestamp_millis() as u64
+                    - Duration::from_secs(interval).as_millis() as u64
+                    > last
+            }
+            None => true,
+        } {
+            Ok(Some(acquired_guard))
+        } else {
+            drop(acquired_guard);
+            Ok(None)
+        }
+    }
+
+    pub async fn acquire_with_guard(
+        &self,
+        meta_key: &str,
+        interval: u64,
+    ) -> Result<Option<PermitGuard>> {
+        if let Some(permit) = self.acquire(meta_key, interval).await? {
+            Ok(Some(PermitGuard::new(
+                permit,
+                Arc::new(TaskMetaHandle {
+                    meta_client: self.meta_client.clone(),
+                    node_id: self.node_id.clone(),
+                }),
+                meta_key.to_string(),
+            )))
+        } else {
+            Ok(None)
+        }
+    }
+
+    pub async fn update_last_execution_timestamp(&self, meta_key: &str) -> Result<()> {
+        self.meta_client
+            .upsert_kv(UpsertKV::new(
+                format!("{}/last_timestamp", meta_key),
+                MatchSeq::Any,
+                Operation::Update(serde_json::to_vec(&chrono::Utc::now().timestamp_millis())?),
+                None,
+            ))
+            .await?;
+        Ok(())
+    }
+}

src/query/service/src/task/mod.rs

@@ -0,0 +1,19 @@
+// 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.
+
+mod meta;
+mod session;
+pub mod task_service;
+
+pub use task_service::TaskService;