feat: flush or compact table and region functions (GreptimeTeam#3363)

* feat: adds Requester to process table flush and compaction request * feat: admin_fn macros for administration functions * test: add query result * feat: impl flush_region, flush_table, compact_region, and flush_region functions * docs: add Arguments to admin_fn macro * chore: apply suggestion Co-authored-by: Zhenchi <[email protected]> * chore: apply suggestion Co-authored-by: Zhenchi <[email protected]> * fix: group_requests_by_peer and adds log * Update src/common/macro/src/admin_fn.rs Co-authored-by: Ruihang Xia <[email protected]> * feat: adds todo for spawan thread * feat: rebase with main --------- Co-authored-by: Zhenchi <[email protected]> Co-authored-by: Ruihang Xia <[email protected]>
evenyag · Feb 27, 2024 · 4b36c28 · 4b36c28
1 parent dbb1ce1
commit 4b36c28
Show file tree

Hide file tree

Showing 44 changed files with 1,596 additions and 395 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/client/src/region.rs b/src/client/src/region.rs
@@ -197,7 +197,7 @@ impl RegionRequester {
 
         check_response_header(header)?;
 
-        Ok(affected_rows)
+        Ok(affected_rows as _)
     }
 
     pub async fn handle(&self, request: RegionRequest) -> Result<AffectedRows> {

diff --git a/src/common/base/src/lib.rs b/src/common/base/src/lib.rs
@@ -21,6 +21,8 @@ pub mod readable_size;
 use core::any::Any;
 use std::sync::{Arc, Mutex, MutexGuard};
 
+pub type AffectedRows = usize;
+
 pub use bit_vec::BitVec;
 
 /// [`Plugins`] is a wrapper of Arc contents.

diff --git a/src/common/function/Cargo.toml b/src/common/function/Cargo.toml
@@ -12,6 +12,7 @@ api.workspace = true
 arc-swap = "1.0"
 async-trait.workspace = true
 chrono-tz = "0.6"
+common-base.workspace = true
 common-catalog.workspace = true
 common-error.workspace = true
 common-macro.workspace = true
@@ -33,6 +34,7 @@ serde_json.workspace = true
 session.workspace = true
 snafu.workspace = true
 statrs = "0.16"
+store-api.workspace = true
 table.workspace = true
 
 [dev-dependencies]

diff --git a/src/common/function/src/handlers.rs b/src/common/function/src/handlers.rs
@@ -15,12 +15,12 @@
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use common_base::AffectedRows;
 use common_meta::rpc::procedure::{MigrateRegionRequest, ProcedureStateResponse};
 use common_query::error::Result;
 use session::context::QueryContextRef;
-use table::requests::{DeleteRequest, InsertRequest};
-
-pub type AffectedRows = usize;
+use store_api::storage::RegionId;
+use table::requests::{CompactTableRequest, DeleteRequest, FlushTableRequest, InsertRequest};
 
 /// A trait for handling table mutations in `QueryEngine`.
 #[async_trait]
@@ -30,6 +30,28 @@ pub trait TableMutationHandler: Send + Sync {
 
     /// Delete rows from the table.
     async fn delete(&self, request: DeleteRequest, ctx: QueryContextRef) -> Result<AffectedRows>;
+
+    /// Trigger a flush task for table.
+    async fn flush(&self, request: FlushTableRequest, ctx: QueryContextRef)
+        -> Result<AffectedRows>;
+
+    /// Trigger a compaction task for table.
+    async fn compact(
+        &self,
+        request: CompactTableRequest,
+        ctx: QueryContextRef,
+    ) -> Result<AffectedRows>;
+
+    /// Trigger a flush task for a table region.
+    async fn flush_region(&self, region_id: RegionId, ctx: QueryContextRef)
+        -> Result<AffectedRows>;
+
+    /// Trigger a compaction task for a table region.
+    async fn compact_region(
+        &self,
+        region_id: RegionId,
+        ctx: QueryContextRef,
+    ) -> Result<AffectedRows>;
 }
 
 /// A trait for handling procedure service requests in `QueryEngine`.

diff --git a/src/common/function/src/helper.rs b/src/common/function/src/helper.rs
@@ -12,8 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use common_query::error::{InvalidInputTypeSnafu, Result};
 use common_query::prelude::{Signature, TypeSignature, Volatility};
 use datatypes::prelude::ConcreteDataType;
+use datatypes::types::cast::cast;
+use datatypes::value::ValueRef;
+use snafu::ResultExt;
 
 /// Create a function signature with oneof signatures of interleaving two arguments.
 pub fn one_of_sigs2(args1: Vec<ConcreteDataType>, args2: Vec<ConcreteDataType>) -> Signature {
@@ -27,3 +31,15 @@ pub fn one_of_sigs2(args1: Vec<ConcreteDataType>, args2: Vec<ConcreteDataType>)
 
     Signature::one_of(sigs, Volatility::Immutable)
 }
+
+/// Cast a [`ValueRef`] to u64, returns `None` if fails
+pub fn cast_u64(value: &ValueRef) -> Result<Option<u64>> {
+    cast((*value).into(), &ConcreteDataType::uint64_datatype())
+        .context(InvalidInputTypeSnafu {
+            err_msg: format!(
+                "Failed to cast input into uint64, actual type: {:#?}",
+                value.data_type(),
+            ),
+        })
+        .map(|v| v.as_u64())
+}
diff --git a/src/common/function/src/state.rs b/src/common/function/src/state.rs
@@ -32,11 +32,19 @@ impl FunctionState {
 
         use api::v1::meta::ProcedureStatus;
         use async_trait::async_trait;
+        use common_base::AffectedRows;
         use common_meta::rpc::procedure::{MigrateRegionRequest, ProcedureStateResponse};
         use common_query::error::Result;
+        use session::context::QueryContextRef;
+        use store_api::storage::RegionId;
+        use table::requests::{
+            CompactTableRequest, DeleteRequest, FlushTableRequest, InsertRequest,
+        };
 
-        use crate::handlers::ProcedureServiceHandler;
+        use crate::handlers::{ProcedureServiceHandler, TableMutationHandler};
         struct MockProcedureServiceHandler;
+        struct MockTableMutationHandler;
+        const ROWS: usize = 42;
 
         #[async_trait]
         impl ProcedureServiceHandler for MockProcedureServiceHandler {
@@ -56,8 +64,59 @@ impl FunctionState {
             }
         }
 
+        #[async_trait]
+        impl TableMutationHandler for MockTableMutationHandler {
+            async fn insert(
+                &self,
+                _request: InsertRequest,
+                _ctx: QueryContextRef,
+            ) -> Result<AffectedRows> {
+                Ok(ROWS)
+            }
+
+            async fn delete(
+                &self,
+                _request: DeleteRequest,
+                _ctx: QueryContextRef,
+            ) -> Result<AffectedRows> {
+                Ok(ROWS)
+            }
+
+            async fn flush(
+                &self,
+                _request: FlushTableRequest,
+                _ctx: QueryContextRef,
+            ) -> Result<AffectedRows> {
+                Ok(ROWS)
+            }
+
+            async fn compact(
+                &self,
+                _request: CompactTableRequest,
+                _ctx: QueryContextRef,
+            ) -> Result<AffectedRows> {
+                Ok(ROWS)
+            }
+
+            async fn flush_region(
+                &self,
+                _region_id: RegionId,
+                _ctx: QueryContextRef,
+            ) -> Result<AffectedRows> {
+                Ok(ROWS)
+            }
+
+            async fn compact_region(
+                &self,
+                _region_id: RegionId,
+                _ctx: QueryContextRef,
+            ) -> Result<AffectedRows> {
+                Ok(ROWS)
+            }
+        }
+
         Self {
-            table_mutation_handler: None,
+            table_mutation_handler: Some(Arc::new(MockTableMutationHandler)),
             procedure_service_handler: Some(Arc::new(MockProcedureServiceHandler)),
         }
     }