chore(query): mark the 'create as select' query as heavy (#18365)

Author: zhang2014

src/query/service/src/sessions/queue_mgr.rs

@@ -174,56 +174,59 @@ impl<Data: QueueData> QueueManager<Data> {
     }
 
     pub async fn acquire(self: &Arc<Self>, data: Data) -> Result<AcquireQueueGuard> {
-        if data.need_acquire_to_queue() {
-            info!(
-                "[QUERY-QUEUE] Preparing to acquire from query queue, current length: {}",
-                self.length()
-            );
-
-            let start_time = SystemTime::now();
-            let instant = Instant::now();
-            let mut timeout = data.timeout();
-            let mut guards = vec![];
-
-            let data = Arc::new(data);
-            if let Some(workload_group) = ThreadTracker::workload_group() {
-                if let Some(QuotaValue::Number(permits)) =
-                    workload_group.meta.get_quota(MAX_CONCURRENCY_QUOTA_KEY)
-                {
-                    let mut workload_group_timeout = timeout;
-
-                    if let Some(QuotaValue::Duration(queue_timeout)) = workload_group
-                        .meta
-                        .get_quota(QUERY_QUEUED_TIMEOUT_QUOTA_KEY)
-                    {
-                        workload_group_timeout =
-                            std::cmp::min(queue_timeout, workload_group_timeout);
-                    }
+        if !data.need_acquire_to_queue() {
+            info!("[QUERY-QUEUE] Non-heavy queries skip the query queue and execute directly.");
+            return Ok(AcquireQueueGuard::create(vec![]));
+        }
 
-                    let workload_queue_guard = self
-                        .acquire_workload_queue(
-                            data.clone(),
-                            workload_group.queue_key.clone(),
-                            permits as u64,
-                            workload_group_timeout,
-                        )
-                        .await?;
-
-                    info!("[QUERY-QUEUE] Successfully acquired from workload group queue. elapsed: {:?}", instant.elapsed());
-                    timeout -= instant.elapsed();
-                    guards.push(workload_queue_guard);
-                }
-            }
+        info!(
+            "[QUERY-QUEUE] Preparing to acquire from query queue, current length: {}",
+            self.length()
+        );
 
-            guards.push(self.acquire_warehouse_queue(data, timeout).await?);
+        let start_time = SystemTime::now();
+        let instant = Instant::now();
+        let mut timeout = data.timeout();
+        let mut guards = vec![];
+
+        let data = Arc::new(data);
+        if let Some(workload_group) = ThreadTracker::workload_group() {
+            if let Some(QuotaValue::Number(permits)) =
+                workload_group.meta.get_quota(MAX_CONCURRENCY_QUOTA_KEY)
+            {
+                let mut workload_group_timeout = timeout;
+
+                if let Some(QuotaValue::Duration(queue_timeout)) = workload_group
+                    .meta
+                    .get_quota(QUERY_QUEUED_TIMEOUT_QUOTA_KEY)
+                {
+                    workload_group_timeout = std::cmp::min(queue_timeout, workload_group_timeout);
+                }
 
-            inc_session_running_acquired_queries();
-            record_session_queue_acquire_duration_ms(start_time.elapsed().unwrap_or_default());
+                let workload_queue_guard = self
+                    .acquire_workload_queue(
+                        data.clone(),
+                        workload_group.queue_key.clone(),
+                        permits as u64,
+                        workload_group_timeout,
+                    )
+                    .await?;
 
-            return Ok(AcquireQueueGuard::create(guards));
+                info!(
+                    "[QUERY-QUEUE] Successfully acquired from workload group queue. elapsed: {:?}",
+                    instant.elapsed()
+                );
+                timeout -= instant.elapsed();
+                guards.push(workload_queue_guard);
+            }
         }
 
-        Ok(AcquireQueueGuard::create(vec![]))
+        guards.push(self.acquire_warehouse_queue(data, timeout).await?);
+
+        inc_session_running_acquired_queries();
+        record_session_queue_acquire_duration_ms(start_time.elapsed().unwrap_or_default());
+
+        Ok(AcquireQueueGuard::create(guards))
     }
 
     async fn acquire_workload_queue(
@@ -581,6 +584,9 @@ impl QueryEntry {
             | Plan::TruncateTable(_) => {
                 return true;
             }
+            Plan::CreateTable(v) if v.as_select.is_some() => {
+                return true;
+            }
             Plan::DropTable(v) if v.all => {
                 return true;
             }

src/query/service/tests/it/sessions/queue_mgr.rs

@@ -382,6 +382,10 @@ async fn test_heavy_actions() -> Result<()> {
             // MERGE INTO
             add_to_queue: true,
         },
+        Query {
+            sql: "CREATE TABLE test_heavy_create AS SELECT 1",
+            add_to_queue: true,
+        }
     ];
 
     let fixture = TestFixture::setup().await?;