chore: add VACUUM TEMPORARY TABLES command (#18349)

* support limit

* new syntax

Author: SkyFan2002

src/query/ast/src/parser/statement.rs

@@ -1904,6 +1904,18 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
         },
     );
 
+    let vacuum_temporary_tables = map(
+        rule! {
+            VACUUM ~ TEMPORARY ~ TABLES ~ ( LIMIT ~ ^#literal_u64 )?
+        },
+        |(_, _, _, opt_limit)| {
+            Statement::Call(CallStmt {
+                name: "fuse_vacuum_temporary_table".to_string(),
+                args: opt_limit.map(|v| v.1.to_string()).into_iter().collect(),
+            })
+        },
+    );
+
     let presign = map(
         rule! {
             PRESIGN ~ ( #presign_action )?
@@ -2726,6 +2738,7 @@ AS
             | #call_procedure : "`CALL PROCEDURE <procedure_name>()`"
         ),
         rule!(#comment),
+        rule!(#vacuum_temporary_tables),
     ))(i)
 }
 

src/query/storages/fuse/src/table_functions/fuse_vacuum_temporary_table.rs

@@ -16,11 +16,14 @@ use std::collections::HashSet;
 use std::sync::Arc;
 
 use databend_common_catalog::plan::DataSourcePlan;
+use databend_common_catalog::table_args::TableArgs;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
+use databend_common_expression::types::NumberScalar;
 use databend_common_expression::types::StringType;
 use databend_common_expression::DataBlock;
 use databend_common_expression::FromData;
+use databend_common_expression::Scalar;
 use databend_common_expression::TableDataType;
 use databend_common_expression::TableField;
 use databend_common_expression::TableSchemaRef;
@@ -33,8 +36,10 @@ use log::info;
 
 use crate::sessions::TableContext;
 use crate::table_functions::SimpleTableFunc;
-use crate::table_functions::TableArgs;
-pub struct FuseVacuumTemporaryTable;
+
+pub struct FuseVacuumTemporaryTable {
+    limit: Option<u64>,
+}
 
 #[async_trait::async_trait]
 impl SimpleTableFunc for FuseVacuumTemporaryTable {
@@ -43,7 +48,11 @@ impl SimpleTableFunc for FuseVacuumTemporaryTable {
     }
 
     fn table_args(&self) -> Option<TableArgs> {
-        None
+        self.limit.map(|limit| {
+            TableArgs::new_positioned(vec![databend_common_catalog::table_args::u64_literal(
+                limit,
+            )])
+        })
     }
 
     fn schema(&self) -> TableSchemaRef {
@@ -62,17 +71,39 @@ impl SimpleTableFunc for FuseVacuumTemporaryTable {
             .await?;
         let client_session_mgr = UserApiProvider::instance().client_session_api(&ctx.get_tenant());
         let mut user_session_ids = HashSet::new();
+        let mut inactive_user_session_ids = HashSet::new();
         while let Some(entry) = lister.try_next().await? {
+            if entry.metadata().is_dir() {
+                continue;
+            }
             let path = entry.path();
             let parts: Vec<_> = path.split('/').collect();
             if parts.len() < 3 {
                 return Err(ErrorCode::Internal(format!(
                     "invalid path for temp table: {path}"
                 )));
             };
-            user_session_ids.insert((parts[1].to_string(), parts[2].to_string()));
+            let user_name = parts[1].to_string();
+            let session_id = parts[2].to_string();
+            if user_session_ids.contains(&(user_name.clone(), session_id.clone())) {
+                continue;
+            }
+            user_session_ids.insert((user_name.clone(), session_id.clone()));
+            if client_session_mgr
+                .get_client_session(&user_name, &session_id)
+                .await?
+                .is_none()
+            {
+                inactive_user_session_ids.insert((user_name, session_id));
+                if inactive_user_session_ids.len() >= self.limit.unwrap_or(u64::MAX) as usize {
+                    break;
+                }
+            }
         }
-        for (user_name, session_id) in user_session_ids {
+
+        let session_num = inactive_user_session_ids.len();
+
+        for (user_name, session_id) in inactive_user_session_ids {
             if client_session_mgr
                 .get_client_session(&user_name, &session_id)
                 .await?
@@ -86,15 +117,43 @@ impl SimpleTableFunc for FuseVacuumTemporaryTable {
                 op.remove_all(&path).await?;
             }
         }
-        let col: Vec<String> = vec!["Ok".to_owned()];
+        let col: Vec<String> = vec![format!(
+            "Ok: processed temporary tables from {} inactive sessions",
+            session_num
+        )];
 
         Ok(Some(DataBlock::new_from_columns(vec![
             StringType::from_data(col),
         ])))
     }
 
-    fn create(_func_name: &str, _table_args: TableArgs) -> Result<Self>
+    fn create(func_name: &str, table_args: TableArgs) -> Result<Self>
     where Self: Sized {
-        Ok(Self)
+        let limit = match table_args.positioned.len() {
+            0 => None,
+            1 => {
+                let args = table_args.expect_all_positioned(func_name, Some(1))?;
+                let limit_val = match &args[0] {
+                    Scalar::Number(NumberScalar::UInt64(val)) => *val,
+                    Scalar::Number(NumberScalar::UInt32(val)) => *val as u64,
+                    Scalar::Number(NumberScalar::UInt16(val)) => *val as u64,
+                    Scalar::Number(NumberScalar::UInt8(val)) => *val as u64,
+                    Scalar::String(val) => val.parse::<u64>()?,
+                    _ => {
+                        return Err(ErrorCode::BadArguments(format!(
+                            "invalid value {:?} expect to be unsigned integer literal.",
+                            args[0]
+                        )))
+                    }
+                };
+                Some(limit_val)
+            }
+            _ => {
+                return Err(ErrorCode::NumberArgumentsNotMatch(
+                    "Expected 0 or 1 arguments".to_string(),
+                ));
+            }
+        };
+        Ok(Self { limit })
     }
 }