diff --git a/Cargo.lock b/Cargo.lock index dcd2345140..97b2fc76a7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6039,7 +6039,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=e7be20ff855b7522efe235383d82238d56e676c3#e7be20ff855b7522efe235383d82238d56e676c3" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=32f467fa2ba2b3588a58381a24af83de09fbb00a#32f467fa2ba2b3588a58381a24af83de09fbb00a" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", @@ -8312,6 +8312,7 @@ dependencies = [ "async-trait", "common-base", "common-error", + "common-event-recorder", "common-grpc", "common-macro", "common-meta", diff --git a/Cargo.toml b/Cargo.toml index 0e6da7888d..63dbc76f38 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -159,7 +159,7 @@ fs2 = "0.4" fst = "0.4.7" futures = "0.3" futures-util = "0.3" -greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "e7be20ff855b7522efe235383d82238d56e676c3" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "32f467fa2ba2b3588a58381a24af83de09fbb00a" } hex = "0.4" http = "1" humantime = "2.1" diff --git a/src/common/event-recorder/src/context.rs b/src/common/event-recorder/src/context.rs index 601276967e..1220c25419 100644 --- a/src/common/event-recorder/src/context.rs +++ b/src/common/event-recorder/src/context.rs @@ -12,9 +12,51 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; + +use api::v1::meta::ProcedureEventContext as PbProcedureEventContext; use serde::{Deserialize, Serialize}; use strum::{AsRefStr, EnumString}; +/// Event metadata supplied when a procedure is submitted. +/// +/// Protocol is deliberately absent. The trusted submission adapter derives it +/// from the typed query channel before serializing the protobuf request. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct ProcedureEventInput { + pub reason: TriggerReason, + pub extensions: HashMap, +} + +impl ProcedureEventInput { + /// Creates procedure event input with a trigger reason and no extensions. + pub fn new(reason: TriggerReason) -> Self { + Self { + reason, + extensions: Default::default(), + } + } +} + +impl From<&ProcedureEventInput> for PbProcedureEventContext { + fn from(input: &ProcedureEventInput) -> Self { + Self { + reason: input.reason.as_ref().to_string(), + protocol: String::new(), + extensions: input.extensions.clone(), + } + } +} + +impl From for ProcedureEventInput { + fn from(context: PbProcedureEventContext) -> Self { + Self { + reason: TriggerReason::from_extension(&context.reason), + extensions: context.extensions, + } + } +} + /// Stable context recorded for a procedure event. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct PersistentEventContext { @@ -42,12 +84,40 @@ impl PersistentEventContext { } } +impl From<(ProcedureEventInput, Option)> for PersistentEventContext { + fn from((input, protocol): (ProcedureEventInput, Option)) -> Self { + Self { + reason: input.reason, + protocol, + extensions: input + .extensions + .into_iter() + .map(|(key, value)| (key, serde_json::Value::String(value))) + .collect(), + } + } +} + impl Default for PersistentEventContext { fn default() -> Self { Self::new(TriggerReason::default()) } } +impl From for PersistentEventContext { + fn from(context: PbProcedureEventContext) -> Self { + Self { + reason: TriggerReason::from_extension(&context.reason), + protocol: (!context.protocol.is_empty()).then_some(context.protocol), + extensions: context + .extensions + .into_iter() + .map(|(key, value)| (key, serde_json::Value::String(value))) + .collect(), + } + } +} + /// Stable classification of a procedure trigger. #[derive( Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, AsRefStr, EnumString, @@ -96,6 +166,55 @@ mod tests { ); } + #[test] + fn test_event_context_from_protobuf() { + let protobuf = PbProcedureEventContext { + reason: "auto_create".to_string(), + protocol: "postgres".to_string(), + extensions: HashMap::from([ + ("source".to_string(), "sql".to_string()), + ("tenant".to_string(), "a".to_string()), + ]), + }; + assert_eq!( + PersistentEventContext::from(protobuf), + PersistentEventContext { + reason: TriggerReason::AutoCreate, + protocol: Some("postgres".to_string()), + extensions: serde_json::Map::from_iter([ + ("source".to_string(), json!("sql")), + ("tenant".to_string(), json!("a")), + ]), + } + ); + } + + #[test] + fn test_event_input_protobuf_has_no_protocol() { + let input = ProcedureEventInput { + reason: TriggerReason::AutoCreate, + extensions: HashMap::from([("source".to_string(), "sql".to_string())]), + }; + + assert_eq!( + PbProcedureEventContext::from(&input), + PbProcedureEventContext { + reason: "auto_create".to_string(), + protocol: String::new(), + extensions: input.extensions.clone(), + } + ); + + assert_eq!( + ProcedureEventInput::from(PbProcedureEventContext { + reason: "auto_create".to_string(), + protocol: "untrusted".to_string(), + extensions: input.extensions.clone(), + }), + input + ); + } + #[test] fn test_trigger_reason_deserializes_unknown_value() { let reason: TriggerReason = serde_json::from_str("\"future_reason\"").unwrap(); diff --git a/src/common/event-recorder/src/lib.rs b/src/common/event-recorder/src/lib.rs index 5b9cb4a7b9..7ee7172222 100644 --- a/src/common/event-recorder/src/lib.rs +++ b/src/common/event-recorder/src/lib.rs @@ -22,5 +22,5 @@ pub mod recorder; #[cfg(any(test, feature = "testing"))] pub mod testing; -pub use context::{PersistentEventContext, TriggerReason}; +pub use context::{PersistentEventContext, ProcedureEventInput, TriggerReason}; pub use recorder::*; diff --git a/src/common/function/src/admin/gc.rs b/src/common/function/src/admin/gc.rs index e9ba33a050..0090140a46 100644 --- a/src/common/function/src/admin/gc.rs +++ b/src/common/function/src/admin/gc.rs @@ -38,17 +38,20 @@ const DEFAULT_FULL_FILE_LISTING: bool = false; )] pub(crate) async fn gc_regions( procedure_service_handler: &ProcedureServiceHandlerRef, - _ctx: &QueryContextRef, + query_ctx: &QueryContextRef, params: &[ValueRef<'_>], ) -> Result { let (region_ids, full_file_listing) = parse_gc_regions_params(params)?; let resp = procedure_service_handler - .gc_regions(GcRegionsRequest { - region_ids, - full_file_listing, - timeout: None, - }) + .gc_regions( + query_ctx.clone(), + GcRegionsRequest { + region_ids, + full_file_listing, + timeout: None, + }, + ) .await?; Ok(Value::from(resp.processed_regions)) @@ -69,13 +72,16 @@ pub(crate) async fn gc_table( parse_gc_table_params(params, query_ctx)?; let resp = procedure_service_handler - .gc_table(GcTableRequest { - catalog_name, - schema_name, - table_name, - full_file_listing, - timeout: None, - }) + .gc_table( + query_ctx.clone(), + GcTableRequest { + catalog_name, + schema_name, + table_name, + full_file_listing, + timeout: None, + }, + ) .await?; Ok(Value::from(resp.processed_regions)) @@ -268,13 +274,17 @@ mod tests { impl ProcedureServiceHandler for MockProcedureServiceHandler { async fn purge_table( &self, - _table_name: table::table_name::TableName, _query_ctx: QueryContextRef, + _table_name: table::table_name::TableName, ) -> Result<()> { unreachable!() } - async fn migrate_region(&self, _request: MigrateRegionRequest) -> Result> { + async fn migrate_region( + &self, + _query_ctx: QueryContextRef, + _request: MigrateRegionRequest, + ) -> Result> { unreachable!() } @@ -297,12 +307,20 @@ mod tests { unreachable!() } - async fn gc_regions(&self, request: GcRegionsRequest) -> Result { + async fn gc_regions( + &self, + _query_ctx: QueryContextRef, + request: GcRegionsRequest, + ) -> Result { *self.gc_regions_request.lock().unwrap() = Some(request); Ok(GcResponse::default()) } - async fn gc_table(&self, request: GcTableRequest) -> Result { + async fn gc_table( + &self, + _query_ctx: QueryContextRef, + request: GcTableRequest, + ) -> Result { *self.gc_table_request.lock().unwrap() = Some(request); Ok(GcResponse::default()) } diff --git a/src/common/function/src/admin/migrate_region.rs b/src/common/function/src/admin/migrate_region.rs index d958b3cce9..5d258e1494 100644 --- a/src/common/function/src/admin/migrate_region.rs +++ b/src/common/function/src/admin/migrate_region.rs @@ -47,7 +47,7 @@ const DEFAULT_TIMEOUT_SECS: u64 = 300; )] pub(crate) async fn migrate_region( procedure_service_handler: &ProcedureServiceHandlerRef, - _ctx: &QueryContextRef, + query_ctx: &QueryContextRef, params: &[ValueRef<'_>], ) -> Result { let (region_id, from_peer, to_peer, timeout) = match params.len() { @@ -82,12 +82,15 @@ pub(crate) async fn migrate_region( match (region_id, from_peer, to_peer, timeout) { (Some(region_id), Some(from_peer), Some(to_peer), Some(timeout)) => { let pid = procedure_service_handler - .migrate_region(MigrateRegionRequest { - region_id, - from_peer, - to_peer, - timeout: Duration::from_secs(timeout), - }) + .migrate_region( + query_ctx.clone(), + MigrateRegionRequest { + region_id, + from_peer, + to_peer, + timeout: Duration::from_secs(timeout), + }, + ) .await?; match pid { diff --git a/src/common/function/src/admin/purge_table.rs b/src/common/function/src/admin/purge_table.rs index 59613e155f..7e5fcd8437 100644 --- a/src/common/function/src/admin/purge_table.rs +++ b/src/common/function/src/admin/purge_table.rs @@ -64,8 +64,8 @@ pub(crate) async fn purge_table( procedure_service_handler .purge_table( - TableName::new(catalog_name, schema_name, table_name), query_ctx.clone(), + TableName::new(catalog_name, schema_name, table_name), ) .await?; Ok(Value::from(0_u64)) @@ -131,8 +131,8 @@ mod tests { impl ProcedureServiceHandler for RecordingHandler { async fn purge_table( &self, - table_name: TableName, query_ctx: QueryContextRef, + table_name: TableName, ) -> Result<()> { if self.fail { return InvalidFuncArgsSnafu { @@ -144,7 +144,11 @@ mod tests { Ok(()) } - async fn migrate_region(&self, _: MigrateRegionRequest) -> Result> { + async fn migrate_region( + &self, + _: QueryContextRef, + _: MigrateRegionRequest, + ) -> Result> { unreachable!() } async fn reconcile(&self, _: ReconcileRequest) -> Result> { @@ -159,10 +163,10 @@ mod tests { fn catalog_manager(&self) -> &CatalogManagerRef { unreachable!() } - async fn gc_regions(&self, _: GcRegionsRequest) -> Result { + async fn gc_regions(&self, _: QueryContextRef, _: GcRegionsRequest) -> Result { unreachable!() } - async fn gc_table(&self, _: GcTableRequest) -> Result { + async fn gc_table(&self, _: QueryContextRef, _: GcTableRequest) -> Result { unreachable!() } } diff --git a/src/common/function/src/handlers.rs b/src/common/function/src/handlers.rs index 24ce115a0d..6ee8064d6c 100644 --- a/src/common/function/src/handlers.rs +++ b/src/common/function/src/handlers.rs @@ -89,10 +89,14 @@ pub trait TableMutationHandler: Send + Sync { #[async_trait] pub trait ProcedureServiceHandler: Send + Sync { /// Permanently purge a dropped table. - async fn purge_table(&self, table_name: TableName, query_ctx: QueryContextRef) -> Result<()>; + async fn purge_table(&self, query_ctx: QueryContextRef, table_name: TableName) -> Result<()>; /// Migrate a region from source peer to target peer, returns the procedure id if success. - async fn migrate_region(&self, request: MigrateRegionRequest) -> Result>; + async fn migrate_region( + &self, + query_ctx: QueryContextRef, + request: MigrateRegionRequest, + ) -> Result>; /// Reconcile a table, database or catalog, returns the procedure id if success. async fn reconcile(&self, request: ReconcileRequest) -> Result>; @@ -107,10 +111,18 @@ pub trait ProcedureServiceHandler: Send + Sync { fn catalog_manager(&self) -> &CatalogManagerRef; /// Manually trigger GC for specific regions. - async fn gc_regions(&self, request: MetaGcRegionsRequest) -> Result; + async fn gc_regions( + &self, + query_ctx: QueryContextRef, + request: MetaGcRegionsRequest, + ) -> Result; /// Manually trigger GC for a table. - async fn gc_table(&self, request: MetaGcTableRequest) -> Result; + async fn gc_table( + &self, + query_ctx: QueryContextRef, + request: MetaGcTableRequest, + ) -> Result; } /// This flow service handler is only use for flush flow for now. diff --git a/src/common/function/src/state.rs b/src/common/function/src/state.rs index b1802b6257..2b9f3a517a 100644 --- a/src/common/function/src/state.rs +++ b/src/common/function/src/state.rs @@ -60,14 +60,15 @@ impl FunctionState { impl ProcedureServiceHandler for MockProcedureServiceHandler { async fn purge_table( &self, - _table_name: table::table_name::TableName, _query_ctx: QueryContextRef, + _table_name: table::table_name::TableName, ) -> Result<()> { Ok(()) } async fn migrate_region( &self, + _ctx: QueryContextRef, _request: MigrateRegionRequest, ) -> Result> { Ok(Some("test_pid".to_string())) @@ -92,7 +93,11 @@ impl FunctionState { Ok(()) } - async fn gc_regions(&self, _request: GcRegionsRequest) -> Result { + async fn gc_regions( + &self, + _context: QueryContextRef, + _request: GcRegionsRequest, + ) -> Result { Ok(GcResponse { processed_regions: 1, need_retry_regions: vec![], @@ -101,7 +106,11 @@ impl FunctionState { }) } - async fn gc_table(&self, _request: GcTableRequest) -> Result { + async fn gc_table( + &self, + _context: QueryContextRef, + _request: GcTableRequest, + ) -> Result { Ok(GcResponse { processed_regions: 1, need_retry_regions: vec![], diff --git a/src/common/meta/src/ddl/tests/event/flow.rs b/src/common/meta/src/ddl/tests/event/flow.rs index aae3613a08..cec0162adf 100644 --- a/src/common/meta/src/ddl/tests/event/flow.rs +++ b/src/common/meta/src/ddl/tests/event/flow.rs @@ -19,9 +19,8 @@ use api::v1::{ColumnSchema, Row, Value}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_event_recorder::Event; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, - PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, - jsonb_value, + CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, jsonb_value, + procedure_event_column_schemas, }; use common_event_recorder::testing::assert_event_contract; use common_procedure::{EventTrigger, ProcedureEvent, ProcedureId, ProcedureState}; @@ -240,12 +239,7 @@ fn assert_procedure_event_contract( trigger: &str, locator: FlowEventLocator<'_>, ) { - let mut schema = vec![ - PROCEDURE_ID_COLUMN.column_schema(), - PROCEDURE_STATE_COLUMN.column_schema(), - PROCEDURE_ERROR_COLUMN.column_schema(), - PROCEDURE_TRIGGER_COLUMN.column_schema(), - ]; + let mut schema = procedure_event_column_schemas(); schema.extend(flow_schema()); schema.push(EVENT_CONTEXT_COLUMN.column_schema()); assert_event_contract( diff --git a/src/common/meta/src/ddl/tests/event/view.rs b/src/common/meta/src/ddl/tests/event/view.rs index 5cb5a6bb7d..6ebb62cda8 100644 --- a/src/common/meta/src/ddl/tests/event/view.rs +++ b/src/common/meta/src/ddl/tests/event/view.rs @@ -17,9 +17,8 @@ use std::sync::Arc; use api::v1::value::ValueData; use api::v1::{ColumnSchema, Row, Value}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, - PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, - VIEW_NAME_COLUMN, jsonb_value, + CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, + VIEW_NAME_COLUMN, jsonb_value, procedure_event_column_schemas, }; use common_event_recorder::testing::assert_event_contract; use common_event_recorder::{Event, EventTypeFilter}; @@ -334,12 +333,7 @@ fn assert_procedure_event_contract( trigger: &str, locator: ViewEventLocator<'_>, ) { - let mut schema = vec![ - PROCEDURE_ID_COLUMN.column_schema(), - PROCEDURE_STATE_COLUMN.column_schema(), - PROCEDURE_ERROR_COLUMN.column_schema(), - PROCEDURE_TRIGGER_COLUMN.column_schema(), - ]; + let mut schema = procedure_event_column_schemas(); schema.extend(view_schema()); schema.push(EVENT_CONTEXT_COLUMN.column_schema()); diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index 2eed4628ac..7ba408ca2d 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -19,9 +19,12 @@ use api::v1::Repartition; use api::v1::alter_table_expr::Kind; use api::v1::repartition::Source as PbRepartitionSource; use common_error::ext::BoxedError; +use common_event_recorder::PersistentEventContext; +#[cfg(feature = "enterprise")] +use common_event_recorder::TriggerReason; use common_procedure::{ - BoxedProcedure, BoxedProcedureLoader, Output, ProcedureId, ProcedureManagerRef, - ProcedureWithId, watcher, + BoxedProcedure, BoxedProcedureLoader, Output, ProcedureContext, ProcedureId, + ProcedureManagerRef, ProcedureWithId, watcher, }; use common_telemetry::tracing_context::{FutureExt, TracingContext}; use common_telemetry::{debug, info, tracing}; @@ -54,7 +57,7 @@ use crate::error::{ PersistRepartitionGcRequirementSnafu, ProcedureOutputSnafu, RegisterProcedureLoaderSnafu, RegisterRepartitionProcedureLoaderSnafu, Result, SubmitProcedureSnafu, TableInfoNotFoundSnafu, TableNotFoundSnafu, TableRouteNotFoundSnafu, UnexpectedLogicalRouteTableSnafu, - WaitProcedureSnafu, + UnsupportedSnafu, WaitProcedureSnafu, }; use crate::key::table_info::TableInfoValue; use crate::key::table_name::TableNameKey; @@ -70,16 +73,14 @@ use crate::rpc::ddl::DdlTask::{ DropTable, DropView, PurgeDroppedTable, TruncateTable, UndropTable, }; #[cfg(feature = "enterprise")] -use crate::rpc::ddl::TriggerReason; -#[cfg(feature = "enterprise")] use crate::rpc::ddl::trigger::CreateTriggerTask; #[cfg(feature = "enterprise")] use crate::rpc::ddl::trigger::DropTriggerTask; use crate::rpc::ddl::{ AlterDatabaseTask, AlterTableTask, CommentOnTask, CreateDatabaseTask, CreateFlowTask, CreateTableTask, CreateViewTask, DropDatabaseTask, DropFlowTask, DropTableTask, DropViewTask, - PersistentEventContext, PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest, - SubmitDdlTaskResponse, TruncateTableTask, UndropTableTask, + PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest, SubmitDdlTaskResponse, + TruncateTableTask, UndropTableTask, }; const MAX_REGION_ROUTE_CHANGE_RETRIES: usize = 3; @@ -122,6 +123,7 @@ pub trait TriggerDdlManager: Send + Sync { procedure_manager: ProcedureManagerRef, ddl_context: DdlContext, query_context: QueryContext, + procedure_context: ProcedureContext, ) -> Result; async fn drop_trigger( @@ -130,6 +132,7 @@ pub trait TriggerDdlManager: Send + Sync { procedure_manager: ProcedureManagerRef, ddl_context: DdlContext, query_context: QueryContext, + procedure_context: ProcedureContext, ) -> Result; fn as_any(&self) -> &dyn std::any::Any; @@ -341,7 +344,7 @@ impl DdlManager { repartition: Repartition, wait: bool, timeout: Duration, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); @@ -384,7 +387,7 @@ impl DdlManager { .await .context(PersistRepartitionGcRequirementSnafu)?; let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); if wait { self.execute_procedure_and_wait(procedure_with_id).await } else { @@ -400,7 +403,7 @@ impl DdlManager { &self, table_id: TableId, alter_table_task: AlterTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ddl_options: DdlOptions, ) -> Result<(ProcedureId, Option)> { // make alter_table_task mutable so we can call .take() on its field @@ -421,7 +424,7 @@ impl DdlManager { repartition, ddl_options.wait, ddl_options.timeout, - event_context, + procedure_context, ) .await; } @@ -460,7 +463,7 @@ impl DdlManager { )?; let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) - .with_event_context(event_context.clone()); + .with_context(procedure_context.clone()); let result = self.execute_procedure_and_wait(procedure_with_id).await?; if result .1 @@ -489,7 +492,7 @@ impl DdlManager { &self, create_table_task: CreateTableTask, query_context: QueryContext, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); @@ -500,7 +503,7 @@ impl DdlManager { )?; let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -510,14 +513,14 @@ impl DdlManager { pub async fn submit_create_view_task( &self, create_view_task: CreateViewTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = CreateViewProcedure::new(create_view_task, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -528,7 +531,7 @@ impl DdlManager { &self, create_table_tasks: Vec, physical_table_id: TableId, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); @@ -536,7 +539,7 @@ impl DdlManager { CreateLogicalTablesProcedure::new(create_table_tasks, physical_table_id, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -547,7 +550,7 @@ impl DdlManager { &self, alter_table_tasks: Vec, physical_table_id: TableId, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); @@ -555,7 +558,7 @@ impl DdlManager { AlterLogicalTablesProcedure::new(alter_table_tasks, physical_table_id, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -565,14 +568,14 @@ impl DdlManager { pub async fn submit_drop_table_task( &self, drop_table_task: DropTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = DropTableProcedure::new(drop_table_task, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -583,7 +586,7 @@ impl DdlManager { pub async fn submit_undrop_table_task( &self, undrop_table_task: UndropTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { #[cfg(not(feature = "enterprise"))] { @@ -611,7 +614,7 @@ impl DdlManager { Some(original_table_name), ); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) - .with_event_context(event_context); + .with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -623,7 +626,7 @@ impl DdlManager { pub async fn submit_purge_dropped_table_task( &self, purge_dropped_table_task: PurgeDroppedTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { #[cfg(not(feature = "enterprise"))] { @@ -639,7 +642,7 @@ impl DdlManager { let context = self.create_context(); let procedure = PurgeDroppedTableProcedure::new(purge_dropped_table_task, context); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) - .with_event_context(event_context); + .with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -665,11 +668,13 @@ impl DdlManager { #[cfg(feature = "enterprise")] { let context = self.create_context(); - let event_context = PersistentEventContext::new(TriggerReason::ScheduledGc); + let procedure_context = ProcedureContext::from_event_context( + PersistentEventContext::new(TriggerReason::ScheduledGc), + ); let procedure = PurgeDroppedTableProcedure::new_if_expired(purge_dropped_table_task, context); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) - .with_event_context(event_context); + .with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -686,7 +691,7 @@ impl DdlManager { options, creator, }: CreateDatabaseTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = CreateDatabaseProcedure::new( @@ -698,7 +703,7 @@ impl DdlManager { context, ); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -712,12 +717,12 @@ impl DdlManager { schema, drop_if_exists, }: DropDatabaseTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = DropDatabaseProcedure::new(catalog, schema, drop_if_exists, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -725,12 +730,12 @@ impl DdlManager { pub async fn submit_alter_database( &self, alter_database_task: AlterDatabaseTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = AlterDatabaseProcedure::new(alter_database_task, context)?; let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -741,12 +746,12 @@ impl DdlManager { &self, create_flow: CreateFlowTask, query_context: QueryContext, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = CreateFlowProcedure::new(create_flow, query_context, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -756,12 +761,12 @@ impl DdlManager { pub async fn submit_drop_flow_task( &self, drop_flow: DropFlowTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = DropFlowProcedure::new(drop_flow, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -771,12 +776,12 @@ impl DdlManager { pub async fn submit_drop_view_task( &self, drop_view: DropViewTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = DropViewProcedure::new(drop_view, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -787,13 +792,13 @@ impl DdlManager { &self, truncate_table_task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = TruncateTableProcedure::new(truncate_table_task, table_info_value, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -803,7 +808,7 @@ impl DdlManager { pub async fn submit_comment_on_task( &self, mut comment_on_task: CommentOnTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); comment_on_task @@ -814,7 +819,7 @@ impl DdlManager { .await?; let procedure = CommentOnProcedure::new(comment_on_task, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -853,89 +858,125 @@ impl DdlManager { pub async fn submit_ddl_task( &self, - ctx: &ExecutorContext, + context: ExecutorContext, request: SubmitDdlTaskRequest, ) -> Result { - let span = ctx - .tracing_context + let ExecutorContext { + tracing_context, + query_context, + actor, + event_input, + } = context; + let query_context = query_context.context(UnsupportedSnafu { + operation: "submit_ddl_task without query context", + })?; + let procedure_context = ProcedureContext { + actor, + event_context: event_input + .map(|input| PersistentEventContext::from((input, query_context.protocol()))), + }; + let span = tracing_context .as_ref() .map(TracingContext::from_w3c) .unwrap_or_else(TracingContext::from_current_span) .attach(tracing::info_span!("DdlManager::submit_ddl_task")); let SubmitDdlTaskRequest { - query_context, wait, timeout, task, } = request; - let event_context = ctx - .event_context - .clone() - .unwrap_or_else(|| crate::rpc::ddl::event_context_from_query_context(&query_context)); let ddl_options = DdlOptions { wait, timeout }; async move { debug!("Submitting Ddl task: {:?}", task); match task { CreateTable(create_table_task) => { - handle_create_table_task(self, create_table_task, query_context, event_context) - .await + handle_create_table_task( + self, + create_table_task, + query_context, + procedure_context, + ) + .await } DropTable(drop_table_task) => { - handle_drop_table_task(self, drop_table_task, event_context).await + handle_drop_table_task(self, drop_table_task, procedure_context).await } UndropTable(undrop_table_task) => { - handle_undrop_table_task(self, undrop_table_task, event_context).await + handle_undrop_table_task(self, undrop_table_task, procedure_context).await } PurgeDroppedTable(purge_dropped_table_task) => { - handle_purge_dropped_table_task(self, purge_dropped_table_task, event_context) - .await + handle_purge_dropped_table_task( + self, + purge_dropped_table_task, + procedure_context, + ) + .await } AlterTable(alter_table_task) => { - handle_alter_table_task(self, alter_table_task, ddl_options, event_context) + handle_alter_table_task(self, alter_table_task, ddl_options, procedure_context) .await } TruncateTable(truncate_table_task) => { - handle_truncate_table_task(self, truncate_table_task, event_context).await + handle_truncate_table_task(self, truncate_table_task, procedure_context).await } CreateLogicalTables(create_table_tasks) => { - handle_create_logical_table_tasks(self, create_table_tasks, event_context).await + handle_create_logical_table_tasks(self, create_table_tasks, procedure_context) + .await } AlterLogicalTables(alter_table_tasks) => { - handle_alter_logical_table_tasks(self, alter_table_tasks, event_context).await + handle_alter_logical_table_tasks(self, alter_table_tasks, procedure_context) + .await } DropLogicalTables(_) => todo!(), CreateDatabase(create_database_task) => { - handle_create_database_task(self, create_database_task, event_context).await + handle_create_database_task(self, create_database_task, procedure_context).await } DropDatabase(drop_database_task) => { - handle_drop_database_task(self, drop_database_task, event_context).await + handle_drop_database_task(self, drop_database_task, procedure_context).await } AlterDatabase(alter_database_task) => { - handle_alter_database_task(self, alter_database_task, event_context).await + handle_alter_database_task(self, alter_database_task, procedure_context).await } CreateFlow(create_flow_task) => { - handle_create_flow_task(self, create_flow_task, query_context, event_context) - .await + handle_create_flow_task( + self, + create_flow_task, + query_context, + procedure_context, + ) + .await } DropFlow(drop_flow_task) => { - handle_drop_flow_task(self, drop_flow_task, event_context).await + handle_drop_flow_task(self, drop_flow_task, procedure_context).await } CreateView(create_view_task) => { - handle_create_view_task(self, create_view_task, event_context).await + handle_create_view_task(self, create_view_task, procedure_context).await } DropView(drop_view_task) => { - handle_drop_view_task(self, drop_view_task, event_context).await + handle_drop_view_task(self, drop_view_task, procedure_context).await } CommentOn(comment_on_task) => { - handle_comment_on_task(self, comment_on_task, event_context).await + handle_comment_on_task(self, comment_on_task, procedure_context).await } #[cfg(feature = "enterprise")] CreateTrigger(create_trigger_task) => { - handle_create_trigger_task(self, create_trigger_task, query_context).await + handle_create_trigger_task( + self, + create_trigger_task, + query_context, + procedure_context, + ) + .await } #[cfg(feature = "enterprise")] DropTrigger(drop_trigger_task) => { - handle_drop_trigger_task(self, drop_trigger_task, query_context).await + handle_drop_trigger_task( + self, + drop_trigger_task, + query_context, + procedure_context, + ) + .await } } } @@ -947,7 +988,7 @@ impl DdlManager { async fn handle_truncate_table_task( ddl_manager: &DdlManager, truncate_table_task: TruncateTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let table_id = truncate_table_task.table_id; let table_metadata_manager = &ddl_manager.table_metadata_manager(); @@ -972,7 +1013,7 @@ async fn handle_truncate_table_task( ); let (id, _) = ddl_manager - .submit_truncate_table_task(truncate_table_task, table_info_value, event_context) + .submit_truncate_table_task(truncate_table_task, table_info_value, procedure_context) .await?; info!("Table: {table_id} is truncated via procedure_id {id:?}"); @@ -987,7 +1028,7 @@ async fn handle_alter_table_task( ddl_manager: &DdlManager, alter_table_task: AlterTableTask, ddl_options: DdlOptions, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let table_ref = alter_table_task.table_ref(); @@ -1020,7 +1061,7 @@ async fn handle_alter_table_task( ); let (id, _) = ddl_manager - .submit_alter_table_task(table_id, alter_table_task, event_context, ddl_options) + .submit_alter_table_task(table_id, alter_table_task, procedure_context, ddl_options) .await?; info!("Table: {table_id} is altered via procedure_id {id:?}"); @@ -1034,11 +1075,11 @@ async fn handle_alter_table_task( async fn handle_drop_table_task( ddl_manager: &DdlManager, drop_table_task: DropTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let table_id = drop_table_task.table_id; let (id, _) = ddl_manager - .submit_drop_table_task(drop_table_task, event_context) + .submit_drop_table_task(drop_table_task, procedure_context) .await?; info!("Table: {table_id} is dropped via procedure_id {id:?}"); @@ -1052,11 +1093,11 @@ async fn handle_drop_table_task( async fn handle_undrop_table_task( ddl_manager: &DdlManager, undrop_table_task: UndropTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let table_id = undrop_table_task.table_id; let (id, _) = ddl_manager - .submit_undrop_table_task(undrop_table_task, event_context) + .submit_undrop_table_task(undrop_table_task, procedure_context) .await?; info!("Table: {table_id} is undropped via procedure_id {id:?}"); @@ -1070,10 +1111,10 @@ async fn handle_undrop_table_task( async fn handle_purge_dropped_table_task( ddl_manager: &DdlManager, purge_dropped_table_task: PurgeDroppedTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, _) = ddl_manager - .submit_purge_dropped_table_task(purge_dropped_table_task, event_context) + .submit_purge_dropped_table_task(purge_dropped_table_task, procedure_context) .await?; info!("Dropped table is purged via procedure_id {id:?}"); @@ -1088,10 +1129,10 @@ async fn handle_create_table_task( ddl_manager: &DdlManager, create_table_task: CreateTableTask, query_context: QueryContext, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, output) = ddl_manager - .submit_create_table_task(create_table_task, query_context, event_context) + .submit_create_table_task(create_table_task, query_context, procedure_context) .await?; let procedure_id = id.to_string(); @@ -1114,7 +1155,7 @@ async fn handle_create_table_task( async fn handle_create_logical_table_tasks( ddl_manager: &DdlManager, create_table_tasks: Vec, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { ensure!( !create_table_tasks.is_empty(), @@ -1130,7 +1171,7 @@ async fn handle_create_logical_table_tasks( let num_logical_tables = create_table_tasks.len(); let (id, output) = ddl_manager - .submit_create_logical_table_tasks(create_table_tasks, physical_table_id, event_context) + .submit_create_logical_table_tasks(create_table_tasks, physical_table_id, procedure_context) .await?; info!( @@ -1159,12 +1200,12 @@ async fn handle_create_logical_table_tasks( async fn handle_create_database_task( ddl_manager: &DdlManager, create_database_task: CreateDatabaseTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let catalog = create_database_task.catalog.clone(); let schema = create_database_task.schema.clone(); let (id, _) = ddl_manager - .submit_create_database(create_database_task, event_context) + .submit_create_database(create_database_task, procedure_context) .await?; let procedure_id = id.to_string(); @@ -1182,10 +1223,10 @@ async fn handle_create_database_task( async fn handle_drop_database_task( ddl_manager: &DdlManager, drop_database_task: DropDatabaseTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, _) = ddl_manager - .submit_drop_database(drop_database_task.clone(), event_context) + .submit_drop_database(drop_database_task.clone(), procedure_context) .await?; let procedure_id = id.to_string(); @@ -1203,10 +1244,10 @@ async fn handle_drop_database_task( async fn handle_alter_database_task( ddl_manager: &DdlManager, alter_database_task: AlterDatabaseTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, _) = ddl_manager - .submit_alter_database(alter_database_task.clone(), event_context) + .submit_alter_database(alter_database_task.clone(), procedure_context) .await?; let procedure_id = id.to_string(); @@ -1225,10 +1266,10 @@ async fn handle_alter_database_task( async fn handle_drop_flow_task( ddl_manager: &DdlManager, drop_flow_task: DropFlowTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, _) = ddl_manager - .submit_drop_flow_task(drop_flow_task.clone(), event_context) + .submit_drop_flow_task(drop_flow_task.clone(), procedure_context) .await?; let procedure_id = id.to_string(); @@ -1248,6 +1289,7 @@ async fn handle_drop_trigger_task( ddl_manager: &DdlManager, drop_trigger_task: DropTriggerTask, query_context: QueryContext, + procedure_context: ProcedureContext, ) -> Result { let Some(m) = ddl_manager.trigger_ddl_manager.as_ref() else { use crate::error::UnsupportedSnafu; @@ -1263,6 +1305,7 @@ async fn handle_drop_trigger_task( ddl_manager.procedure_manager.clone(), ddl_manager.ddl_context.clone(), query_context, + procedure_context, ) .await } @@ -1270,10 +1313,10 @@ async fn handle_drop_trigger_task( async fn handle_drop_view_task( ddl_manager: &DdlManager, drop_view_task: DropViewTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, _) = ddl_manager - .submit_drop_view_task(drop_view_task.clone(), event_context) + .submit_drop_view_task(drop_view_task.clone(), procedure_context) .await?; let procedure_id = id.to_string(); @@ -1293,10 +1336,10 @@ async fn handle_create_flow_task( ddl_manager: &DdlManager, create_flow_task: CreateFlowTask, query_context: QueryContext, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, output) = ddl_manager - .submit_create_flow_task(create_flow_task.clone(), query_context, event_context) + .submit_create_flow_task(create_flow_task.clone(), query_context, procedure_context) .await?; let procedure_id = id.to_string(); @@ -1331,6 +1374,7 @@ async fn handle_create_trigger_task( ddl_manager: &DdlManager, create_trigger_task: CreateTriggerTask, query_context: QueryContext, + procedure_context: ProcedureContext, ) -> Result { let Some(m) = ddl_manager.trigger_ddl_manager.as_ref() else { use crate::error::UnsupportedSnafu; @@ -1346,6 +1390,7 @@ async fn handle_create_trigger_task( ddl_manager.procedure_manager.clone(), ddl_manager.ddl_context.clone(), query_context, + procedure_context, ) .await } @@ -1353,7 +1398,7 @@ async fn handle_create_trigger_task( async fn handle_alter_logical_table_tasks( ddl_manager: &DdlManager, alter_table_tasks: Vec, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { ensure!( !alter_table_tasks.is_empty(), @@ -1373,7 +1418,7 @@ async fn handle_alter_logical_table_tasks( let num_logical_tables = alter_table_tasks.len(); let (id, _) = ddl_manager - .submit_alter_logical_table_tasks(alter_table_tasks, physical_table_id, event_context) + .submit_alter_logical_table_tasks(alter_table_tasks, physical_table_id, procedure_context) .await?; info!( @@ -1392,10 +1437,10 @@ async fn handle_alter_logical_table_tasks( async fn handle_create_view_task( ddl_manager: &DdlManager, create_view_task: CreateViewTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, output) = ddl_manager - .submit_create_view_task(create_view_task, event_context) + .submit_create_view_task(create_view_task, procedure_context) .await?; let procedure_id = id.to_string(); @@ -1418,10 +1463,10 @@ async fn handle_create_view_task( async fn handle_comment_on_task( ddl_manager: &DdlManager, comment_on_task: CommentOnTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, _) = ddl_manager - .submit_comment_on_task(comment_on_task.clone(), event_context) + .submit_comment_on_task(comment_on_task.clone(), procedure_context) .await?; let procedure_id = id.to_string(); @@ -1439,16 +1484,22 @@ async fn handle_comment_on_task( #[cfg(test)] mod tests { use std::sync::Arc; + #[cfg(feature = "enterprise")] + use std::sync::Mutex; use std::time::Duration; + #[cfg(feature = "enterprise")] + use common_base::protocol::Channel; use common_error::ext::BoxedError; #[cfg(feature = "enterprise")] use common_error::ext::ErrorExt; #[cfg(feature = "enterprise")] use common_error::status_code::StatusCode; + #[cfg(feature = "enterprise")] + use common_event_recorder::{PersistentEventContext, ProcedureEventInput, TriggerReason}; use common_procedure::local::LocalManager; use common_procedure::test_util::InMemoryPoisonStore; - use common_procedure::{BoxedProcedure, ProcedureManagerRef}; + use common_procedure::{BoxedProcedure, ProcedureContext, ProcedureManagerRef}; use store_api::storage::TableId; use table::table_name::TableName; @@ -1470,13 +1521,16 @@ mod tests { use crate::kv_backend::memory::MemoryKvBackend; use crate::node_manager::{DatanodeManager, DatanodeRef, FlownodeManager, FlownodeRef}; use crate::peer::Peer; - #[cfg(not(feature = "enterprise"))] use crate::procedure_executor::ExecutorContext; use crate::region_keeper::MemoryRegionKeeper; use crate::region_registry::LeaderRegionRegistry; - use crate::rpc::ddl::{CreatorGrantIntent, PersistentEventContext, UndropTableTask}; + #[cfg(feature = "enterprise")] + use crate::rpc::ddl::trigger::{CreateTriggerTask, DropTriggerTask}; + use crate::rpc::ddl::{CreatorGrantIntent, UndropTableTask}; #[cfg(not(feature = "enterprise"))] use crate::rpc::ddl::{DdlTask, PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest}; + #[cfg(feature = "enterprise")] + use crate::rpc::ddl::{DdlTask, QueryContext, SubmitDdlTaskRequest}; use crate::sequence::SequenceBuilder; use crate::state_store::KvStateStore; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -1543,6 +1597,50 @@ mod tests { } } + #[cfg(feature = "enterprise")] + #[derive(Default)] + struct RecordingTriggerDdlManager { + procedure_contexts: Mutex>, + } + + #[cfg(feature = "enterprise")] + #[async_trait::async_trait] + impl super::TriggerDdlManager for RecordingTriggerDdlManager { + async fn create_trigger( + &self, + _create_trigger_task: CreateTriggerTask, + _procedure_manager: ProcedureManagerRef, + _ddl_context: DdlContext, + _query_context: QueryContext, + procedure_context: ProcedureContext, + ) -> crate::error::Result { + self.procedure_contexts + .lock() + .unwrap() + .push(procedure_context); + Ok(Default::default()) + } + + async fn drop_trigger( + &self, + _drop_trigger_task: DropTriggerTask, + _procedure_manager: ProcedureManagerRef, + _ddl_context: DdlContext, + _query_context: QueryContext, + procedure_context: ProcedureContext, + ) -> crate::error::Result { + self.procedure_contexts + .lock() + .unwrap() + .push(procedure_context); + Ok(Default::default()) + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + } + #[test] fn test_generic_loader_captures_configured_committer() { let mut context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); @@ -1678,6 +1776,67 @@ mod tests { ) } + #[cfg(feature = "enterprise")] + #[tokio::test] + async fn test_trigger_ddl_forwards_procedure_context() { + let trigger_ddl_manager = Arc::new(RecordingTriggerDdlManager::default()); + let ddl_manager = build_soft_drop_test_ddl_manager() + .with_trigger_ddl_manager(trigger_ddl_manager.clone()); + let procedure_context = ProcedureContext { + actor: Some("test-user".to_string()), + event_context: Some( + PersistentEventContext::new(TriggerReason::Manual).with_protocol("mysql"), + ), + }; + let executor_context = || ExecutorContext { + query_context: Some(QueryContext { + channel: Channel::Mysql as u8, + ..Default::default() + }), + actor: Some("test-user".to_string()), + event_input: Some(ProcedureEventInput::new(TriggerReason::Manual)), + ..Default::default() + }; + + ddl_manager + .submit_ddl_task( + executor_context(), + SubmitDdlTaskRequest::new(DdlTask::CreateTrigger(CreateTriggerTask { + catalog_name: "greptime".to_string(), + trigger_name: "test_trigger".to_string(), + if_not_exists: false, + sql: "SELECT 1".to_string(), + channels: vec![], + labels: Default::default(), + annotations: Default::default(), + interval: Duration::from_secs(1), + raw_interval_expr: None, + r#for: None, + for_raw_expr: None, + keep_firing_for: None, + keep_firing_for_raw_expr: None, + })), + ) + .await + .unwrap(); + ddl_manager + .submit_ddl_task( + executor_context(), + SubmitDdlTaskRequest::new(DdlTask::DropTrigger(DropTriggerTask { + catalog_name: "greptime".to_string(), + trigger_name: "test_trigger".to_string(), + drop_if_exists: false, + })), + ) + .await + .unwrap(); + + assert_eq!( + *trigger_ddl_manager.procedure_contexts.lock().unwrap(), + vec![procedure_context.clone(), procedure_context] + ); + } + #[cfg(feature = "enterprise")] #[tokio::test] async fn test_submit_undrop_missing_tombstone_returns_table_not_found_directly() { @@ -1686,7 +1845,7 @@ mod tests { let err = ddl_manager .submit_undrop_table_task( UndropTableTask { table_id: 1024 }, - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap_err(); @@ -1703,7 +1862,7 @@ mod tests { let err = ddl_manager .submit_undrop_table_task( UndropTableTask { table_id: 1024 }, - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap_err(); @@ -1712,7 +1871,7 @@ mod tests { let err = ddl_manager .submit_purge_dropped_table_task( PurgeDroppedTableTask { table_id: 1024 }, - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap_err(); @@ -1730,8 +1889,11 @@ mod tests { ] { let err = ddl_manager .submit_ddl_task( - &ExecutorContext::default(), - SubmitDdlTaskRequest::new(QueryContext::default(), task), + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + SubmitDdlTaskRequest::new(task), ) .await .unwrap_err(); diff --git a/src/common/meta/src/procedure_executor.rs b/src/common/meta/src/procedure_executor.rs index b812aeb6aa..a873c26b39 100644 --- a/src/common/meta/src/procedure_executor.rs +++ b/src/common/meta/src/procedure_executor.rs @@ -15,7 +15,7 @@ use std::sync::Arc; use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse}; -use common_event_recorder::PersistentEventContext; +use common_event_recorder::ProcedureEventInput; use common_procedure::{ProcedureId, ProcedureManagerRef}; use common_telemetry::tracing_context::W3cTrace; use snafu::{OptionExt, ResultExt}; @@ -24,7 +24,7 @@ use crate::ddl_manager::DdlManagerRef; use crate::error::{ ParseProcedureIdSnafu, ProcedureNotFoundSnafu, QueryProcedureSnafu, Result, UnsupportedSnafu, }; -use crate::rpc::ddl::{SubmitDdlTaskRequest, SubmitDdlTaskResponse}; +use crate::rpc::ddl::{QueryContext, SubmitDdlTaskRequest, SubmitDdlTaskResponse}; use crate::rpc::procedure::{ self, GcRegionsRequest, GcResponse, GcTableRequest, ManageRegionFollowerRequest, MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse, @@ -34,8 +34,14 @@ use crate::rpc::procedure::{ #[derive(Debug, Clone, Default)] pub struct ExecutorContext { pub tracing_context: Option, + /// Query execution data available at the frontend/standalone submission boundary. + /// + /// Only DDL serializes the full context. Migration and GC use its typed + /// channel to derive protocol without sending the query context itself. + pub query_context: Option, pub actor: Option, - pub event_context: Option, + /// Caller-supplied event metadata. Protocol is derived by the submission adapter. + pub event_input: Option, } /// The procedure executor that accepts ddl, region migration task etc. @@ -44,7 +50,7 @@ pub trait ProcedureExecutor: Send + Sync { /// Submit a ddl task async fn submit_ddl_task( &self, - ctx: &ExecutorContext, + ctx: ExecutorContext, request: SubmitDdlTaskRequest, ) -> Result; @@ -129,7 +135,7 @@ impl LocalProcedureExecutor { impl ProcedureExecutor for LocalProcedureExecutor { async fn submit_ddl_task( &self, - ctx: &ExecutorContext, + ctx: ExecutorContext, request: SubmitDdlTaskRequest, ) -> Result { self.ddl_manager.submit_ddl_task(ctx, request).await diff --git a/src/common/meta/src/rpc/ddl.rs b/src/common/meta/src/rpc/ddl.rs index d796680473..47bd537b9b 100644 --- a/src/common/meta/src/rpc/ddl.rs +++ b/src/common/meta/src/rpc/ddl.rs @@ -45,7 +45,7 @@ use base64::engine::general_purpose; use common_base::protocol::Channel; use common_catalog::{format_full_flow_name, format_full_table_name}; use common_error::ext::BoxedError; -pub use common_event_recorder::{PersistentEventContext, TriggerReason}; +pub use common_event_recorder::TriggerReason; use common_time::{DatabaseTimeToLive, Timestamp}; use prost::Message; use serde::{Deserialize, Serialize}; @@ -68,8 +68,6 @@ use crate::key::table_name::{TableNameKey, TableNameManager}; /// Reserved query-context extension key for the frontend peer address that submitted a DDL request. pub const ORIGIN_FRONTEND_ADDR_EXTENSION_KEY: &str = "__greptime_origin_frontend.addr"; -/// Reserved query-context extension key for the trigger reason supplied by frontend. -pub const TRIGGER_REASON_EXTENSION_KEY: &str = "__greptime_event.trigger_reason"; /// Reserved query-context extension key for the authenticated database creator. pub const CREATE_DATABASE_CREATOR_EXTENSION_KEY: &str = "__greptime_create_database.creator"; /// Internal gRPC metadata key for the authenticated database creator. @@ -336,7 +334,6 @@ impl TryFrom for DdlTask { #[derive(Clone)] pub struct SubmitDdlTaskRequest { - pub query_context: QueryContext, pub wait: bool, pub timeout: Duration, pub task: DdlTask, @@ -344,9 +341,8 @@ pub struct SubmitDdlTaskRequest { impl SubmitDdlTaskRequest { /// The default constructor for [`SubmitDdlTaskRequest`]. - pub fn new(query_context: QueryContext, task: DdlTask) -> Self { + pub fn new(task: DdlTask) -> Self { Self { - query_context, wait: Self::default_wait(), timeout: Self::default_timeout(), task, @@ -377,26 +373,11 @@ impl TryFrom for PbDdlTaskRequest { fn try_from(request: SubmitDdlTaskRequest) -> Result { let SubmitDdlTaskRequest { - mut query_context, wait, timeout, task, } = request; - query_context - .extensions - .remove(CREATE_DATABASE_CREATOR_EXTENSION_KEY); - if let DdlTask::CreateDatabase(CreateDatabaseTask { - creator: Some(creator), - .. - }) = &task - { - query_context.extensions.insert( - CREATE_DATABASE_CREATOR_EXTENSION_KEY.to_string(), - serde_json::to_string(creator).context(error::SerdeJsonSnafu)?, - ); - } - let task = match task { DdlTask::CreateTable(task) => Task::CreateTableTask(task.try_into()?), DdlTask::DropTable(task) => Task::DropTableTask(task.into()), @@ -444,10 +425,12 @@ impl TryFrom for PbDdlTaskRequest { Ok(Self { header: None, - query_context: Some(query_context.into()), + query_context: None, timeout_secs: ddl_timeout_secs(timeout), wait, task: Some(task), + event_context: None, + actor: None, }) } } @@ -1668,22 +1651,6 @@ pub struct QueryContext { pub sst_min_sequences: HashMap, } -/// Builds event context from frontend query context metadata. -pub fn event_context_from_query_context(query_context: &QueryContext) -> PersistentEventContext { - let reason = query_context - .extensions - .get(TRIGGER_REASON_EXTENSION_KEY) - .map(|reason| TriggerReason::from_extension(reason)) - .unwrap_or_default(); - let channel = Channel::from(u32::from(query_context.channel)); - let context = PersistentEventContext::new(reason); - if channel == Channel::Unknown { - context - } else { - context.with_protocol(channel.as_ref()) - } -} - impl QueryContext { /// Get the current catalog pub fn current_catalog(&self) -> &str { @@ -1710,6 +1677,12 @@ impl QueryContext { self.channel } + /// Returns the protocol derived from the typed query channel. + pub fn protocol(&self) -> Option { + let channel = Channel::from(u32::from(self.channel)); + (channel != Channel::Unknown).then(|| channel.as_ref().to_string()) + } + pub fn snapshot_seqs(&self) -> &HashMap { &self.snapshot_seqs } @@ -1922,10 +1895,7 @@ mod tests { #[test] fn test_undrop_table_task_pb_roundtrip() { let expected = UndropTableTask { table_id: 1024 }; - let request = SubmitDdlTaskRequest::new( - QueryContext::default(), - DdlTask::UndropTable(expected.clone()), - ); + let request = SubmitDdlTaskRequest::new(DdlTask::UndropTable(expected.clone())); let pb = PbDdlTaskRequest::try_from(request).unwrap(); let pb_task = pb.task.unwrap(); @@ -1937,10 +1907,7 @@ mod tests { #[test] fn test_purge_dropped_table_task_pb_roundtrip() { let expected = PurgeDroppedTableTask { table_id: 1024 }; - let request = SubmitDdlTaskRequest::new( - QueryContext::default(), - DdlTask::PurgeDroppedTable(expected.clone()), - ); + let request = SubmitDdlTaskRequest::new(DdlTask::PurgeDroppedTable(expected.clone())); let pb = PbDdlTaskRequest::try_from(request).unwrap(); let pb_task = pb.task.unwrap(); @@ -1949,53 +1916,6 @@ mod tests { assert!(matches!(de, DdlTask::PurgeDroppedTable(task) if task == expected)); } - #[test] - fn test_create_database_creator_transport_sanitizes_extension() { - let request = |creator| { - let mut query_context = QueryContext::default(); - query_context.extensions.insert( - CREATE_DATABASE_CREATOR_EXTENSION_KEY.to_string(), - "spoofed".to_string(), - ); - SubmitDdlTaskRequest::new( - query_context, - DdlTask::new_create_database( - "greptime".to_string(), - "test".to_string(), - false, - HashMap::new(), - creator, - ), - ) - }; - let creator = CreatorGrantIntent { - username: "alice".to_string(), - created_at_ns: 42, - }; - - let mut pb = PbDdlTaskRequest::try_from(request(Some(creator.clone()))).unwrap(); - let encoded = pb - .query_context - .as_ref() - .unwrap() - .extensions - .get(CREATE_DATABASE_CREATOR_EXTENSION_KEY) - .unwrap(); - assert_eq!(creator, serde_json::from_str(encoded).unwrap()); - assert!(matches!( - DdlTask::try_from(pb.task.take().unwrap()).unwrap(), - DdlTask::CreateDatabase(CreateDatabaseTask { creator: None, .. }) - )); - - let pb = PbDdlTaskRequest::try_from(request(None)).unwrap(); - assert!( - !pb.query_context - .unwrap() - .extensions - .contains_key(CREATE_DATABASE_CREATOR_EXTENSION_KEY) - ); - } - #[test] fn test_undrop_table_task_json_roundtrip() { let task = UndropTableTask { table_id: 1024 }; @@ -2262,44 +2182,4 @@ mod tests { let reason: TriggerReason = serde_json::from_str("\"future_reason\"").unwrap(); assert_eq!(TriggerReason::Unknown, reason); } - - #[test] - fn test_event_context_serialization() { - let context = PersistentEventContext::new(TriggerReason::Manual).with_protocol("mysql"); - - assert_eq!( - serde_json::json!({ - "reason": "manual", - "protocol": "mysql", - }), - serde_json::to_value(context).unwrap() - ); - - assert_eq!( - serde_json::json!({ "reason": "manual" }), - serde_json::to_value(PersistentEventContext::new(TriggerReason::Manual)).unwrap() - ); - } - - #[test] - fn test_event_context_from_query_context_preserves_extensions() { - let mut query_context = QueryContext::default(); - query_context.extensions.insert( - TRIGGER_REASON_EXTENSION_KEY.to_string(), - TriggerReason::AutoCreate.as_ref().to_string(), - ); - query_context.channel = 4; - - let event_context = event_context_from_query_context(&query_context); - - assert_eq!(event_context.reason, TriggerReason::AutoCreate); - assert_eq!(event_context.protocol.as_deref(), Some("prometheus")); - assert_eq!( - query_context - .extensions - .get(TRIGGER_REASON_EXTENSION_KEY) - .map(String::as_str), - Some("auto_create") - ); - } } diff --git a/src/common/procedure/src/local.rs b/src/common/procedure/src/local.rs index 8fdd316758..c870cc39a6 100644 --- a/src/common/procedure/src/local.rs +++ b/src/common/procedure/src/local.rs @@ -1062,7 +1062,7 @@ mod tests { use super::*; use crate::error::{self, Error}; - use crate::store::state_store::ObjectStateStore; + use crate::store::state_store::{ObjectStateStore, StateStore}; use crate::test_util::InMemoryPoisonStore; use crate::{ ChildSubmissionOutcome, Context, EventContext, EventTrigger, Procedure, ProcedureEvent, @@ -1721,7 +1721,63 @@ mod tests { } #[tokio::test] - async fn test_recovery_emits_recovered_event() { + async fn test_recovery_loads_legacy_procedure_message() { + let dir = create_temp_dir("legacy_procedure_message_recovery"); + let object_store = test_util::new_object_store(&dir); + let state_store = Arc::new(ObjectStateStore::new(object_store)); + let poison_manager = Arc::new(InMemoryPoisonStore::new()); + let event_recorder = Arc::new(CapturingEventRecorder::default()); + let manager = LocalManager::new( + ManagerConfig { + parent_path: "data/".to_string(), + ..Default::default() + }, + state_store.clone(), + poison_manager, + None, + None, + ); + manager + .event_recorder_handle() + .install(event_recorder.clone()); + manager.manager_ctx.start(); + manager + .register_loader("ProcedureToLoad", ProcedureToLoad::loader()) + .unwrap(); + + let procedure = ProcedureToLoad::new("legacy recovered submission"); + let procedure_id = ProcedureId::random(); + let key = format!("data/procedure/{procedure_id}/0000000000.step"); + // This is the persisted JSON shape written before ProcedureMessage had `context`. + let legacy_message = serde_json::json!({ + "type_name": procedure.type_name(), + "data": procedure.dump().unwrap(), + "parent_id": null, + "step": 0, + }); + state_store + .put(&key, serde_json::to_vec(&legacy_message).unwrap()) + .await + .unwrap(); + + manager.recover().await.unwrap(); + + assert!( + manager + .procedure_state(procedure_id) + .await + .unwrap() + .is_some() + ); + wait_for_trigger(&event_recorder, EventTrigger::Recovered).await; + assert_eq!( + event_recorder.procedure_context(procedure_id), + Some(ProcedureContext::default()) + ); + } + + #[tokio::test] + async fn test_recovery_preserves_procedure_context() { let dir = create_temp_dir("recovery_submission_event"); let object_store = test_util::new_object_store(&dir); let state_store = Arc::new(ObjectStateStore::new(object_store.clone())); @@ -1747,6 +1803,16 @@ mod tests { let procedure = ProcedureToLoad::new("recovered submission"); let procedure_id = ProcedureId::random(); + let mut event_context = + PersistentEventContext::new(TriggerReason::AutoRebalance).with_protocol("mysql"); + event_context.extensions.insert( + "source".to_string(), + serde_json::Value::String("test".to_string()), + ); + let procedure_context = ProcedureContext { + actor: Some("alice".to_string()), + event_context: Some(event_context), + }; ProcedureStore::from_object_store(object_store) .store_procedure_with_context( procedure_id, @@ -1754,9 +1820,7 @@ mod tests { procedure.type_name().to_string(), procedure.dump().unwrap(), None, - ProcedureContext::from_event_context(PersistentEventContext::new( - TriggerReason::AutoRebalance, - )), + procedure_context.clone(), ) .await .unwrap(); @@ -1773,13 +1837,8 @@ mod tests { wait_for_trigger(&event_recorder, EventTrigger::Recovered).await; assert!(!event_recorder.triggers().contains(&EventTrigger::Submitted)); assert_eq!( - event_recorder - .procedure_context(procedure_id) - .unwrap() - .event_context - .unwrap() - .reason, - TriggerReason::AutoRebalance + event_recorder.procedure_context(procedure_id), + Some(procedure_context) ); } diff --git a/src/common/procedure/src/procedure.rs b/src/common/procedure/src/procedure.rs index cd11aea921..ba6a5d275c 100644 --- a/src/common/procedure/src/procedure.rs +++ b/src/common/procedure/src/procedure.rs @@ -35,6 +35,9 @@ pub type Output = Arc; /// Context attached to a procedure submission and inherited by its children. #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] pub struct ProcedureContext { + /// Effective user that submitted the procedure. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub actor: Option, /// Context describing why and how the procedure was submitted. #[serde(default, skip_serializing_if = "Option::is_none")] pub event_context: Option, @@ -43,12 +46,13 @@ pub struct ProcedureContext { impl ProcedureContext { pub fn from_event_context(event_context: PersistentEventContext) -> Self { Self { + actor: None, event_context: Some(event_context), } } pub fn is_empty(&self) -> bool { - self.event_context.is_none() + self.actor.is_none() && self.event_context.is_none() } } @@ -498,9 +502,9 @@ impl ProcedureWithId { } } - /// Attaches event context to this procedure submission. - pub fn with_event_context(mut self, event_context: PersistentEventContext) -> Self { - self.context.event_context = Some(event_context); + /// Attaches a complete persisted context to this procedure submission. + pub fn with_context(mut self, context: ProcedureContext) -> Self { + self.context = context; self } } diff --git a/src/common/procedure/src/store.rs b/src/common/procedure/src/store.rs index 4867046f8d..b2f19b2f94 100644 --- a/src/common/procedure/src/store.rs +++ b/src/common/procedure/src/store.rs @@ -514,6 +514,14 @@ mod tests { r#"{"type_name":"TestMessage","data":"no parent id","parent_id":"9f805a1f-05f7-490c-9f91-bd56e3cc54c1","step":4,"context":{"event_context":{"reason":"auto_rebalance"}}}"# ); + message.context.actor = Some("alice".to_string()); + let json = serde_json::to_string(&message).unwrap(); + assert!(json.contains(r#""actor":"alice""#)); + assert_eq!( + serde_json::from_str::(&json).unwrap(), + message + ); + let legacy: ProcedureMessage = serde_json::from_str( r#"{"type_name":"TestMessage","data":"legacy","parent_id":null,"step":1}"#, ) diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 0c300835e0..78b701a344 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -2281,7 +2281,7 @@ mod tests { impl ProcedureExecutor for NoopProcedureExecutor { async fn submit_ddl_task( &self, - _ctx: &ExecutorContext, + _ctx: ExecutorContext, _request: SubmitDdlTaskRequest, ) -> common_meta::error::Result { common_meta::error::UnsupportedSnafu { @@ -2358,7 +2358,7 @@ mod tests { impl ProcedureExecutor for MockProcedureExecutor { async fn submit_ddl_task( &self, - _ctx: &ExecutorContext, + _ctx: ExecutorContext, request: SubmitDdlTaskRequest, ) -> common_meta::error::Result { self.submitted.lock().unwrap().push(request.task.clone()); diff --git a/src/meta-client/Cargo.toml b/src/meta-client/Cargo.toml index 28c03b3f0d..984f434cc5 100644 --- a/src/meta-client/Cargo.toml +++ b/src/meta-client/Cargo.toml @@ -30,6 +30,7 @@ tokio-stream = { workspace = true, features = ["net"] } tonic.workspace = true [dev-dependencies] +common-event-recorder.workspace = true datatypes.workspace = true futures.workspace = true meta-srv = { workspace = true, features = ["mock"] } diff --git a/src/meta-client/src/client.rs b/src/meta-client/src/client.rs index 4fc6020fff..e0932af5a6 100644 --- a/src/meta-client/src/client.rs +++ b/src/meta-client/src/client.rs @@ -49,7 +49,10 @@ use common_meta::peer::PeerDiscovery; use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor}; use common_meta::range_stream::PaginationStream; use common_meta::rpc::KeyValue; -use common_meta::rpc::ddl::{SubmitDdlTaskRequest, SubmitDdlTaskResponse}; +use common_meta::rpc::ddl::{ + CREATE_DATABASE_CREATOR_EXTENSION_KEY, CreateDatabaseTask, DdlTask, SubmitDdlTaskRequest, + SubmitDdlTaskResponse, +}; use common_meta::rpc::procedure::{ AddRegionFollowerRequest, AddTableFollowerRequest, GcRegionsRequest, GcResponse, GcTableRequest, ManageRegionFollowerRequest, MigrateRegionRequest, MigrateRegionResponse, @@ -66,7 +69,7 @@ use common_time::util::DefaultSystemTimer; use config::Client as ConfigClient; use futures::TryStreamExt; use heartbeat::{Client as HeartbeatClient, HeartbeatConfig}; -use procedure::Client as ProcedureClient; +use procedure::{Client as ProcedureClient, procedure_actor, procedure_event_context}; use serde::de::DeserializeOwned; use snafu::{OptionExt, ResultExt}; use store::Client as StoreClient; @@ -75,7 +78,7 @@ pub use self::heartbeat::{HeartbeatSender, HeartbeatStream}; use crate::client::ask_leader::{LeaderProviderFactoryImpl, LeaderProviderFactoryRef}; use crate::error::{ ConvertMetaConfigSnafu, ConvertMetaRequestSnafu, ConvertMetaResponseSnafu, Error, - GetFlowStatSnafu, NotStartedSnafu, Result, + GetFlowStatSnafu, MissingQueryContextSnafu, NotStartedSnafu, Result, }; pub type Id = u64; @@ -334,10 +337,10 @@ pub trait RegionFollowerClient: Sync + Send + Debug { impl ProcedureExecutor for MetaClient { async fn submit_ddl_task( &self, - _ctx: &ExecutorContext, + ctx: ExecutorContext, request: SubmitDdlTaskRequest, ) -> MetaResult { - self.submit_ddl_task(request) + MetaClient::submit_ddl_task(self, ctx, request) .await .map_err(BoxedError::new) .context(meta_error::ExternalSnafu) @@ -345,10 +348,10 @@ impl ProcedureExecutor for MetaClient { async fn migrate_region( &self, - _ctx: &ExecutorContext, + ctx: &ExecutorContext, request: MigrateRegionRequest, ) -> MetaResult { - self.migrate_region(request) + self.migrate_region(ctx, request) .await .map_err(BoxedError::new) .context(meta_error::ExternalSnafu) @@ -418,10 +421,10 @@ impl ProcedureExecutor for MetaClient { async fn gc_regions( &self, - _ctx: &ExecutorContext, + ctx: &ExecutorContext, request: GcRegionsRequest, ) -> MetaResult { - self.gc_regions(request) + self.gc_regions(ctx, request) .await .map_err(BoxedError::new) .context(meta_error::ExternalSnafu) @@ -429,10 +432,10 @@ impl ProcedureExecutor for MetaClient { async fn gc_table( &self, - _ctx: &ExecutorContext, + ctx: &ExecutorContext, request: GcTableRequest, ) -> MetaResult { - self.gc_table(request) + self.gc_table(ctx, request) .await .map_err(BoxedError::new) .context(meta_error::ExternalSnafu) @@ -807,10 +810,12 @@ impl MetaClient { /// Submit a region migration task. pub async fn migrate_region( &self, + context: &ExecutorContext, request: MigrateRegionRequest, ) -> Result { self.procedure_client()? .migrate_region( + context, request.region_id, request.from_peer, request.to_peer, @@ -825,28 +830,57 @@ impl MetaClient { } /// Manually trigger GC for specific regions. - pub async fn gc_regions(&self, request: GcRegionsRequest) -> Result { - self.procedure_client()?.gc_regions(request).await + pub async fn gc_regions( + &self, + context: &ExecutorContext, + request: GcRegionsRequest, + ) -> Result { + self.procedure_client()?.gc_regions(context, request).await } /// Manually trigger GC for a table (all its regions). - pub async fn gc_table(&self, request: GcTableRequest) -> Result { - self.procedure_client()?.gc_table(request).await + pub async fn gc_table( + &self, + context: &ExecutorContext, + request: GcTableRequest, + ) -> Result { + self.procedure_client()?.gc_table(context, request).await } - /// Submit a DDL task + /// Submit a DDL task. pub async fn submit_ddl_task( &self, - req: SubmitDdlTaskRequest, + context: ExecutorContext, + request: SubmitDdlTaskRequest, ) -> Result { - let res = self - .procedure_client()? - .submit_ddl_task(req.try_into().context(ConvertMetaRequestSnafu)?) + let event_context = procedure_event_context(&context); + let actor = procedure_actor(&context); + let mut query_context = context.query_context.context(MissingQueryContextSnafu)?; + query_context + .extensions + .remove(CREATE_DATABASE_CREATOR_EXTENSION_KEY); + if let DdlTask::CreateDatabase(CreateDatabaseTask { + creator: Some(creator), + .. + }) = &request.task + { + query_context.extensions.insert( + CREATE_DATABASE_CREATOR_EXTENSION_KEY.to_string(), + serde_json::to_string(creator).context(ConvertMetaConfigSnafu)?, + ); + } + + let mut request: api::v1::meta::DdlTaskRequest = + request.try_into().context(ConvertMetaRequestSnafu)?; + request.query_context = Some(api::v1::QueryContext::from(query_context)); + request.event_context = event_context; + request.actor = actor; + + self.procedure_client()? + .submit_ddl_task(request) .await? .try_into() - .context(ConvertMetaResponseSnafu)?; - - Ok(res) + .context(ConvertMetaResponseSnafu) } pub fn heartbeat_client(&self) -> Result { diff --git a/src/meta-client/src/client/procedure.rs b/src/meta-client/src/client/procedure.rs index d7d5ccfaff..ccf2c28188 100644 --- a/src/meta-client/src/client/procedure.rs +++ b/src/meta-client/src/client/procedure.rs @@ -20,11 +20,13 @@ use api::v1::meta::ddl_task_request::Task; use api::v1::meta::procedure_service_client::ProcedureServiceClient; use api::v1::meta::{ DdlTaskRequest, DdlTaskResponse, GcRegionsRequest, GcRegionsResponse, GcTableRequest, - GcTableResponse, MigrateRegionRequest, MigrateRegionResponse, ProcedureDetailRequest, - ProcedureDetailResponse, ProcedureId, ProcedureStateResponse, QueryProcedureRequest, - ReconcileRequest, ReconcileResponse, RequestHeader, ResponseHeader, Role, + GcTableResponse, MigrateRegionRequest, MigrateRegionResponse, ProcedureActor, + ProcedureDetailRequest, ProcedureDetailResponse, ProcedureEventContext, ProcedureId, + ProcedureStateResponse, QueryProcedureRequest, ReconcileRequest, ReconcileResponse, + RequestHeader, ResponseHeader, Role, }; use common_grpc::channel_manager::ChannelManager; +use common_meta::procedure_executor::ExecutorContext; use common_meta::rpc::ddl::{ CREATE_DATABASE_CREATOR_EXTENSION_KEY, CREATE_DATABASE_CREATOR_METADATA_KEY, }; @@ -43,6 +45,33 @@ use crate::client::{Id, LeaderProviderRef, util}; use crate::error; use crate::error::Result; +/// Builds the event context transported by a procedure RPC. +/// +/// The caller can only supply reason/extensions. Protocol is derived here from +/// the trusted, typed query channel held locally in the executor context. +pub(crate) fn procedure_event_context(context: &ExecutorContext) -> Option { + context.event_input.as_ref().map(|input| { + let mut event_context = ProcedureEventContext::from(input); + event_context.protocol = context + .query_context + .as_ref() + .and_then(|query_context| query_context.protocol()) + .unwrap_or_default(); + event_context + }) +} + +/// Builds the optional procedure actor transported by a procedure RPC. +pub(crate) fn procedure_actor(context: &ExecutorContext) -> Option { + context + .actor + .as_deref() + .filter(|username| !username.is_empty()) + .map(|username| ProcedureActor { + username: username.to_string(), + }) +} + #[derive(Clone, Debug)] pub struct Client { inner: Arc>, @@ -92,6 +121,7 @@ impl Client { /// - `timeout`: timeout for downgrading region and upgrading region operations pub async fn migrate_region( &self, + context: &ExecutorContext, region_id: u64, from_peer: u64, to_peer: u64, @@ -99,7 +129,7 @@ impl Client { ) -> Result { let inner = self.inner.read().await; inner - .migrate_region(region_id, from_peer, to_peer, timeout) + .migrate_region(context, region_id, from_peer, to_peer, timeout) .await } @@ -114,14 +144,22 @@ impl Client { inner.list_procedures().await } - pub async fn gc_regions(&self, request: MetaGcRegionsRequest) -> Result { + pub async fn gc_regions( + &self, + context: &ExecutorContext, + request: MetaGcRegionsRequest, + ) -> Result { let inner = self.inner.read().await; - inner.gc_regions(request).await + inner.gc_regions(context, request).await } - pub async fn gc_table(&self, request: MetaGcTableRequest) -> Result { + pub async fn gc_table( + &self, + context: &ExecutorContext, + request: MetaGcTableRequest, + ) -> Result { let inner = self.inner.read().await; - inner.gc_table(request).await + inner.gc_table(context, request).await } } @@ -225,6 +263,7 @@ impl Inner { async fn migrate_region( &self, + context: &ExecutorContext, region_id: u64, from_peer: u64, to_peer: u64, @@ -235,6 +274,8 @@ impl Inner { from_peer, to_peer, timeout_secs: timeout.as_secs() as u32, + event_context: procedure_event_context(context), + actor: procedure_actor(context), ..Default::default() }; @@ -278,7 +319,11 @@ impl Inner { .await } - async fn gc_regions(&self, request: MetaGcRegionsRequest) -> Result { + async fn gc_regions( + &self, + context: &ExecutorContext, + request: MetaGcRegionsRequest, + ) -> Result { let timeout = request.timeout; let req = GcRegionsRequest { header: Some(RequestHeader { @@ -290,6 +335,8 @@ impl Inner { region_ids: request.region_ids, full_file_listing: request.full_file_listing, timeout_secs: gc_timeout_secs(timeout), + event_context: procedure_event_context(context), + actor: procedure_actor(context), }; let resp: GcRegionsResponse = self @@ -315,7 +362,11 @@ impl Inner { }) } - async fn gc_table(&self, request: MetaGcTableRequest) -> Result { + async fn gc_table( + &self, + context: &ExecutorContext, + request: MetaGcTableRequest, + ) -> Result { let timeout = request.timeout; let req = GcTableRequest { header: Some(RequestHeader { @@ -329,6 +380,8 @@ impl Inner { table_name: request.table_name, full_file_listing: request.full_file_listing, timeout_secs: gc_timeout_secs(timeout), + event_context: procedure_event_context(context), + actor: procedure_actor(context), }; let resp: GcTableResponse = self @@ -460,11 +513,14 @@ mod tests { ReconcileResponse, ResponseHeader, Role, }; use async_trait::async_trait; + use common_base::protocol::Channel; use common_error::status_code::StatusCode; + use common_event_recorder::{PersistentEventContext, ProcedureEventInput}; + use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor}; use common_meta::rpc::ddl::{ CREATE_DATABASE_CREATOR_EXTENSION_KEY, CREATE_DATABASE_CREATOR_METADATA_KEY, CommentObjectType, CommentOnTask, CreatorGrantIntent, DdlTask, QueryContext, - SubmitDdlTaskRequest, + SubmitDdlTaskRequest, TriggerReason, }; use common_telemetry::common_error::ext::ErrorExt; use common_telemetry::info; @@ -474,8 +530,8 @@ mod tests { use tonic::codec::CompressionEncoding; use tonic::{Request, Response, Status}; - use super::gc_timeout_secs; use crate::client::MetaClientBuilder; + use crate::client::procedure::{gc_timeout_secs, procedure_actor, procedure_event_context}; #[test] fn test_gc_timeout_secs() { @@ -486,6 +542,57 @@ mod tests { assert_eq!(gc_timeout_secs(Some(Duration::from_secs(10))), 10); } + #[test] + fn test_procedure_event_context_derives_protocol_from_query_context() { + let context = ExecutorContext { + query_context: Some(QueryContext { + channel: Channel::Postgres as u8, + ..Default::default() + }), + event_input: Some(ProcedureEventInput::new(TriggerReason::Manual)), + ..Default::default() + }; + + assert_eq!( + procedure_event_context(&context).map(PersistentEventContext::from), + Some(PersistentEventContext::new(TriggerReason::Manual).with_protocol("postgres")) + ); + + let automatic_context = ExecutorContext { + event_input: Some(ProcedureEventInput::new(TriggerReason::ScheduledGc)), + ..Default::default() + }; + assert_eq!( + procedure_event_context(&automatic_context).map(PersistentEventContext::from), + Some(PersistentEventContext::new(TriggerReason::ScheduledGc)) + ); + + let unknown_channel_context = ExecutorContext { + query_context: Some(QueryContext::default()), + event_input: Some(ProcedureEventInput::new(TriggerReason::Manual)), + ..Default::default() + }; + assert_eq!( + procedure_event_context(&unknown_channel_context).map(PersistentEventContext::from), + Some(PersistentEventContext::new(TriggerReason::Manual)) + ); + } + + #[test] + fn test_procedure_actor() { + let context = ExecutorContext { + actor: Some(String::new()), + ..Default::default() + }; + assert!(procedure_actor(&context).is_none()); + + let context = ExecutorContext { + actor: Some("alice".to_string()), + ..Default::default() + }; + assert_eq!(procedure_actor(&context).unwrap().username, "alice"); + } + #[derive(Clone)] struct MockHeartbeat { leader_addr: String, @@ -621,19 +728,28 @@ mod tests { username: "alice".to_string(), created_at_ns: 42, }; - client - .submit_ddl_task(SubmitDdlTaskRequest::new( - QueryContext::default(), - DdlTask::new_create_database( - "greptime".to_string(), - "metrics".to_string(), - false, - Default::default(), - Some(creator.clone()), - ), - )) - .await - .unwrap(); + let executor_context = |actor: String| ExecutorContext { + query_context: Some(QueryContext { + channel: Channel::Postgres as u8, + ..Default::default() + }), + actor: Some(actor), + event_input: Some(ProcedureEventInput::new(TriggerReason::Manual)), + ..Default::default() + }; + ProcedureExecutor::submit_ddl_task( + &client, + executor_context("effective-user".to_string()), + SubmitDdlTaskRequest::new(DdlTask::new_create_database( + "greptime".to_string(), + "metrics".to_string(), + false, + Default::default(), + Some(creator.clone()), + )), + ) + .await + .unwrap(); let request = request_rx.recv().await.unwrap(); let encoded = serde_json::to_string(&creator).unwrap(); @@ -647,9 +763,36 @@ mod tests { .as_ref(), encoded.as_bytes() ); - let extensions = &request.into_inner().query_context.unwrap().extensions; + let request = request.into_inner(); + assert_eq!(request.actor.unwrap().username, "effective-user"); + assert_eq!( + PersistentEventContext::from(request.event_context.unwrap()), + PersistentEventContext::new(TriggerReason::Manual).with_protocol("postgres") + ); + let extensions = &request.query_context.unwrap().extensions; assert_eq!(extensions[CREATE_DATABASE_CREATOR_EXTENSION_KEY], encoded); + ProcedureExecutor::submit_ddl_task( + &client, + executor_context(String::new()), + SubmitDdlTaskRequest::new(DdlTask::new_drop_database( + "greptime".to_string(), + "metrics".to_string(), + false, + )), + ) + .await + .unwrap(); + assert!( + request_rx + .recv() + .await + .unwrap() + .into_inner() + .actor + .is_none() + ); + server_handle.abort(); } @@ -689,22 +832,28 @@ mod tests { .build(); client.start(&[addr_str.as_str()]).await.unwrap(); - let mut request = SubmitDdlTaskRequest::new( - QueryContext::default(), - DdlTask::new_comment_on(CommentOnTask { - catalog_name: "greptime".to_string(), - schema_name: "public".to_string(), - object_type: CommentObjectType::Table, - object_name: "test_table".to_string(), - column_name: None, - object_id: None, - comment: Some("timeout".to_string()), - }), - ); + let mut request = SubmitDdlTaskRequest::new(DdlTask::new_comment_on(CommentOnTask { + catalog_name: "greptime".to_string(), + schema_name: "public".to_string(), + object_type: CommentObjectType::Table, + object_name: "test_table".to_string(), + column_name: None, + object_id: None, + comment: Some("timeout".to_string()), + })); request.timeout = Duration::from_secs(1); let now = Instant::now(); - let err = client.submit_ddl_task(request).await.unwrap_err(); + let err = client + .submit_ddl_task( + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + request, + ) + .await + .unwrap_err(); let elapsed = now.elapsed(); // The request should be cancelled within 1 second. assert!(elapsed < Duration::from_secs(2)); diff --git a/src/meta-client/src/error.rs b/src/meta-client/src/error.rs index 5121bba847..4623cfee33 100644 --- a/src/meta-client/src/error.rs +++ b/src/meta-client/src/error.rs @@ -67,6 +67,12 @@ pub enum Error { location: Location, }, + #[snafu(display("Procedure submission requires query context"))] + MissingQueryContext { + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to send heartbeat: {}", err_msg))] SendHeartbeat { err_msg: String, @@ -141,6 +147,7 @@ impl ErrorExt for Error { | Error::NoLeader { .. } | Error::AskLeaderTimeout { .. } | Error::NotStarted { .. } + | Error::MissingQueryContext { .. } | Error::SendHeartbeat { .. } | Error::CreateHeartbeatStream { .. } | Error::CreateChannel { .. } diff --git a/src/meta-srv/src/gc/ctx.rs b/src/meta-srv/src/gc/ctx.rs index c9cc804986..4b46a941e2 100644 --- a/src/meta-srv/src/gc/ctx.rs +++ b/src/meta-srv/src/gc/ctx.rs @@ -29,10 +29,9 @@ use common_meta::key::DroppedTableName; use common_meta::key::TableMetadataManagerRef; use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; -use common_meta::rpc::ddl::PersistentEventContext; #[cfg(feature = "enterprise")] use common_meta::rpc::ddl::PurgeDroppedTableTask; -use common_procedure::{ProcedureManagerRef, ProcedureWithId, watcher}; +use common_procedure::{ProcedureContext, ProcedureManagerRef, ProcedureWithId, watcher}; use common_telemetry::debug; use snafu::{OptionExt as _, ResultExt as _}; use store_api::storage::{GcReport, RegionId}; @@ -68,7 +67,7 @@ pub(crate) trait SchedulerCtx: Send + Sync { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result; #[cfg(feature = "enterprise")] @@ -256,14 +255,14 @@ impl SchedulerCtx for DefaultGcSchedulerCtx { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { self.gc_regions_inner( region_ids, full_file_listing, timeout, region_routes_override, - event_context, + procedure_context, ) .await } @@ -310,7 +309,7 @@ impl DefaultGcSchedulerCtx { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { debug!( "Sending GC instruction for {} regions (full_file_listing: {})", @@ -328,7 +327,7 @@ impl DefaultGcSchedulerCtx { region_routes_override, ); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); let id = procedure_with_id.id; diff --git a/src/meta-srv/src/gc/handler.rs b/src/meta-srv/src/gc/handler.rs index 16b29a39c7..eda87fbb14 100644 --- a/src/meta-srv/src/gc/handler.rs +++ b/src/meta-srv/src/gc/handler.rs @@ -18,7 +18,7 @@ use std::time::Instant; use common_catalog::consts::MITO_ENGINE; use common_meta::datanode::{RegionManifestInfo, RegionStat}; use common_meta::peer::Peer; -use common_meta::rpc::ddl::PersistentEventContext; +use common_procedure::ProcedureContext; use common_telemetry::tracing::Instrument as _; use common_telemetry::{debug, error, info, warn}; use futures::StreamExt; @@ -39,7 +39,7 @@ use crate::metrics::METRIC_META_GC_CANDIDATE_REGIONS; impl GcScheduler { pub(crate) async fn trigger_gc( &self, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let start_time = Instant::now(); info!("Starting GC cycle"); @@ -111,7 +111,7 @@ impl GcScheduler { datanode_to_candidates, dropped_assignment.force_full_listing, dropped_assignment.region_routes_override, - event_context, + procedure_context, ) .instrument(common_telemetry::tracing::info_span!( "meta_gc_dispatch_to_datanodes" @@ -215,7 +215,7 @@ impl GcScheduler { datanode_to_candidates: HashMap>, force_full_listing_by_peer: HashMap>, region_routes_override_by_peer: HashMap, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> GcJobReport { let mut per_datanode_reports = HashMap::new(); let mut failed_datanodes: HashMap<_, Vec<_>> = HashMap::new(); @@ -237,7 +237,7 @@ impl GcScheduler { .get(&peer) .cloned() .unwrap_or_default(); - let event_context = event_context.clone(); + let procedure_context = procedure_context.clone(); async move { ( peer, @@ -247,7 +247,7 @@ impl GcScheduler { candidates, force_full_listing, region_routes_override, - event_context, + procedure_context, ) .await, ) @@ -286,7 +286,7 @@ impl GcScheduler { candidates: Vec<(TableId, GcCandidate)>, force_full_listing: HashSet, region_routes_override: Region2Peers, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { info!( "Starting GC for datanode {} with {} candidate regions", @@ -336,7 +336,7 @@ impl GcScheduler { false, self.config.mailbox_timeout, region_routes_override.clone(), - event_context.clone(), + procedure_context.clone(), ) .instrument(common_telemetry::tracing::info_span!( "meta_gc_call_datanode", @@ -369,7 +369,7 @@ impl GcScheduler { true, self.config.mailbox_timeout, region_routes_override, - event_context, + procedure_context, ) .instrument(common_telemetry::tracing::info_span!( "meta_gc_call_datanode", diff --git a/src/meta-srv/src/gc/mock.rs b/src/meta-srv/src/gc/mock.rs index 94749dfb46..dcf913462c 100644 --- a/src/meta-srv/src/gc/mock.rs +++ b/src/meta-srv/src/gc/mock.rs @@ -26,12 +26,13 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use common_catalog::consts::MITO_ENGINE; +use common_event_recorder::PersistentEventContext; use common_meta::datanode::{RegionManifestInfo, RegionStat}; use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; use common_meta::peer::Peer; -use common_meta::rpc::ddl::PersistentEventContext; use common_meta::rpc::router::{Region, RegionRoute}; +use common_procedure::ProcedureContext; use common_telemetry::debug; use ordered_float::OrderedFloat; use store_api::region_engine::RegionRole; @@ -214,10 +215,13 @@ impl SchedulerCtx for MockSchedulerCtx { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { *self.gc_regions_calls.lock().unwrap() += 1; - self.gc_event_contexts.lock().unwrap().push(event_context); + self.gc_event_contexts + .lock() + .unwrap() + .push(procedure_context.event_context.unwrap_or_default()); // Check per-region error injection first (for any region) for ®ion_id in region_ids { diff --git a/src/meta-srv/src/gc/mock/basic.rs b/src/meta-srv/src/gc/mock/basic.rs index ed051cc839..70f7a10b1a 100644 --- a/src/meta-srv/src/gc/mock/basic.rs +++ b/src/meta-srv/src/gc/mock/basic.rs @@ -16,8 +16,10 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex}; use std::time::Instant; +use common_event_recorder::PersistentEventContext; use common_meta::peer::Peer; -use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; +use common_meta::rpc::ddl::TriggerReason; +use common_procedure::ProcedureContext; use common_telemetry::init_default_ut_logging; use store_api::region_engine::RegionRole; use store_api::storage::{FileId, FileRefsManifest, GcReport, RegionId}; @@ -36,7 +38,7 @@ async fn test_parallel_process_datanodes_empty() { HashMap::new(), HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; @@ -105,7 +107,7 @@ async fn test_parallel_process_datanodes_with_candidates() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; @@ -245,7 +247,18 @@ async fn test_handle_manual_gc_without_regions_records_manual_event_context() { last_tracker_cleanup: Arc::new(tokio::sync::Mutex::new(Instant::now())), }; - scheduler.handle_manual_gc(None, None, None).await.unwrap(); + scheduler + .handle_manual_gc( + None, + None, + None, + ProcedureContext { + event_context: Some(PersistentEventContext::new(TriggerReason::Manual)), + ..Default::default() + }, + ) + .await + .unwrap(); assert_eq!( ctx.gc_event_contexts.lock().unwrap().as_slice(), diff --git a/src/meta-srv/src/gc/mock/concurrent.rs b/src/meta-srv/src/gc/mock/concurrent.rs index 5bce9a1be8..bfb7ab1d26 100644 --- a/src/meta-srv/src/gc/mock/concurrent.rs +++ b/src/meta-srv/src/gc/mock/concurrent.rs @@ -18,8 +18,8 @@ use std::time::{Duration, Instant}; use common_meta::key::table_route::PhysicalTableRouteValue; use common_meta::peer::Peer; -use common_meta::rpc::ddl::PersistentEventContext; use common_meta::rpc::router::{Region, RegionRoute}; +use common_procedure::ProcedureContext; use common_telemetry::{info, init_default_ut_logging}; use store_api::region_engine::RegionRole; use store_api::storage::{FileId, FileRefsManifest, GcReport, RegionId}; @@ -102,7 +102,7 @@ async fn test_concurrent_table_processing_limits() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; @@ -190,7 +190,7 @@ async fn test_datanode_processes_tables_with_partial_gc_failures() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; @@ -296,7 +296,7 @@ async fn test_region_gc_concurrency_limit() { candidates.into_iter().map(|c| (table_id, c)).collect(), HashSet::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap(); @@ -413,7 +413,7 @@ async fn test_region_gc_concurrency_with_partial_failures() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; @@ -557,7 +557,7 @@ async fn test_region_gc_concurrency_with_retryable_errors() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; diff --git a/src/meta-srv/src/gc/mock/full_list.rs b/src/meta-srv/src/gc/mock/full_list.rs index 8193a9094e..843e45134d 100644 --- a/src/meta-srv/src/gc/mock/full_list.rs +++ b/src/meta-srv/src/gc/mock/full_list.rs @@ -17,7 +17,7 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use common_meta::peer::Peer; -use common_meta::rpc::ddl::PersistentEventContext; +use common_procedure::ProcedureContext; use common_telemetry::init_default_ut_logging; use store_api::region_engine::RegionRole; use store_api::storage::{FileId, FileRefsManifest, GcReport, RegionId}; @@ -83,7 +83,7 @@ async fn test_full_file_listing_first_time_gc() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap(); @@ -157,7 +157,7 @@ async fn test_full_file_listing_interval_enforcement() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap(); @@ -184,7 +184,7 @@ async fn test_full_file_listing_interval_enforcement() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap(); @@ -260,7 +260,7 @@ async fn test_full_file_listing_no_interval_passed() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap(); @@ -284,7 +284,7 @@ async fn test_full_file_listing_no_interval_passed() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap(); diff --git a/src/meta-srv/src/gc/mock/misc.rs b/src/meta-srv/src/gc/mock/misc.rs index 575a5c5fc0..fabf5bcae4 100644 --- a/src/meta-srv/src/gc/mock/misc.rs +++ b/src/meta-srv/src/gc/mock/misc.rs @@ -17,7 +17,7 @@ use std::sync::{Arc, Mutex}; use std::time::Instant; use common_meta::peer::Peer; -use common_meta::rpc::ddl::PersistentEventContext; +use common_procedure::ProcedureContext; use common_telemetry::init_default_ut_logging; use store_api::storage::{FileRefsManifest, GcReport, RegionId}; @@ -76,7 +76,7 @@ async fn test_empty_file_refs_manifest() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; @@ -169,7 +169,7 @@ async fn test_multiple_regions_per_table() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; diff --git a/src/meta-srv/src/gc/scheduler.rs b/src/meta-srv/src/gc/scheduler.rs index b89c0d06be..857d4deb57 100644 --- a/src/meta-srv/src/gc/scheduler.rs +++ b/src/meta-srv/src/gc/scheduler.rs @@ -16,9 +16,11 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::time::{Duration, Instant}; +use common_event_recorder::PersistentEventContext; use common_meta::DatanodeId; use common_meta::key::runtime_switch::RuntimeSwitchManagerRef; -use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; +use common_meta::rpc::ddl::TriggerReason; +use common_procedure::ProcedureContext; use common_telemetry::tracing::Instrument as _; use common_telemetry::{error, info}; use snafu::ResultExt; @@ -104,6 +106,7 @@ pub enum Event { full_file_listing: Option, /// Optional override for timeout. If None, uses scheduler config. timeout: Option, + procedure_context: ProcedureContext, }, } @@ -182,12 +185,13 @@ impl GcScheduler { region_ids, full_file_listing, timeout, + procedure_context, } => { info!("Received manually gc request"); let span = common_telemetry::tracing::info_span!("meta_gc_tick", trigger = "manual"); let result = self - .handle_manual_gc(region_ids, full_file_listing, timeout) + .handle_manual_gc(region_ids, full_file_listing, timeout, procedure_context) .instrument(span) .await; if let Err(e) = &result { @@ -218,7 +222,10 @@ impl GcScheduler { } let span = common_telemetry::tracing::info_span!("meta_gc_handle_tick"); let report = self - .trigger_gc(PersistentEventContext::new(TriggerReason::ScheduledGc)) + .trigger_gc(ProcedureContext { + actor: None, + event_context: Some(PersistentEventContext::new(TriggerReason::ScheduledGc)), + }) .instrument(span) .await?; @@ -294,6 +301,7 @@ impl GcScheduler { region_ids: Option>, full_file_listing: Option, timeout: Option, + procedure_context: ProcedureContext, ) -> Result { info!("Start to handle manual gc request"); @@ -304,9 +312,7 @@ impl GcScheduler { // No specific regions, use default tick behavior let Some(regions) = region_ids else { - let report = self - .trigger_gc(PersistentEventContext::new(TriggerReason::Manual)) - .await?; + let report = self.trigger_gc(procedure_context).await?; info!("Finished manual gc request"); return Ok(report); }; @@ -353,7 +359,7 @@ impl GcScheduler { full_listing, gc_timeout, Region2Peers::new(), - PersistentEventContext::new(TriggerReason::Manual), + procedure_context.clone(), ) .await?; combined_report.merge(report); @@ -367,7 +373,7 @@ impl GcScheduler { true, gc_timeout, dropped_routes_override, - PersistentEventContext::new(TriggerReason::Manual), + procedure_context, ) .await?; combined_report.merge(report); @@ -490,7 +496,7 @@ mod tests { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - _event_context: PersistentEventContext, + _procedure_context: ProcedureContext, ) -> Result { self.gc_regions_calls.fetch_add(1, Ordering::Relaxed); panic!("gc_regions should not be called in maintenance mode") @@ -564,7 +570,7 @@ mod tests { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - _event_context: PersistentEventContext, + _procedure_context: ProcedureContext, ) -> Result { Ok(GcReport::default()) } @@ -861,7 +867,10 @@ mod tests { *ctx.dropped_tables.lock().unwrap() = vec![dropped_table(1, Some(i64::MIN))]; let scheduler = soft_drop_scheduler(ctx.clone()); - scheduler.handle_manual_gc(None, None, None).await.unwrap(); + scheduler + .handle_manual_gc(None, None, None, ProcedureContext::default()) + .await + .unwrap(); assert!(ctx.purge_attempts.lock().unwrap().is_empty()); } @@ -898,7 +907,7 @@ mod tests { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - _event_context: PersistentEventContext, + _procedure_context: ProcedureContext, ) -> Result { crate::error::UnexpectedSnafu { violated: "mock gc failure".to_string(), @@ -949,6 +958,7 @@ mod tests { Some(vec![RegionId::new(1, 0)]), Some(false), Some(Duration::from_secs(1)), + ProcedureContext::default(), ) .await; @@ -977,6 +987,7 @@ mod tests { Some(vec![RegionId::new(1, 0)]), Some(false), Some(Duration::from_secs(1)), + ProcedureContext::default(), ) .await; diff --git a/src/meta-srv/src/procedure/region_migration.rs b/src/meta-srv/src/procedure/region_migration.rs index ad1553c92d..75ab5c7939 100644 --- a/src/meta-srv/src/procedure/region_migration.rs +++ b/src/meta-srv/src/procedure/region_migration.rs @@ -32,7 +32,7 @@ use std::fmt::{Debug, Display}; use std::time::Duration; use common_error::ext::BoxedError; -use common_event_recorder::Event; +use common_event_recorder::{Event, PersistentEventContext}; use common_meta::cache_invalidator::CacheInvalidatorRef; use common_meta::ddl::RegionFailureDetectorControllerRef; use common_meta::instruction::CacheIdent; @@ -45,7 +45,7 @@ use common_meta::kv_backend::{KvBackendRef, ResettableKvBackendRef}; use common_meta::lock_key::{CatalogLock, RegionLock, SchemaLock, TableLock}; use common_meta::peer::Peer; use common_meta::region_keeper::{MemoryRegionKeeperRef, OperatingRegionGuard}; -use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; +use common_meta::rpc::ddl::TriggerReason; use common_procedure::error::{ Error as ProcedureError, FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu, }; diff --git a/src/meta-srv/src/procedure/region_migration/manager.rs b/src/meta-srv/src/procedure/region_migration/manager.rs index e1d719d974..c21de50624 100644 --- a/src/meta-srv/src/procedure/region_migration/manager.rs +++ b/src/meta-srv/src/procedure/region_migration/manager.rs @@ -18,12 +18,14 @@ use std::fmt::Display; use std::sync::{Arc, RwLock}; use std::time::Duration; +use common_event_recorder::PersistentEventContext; use common_meta::key::table_info::TableInfoValue; use common_meta::key::table_route::TableRouteValue; use common_meta::peer::Peer; -use common_meta::rpc::ddl::PersistentEventContext; use common_meta::rpc::router::RegionRoute; -use common_procedure::{ProcedureId, ProcedureManagerRef, ProcedureWithId, watcher}; +use common_procedure::{ + ProcedureContext, ProcedureId, ProcedureManagerRef, ProcedureWithId, watcher, +}; use common_telemetry::{error, info, warn}; use serde::{Deserialize, Serialize}; use snafu::{OptionExt, ResultExt, ensure}; @@ -458,8 +460,20 @@ impl RegionMigrationManager { } let submitting_region_ids = task.region_ids.clone(); + // TODO(weny): Pass the actor when region migration task submission propagates it. + let procedure_context = ProcedureContext { + actor: None, + event_context: Some(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + }; let procedure_id = self - .submit_procedure_inner(task, procedure_guards, catalog_and_schema) + .submit_procedure_inner( + procedure_context, + task, + procedure_guards, + catalog_and_schema, + ) .await?; Ok(SubmitRegionMigrationTaskResult { migrated, @@ -475,11 +489,11 @@ impl RegionMigrationManager { async fn submit_procedure_inner( &self, + procedure_context: ProcedureContext, task: RegionMigrationTaskBatch, procedure_guards: Vec, catalog_and_schema: Vec<(String, String)>, ) -> Result { - let event_context = PersistentEventContext::new(task.trigger_reason.to_trigger_reason()); let procedure = RegionMigrationProcedure::new( PersistentContext::new( catalog_and_schema, @@ -493,7 +507,7 @@ impl RegionMigrationManager { procedure_guards, ); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); let procedure_id = procedure_with_id.id; info!("Starting region migration procedure {procedure_id} for {task}"); let procedure_manager = self.procedure_manager.clone(); @@ -529,8 +543,14 @@ impl RegionMigrationManager { /// Submits a new region migration procedure. pub async fn submit_procedure( &self, + procedure_context: ProcedureContext, mut task: RegionMigrationProcedureTask, ) -> Result> { + if let Some(event_context) = procedure_context.event_context.as_ref() { + task.trigger_reason = + RegionMigrationTriggerReason::from_trigger_reason(event_context.reason); + } + let Some(guard) = self.insert_running_procedure(&task) else { return error::MigrationRunningSnafu { region_id: task.region_id, @@ -577,7 +597,6 @@ impl RegionMigrationManager { timeout, trigger_reason, } = task.clone(); - let event_context = PersistentEventContext::new(trigger_reason.to_trigger_reason()); let procedure = RegionMigrationProcedure::new( PersistentContext::new( vec![(catalog_name, schema_name)], @@ -591,7 +610,7 @@ impl RegionMigrationManager { vec![guard], ); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); let procedure_id = procedure_with_id.id; info!("Starting region migration procedure {procedure_id} for {task}"); let procedure_manager = self.procedure_manager.clone(); @@ -655,7 +674,15 @@ mod test { .unwrap() .insert(region_id, task.clone()); - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::MigrationRunning { .. }); } @@ -673,7 +700,15 @@ mod test { trigger_reason: RegionMigrationTriggerReason::Manual, }; - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::InvalidArguments { .. }); } @@ -691,7 +726,15 @@ mod test { trigger_reason: RegionMigrationTriggerReason::Manual, }; - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::TableRouteNotFound { .. }); } @@ -719,7 +762,15 @@ mod test { env.create_physical_table_metadata(table_info, region_routes) .await; - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::RegionRouteNotFound { .. }); } @@ -747,7 +798,15 @@ mod test { env.create_physical_table_metadata(table_info, region_routes) .await; - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::LeaderPeerChanged { .. }); assert_eq!( err.to_string(), @@ -780,7 +839,15 @@ mod test { env.create_physical_table_metadata(table_info, region_routes) .await; - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::InvalidArguments { .. }); assert_eq!( err.to_string(), @@ -813,7 +880,15 @@ mod test { env.create_physical_table_metadata(table_info, region_routes) .await; - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::RegionMigrated { .. }); } diff --git a/src/meta-srv/src/service/procedure.rs b/src/meta-srv/src/service/procedure.rs index 1b5c4c42db..9e226bc18f 100644 --- a/src/meta-srv/src/service/procedure.rs +++ b/src/meta-srv/src/service/procedure.rs @@ -18,10 +18,12 @@ use api::v1::meta::reconcile_request::Target; use api::v1::meta::{ DdlTaskRequest as PbDdlTaskRequest, DdlTaskResponse as PbDdlTaskResponse, GcRegionsRequest, GcRegionsResponse, GcStats, GcTableRequest, GcTableResponse, MigrateRegionRequest, - MigrateRegionResponse, ProcedureDetailRequest, ProcedureDetailResponse, ProcedureStateResponse, + MigrateRegionResponse, ProcedureActor, ProcedureDetailRequest, ProcedureDetailResponse, + ProcedureEventContext as PbProcedureEventContext, ProcedureStateResponse, QueryProcedureRequest, ReconcileCatalog, ReconcileDatabase, ReconcileRequest, ReconcileResponse, ReconcileTable, ResolveStrategy, Role, procedure_service_server, }; +use common_event_recorder::{PersistentEventContext, ProcedureEventInput}; use common_meta::key::TableMetadataManagerRef; use common_meta::key::table_name::TableNameKey; use common_meta::procedure_executor::ExecutorContext; @@ -33,6 +35,7 @@ use common_meta::rpc::procedure::{ self, GcRegionsRequest as MetaGcRegionsRequest, GcResponse, GcTableRequest as MetaGcTableRequest, }; +use common_procedure::ProcedureContext; use snafu::{OptionExt, ResultExt}; use store_api::storage::RegionId; use table::table_reference::TableReference; @@ -47,6 +50,45 @@ use crate::procedure::region_migration::manager::{ use crate::service::GrpcResult; use crate::{check_leader, error, gc}; +struct ProcedureSubmission { + actor: Option, + event_context: Option, +} + +impl TryFrom<(Option, Option)> for ProcedureSubmission { + type Error = Status; + + fn try_from( + (actor, event_context): (Option, Option), + ) -> Result { + let actor = actor + .map(|actor| { + if actor.username.is_empty() { + Err(Status::invalid_argument( + "Procedure actor must not be empty", + )) + } else { + Ok(actor.username) + } + }) + .transpose()?; + + Ok(Self { + actor, + event_context, + }) + } +} + +impl From for ProcedureContext { + fn from(context: ProcedureSubmission) -> Self { + Self { + actor: context.actor, + event_context: context.event_context.map(PersistentEventContext::from), + } + } +} + #[async_trait::async_trait] impl procedure_service_server::ProcedureService for Metasrv { async fn query( @@ -89,9 +131,15 @@ impl procedure_service_server::ProcedureService for Metasrv { task, wait, timeout_secs, + event_context, + actor, } = request; let header = header.context(error::MissingRequestHeaderSnafu)?; + let ProcedureSubmission { + actor, + event_context, + } = ProcedureSubmission::try_from((actor, event_context))?; let mut query_context = query_context .context(error::MissingRequiredParameterSnafu { param: "query_context", @@ -102,15 +150,17 @@ impl procedure_service_server::ProcedureService for Metasrv { .try_into() .context(error::ConvertProtoDataSnafu)?; restore_create_database_creator(&metadata, header.role, &mut task, &mut query_context)?; + let executor_context = ExecutorContext { + tracing_context: Some(header.tracing_context), + query_context: Some(query_context), + actor, + event_input: event_context.map(ProcedureEventInput::from), + }; let resp = self .ddl_manager() .submit_ddl_task( - &ExecutorContext { - tracing_context: Some(header.tracing_context), - ..Default::default() - }, + executor_context, SubmitDdlTaskRequest { - query_context, wait, timeout: Duration::from_secs(timeout_secs.into()), task, @@ -135,9 +185,13 @@ impl procedure_service_server::ProcedureService for Metasrv { from_peer, to_peer, timeout_secs, + event_context, + actor, } = request.into_inner(); let _header = header.context(error::MissingRequestHeaderSnafu)?; + let procedure_context = + ProcedureContext::from(ProcedureSubmission::try_from((actor, event_context))?); let from_peer = self .lookup_datanode_peer(from_peer) .await? @@ -149,13 +203,16 @@ impl procedure_service_server::ProcedureService for Metasrv { let pid = self .region_migration_manager() - .submit_procedure(RegionMigrationProcedureTask { - region_id: region_id.into(), - from_peer, - to_peer, - timeout: Duration::from_secs(timeout_secs.into()), - trigger_reason: RegionMigrationTriggerReason::Manual, - }) + .submit_procedure( + procedure_context, + RegionMigrationProcedureTask { + region_id: region_id.into(), + from_peer, + to_peer, + timeout: Duration::from_secs(timeout_secs.into()), + trigger_reason: RegionMigrationTriggerReason::Manual, + }, + ) .await? .map(procedure::pid_to_pb_pid); @@ -266,16 +323,23 @@ impl procedure_service_server::ProcedureService for Metasrv { region_ids, full_file_listing, timeout_secs, + event_context, + actor, } = request.into_inner(); let _header = header.context(error::MissingRequestHeaderSnafu)?; + let procedure_context = + ProcedureContext::from(ProcedureSubmission::try_from((actor, event_context))?); let response = self - .handle_gc_regions(MetaGcRegionsRequest { - region_ids, - full_file_listing, - timeout: Self::normalize_gc_timeout(Duration::from_secs(timeout_secs as u64)), - }) + .handle_gc_regions( + procedure_context, + MetaGcRegionsRequest { + region_ids, + full_file_listing, + timeout: Self::normalize_gc_timeout(Duration::from_secs(timeout_secs as u64)), + }, + ) .await?; Ok(Response::new(gc_response_to_regions_pb(response))) @@ -291,18 +355,25 @@ impl procedure_service_server::ProcedureService for Metasrv { table_name, full_file_listing, timeout_secs, + event_context, + actor, } = request.into_inner(); let _header = header.context(error::MissingRequestHeaderSnafu)?; + let procedure_context = + ProcedureContext::from(ProcedureSubmission::try_from((actor, event_context))?); let response = self - .handle_gc_table(MetaGcTableRequest { - catalog_name, - schema_name, - table_name, - full_file_listing, - timeout: Self::normalize_gc_timeout(Duration::from_secs(timeout_secs as u64)), - }) + .handle_gc_table( + procedure_context, + MetaGcTableRequest { + catalog_name, + schema_name, + table_name, + full_file_listing, + timeout: Self::normalize_gc_timeout(Duration::from_secs(timeout_secs as u64)), + }, + ) .await?; Ok(Response::new(gc_response_to_table_pb(response))) @@ -371,17 +442,30 @@ impl Metasrv { } } - async fn handle_gc_regions(&self, request: MetaGcRegionsRequest) -> error::Result { + async fn handle_gc_regions( + &self, + procedure_context: ProcedureContext, + request: MetaGcRegionsRequest, + ) -> error::Result { let region_ids: Vec = request .region_ids .into_iter() .map(RegionId::from_u64) .collect(); - self.trigger_gc_for_regions(region_ids, request.full_file_listing, request.timeout) - .await + self.trigger_gc_for_regions( + procedure_context, + region_ids, + request.full_file_listing, + request.timeout, + ) + .await } - async fn handle_gc_table(&self, request: MetaGcTableRequest) -> error::Result { + async fn handle_gc_table( + &self, + procedure_context: ProcedureContext, + request: MetaGcTableRequest, + ) -> error::Result { let table_name_key = TableNameKey::new( &request.catalog_name, &request.schema_name, @@ -406,13 +490,19 @@ impl Metasrv { .context(TableMetadataManagerSnafu)?; let region_ids: Vec = route.region_routes.iter().map(|r| r.region.id).collect(); - self.trigger_gc_for_regions(region_ids, request.full_file_listing, request.timeout) - .await + self.trigger_gc_for_regions( + procedure_context, + region_ids, + request.full_file_listing, + request.timeout, + ) + .await } /// Triggers manual GC for specified regions and returns the GC response. async fn trigger_gc_for_regions( &self, + procedure_context: ProcedureContext, region_ids: Vec, full_file_listing: bool, timeout: Option, @@ -429,6 +519,7 @@ impl Metasrv { region_ids: Some(region_ids), full_file_listing: Some(full_file_listing), timeout, + procedure_context, }) .await .map_err(|_| { @@ -502,14 +593,15 @@ fn gc_response_to_table_pb(resp: GcResponse) -> GcTableResponse { mod tests { use std::time::Duration; - use api::v1::meta::Role; + use api::v1::meta::{ProcedureActor, ProcedureEventContext, Role}; use common_meta::rpc::ddl::{ CREATE_DATABASE_CREATOR_EXTENSION_KEY, CREATE_DATABASE_CREATOR_METADATA_KEY, CreatorGrantIntent, DdlTask, QueryContext, }; + use common_procedure::ProcedureContext; use tonic::metadata::{MetadataMap, MetadataValue}; - use super::{Metasrv, restore_create_database_creator}; + use super::{Metasrv, ProcedureSubmission, restore_create_database_creator}; fn create_database_task() -> DdlTask { DdlTask::new_create_database( @@ -646,4 +738,41 @@ mod tests { assert_eq!(task.creator, Some(creator)); assert!(trusted.extensions.is_empty()); } + + #[test] + fn test_procedure_submission_rejects_empty_actor() { + let submission = ProcedureSubmission::try_from(( + Some(ProcedureActor { + username: "alice".to_string(), + }), + Some(ProcedureEventContext { + reason: "manual".to_string(), + protocol: "postgres".to_string(), + extensions: Default::default(), + }), + )) + .unwrap(); + assert_eq!(submission.actor.as_deref(), Some("alice")); + assert_eq!( + submission + .event_context + .as_ref() + .map(|context| context.protocol.as_str()), + Some("postgres") + ); + + assert!( + ProcedureSubmission::try_from(( + Some(ProcedureActor { + username: String::new(), + }), + None, + )) + .is_err() + ); + + let procedure_context = + ProcedureContext::from(ProcedureSubmission::try_from((None, None)).unwrap()); + assert_eq!(procedure_context.event_context, None); + } } diff --git a/src/operator/src/procedure.rs b/src/operator/src/procedure.rs index 1a849afaf7..03fb4f909c 100644 --- a/src/operator/src/procedure.rs +++ b/src/operator/src/procedure.rs @@ -27,11 +27,12 @@ use common_meta::rpc::procedure::{ }; use common_query::error as query_error; use common_query::error::Result as QueryResult; +use session::context::QueryContextRef; use snafu::ResultExt; use table::table_name::TableName; use crate::error; -use crate::utils::to_meta_query_context_with_trigger_reason; +use crate::utils::to_executor_context; /// The operator for procedures which implements [`ProcedureServiceHandler`]. #[derive(Clone)] @@ -59,8 +60,8 @@ impl ProcedureServiceOperator { impl ProcedureServiceHandler for ProcedureServiceOperator { async fn purge_table( &self, + query_ctx: QueryContextRef, table_name: TableName, - query_ctx: session::context::QueryContextRef, ) -> QueryResult<()> { let dropped = self .table_metadata_manager @@ -76,24 +77,25 @@ impl ProcedureServiceHandler for ProcedureServiceOperator { }) .map_err(BoxedError::new) .context(query_error::ProcedureServiceSnafu)?; - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_ctx, TriggerReason::Manual); - let request = SubmitDdlTaskRequest::new( - meta_query_context, - DdlTask::new_purge_dropped_table(dropped.table_id), - ); + let executor_context = to_executor_context(query_ctx, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_purge_dropped_table(dropped.table_id)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .map_err(BoxedError::new) .context(query_error::ProcedureServiceSnafu)?; Ok(()) } - async fn migrate_region(&self, request: MigrateRegionRequest) -> QueryResult> { + async fn migrate_region( + &self, + query_ctx: QueryContextRef, + request: MigrateRegionRequest, + ) -> QueryResult> { + let executor_context = to_executor_context(query_ctx, TriggerReason::Manual); Ok(self .procedure_executor - .migrate_region(&ExecutorContext::default(), request) + .migrate_region(&executor_context, request) .await .map_err(BoxedError::new) .context(query_error::ProcedureServiceSnafu)? @@ -135,17 +137,27 @@ impl ProcedureServiceHandler for ProcedureServiceOperator { &self.catalog_manager } - async fn gc_regions(&self, request: MetaGcRegionsRequest) -> QueryResult { + async fn gc_regions( + &self, + query_ctx: QueryContextRef, + request: MetaGcRegionsRequest, + ) -> QueryResult { + let executor_context = to_executor_context(query_ctx, TriggerReason::Manual); self.procedure_executor - .gc_regions(&ExecutorContext::default(), request) + .gc_regions(&executor_context, request) .await .map_err(BoxedError::new) .context(query_error::ProcedureServiceSnafu) } - async fn gc_table(&self, request: MetaGcTableRequest) -> QueryResult { + async fn gc_table( + &self, + query_ctx: QueryContextRef, + request: MetaGcTableRequest, + ) -> QueryResult { + let executor_context = to_executor_context(query_ctx, TriggerReason::Manual); self.procedure_executor - .gc_table(&ExecutorContext::default(), request) + .gc_table(&executor_context, request) .await .map_err(BoxedError::new) .context(query_error::ProcedureServiceSnafu) @@ -162,10 +174,10 @@ mod tests { use common_meta::key::table_route::TableRouteValue; use common_meta::key::test_utils::new_test_table_info_with_name; use common_meta::kv_backend::memory::MemoryKvBackend; - use common_meta::procedure_executor::{ProcedureExecutor, ProcedureExecutorRef}; - use common_meta::rpc::ddl::{ - DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse, TRIGGER_REASON_EXTENSION_KEY, + use common_meta::procedure_executor::{ + ExecutorContext, ProcedureExecutor, ProcedureExecutorRef, }; + use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse}; use common_meta::rpc::procedure::{MigrateRegionResponse, ProcedureStateResponse}; use session::context::QueryContextBuilder; use table::table_name::TableName; @@ -175,6 +187,7 @@ mod tests { #[derive(Default)] struct RecordingProcedureExecutor { requests: Mutex>, + contexts: Mutex>, fail: bool, } @@ -182,9 +195,10 @@ mod tests { impl ProcedureExecutor for RecordingProcedureExecutor { async fn submit_ddl_task( &self, - _: &ExecutorContext, + context: ExecutorContext, request: SubmitDdlTaskRequest, ) -> common_meta::error::Result { + self.contexts.lock().unwrap().push(context); self.requests.lock().unwrap().push(request); if self.fail { return common_meta::error::UnsupportedSnafu { @@ -273,20 +287,19 @@ mod tests { .build() .into(); - operator.purge_table(name.clone(), query_ctx).await.unwrap(); + operator.purge_table(query_ctx, name.clone()).await.unwrap(); let requests = executor.requests.lock().unwrap(); assert!( matches!(&requests[0].task, DdlTask::PurgeDroppedTable(task) if task.table_id == 42) ); - assert_eq!(requests[0].query_context.current_catalog, "catalog"); - assert_eq!(requests[0].query_context.current_schema, "schema"); + let contexts = executor.contexts.lock().unwrap(); + let query_context = contexts[0].query_context.as_ref().unwrap(); + assert_eq!(query_context.current_catalog, "catalog"); + assert_eq!(query_context.current_schema, "schema"); assert_eq!( - Some(&TriggerReason::Manual.as_ref().to_string()), - requests[0] - .query_context - .extensions - .get(TRIGGER_REASON_EXTENSION_KEY) + contexts[0].event_input.as_ref().map(|input| input.reason), + Some(TriggerReason::Manual) ); } @@ -303,12 +316,12 @@ mod tests { ); let error = operator .purge_table( - TableName::new("catalog", "schema", "missing"), QueryContextBuilder::default() .current_catalog("catalog".to_string()) .current_schema("schema".to_string()) .build() .into(), + TableName::new("catalog", "schema", "missing"), ) .await .unwrap_err(); @@ -326,12 +339,12 @@ mod tests { assert!( operator .purge_table( - name, QueryContextBuilder::default() .current_catalog("catalog".to_string()) .current_schema("schema".to_string()) .build() - .into() + .into(), + name, ) .await .is_err() diff --git a/src/operator/src/statement/comment.rs b/src/operator/src/statement/comment.rs index 824287fc80..c55e1be9f7 100644 --- a/src/operator/src/statement/comment.rs +++ b/src/operator/src/statement/comment.rs @@ -15,7 +15,6 @@ use api::v1::CommentOnExpr; use common_error::ext::BoxedError; use common_meta::cache_invalidator::Context; -use common_meta::procedure_executor::ExecutorContext; use common_meta::rpc::ddl::{ CommentObjectType, CommentOnTask, DdlTask, SubmitDdlTaskRequest, TriggerReason, }; @@ -30,7 +29,7 @@ use crate::error::{ self, ExecuteDdlSnafu, ExternalSnafu, InvalidSqlSnafu, Result, TableMetadataManagerSnafu, }; use crate::statement::StatementExecutor; -use crate::utils::to_meta_query_context_with_trigger_reason; +use crate::utils::to_executor_context; impl StatementExecutor { /// Adds a comment to a database object (table, column, or flow). @@ -54,13 +53,11 @@ impl StatementExecutor { .context(TableMetadataManagerSnafu)?; let cache_idents = comment_on_task.cache_idents(); - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_ctx, TriggerReason::Manual); - let request = - SubmitDdlTaskRequest::new(meta_query_context, DdlTask::new_comment_on(comment_on_task)); + let executor_context = to_executor_context(query_ctx, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_comment_on(comment_on_task)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(ExecuteDdlSnafu)?; @@ -88,13 +85,11 @@ impl StatementExecutor { .context(TableMetadataManagerSnafu)?; let cache_idents = comment_on_task.cache_idents(); - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_ctx, TriggerReason::Manual); - let request = - SubmitDdlTaskRequest::new(meta_query_context, DdlTask::new_comment_on(comment_on_task)); + let executor_context = to_executor_context(query_ctx, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_comment_on(comment_on_task)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(ExecuteDdlSnafu)?; diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index 256b849123..3c10186ca4 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -47,7 +47,6 @@ use common_meta::instruction::CacheIdent; #[cfg(feature = "enterprise")] use common_meta::key::TableMetadataManagerRef; use common_meta::key::schema_name::{SchemaName, SchemaNameKey}; -use common_meta::procedure_executor::ExecutorContext; #[cfg(feature = "enterprise")] use common_meta::procedure_executor::ProcedureExecutorRef; #[cfg(feature = "enterprise")] @@ -119,10 +118,7 @@ use crate::error::{ use crate::expr_helper::{self, RepartitionRequest, RepartitionSource}; use crate::statement::StatementExecutor; use crate::statement::show::create_partitions_stmt; -use crate::utils::{ - to_meta_query_context_with_origin_frontend_and_trigger_reason, - to_meta_query_context_with_trigger_reason, -}; +use crate::utils::{to_executor_context, to_executor_context_with_origin_frontend}; #[derive(Debug, Clone, Copy)] struct DdlSubmitOptions { @@ -729,13 +725,11 @@ impl StatementExecutor { }) .context(error::InvalidExprSnafu)?; - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_create_trigger(task), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_create_trigger(task)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -793,13 +787,11 @@ impl StatementExecutor { create_flow: Some(expr), }) .context(error::InvalidExprSnafu)?; - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_create_flow(task), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_create_flow(task)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -1109,14 +1101,12 @@ impl StatementExecutor { table_type: TableType::View, }; - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(ctx, TriggerReason::Manual), - DdlTask::new_create_view(expr, view_info.clone()), - ); + let executor_context = to_executor_context(ctx, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_create_view(expr, view_info.clone())); let resp = self .procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu)?; @@ -1195,13 +1185,11 @@ impl StatementExecutor { expr: DropFlowTask, query_context: QueryContextRef, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_drop_flow(expr), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_drop_flow(expr)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -1230,13 +1218,11 @@ impl StatementExecutor { expr: DropTriggerTask, query_context: QueryContextRef, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_drop_trigger(expr), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_drop_trigger(expr)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -1311,13 +1297,11 @@ impl StatementExecutor { expr: DropViewTask, query_context: QueryContextRef, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_drop_view(expr), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_drop_view(expr)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -1780,17 +1764,13 @@ impl StatementExecutor { source: Some(source), ..Default::default() }; - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_context.clone(), TriggerReason::Manual); - let mut req = SubmitDdlTaskRequest::new( - meta_query_context, - DdlTask::new_alter_table(AlterTableExpr { - catalog_name: request.catalog_name.clone(), - schema_name: request.schema_name.clone(), - table_name: request.table_name.clone(), - kind: Some(Kind::Repartition(repartition)), - }), - ); + let executor_context = to_executor_context(query_context.clone(), TriggerReason::Manual); + let mut req = SubmitDdlTaskRequest::new(DdlTask::new_alter_table(AlterTableExpr { + catalog_name: request.catalog_name.clone(), + schema_name: request.schema_name.clone(), + table_name: request.table_name.clone(), + kind: Some(Kind::Repartition(repartition)), + })); req.wait = ddl_options.wait; req.timeout = ddl_options.timeout; @@ -1806,7 +1786,7 @@ impl StatementExecutor { let response = self .procedure_executor - .submit_ddl_task(&ExecutorContext::default(), req) + .submit_ddl_task(executor_context, req) .await .context(error::ExecuteDdlSnafu)?; @@ -1888,12 +1868,11 @@ impl StatementExecutor { .await .context(TableMetadataManagerSnafu)?; - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_context, trigger_reason); + let executor_context = to_executor_context(query_context, trigger_reason); let (req, invalidate_keys) = if physical_table_id == table_id { // This is physical table - let req = SubmitDdlTaskRequest::new(meta_query_context, DdlTask::new_alter_table(expr)); + let req = SubmitDdlTaskRequest::new(DdlTask::new_alter_table(expr)); let invalidate_keys = vec![ CacheIdent::TableId(table_id), @@ -1903,10 +1882,7 @@ impl StatementExecutor { (req, invalidate_keys) } else { // This is logical table - let req = SubmitDdlTaskRequest::new( - meta_query_context, - DdlTask::new_alter_logical_tables(vec![expr]), - ); + let req = SubmitDdlTaskRequest::new(DdlTask::new_alter_logical_tables(vec![expr])); let mut invalidate_keys = vec![ CacheIdent::TableId(physical_table_id), @@ -1934,7 +1910,7 @@ impl StatementExecutor { }; self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), req) + .submit_ddl_task(executor_context, req) .await .context(error::ExecuteDdlSnafu)?; @@ -2025,17 +2001,19 @@ impl StatementExecutor { .map(|expr| expr.as_pb_partition().context(PartitionExprToPbSnafu)) .collect::>>()?; - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_origin_frontend_and_trigger_reason( - query_context, - &self.origin_frontend_addr, - trigger_reason, - ), - DdlTask::new_create_table(create_table, partitions, table_info), + let executor_context = to_executor_context_with_origin_frontend( + query_context, + &self.origin_frontend_addr, + trigger_reason, ); + let request = SubmitDdlTaskRequest::new(DdlTask::new_create_table( + create_table, + partitions, + table_info, + )); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2046,17 +2024,15 @@ impl StatementExecutor { query_context: QueryContextRef, trigger_reason: TriggerReason, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_origin_frontend_and_trigger_reason( - query_context, - &self.origin_frontend_addr, - trigger_reason, - ), - DdlTask::new_create_logical_tables(tables_data), + let executor_context = to_executor_context_with_origin_frontend( + query_context, + &self.origin_frontend_addr, + trigger_reason, ); + let request = SubmitDdlTaskRequest::new(DdlTask::new_create_logical_tables(tables_data)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2067,15 +2043,11 @@ impl StatementExecutor { query_context: QueryContextRef, trigger_reason: TriggerReason, ) -> Result { - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_context, trigger_reason); - let request = SubmitDdlTaskRequest::new( - meta_query_context, - DdlTask::new_alter_logical_tables(tables_data), - ); + let executor_context = to_executor_context(query_context, trigger_reason); + let request = SubmitDdlTaskRequest::new(DdlTask::new_alter_logical_tables(tables_data)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2087,19 +2059,17 @@ impl StatementExecutor { drop_if_exists: bool, query_context: QueryContextRef, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_drop_table( - table_name.catalog_name.clone(), - table_name.schema_name.clone(), - table_name.table_name.clone(), - table_id, - drop_if_exists, - ), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_drop_table( + table_name.catalog_name.clone(), + table_name.schema_name.clone(), + table_name.table_name.clone(), + table_id, + drop_if_exists, + )); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2111,13 +2081,12 @@ impl StatementExecutor { drop_if_exists: bool, query_context: QueryContextRef, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_drop_database(catalog, schema, drop_if_exists), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = + SubmitDdlTaskRequest::new(DdlTask::new_drop_database(catalog, schema, drop_if_exists)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2127,13 +2096,11 @@ impl StatementExecutor { alter_expr: AlterDatabaseExpr, query_context: QueryContextRef, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_alter_database(alter_expr), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_alter_database(alter_expr)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2145,21 +2112,17 @@ impl StatementExecutor { time_ranges: Vec<(Timestamp, Timestamp)>, query_context: QueryContextRef, ) -> Result { - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual); - let request = SubmitDdlTaskRequest::new( - meta_query_context, - DdlTask::new_truncate_table( - table_name.catalog_name.clone(), - table_name.schema_name.clone(), - table_name.table_name.clone(), - table_id, - time_ranges, - ), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_truncate_table( + table_name.catalog_name.clone(), + table_name.schema_name.clone(), + table_name.table_name.clone(), + table_id, + time_ranges, + )); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2239,13 +2202,17 @@ impl StatementExecutor { query_context: QueryContextRef, creator: Option, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_create_database(catalog, database, create_if_not_exists, options, creator), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_create_database( + catalog, + database, + create_if_not_exists, + options, + creator, + )); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2876,14 +2843,10 @@ async fn execute_undrop_table( table_name: table_name.to_string(), })?; - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual); - let request = SubmitDdlTaskRequest::new( - meta_query_context, - DdlTask::new_undrop_table(dropped.table_id), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_undrop_table(dropped.table_id)); procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu)?; @@ -2931,9 +2894,7 @@ mod test { ExecutorContext, ProcedureExecutor, ProcedureExecutorRef, }; #[cfg(feature = "enterprise")] - use common_meta::rpc::ddl::{ - DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse, TRIGGER_REASON_EXTENSION_KEY, - }; + use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse}; #[cfg(feature = "enterprise")] use common_meta::rpc::procedure::{ MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse, @@ -2951,6 +2912,7 @@ mod test { #[derive(Default)] struct RecordingProcedureExecutor { requests: Mutex>, + contexts: Mutex>, fail: bool, } @@ -2959,9 +2921,10 @@ mod test { impl ProcedureExecutor for RecordingProcedureExecutor { async fn submit_ddl_task( &self, - _ctx: &ExecutorContext, + ctx: ExecutorContext, request: SubmitDdlTaskRequest, ) -> common_meta::error::Result { + self.contexts.lock().unwrap().push(ctx); self.requests.lock().unwrap().push(request); if self.fail { return common_meta::error::UnsupportedSnafu { @@ -3078,12 +3041,10 @@ mod test { let requests = procedure.requests.lock().unwrap(); assert!(matches!(&requests[0].task, DdlTask::UndropTable(task) if task.table_id == 42)); + let contexts = procedure.contexts.lock().unwrap(); assert_eq!( - requests[0] - .query_context - .extensions - .get(TRIGGER_REASON_EXTENSION_KEY), - Some(&TriggerReason::Manual.as_ref().to_string()) + contexts[0].event_input.as_ref().map(|input| input.reason), + Some(TriggerReason::Manual) ); assert_eq!( cache.invalidations.lock().unwrap()[0], diff --git a/src/operator/src/utils.rs b/src/operator/src/utils.rs index 03649713b0..160d712072 100644 --- a/src/operator/src/utils.rs +++ b/src/operator/src/utils.rs @@ -14,9 +14,9 @@ use std::sync::{Arc, RwLock}; -use common_meta::rpc::ddl::{ - ORIGIN_FRONTEND_ADDR_EXTENSION_KEY, TRIGGER_REASON_EXTENSION_KEY, TriggerReason, -}; +use common_event_recorder::ProcedureEventInput; +use common_meta::procedure_executor::ExecutorContext; +use common_meta::rpc::ddl::{ORIGIN_FRONTEND_ADDR_EXTENSION_KEY, TriggerReason}; use common_time::Timezone; use session::context::{QueryContextBuilder, QueryContextRef}; use snafu::ResultExt; @@ -37,31 +37,34 @@ pub fn to_meta_query_context( } } -/// Converts a session query context and records the DDL trigger reason. -pub fn to_meta_query_context_with_trigger_reason( +/// Builds the execution envelope for a frontend-originated procedure. +pub(crate) fn to_executor_context( query_context: QueryContextRef, trigger_reason: TriggerReason, -) -> common_meta::rpc::ddl::QueryContext { - let mut meta_query_context = to_meta_query_context(query_context); - meta_query_context.extensions.insert( - TRIGGER_REASON_EXTENSION_KEY.to_string(), - trigger_reason.as_ref().to_string(), - ); - meta_query_context +) -> ExecutorContext { + let actor = query_context.current_user().username().to_string(); + ExecutorContext { + query_context: Some(to_meta_query_context(query_context)), + actor: Some(actor), + event_input: Some(ProcedureEventInput::new(trigger_reason)), + ..Default::default() + } } -pub fn to_meta_query_context_with_origin_frontend_and_trigger_reason( +/// Builds a frontend-originated procedure envelope with the trusted frontend origin. +pub(crate) fn to_executor_context_with_origin_frontend( query_context: QueryContextRef, origin_frontend_addr: &str, trigger_reason: TriggerReason, -) -> common_meta::rpc::ddl::QueryContext { - let mut meta_query_context = - to_meta_query_context_with_trigger_reason(query_context, trigger_reason); - meta_query_context.extensions.insert( - ORIGIN_FRONTEND_ADDR_EXTENSION_KEY.to_string(), - origin_frontend_addr.to_string(), - ); - meta_query_context +) -> ExecutorContext { + let mut executor_context = to_executor_context(query_context, trigger_reason); + if let Some(query_context) = &mut executor_context.query_context { + query_context.extensions.insert( + ORIGIN_FRONTEND_ADDR_EXTENSION_KEY.to_string(), + origin_frontend_addr.to_string(), + ); + } + executor_context } pub fn try_to_session_query_context( @@ -87,14 +90,12 @@ mod tests { use std::collections::HashMap; use std::sync::{Arc, RwLock}; - use common_meta::rpc::ddl::{ - ORIGIN_FRONTEND_ADDR_EXTENSION_KEY, TRIGGER_REASON_EXTENSION_KEY, TriggerReason, - }; + use common_meta::rpc::ddl::{ORIGIN_FRONTEND_ADDR_EXTENSION_KEY, TriggerReason}; use common_time::Timezone; use session::context::QueryContextBuilder; use super::{ - to_meta_query_context, to_meta_query_context_with_origin_frontend_and_trigger_reason, + to_executor_context_with_origin_frontend, to_meta_query_context, try_to_session_query_context, }; @@ -122,7 +123,7 @@ mod tests { } #[test] - fn test_meta_query_context_with_origin_frontend_overrides_reserved_key() { + fn test_executor_context_with_origin_frontend_overrides_reserved_key() { let session_ctx = Arc::new( QueryContextBuilder::default() .set_extension( @@ -132,11 +133,12 @@ mod tests { .build(), ); - let meta_ctx = to_meta_query_context_with_origin_frontend_and_trigger_reason( + let executor_context = to_executor_context_with_origin_frontend( session_ctx, "127.0.0.1:4000", TriggerReason::Manual, ); + let meta_ctx = executor_context.query_context.as_ref().unwrap(); assert_eq!( meta_ctx @@ -146,11 +148,8 @@ mod tests { Some("127.0.0.1:4000") ); assert_eq!( - meta_ctx - .extensions - .get(TRIGGER_REASON_EXTENSION_KEY) - .map(String::as_str), - Some("manual") + executor_context.event_input.unwrap().reason, + TriggerReason::Manual ); } } diff --git a/tests-integration/tests/gc_event.rs b/tests-integration/tests/gc_event.rs index 9ecedc7777..54a0787d12 100644 --- a/tests-integration/tests/gc_event.rs +++ b/tests-integration/tests/gc_event.rs @@ -15,8 +15,8 @@ use std::sync::Arc; use std::time::Duration; -use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; -use common_procedure::{ProcedureId, ProcedureWithId, watcher}; +use common_event_recorder::{PersistentEventContext, TriggerReason}; +use common_procedure::{ProcedureContext, ProcedureId, ProcedureWithId, watcher}; use common_test_util::temp_dir::create_temp_dir; use meta_srv::gc::{BatchGcProcedure, GcSchedulerOptions}; use mito2::gc::GcConfig; @@ -137,9 +137,9 @@ async fn test_batch_gc_event() { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), - context: common_procedure::ProcedureContext::from_event_context( - PersistentEventContext::new(TriggerReason::Manual), - ), + context: ProcedureContext::from_event_context(PersistentEventContext::new( + TriggerReason::Manual, + )), }) .await .unwrap(); diff --git a/tests-integration/tests/region_migration.rs b/tests-integration/tests/region_migration.rs index 3d27dabc3f..d67180249c 100644 --- a/tests-integration/tests/region_migration.rs +++ b/tests-integration/tests/region_migration.rs @@ -23,10 +23,11 @@ use common_event_recorder::event_table::{ }; use common_event_recorder::{ DEFAULT_EVENTS_TABLE_NAME, DEFAULT_FLUSH_INTERVAL_SECONDS, EVENTS_TABLE_TIMESTAMP_COLUMN_NAME, - EVENTS_TABLE_TYPE_COLUMN_NAME, + EVENTS_TABLE_TYPE_COLUMN_NAME, PersistentEventContext, TriggerReason, }; use common_meta::key::{RegionDistribution, RegionRoleSet, TableMetadataManagerRef}; use common_meta::peer::Peer; +use common_procedure::ProcedureContext; use common_procedure::event::{ EVENTS_TABLE_PROCEDURE_ID_COLUMN_NAME, EVENTS_TABLE_PROCEDURE_STATE_COLUMN_NAME, }; @@ -183,13 +184,18 @@ pub async fn test_region_migration(store_type: StorageType, endpoints: Vec