chore: add license for private task

Author: KKould

src/common/license/src/license.rs

@@ -83,6 +83,8 @@ pub enum Feature {
     SystemHistory,
     #[serde(alias = "vector_index", alias = "VECTOR_INDEX")]
     VectorIndex,
+    #[serde(alias = "private_task", alias = "PRIVATE_TASK")]
+    PrivateTask,
     #[serde(other)]
     Unknown,
 }
@@ -134,6 +136,7 @@ impl fmt::Display for Feature {
             Feature::WorkloadGroup => write!(f, "workload_group"),
             Feature::SystemHistory => write!(f, "system_history"),
             Feature::VectorIndex => write!(f, "vector_index"),
+            Feature::PrivateTask => write!(f, "private_task"),
             Feature::Unknown => write!(f, "unknown"),
         }
     }
@@ -372,6 +375,11 @@ mod tests {
             serde_json::from_str::<Feature>("\"VectorIndex\"").unwrap()
         );
 
+        assert_eq!(
+            Feature::PrivateTask,
+            serde_json::from_str::<Feature>("\"private_task\"").unwrap()
+        );
+
         assert_eq!(
             Feature::Unknown,
             serde_json::from_str::<Feature>("\"ssss\"").unwrap()
@@ -408,6 +416,7 @@ mod tests {
                 Feature::NgramIndex,
                 Feature::WorkloadGroup,
                 Feature::SystemHistory,
+                Feature::PrivateTask,
             ]),
         };
 

src/query/service/src/interpreters/interpreter_task_alter.rs

@@ -14,6 +14,7 @@
 
 use std::sync::Arc;
 
+use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::Result;
 use databend_common_sql::plans::AlterTaskPlan;
 
@@ -50,7 +51,7 @@ impl Interpreter for AlterTaskInterpreter {
     #[fastrace::trace]
     #[async_backtrace::framed]
     async fn execute2(&self) -> Result<PipelineBuildResult> {
-        TaskInterpreterFactory::build()
+        TaskInterpreterFactory::build(&self.ctx)
             .alter_task(&self.ctx, &self.plan)
             .await?;
 

src/query/service/src/interpreters/interpreter_task_create.rs

@@ -48,7 +48,7 @@ impl Interpreter for CreateTaskInterpreter {
     #[fastrace::trace]
     #[async_backtrace::framed]
     async fn execute2(&self) -> Result<PipelineBuildResult> {
-        TaskInterpreterFactory::build()
+        TaskInterpreterFactory::build(&self.ctx)
             .create_task(&self.ctx, &self.plan)
             .await?;
 

src/query/service/src/interpreters/interpreter_task_describe.rs

@@ -49,7 +49,7 @@ impl Interpreter for DescribeTaskInterpreter {
     #[fastrace::trace]
     #[async_backtrace::framed]
     async fn execute2(&self) -> Result<PipelineBuildResult> {
-        let Some(task) = TaskInterpreterFactory::build()
+        let Some(task) = TaskInterpreterFactory::build(&self.ctx)
             .describe_task(&self.ctx, &self.plan)
             .await?
         else {

src/query/service/src/interpreters/interpreter_task_drop.rs

@@ -48,7 +48,7 @@ impl Interpreter for DropTaskInterpreter {
     #[fastrace::trace]
     #[async_backtrace::framed]
     async fn execute2(&self) -> Result<PipelineBuildResult> {
-        TaskInterpreterFactory::build()
+        TaskInterpreterFactory::build(&self.ctx)
             .drop_task(&self.ctx, &self.plan)
             .await?;
 

src/query/service/src/interpreters/interpreter_task_execute.rs

@@ -48,7 +48,7 @@ impl Interpreter for ExecuteTaskInterpreter {
     #[fastrace::trace]
     #[async_backtrace::framed]
     async fn execute2(&self) -> Result<PipelineBuildResult> {
-        TaskInterpreterFactory::build()
+        TaskInterpreterFactory::build(&self.ctx)
             .execute_task(&self.ctx, &self.plan)
             .await?;
 

src/query/service/src/interpreters/interpreter_tasks_show.rs

@@ -51,7 +51,7 @@ impl Interpreter for ShowTasksInterpreter {
     #[fastrace::trace]
     #[async_backtrace::framed]
     async fn execute2(&self) -> Result<PipelineBuildResult> {
-        let tasks = TaskInterpreterFactory::build()
+        let tasks = TaskInterpreterFactory::build(&self.ctx)
             .show_tasks(&self.ctx, &self.plan)
             .await?;
 

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

@@ -14,9 +14,13 @@
 
 use std::sync::Arc;
 
+use databend_common_catalog::table_context::TableContext;
 use databend_common_cloud_control::task_utils;
 use databend_common_config::GlobalConfig;
 use databend_common_exception::Result;
+use databend_common_license::license::Feature;
+use databend_common_license::license_manager::LicenseManagerSwitch;
+use databend_common_settings::Settings;
 use databend_common_sql::plans::AlterTaskPlan;
 use databend_common_sql::plans::CreateTaskPlan;
 use databend_common_sql::plans::DescribeTaskPlan;
@@ -34,8 +38,12 @@ mod private;
 pub(crate) struct TaskInterpreterFactory;
 
 impl TaskInterpreterFactory {
-    pub fn build() -> TaskInterpreterImpl {
+    pub fn build(ctx: &QueryContext) -> TaskInterpreterImpl {
         if GlobalConfig::instance().task.on {
+            LicenseManagerSwitch::instance().check_enterprise_enabled(
+                ctx.get_settings().get_enterprise_license(),
+                Feature::PrivateTask,
+            )?;
             return TaskInterpreterImpl::Private(PrivateTaskInterpreter);
         }
         TaskInterpreterImpl::Cloud(CloudTaskInterpreter)

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

@@ -382,7 +382,15 @@ impl TaskService {
                                                     .ok_or_else(|| {
                                                         ErrorCode::UnknownTask(next_task)
                                                     })?;
-                                                if let Some(_guard) = fn_lock(&TaskService::instance(), &TaskMessageIdent::new(tenant.clone(), format!("check_{}", task_name))).await? {
+                                                if let Some(_guard) = fn_lock(
+                                                    &TaskService::instance(),
+                                                    &TaskMessageIdent::new(
+                                                        tenant.clone(),
+                                                        format!("check_{}", task_name),
+                                                    ),
+                                                )
+                                                .await?
+                                                {
                                                     task_mgr
                                                         .send(TaskMessage::ExecuteTask(next_task))
                                                         .await?;