From 05a25ff0d864c84a74ede941d5b51446a7df2b44 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Wed, 12 Aug 2026 23:59:55 +0800 Subject: [PATCH] feat(flow): enterprise-gated incremental flow DDL hook Add a generic enterprise-gated FlowExtension seam that lets an injected plugin validate and normalize extension-owned CREATE FLOW options, then either handle the DDL task or pass a cleaned task to the ordinary procedure. Wire the same extension reference into the frontend statement executor and common-meta DDL manager. Preserve the ordinary OSS path when no extension is installed, reject unknown private options without a handler, and propagate the existing query and persistent event contexts across interception. An absent event input is represented by the neutral Unknown reason. Keep extension-owned option names and incremental-flow implementation details outside OSS. Cover normalization invariants, fail-closed errors, handled and passthrough outcomes, and present/absent event-context behavior. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- Cargo.lock | 1 + src/common/meta/src/ddl/create_flow.rs | 44 +++ src/common/meta/src/ddl_manager.rs | 403 +++++++++++++++++++++++++ src/frontend/src/instance/builder.rs | 9 + src/operator/Cargo.toml | 1 + src/operator/src/statement.rs | 12 + src/operator/src/statement/ddl.rs | 354 ++++++++++++++++++++-- 7 files changed, 797 insertions(+), 27 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ece381b7b84..115b2ee0ccc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9947,6 +9947,7 @@ dependencies = [ "common-grpc-expr", "common-macro", "common-meta", + "common-procedure", "common-query", "common-recordbatch", "common-runtime", diff --git a/src/common/meta/src/ddl/create_flow.rs b/src/common/meta/src/ddl/create_flow.rs index a2f306306d8..cec772f7b0a 100644 --- a/src/common/meta/src/ddl/create_flow.rs +++ b/src/common/meta/src/ddl/create_flow.rs @@ -839,3 +839,47 @@ impl From<&CreateFlowData> for (FlowInfoValue, Vec<(FlowPartitionId, FlowRouteVa (flow_info, flow_routes) } } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use table::table_name::TableName; + + use super::validate_flow_options; + use crate::rpc::ddl::CreateFlowTask; + + fn test_create_flow_task() -> CreateFlowTask { + CreateFlowTask { + catalog_name: "greptime".to_string(), + flow_name: "flow".to_string(), + source_table_names: vec![], + sink_table_name: TableName::new("greptime", "public", "sink"), + or_replace: false, + create_if_not_exists: false, + expire_after: None, + eval_interval_secs: None, + comment: String::new(), + sql: "select 1".to_string(), + flow_options: HashMap::new(), + eval_schedule: None, + } + } + + #[test] + fn test_validate_flow_options_rejects_ee_option() { + // A fake enterprise-private option is not known to OSS and must be + // rejected as an unknown option in every build. Enterprise builds + // intercept such options via the `FlowExtension` seam before the + // ordinary `CreateFlowProcedure` validation runs. + let mut task = test_create_flow_task(); + task.flow_options + .insert("ee_test_option".to_string(), "value".to_string()); + + let err = validate_flow_options(&task).unwrap_err(); + assert!( + err.to_string() + .contains("Unknown flow option 'ee_test_option'") + ); + } +} diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index 7ba408ca2d9..3161f433b51 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +#[cfg(feature = "enterprise")] +use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; @@ -52,6 +54,8 @@ use crate::ddl::truncate_table::TruncateTableProcedure; #[cfg(feature = "enterprise")] use crate::ddl::undrop_table::UndropTableProcedure; use crate::ddl::{DdlContext, utils}; +#[cfg(feature = "enterprise")] +use crate::error::ExternalSnafu; use crate::error::{ self, CreateRepartitionProcedureSnafu, EmptyDdlTasksSnafu, PersistRepartitionGcRequirementSnafu, ProcedureOutputSnafu, RegisterProcedureLoaderSnafu, @@ -110,6 +114,8 @@ pub struct DdlManager { repartition_procedure_factory: RepartitionProcedureFactoryRef, #[cfg(feature = "enterprise")] trigger_ddl_manager: Option, + #[cfg(feature = "enterprise")] + flow_extension: Option, } /// This trait is responsible for handling DDL tasks about triggers. e.g., @@ -141,6 +147,59 @@ pub trait TriggerDdlManager: Send + Sync { #[cfg(feature = "enterprise")] pub type TriggerDdlManagerRef = Arc; +/// A generic extension seam for the CREATE FLOW DDL path. +/// +/// Enterprise builds implement this trait to own enterprise-private flow +/// option keys (batch validation/normalization) and to decide whether a CREATE +/// FLOW task should be intercepted instead of following the ordinary OSS +/// `CreateFlowProcedure` path. OSS never reads enterprise-private option keys. +#[cfg(feature = "enterprise")] +#[async_trait::async_trait] +pub trait FlowExtension: Send + Sync { + /// Batch-validates and normalizes the enterprise-owned flow options. + /// + /// `oss_options` contains the already-canonicalized OSS-owned options; + /// `extension_options` contains every option key that OSS does not + /// recognize. The caller verifies that the returned key set exactly + /// matches the input `extension_options` key set (fail-closed on missing, + /// extra or renamed keys), which also guarantees the extension can never + /// overwrite an OSS-owned key. + fn validate_and_normalize_options( + &self, + oss_options: &HashMap, + extension_options: HashMap, + ) -> std::result::Result, BoxedError>; + + /// Attempts to intercept a CREATE FLOW task. + /// + /// - `Passthrough(task)`: the ordinary OSS `CreateFlowProcedure` path runs + /// on the returned (possibly cleaned) task. + /// - `Handled(response)`: the response is returned to the caller as-is. + /// - `Err(error)`: fail closed; the error is propagated and the task is + /// never submitted through the ordinary path. + async fn intercept_create_flow( + &self, + task: CreateFlowTask, + procedure_manager: &ProcedureManagerRef, + ddl_context: &DdlContext, + query_context: &QueryContext, + event_context: &PersistentEventContext, + ) -> std::result::Result; +} + +/// The outcome of [`FlowExtension::intercept_create_flow`]. +#[cfg(feature = "enterprise")] +pub enum CreateFlowExtensionOutcome { + /// Continue with the ordinary OSS `CreateFlowProcedure` path, using the + /// returned (possibly cleaned) task. + Passthrough(CreateFlowTask), + /// The extension fully handled the task; return the response as-is. + Handled(SubmitDdlTaskResponse), +} + +#[cfg(feature = "enterprise")] +pub type FlowExtensionRef = Arc; + macro_rules! procedure_loader_entry { ($procedure:ident) => { ( @@ -235,6 +294,8 @@ impl DdlManager { repartition_procedure_factory, #[cfg(feature = "enterprise")] trigger_ddl_manager: None, + #[cfg(feature = "enterprise")] + flow_extension: None, } } @@ -244,6 +305,12 @@ impl DdlManager { self } + #[cfg(feature = "enterprise")] + pub fn with_flow_extension(mut self, flow_extension: FlowExtensionRef) -> Self { + self.flow_extension = Some(flow_extension); + self + } + pub fn with_create_database_metadata_committer( mut self, committer: CreateDatabaseMetadataCommitterRef, @@ -1338,6 +1405,38 @@ async fn handle_create_flow_task( query_context: QueryContext, procedure_context: ProcedureContext, ) -> Result { + // Direct DDL submissions bypass the operator layer. If an enterprise + // `FlowExtension` is injected, give it the first chance to intercept the + // task: `Handled` returns the response as-is, `Passthrough` continues with + // the ordinary `CreateFlowProcedure` path on the returned (possibly + // cleaned) task, and `Err` fails closed. Without an extension the ordinary + // path runs unchanged. + #[cfg(feature = "enterprise")] + // Direct DDL submissions may carry no event context; fall back to the + // neutral default (`TriggerReason::Unknown`) instead of misclassifying + // the absent context as a manual trigger. + let event_context = procedure_context.event_context.clone().unwrap_or_default(); + + #[cfg(feature = "enterprise")] + let create_flow_task = if let Some(extension) = ddl_manager.flow_extension.as_ref() { + match extension + .intercept_create_flow( + create_flow_task, + &ddl_manager.procedure_manager, + &ddl_manager.ddl_context, + &query_context, + &event_context, + ) + .await + .context(ExternalSnafu)? + { + CreateFlowExtensionOutcome::Handled(response) => return Ok(response), + CreateFlowExtensionOutcome::Passthrough(task) => task, + } + } else { + create_flow_task + }; + let (id, output) = ddl_manager .submit_create_flow_task(create_flow_task.clone(), query_context, procedure_context) .await?; @@ -1483,9 +1582,13 @@ async fn handle_comment_on_task( #[cfg(test)] mod tests { + #[cfg(feature = "enterprise")] + use std::collections::HashMap; use std::sync::Arc; #[cfg(feature = "enterprise")] use std::sync::Mutex; + #[cfg(feature = "enterprise")] + use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; #[cfg(feature = "enterprise")] @@ -1497,6 +1600,8 @@ mod tests { use common_error::status_code::StatusCode; #[cfg(feature = "enterprise")] use common_event_recorder::{PersistentEventContext, ProcedureEventInput, TriggerReason}; + #[cfg(feature = "enterprise")] + use common_procedure::ProcedureManager; use common_procedure::local::LocalManager; use common_procedure::test_util::InMemoryPoisonStore; use common_procedure::{BoxedProcedure, ProcedureContext, ProcedureManagerRef}; @@ -1504,6 +1609,8 @@ mod tests { use table::table_name::TableName; use super::DdlManager; + #[cfg(feature = "enterprise")] + use super::{CreateFlowExtensionOutcome, FlowExtension, FlowExtensionRef}; use crate::cache_invalidator::DummyCacheInvalidator; use crate::ddl::alter_table::AlterTableProcedure; use crate::ddl::create_database::{ @@ -1513,6 +1620,8 @@ mod tests { use crate::ddl::drop_table::DropTableProcedure; use crate::ddl::flow_meta::FlowMetadataAllocator; use crate::ddl::table_meta::TableMetadataAllocator; + #[cfg(feature = "enterprise")] + use crate::ddl::test_util::flownode_handler::NaiveFlownodeHandler; use crate::ddl::truncate_table::TruncateTableProcedure; use crate::ddl::{DdlContext, NoopRegionFailureDetectorControl}; use crate::ddl_manager::{RepartitionProcedureFactory, RepartitionSource}; @@ -1526,6 +1635,8 @@ mod tests { use crate::region_registry::LeaderRegionRegistry; #[cfg(feature = "enterprise")] use crate::rpc::ddl::trigger::{CreateTriggerTask, DropTriggerTask}; + #[cfg(feature = "enterprise")] + use crate::rpc::ddl::{CreateFlowTask, SubmitDdlTaskResponse}; use crate::rpc::ddl::{CreatorGrantIntent, UndropTableTask}; #[cfg(not(feature = "enterprise"))] use crate::rpc::ddl::{DdlTask, PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest}; @@ -1533,6 +1644,8 @@ mod tests { use crate::rpc::ddl::{DdlTask, QueryContext, SubmitDdlTaskRequest}; use crate::sequence::SequenceBuilder; use crate::state_store::KvStateStore; + #[cfg(feature = "enterprise")] + use crate::test_util::MockFlownodeManager; use crate::test_util::{MockDatanodeManager, new_ddl_context}; use crate::wal_provider::WalProvider; @@ -1900,4 +2013,294 @@ mod tests { assert!(matches!(err, crate::error::Error::Unsupported { .. })); } } + + #[cfg(feature = "enterprise")] + #[derive(Clone, Copy, PartialEq)] + enum MockFlowExtensionOutcome { + Passthrough, + Handled, + Err, + } + + #[cfg(feature = "enterprise")] + struct MockFlowExtension { + calls: Arc, + last_task: Arc>>, + last_event_context: Arc>>, + outcome: MockFlowExtensionOutcome, + } + + #[cfg(feature = "enterprise")] + impl MockFlowExtension { + fn new(outcome: MockFlowExtensionOutcome) -> Self { + Self { + calls: Arc::new(AtomicUsize::new(0)), + last_task: Arc::new(Mutex::new(None)), + last_event_context: Arc::new(Mutex::new(None)), + outcome, + } + } + } + + #[cfg(feature = "enterprise")] + #[async_trait::async_trait] + impl FlowExtension for MockFlowExtension { + fn validate_and_normalize_options( + &self, + _oss_options: &HashMap, + extension_options: HashMap, + ) -> std::result::Result, BoxedError> { + Ok(extension_options) + } + + async fn intercept_create_flow( + &self, + create_flow_task: CreateFlowTask, + _procedure_manager: &ProcedureManagerRef, + _ddl_context: &DdlContext, + _query_context: &QueryContext, + event_context: &PersistentEventContext, + ) -> std::result::Result { + self.calls.fetch_add(1, Ordering::SeqCst); + *self.last_task.lock().unwrap() = Some(create_flow_task.clone()); + *self.last_event_context.lock().unwrap() = Some(event_context.clone()); + match self.outcome { + MockFlowExtensionOutcome::Passthrough => { + // Strip the fake enterprise option so the ordinary + // `CreateFlowProcedure` validation accepts the task. + let mut task = create_flow_task; + task.flow_options.remove("ee_test_option"); + Ok(CreateFlowExtensionOutcome::Passthrough(task)) + } + MockFlowExtensionOutcome::Handled => Ok(CreateFlowExtensionOutcome::Handled( + SubmitDdlTaskResponse::default(), + )), + MockFlowExtensionOutcome::Err => Err(BoxedError::new( + crate::error::UnexpectedSnafu { + err_msg: "mock extension error".to_string(), + } + .build(), + )), + } + } + } + + #[cfg(feature = "enterprise")] + fn flow_extension_test_task(or_replace: bool, create_if_not_exists: bool) -> CreateFlowTask { + CreateFlowTask { + catalog_name: "greptime".to_string(), + flow_name: "flow".to_string(), + source_table_names: vec![], + sink_table_name: TableName::new("greptime", "public", "sink"), + or_replace, + create_if_not_exists, + expire_after: None, + eval_interval_secs: None, + comment: String::new(), + sql: "select 1".to_string(), + flow_options: HashMap::new(), + eval_schedule: None, + } + } + + #[cfg(feature = "enterprise")] + async fn build_flow_extension_test_ddl_manager( + extension: Option, + ) -> DdlManager { + let state_store = Arc::new(KvStateStore::new(Arc::new(MemoryKvBackend::new()))); + let poison_manager = Arc::new(InMemoryPoisonStore::default()); + let procedure_manager = Arc::new(LocalManager::new( + Default::default(), + state_store, + poison_manager, + None, + None, + )); + // The ordinary `CreateFlowProcedure` path submits through the procedure + // manager, so it must be started before any DDL task is executed. + procedure_manager.start().await.unwrap(); + + let ddl_manager = DdlManager::new( + new_ddl_context(Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler))), + procedure_manager, + Arc::new(DummyRepartitionProcedureFactory), + ); + + match extension { + Some(extension) => ddl_manager.with_flow_extension(extension), + None => ddl_manager, + } + } + + #[cfg(feature = "enterprise")] + #[tokio::test] + async fn test_create_flow_without_extension_rejects_unknown_ee_option() { + // No extension injected: a fake enterprise-private option must be + // rejected by the ordinary `CreateFlowProcedure` validation as an + // unknown option (fail-closed), not by an enterprise-specific error. + let ddl_manager = build_flow_extension_test_ddl_manager(None).await; + let mut task = flow_extension_test_task(false, false); + task.flow_options + .insert("ee_test_option".to_string(), "value".to_string()); + + let err = ddl_manager + .submit_ddl_task( + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + SubmitDdlTaskRequest::new(DdlTask::new_create_flow(task)), + ) + .await + .unwrap_err(); + assert!( + err.output_msg() + .contains("Unknown flow option 'ee_test_option'"), + "unexpected error: {err}" + ); + } + + #[cfg(feature = "enterprise")] + #[tokio::test] + async fn test_create_flow_extension_passthrough_runs_ordinary_path() { + // The extension returns `Passthrough` with a cleaned task: the ordinary + // `CreateFlowProcedure` path must run on that task. The fake option can + // only be cleaned by the extension; if the returned task were ignored, + // the ordinary path would reject it as an unknown option. + let extension = Arc::new(MockFlowExtension::new( + MockFlowExtensionOutcome::Passthrough, + )); + let ddl_manager = build_flow_extension_test_ddl_manager(Some(extension.clone())).await; + let mut task = flow_extension_test_task(false, false); + task.flow_options + .insert("ee_test_option".to_string(), "value".to_string()); + + let resp = ddl_manager + .submit_ddl_task( + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + SubmitDdlTaskRequest::new(DdlTask::new_create_flow(task)), + ) + .await + .unwrap(); + + assert_eq!(extension.calls.load(Ordering::SeqCst), 1); + let received = extension.last_task.lock().unwrap().take().unwrap(); + assert!(!received.or_replace); + assert!(!received.create_if_not_exists); + assert_eq!( + received + .flow_options + .get("ee_test_option") + .map(String::as_str), + Some("value") + ); + // The ordinary path ran on the cleaned task and filled in a procedure id. + assert!(!resp.key.is_empty()); + } + + #[cfg(feature = "enterprise")] + #[tokio::test] + async fn test_create_flow_extension_handled_skips_ordinary_path() { + // The extension fully handles the task (`Handled`): its response is + // returned as-is and the ordinary procedure path never runs. + let extension = Arc::new(MockFlowExtension::new(MockFlowExtensionOutcome::Handled)); + let ddl_manager = build_flow_extension_test_ddl_manager(Some(extension.clone())).await; + let task = flow_extension_test_task(true, true); + + let resp = ddl_manager + .submit_ddl_task( + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + SubmitDdlTaskRequest::new(DdlTask::new_create_flow(task)), + ) + .await + .unwrap(); + + assert_eq!(extension.calls.load(Ordering::SeqCst), 1); + let received = extension.last_task.lock().unwrap().take().unwrap(); + assert!(received.or_replace); + assert!(received.create_if_not_exists); + // The mock extension returns a default response, unlike the normal path + // which would fill in a procedure id. + assert!(resp.key.is_empty()); + } + + #[cfg(feature = "enterprise")] + #[tokio::test] + async fn test_create_flow_extension_error_does_not_fall_back() { + // The extension fails (`Err`): the error must propagate and the task + // must never be submitted through the ordinary path. + let extension = Arc::new(MockFlowExtension::new(MockFlowExtensionOutcome::Err)); + let ddl_manager = build_flow_extension_test_ddl_manager(Some(extension.clone())).await; + let task = flow_extension_test_task(false, false); + + let err = ddl_manager + .submit_ddl_task( + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + SubmitDdlTaskRequest::new(DdlTask::new_create_flow(task)), + ) + .await + .unwrap_err(); + + assert_eq!(extension.calls.load(Ordering::SeqCst), 1); + assert!(matches!(err, crate::error::Error::External { .. })); + } + + #[cfg(feature = "enterprise")] + #[tokio::test] + async fn test_create_flow_extension_observes_event_context() { + // Minimal assertion that the extension receives the persistent event + // context of the DDL task. + let extension = Arc::new(MockFlowExtension::new(MockFlowExtensionOutcome::Handled)); + let ddl_manager = build_flow_extension_test_ddl_manager(Some(extension.clone())).await; + let task = flow_extension_test_task(false, false); + + ddl_manager + .submit_ddl_task( + ExecutorContext { + query_context: Some(QueryContext::default()), + event_input: Some(ProcedureEventInput::new(TriggerReason::ScheduledGc)), + ..Default::default() + }, + SubmitDdlTaskRequest::new(DdlTask::new_create_flow(task)), + ) + .await + .unwrap(); + + let received = extension.last_event_context.lock().unwrap().take().unwrap(); + assert_eq!(received.reason, TriggerReason::ScheduledGc); + } + + #[cfg(feature = "enterprise")] + #[tokio::test] + async fn test_create_flow_extension_observes_unknown_reason_without_event_input() { + // An `ExecutorContext` without `event_input` must not be misclassified + // as a manual trigger: the adapter falls back to the neutral default + // persistent event context whose reason is `TriggerReason::Unknown`. + let extension = Arc::new(MockFlowExtension::new(MockFlowExtensionOutcome::Handled)); + let ddl_manager = build_flow_extension_test_ddl_manager(Some(extension.clone())).await; + let task = flow_extension_test_task(false, false); + + ddl_manager + .submit_ddl_task( + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + SubmitDdlTaskRequest::new(DdlTask::new_create_flow(task)), + ) + .await + .unwrap(); + + let received = extension.last_event_context.lock().unwrap().take().unwrap(); + assert_eq!(received.reason, TriggerReason::Unknown); + } } diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index c49a8394beb..50a7837459b 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -26,6 +26,8 @@ use common_datasource::object_store::LocalFileAccess; use common_event_recorder::{EventRecorderImpl, EventRecorderRef}; use common_meta::cache::{LayeredCacheRegistryRef, TableRouteCacheRef}; use common_meta::cache_invalidator::{CacheInvalidatorRef, DummyCacheInvalidator}; +#[cfg(feature = "enterprise")] +use common_meta::ddl_manager::FlowExtensionRef; use common_meta::key::TableMetadataManager; use common_meta::key::flow::FlowMetadataManager; use common_meta::kv_backend::KvBackendRef; @@ -323,6 +325,13 @@ impl FrontendBuilder { statement_executor }; + #[cfg(feature = "enterprise")] + let statement_executor = if let Some(extension) = plugins.get::() { + statement_executor.with_flow_extension(extension) + } else { + statement_executor + }; + let admin_event_recorder = AdminEventRecorderHandle::default(); let statement_executor = statement_executor.with_admin_function_layer(Arc::new( AdminFunctionRecordingLayer::new(admin_event_recorder.clone()), diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml index c9258df2c0f..07af4d8c43e 100644 --- a/src/operator/Cargo.toml +++ b/src/operator/Cargo.toml @@ -87,5 +87,6 @@ tracing.workspace = true [dev-dependencies] catalog = { workspace = true, features = ["testing"] } common-meta = { workspace = true, features = ["testing"] } +common-procedure.workspace = true common-test-util.workspace = true path-slash = "0.2" diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 478f06e83f7..71a8a3259c1 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -41,6 +41,8 @@ use client::inserter::{InsertOptions, Inserter}; use common_datasource::object_store::LocalFileAccess; use common_error::ext::BoxedError; use common_meta::cache_invalidator::CacheInvalidatorRef; +#[cfg(feature = "enterprise")] +use common_meta::ddl_manager::FlowExtensionRef; use common_meta::key::flow::{FlowMetadataManager, FlowMetadataManagerRef}; use common_meta::key::schema_name::SchemaNameKey; use common_meta::key::view_info::{ViewInfoManager, ViewInfoManagerRef}; @@ -150,6 +152,8 @@ pub struct StatementExecutor { create_database_handler: Option, #[cfg(feature = "enterprise")] trigger_querier: Option, + #[cfg(feature = "enterprise")] + flow_extension: Option, } pub type StatementExecutorRef = Arc; @@ -205,6 +209,8 @@ impl StatementExecutor { create_database_handler: None, #[cfg(feature = "enterprise")] trigger_querier: None, + #[cfg(feature = "enterprise")] + flow_extension: None, } } @@ -228,6 +234,12 @@ impl StatementExecutor { self } + #[cfg(feature = "enterprise")] + pub fn with_flow_extension(mut self, extension: FlowExtensionRef) -> Self { + self.flow_extension = Some(extension); + self + } + #[cfg(feature = "testing")] pub async fn execute_stmt( &self, diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index 3c10186ca4e..98cb483d941 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -43,6 +43,8 @@ use common_meta::cache_invalidator::Context; use common_meta::ddl::create_flow::{ DEFER_ON_MISSING_SOURCE_KEY, FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY, FlowType, }; +#[cfg(feature = "enterprise")] +use common_meta::ddl_manager::FlowExtensionRef; use common_meta::instruction::CacheIdent; #[cfg(feature = "enterprise")] use common_meta::key::TableMetadataManagerRef; @@ -126,7 +128,7 @@ struct DdlSubmitOptions { timeout: Duration, } -const ALLOWED_FLOW_OPTIONS: [&str; 2] = [ +const ALLOWED_FLOW_OPTIONS: &[&str] = &[ DEFER_ON_MISSING_SOURCE_KEY, FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY, ]; @@ -188,9 +190,20 @@ fn normalize_flow_bool_option(key: &str, value: &str) -> Result { }) } +fn unknown_flow_option_error(key: &str) -> crate::error::Error { + InvalidSqlSnafu { + err_msg: format!( + "unknown flow option '{key}', supported options: {}", + supported_flow_options() + ), + } + .build() +} + fn validate_and_normalize_flow_options( options: HashMap, eval_interval: Option, + #[cfg(feature = "enterprise")] flow_extension: Option<&FlowExtensionRef>, ) -> Result> { // Reject non-positive eval_interval (zero or negative). if let Some(secs) = eval_interval @@ -202,34 +215,67 @@ fn validate_and_normalize_flow_options( .fail(); } - options - .into_iter() - .map(|(key, value)| { - if key == FlowType::FLOW_TYPE_KEY { - return InvalidSqlSnafu { - err_msg: format!("flow option '{key}' is reserved for internal use"), - } - .fail(); + let mut oss_normalized = HashMap::new(); + let mut extension_options = HashMap::new(); + + for (key, value) in options { + if key == FlowType::FLOW_TYPE_KEY { + return InvalidSqlSnafu { + err_msg: format!("flow option '{key}' is reserved for internal use"), } + .fail(); + } - let normalized_value = match key.as_str() { - DEFER_ON_MISSING_SOURCE_KEY | FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY => { - normalize_flow_bool_option(&key, &value)? - } - _ => { - return InvalidSqlSnafu { - err_msg: format!( - "unknown flow option '{key}', supported options: {}", - supported_flow_options() - ), - } - .fail(); - } - }; + match key.as_str() { + DEFER_ON_MISSING_SOURCE_KEY | FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY => { + let normalized = normalize_flow_bool_option(&key, &value)?; + oss_normalized.insert(key, normalized); + } + // Unknown keys may be owned by an injected enterprise + // `FlowExtension`: they are batch-validated/normalized through it; + // otherwise they are rejected as unknown options below. + _ => { + extension_options.insert(key, value); + } + } + } - Ok((key, normalized_value)) - }) - .collect() + if extension_options.is_empty() { + return Ok(oss_normalized); + } + + #[cfg(feature = "enterprise")] + if let Some(extension) = flow_extension { + let normalized = extension + .validate_and_normalize_options(&oss_normalized, extension_options.clone()) + .context(ExternalSnafu)?; + + // Key-set contract: the returned key set must exactly match the input + // extension key set. Missing, extra or renamed keys fail closed, and + // the extension can never overwrite an OSS-owned key. + let input_keys: HashSet<&String> = extension_options.keys().collect(); + let returned_keys: HashSet<&String> = normalized.keys().collect(); + ensure!( + input_keys == returned_keys, + error::UnexpectedSnafu { + violated: format!( + "flow extension returned mismatched option keys: input {:?}, got {:?}", + input_keys.iter().collect::>(), + returned_keys.iter().collect::>(), + ) + } + ); + + oss_normalized.extend(normalized); + return Ok(oss_normalized); + } + + // No extension (or non-enterprise build): reject the first unknown key + // with the standard unknown-option error. + let Some(unknown) = extension_options.keys().next() else { + return Ok(oss_normalized); + }; + Err(unknown_flow_option_error(unknown)) } fn determine_flow_type_for_source_state( @@ -772,8 +818,21 @@ impl StatementExecutor { .fail(); } - expr.flow_options = + #[cfg(feature = "enterprise")] + let normalized_flow_options = match self.flow_extension.as_ref() { + Some(extension) => validate_and_normalize_flow_options( + expr.flow_options, + eval_interval_secs, + Some(extension), + )?, + None => { + validate_and_normalize_flow_options(expr.flow_options, eval_interval_secs, None)? + } + }; + #[cfg(not(feature = "enterprise"))] + let normalized_flow_options = validate_and_normalize_flow_options(expr.flow_options, eval_interval_secs)?; + expr.flow_options = normalized_flow_options; let flow_type = self .determine_flow_type(&expr, query_context.clone()) @@ -2871,6 +2930,8 @@ async fn execute_undrop_table( #[cfg(test)] mod test { + #[cfg(feature = "enterprise")] + use std::sync::atomic::{AtomicUsize, Ordering}; #[cfg(feature = "enterprise")] use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -2880,6 +2941,8 @@ mod test { #[cfg(feature = "enterprise")] use common_meta::cache_invalidator::{CacheInvalidator, CacheInvalidatorRef}; #[cfg(feature = "enterprise")] + use common_meta::ddl_manager::{CreateFlowExtensionOutcome, FlowExtension, FlowExtensionRef}; + #[cfg(feature = "enterprise")] use common_meta::instruction::CacheIdent; #[cfg(feature = "enterprise")] use common_meta::key::table_route::TableRouteValue; @@ -2899,6 +2962,8 @@ mod test { use common_meta::rpc::procedure::{ MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse, }; + #[cfg(feature = "enterprise")] + use common_procedure::ProcedureManagerRef; use session::context::{QueryContext, QueryContextBuilder}; use sql::dialect::GreptimeDbDialect; use sql::parser::{ParseOptions, ParserContext}; @@ -3174,6 +3239,103 @@ mod test { assert!(matches!(binary_id, error::Error::InvalidSql { .. })); } + /// 2-arg wrapper over the real [`validate_and_normalize_flow_options`] + /// used by OSS tests. The real function takes an additional enterprise-gated + /// extension parameter that only exists in enterprise builds; this wrapper + /// hides that difference so the OSS tests compile in both builds. + fn validate_and_normalize_flow_options( + options: HashMap, + eval_interval: Option, + ) -> Result> { + #[cfg(feature = "enterprise")] + { + super::validate_and_normalize_flow_options(options, eval_interval, None) + } + #[cfg(not(feature = "enterprise"))] + { + super::validate_and_normalize_flow_options(options, eval_interval) + } + } + + /// How a generic mock [`FlowExtension`] responds to batch option + /// validation, used to exercise the extension seam. + #[cfg(feature = "enterprise")] + #[derive(Clone, Copy)] + enum MockExtensionMode { + /// Normalize `key` to `to`, pass every other extension key through. + Normalize { key: &'static str, to: &'static str }, + /// Return an empty map (drops all extension keys -> key subset). + DropKeys, + /// Add a key the caller never sent (extra key). + AddExtraKey, + /// Rename every returned key (renamed keys). + RenameKey, + /// Always fail. + Fail, + } + + #[cfg(feature = "enterprise")] + struct MockFlowExtension { + mode: MockExtensionMode, + calls: Arc, + } + + #[cfg(feature = "enterprise")] + impl MockFlowExtension { + fn new(mode: MockExtensionMode) -> Self { + Self { + mode, + calls: Arc::new(AtomicUsize::new(0)), + } + } + } + + #[cfg(feature = "enterprise")] + #[async_trait::async_trait] + impl FlowExtension for MockFlowExtension { + fn validate_and_normalize_options( + &self, + _oss_options: &HashMap, + mut extension_options: HashMap, + ) -> std::result::Result, BoxedError> { + self.calls.fetch_add(1, Ordering::SeqCst); + match self.mode { + MockExtensionMode::Normalize { key, to } => { + if let Some(value) = extension_options.get_mut(key) { + *value = to.to_string(); + } + Ok(extension_options) + } + MockExtensionMode::DropKeys => Ok(HashMap::new()), + MockExtensionMode::AddExtraKey => { + extension_options.insert("extra_key".to_string(), "value".to_string()); + Ok(extension_options) + } + MockExtensionMode::RenameKey => Ok(extension_options + .drain() + .map(|(key, value)| (format!("renamed_{key}"), value)) + .collect()), + MockExtensionMode::Fail => Err(BoxedError::new( + error::UnexpectedSnafu { + violated: "mock extension failure".to_string(), + } + .build(), + )), + } + } + + async fn intercept_create_flow( + &self, + task: common_meta::rpc::ddl::CreateFlowTask, + _procedure_manager: &ProcedureManagerRef, + _ddl_context: &common_meta::ddl::DdlContext, + _query_context: &common_meta::rpc::ddl::QueryContext, + _event_context: &common_event_recorder::PersistentEventContext, + ) -> std::result::Result { + Ok(CreateFlowExtensionOutcome::Passthrough(task)) + } + } + #[test] fn test_validate_and_normalize_flow_options_empty() { assert!( @@ -3253,6 +3415,144 @@ mod test { ); } + #[test] + fn test_flow_option_enterprise_private_rejected_without_extension() { + // A fake enterprise-private option is not in the OSS allowlist and + // must be rejected as unknown in every build when no extension owns it. + let err = validate_and_normalize_flow_options( + HashMap::from([("ee_test_option".to_string(), "true".to_string())]), + None, + ) + .unwrap_err(); + + assert!( + err.to_string() + .contains("unknown flow option 'ee_test_option'") + ); + } + + #[cfg(feature = "enterprise")] + #[test] + fn test_flow_option_extension_normalizes_batch() { + // Batch validation through the extension: multiple unknown keys are + // normalized in one call and the returned key set matches the input. + let extension = Arc::new(MockFlowExtension::new(MockExtensionMode::Normalize { + key: "ee_test_option", + to: "normalized", + })); + let extension_ref: FlowExtensionRef = extension.clone(); + + let normalized = super::validate_and_normalize_flow_options( + HashMap::from([ + ("ee_test_option".to_string(), "raw value".to_string()), + ("ee_other_option".to_string(), "v".to_string()), + ]), + None, + Some(&extension_ref), + ) + .unwrap(); + + assert_eq!( + normalized, + HashMap::from([ + ("ee_test_option".to_string(), "normalized".to_string()), + ("ee_other_option".to_string(), "v".to_string()), + ]) + ); + assert_eq!(extension.calls.load(Ordering::SeqCst), 1); + } + + #[cfg(feature = "enterprise")] + #[test] + fn test_flow_option_extension_key_subset_rejected() { + // The extension drops keys (`DropKeys`): the returned key set is a + // strict subset of the input and must fail closed. + let extension: FlowExtensionRef = + Arc::new(MockFlowExtension::new(MockExtensionMode::DropKeys)); + + let err = super::validate_and_normalize_flow_options( + HashMap::from([("ee_test_option".to_string(), "value".to_string())]), + None, + Some(&extension), + ) + .unwrap_err(); + + assert!( + matches!(err, error::Error::Unexpected { .. }), + "unexpected error: {err}" + ); + } + + #[cfg(feature = "enterprise")] + #[test] + fn test_flow_option_extension_extra_or_renamed_keys_rejected() { + // The extension adds a key it never received (`AddExtraKey`). + let extension: FlowExtensionRef = + Arc::new(MockFlowExtension::new(MockExtensionMode::AddExtraKey)); + + let err = super::validate_and_normalize_flow_options( + HashMap::from([("ee_test_option".to_string(), "value".to_string())]), + None, + Some(&extension), + ) + .unwrap_err(); + assert!(matches!(err, error::Error::Unexpected { .. })); + + // The extension renames the key it received (`RenameKey`). + let extension: FlowExtensionRef = + Arc::new(MockFlowExtension::new(MockExtensionMode::RenameKey)); + + let err = super::validate_and_normalize_flow_options( + HashMap::from([("ee_test_option".to_string(), "value".to_string())]), + None, + Some(&extension), + ) + .unwrap_err(); + assert!(matches!(err, error::Error::Unexpected { .. })); + } + + #[cfg(feature = "enterprise")] + #[test] + fn test_flow_option_extension_error_propagates() { + let extension: FlowExtensionRef = Arc::new(MockFlowExtension::new(MockExtensionMode::Fail)); + + let err = super::validate_and_normalize_flow_options( + HashMap::from([("ee_test_option".to_string(), "value".to_string())]), + None, + Some(&extension), + ) + .unwrap_err(); + + assert!( + matches!(err, error::Error::External { .. }), + "unexpected error: {err}" + ); + } + + #[cfg(feature = "enterprise")] + #[test] + fn test_flow_option_oss_keys_do_not_call_extension() { + // All keys are OSS-owned: the extension must not be consulted at all. + let extension = Arc::new(MockFlowExtension::new(MockExtensionMode::Normalize { + key: "ee_test_option", + to: "normalized", + })); + let extension_ref: FlowExtensionRef = extension.clone(); + + let normalized = super::validate_and_normalize_flow_options( + HashMap::from([(DEFER_ON_MISSING_SOURCE_KEY.to_string(), "TRUE".to_string())]), + None, + Some(&extension_ref), + ) + .unwrap(); + + assert_eq!(extension.calls.load(Ordering::SeqCst), 0); + assert_eq!( + normalized, + HashMap::from([(DEFER_ON_MISSING_SOURCE_KEY.to_string(), "true".to_string())]) + ); + } + #[test] fn test_validate_and_normalize_flow_options_rejects_redacted_invalid_input() { let sql = r"