From 0d5585ba332fe63374ae72e5745643465bb78fc4 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 11 Aug 2026 20:50:50 +0800 Subject: [PATCH] feat(event): record admin function executions (#8835) * feat(event): record admin function executions Signed-off-by: WenyXu * fix(event): handle admin function recording edge cases Signed-off-by: WenyXu * feat(event): record actor for admin functions Signed-off-by: WenyXu * fix(event): preserve admin function event values Signed-off-by: WenyXu * fix(event): preserve non-finite admin results Signed-off-by: WenyXu --------- Signed-off-by: WenyXu (cherry picked from commit 943eee852f292bfffd7240867dc27248216c353f) --- Cargo.lock | 1 + config/config.md | 3 +- config/frontend.example.toml | 4 + config/standalone.example.toml | 2 +- src/cmd/tests/load_config_test.rs | 5 + src/common/event-recorder/src/event_table.rs | 63 ++ src/frontend/src/frontend.rs | 6 +- src/frontend/src/instance/builder.rs | 13 +- src/operator/Cargo.toml | 1 + src/operator/src/error.rs | 4 + src/operator/src/statement.rs | 16 + src/operator/src/statement/admin.rs | 116 +++- src/operator/src/statement/admin/event.rs | 569 ++++++++++++++++++ src/operator/src/statement/admin/layer.rs | 440 ++++++++++++++ .../tests/admin_function_event.rs | 193 ++++++ tests-integration/tests/main.rs | 1 + 16 files changed, 1421 insertions(+), 16 deletions(-) create mode 100644 src/operator/src/statement/admin/event.rs create mode 100644 src/operator/src/statement/admin/layer.rs create mode 100644 tests-integration/tests/admin_function_event.rs diff --git a/Cargo.lock b/Cargo.lock index d2c45a28fe9..7454707d1c3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9949,6 +9949,7 @@ dependencies = [ "common-catalog", "common-datasource", "common-error", + "common-event-recorder", "common-frontend", "common-function", "common-grpc", diff --git a/config/config.md b/config/config.md index 90c346d09ee..7f9064300a6 100644 --- a/config/config.md +++ b/config/config.md @@ -233,7 +233,7 @@ | `tracing.tokio_console_addr` | String | Unset | The tokio console address. | | `event_recorder` | -- | -- | Configuration options for the event recorder. | | `event_recorder.ttl` | String | `90d` | TTL for the events table that will be used to store the events. Default is `90d`. | -| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `create_database`,
`alter_database`, `drop_database`, `create_flow`, `drop_flow`,
`create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,
`drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,
`create_view`, `drop_view`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | +| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `create_database`,
`alter_database`, `drop_database`, `create_flow`, `drop_flow`,
`create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,
`drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,
`create_view`, `drop_view`, `admin_function`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | | `memory` | -- | -- | The memory options. | | `memory.enable_heap_profiling` | Bool | `true` | Whether to enable heap profiling activation during startup.
When enabled, heap profiling will be activated if the `MALLOC_CONF` environment variable
is set to "prof:true,prof_active:false". The official image adds this env variable.
Default is true. | @@ -370,6 +370,7 @@ | `memory.enable_heap_profiling` | Bool | `true` | Whether to enable heap profiling activation during startup.
When enabled, heap profiling will be activated if the `MALLOC_CONF` environment variable
is set to "prof:true,prof_active:false". The official image adds this env variable.
Default is true. | | `event_recorder` | -- | -- | Configuration options for the event recorder. | | `event_recorder.ttl` | String | `90d` | TTL for the events table that will be used to store the events. Default is `90d`. | +| `event_recorder.event_types` | Array | -- | Event types to record. Current available event type: `admin_function`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | ### Metasrv diff --git a/config/frontend.example.toml b/config/frontend.example.toml index f495f3b19f1..30e5f1ff392 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -402,3 +402,7 @@ enable_heap_profiling = true [event_recorder] ## TTL for the events table that will be used to store the events. Default is `90d`. ttl = "90d" +## Event types to record. Current available event type: `admin_function`. +## When omitted, all current and future event types are recorded. +## Set to an empty array to disable event recording. +#+ event_types = ["admin_function"] diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 4abbf815c06..8ced8b17906 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -919,7 +919,7 @@ ttl = "90d" ## `alter_database`, `drop_database`, `create_flow`, `drop_flow`, ## `create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`, ## `drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`, -## `create_view`, `drop_view`. +## `create_view`, `drop_view`, `admin_function`. ## When omitted, all current and future event types are recorded. ## Set to an empty array to disable event recording. #+ event_types = ["create_database"] diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index bdbaf720b7d..fe06d700117 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -445,6 +445,11 @@ fn test_load_event_types_from_env() { .component .event_recorder .event_types, + GreptimeOptions::::load_layered_options(None, env_prefix) + .unwrap() + .component + .event_recorder + .event_types, ] { assert!(event_types.allows("region_migration")); assert!(!event_types.allows("other_event")); diff --git a/src/common/event-recorder/src/event_table.rs b/src/common/event-recorder/src/event_table.rs index 0ec38c8eabc..1ea4369a00a 100644 --- a/src/common/event-recorder/src/event_table.rs +++ b/src/common/event-recorder/src/event_table.rs @@ -88,6 +88,9 @@ pub const TIMESTAMP_COLUMN: EventTableColumn = EventTableColumn::new( ColumnDataType::TimestampNanosecond, SemanticType::Timestamp, ); +/// The effective database user that initiated an event. +pub const ACTOR_COLUMN: EventTableColumn = + EventTableColumn::new("actor", ColumnDataType::String, SemanticType::Field); /// The canonical procedure identifier envelope column. pub const PROCEDURE_ID_COLUMN: EventTableColumn = EventTableColumn::new("procedure_id", ColumnDataType::String, SemanticType::Field); @@ -144,6 +147,25 @@ pub const LATEST_OFFSET_COLUMN: EventTableColumn = pub const GC_REPORT_COLUMN: EventTableColumn = EventTableColumn::json_binary("gc_report", ColumnDataType::Binary, SemanticType::Field); +/// The canonical ADMIN function name field. +pub const ADMIN_FUNCTION_NAME_COLUMN: EventTableColumn = EventTableColumn::new( + "admin_function_name", + ColumnDataType::String, + SemanticType::Field, +); +/// The canonical ADMIN function execution status field. +pub const ADMIN_FUNCTION_STATUS_COLUMN: EventTableColumn = EventTableColumn::new( + "admin_function_status", + ColumnDataType::String, + SemanticType::Field, +); +/// The canonical ADMIN function output field. +pub const ADMIN_FUNCTION_OUTPUT_COLUMN: EventTableColumn = EventTableColumn::json_binary( + "admin_function_output", + ColumnDataType::Binary, + SemanticType::Field, +); + /// The canonical table name field for table DDL events. pub const TABLE_NAME_COLUMN: EventTableColumn = EventTableColumn::new("table_name", ColumnDataType::String, SemanticType::Field); @@ -389,6 +411,47 @@ mod tests { ); } + #[test] + fn admin_function_schema_preserves_names_types_semantics_extensions_and_order() { + assert_eq!( + column_schemas([ + &ACTOR_COLUMN, + &ADMIN_FUNCTION_NAME_COLUMN, + &ADMIN_FUNCTION_STATUS_COLUMN, + &ADMIN_FUNCTION_OUTPUT_COLUMN, + ]), + vec![ + ColumnSchema { + column_name: "actor".to_string(), + datatype: ColumnDataType::String.into(), + semantic_type: SemanticType::Field.into(), + ..Default::default() + }, + ColumnSchema { + column_name: "admin_function_name".to_string(), + datatype: ColumnDataType::String.into(), + semantic_type: SemanticType::Field.into(), + ..Default::default() + }, + ColumnSchema { + column_name: "admin_function_status".to_string(), + datatype: ColumnDataType::String.into(), + semantic_type: SemanticType::Field.into(), + ..Default::default() + }, + ColumnSchema { + column_name: "admin_function_output".to_string(), + datatype: ColumnDataType::Binary.into(), + semantic_type: SemanticType::Field.into(), + datatype_extension: Some(ColumnDataTypeExtension { + type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())), + }), + ..Default::default() + }, + ] + ); + } + #[test] fn table_dimension_schema_preserves_names_types_semantics_and_order() { assert_eq!( diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index 30fcd394922..17435c01a14 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -116,7 +116,11 @@ impl Default for FrontendOptions { impl Configurable for FrontendOptions { fn env_list_keys() -> Option<&'static [&'static str]> { - Some(&["heartbeat_env_vars", "meta_client.metasrv_addrs"]) + Some(&[ + "heartbeat_env_vars", + "meta_client.metasrv_addrs", + "event_recorder.event_types", + ]) } } diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index 42d234fae41..b8e3a9f2828 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -20,7 +20,7 @@ use catalog::CatalogManagerRef; use catalog::process_manager::ProcessManagerRef; use common_base::Plugins; use common_datasource::object_store::LocalFileAccess; -use common_event_recorder::EventRecorderImpl; +use common_event_recorder::{EventRecorderImpl, EventRecorderRef}; use common_meta::cache::{LayeredCacheRegistryRef, TableRouteCacheRef}; use common_meta::cache_invalidator::{CacheInvalidatorRef, DummyCacheInvalidator}; use common_meta::key::TableMetadataManager; @@ -37,8 +37,8 @@ use operator::request::Requester; #[cfg(feature = "enterprise")] use operator::statement::CreateDatabaseHandlerRef; use operator::statement::{ - ExecutorConfigureContext, StatementExecutor, StatementExecutorConfiguratorRef, - StatementExecutorRef, + AdminEventRecorderHandle, AdminFunctionRecordingLayer, ExecutorConfigureContext, + StatementExecutor, StatementExecutorConfiguratorRef, StatementExecutorRef, }; use operator::table::TableMutationOperator; use partition::cache::PartitionInfoCacheRef; @@ -302,6 +302,10 @@ impl FrontendBuilder { statement_executor }; + let admin_event_recorder = AdminEventRecorderHandle::default(); + let statement_executor = statement_executor.with_admin_function_layer(Arc::new( + AdminFunctionRecordingLayer::new(admin_event_recorder.clone()), + )); let statement_executor = Arc::new(statement_executor); let pipeline_operator = Arc::new(PipelineOperator::new( @@ -316,13 +320,14 @@ impl FrontendBuilder { let slow_query_recorder = Arc::new(EventRecorderImpl::new(Box::new( EventHandlerImpl::new(statement_executor.clone(), self.options.slow_query.ttl), ))); - let event_recorder = Arc::new(EventRecorderImpl::with_event_type_filter( + let event_recorder: EventRecorderRef = Arc::new(EventRecorderImpl::with_event_type_filter( Box::new(EventHandlerImpl::new( statement_executor.clone(), self.options.event_recorder.ttl, )), self.options.event_recorder.event_types.clone(), )); + admin_event_recorder.install(&event_recorder); Ok(Instance { frontend_peer_addr, diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml index bd34e21ed24..087ea461f10 100644 --- a/src/operator/Cargo.toml +++ b/src/operator/Cargo.toml @@ -33,6 +33,7 @@ common-base.workspace = true common-catalog.workspace = true common-datasource.workspace = true common-error.workspace = true +common-event-recorder.workspace = true common-frontend.workspace = true common-function.workspace = true common-grpc.workspace = true diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs index eabd6477737..244776ddc9d 100644 --- a/src/operator/src/error.rs +++ b/src/operator/src/error.rs @@ -65,6 +65,9 @@ pub enum Error { location: Location, }, + #[snafu(display("Admin function execution was cancelled"))] + AdminFunctionCancelled, + #[snafu(display("Expected {expected} args, but actual {actual}"))] FunctionArityMismatch { expected: usize, actual: usize }, @@ -997,6 +1000,7 @@ impl ErrorExt for Error { | Error::EncodeJson { .. } | Error::DeserializePartitionExpr { .. } | Error::SerializePartitionExpr { .. } => StatusCode::Unexpected, + Error::AdminFunctionCancelled => StatusCode::Cancelled, Error::ViewNotFound { .. } | Error::ViewInfoNotFound { .. } | Error::TableNotFound { .. } => StatusCode::TableNotFound, diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 320e7c2e642..1143b1c68ce 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -80,6 +80,11 @@ use table::requests::{CopyDatabaseRequest, CopyDirection, CopyQueryToRequest, Co use table::table_name::TableName; use table::table_reference::TableReference; +pub use self::admin::{ + AdminEventRecorderHandle, AdminFunctionLayer, AdminFunctionLayerRef, + AdminFunctionRecordingLayer, AdminFunctionRequest, AdminFunctionResponse, AdminFunctionService, + AdminFunctionServiceRef, +}; use self::set::{ set_bytea_output, set_datestyle, set_intervalstyle, set_timezone, validate_client_encoding, }; @@ -138,6 +143,7 @@ pub struct StatementExecutor { inserter: InserterRef, process_manager: Option, origin_frontend_addr: String, + admin_function_service: AdminFunctionServiceRef, pub(crate) local_file_access: LocalFileAccess, #[cfg(feature = "enterprise")] create_database_handler: Option, @@ -179,6 +185,7 @@ impl StatementExecutor { origin_frontend_addr: String, local_file_access: LocalFileAccess, ) -> Self { + let admin_function_service = admin::new_admin_function_service(query_engine.clone()); Self { catalog_manager, query_engine, @@ -191,6 +198,7 @@ impl StatementExecutor { inserter, process_manager, origin_frontend_addr, + admin_function_service, local_file_access, #[cfg(feature = "enterprise")] create_database_handler: None, @@ -199,6 +207,14 @@ impl StatementExecutor { } } + /// Adds a layer around the ADMIN function execution service. + /// + /// The last added layer is the outermost layer. + pub fn with_admin_function_layer(mut self, layer: AdminFunctionLayerRef) -> Self { + self.admin_function_service = layer.layer(self.admin_function_service); + self + } + #[cfg(feature = "enterprise")] pub fn with_trigger_querier(mut self, querier: TriggerQuerierRef) -> Self { self.trigger_querier = Some(querier); diff --git a/src/operator/src/statement/admin.rs b/src/operator/src/statement/admin.rs index 7494977b214..2b4a909e3cd 100644 --- a/src/operator/src/statement/admin.rs +++ b/src/operator/src/statement/admin.rs @@ -12,6 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod event; +mod layer; + use std::sync::Arc; use common_function::function::FunctionContext; @@ -28,6 +31,10 @@ use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::value::Value; use datatypes::vectors::VectorRef; +pub use layer::{ + AdminEventRecorderHandle, AdminFunctionLayer, AdminFunctionLayerRef, + AdminFunctionRecordingLayer, +}; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt, ensure}; use sql::ast::{Expr, FunctionArg, FunctionArgExpr, FunctionArguments, Value as SqlValue}; @@ -38,14 +45,49 @@ use crate::statement::StatementExecutor; const DUMMY_COLUMN: &str = ""; -impl StatementExecutor { - /// Execute the [`Admin`] statement and returns the output. - #[tracing::instrument(skip_all)] - pub(super) async fn execute_admin_command( - &self, - stmt: Admin, - query_ctx: QueryContextRef, - ) -> Result { +/// A request to execute one ADMIN function statement. +#[derive(Clone)] +pub struct AdminFunctionRequest { + /// The parsed ADMIN statement to execute. + pub statement: Admin, + /// The query context of the request. + pub query_ctx: QueryContextRef, +} + +/// The client output and immediate typed result of one ADMIN function execution. +pub struct AdminFunctionResponse { + /// The output returned to the client. + pub output: Output, + /// The typed immediate result exposed to outer layers. + pub immediate_result: Option, +} + +/// Executes an ADMIN function request. +#[async_trait::async_trait] +pub trait AdminFunctionService: Send + Sync { + /// Executes an ADMIN function request. + async fn call(&self, request: AdminFunctionRequest) -> Result; +} + +/// A shared ADMIN function service. +pub type AdminFunctionServiceRef = Arc; + +#[derive(Clone)] +struct CoreAdminFunctionService { + query_engine: query::QueryEngineRef, +} + +impl CoreAdminFunctionService { + fn new(query_engine: query::QueryEngineRef) -> Self { + Self { query_engine } + } + + async fn execute(&self, request: AdminFunctionRequest) -> Result { + let AdminFunctionRequest { + statement: stmt, + query_ctx, + } = request; + let Admin::Func(func) = &stmt; // the function name should be in lower case. let func_name = func.name.to_string().to_lowercase(); @@ -141,6 +183,7 @@ impl StatementExecutor { (&result_columnar).try_into().context(CastSnafu)?; let result_vector: VectorRef = result_columnar.try_into_vector(1).context(CastSnafu)?; + let immediate_result = immediate_result(&result_vector); let column_schemas = vec![ColumnSchema::new( // Use statement as the result column name @@ -154,7 +197,46 @@ impl StatementExecutor { let batches = RecordBatches::try_new(schema, vec![batch]).context(error::BuildRecordBatchSnafu)?; - Ok(Output::new_with_record_batches(batches)) + Ok(AdminFunctionResponse { + output: Output::new_with_record_batches(batches), + immediate_result, + }) + } +} + +fn immediate_result(result_vector: &VectorRef) -> Option { + (!result_vector.is_empty()).then(|| result_vector.get(0)) +} + +#[async_trait::async_trait] +impl AdminFunctionService for CoreAdminFunctionService { + async fn call(&self, request: AdminFunctionRequest) -> Result { + self.execute(request).await + } +} + +/// Creates the core ADMIN function service. +pub(crate) fn new_admin_function_service( + query_engine: query::QueryEngineRef, +) -> AdminFunctionServiceRef { + Arc::new(CoreAdminFunctionService::new(query_engine)) +} + +impl StatementExecutor { + /// Executes the [`Admin`] statement and returns the output. + #[tracing::instrument(skip_all)] + pub(crate) async fn execute_admin_command( + &self, + stmt: Admin, + query_ctx: QueryContextRef, + ) -> Result { + self.admin_function_service + .call(AdminFunctionRequest { + statement: stmt, + query_ctx, + }) + .await + .map(|response| response.output) } } @@ -306,3 +388,19 @@ fn try_get_data_type_for_sql_value(value: &SqlValue) -> Result { .fail(), } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datatypes::vectors::{Int32Vector, VectorRef}; + + use crate::statement::admin::immediate_result; + + #[test] + fn empty_admin_function_result_has_no_immediate_value() { + let result_vector: VectorRef = Arc::new(Int32Vector::from(vec![])); + + assert_eq!(immediate_result(&result_vector), None); + } +} diff --git a/src/operator/src/statement/admin/event.rs b/src/operator/src/statement/admin/event.rs new file mode 100644 index 00000000000..bafa2ad8824 --- /dev/null +++ b/src/operator/src/statement/admin/event.rs @@ -0,0 +1,569 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::any::Any; + +use api::v1::value::ValueData; +use api::v1::{ColumnSchema, Row}; +use common_event_recorder::Event; +use common_event_recorder::error::Result as EventResult; +use common_event_recorder::event_table::{ + ACTOR_COLUMN, ADMIN_FUNCTION_NAME_COLUMN, ADMIN_FUNCTION_OUTPUT_COLUMN, + ADMIN_FUNCTION_STATUS_COLUMN, column_schemas, jsonb_value, +}; +use datatypes::value::Value; +use serde_json::{Value as JsonValue, json}; +use sql::ast::{Expr, FunctionArg, FunctionArgExpr, FunctionArguments, Value as SqlValue}; +use sql::statements::admin::Admin; + +use crate::error::Error; +use crate::statement::admin::AdminFunctionRequest; + +/// Event type emitted for an ADMIN function execution. +pub(crate) const ADMIN_FUNCTION_EVENT_TYPE: &str = "admin_function"; + +const PAYLOAD_VERSION: u8 = 1; +const VERSION_KEY: &str = "version"; +const ARGUMENTS_KEY: &str = "arguments"; +const ARGUMENT_TYPE_KEY: &str = "type"; +const ARGUMENT_VALUE_KEY: &str = "value"; +const LIST_ARGUMENT_TYPE: &str = "list"; +const SUBQUERY_ARGUMENT_TYPE: &str = "subquery"; +const RESULT_KEY: &str = "result"; +const ERROR_KEY: &str = "error"; +const SUCCEEDED_STATUS: &str = "Succeeded"; +const FAILED_STATUS: &str = "Failed"; +const UNSUPPORTED: &str = ""; + +/// The ADMIN function input captured before execution. +#[derive(Debug, Clone)] +pub(crate) struct AdminFunctionEventInput { + actor: String, + function: String, + arguments: Option, +} + +impl AdminFunctionEventInput { + /// Creates an event input from an ADMIN function request. + pub(crate) fn from_request(request: &AdminFunctionRequest) -> Self { + let Admin::Func(function) = &request.statement; + let function_name = function.name.to_string().to_lowercase(); + + let arguments = match &function.args { + FunctionArguments::List(arguments) => Some(json!({ + (ARGUMENT_TYPE_KEY): LIST_ARGUMENT_TYPE, + (ARGUMENT_VALUE_KEY): arguments.args.iter().map(argument_to_json).collect::>(), + })), + FunctionArguments::Subquery(query) => Some(json!({ + (ARGUMENT_TYPE_KEY): SUBQUERY_ARGUMENT_TYPE, + (ARGUMENT_VALUE_KEY): query.to_string(), + })), + FunctionArguments::None => None, + }; + + Self { + actor: request.query_ctx.current_user().username().to_string(), + function: function_name, + arguments, + } + } +} + +/// An event describing the outcome of an ADMIN function execution. +#[derive(Debug)] +pub(crate) struct AdminFunctionEvent { + actor: String, + function_name: String, + status: &'static str, + output: JsonValue, + payload: JsonValue, +} + +impl AdminFunctionEvent { + /// Creates a successful ADMIN function event. + pub(crate) fn success(input: AdminFunctionEventInput, result: Option<&Value>) -> Self { + let AdminFunctionEventInput { + actor, + function, + arguments, + } = input; + Self { + actor, + function_name: function, + status: SUCCEEDED_STATUS, + output: result.map_or_else( + || json!({}), + |result| { + json!({ + (RESULT_KEY): value_to_json(result), + }) + }, + ), + payload: input_payload(arguments), + } + } + + /// Creates a failed ADMIN function event with the debug representation of the error. + pub(crate) fn failure(input: AdminFunctionEventInput, error: &Error) -> Self { + let AdminFunctionEventInput { + actor, + function, + arguments, + } = input; + Self { + actor, + function_name: function, + status: FAILED_STATUS, + output: json!({ + (ERROR_KEY): format!("{error:?}"), + }), + payload: input_payload(arguments), + } + } + + /// Creates a failed ADMIN function event for a cancelled execution. + pub(crate) fn cancelled(input: AdminFunctionEventInput) -> Self { + Self::failure(input, &Error::AdminFunctionCancelled) + } +} + +impl Event for AdminFunctionEvent { + fn event_type(&self) -> &str { + ADMIN_FUNCTION_EVENT_TYPE + } + + fn json_payload(&self) -> EventResult { + Ok(self.payload.clone()) + } + + fn extra_schema(&self) -> Vec { + column_schemas([ + &ACTOR_COLUMN, + &ADMIN_FUNCTION_NAME_COLUMN, + &ADMIN_FUNCTION_STATUS_COLUMN, + &ADMIN_FUNCTION_OUTPUT_COLUMN, + ]) + } + + fn extra_rows(&self) -> EventResult> { + Ok(vec![Row { + values: vec![ + ValueData::StringValue(self.actor.clone()).into(), + ValueData::StringValue(self.function_name.clone()).into(), + ValueData::StringValue(self.status.to_string()).into(), + jsonb_value(&self.output), + ], + }]) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +fn input_payload(arguments: Option) -> JsonValue { + let mut payload = json!({ VERSION_KEY: PAYLOAD_VERSION }); + if let Some(arguments) = arguments { + payload[ARGUMENTS_KEY] = arguments; + } + payload +} + +fn argument_to_json(argument: &FunctionArg) -> JsonValue { + match argument { + FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(value))) => { + sql_value_to_json(&value.value) + } + _ => JsonValue::String(argument.to_string()), + } +} + +fn sql_value_to_json(value: &SqlValue) -> JsonValue { + match value { + SqlValue::Number(value, _) => { + serde_json::from_str(value).unwrap_or_else(|_| JsonValue::String(value.clone())) + } + SqlValue::Boolean(value) => JsonValue::Bool(*value), + SqlValue::Null => JsonValue::Null, + SqlValue::SingleQuotedString(value) | SqlValue::DoubleQuotedString(value) => { + JsonValue::String(value.clone()) + } + SqlValue::Placeholder(value) => JsonValue::String(value.clone()), + _ => value + .clone() + .into_string() + .map(JsonValue::String) + .unwrap_or_else(|| JsonValue::String(value.to_string())), + } +} + +fn value_to_json(value: &Value) -> JsonValue { + match value { + Value::Float32(value) if !value.0.is_finite() => JsonValue::String(value.to_string()), + Value::Float64(value) if !value.0.is_finite() => JsonValue::String(value.to_string()), + _ => JsonValue::try_from(value.clone()) + .unwrap_or_else(|_| JsonValue::String(UNSUPPORTED.to_string())), + } +} + +#[cfg(test)] +mod tests { + use api::v1::Row; + use api::v1::value::ValueData; + use common_event_recorder::Event; + use common_event_recorder::event_table::{ + ACTOR_COLUMN, ADMIN_FUNCTION_NAME_COLUMN, ADMIN_FUNCTION_OUTPUT_COLUMN, + ADMIN_FUNCTION_STATUS_COLUMN, column_schemas, jsonb_value, + }; + use datatypes::value::Value; + use serde_json::json; + use session::context::QueryContext; + use sql::ast::{ + Expr, FunctionArg, FunctionArgExpr, FunctionArguments, Ident, Value as SqlValue, + }; + use sql::dialect::GreptimeDbDialect; + use sql::parser::{ParseOptions, ParserContext}; + use sql::statements::admin::Admin; + use sql::statements::statement::Statement; + use sqlparser::ast::{DollarQuotedString, FunctionArgOperator, FunctionArgumentList}; + + use crate::error::Error; + use crate::statement::admin::AdminFunctionRequest; + use crate::statement::admin::event::{ + ADMIN_FUNCTION_EVENT_TYPE, AdminFunctionEvent, AdminFunctionEventInput, sql_value_to_json, + value_to_json, + }; + + fn request(sql: &str) -> AdminFunctionRequest { + let Statement::Admin(statement) = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap() + .remove(0) + else { + panic!("expected ADMIN statement") + }; + AdminFunctionRequest { + statement, + query_ctx: QueryContext::arc(), + } + } + + fn request_with_arguments(arguments: FunctionArguments) -> AdminFunctionRequest { + let mut request = request("ADMIN plugin_function()"); + let Admin::Func(function) = &mut request.statement; + function.args = arguments; + request + } + + fn assert_admin_columns( + event: &AdminFunctionEvent, + function_name: &str, + status: &str, + output: serde_json::Value, + ) { + assert_eq!(event.event_type(), ADMIN_FUNCTION_EVENT_TYPE); + assert_eq!( + event.extra_schema(), + column_schemas([ + &ACTOR_COLUMN, + &ADMIN_FUNCTION_NAME_COLUMN, + &ADMIN_FUNCTION_STATUS_COLUMN, + &ADMIN_FUNCTION_OUTPUT_COLUMN, + ]) + ); + assert_eq!( + event.extra_rows().unwrap(), + vec![Row { + values: vec![ + ValueData::StringValue("greptime".to_string()).into(), + ValueData::StringValue(function_name.to_string()).into(), + ValueData::StringValue(status.to_string()).into(), + jsonb_value(&output), + ], + }] + ); + } + + #[test] + fn success_uses_separate_output_columns() { + let input = AdminFunctionEventInput::from_request(&request( + "ADMIN flush_table('greptime.public.demo')", + )); + let event = AdminFunctionEvent::success(input, Some(&Value::UInt64(3))); + + assert_admin_columns(&event, "flush_table", "Succeeded", json!({"result": 3})); + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "list", "value": ["greptime.public.demo"]}, + }) + ); + } + + #[test] + fn successful_null_is_explicit() { + let input = AdminFunctionEventInput::from_request(&request( + "ADMIN migrate_region(NULL, NULL, NULL)", + )); + let event = AdminFunctionEvent::success(input, Some(&Value::Null)); + + assert_admin_columns( + &event, + "migrate_region", + "Succeeded", + json!({"result": null}), + ); + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "list", "value": [null, null, null]}, + }) + ); + } + + #[test] + fn unknown_function_values_are_recorded() { + let input = AdminFunctionEventInput::from_request(&request( + "ADMIN plugin_function('plugin-value', NULL)", + )); + let event = AdminFunctionEvent::success(input, Some(&Value::from("plugin-result"))); + + assert_admin_columns( + &event, + "plugin_function", + "Succeeded", + json!({"result": "plugin-result"}), + ); + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "list", "value": ["plugin-value", null]}, + }) + ); + } + + #[test] + fn failure_payload_has_error_only() { + let input = AdminFunctionEventInput::from_request(&request( + "ADMIN flush_table('greptime.public.demo')", + )); + let error = Error::BuildAdminFunctionArgs { + msg: "invalid input".to_string(), + }; + let expected_error = format!("{error:?}"); + let event = AdminFunctionEvent::failure(input, &error); + + assert_admin_columns( + &event, + "flush_table", + "Failed", + json!({"error": expected_error}), + ); + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "list", "value": ["greptime.public.demo"]}, + }) + ); + } + + #[test] + fn cancellation_uses_debug_error_format() { + let input = AdminFunctionEventInput::from_request(&request("ADMIN flush_table('demo')")); + let event = AdminFunctionEvent::cancelled(input); + + assert_admin_columns( + &event, + "flush_table", + "Failed", + json!({"error": format!("{:?}", Error::AdminFunctionCancelled)}), + ); + } + + #[test] + fn subquery_arguments_use_display() { + let Statement::Query(query) = ParserContext::create_with_dialect( + "SELECT 1", + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap() + .remove(0) else { + panic!("expected query statement") + }; + let input = AdminFunctionEventInput::from_request(&request_with_arguments( + FunctionArguments::Subquery(Box::new(query.inner)), + )); + let event = AdminFunctionEvent::failure( + input, + &Error::BuildAdminFunctionArgs { + msg: "subquery is not executable".to_string(), + }, + ); + + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "subquery", "value": "SELECT 1"}, + }) + ); + } + + #[test] + fn no_arguments_omits_arguments_field_and_empty_result() { + let input = + AdminFunctionEventInput::from_request(&request_with_arguments(FunctionArguments::None)); + let event = AdminFunctionEvent::success(input, None); + + assert_admin_columns(&event, "plugin_function", "Succeeded", json!({})); + assert_eq!(event.json_payload().unwrap(), json!({"version": 1})); + } + + #[test] + fn empty_argument_list_is_recorded() { + let input = AdminFunctionEventInput::from_request(&request("ADMIN plugin_function()")); + let event = AdminFunctionEvent::success(input, Some(&Value::UInt64(0))); + + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "list", "value": []}, + }) + ); + } + + #[test] + fn records_extended_sql_values_and_non_literal_arguments() { + let input = AdminFunctionEventInput::from_request(&request( + "ADMIN plugin_function(1, true, NULL, X'48656c6c6f', table_name)", + )); + let event = AdminFunctionEvent::success(input, Some(&Value::UInt64(0))); + + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": { + "type": "list", + "value": [1, true, null, "48656c6c6f", "table_name"], + }, + }) + ); + } + + #[test] + fn string_literals_preserve_logical_content() { + let values = [ + SqlValue::SingleQuotedString("single".to_string()), + SqlValue::DoubleQuotedString("double".to_string()), + SqlValue::DollarQuotedString(DollarQuotedString { + value: "dollar".to_string(), + tag: None, + }), + SqlValue::SingleQuotedRawStringLiteral("raw".to_string()), + SqlValue::SingleQuotedByteStringLiteral("bytes".to_string()), + SqlValue::NationalStringLiteral("national".to_string()), + SqlValue::UnicodeStringLiteral("unicode".to_string()), + SqlValue::HexStringLiteral("48656c6c6f".to_string()), + ]; + + let expected = [ + "single", + "double", + "dollar", + "raw", + "bytes", + "national", + "unicode", + "48656c6c6f", + ]; + for (value, expected) in values.iter().zip(expected) { + assert_eq!(sql_value_to_json(value), json!(expected)); + } + } + + #[test] + fn non_literal_arguments_use_sql_display() { + let arguments = FunctionArgumentList { + duplicate_treatment: None, + args: vec![ + FunctionArg::Named { + name: Ident::new("named"), + arg: FunctionArgExpr::Expr(Expr::Identifier(Ident::new("value"))), + operator: FunctionArgOperator::RightArrow, + }, + FunctionArg::Unnamed(FunctionArgExpr::Wildcard), + ], + clauses: vec![], + }; + let input = AdminFunctionEventInput::from_request(&request_with_arguments( + FunctionArguments::List(arguments), + )); + let event = AdminFunctionEvent::success(input, Some(&Value::UInt64(0))); + + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "list", "value": ["named => value", "*"]}, + }) + ); + assert_eq!( + sql_value_to_json(&SqlValue::Placeholder("$1".to_string())), + json!("$1") + ); + } + + #[test] + fn serializes_binary_result() { + assert_eq!( + value_to_json(&Value::from(vec![1_u8, 2, 3])), + json!([1, 2, 3]) + ); + } + + #[test] + fn serializes_non_finite_float_results_as_strings() { + assert_eq!( + value_to_json(&Value::Float32(f32::NAN.into())), + json!("NaN") + ); + assert_eq!( + value_to_json(&Value::Float32(f32::INFINITY.into())), + json!("inf") + ); + assert_eq!( + value_to_json(&Value::Float32(f32::NEG_INFINITY.into())), + json!("-inf") + ); + assert_eq!( + value_to_json(&Value::Float64(f64::NAN.into())), + json!("NaN") + ); + assert_eq!( + value_to_json(&Value::Float64(f64::INFINITY.into())), + json!("inf") + ); + assert_eq!( + value_to_json(&Value::Float64(f64::NEG_INFINITY.into())), + json!("-inf") + ); + } +} diff --git a/src/operator/src/statement/admin/layer.rs b/src/operator/src/statement/admin/layer.rs new file mode 100644 index 00000000000..195d47d48aa --- /dev/null +++ b/src/operator/src/statement/admin/layer.rs @@ -0,0 +1,440 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::{Arc, RwLock, Weak}; + +use common_event_recorder::{EventRecorder, EventRecorderRef}; + +use crate::error::Result; +use crate::statement::admin::event::{ + ADMIN_FUNCTION_EVENT_TYPE, AdminFunctionEvent, AdminFunctionEventInput, +}; +use crate::statement::admin::{ + AdminFunctionRequest, AdminFunctionResponse, AdminFunctionService, AdminFunctionServiceRef, +}; + +/// A composable layer over ADMIN function execution. +pub trait AdminFunctionLayer: Send + Sync { + /// Wraps `inner` and returns the resulting ADMIN function service. + fn layer(&self, inner: AdminFunctionServiceRef) -> AdminFunctionServiceRef; +} + +/// A shared composable ADMIN function layer. +pub type AdminFunctionLayerRef = Arc; + +/// A late-bound weak reference to the event recorder. +/// +/// The weak reference avoids retaining the frontend event handler, which owns a +/// statement executor itself. +#[derive(Clone, Default)] +pub struct AdminEventRecorderHandle { + recorder: Arc>>>, +} + +impl AdminEventRecorderHandle { + /// Installs the event recorder as a weak reference. + pub fn install(&self, recorder: &EventRecorderRef) { + *self.recorder.write().unwrap_or_else(|e| e.into_inner()) = Some(Arc::downgrade(recorder)); + } + + fn upgrade(&self) -> Option { + self.recorder + .read() + .unwrap_or_else(|e| e.into_inner()) + .as_ref() + .and_then(Weak::upgrade) + } +} + +/// A layer that records ADMIN function inputs and outcomes as events. +#[derive(Clone)] +pub struct AdminFunctionRecordingLayer { + recorder: AdminEventRecorderHandle, +} + +impl AdminFunctionRecordingLayer { + /// Creates a recording layer backed by the late-bound recorder handle. + pub fn new(recorder: AdminEventRecorderHandle) -> Self { + Self { recorder } + } +} + +impl AdminFunctionLayer for AdminFunctionRecordingLayer { + fn layer(&self, inner: AdminFunctionServiceRef) -> AdminFunctionServiceRef { + Arc::new(AdminFunctionRecordingService { + inner, + recorder: self.recorder.clone(), + }) + } +} + +#[derive(Clone)] +struct AdminFunctionRecordingService { + inner: AdminFunctionServiceRef, + recorder: AdminEventRecorderHandle, +} + +#[async_trait::async_trait] +impl AdminFunctionService for AdminFunctionRecordingService { + async fn call(&self, request: AdminFunctionRequest) -> Result { + // Avoid building the event input when this event type is disabled. Do + // not retain the recorder across ADMIN execution because its handler can + // own the statement executor that contains this layer. + let enabled = self.recorder.upgrade().is_some_and(|recorder| { + recorder + .event_type_filter() + .allows(ADMIN_FUNCTION_EVENT_TYPE) + }); + let mut recording = AdminFunctionRecordingGuard::new( + enabled.then(|| AdminFunctionEventInput::from_request(&request)), + self.recorder.clone(), + ); + let result = self.inner.call(request).await; + recording.record_result(&result); + result + } +} + +/// Records an ADMIN event after completion or when an in-flight call is dropped. +struct AdminFunctionRecordingGuard { + input: Option, + recorder: AdminEventRecorderHandle, +} + +impl AdminFunctionRecordingGuard { + fn new(input: Option, recorder: AdminEventRecorderHandle) -> Self { + Self { input, recorder } + } + + fn record_result(&mut self, result: &Result) { + let Some(input) = self.input.take() else { + return; + }; + let Some(recorder) = self.recorder.upgrade() else { + return; + }; + let event = match result { + Ok(response) => AdminFunctionEvent::success(input, response.immediate_result.as_ref()), + Err(error) => AdminFunctionEvent::failure(input, error), + }; + recorder.record(Box::new(event)); + } +} + +impl Drop for AdminFunctionRecordingGuard { + fn drop(&mut self) { + let Some(input) = self.input.take() else { + return; + }; + let Some(recorder) = self.recorder.upgrade() else { + return; + }; + recorder.record(Box::new(AdminFunctionEvent::cancelled(input))); + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + use std::fmt::{Debug, Formatter}; + use std::sync::{Arc, Mutex}; + use std::time::Duration; + + use common_event_recorder::event_table::jsonb_value; + use common_event_recorder::{ + Event, EventRecorder, EventRecorderRef, EventTypeFilter, EventTypeFilterRef, + }; + use common_query::{Output, OutputData}; + use datatypes::value::Value; + use serde_json::json; + use session::context::QueryContext; + use sql::dialect::GreptimeDbDialect; + use sql::parser::{ParseOptions, ParserContext}; + use sql::statements::statement::Statement; + use tokio::sync::Notify; + + use crate::error::Result; + use crate::statement::admin::layer::{AdminEventRecorderHandle, AdminFunctionRecordingLayer}; + use crate::statement::admin::{ + AdminFunctionLayer, AdminFunctionRequest, AdminFunctionResponse, AdminFunctionService, + AdminFunctionServiceRef, + }; + + #[derive(Clone)] + struct TraceLayer { + trace: Arc>>, + before: &'static str, + after: &'static str, + } + + impl AdminFunctionLayer for TraceLayer { + fn layer(&self, inner: AdminFunctionServiceRef) -> AdminFunctionServiceRef { + Arc::new(TraceService { + inner, + trace: self.trace.clone(), + before: self.before, + after: self.after, + }) + } + } + + struct TraceService { + inner: AdminFunctionServiceRef, + trace: Arc>>, + before: &'static str, + after: &'static str, + } + + #[async_trait::async_trait] + impl AdminFunctionService for TraceService { + async fn call(&self, request: AdminFunctionRequest) -> Result { + self.trace.lock().unwrap().push(self.before); + let result = self.inner.call(request).await; + self.trace.lock().unwrap().push(self.after); + result + } + } + + struct TestService { + trace: Option>>>, + immediate_result: Option, + } + + #[async_trait::async_trait] + impl AdminFunctionService for TestService { + async fn call(&self, _request: AdminFunctionRequest) -> Result { + if let Some(trace) = &self.trace { + trace.lock().unwrap().push("core"); + } + Ok(response(self.immediate_result.clone())) + } + } + + struct PendingService { + started: Arc, + } + + #[async_trait::async_trait] + impl AdminFunctionService for PendingService { + async fn call(&self, _request: AdminFunctionRequest) -> Result { + self.started.notify_one(); + std::future::pending().await + } + } + + struct TestRecorder { + events: Mutex>>, + filter: EventTypeFilterRef, + } + + impl TestRecorder { + fn new(filter: EventTypeFilter) -> Self { + Self { + events: Mutex::new(Vec::new()), + filter: Arc::new(filter), + } + } + } + + impl Debug for TestRecorder { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TestRecorder").finish() + } + } + + impl EventRecorder for TestRecorder { + fn record(&self, event: Box) { + self.events.lock().unwrap().push(event); + } + + fn event_type_filter(&self) -> EventTypeFilterRef { + self.filter.clone() + } + + fn close(&self) {} + } + + #[tokio::test] + async fn last_added_layer_runs_outermost() { + let trace = Arc::new(Mutex::new(Vec::new())); + let mut service: AdminFunctionServiceRef = Arc::new(TestService { + trace: Some(trace.clone()), + immediate_result: Some(Value::UInt64(0)), + }); + for (before, after) in [("a_before", "a_after"), ("b_before", "b_after")] { + service = TraceLayer { + trace: trace.clone(), + before, + after, + } + .layer(service); + } + + service.call(request()).await.unwrap(); + assert_eq!( + *trace.lock().unwrap(), + ["b_before", "a_before", "core", "a_after", "b_after"] + ); + } + + #[tokio::test] + async fn recording_respects_filter_and_weak_lifetime() { + let enabled = Arc::new(TestRecorder::new(EventTypeFilter::All)); + let enabled_ref: EventRecorderRef = enabled.clone(); + let handle = AdminEventRecorderHandle::default(); + handle.install(&enabled_ref); + recording_service(handle.clone()) + .call(request()) + .await + .unwrap(); + assert_eq!(enabled.events.lock().unwrap().len(), 1); + + let disabled = Arc::new(TestRecorder::new(EventTypeFilter::Only(HashSet::new()))); + let disabled_ref: EventRecorderRef = disabled.clone(); + handle.install(&disabled_ref); + recording_service(handle.clone()) + .call(request()) + .await + .unwrap(); + assert!(disabled.events.lock().unwrap().is_empty()); + + drop(disabled_ref); + drop(disabled); + assert!(handle.upgrade().is_none()); + recording_service(handle).call(request()).await.unwrap(); + + let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All)); + let recorder_ref: EventRecorderRef = recorder.clone(); + let handle = AdminEventRecorderHandle::default(); + handle.install(&recorder_ref); + let service = recording_service(handle.clone()); + let future = service.call(request()); + drop(recorder_ref); + drop(recorder); + assert!(handle.upgrade().is_none()); + future.await.unwrap(); + } + + #[tokio::test] + async fn recording_does_not_retain_recorder_while_pending() { + let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All)); + let recorder_ref: EventRecorderRef = recorder.clone(); + let handle = AdminEventRecorderHandle::default(); + handle.install(&recorder_ref); + let started = Arc::new(Notify::new()); + let service = + AdminFunctionRecordingLayer::new(handle.clone()).layer(Arc::new(PendingService { + started: started.clone(), + })); + + let task = tokio::spawn(async move { service.call(request()).await }); + started.notified().await; + drop(recorder_ref); + drop(recorder); + + assert!(handle.upgrade().is_none()); + task.abort(); + assert!(matches!(task.await, Err(error) if error.is_cancelled())); + } + + #[tokio::test] + async fn dropping_pending_call_records_failure() { + let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All)); + let recorder_ref: EventRecorderRef = recorder.clone(); + let handle = AdminEventRecorderHandle::default(); + handle.install(&recorder_ref); + let started = Arc::new(Notify::new()); + let service = AdminFunctionRecordingLayer::new(handle).layer(Arc::new(PendingService { + started: started.clone(), + })); + + let task = tokio::spawn(async move { service.call(request()).await }); + started.notified().await; + task.abort(); + assert!(matches!(task.await, Err(error) if error.is_cancelled())); + assert_eq!(recorder.events.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn timing_out_pending_call_records_failure() { + let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All)); + let recorder_ref: EventRecorderRef = recorder.clone(); + let handle = AdminEventRecorderHandle::default(); + handle.install(&recorder_ref); + let service = AdminFunctionRecordingLayer::new(handle).layer(Arc::new(PendingService { + started: Arc::new(Notify::new()), + })); + + assert!( + tokio::time::timeout(Duration::from_millis(10), service.call(request())) + .await + .is_err() + ); + assert_eq!(recorder.events.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn empty_immediate_result_preserves_success_output_and_records_no_result() { + let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All)); + let recorder_ref: EventRecorderRef = recorder.clone(); + let handle = AdminEventRecorderHandle::default(); + handle.install(&recorder_ref); + let service = AdminFunctionRecordingLayer::new(handle).layer(Arc::new(TestService { + trace: None, + immediate_result: None, + })); + + let response = service.call(request()).await.unwrap(); + assert!(matches!(response.output.data, OutputData::AffectedRows(0))); + assert_eq!(response.immediate_result, None); + + let events = recorder.events.lock().unwrap(); + assert_eq!(events.len(), 1); + assert_eq!( + events[0].extra_rows().unwrap()[0].values[3], + jsonb_value(&json!({})) + ); + } + + fn recording_service(handle: AdminEventRecorderHandle) -> AdminFunctionServiceRef { + let core = Arc::new(TestService { + trace: None, + immediate_result: Some(Value::UInt64(0)), + }); + AdminFunctionRecordingLayer::new(handle).layer(core) + } + + fn response(immediate_result: Option) -> AdminFunctionResponse { + AdminFunctionResponse { + output: Output::new_with_affected_rows(0), + immediate_result, + } + } + + fn request() -> AdminFunctionRequest { + let Statement::Admin(statement) = ParserContext::create_with_dialect( + "ADMIN flush_table('demo')", + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap() + .remove(0) else { + panic!("expected ADMIN statement") + }; + AdminFunctionRequest { + statement, + query_ctx: QueryContext::arc(), + } + } +} diff --git a/tests-integration/tests/admin_function_event.rs b/tests-integration/tests/admin_function_event.rs new file mode 100644 index 00000000000..d669d2900ad --- /dev/null +++ b/tests-integration/tests/admin_function_event.rs @@ -0,0 +1,193 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use common_test_util::temp_dir::create_temp_dir; +use frontend::instance::Instance; +use servers::query_handler::sql::SqlQueryHandler; +use session::context::QueryContext; +use sql::ast::FunctionArguments; +use sql::dialect::GreptimeDbDialect; +use sql::parser::{ParseOptions, ParserContext}; +use sql::statements::admin::Admin; +use sql::statements::statement::Statement; +use tests_integration::cluster::GreptimeDbClusterBuilder; +use tests_integration::standalone::GreptimeDbStandaloneBuilder; +use tests_integration::test_util::{StorageType, get_test_store_config}; + +use crate::event_recorder_test_util::assert_eventually_eq; + +const TABLE_NAME: &str = "admin_function_event_test"; + +#[tokio::test(flavor = "multi_thread")] +async fn test_standalone_admin_function_events() { + common_telemetry::init_default_ut_logging(); + let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_admin_function_events") + .build() + .await; + assert_admin_function_events(standalone.fe_instance()).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_distributed_admin_function_events() { + let store_type = StorageType::File; + if !store_type.test_on() { + return; + } + + common_telemetry::init_default_ut_logging(); + let (store_config, _guard) = get_test_store_config(&store_type); + let home_dir = create_temp_dir("test_distributed_admin_function_events_data_home"); + let cluster = GreptimeDbClusterBuilder::new("test_distributed_admin_function_events") + .await + .with_datanodes(1) + .with_store_config(store_config) + .with_shared_home_dir(Arc::new(home_dir)) + .build(true) + .await; + assert_admin_function_events(cluster.fe_instance()).await; +} + +async fn assert_admin_function_events(instance: &Arc) { + execute( + instance, + &format!("CREATE TABLE {TABLE_NAME} (ts TIMESTAMP TIME INDEX)"), + ) + .await + .unwrap(); + execute(instance, &format!("ADMIN FLUSH_TABLE('{TABLE_NAME}')")) + .await + .unwrap(); + + assert_admin_event( + instance, + "flush_table", + "\ ++----------------+----------+---------------------+-----------+---------------------------------------------------------------------------------+--------------+ +| type | actor | admin_function_name | status | payload | output | ++----------------+----------+---------------------+-----------+---------------------------------------------------------------------------------+--------------+ +| admin_function | greptime | flush_table | Succeeded | {\"arguments\":{\"type\":\"list\",\"value\":[\"admin_function_event_test\"]},\"version\":1} | {\"result\":0} | ++----------------+----------+---------------------+-----------+---------------------------------------------------------------------------------+--------------+", + ) + .await; + + let failure = execute(instance, "ADMIN missing_admin_function('missing-input')").await; + assert!(failure.is_err()); + + assert_admin_event( + instance, + "missing_admin_function", + "\ ++----------------+----------+------------------------+--------+---------------------------------------------------------------------+-----------------------------------------------------------------+ +| type | actor | admin_function_name | status | payload | output | ++----------------+----------+------------------------+--------+---------------------------------------------------------------------+-----------------------------------------------------------------+ +| admin_function | greptime | missing_admin_function | Failed | {\"arguments\":{\"type\":\"list\",\"value\":[\"missing-input\"]},\"version\":1} | {\"error\":\"0: Admin function not found: missing_admin_function\"} | ++----------------+----------+------------------------+--------+---------------------------------------------------------------------+-----------------------------------------------------------------+", + ) + .await; + + let failure = + execute_with_arguments(instance, "missing_no_arguments", FunctionArguments::None).await; + assert!(failure.is_err()); + + assert_admin_event( + instance, + "missing_no_arguments", + "\ ++----------------+----------+----------------------+--------+---------------+---------------------------------------------------------------+ +| type | actor | admin_function_name | status | payload | output | ++----------------+----------+----------------------+--------+---------------+---------------------------------------------------------------+ +| admin_function | greptime | missing_no_arguments | Failed | {\"version\":1} | {\"error\":\"0: Admin function not found: missing_no_arguments\"} | ++----------------+----------+----------------------+--------+---------------+---------------------------------------------------------------+", + ) + .await; + + let Statement::Query(query) = ParserContext::create_with_dialect( + "SELECT 1", + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap() + .remove(0) else { + panic!("expected query statement") + }; + let failure = execute_with_arguments( + instance, + "missing_subquery", + FunctionArguments::Subquery(Box::new(query.inner)), + ) + .await; + assert!(failure.is_err()); + + assert_admin_event( + instance, + "missing_subquery", + "\ ++----------------+----------+---------------------+--------+------------------------------------------------------------------+-----------------------------------------------------------+ +| type | actor | admin_function_name | status | payload | output | ++----------------+----------+---------------------+--------+------------------------------------------------------------------+-----------------------------------------------------------+ +| admin_function | greptime | missing_subquery | Failed | {\"arguments\":{\"type\":\"subquery\",\"value\":\"SELECT 1\"},\"version\":1} | {\"error\":\"0: Admin function not found: missing_subquery\"} | ++----------------+----------+---------------------+--------+------------------------------------------------------------------+-----------------------------------------------------------+", + ) + .await; +} + +async fn assert_admin_event(instance: &Arc, function: &str, expected: &str) { + let query = format!( + r#"SELECT + type, + actor, + admin_function_name, + admin_function_status AS status, + json_to_string(payload) AS payload, + json_to_string(admin_function_output) AS output +FROM greptime_private.events +WHERE type = 'admin_function' + AND admin_function_name = '{function}' +ORDER BY timestamp DESC +LIMIT 1"# + ); + assert_eventually_eq(instance, &query, expected).await; +} + +async fn execute( + instance: &Arc, + sql: &str, +) -> servers::error::Result { + instance.do_query(sql, QueryContext::arc()).await.remove(0) +} + +async fn execute_with_arguments( + instance: &Arc, + function: &str, + arguments: FunctionArguments, +) -> operator::error::Result { + let Statement::Admin(mut statement) = ParserContext::create_with_dialect( + &format!("ADMIN {function}()"), + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap() + .remove(0) else { + panic!("expected ADMIN statement") + }; + let Admin::Func(admin_function) = &mut statement; + admin_function.args = arguments; + + instance + .statement_executor() + .execute_sql(Statement::Admin(statement), QueryContext::arc()) + .await +} diff --git a/tests-integration/tests/main.rs b/tests-integration/tests/main.rs index 1925b46d453..3112cb72c32 100644 --- a/tests-integration/tests/main.rs +++ b/tests-integration/tests/main.rs @@ -14,6 +14,7 @@ #![recursion_limit = "256"] +mod admin_function_event; mod database_ddl_event; mod event_recorder_test_util; #[macro_use]