refactor(event): separate procedure submission context (#8856)

* refactor(event): separate procedure submission context

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix(event): map extensions and forward GC context

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix(gc): initialize integration test context

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix(test): pass procedure context to DDL helpers

Signed-off-by: WenyXu <wenymedia@gmail.com>

* refactor: simplify procedure submission contexts

Signed-off-by: WenyXu <wenymedia@gmail.com>

* refactor(event): separate procedure and query contexts

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix(event): clarify procedure context propagation

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix(event): preserve procedure submission context

Signed-off-by: WenyXu <wenymedia@gmail.com>

* refactor(event): move DDL context by value

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix(test): retain manual GC event context

Signed-off-by: WenyXu <wenymedia@gmail.com>

* refactor(event): tighten procedure context API

Signed-off-by: WenyXu <wenymedia@gmail.com>

* chore: update greptime-proto

Signed-off-by: WenyXu <wenymedia@gmail.com>

---------

Signed-off-by: WenyXu <wenymedia@gmail.com>
This commit is contained in:
Weny Xu
2026-08-13 07:29:00 +00:00
committed by GitHub
parent d0fecdd6b0
commit 1af4c33524
41 changed files with 1487 additions and 770 deletions
Generated
+2 -1
View File
@@ -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",
+1 -1
View File
@@ -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"
+119
View File
@@ -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<String, String>,
}
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<PbProcedureEventContext> 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<String>)> for PersistentEventContext {
fn from((input, protocol): (ProcedureEventInput, Option<String>)) -> 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<PbProcedureEventContext> 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();
+1 -1
View File
@@ -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::*;
+35 -17
View File
@@ -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<Value> {
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<Option<String>> {
async fn migrate_region(
&self,
_query_ctx: QueryContextRef,
_request: MigrateRegionRequest,
) -> Result<Option<String>> {
unreachable!()
}
@@ -297,12 +307,20 @@ mod tests {
unreachable!()
}
async fn gc_regions(&self, request: GcRegionsRequest) -> Result<GcResponse> {
async fn gc_regions(
&self,
_query_ctx: QueryContextRef,
request: GcRegionsRequest,
) -> Result<GcResponse> {
*self.gc_regions_request.lock().unwrap() = Some(request);
Ok(GcResponse::default())
}
async fn gc_table(&self, request: GcTableRequest) -> Result<GcResponse> {
async fn gc_table(
&self,
_query_ctx: QueryContextRef,
request: GcTableRequest,
) -> Result<GcResponse> {
*self.gc_table_request.lock().unwrap() = Some(request);
Ok(GcResponse::default())
}
@@ -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<Value> {
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 {
+9 -5
View File
@@ -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<Option<String>> {
async fn migrate_region(
&self,
_: QueryContextRef,
_: MigrateRegionRequest,
) -> Result<Option<String>> {
unreachable!()
}
async fn reconcile(&self, _: ReconcileRequest) -> Result<Option<String>> {
@@ -159,10 +163,10 @@ mod tests {
fn catalog_manager(&self) -> &CatalogManagerRef {
unreachable!()
}
async fn gc_regions(&self, _: GcRegionsRequest) -> Result<GcResponse> {
async fn gc_regions(&self, _: QueryContextRef, _: GcRegionsRequest) -> Result<GcResponse> {
unreachable!()
}
async fn gc_table(&self, _: GcTableRequest) -> Result<GcResponse> {
async fn gc_table(&self, _: QueryContextRef, _: GcTableRequest) -> Result<GcResponse> {
unreachable!()
}
}
+16 -4
View File
@@ -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<Option<String>>;
async fn migrate_region(
&self,
query_ctx: QueryContextRef,
request: MigrateRegionRequest,
) -> Result<Option<String>>;
/// Reconcile a table, database or catalog, returns the procedure id if success.
async fn reconcile(&self, request: ReconcileRequest) -> Result<Option<String>>;
@@ -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<MetaGcResponse>;
async fn gc_regions(
&self,
query_ctx: QueryContextRef,
request: MetaGcRegionsRequest,
) -> Result<MetaGcResponse>;
/// Manually trigger GC for a table.
async fn gc_table(&self, request: MetaGcTableRequest) -> Result<MetaGcResponse>;
async fn gc_table(
&self,
query_ctx: QueryContextRef,
request: MetaGcTableRequest,
) -> Result<MetaGcResponse>;
}
/// This flow service handler is only use for flush flow for now.
+12 -3
View File
@@ -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<Option<String>> {
Ok(Some("test_pid".to_string()))
@@ -92,7 +93,11 @@ impl FunctionState {
Ok(())
}
async fn gc_regions(&self, _request: GcRegionsRequest) -> Result<GcResponse> {
async fn gc_regions(
&self,
_context: QueryContextRef,
_request: GcRegionsRequest,
) -> Result<GcResponse> {
Ok(GcResponse {
processed_regions: 1,
need_retry_regions: vec![],
@@ -101,7 +106,11 @@ impl FunctionState {
})
}
async fn gc_table(&self, _request: GcTableRequest) -> Result<GcResponse> {
async fn gc_table(
&self,
_context: QueryContextRef,
_request: GcTableRequest,
) -> Result<GcResponse> {
Ok(GcResponse {
processed_regions: 1,
need_retry_regions: vec![],
+3 -9
View File
@@ -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(
+3 -9
View File
@@ -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());
+275 -113
View File
@@ -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<SubmitDdlTaskResponse>;
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<SubmitDdlTaskResponse>;
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<Output>)> {
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<Output>)> {
// 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<Output>)> {
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<Output>)> {
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<CreateTableTask>,
physical_table_id: TableId,
event_context: PersistentEventContext,
procedure_context: ProcedureContext,
) -> Result<(ProcedureId, Option<Output>)> {
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<AlterTableTask>,
physical_table_id: TableId,
event_context: PersistentEventContext,
procedure_context: ProcedureContext,
) -> Result<(ProcedureId, Option<Output>)> {
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<Output>)> {
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<Output>)> {
#[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<Output>)> {
#[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<Output>)> {
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<Output>)> {
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<Output>)> {
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<Output>)> {
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<Output>)> {
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<Output>)> {
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<TableInfoValue>,
event_context: PersistentEventContext,
procedure_context: ProcedureContext,
) -> Result<(ProcedureId, Option<Output>)> {
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<Output>)> {
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<CreateTableTask>,
event_context: PersistentEventContext,
procedure_context: ProcedureContext,
) -> Result<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<AlterTableTask>,
event_context: PersistentEventContext,
procedure_context: ProcedureContext,
) -> Result<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<Vec<ProcedureContext>>,
}
#[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<crate::rpc::ddl::SubmitDdlTaskResponse> {
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<crate::rpc::ddl::SubmitDdlTaskResponse> {
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();
+11 -5
View File
@@ -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<W3cTrace>,
/// 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<QueryContext>,
pub actor: Option<String>,
pub event_context: Option<PersistentEventContext>,
/// Caller-supplied event metadata. Protocol is derived by the submission adapter.
pub event_input: Option<ProcedureEventInput>,
}
/// 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<SubmitDdlTaskResponse>;
@@ -129,7 +135,7 @@ impl LocalProcedureExecutor {
impl ProcedureExecutor for LocalProcedureExecutor {
async fn submit_ddl_task(
&self,
ctx: &ExecutorContext,
ctx: ExecutorContext,
request: SubmitDdlTaskRequest,
) -> Result<SubmitDdlTaskResponse> {
self.ddl_manager.submit_ddl_task(ctx, request).await
+13 -133
View File
@@ -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<Task> 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<SubmitDdlTaskRequest> for PbDdlTaskRequest {
fn try_from(request: SubmitDdlTaskRequest) -> Result<Self> {
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<SubmitDdlTaskRequest> 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<u64, u64>,
}
/// 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<String> {
let channel = Channel::from(u32::from(self.channel));
(channel != Channel::Unknown).then(|| channel.as_ref().to_string())
}
pub fn snapshot_seqs(&self) -> &HashMap<u64, u64> {
&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")
);
}
}
+71 -12
View File
@@ -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)
);
}
+8 -4
View File
@@ -35,6 +35,9 @@ pub type Output = Arc<dyn Any + Send + Sync>;
/// 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<String>,
/// Context describing why and how the procedure was submitted.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub event_context: Option<PersistentEventContext>,
@@ -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
}
}
+8
View File
@@ -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::<ProcedureMessage>(&json).unwrap(),
message
);
let legacy: ProcedureMessage = serde_json::from_str(
r#"{"type_name":"TestMessage","data":"legacy","parent_id":null,"step":1}"#,
)
+2 -2
View File
@@ -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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
self.submitted.lock().unwrap().push(request.task.clone());
+1
View File
@@ -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"] }
+57 -23
View File
@@ -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<SubmitDdlTaskResponse> {
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<MigrateRegionResponse> {
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<GcResponse> {
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<GcResponse> {
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<MigrateRegionResponse> {
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<GcResponse> {
self.procedure_client()?.gc_regions(request).await
pub async fn gc_regions(
&self,
context: &ExecutorContext,
request: GcRegionsRequest,
) -> Result<GcResponse> {
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<GcResponse> {
self.procedure_client()?.gc_table(request).await
pub async fn gc_table(
&self,
context: &ExecutorContext,
request: GcTableRequest,
) -> Result<GcResponse> {
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<SubmitDdlTaskResponse> {
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<HeartbeatClient> {
+188 -39
View File
@@ -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<ProcedureEventContext> {
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<ProcedureActor> {
context
.actor
.as_deref()
.filter(|username| !username.is_empty())
.map(|username| ProcedureActor {
username: username.to_string(),
})
}
#[derive(Clone, Debug)]
pub struct Client {
inner: Arc<RwLock<Inner>>,
@@ -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<MigrateRegionResponse> {
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<MetaGcResponse> {
pub async fn gc_regions(
&self,
context: &ExecutorContext,
request: MetaGcRegionsRequest,
) -> Result<MetaGcResponse> {
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<MetaGcResponse> {
pub async fn gc_table(
&self,
context: &ExecutorContext,
request: MetaGcTableRequest,
) -> Result<MetaGcResponse> {
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<MetaGcResponse> {
async fn gc_regions(
&self,
context: &ExecutorContext,
request: MetaGcRegionsRequest,
) -> Result<MetaGcResponse> {
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<MetaGcResponse> {
async fn gc_table(
&self,
context: &ExecutorContext,
request: MetaGcTableRequest,
) -> Result<MetaGcResponse> {
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));
+7
View File
@@ -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 { .. }
+6 -7
View File
@@ -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<GcReport>;
#[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<GcReport> {
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<GcReport> {
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;
+9 -9
View File
@@ -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<GcJobReport> {
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<Peer, Vec<(TableId, GcCandidate)>>,
force_full_listing_by_peer: HashMap<Peer, HashSet<RegionId>>,
region_routes_override_by_peer: HashMap<Peer, Region2Peers>,
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<RegionId>,
region_routes_override: Region2Peers,
event_context: PersistentEventContext,
procedure_context: ProcedureContext,
) -> Result<GcReport> {
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",
+7 -3
View File
@@ -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<GcReport> {
*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 &region_id in region_ids {
+17 -4
View File
@@ -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(),
+6 -6
View File
@@ -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;
+6 -6
View File
@@ -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();
+3 -3
View File
@@ -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;
+23 -12
View File
@@ -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<bool>,
/// Optional override for timeout. If None, uses scheduler config.
timeout: Option<Duration>,
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<Vec<RegionId>>,
full_file_listing: Option<bool>,
timeout: Option<Duration>,
procedure_context: ProcedureContext,
) -> Result<GcJobReport> {
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<GcReport> {
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<GcReport> {
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<GcReport> {
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;
@@ -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,
};
@@ -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<RegionMigrationProcedureGuard>,
catalog_and_schema: Vec<(String, String)>,
) -> Result<ProcedureId> {
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<Option<ProcedureId>> {
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 { .. });
}
+162 -33
View File
@@ -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<String>,
event_context: Option<PbProcedureEventContext>,
}
impl TryFrom<(Option<ProcedureActor>, Option<PbProcedureEventContext>)> for ProcedureSubmission {
type Error = Status;
fn try_from(
(actor, event_context): (Option<ProcedureActor>, Option<PbProcedureEventContext>),
) -> Result<Self, Self::Error> {
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<ProcedureSubmission> 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<GcResponse> {
async fn handle_gc_regions(
&self,
procedure_context: ProcedureContext,
request: MetaGcRegionsRequest,
) -> error::Result<GcResponse> {
let region_ids: Vec<RegionId> = 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<GcResponse> {
async fn handle_gc_table(
&self,
procedure_context: ProcedureContext,
request: MetaGcTableRequest,
) -> error::Result<GcResponse> {
let table_name_key = TableNameKey::new(
&request.catalog_name,
&request.schema_name,
@@ -406,13 +490,19 @@ impl Metasrv {
.context(TableMetadataManagerSnafu)?;
let region_ids: Vec<RegionId> = 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<RegionId>,
full_file_listing: bool,
timeout: Option<Duration>,
@@ -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);
}
}
+43 -30
View File
@@ -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<Option<String>> {
async fn migrate_region(
&self,
query_ctx: QueryContextRef,
request: MigrateRegionRequest,
) -> QueryResult<Option<String>> {
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<MetaGcResponse> {
async fn gc_regions(
&self,
query_ctx: QueryContextRef,
request: MetaGcRegionsRequest,
) -> QueryResult<MetaGcResponse> {
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<MetaGcResponse> {
async fn gc_table(
&self,
query_ctx: QueryContextRef,
request: MetaGcTableRequest,
) -> QueryResult<MetaGcResponse> {
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<Vec<SubmitDdlTaskRequest>>,
contexts: Mutex<Vec<ExecutorContext>>,
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<SubmitDdlTaskResponse> {
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()
+7 -12
View File
@@ -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)?;
+94 -133
View File
@@ -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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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::<Result<Vec<_>>>()?;
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<SubmitDdlTaskResponse> {
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<CreatorGrantIntent>,
) -> Result<SubmitDdlTaskResponse> {
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<Vec<SubmitDdlTaskRequest>>,
contexts: Mutex<Vec<ExecutorContext>>,
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<SubmitDdlTaskResponse> {
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],
+31 -32
View File
@@ -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
);
}
}
+5 -5
View File
@@ -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();
+110 -64
View File
@@ -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<Strin
let region_id = RegionId::new(table_id, from_regions.leader_regions[0]);
// Trigger region migration.
let procedure = region_migration_manager
.submit_procedure(RegionMigrationProcedureTask::new(
region_id,
peer_factory(from_peer_id),
peer_factory(to_peer_id),
Duration::from_millis(1000),
RegionMigrationTriggerReason::Manual,
))
.submit_procedure(
ProcedureContext::from_event_context(PersistentEventContext::new(
TriggerReason::Manual,
)),
RegionMigrationProcedureTask::new(
region_id,
peer_factory(from_peer_id),
peer_factory(to_peer_id),
Duration::from_millis(1000),
RegionMigrationTriggerReason::Manual,
),
)
.await
.unwrap();
info!("Started region procedure: {}!", procedure.unwrap());
@@ -235,13 +241,18 @@ pub async fn test_region_migration(store_type: StorageType, endpoints: Vec<Strin
// Triggers again.
let err = region_migration_manager
.submit_procedure(RegionMigrationProcedureTask::new(
region_id,
peer_factory(from_peer_id),
peer_factory(to_peer_id),
Duration::from_millis(1000),
RegionMigrationTriggerReason::Manual,
))
.submit_procedure(
ProcedureContext::from_event_context(PersistentEventContext::new(
TriggerReason::Manual,
)),
RegionMigrationProcedureTask::new(
region_id,
peer_factory(from_peer_id),
peer_factory(to_peer_id),
Duration::from_millis(1000),
RegionMigrationTriggerReason::Manual,
),
)
.await
.unwrap_err();
assert!(matches!(err, error::Error::RegionMigrated { .. }));
@@ -523,13 +534,18 @@ pub async fn test_region_migration_by_sql(store_type: StorageType, endpoints: Ve
// Triggers again.
let err = region_migration_manager
.submit_procedure(RegionMigrationProcedureTask::new(
region_id,
peer_factory(from_peer_id),
peer_factory(to_peer_id),
Duration::from_millis(1000),
RegionMigrationTriggerReason::Manual,
))
.submit_procedure(
ProcedureContext::from_event_context(PersistentEventContext::new(
TriggerReason::Manual,
)),
RegionMigrationProcedureTask::new(
region_id,
peer_factory(from_peer_id),
peer_factory(to_peer_id),
Duration::from_millis(1000),
RegionMigrationTriggerReason::Manual,
),
)
.await
.unwrap_err();
assert!(matches!(err, error::Error::RegionMigrated { .. }));
@@ -628,13 +644,18 @@ pub async fn test_region_migration_multiple_regions(
let region_id = RegionId::new(table_id, from_regions.leader_regions[0]);
// Trigger region migration.
let procedure = region_migration_manager
.submit_procedure(RegionMigrationProcedureTask::new(
region_id,
peer_factory(from_peer_id),
peer_factory(to_peer_id),
Duration::from_millis(1000),
RegionMigrationTriggerReason::Manual,
))
.submit_procedure(
ProcedureContext::from_event_context(PersistentEventContext::new(
TriggerReason::Manual,
)),
RegionMigrationProcedureTask::new(
region_id,
peer_factory(from_peer_id),
peer_factory(to_peer_id),
Duration::from_millis(1000),
RegionMigrationTriggerReason::Manual,
),
)
.await
.unwrap();
info!("Started region procedure: {}!", procedure.unwrap());
@@ -689,13 +710,18 @@ pub async fn test_region_migration_multiple_regions(
// Triggers again.
let err = region_migration_manager
.submit_procedure(RegionMigrationProcedureTask::new(
region_id,
peer_factory(from_peer_id),
peer_factory(to_peer_id),
Duration::from_millis(1000),
RegionMigrationTriggerReason::Manual,
))
.submit_procedure(
ProcedureContext::from_event_context(PersistentEventContext::new(
TriggerReason::Manual,
)),
RegionMigrationProcedureTask::new(
region_id,
peer_factory(from_peer_id),
peer_factory(to_peer_id),
Duration::from_millis(1000),
RegionMigrationTriggerReason::Manual,
),
)
.await
.unwrap_err();
assert!(matches!(err, error::Error::RegionMigrated { .. }));
@@ -779,13 +805,18 @@ pub async fn test_region_migration_all_regions(store_type: StorageType, endpoint
let region_id = RegionId::new(table_id, from_regions.leader_regions[0]);
// Trigger region migration.
let procedure = region_migration_manager
.submit_procedure(RegionMigrationProcedureTask::new(
region_id,
peer_factory(from_peer_id),
peer_factory(to_peer_id),
Duration::from_millis(1000),
RegionMigrationTriggerReason::Manual,
))
.submit_procedure(
ProcedureContext::from_event_context(PersistentEventContext::new(
TriggerReason::Manual,
)),
RegionMigrationProcedureTask::new(
region_id,
peer_factory(from_peer_id),
peer_factory(to_peer_id),
Duration::from_millis(1000),
RegionMigrationTriggerReason::Manual,
),
)
.await
.unwrap();
info!("Started region procedure: {}!", procedure.unwrap());
@@ -838,13 +869,18 @@ pub async fn test_region_migration_all_regions(store_type: StorageType, endpoint
// Triggers again.
let err = region_migration_manager
.submit_procedure(RegionMigrationProcedureTask::new(
region_id,
peer_factory(from_peer_id),
peer_factory(to_peer_id),
Duration::from_millis(1000),
RegionMigrationTriggerReason::Manual,
))
.submit_procedure(
ProcedureContext::from_event_context(PersistentEventContext::new(
TriggerReason::Manual,
)),
RegionMigrationProcedureTask::new(
region_id,
peer_factory(from_peer_id),
peer_factory(to_peer_id),
Duration::from_millis(1000),
RegionMigrationTriggerReason::Manual,
),
)
.await
.unwrap_err();
assert!(matches!(err, error::Error::RegionMigrated { .. }));
@@ -917,13 +953,18 @@ pub async fn test_region_migration_incorrect_from_peer(
// Trigger region migration.
let err = region_migration_manager
.submit_procedure(RegionMigrationProcedureTask::new(
region_id,
peer_factory(5),
peer_factory(1),
Duration::from_millis(1000),
RegionMigrationTriggerReason::Manual,
))
.submit_procedure(
ProcedureContext::from_event_context(PersistentEventContext::new(
TriggerReason::Manual,
)),
RegionMigrationProcedureTask::new(
region_id,
peer_factory(5),
peer_factory(1),
Duration::from_millis(1000),
RegionMigrationTriggerReason::Manual,
),
)
.await
.unwrap_err();
@@ -1000,13 +1041,18 @@ pub async fn test_region_migration_incorrect_region_id(
// Trigger region migration.
let err = region_migration_manager
.submit_procedure(RegionMigrationProcedureTask::new(
region_id,
peer_factory(2),
peer_factory(1),
Duration::from_millis(1000),
RegionMigrationTriggerReason::Manual,
))
.submit_procedure(
ProcedureContext::from_event_context(PersistentEventContext::new(
TriggerReason::Manual,
)),
RegionMigrationProcedureTask::new(
region_id,
peer_factory(2),
peer_factory(1),
Duration::from_millis(1000),
RegionMigrationTriggerReason::Manual,
),
)
.await
.unwrap_err();
+5 -1
View File
@@ -18,8 +18,9 @@ use std::time::Duration;
use client::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
use common_error::root_source;
use common_event_recorder::{PersistentEventContext, TriggerReason};
use common_meta::key::table_name::TableNameKey;
use common_procedure::{ProcedureWithId, watcher};
use common_procedure::{ProcedureContext, ProcedureWithId, watcher};
use common_query::Output;
use common_telemetry::info;
use common_test_util::recordbatch::check_output_stream;
@@ -854,6 +855,9 @@ async fn trigger_full_gc(ticker: &GcTickerRef) {
region_ids: None,
full_file_listing: None,
timeout: None,
procedure_context: ProcedureContext::from_event_context(PersistentEventContext::new(
TriggerReason::Manual,
)),
})
.await
.unwrap();
+10 -6
View File
@@ -29,6 +29,10 @@ use common_event_recorder::event_table::{
CATALOG_NAME_COLUMN, PAYLOAD_COLUMN, PHYSICAL_TABLE_ID_COLUMN, PROCEDURE_ID_COLUMN,
PROCEDURE_TRIGGER_COLUMN, SCHEMA_NAME_COLUMN, TABLE_ID_COLUMN, TABLE_NAME_COLUMN, TYPE_COLUMN,
};
#[cfg(feature = "enterprise")]
use common_event_recorder::{PersistentEventContext, TriggerReason};
#[cfg(feature = "enterprise")]
use common_procedure::ProcedureContext;
use common_test_util::temp_dir::create_temp_dir;
use frontend::instance::Instance;
use meta_srv::gc::GcSchedulerOptions;
@@ -377,9 +381,9 @@ async fn test_table_ddl_procedure_events() {
.ddl_manager()
.submit_undrop_table_task(
common_meta::rpc::ddl::UndropTableTask { table_id },
common_meta::rpc::ddl::PersistentEventContext::new(
common_meta::rpc::ddl::TriggerReason::Manual,
),
ProcedureContext::from_event_context(PersistentEventContext::new(
TriggerReason::Manual,
)),
)
.await
.unwrap();
@@ -419,9 +423,9 @@ async fn test_table_ddl_procedure_events() {
.ddl_manager()
.submit_purge_dropped_table_task(
common_meta::rpc::ddl::PurgeDroppedTableTask { table_id },
common_meta::rpc::ddl::PersistentEventContext::new(
common_meta::rpc::ddl::TriggerReason::Manual,
),
ProcedureContext::from_event_context(PersistentEventContext::new(
TriggerReason::Manual,
)),
)
.await
.unwrap();