From b3a61c1006e988872ca3303c2ea8e130cd780d8a Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" Date: Tue, 18 Aug 2026 18:20:36 +0800 Subject: [PATCH] feat: add `InsertLimitInterceptor` hook to `Inserter` Signed-off-by: Lei, HUANG --- src/frontend/src/instance/builder.rs | 19 +-- src/operator/src/insert.rs | 196 ++++++++++++++++++++++++++- 2 files changed, 205 insertions(+), 10 deletions(-) diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index c49a8394beb..5ae44a6ccfe 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -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::()), + ); let deleter = Arc::new(Deleter::new( self.catalog_manager.clone(), partition_manager.clone(), diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index fabc1b3778b..38912f5b111 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -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; + 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, } pub type InserterRef = Arc; @@ -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, + ) -> 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>, + } + + #[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, + ) -> 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(