rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-28 22:56:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: Support procedure_state() using the last procedure ID (#9876)

Browse Source 
* feat: Support procedure_state() using the last procedure ID

Signed-off-by: hahaha <a310608318@gmail.com>

* feat: Fix and add tests according to review comments

Signed-off-by: hahaha <a310608318@gmail.com>

* fix: the zero-argument signature in ADMIN swallows explicit arguments

Signed-off-by: hahaha <a310608318@gmail.com>

* fix: fix the CI failures caused by SQL version issues

Signed-off-by: hahaha <a310608318@gmail.com>

* fix: apply the changes based on the review

Signed-off-by: hahaha <a310608318@gmail.com>

---------

Signed-off-by: hahaha <a310608318@gmail.com>
This commit is contained in:
Chenglong Zheng
2026-10-27 07:27:55 +00:00
committed by GitHub
parent e6b60644b3
commit 2117f8b059
15 changed files with 788 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/common/function/src/admin/migrate_region.rs
+9 -2
View File
@@ -94,7 +94,10 @@ pub(crate) async fn migrate_region(
.await?;
match pid {
Some(pid) => Ok(Value::from(pid)),
Some(pid) => {
query_ctx.set_last_procedure_id(pid.clone());
Ok(Value::from(pid))
}
None => Ok(Value::Null),
}
}
@@ -184,7 +187,10 @@ mod tests {
#[tokio::test]
async fn test_migrate_region() {
let factory: ScalarFunctionFactory = MigrateRegionFunction::factory().into();
let provider = factory.provide(FunctionContext::mock());
let context = FunctionContext::mock();
let query_ctx = context.query_ctx.clone();
query_ctx.set_last_procedure_id("previous-id".to_string());
let provider = factory.provide(context);
let f = provider.as_async().unwrap();
let func_args = datafusion::logical_expr::ScalarFunctionArgs {
@@ -203,6 +209,7 @@ mod tests {
config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
};
let result = f.invoke_async_with_args(func_args).await.unwrap();
assert_eq!(query_ctx.last_procedure_id().as_deref(), Some("test_pid"));
match result {
ColumnarValue::Array(array) => {
src/common/function/src/admin/reconcile_catalog.rs
+9 -2
View File
@@ -96,7 +96,10 @@ pub(crate) async fn reconcile_catalog(
})
.await?;
match pid {
Some(pid) => Ok(Value::from(pid)),
Some(pid) => {
query_ctx.set_last_procedure_id(pid.clone());
Ok(Value::from(pid))
}
None => Ok(Value::Null),
}
}
@@ -138,7 +141,10 @@ mod tests {
// reconcile_catalog()
let factory: ScalarFunctionFactory = ReconcileCatalogFunction::factory().into();
let provider = factory.provide(FunctionContext::mock());
let context = FunctionContext::mock();
let query_ctx = context.query_ctx.clone();
query_ctx.set_last_procedure_id("previous-id".to_string());
let provider = factory.provide(context);
let f = provider.as_async().unwrap();
let func_args = datafusion::logical_expr::ScalarFunctionArgs {
@@ -150,6 +156,7 @@ mod tests {
};
let result = f.invoke_async_with_args(func_args).await.unwrap();
assert_eq!(query_ctx.last_procedure_id().as_deref(), Some("test_pid"));
match result {
ColumnarValue::Array(array) => {
let result_array = array.as_any().downcast_ref::<StringArray>().unwrap();
src/common/function/src/admin/reconcile_database.rs
+9 -2
View File
@@ -105,7 +105,10 @@ pub(crate) async fn reconcile_database(
})
.await?;
match pid {
Some(pid) => Ok(Value::from(pid)),
Some(pid) => {
query_ctx.set_last_procedure_id(pid.clone());
Ok(Value::from(pid))
}
None => Ok(Value::Null),
}
}
@@ -148,7 +151,10 @@ mod tests {
// reconcile_database(database_name)
let factory: ScalarFunctionFactory = ReconcileDatabaseFunction::factory().into();
let provider = factory.provide(FunctionContext::mock());
let context = FunctionContext::mock();
let query_ctx = context.query_ctx.clone();
query_ctx.set_last_procedure_id("previous-id".to_string());
let provider = factory.provide(context);
let f = provider.as_async().unwrap();
let func_args = datafusion::logical_expr::ScalarFunctionArgs {
@@ -161,6 +167,7 @@ mod tests {
config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
};
let result = f.invoke_async_with_args(func_args).await.unwrap();
assert_eq!(query_ctx.last_procedure_id().as_deref(), Some("test_pid"));
match result {
ColumnarValue::Array(array) => {
let result_array = array.as_any().downcast_ref::<StringArray>().unwrap();
src/common/function/src/admin/reconcile_table.rs
+9 -2
View File
@@ -86,7 +86,10 @@ pub(crate) async fn reconcile_table(
})
.await?;
match pid {
Some(pid) => Ok(Value::from(pid)),
Some(pid) => {
query_ctx.set_last_procedure_id(pid.clone());
Ok(Value::from(pid))
}
None => Ok(Value::Null),
}
}
@@ -121,7 +124,10 @@ mod tests {
// reconcile_table(table_name)
let factory: ScalarFunctionFactory = ReconcileTableFunction::factory().into();
let provider = factory.provide(FunctionContext::mock());
let context = FunctionContext::mock();
let query_ctx = context.query_ctx.clone();
query_ctx.set_last_procedure_id("previous-id".to_string());
let provider = factory.provide(context);
let f = provider.as_async().unwrap();
let func_args = datafusion::logical_expr::ScalarFunctionArgs {
@@ -134,6 +140,7 @@ mod tests {
config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
};
let result = f.invoke_async_with_args(func_args).await.unwrap();
assert_eq!(query_ctx.last_procedure_id().as_deref(), Some("test_pid"));
match result {
ColumnarValue::Array(array) => {
let result_array = array.as_any().downcast_ref::<StringArray>().unwrap();
src/common/function/src/state.rs
+14 -2
View File
@@ -34,7 +34,16 @@ impl FunctionState {
/// Create a mock [`FunctionState`] for test.
#[cfg(any(test, feature = "testing"))]
pub fn mock() -> Self {
Self::mock_with_procedure_state_counter(Default::default())
}
/// Creates a mock state that counts procedure state queries.
#[cfg(any(test, feature = "testing"))]
pub fn mock_with_procedure_state_counter(
calls: std::sync::Arc<std::sync::atomic::AtomicUsize>,
) -> Self {
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use api::v1::meta::{ProcedureStatus, ReconcileRequest};
use async_trait::async_trait;
@@ -55,7 +64,9 @@ impl FunctionState {
use table::table_name::TableName;
use crate::handlers::{FlowServiceHandler, ProcedureServiceHandler, TableMutationHandler};
struct MockProcedureServiceHandler;
struct MockProcedureServiceHandler {
calls: Arc<AtomicUsize>,
}
struct MockTableMutationHandler;
struct MockFlowServiceHandler;
const ROWS: usize = 42;
@@ -83,6 +94,7 @@ impl FunctionState {
}
async fn query_procedure_state(&self, _pid: &str) -> Result<ProcedureStateResponse> {
self.calls.fetch_add(1, Ordering::Relaxed);
Ok(ProcedureStateResponse {
status: ProcedureStatus::Done.into(),
error: "OK".to_string(),
@@ -219,7 +231,7 @@ impl FunctionState {
Self {
plugins: Plugins::default(),
table_mutation_handler: Some(Arc::new(MockTableMutationHandler)),
procedure_service_handler: Some(Arc::new(MockProcedureServiceHandler)),
procedure_service_handler: Some(Arc::new(MockProcedureServiceHandler { calls })),
flow_service_handler: Some(Arc::new(MockFlowServiceHandler)),
}
}
src/common/function/src/system/procedure_state.rs
+103 -19
View File
@@ -20,11 +20,11 @@ use common_query::error::{
InvalidFuncArgsSnafu, MissingProcedureServiceHandlerSnafu, Result,
UnsupportedInputDataTypeSnafu,
};
use datafusion_expr::{Signature, Volatility};
use datafusion_expr::{Signature, TypeSignature, Volatility};
use datatypes::prelude::*;
use serde::Serialize;
use session::context::QueryContextRef;
use snafu::ensure;
use snafu::{OptionExt, ensure};
use crate::handlers::ProcedureServiceHandlerRef;
@@ -35,8 +35,10 @@ struct ProcedureStateJson {
error: Option<String>,
}
/// A function to query procedure state by its id.
/// Such as `procedure_state(pid)`.
/// Queries procedure state using an explicit procedure ID, or the latest
/// procedure ID recorded in the Session.
/// The zero-argument form reads the procedure ID from the Session at execution
/// time. State lookups do not update the recorded procedure ID.
#[admin_fn(
name = ProcedureStateFunction,
display_name = procedure_state,
@@ -45,25 +47,35 @@ struct ProcedureStateJson {
)]
pub(crate) async fn procedure_state(
procedure_service_handler: &ProcedureServiceHandlerRef,
_ctx: &QueryContextRef,
ctx: &QueryContextRef,
params: &[ValueRef<'_>],
) -> Result<Value> {
ensure!(
params.len() == 1,
params.len() <= 1,
InvalidFuncArgsSnafu {
err_msg: format!(
"The length of the args is not correct, expect 1, have: {}",
"The length of the args is not correct, expect 0 or 1, have: {}",
params.len()
),
}
);
let ValueRef::String(pid) = params[0] else {
return UnsupportedInputDataTypeSnafu {
function: "procedure_state",
datatypes: params.iter().map(|v| v.data_type()).collect::<Vec<_>>(),
let last_procedure_id;
let pid = match params {
[] => {
last_procedure_id = ctx.last_procedure_id().context(InvalidFuncArgsSnafu {
err_msg: "No procedure ID has been returned in this session; pass an explicit procedure ID",
})?;
last_procedure_id.as_str()
}
[ValueRef::String(pid)] => *pid,
_ => {
return UnsupportedInputDataTypeSnafu {
function: "procedure_state",
datatypes: params.iter().map(|v| v.data_type()).collect::<Vec<_>>(),
}
.fail();
}
.fail();
};
let ProcedureStateResponse { status, error, .. } =
@@ -82,20 +94,29 @@ pub(crate) async fn procedure_state(
}
fn signature() -> Signature {
Signature::uniform(1, vec![ArrowDataType::Utf8], Volatility::Immutable)
Signature::one_of(
vec![
TypeSignature::Exact(vec![]),
TypeSignature::Exact(vec![ArrowDataType::Utf8]),
],
Volatility::Volatile,
)
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use arrow::array::StringArray;
use arrow::datatypes::{DataType, Field};
use datafusion_expr::ColumnarValue;
use datafusion_expr::type_coercion::functions::fields_with_udf;
use super::*;
use crate::function::FunctionContext;
use crate::function_factory::ScalarFunctionFactory;
use crate::state::FunctionState;
#[test]
fn test_procedure_state_misc() {
@@ -103,12 +124,75 @@ mod tests {
let f = factory.provide(FunctionContext::mock());
assert_eq!("procedure_state", f.name());
assert_eq!(DataType::Utf8, f.return_type(&[]).unwrap());
assert!(matches!(f.signature(),
datafusion_expr::Signature {
type_signature: datafusion_expr::TypeSignature::Uniform(1, valid_types),
volatility: datafusion_expr::Volatility::Immutable,
..
} if valid_types == &vec![ArrowDataType::Utf8]));
assert_eq!(
f.signature(),
&Signature::one_of(
vec![
TypeSignature::Exact(vec![]),
TypeSignature::Exact(vec![ArrowDataType::Utf8]),
],
Volatility::Volatile,
)
);
let field = Arc::new(Field::new("pid", DataType::Utf8, false));
assert!(fields_with_udf(&[field.clone(), field], &f).is_err());
}
#[tokio::test]
async fn test_procedure_state_args() {
let context = FunctionContext::mock();
let handler = context.state.procedure_service_handler.as_ref().unwrap();
let ctx = &context.query_ctx;
let err = procedure_state(handler, ctx, &[]).await.unwrap_err();
assert!(
err.to_string()
.contains("No procedure ID has been returned")
);
ctx.set_last_procedure_id("session-id".to_string());
for (params, succeeds) in [
(vec![], true),
(vec![ValueRef::String("explicit-id")], true),
(vec![ValueRef::Null], false),
(vec![ValueRef::String("a"), ValueRef::String("b")], false),
] {
let result = procedure_state(handler, ctx, &params).await;
assert_eq!(result.is_ok(), succeeds, "{params:?}: {result:?}");
assert_eq!(ctx.last_procedure_id().as_deref(), Some("session-id"));
}
}
#[tokio::test]
async fn test_procedure_state_zero_arg_rows() {
let calls = Arc::new(AtomicUsize::new(0));
let context = FunctionContext {
state: Arc::new(FunctionState::mock_with_procedure_state_counter(
calls.clone(),
)),
..Default::default()
};
context
.query_ctx
.set_last_procedure_id("session-id".to_string());
let factory: ScalarFunctionFactory = ProcedureStateFunction::factory().into();
let function = factory.provide(context);
let f = function.as_async().unwrap();
for number_rows in [0, 1, 2] {
calls.store(0, Ordering::Relaxed);
let args = datafusion_expr::ScalarFunctionArgs {
args: vec![],
arg_fields: vec![],
return_field: Arc::new(Field::new("result", DataType::Utf8, true)),
number_rows,
config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
};
let ColumnarValue::Array(result) = f.invoke_async_with_args(args).await.unwrap() else {
panic!("expected array output");
};
assert_eq!(result.len(), number_rows);
assert_eq!(calls.load(Ordering::Relaxed), number_rows);
}
}
#[tokio::test]
src/common/macro/src/admin_fn.rs
+8 -19
View File
@@ -311,12 +311,8 @@ fn build_struct(
// Safety check: Ensure under the `greptime` catalog for security
#user_path::ensure_greptime!(self.func_ctx);
let columns_num = columns.len();
let rows_num = if columns.is_empty() {
1
} else {
columns[0].len()
};
// Zero-argument functions also execute once per input row.
let rows_num = columns.first().map_or(args.number_rows, |column| column.len());
#validate_rows
use snafu::{OptionExt, ResultExt};
@@ -333,22 +329,15 @@ fn build_struct(
let mut builder = store_api::storage::ConcreteDataType::#ret()
.create_mutable_vector(rows_num);
if columns_num == 0 {
let result = #fn_name(handler, query_ctx, &[]).await
for i in 0..rows_num {
let args: Vec<_> = columns.iter()
.map(|vector| vector.get_ref(i))
.collect();
let result = #fn_name(handler, query_ctx, &args).await
.map_err(|e| datafusion_common::DataFusionError::Execution(e.output_msg()))?;
builder.push_value_ref(&result.as_value_ref());
} else {
for i in 0..rows_num {
let args: Vec<_> = columns.iter()
.map(|vector| vector.get_ref(i))
.collect();
let result = #fn_name(handler, query_ctx, &args).await
.map_err(|e| datafusion_common::DataFusionError::Execution(e.output_msg()))?;
builder.push_value_ref(&result.as_value_ref());
}
}
let result_vector = builder.to_vector();
src/operator/src/statement/admin.rs
+71 -1
View File
@@ -306,6 +306,17 @@ fn args_to_vector(
let tz = query_ctx.timezone();
match type_signature {
TypeSignature::Nullary => {
ensure!(
args.is_empty(),
error::FunctionArityMismatchSnafu {
actual: args.len(),
expected: 0_usize,
}
);
Ok(vec![])
}
TypeSignature::Variadic(valid_types) => {
values_to_vectors_by_valid_types(valid_types, args, Some(&tz))
}
@@ -323,6 +334,13 @@ fn args_to_vector(
}
TypeSignature::Exact(data_types) => {
ensure!(
args.len() == data_types.len(),
error::FunctionArityMismatchSnafu {
actual: args.len(),
expected: data_types.len(),
}
);
values_to_vectors_by_exact_types(data_types, args, Some(&tz))
}
@@ -450,9 +468,61 @@ fn try_get_data_type_for_sql_value(value: &SqlValue) -> Result<ArrowDataType> {
mod tests {
use std::sync::Arc;
use datafusion_expr::TypeSignature;
use datatypes::arrow::datatypes::DataType as ArrowDataType;
use datatypes::value::Value;
use datatypes::vectors::{Int32Vector, VectorRef};
use session::context::QueryContext;
use sql::ast::Value as SqlValue;
use sql::dialect::GreptimeDbDialect;
use sql::parser::{ParseOptions, ParserContext};
use sql::statements::statement::Statement;
use crate::statement::admin::immediate_result;
use crate::error::Error;
use crate::statement::admin::{admin_output_schema, args_to_vector, immediate_result};
#[test]
fn test_args_to_vector_arity() {
let query_ctx = QueryContext::arc();
let value = SqlValue::SingleQuotedString("explicit-id".to_string());
for (signature, expected) in [
(TypeSignature::Nullary, 0),
(TypeSignature::Exact(vec![]), 0),
(TypeSignature::Exact(vec![ArrowDataType::Utf8]), 1),
(TypeSignature::Exact(vec![ArrowDataType::Utf8; 2]), 2),
] {
for actual in 0..=3 {
let result = args_to_vector(&signature, &vec![&value; actual], &query_ctx);
if actual == expected {
let vectors = result.unwrap();
assert_eq!(vectors.len(), expected);
for vector in vectors {
assert_eq!(vector.get(0), Value::from("explicit-id"));
}
} else {
assert!(matches!(
result,
Err(Error::FunctionArityMismatch { actual: a, expected: e })
if a == actual && e == expected
));
}
}
}
}
#[test]
fn test_admin_reconcile_catalog_nullary() {
let Statement::Admin(statement) = ParserContext::create_with_dialect(
"ADMIN reconcile_catalog()",
&GreptimeDbDialect {},
ParseOptions::default(),
)
.unwrap()
.remove(0) else {
panic!("expected ADMIN statement");
};
assert!(admin_output_schema(&statement, &QueryContext::arc()).is_some());
}
#[test]
fn empty_admin_function_result_has_no_immediate_value() {
src/operator/src/statement/ddl.rs
+10 -5
View File
@@ -146,8 +146,7 @@ const ALLOWED_FLOW_OPTIONS: [&str; 2] = [
FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY,
];
fn build_procedure_id_output(procedure_id: Vec<u8>) -> Result<Output> {
let procedure_id = String::from_utf8_lossy(&procedure_id).to_string();
fn build_procedure_id_output(procedure_id: &str) -> Result<Output> {
let vector: VectorRef = Arc::new(StringVector::from(vec![procedure_id]));
let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
"Procedure ID",
@@ -1775,12 +1774,15 @@ impl StatementExecutor {
table_id,
time_ranges,
ddl_options,
query_context,
query_context.clone(),
)
.await?;
if !ddl_options.wait {
return build_procedure_id_output(response.key);
let procedure_id = String::from_utf8_lossy(&response.key).into_owned();
let output = build_procedure_id_output(&procedure_id)?;
query_context.set_last_procedure_id(procedure_id);
return Ok(output);
}
Ok(Output::new_with_affected_rows(0))
@@ -2062,7 +2064,10 @@ impl StatementExecutor {
.context(error::ExecuteDdlSnafu)?;
if !ddl_options.wait {
return build_procedure_id_output(response.key);
let procedure_id = String::from_utf8_lossy(&response.key).into_owned();
let output = build_procedure_id_output(&procedure_id)?;
query_context.set_last_procedure_id(procedure_id);
return Ok(output);
}
// Only invalidate cache if wait is true.
src/session/src/context.rs
+18
View File
@@ -437,6 +437,24 @@ impl QueryContext {
self.mutable_session_data.write().unwrap().read_preference = read_preference;
}
/// Returns the latest procedure ID recorded in this Session.
pub fn last_procedure_id(&self) -> Option<String> {
self.mutable_session_data
.read()
.unwrap()
.last_procedure_id
.clone()
}
/// Records a procedure ID after successful statement execution and
/// construction of its ID result.
///
/// Recording precedes `post_execute` hooks and protocol transmission, which
/// may still fail. It does not confirm that the client received the ID.
pub fn set_last_procedure_id(&self, procedure_id: String) {
self.mutable_session_data.write().unwrap().last_procedure_id = Some(procedure_id);
}
pub fn current_user(&self) -> UserInfoRef {
self.mutable_session_data.read().unwrap().user_info.clone()
}
src/session/src/lib.rs
+4
View File
@@ -61,6 +61,9 @@ pub(crate) struct MutableInner {
timezone: Timezone,
query_timeout: Option<Duration>,
read_preference: ReadPreference,
/// Latest procedure ID recorded during statement execution, before
/// `post_execute` hooks and protocol transmission.
last_procedure_id: Option<String>,
/// Request-level WAL policy for ordinary inserts.
skip_wal: bool,
#[debug(skip)]
@@ -77,6 +80,7 @@ impl Default for MutableInner {
timezone: get_timezone(None).clone(),
query_timeout: None,
read_preference: ReadPreference::Leader,
last_procedure_id: None,
skip_wal: false,
cursors: HashMap::with_capacity(0),
warnings: VecDeque::new(),
tests-integration/tests/repartition.rs
+51 -7
View File
@@ -16,23 +16,25 @@ use std::collections::BTreeSet;
use std::sync::Arc;
use std::time::Duration;
use client::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
use client::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, OutputData};
use common_error::root_source;
use common_event_recorder::{PersistentEventContext, TriggerReason};
use common_meta::key::table_name::TableNameKey;
use common_procedure::{ProcedureContext, ProcedureWithId, watcher};
use common_procedure::{ProcedureContext, ProcedureId, ProcedureWithId, watcher};
use common_query::Output;
use common_telemetry::info;
use common_test_util::recordbatch::check_output_stream;
use common_test_util::temp_dir::{TempDir, create_temp_dir};
use common_wal::config::DatanodeWalConfig;
use datatypes::arrow::array::AsArray;
use frontend::instance::Instance;
use meta_srv::gc::{self, BatchGcProcedure, GcSchedulerOptions, GcTickerRef};
use meta_srv::metasrv::Metasrv;
use mito2::gc::GcConfig;
use servers::error::Result as ServerResult;
use servers::query_handler::sql::SqlQueryHandler;
use session::context::{QueryContext, QueryContextRef};
use session::Session;
use session::context::{Channel, QueryContext, QueryContextRef};
use store_api::codec::PrimaryKeyEncoding;
use store_api::storage::{RegionId, TableId};
use tests_integration::cluster::{GreptimeDbCluster, GreptimeDbClusterBuilder};
@@ -149,6 +151,48 @@ macro_rules! repartition_tests {
};
}
#[tokio::test(flavor = "multi_thread")]
async fn test_repartition_procedure_id_session() {
let (cluster, _home_guard) = repartition_cluster("repartition_procedure_id_session").await;
let instance = cluster.fe_instance();
let session = Session::new(None, Channel::Postgres, Default::default(), 0);
let query_ctx = session.new_query_context();
run_sql(
instance,
r#"CREATE TABLE repartition_session_test ("id" INT PRIMARY KEY, ts TIMESTAMP TIME INDEX)
PARTITION ON COLUMNS ("id") ("id" < 10, "id" >= 10)"#,
session.new_query_context(),
)
.await
.unwrap();
query_ctx.set_last_procedure_id("previous-id".to_string());
let sql = r#"ALTER TABLE repartition_session_test REPARTITION ("id" < 10)
INTO ("id" < 5, "id" >= 5 AND "id" < 10) WITH (wait=false)"#;
let OutputData::RecordBatches(batches) = run_sql(instance, sql, session.new_query_context())
.await
.unwrap()
.data
else {
panic!("expected Procedure ID result");
};
let batch = batches.iter().next().unwrap();
let id = batch
.column_by_name("Procedure ID")
.unwrap()
.as_string::<i32>()
.value(0);
assert_eq!(query_ctx.last_procedure_id().as_deref(), Some(id));
let mut watcher = cluster
.metasrv
.procedure_manager()
.procedure_watcher(ProcedureId::parse_str(id).unwrap())
.unwrap();
tokio::time::timeout(Duration::from_secs(60), watcher::wait(&mut watcher))
.await
.unwrap()
.unwrap();
}
#[tokio::test(flavor = "multi_thread")]
async fn test_repartition_physical_metric_without_logical_table() {
common_telemetry::init_default_ut_logging();
@@ -247,7 +291,7 @@ async fn test_repartition_physical_metric_without_logical_table() {
/// COUNT must use visible rows rather than the full row counts of shared SSTs.
#[tokio::test(flavor = "multi_thread")]
async fn test_repartition_append_count_file() {
let (cluster, _home_guard) = append_count_cluster("repartition_append_count").await;
let (cluster, _home_guard) = repartition_cluster("repartition_append_count").await;
let instance = cluster.fe_instance();
let table = "count_repartition";
prepare_append_count_table(instance, table, "").await;
@@ -265,7 +309,7 @@ async fn test_repartition_append_count_file() {
#[tokio::test(flavor = "multi_thread")]
async fn test_split_append_count_file() {
let (cluster, _home_guard) = append_count_cluster("split_append_count").await;
let (cluster, _home_guard) = repartition_cluster("split_append_count").await;
let instance = cluster.fe_instance();
let table = "count_split";
prepare_append_count_table(
@@ -288,7 +332,7 @@ async fn test_split_append_count_file() {
#[tokio::test(flavor = "multi_thread")]
async fn test_repartition_append_count_memtable_file() {
let (cluster, home_guard) = append_count_cluster("repartition_append_count_memtable").await;
let (cluster, home_guard) = repartition_cluster("repartition_append_count_memtable").await;
assert!(home_guard.path().is_dir());
let instance = cluster.fe_instance();
let table = "count_repartition_memtable";
@@ -305,7 +349,7 @@ async fn test_repartition_append_count_memtable_file() {
assert!(home_guard.path().is_dir());
}
async fn append_count_cluster(name: &str) -> (GreptimeDbCluster, Arc<TempDir>) {
async fn repartition_cluster(name: &str) -> (GreptimeDbCluster, Arc<TempDir>) {
common_telemetry::init_default_ut_logging();
let (store_config, _guard) = get_test_store_config(&StorageType::File);
let home_dir = Arc::new(create_temp_dir(name));
tests-integration/tests/truncate.rs
+288 -20
View File
@@ -17,15 +17,23 @@ use std::time::Duration;
use client::OutputData;
use common_error::ext::ErrorExt;
use common_error::snafu::ResultExt;
use common_error::status_code::StatusCode;
use common_meta::lock_key::TableLock;
use common_procedure::{
LockKey, Procedure, ProcedureId, ProcedureState, ProcedureWithId, Status, watcher,
};
use common_recordbatch::RecordBatches;
use datatypes::arrow::array::AsArray;
use frontend::instance::Instance;
use servers::error::{CollectRecordbatchSnafu, Result as ServerResult};
use servers::query_handler::sql::SqlQueryHandler;
use session::Session;
use session::context::Channel;
use sqlx::{Connection, Executor, Row, Statement};
use tests_integration::cluster::GreptimeDbClusterBuilder;
use tests_integration::standalone::GreptimeDbStandaloneBuilder;
use tests_integration::test_util::{execute_sql, execute_sql_and_expect, try_execute_sql};
use tests_integration::test_util::{execute_sql, execute_sql_and_expect};
use tokio::sync::Notify;
/// Holds the table lock until explicitly released, keeping truncate pending.
@@ -126,37 +134,41 @@ async fn test_truncate_wait_and_timeout(distributed: bool) {
.await
.unwrap();
let output = execute_sql(
let session = Session::new(None, Channel::Mysql, Default::default(), 0);
let query_ctx = session.new_query_context();
let output = run_sql(
&frontend,
"TRUNCATE TABLE truncate_wait_test WITH (wait=false, timeout='10s')",
&session,
)
.await;
let OutputData::RecordBatches(batches) = output.data else {
panic!("expected Procedure ID result");
};
let batch = batches.iter().next().unwrap();
let id = batch
.column_by_name("Procedure ID")
.unwrap()
.as_string::<i32>()
.value(0);
let id = ProcedureId::parse_str(id).unwrap();
let state = execute_sql(&frontend, &format!("SELECT procedure_state('{id}')"))
.await
.data
.pretty_print()
.await;
assert!(state.contains(r#""status":"Running""#), "{state}");
.await
.unwrap();
let id = procedure_id_from_output(output).to_string();
assert_eq!(query_ctx.last_procedure_id(), Some(id.clone()));
for sql in [
format!("SELECT procedure_state('{id}')"),
"SELECT procedure_state()".to_string(),
] {
let state = run_sql(&frontend, &sql, &session)
.await
.unwrap()
.pretty_print()
.await;
assert!(state.contains(r#""status":"Running""#), "{state}");
assert_eq!(query_ctx.last_procedure_id(), Some(id.clone()));
}
// The same table lock holds a synchronous truncate past its wait deadline.
// Give metadata reads and submission a generous budget. Using the same budget
// also verifies that the async procedure keeps running after it expires.
let err = try_execute_sql(
let err = run_sql(
&frontend,
"TRUNCATE TABLE truncate_wait_test WITH (wait=true, timeout='10s')",
&session,
)
.await
.unwrap_err();
assert_eq!(query_ctx.last_procedure_id(), Some(id.clone()));
if distributed {
assert!(matches!(
err.status_code(),
@@ -222,4 +234,260 @@ async fn test_truncate_wait_and_timeout(distributed: bool) {
.pretty_print()
.await
);
for (sql, succeeds) in [
("SELECT 1", true),
("TRUNCATE TABLE truncate_wait_test WITH (wait=true)", true),
(
"TRUNCATE TABLE nonexistent_session_table WITH (wait=false)",
false,
),
("SELECT procedure_state(NULL)", false),
("SELECT procedure_state('')", false),
("SELECT procedure_state('invalid-id')", false),
] {
let result = run_sql(&frontend, sql, &session).await;
assert_eq!(result.is_ok(), succeeds, "{sql}: {result:?}");
assert_eq!(query_ctx.last_procedure_id(), Some(id.clone()));
}
let other_session = Session::new(None, Channel::Mysql, Default::default(), 1);
let state = run_sql(
&frontend,
&format!("SELECT procedure_state('{id}')"),
&other_session,
)
.await
.unwrap()
.pretty_print()
.await;
assert!(state.contains(r#""status":"Done""#), "{state}");
let err = run_sql(&frontend, "SELECT procedure_state()", &other_session)
.await
.unwrap_err();
assert!(
err.output_msg()
.contains("No procedure ID has been returned")
);
let output = run_sql(
&frontend,
"TRUNCATE TABLE truncate_wait_test WITH (wait=false)",
&session,
)
.await
.unwrap();
let new_id = procedure_id_from_output(output);
assert_ne!(new_id.to_string(), id);
assert_eq!(query_ctx.last_procedure_id(), Some(new_id.to_string()));
let mut watcher = manager.procedure_watcher(new_id).unwrap();
tokio::time::timeout(Duration::from_secs(10), watcher::wait(&mut watcher))
.await
.unwrap()
.unwrap();
}
#[rstest::rstest]
#[case::mysql("MYSQL_SERVER", "mysql")]
#[case::postgres("POSTGRES_SERVER", "postgres")]
#[tokio::test(flavor = "multi_thread")]
async fn test_procedure_state_prepared_and_admin(
#[case] server_name: &str,
#[case] protocol: &str,
) {
let instance = GreptimeDbStandaloneBuilder::new("procedure_state_prepared")
.build()
.await;
let mut opts = instance.opts.clone();
opts.http.addr = "127.0.0.1:0".into();
opts.grpc.bind_addr = "127.0.0.1:0".into();
opts.mysql.addr = "127.0.0.1:0".into();
opts.postgres.addr = "127.0.0.1:0".into();
let mut servers =
frontend::server::Services::new(opts, instance.fe_instance().clone(), Default::default())
.build()
.unwrap();
servers.start_all().await.unwrap();
sqlx::any::install_default_drivers();
let url = format!("{protocol}://{}/public", servers.addr(server_name).unwrap());
let mut connection = sqlx::AnyConnection::connect(&url).await.unwrap();
for table in ["prepared_async", "prepared_sync"] {
// Safety: the table names are fixed test-local identifiers.
connection
.execute(sqlx::AssertSqlSafe(format!(
"CREATE TABLE {table} (ts TIMESTAMP TIME INDEX)"
)))
.await
.unwrap();
}
// Prepare once, before any ID exists, and reuse the same statement below.
let statement = connection
.prepare(sqlx::SqlStr::from_static(
"SELECT procedure_state() AS state",
))
.await
.unwrap();
let manager = &instance.procedure_manager;
let id: String = connection
.fetch_one("TRUNCATE TABLE prepared_async WITH (wait=false)")
.await
.unwrap()
.get(0);
let mut completed = manager
.procedure_watcher(ProcedureId::parse_str(&id).unwrap())
.unwrap();
tokio::time::timeout(Duration::from_secs(10), watcher::wait(&mut completed))
.await
.unwrap()
.unwrap();
assert_eq!(prepared_state(&mut connection, &statement).await, "Done");
// An explicit ADMIN ID works without history and does not record that ID.
let mut other_connection = sqlx::AnyConnection::connect(&url).await.unwrap();
for explicit in [false, true, false] {
let sql = if explicit {
format!("ADMIN procedure_state('{id}')")
} else {
"ADMIN procedure_state()".to_string()
};
// Safety: the SQL contains only a procedure ID returned by this test.
let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str()))
.fetch_one(&mut other_connection)
.await;
if explicit {
assert_eq!(result.unwrap().get::<String, _>(0), r#"{"status":"Done"}"#);
} else {
let err = result.err().expect("expected missing procedure ID error");
assert!(
err.to_string()
.contains("No procedure ID has been returned"),
"{err}"
);
}
}
other_connection.close().await.unwrap();
let table_id = instance
.fe_instance()
.catalog_manager()
.table("greptime", "public", "prepared_async", None)
.await
.unwrap()
.unwrap()
.table_info()
.table_id();
let started = Arc::new(Notify::new());
let release = Arc::new(Notify::new());
let mut blocker = manager
.submit(ProcedureWithId::with_random_id(Box::new(
HoldTruncateTableLock {
table_id,
started: started.clone(),
release: release.clone(),
},
)))
.await
.unwrap();
tokio::time::timeout(Duration::from_secs(10), started.notified())
.await
.unwrap();
let new_id: String = connection
.fetch_one("TRUNCATE TABLE prepared_async WITH (wait=false)")
.await
.unwrap()
.get(0);
assert_ne!(id, new_id);
let mut pending = manager
.procedure_watcher(ProcedureId::parse_str(&new_id).unwrap())
.unwrap();
// A is Done and B is Running, so retaining A cannot pass this assertion.
assert_eq!(prepared_state(&mut connection, &statement).await, "Running");
for (sql, expected) in [
("ADMIN procedure_state()".to_string(), "Running"),
(format!("ADMIN procedure_state('{id}')"), "Done"),
("ADMIN procedure_state()".to_string(), "Running"),
] {
// Safety: the SQL contains only a procedure ID returned by this test.
let state: String = connection
.fetch_one(sqlx::AssertSqlSafe(sql.as_str()))
.await
.unwrap()
.get(0);
assert_eq!(state, format!(r#"{{"status":"{expected}"}}"#), "{sql}");
}
for sql in [
"ADMIN procedure_state('a', 'b')",
"ADMIN procedure_state(NULL)",
"ADMIN procedure_state('')",
"ADMIN procedure_state('invalid-id')",
] {
assert!(connection.fetch_one(sql).await.is_err(), "{sql}");
}
connection
.execute("TRUNCATE TABLE prepared_sync WITH (wait=true)")
.await
.unwrap();
assert_eq!(prepared_state(&mut connection, &statement).await, "Running");
release.notify_one();
tokio::time::timeout(Duration::from_secs(10), async {
watcher::wait(&mut blocker).await.unwrap();
watcher::wait(&mut pending).await.unwrap();
})
.await
.unwrap();
// The ID stays B; only its state changes between executions.
assert_eq!(prepared_state(&mut connection, &statement).await, "Done");
connection.close().await.unwrap();
servers.shutdown_all().await.unwrap();
}
async fn prepared_state(
connection: &mut sqlx::AnyConnection,
statement: &sqlx::any::AnyStatement,
) -> String {
let state: String = statement
.query()
.fetch_one(connection)
.await
.unwrap()
.get(0);
serde_json::from_str::<serde_json::Value>(&state).unwrap()["status"]
.as_str()
.unwrap()
.to_string()
}
async fn run_sql(
frontend: &Arc<Instance>,
sql: &str,
session: &Session,
) -> ServerResult<OutputData> {
let data = frontend
.do_query(sql, session.new_query_context())
.await
.remove(0)?
.data;
match data {
OutputData::Stream(stream) => Ok(OutputData::RecordBatches(
RecordBatches::try_collect(stream)
.await
.context(CollectRecordbatchSnafu)?,
)),
data => Ok(data),
}
}
fn procedure_id_from_output(output: OutputData) -> ProcedureId {
let OutputData::RecordBatches(batches) = output else {
panic!("expected Procedure ID result");
};
let batch = batches.iter().next().unwrap();
let id = batch
.column_by_name("Procedure ID")
.unwrap()
.as_string::<i32>()
.value(0);
ProcedureId::parse_str(id).unwrap()
}
tests/cases/standalone/common/function/procedure_state_session.result
+124
View File
@@ -0,0 +1,124 @@
-- Each protocol reuses its own connection across requests.
-- SQLNESS PROTOCOL MYSQL
SELECT procedure_state();
Failed to parse query result, err: MySqlError { ERROR 1815 (HY000): (EngineExecuteQuery): Execution error: Invalid function args: No procedure ID has been returned in this session; pass an explicit procedure ID }
-- SQLNESS PROTOCOL POSTGRES
SELECT procedure_state();
Failed to execute query, encountered: Error { kind: Db, cause: Some(DbError { severity: "ERROR", parsed_severity: None, code: SqlState(EP0000), message: "Execution error: Invalid function args: No procedure ID has been returned in this session; pass an explicit procedure ID", detail: None, hint: None, position: None, where_: None, schema: None, table: None, column: None, datatype: None, constraint: None, file: None, line: None, routine: None }) }
CREATE TABLE procedure_state_session (ts TIMESTAMP TIME INDEX);
Affected Rows: 0
-- Synchronous DDL must not record a Procedure ID in the Session.
-- SQLNESS PROTOCOL MYSQL
TRUNCATE TABLE procedure_state_session WITH (wait=true);
affected_rows: 0
-- SQLNESS PROTOCOL MYSQL
SELECT procedure_state();
Failed to parse query result, err: MySqlError { ERROR 1815 (HY000): (EngineExecuteQuery): Execution error: Invalid function args: No procedure ID has been returned in this session; pass an explicit procedure ID }
-- SQLNESS PROTOCOL POSTGRES
TRUNCATE TABLE procedure_state_session WITH (wait=true);
Affected Rows: 0
-- SQLNESS PROTOCOL POSTGRES
SELECT procedure_state();
Failed to execute query, encountered: Error { kind: Db, cause: Some(DbError { severity: "ERROR", parsed_severity: None, code: SqlState(EP0000), message: "Execution error: Invalid function args: No procedure ID has been returned in this session; pass an explicit procedure ID", detail: None, hint: None, position: None, where_: None, schema: None, table: None, column: None, datatype: None, constraint: None, file: None, line: None, routine: None }) }
-- Synchronous DDL preserves the ID returned by the previous request.
-- SQLNESS PROTOCOL MYSQL
-- SQLNESS REPLACE [0-9a-f]{8}-[0-9a-f-]{27} PROC_ID
TRUNCATE TABLE procedure_state_session WITH (wait=false);
+--------------------------------------+
| Procedure ID |
+--------------------------------------+
| PROC_ID |
+--------------------------------------+
-- SQLNESS PROTOCOL MYSQL
TRUNCATE TABLE procedure_state_session WITH (wait=true);
affected_rows: 0
-- SQLNESS PROTOCOL MYSQL
SELECT procedure_state() IN ('{"status":"Running"}', '{"status":"Done"}') AS has_state;
+-----------+
| has_state |
+-----------+
| 1 |
+-----------+
-- A different physical connection still has no history.
-- SQLNESS PROTOCOL POSTGRES
SELECT procedure_state();
Failed to execute query, encountered: Error { kind: Db, cause: Some(DbError { severity: "ERROR", parsed_severity: None, code: SqlState(EP0000), message: "Execution error: Invalid function args: No procedure ID has been returned in this session; pass an explicit procedure ID", detail: None, hint: None, position: None, where_: None, schema: None, table: None, column: None, datatype: None, constraint: None, file: None, line: None, routine: None }) }
-- SQLNESS PROTOCOL POSTGRES
-- SQLNESS REPLACE [0-9a-f]{8}-[0-9a-f-]{27} PROC_ID
TRUNCATE TABLE procedure_state_session WITH (wait=false);
+--------------------------------------+
| Procedure ID |
+--------------------------------------+
| PROC_ID |
+--------------------------------------+
-- SQLNESS PROTOCOL POSTGRES
TRUNCATE TABLE procedure_state_session WITH (wait=true);
Affected Rows: 0
-- SQLNESS PROTOCOL POSTGRES
SELECT procedure_state() IN ('{"status":"Running"}', '{"status":"Done"}') AS has_state;
+-----------+
| has_state |
+-----------+
| t |
+-----------+
-- Failed submissions must preserve history.
-- SQLNESS PROTOCOL MYSQL
TRUNCATE TABLE nonexistent_session_table WITH (wait=false);
Failed to execute query, err: MySqlError { ERROR 1146 (42S02): (TableNotFound): Table not found: greptime.public.nonexistent_session_table }
-- SQLNESS PROTOCOL MYSQL
SELECT procedure_state() IN ('{"status":"Running"}', '{"status":"Done"}') AS has_state;
+-----------+
| has_state |
+-----------+
| 1 |
+-----------+
-- SQLNESS PROTOCOL POSTGRES
TRUNCATE TABLE nonexistent_session_table WITH (wait=false);
Failed to execute query, encountered: Error { kind: Db, cause: Some(DbError { severity: "ERROR", parsed_severity: None, code: SqlState(E42P01), message: "Table not found: greptime.public.nonexistent_session_table", detail: None, hint: None, position: None, where_: None, schema: None, table: None, column: None, datatype: None, constraint: None, file: None, line: None, routine: None }) }
-- SQLNESS PROTOCOL POSTGRES
SELECT procedure_state() IN ('{"status":"Running"}', '{"status":"Done"}') AS has_state;
+-----------+
| has_state |
+-----------+
| t |
+-----------+
DROP TABLE procedure_state_session;
Affected Rows: 0
tests/cases/standalone/common/function/procedure_state_session.sql
+61
View File
@@ -0,0 +1,61 @@
-- Each protocol reuses its own connection across requests.
-- SQLNESS PROTOCOL MYSQL
SELECT procedure_state();
-- SQLNESS PROTOCOL POSTGRES
SELECT procedure_state();
CREATE TABLE procedure_state_session (ts TIMESTAMP TIME INDEX);
-- Synchronous DDL must not record a Procedure ID in the Session.
-- SQLNESS PROTOCOL MYSQL
TRUNCATE TABLE procedure_state_session WITH (wait=true);
-- SQLNESS PROTOCOL MYSQL
SELECT procedure_state();
-- SQLNESS PROTOCOL POSTGRES
TRUNCATE TABLE procedure_state_session WITH (wait=true);
-- SQLNESS PROTOCOL POSTGRES
SELECT procedure_state();
-- Synchronous DDL preserves the ID returned by the previous request.
-- SQLNESS PROTOCOL MYSQL
-- SQLNESS REPLACE [0-9a-f]{8}-[0-9a-f-]{27} PROC_ID
TRUNCATE TABLE procedure_state_session WITH (wait=false);
-- SQLNESS PROTOCOL MYSQL
TRUNCATE TABLE procedure_state_session WITH (wait=true);
-- SQLNESS PROTOCOL MYSQL
SELECT procedure_state() IN ('{"status":"Running"}', '{"status":"Done"}') AS has_state;
-- A different physical connection still has no history.
-- SQLNESS PROTOCOL POSTGRES
SELECT procedure_state();
-- SQLNESS PROTOCOL POSTGRES
-- SQLNESS REPLACE [0-9a-f]{8}-[0-9a-f-]{27} PROC_ID
TRUNCATE TABLE procedure_state_session WITH (wait=false);
-- SQLNESS PROTOCOL POSTGRES
TRUNCATE TABLE procedure_state_session WITH (wait=true);
-- SQLNESS PROTOCOL POSTGRES
SELECT procedure_state() IN ('{"status":"Running"}', '{"status":"Done"}') AS has_state;
-- Failed submissions must preserve history.
-- SQLNESS PROTOCOL MYSQL
TRUNCATE TABLE nonexistent_session_table WITH (wait=false);
-- SQLNESS PROTOCOL MYSQL
SELECT procedure_state() IN ('{"status":"Running"}', '{"status":"Done"}') AS has_state;
-- SQLNESS PROTOCOL POSTGRES
TRUNCATE TABLE nonexistent_session_table WITH (wait=false);
-- SQLNESS PROTOCOL POSTGRES
SELECT procedure_state() IN ('{"status":"Running"}', '{"status":"Done"}') AS has_state;
DROP TABLE procedure_state_session;