rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-13 15:32:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add InsertLimitInterceptor hook to Inserter

Browse Source 
Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>
This commit is contained in:
Lei, HUANG
2026-08-18 18:20:36 +08:00
parent 782c495587
commit b3a61c1006
2 changed files with 205 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/frontend/src/instance/builder.rs
+11 -8
View File
@@ -34,7 +34,7 @@ use common_meta::procedure_executor::ProcedureExecutorRef;
use dashmap::DashMap;
use operator::delete::Deleter;
use operator::flow::FlowServiceOperator;
use operator::insert::Inserter;
use operator::insert::{InsertLimitInterceptorRef, Inserter};
use operator::procedure::ProcedureServiceOperator;
use operator::request::Requester;
#[cfg(feature = "enterprise")]
@@ -222,13 +222,16 @@ impl FrontendBuilder {
name: TABLE_FLOWNODE_SET_CACHE_NAME,
})?;
let inserter = Arc::new(Inserter::new(
self.catalog_manager.clone(),
partition_manager.clone(),
node_manager.clone(),
table_flownode_cache,
self.options.auto_create_table,
));
let inserter = Arc::new(
Inserter::new(
self.catalog_manager.clone(),
partition_manager.clone(),
node_manager.clone(),
table_flownode_cache,
self.options.auto_create_table,
)
.with_insert_limit_interceptor(plugins.get::<InsertLimitInterceptorRef>()),
);
let deleter = Arc::new(Deleter::new(
self.catalog_manager.clone(),
partition_manager.clone(),
src/operator/src/insert.rs
+194 -2
View File
@@ -33,6 +33,7 @@ use common_catalog::consts::{
TRACE_TABLE_NAME_SESSION_KEY, default_engine, is_ddl_reserved_table,
trace_operations_table_name, trace_services_table_name,
};
use common_error::ext::BoxedError;
use common_grpc_expr::util::ColumnExpr;
use common_meta::cache::TableFlownodeSetCacheRef;
use common_meta::node_manager::{AffectedRows, NodeManagerRef};
@@ -82,6 +83,23 @@ use crate::req_convert::insert::{
};
use crate::statement::StatementExecutor;
/// Intercepts row inserts right before they are dispatched to datanodes,
/// e.g. to enforce database-level ingest rate limits.
/// Registered via `Plugins`; `None` means no interception (OSS default).
#[async_trait::async_trait]
pub trait InsertLimitInterceptor: Send + Sync {
/// Checks whether `rows` rows may be ingested into `catalog.schema`.
/// Returning `Err` rejects the whole request.
async fn check_ingest(
&self,
catalog: &str,
schema: &str,
rows: u64,
) -> std::result::Result<(), BoxedError>;
}
pub type InsertLimitInterceptorRef = Arc<dyn InsertLimitInterceptor>;
pub struct Inserter {
catalog_manager: CatalogManagerRef,
pub(crate) partition_manager: PartitionRuleManagerRef,
@@ -91,6 +109,7 @@ pub struct Inserter {
/// When `false`, missing tables are never auto-created regardless of the
/// per-request `auto_create_table` hint. When `true`, the hint still applies.
auto_create_table: bool,
insert_limit_interceptor: Option<InsertLimitInterceptorRef>,
}
pub type InserterRef = Arc<Inserter>;
@@ -160,9 +179,20 @@ impl Inserter {
node_manager,
table_flownode_set_cache,
auto_create_table,
insert_limit_interceptor: None,
}
}
/// Sets the [InsertLimitInterceptor] invoked before insert requests are
/// dispatched to datanodes. `None` (the default) disables interception.
pub fn with_insert_limit_interceptor(
mut self,
interceptor: Option<InsertLimitInterceptorRef>,
) -> Self {
self.insert_limit_interceptor = interceptor;
self
}
pub async fn handle_column_inserts(
&self,
requests: InsertRequests,
@@ -385,6 +415,15 @@ impl Inserter {
}
}
/// Counts the total number of rows in a [RegionInsertRequests].
fn count_region_insert_rows(requests: &RegionInsertRequests) -> u64 {
requests
.requests
.iter()
.map(|r| r.rows.as_ref().map(|rows| rows.rows.len()).unwrap_or(0) as u64)
.sum()
}
impl Inserter {
async fn do_request(
&self,
@@ -395,6 +434,18 @@ impl Inserter {
// Fill impure default values in the request
let requests = fill_reqs_with_impure_default(table_infos, requests)?;
if let Some(interceptor) = &self.insert_limit_interceptor {
let rows = count_region_insert_rows(&requests.normal_requests)
+ count_region_insert_rows(&requests.instant_requests);
interceptor
.check_ingest(ctx.current_catalog(), &ctx.current_schema(), rows)
.await
.map_err(|source| crate::error::Error::External {
source,
location: snafu::Location::default(),
})?;
}
let write_cost = write_meter!(
ctx.current_catalog(),
ctx.current_schema(),
@@ -1494,12 +1545,14 @@ impl FlowMirrorTask {
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::sync::{Arc, Mutex};
use api::helper::ColumnDataTypeWrapper;
use api::v1::helper::{field_column_schema, time_index_column_schema};
use api::v1::{RowInsertRequest, Rows, Value};
use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
use common_error::ext::{BoxedError, ErrorExt, PlainError};
use common_error::status_code::StatusCode;
use common_meta::cache::new_table_flownode_set_cache;
use common_meta::ddl::test_util::datanode_handler::NaiveDatanodeHandler;
use common_meta::test_util::MockDatanodeManager;
@@ -1514,7 +1567,9 @@ mod tests {
use table::metadata::{TableInfoBuilder, TableMetaBuilder, TableType};
use super::*;
use crate::tests::{create_partition_rule_manager, prepare_mocked_backend};
use crate::tests::{
create_partition_rule_manager, new_test_table_info, prepare_mocked_backend,
};
fn make_table_ref_with_schema(
ts_name: &str,
@@ -1702,6 +1757,143 @@ mod tests {
assert!(table_is_native_histogram(&table));
}
struct RejectAllInterceptor;
#[async_trait::async_trait]
impl InsertLimitInterceptor for RejectAllInterceptor {
async fn check_ingest(
&self,
_catalog: &str,
_schema: &str,
_rows: u64,
) -> std::result::Result<(), BoxedError> {
Err(BoxedError::new(PlainError::new(
"rate limited for test".to_string(),
StatusCode::RateLimited,
)))
}
}
#[derive(Default)]
struct RecordingInterceptor {
calls: Mutex<Vec<(String, String, u64)>>,
}
#[async_trait::async_trait]
impl InsertLimitInterceptor for RecordingInterceptor {
async fn check_ingest(
&self,
catalog: &str,
schema: &str,
rows: u64,
) -> std::result::Result<(), BoxedError> {
self.calls
.lock()
.unwrap()
.push((catalog.to_string(), schema.to_string(), rows));
Ok(())
}
}
async fn new_inserter_with_interceptor(
interceptor: Option<InsertLimitInterceptorRef>,
) -> Inserter {
let kv_backend = prepare_mocked_backend().await;
Inserter::new(
catalog::memory::MemoryCatalogManager::new(),
create_partition_rule_manager(kv_backend.clone()).await,
Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)),
Arc::new(new_table_flownode_set_cache(
String::new(),
Cache::new(100),
kv_backend.clone(),
)),
true,
)
.with_insert_limit_interceptor(interceptor)
}
fn new_region_insert_requests(
table_id: u32,
region_number: u32,
num_normal_rows: usize,
num_instant_rows: usize,
) -> InstantAndNormalInsertRequests {
let build = |num_rows: usize| RegionInsertRequests {
requests: vec![RegionInsertRequest {
region_id: RegionId::new(table_id, region_number).as_u64(),
rows: Some(Rows {
schema: vec![],
rows: vec![api::v1::Row { values: vec![] }; num_rows],
}),
partition_expr_version: None,
}],
};
InstantAndNormalInsertRequests {
normal_requests: build(num_normal_rows),
instant_requests: build(num_instant_rows),
}
}
#[tokio::test]
async fn test_insert_rejected_by_limit_interceptor() {
let inserter = new_inserter_with_interceptor(Some(Arc::new(RejectAllInterceptor))).await;
let table_id = 1;
let table_infos = HashMap::from_iter([(
table_id,
Arc::new(new_test_table_info(
table_id,
"test_table",
vec![0].into_iter(),
)),
)]);
let requests = new_region_insert_requests(table_id, 0, 3, 0);
let ctx = Arc::new(QueryContext::with(
DEFAULT_CATALOG_NAME,
DEFAULT_SCHEMA_NAME,
));
let err = inserter
.do_request(requests, &table_infos, &ctx)
.await
.unwrap_err();
assert_eq!(err.status_code(), StatusCode::RateLimited);
}
#[tokio::test]
async fn test_insert_limit_interceptor_receives_total_rows() {
let interceptor = Arc::new(RecordingInterceptor::default());
let inserter = new_inserter_with_interceptor(Some(interceptor.clone())).await;
let table_id = 1;
let table_infos = HashMap::from_iter([(
table_id,
Arc::new(new_test_table_info(
table_id,
"test_table",
vec![0].into_iter(),
)),
)]);
let requests = new_region_insert_requests(table_id, 0, 3, 2);
let ctx = Arc::new(QueryContext::with(
DEFAULT_CATALOG_NAME,
DEFAULT_SCHEMA_NAME,
));
// The interceptor passes; the request then fails downstream because the
// mocked backend has no table route. Only the interception is asserted.
let _ = inserter.do_request(requests, &table_infos, &ctx).await;
let calls = interceptor.calls.lock().unwrap();
assert_eq!(
calls.as_slice(),
&[(
DEFAULT_CATALOG_NAME.to_string(),
DEFAULT_SCHEMA_NAME.to_string(),
5
)]
);
}
#[test]
fn test_last_non_null_create_options_preserve_default_without_append_mode() {
let ctx = Arc::new(QueryContext::with(