feat: support vacuum temporary table

src/query/catalog/src/catalog/interface.rs

@@ -225,7 +225,7 @@ pub trait Catalog: DynClone + Send + Sync + Debug {
     /// Get the table meta by table id.
     async fn get_table_meta_by_id(&self, table_id: u64) -> Result<Option<SeqV<TableMeta>>>;
 
-    /// List the tables name by meta ids.
+    /// List the tables name by meta ids. This function should not be used to list temporary tables.
     async fn mget_table_names_by_ids(
         &self,
         tenant: &Tenant,
@@ -253,7 +253,15 @@ pub trait Catalog: DynClone + Send + Sync + Debug {
         table_name: &str,
     ) -> Result<Arc<dyn Table>>;
 
+    /// List all tables in a database.This will not list temporary tables.
     async fn list_tables(&self, tenant: &Tenant, db_name: &str) -> Result<Vec<Arc<dyn Table>>>;
+
+    fn list_temporary_tables(&self) -> Result<Vec<TableInfo>> {
+        Err(ErrorCode::Unimplemented(
+            "'list_temporary_tables' not implemented",
+        ))
+    }
+
     async fn list_tables_history(
         &self,
         tenant: &Tenant,

src/query/catalog/src/table_context.rs

@@ -363,7 +363,7 @@ pub trait TableContext: Send + Sync {
         lock_opt: &LockTableOption,
     ) -> Result<Option<Arc<LockGuard>>>;
 
-    fn get_session_id(&self) -> String;
+    fn get_session_id(&self) -> Result<String>;
 
     fn session_state(&self) -> SessionState;
 

src/query/management/src/client_session.rs

@@ -117,6 +117,18 @@ impl ClientSessionMgr {
         Ok(res.prev.is_none())
     }
 
+    #[async_backtrace::framed]
+    #[fastrace::trace]
+    pub async fn get_client_session(
+        &self,
+        client_session_id: &str,
+    ) -> Result<Option<ClientSession>> {
+        let ident = self.session_ident(client_session_id);
+        let res = self.kv_api.get_pb(&ident).await?;
+
+        Ok(res.map(|r| r.data))
+    }
+
     #[async_backtrace::framed]
     #[fastrace::trace]
     pub async fn drop_client_session_id(&self, client_session_id: &str) -> Result<()> {

src/query/service/src/catalogs/default/database_catalog.rs

@@ -416,6 +416,10 @@ impl Catalog for DatabaseCatalog {
         }
     }
 
+    fn list_temporary_tables(&self) -> Result<Vec<TableInfo>> {
+        self.mutable_catalog.list_temporary_tables()
+    }
+
     #[async_backtrace::framed]
     async fn list_tables_history(
         &self,

src/query/service/src/catalogs/default/session_catalog.rs

@@ -365,6 +365,10 @@ impl Catalog for SessionCatalog {
         self.inner.list_tables(tenant, db_name).await
     }
 
+    fn list_temporary_tables(&self) -> Result<Vec<TableInfo>> {
+        self.temp_tbl_mgr.lock().list_tables()
+    }
+
     async fn list_tables_history(
         &self,
         tenant: &Tenant,

src/query/service/src/databases/system/system_database.rs

@@ -60,6 +60,7 @@ use databend_common_storages_system::TablesTableWithoutHistory;
 use databend_common_storages_system::TaskHistoryTable;
 use databend_common_storages_system::TasksTable;
 use databend_common_storages_system::TempFilesTable;
+use databend_common_storages_system::TemporaryTablesTable;
 use databend_common_storages_system::TerseStreamsTable;
 use databend_common_storages_system::UserFunctionsTable;
 use databend_common_storages_system::UsersTable;
@@ -140,6 +141,7 @@ impl SystemDatabase {
             NotificationHistoryTable::create(sys_db_meta.next_table_id()),
             ViewsTableWithHistory::create(sys_db_meta.next_table_id()),
             ViewsTableWithoutHistory::create(sys_db_meta.next_table_id()),
+            TemporaryTablesTable::create(sys_db_meta.next_table_id()),
         ];
 
         let disable_tables = Self::disable_system_tables();

src/query/service/src/sessions/query_ctx.rs

@@ -1335,8 +1335,17 @@ impl TableContext for QueryContext {
         Ok(lock_guard)
     }
 
-    fn get_session_id(&self) -> String {
-        self.shared.session.id.clone()
+    fn get_session_id(&self) -> Result<String> {
+        let session_type = self.shared.session.get_type();
+        match session_type {
+            SessionType::MySQL => Ok(self.shared.session.id.clone()),
+            _ => self
+                .shared
+                .session
+                .session_ctx
+                .get_client_session_id()
+                .ok_or(ErrorCode::Internal("No client session id".to_string())),
+        }
     }
 
     fn is_temp_table(&self, catalog_name: &str, database_name: &str, table_name: &str) -> bool {

src/query/service/src/table_functions/table_function_factory.rs

@@ -25,6 +25,7 @@ use databend_common_storages_fuse::table_functions::FuseBlockFunc;
 use databend_common_storages_fuse::table_functions::FuseColumnFunc;
 use databend_common_storages_fuse::table_functions::FuseEncodingFunc;
 use databend_common_storages_fuse::table_functions::FuseStatisticsFunc;
+use databend_common_storages_fuse::table_functions::FuseVacuumTemporaryTable;
 use databend_common_storages_fuse::table_functions::TableFunctionTemplate;
 use databend_common_storages_stream::stream_status_table_func::StreamStatusTable;
 use databend_storages_common_table_meta::table_id_ranges::SYS_TBL_FUC_ID_END;
@@ -182,6 +183,14 @@ impl TableFunctionFactory {
             ),
         );
 
+        creators.insert(
+            "fuse_vacuum_temporary_table".to_string(),
+            (
+                next_id(),
+                Arc::new(TableFunctionTemplate::<FuseVacuumTemporaryTable>::create),
+            ),
+        );
+
         creators.insert(
             "stream_status".to_string(),
             (next_id(), Arc::new(StreamStatusTable::create)),

src/query/service/tests/it/sql/exec/get_table_bind_test.rs

@@ -977,7 +977,7 @@ impl TableContext for CtxDelegation {
         todo!()
     }
 
-    fn get_session_id(&self) -> String {
+    fn get_session_id(&self) -> Result<String> {
         todo!()
     }
 

src/query/service/tests/it/storages/fuse/operations/commit.rs

@@ -871,7 +871,7 @@ impl TableContext for CtxDelegation {
         todo!()
     }
 
-    fn get_session_id(&self) -> String {
+    fn get_session_id(&self) -> Result<String> {
         todo!()
     }
 

src/query/service/tests/it/storages/testdata/columns_table.txt

@@ -118,6 +118,7 @@ DB.Table: 'system'.'columns', Table: columns-table_id:1, ver:0, Engine: SystemCo
 | 'database'                        | 'system'             | 'streams_terse'        | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
 | 'database'                        | 'system'             | 'tables'               | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
 | 'database'                        | 'system'             | 'tables_with_history'  | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
+| 'database'                        | 'system'             | 'temporary_tables'     | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
 | 'database'                        | 'system'             | 'views'                | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
 | 'database'                        | 'system'             | 'views_with_history'   | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
 | 'database'                        | 'system'             | 'virtual_columns'      | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
@@ -157,6 +158,7 @@ DB.Table: 'system'.'columns', Table: columns-table_id:1, ver:0, Engine: SystemCo
 | 'engine'                          | 'information_schema' | 'tables'               | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
 | 'engine'                          | 'system'             | 'tables'               | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
 | 'engine'                          | 'system'             | 'tables_with_history'  | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
+| 'engine'                          | 'system'             | 'temporary_tables'     | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
 | 'engine'                          | 'system'             | 'views'                | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
 | 'engine'                          | 'system'             | 'views_with_history'   | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
 | 'engine_full'                     | 'system'             | 'tables'               | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
@@ -274,6 +276,7 @@ DB.Table: 'system'.'columns', Table: columns-table_id:1, ver:0, Engine: SystemCo
 | 'name'                            | 'system'             | 'tables_with_history'  | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
 | 'name'                            | 'system'             | 'task_history'         | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
 | 'name'                            | 'system'             | 'tasks'                | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
+| 'name'                            | 'system'             | 'temporary_tables'     | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
 | 'name'                            | 'system'             | 'user_functions'       | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
 | 'name'                            | 'system'             | 'users'                | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
 | 'name'                            | 'system'             | 'views'                | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
@@ -412,6 +415,7 @@ DB.Table: 'system'.'columns', Table: columns-table_id:1, ver:0, Engine: SystemCo
 | 'table_id'                        | 'system'             | 'streams'              | 'Nullable(UInt64)'    | 'BIGINT UNSIGNED'   | ''       | ''       | 'YES'    | ''       |
 | 'table_id'                        | 'system'             | 'tables'               | 'UInt64'              | 'BIGINT UNSIGNED'   | ''       | ''       | 'NO'     | ''       |
 | 'table_id'                        | 'system'             | 'tables_with_history'  | 'UInt64'              | 'BIGINT UNSIGNED'   | ''       | ''       | 'NO'     | ''       |
+| 'table_id'                        | 'system'             | 'temporary_tables'     | 'UInt64'              | 'BIGINT UNSIGNED'   | ''       | ''       | 'NO'     | ''       |
 | 'table_id'                        | 'system'             | 'views'                | 'UInt64'              | 'BIGINT UNSIGNED'   | ''       | ''       | 'NO'     | ''       |
 | 'table_id'                        | 'system'             | 'views_with_history'   | 'UInt64'              | 'BIGINT UNSIGNED'   | ''       | ''       | 'NO'     | ''       |
 | 'table_name'                      | 'information_schema' | 'columns'              | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |

src/query/sql/src/planner/binder/ddl/table.rs

@@ -497,7 +497,7 @@ impl Binder {
                         "Temporary table is only supported for FUSE and MEMORY engine",
                     ));
                 }
-                let _ = options.insert(OPT_KEY_TEMP_PREFIX.to_string(), self.ctx.get_session_id());
+                let _ = options.insert(OPT_KEY_TEMP_PREFIX.to_string(), self.ctx.get_session_id()?);
             }
         };
 

src/query/storages/common/session/src/temp_table.rs

@@ -45,7 +45,7 @@ use databend_storages_common_table_meta::meta::parse_storage_prefix;
 use databend_storages_common_table_meta::table::OPT_KEY_DATABASE_ID;
 use databend_storages_common_table_meta::table_id_ranges::is_temp_table_id;
 use databend_storages_common_table_meta::table_id_ranges::TEMP_TBL_ID_BEGIN;
-use log::info;
+use log::debug;
 use parking_lot::Mutex;
 
 #[derive(Debug, Clone)]
@@ -198,7 +198,7 @@ impl TempTblMgr {
         let desc = format!("{}.{}", database_name, table_name);
         let id = self.name_to_id.get(&desc);
         let Some(id) = id else {
-            info!(
+            debug!(
                 "Table {}.{} not found in temp table manager {:?}",
                 database_name, table_name, self
             );
@@ -218,6 +218,21 @@ impl TempTblMgr {
         Ok(Some(table_info))
     }
 
+    pub fn list_tables(&self) -> Result<Vec<TableInfo>> {
+        Ok(self
+            .id_to_table
+            .iter()
+            .map(|(id, t)| {
+                TableInfo::new(
+                    &t.db_name,
+                    &t.table_name,
+                    TableIdent::new(*id, 0),
+                    t.meta.clone(),
+                )
+            })
+            .collect())
+    }
+
     pub fn update_multi_table_meta(&mut self, req: Vec<UpdateTempTableReq>) {
         for r in req {
             let UpdateTempTableReq {
@@ -228,7 +243,7 @@ impl TempTblMgr {
             } = r;
             let table = self.id_to_table.get_mut(&table_id).unwrap();
             table.meta = new_table_meta;
-            table.copied_files = copied_files;
+            table.copied_files.extend(copied_files);
         }
     }
 

src/query/storages/common/table_meta/src/meta/mod.rs

@@ -38,6 +38,7 @@ pub use statistics::*;
 // currently, used by versioned readers only
 pub(crate) use testing::*;
 pub use utils::parse_storage_prefix;
+pub use utils::TEMP_TABLE_STORAGE_PREFIX;
 pub(crate) use utils::*;
 pub use versions::testify_version;
 pub use versions::SegmentInfoVersion;

src/query/storages/common/table_meta/src/meta/utils.rs

@@ -28,7 +28,7 @@ use crate::table::OPT_KEY_DATABASE_ID;
 use crate::table::OPT_KEY_STORAGE_PREFIX;
 use crate::table::OPT_KEY_TEMP_PREFIX;
 
-const TEMP_TABLE_STORAGE_PREFIX: &str = "_tmp_tbl";
+pub const TEMP_TABLE_STORAGE_PREFIX: &str = "_tmp_tbl";
 
 pub fn trim_timestamp_to_micro_second(ts: DateTime<Utc>) -> DateTime<Utc> {
     Utc.with_ymd_and_hms(

src/query/storages/fuse/Cargo.toml

@@ -39,6 +39,7 @@ databend-common-pipeline-transforms = { workspace = true }
 databend-common-sharing = { workspace = true }
 databend-common-sql = { workspace = true }
 databend-common-storage = { workspace = true }
+databend-common-users = { workspace = true }
 databend-enterprise-fail-safe = { workspace = true }
 databend-storages-common-blocks = { workspace = true }
 databend-storages-common-cache = { workspace = true }

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

@@ -0,0 +1,103 @@
+// 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::collections::HashSet;
+use std::sync::Arc;
+
+use databend_common_catalog::plan::DataSourcePlan;
+use databend_common_exception::ErrorCode;
+use databend_common_exception::Result;
+use databend_common_expression::types::StringType;
+use databend_common_expression::DataBlock;
+use databend_common_expression::FromData;
+use databend_common_expression::TableDataType;
+use databend_common_expression::TableField;
+use databend_common_expression::TableSchemaRef;
+use databend_common_expression::TableSchemaRefExt;
+use databend_common_storage::DataOperator;
+use databend_common_users::UserApiProvider;
+use databend_storages_common_table_meta::meta::TEMP_TABLE_STORAGE_PREFIX;
+use futures_util::TryStreamExt;
+use log::debug;
+use log::info;
+use uuid::Uuid;
+
+use crate::sessions::TableContext;
+use crate::table_functions::SimpleTableFunc;
+use crate::table_functions::TableArgs;
+pub struct FuseVacuumTemporaryTable;
+
+#[async_trait::async_trait]
+impl SimpleTableFunc for FuseVacuumTemporaryTable {
+    fn get_engine_name(&self) -> String {
+        "fuse_vacuum_temporary_table".to_owned()
+    }
+
+    fn table_args(&self) -> Option<TableArgs> {
+        None
+    }
+
+    fn schema(&self) -> TableSchemaRef {
+        TableSchemaRefExt::create(vec![TableField::new("result", TableDataType::String)])
+    }
+
+    async fn apply(
+        &self,
+        ctx: &Arc<dyn TableContext>,
+        _plan: &DataSourcePlan,
+    ) -> Result<Option<DataBlock>> {
+        let op = DataOperator::instance().operator();
+        let mut lister = op
+            .lister_with(TEMP_TABLE_STORAGE_PREFIX)
+            .recursive(true)
+            .await?;
+        let client_session_mgr = UserApiProvider::instance().client_session_api(&ctx.get_tenant());
+        let mut session_ids = HashSet::new();
+        while let Some(entry) = lister.try_next().await? {
+            let path = entry.path();
+            debug!("path: {}", path);
+            if let Some(session_id) = path.split('/').nth(1) {
+                if session_id.is_empty() {
+                    continue;
+                }
+                // check if session_id is a valid uuid
+                let _ = Uuid::parse_str(session_id)
+                    .map_err(|e| ErrorCode::Internal(format!("Invalid session_id: {}", e)))?;
+                debug!("session_id: {}", session_id);
+                session_ids.insert(session_id.to_string());
+            }
+        }
+        for session_id in session_ids {
+            if client_session_mgr
+                .get_client_session(&session_id)
+                .await?
+                .is_none()
+            {
+                let path = format!("{}/{}", TEMP_TABLE_STORAGE_PREFIX, session_id);
+                info!("Removing temporary table: {}", path);
+                op.remove_all(&path).await?;
+            }
+        }
+        let col: Vec<String> = vec!["Ok".to_owned()];
+
+        Ok(Some(DataBlock::new_from_columns(vec![
+            StringType::from_data(col),
+        ])))
+    }
+
+    fn create(_func_name: &str, _table_args: TableArgs) -> Result<Self>
+    where Self: Sized {
+        Ok(Self)
+    }
+}

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

@@ -22,6 +22,7 @@ mod fuse_encoding;
 mod fuse_segment;
 mod fuse_snapshot;
 mod fuse_statistic;
+mod fuse_vacuum_temporary_table;
 mod table_args;
 
 pub use clustering_information::ClusteringInformationFunc;
@@ -38,4 +39,5 @@ pub use fuse_encoding::FuseEncodingFunc;
 pub use fuse_segment::FuseSegmentFunc;
 pub use fuse_snapshot::FuseSnapshotFunc;
 pub use fuse_statistic::FuseStatisticsFunc;
+pub use fuse_vacuum_temporary_table::FuseVacuumTemporaryTable;
 pub use table_args::*;