From 1851f6bf4de7b6ae94fc9b0cb3048aa24877aea9 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 25 Aug 2026 12:34:53 +0000 Subject: [PATCH 01/22] fix(query): keep INSERT timestamp conversion out of the source query (#8911) * fix(query): keep INSERT timestamp conversion out of the source query Interpreting an INSERT's string timestamps used to work by pushing the conversion down into the source query, which changed what that query means. Two consequences: - Pushing through a UNION's DISTINCT moved the dedup key from the raw strings to parsed instants, so rows spelling the same instant differently collapsed into one. On an append-only table that is a silently dropped row. - A UNION branch that needed no conversion (a NULL, or an explicit cast) made the whole column give up, leaving sibling branches on UTC while the rest of the row used the session timezone. Convert at the assignment instead, by routing its cast through a timezone-carrying timestamp type and back. Arrow applies the timezone when a cast target carries one, and stripping it afterwards preserves the value. The source query is no longer touched, so both cases go away and the tree-walking rewrite (roughly 160 lines) is deleted. The rewrite reads source types, so it now runs TypeCoercion first: a UNION still carries its loose per-branch schema before coercion, and retargeting a cast whose input later becomes a timestamp would shift the value rather than reinterpret it. Signed-off-by: Dennis Zhuang * fix(query): address review on INSERT assignment rewrite - Clone the input `Arc` instead of the whole subtree, and only rebuild it when a `Values` row actually changes. - Defer cloning the cast source until the literal-folding path has been ruled out. - Move the UTC check onto `Timezone::is_utc`, replacing a bare string compare. - Cover a prepared `INSERT ... VALUES (?)`: an untyped placeholder types as `Null`, so the assignment cast is left for parameter substitution. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/common/time/src/timezone.rs | 9 + src/query/src/optimizer.rs | 1 + src/query/src/optimizer/insert_assignment.rs | 307 ++++++++++++++++ src/query/src/optimizer/type_conversion.rs | 11 +- .../type_conversion/insert_assignment.rs | 336 ------------------ src/query/src/planner.rs | 126 +++---- .../insert/insert_default_timezone.result | 58 +++ .../common/insert/insert_default_timezone.sql | 28 ++ 8 files changed, 472 insertions(+), 404 deletions(-) create mode 100644 src/query/src/optimizer/insert_assignment.rs delete mode 100644 src/query/src/optimizer/type_conversion/insert_assignment.rs diff --git a/src/common/time/src/timezone.rs b/src/common/time/src/timezone.rs index 41cc1f78428..dbd8dffcd49 100644 --- a/src/common/time/src/timezone.rs +++ b/src/common/time/src/timezone.rs @@ -108,6 +108,15 @@ impl Timezone { } } + /// A named zone merely sitting at +00:00 today is not UTC: it may have been + /// elsewhere at the timestamp being converted. + pub fn is_utc(&self) -> bool { + match self { + Self::Offset(offset) => offset.local_minus_utc() == 0, + Self::Named(tz) => matches!(tz, Tz::UTC), + } + } + /// Returns the number of seconds to add to convert from UTC to the local time. pub fn local_minus_utc(&self) -> i64 { match self { diff --git a/src/query/src/optimizer.rs b/src/query/src/optimizer.rs index 8827d48ed63..a6b36e6fdae 100644 --- a/src/query/src/optimizer.rs +++ b/src/query/src/optimizer.rs @@ -17,6 +17,7 @@ pub mod constant_term; pub mod count_nest_aggr; pub mod count_wildcard; pub mod global_limit; +pub(crate) mod insert_assignment; pub(crate) mod json_type_concretize; pub mod parallelize_scan; pub mod pass_distribution; diff --git a/src/query/src/optimizer/insert_assignment.rs b/src/query/src/optimizer/insert_assignment.rs new file mode 100644 index 00000000000..ad77736ce0a --- /dev/null +++ b/src/query/src/optimizer/insert_assignment.rs @@ -0,0 +1,307 @@ +// 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_time::Timezone; +use datafusion::config::ConfigOptions; +use datafusion_common::{DFSchemaRef, Result, ScalarValue}; +use datafusion_expr::expr::{Alias, Cast}; +use datafusion_expr::{Distinct, Expr, ExprSchemable, LogicalPlan, Projection, Values}; +use datafusion_optimizer::analyzer::AnalyzerRule; +use datafusion_optimizer::analyzer::type_coercion::TypeCoercion; +use datatypes::arrow::datatypes::{DataType, TimeUnit}; +use session::context::QueryContextRef; + +use crate::optimizer::type_conversion::cast_string_to_timestamp; + +/// Interprets strings assigned to timestamp columns at an `INSERT` boundary +/// using the session timezone. `plan` is the assignment projection under a +/// `WriteOp::Insert`. +/// +/// DataFusion plans `INSERT` as a projection casting each source column to its +/// target column type, and that cast reads a naive string as UTC. Arrow does +/// apply a timezone when the cast target carries one, so the assignment is +/// routed through `Timestamp(unit, Some(tz))` and back. Stripping the timezone +/// afterwards is value-preserving — arrow only shifts values in the opposite +/// direction. +/// +/// The source query is left untouched. Reinterpreting a value where it is +/// *produced* would change what the source query means: pushing the conversion +/// below a `UNION`'s `DISTINCT`, for instance, moves the dedup key from the raw +/// strings to parsed instants and silently drops rows. +/// +/// # Why `TypeCoercion` runs here +/// +/// The rewrite reads source types, and those are only settled once a `UNION`'s +/// branch types have been reconciled: before coercion a union carries its loose +/// schema (the first branch's types), so a mixed +/// `SELECT 'string' UNION ALL SELECT CAST(.. AS TIMESTAMP)` still looks like a +/// string. Retargeting that cast would leave `Timestamp(None) -> +/// Timestamp(Some(tz))` behind once coercion retypes the union — the one +/// direction in which arrow shifts the value instead of relabelling it. +/// +/// Coercing here rather than deferring to the analyzer is forced by where an +/// INSERT is still identifiable: `exec_dml_statement` strips the `Dml` node and +/// executes its input, so by the time the analyzer runs, an assignment +/// projection is indistinguishable from any other projection. +/// +/// Explicit casts stay out of this: the SQL layer turns a user's +/// `CAST(x AS TIMESTAMP)` into an `arrow_cast` call, which only becomes an +/// `Expr::Cast` in the optimizer's `SimplifyExpressions`. Assignment casts are +/// therefore the only `Expr::Cast` reaching a timestamp column here. +/// +/// # Reach +/// +/// The emitted cast only carries its timezone where the expression is evaluated +/// on this node. Substrait drops the timezone name when a plan is pushed down — +/// it encodes any zoned timestamp as `PrecisionTimestampTz` and decodes it back +/// as UTC — so a source reading from a table falls back to UTC, the behaviour it +/// had before this rule existed. Sources that never leave this node (literals, +/// `VALUES`, and `UNION`s of them) keep the session timezone, and those are what +/// an INSERT's timestamp assignment is in practice. +pub(crate) fn rewrite_insert_assignments( + plan: LogicalPlan, + query_ctx: &QueryContextRef, + config: &ConfigOptions, +) -> Result { + let Some(timezone) = session_timezone(query_ctx) else { + return Ok(plan); + }; + + let plan = TypeCoercion::new().analyze(plan, config)?; + rewrite_assignment(plan, &timezone) +} + +/// Session timezone, in both forms the rewrite needs. +struct SessionTimezone { + /// Parses literals, matching the plain `INSERT ... VALUES` path. + parsed: Timezone, + /// Names the intermediate arrow cast target. + name: Arc, +} + +fn session_timezone(query_ctx: &QueryContextRef) -> Option { + let parsed = query_ctx.timezone(); + + // A UTC session already gets UTC semantics from the plain assignment cast. + if parsed.is_utc() { + return None; + } + + Some(SessionTimezone { + name: Arc::from(parsed.to_string()), + parsed, + }) +} + +fn rewrite_assignment(plan: LogicalPlan, timezone: &SessionTimezone) -> Result { + let LogicalPlan::Projection(assignment) = plan else { + return Ok(plan); + }; + + let mut exprs = assignment.expr.clone(); + let mut changed = false; + for expr in &mut exprs { + changed |= retarget_assignment_cast( + expr, + assignment.input.schema(), + Some(assignment.input.as_ref()), + timezone, + )?; + } + + // The planner types `VALUES` against the target table, so the assignment + // cast lands inside the `Values` rows instead of on the projection above. + let mut input = assignment.input.clone(); + if let LogicalPlan::Values(values) = assignment.input.as_ref() + && let Some(rewritten) = rewrite_values(values, timezone)? + { + input = Arc::new(LogicalPlan::Values(rewritten)); + changed = true; + } + + if !changed { + return Ok(LogicalPlan::Projection(assignment)); + } + Projection::try_new(exprs, input).map(LogicalPlan::Projection) +} + +fn rewrite_values(values: &Values, timezone: &SessionTimezone) -> Result> { + let mut rewritten = values.clone(); + let mut changed = false; + for row in &mut rewritten.values { + for expr in row.iter_mut() { + changed |= retarget_assignment_cast(expr, &values.schema, None, timezone)?; + } + } + + Ok(changed.then_some(rewritten)) +} + +/// Reinterprets one assignment cast, returning whether it was rewritten. +/// +/// `source_plan` is the projection's input, used to resolve a literal behind a +/// column reference; `Values` rows carry their expression inline and pass `None`. +fn retarget_assignment_cast( + expr: &mut Expr, + schema: &DFSchemaRef, + source_plan: Option<&LogicalPlan>, + timezone: &SessionTimezone, +) -> Result { + let expr = unalias_mut(expr); + let Expr::Cast(Cast { + expr: source, + data_type: DataType::Timestamp(unit, None), + }) = expr + else { + return Ok(false); + }; + let unit = *unit; + + if !matches!( + source.get_type(schema)?, + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + ) { + return Ok(false); + } + + // Fold literals with the same parser the plain `INSERT ... VALUES` path + // uses, so a given string means the same thing however it reaches a column. + // The parsers disagree on ambiguous local times: this one resolves them, + // arrow rejects them. + let folded = source_literal(source.as_ref(), source_plan) + .and_then(|literal| convert_literal(&literal, unit, &timezone.parsed)); + if let Some(folded) = folded { + *expr = folded; + return Ok(true); + } + + let source = source.as_ref().clone(); + *expr = Expr::Cast(Cast::new( + Box::new(Expr::Cast(Cast::new( + Box::new(source), + DataType::Timestamp(unit, Some(timezone.name.clone())), + ))), + DataType::Timestamp(unit, None), + )); + Ok(true) +} + +fn source_literal(source: &Expr, source_plan: Option<&LogicalPlan>) -> Option { + match source { + Expr::Literal(value, _) => Some(value.clone()), + Expr::Column(column) => { + let plan = source_plan?; + let index = plan.schema().maybe_index_of_column(column)?; + lineage_literal(plan, index).cloned() + } + _ => None, + } +} + +/// Resolves a literal when every row carries the same value at `output_idx`. +/// +/// Read-only: the literal is folded into the assignment above, so nodes that +/// drop, reorder or deduplicate rows can be traversed — none of them changes +/// the value a surviving row carries, and folding above them leaves their keys +/// on the original strings. +fn lineage_literal(plan: &LogicalPlan, output_idx: usize) -> Option<&ScalarValue> { + if output_idx >= plan.schema().fields().len() { + return None; + } + + match plan { + LogicalPlan::Projection(projection) => match unalias(&projection.expr[output_idx]) { + Expr::Literal(value, _) => Some(value), + Expr::Column(column) => { + let input_idx = projection.input.schema().maybe_index_of_column(column)?; + lineage_literal(projection.input.as_ref(), input_idx) + } + _ => None, + }, + LogicalPlan::Filter(_) + | LogicalPlan::Sort(_) + | LogicalPlan::Limit(_) + | LogicalPlan::SubqueryAlias(_) + | LogicalPlan::Distinct(Distinct::All(_)) => { + let inputs = plan.inputs(); + let [input] = inputs.as_slice() else { + return None; + }; + lineage_literal(input, output_idx) + } + _ => None, + } +} + +fn convert_literal(value: &ScalarValue, unit: TimeUnit, timezone: &Timezone) -> Option { + let ScalarValue::Utf8(Some(value)) = value else { + return None; + }; + cast_string_to_timestamp(value, &DataType::Timestamp(unit, None), Some(timezone)) + .ok() + .filter(|value| !value.is_null()) + .map(|value| Expr::Literal(value, None)) +} + +fn unalias(expr: &Expr) -> &Expr { + match expr { + Expr::Alias(Alias { expr, .. }) => unalias(expr), + expr => expr, + } +} + +fn unalias_mut(expr: &mut Expr) -> &mut Expr { + match expr { + Expr::Alias(Alias { expr, .. }) => unalias_mut(expr), + expr => expr, + } +} + +#[cfg(test)] +mod tests { + use datafusion_common::DFSchema; + use datafusion_expr::expr::Placeholder; + + use super::*; + + fn shanghai() -> SessionTimezone { + let parsed = Timezone::from_tz_string("Asia/Shanghai").unwrap(); + SessionTimezone { + name: Arc::from(parsed.to_string()), + parsed, + } + } + + /// A prepared `INSERT ... VALUES (?)` arrives here as a cast over an untyped + /// placeholder, which must survive for parameter substitution. + #[test] + fn test_untyped_placeholder_assignment_is_left_alone() { + let schema = Arc::new(DFSchema::empty()); + let mut expr = Expr::Cast(Cast::new( + Box::new(Expr::Placeholder(Placeholder::new_with_field( + "$1".to_string(), + None, + ))), + DataType::Timestamp(TimeUnit::Millisecond, None), + )); + let original = expr.clone(); + + let changed = retarget_assignment_cast(&mut expr, &schema, None, &shanghai()).unwrap(); + + assert!(!changed); + assert_eq!(expr, original); + } +} diff --git a/src/query/src/optimizer/type_conversion.rs b/src/query/src/optimizer/type_conversion.rs index 3941fedfe8a..40bb43e791e 100644 --- a/src/query/src/optimizer/type_conversion.rs +++ b/src/query/src/optimizer/type_conversion.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -mod insert_assignment; - use std::sync::Arc; use common_time::Timezone; @@ -31,7 +29,7 @@ use session::context::QueryContextRef; use crate::QueryEngineContext; use crate::optimizer::ExtensionAnalyzerRule; -use crate::optimizer::type_conversion::insert_assignment::rewrite_insert_assignments; +use crate::optimizer::insert_assignment::rewrite_insert_assignments; use crate::plan::ExtractExpr; /// TypeConversionRule converts some literal values in logical plan to other types according @@ -46,7 +44,7 @@ impl ExtensionAnalyzerRule for TypeConversionRule { &self, plan: LogicalPlan, ctx: &QueryEngineContext, - _config: &ConfigOptions, + config: &ConfigOptions, ) -> Result { plan.transform_up_with_subqueries(|plan| match plan { LogicalPlan::Filter(filter) => { @@ -129,7 +127,8 @@ impl ExtensionAnalyzerRule for TypeConversionRule { LogicalPlan::Dml(mut dml) if matches!(dml.op, WriteOp::Insert(_)) => { dml.input = Arc::new(rewrite_insert_assignments( dml.input.as_ref().clone(), - ctx.query_ctx(), + &ctx.query_ctx(), + config, )?); Ok(Transformed::yes(LogicalPlan::Dml(dml))) } @@ -324,7 +323,7 @@ fn timestamp_to_timestamp_ms_expr(val: i64, unit: TimeUnit) -> Expr { ) } -fn cast_string_to_timestamp( +pub(crate) fn cast_string_to_timestamp( string: &str, target_type: &DataType, timezone: Option<&Timezone>, diff --git a/src/query/src/optimizer/type_conversion/insert_assignment.rs b/src/query/src/optimizer/type_conversion/insert_assignment.rs deleted file mode 100644 index 11ab7e7829e..00000000000 --- a/src/query/src/optimizer/type_conversion/insert_assignment.rs +++ /dev/null @@ -1,336 +0,0 @@ -// 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 datafusion_common::{Column, Result, ScalarValue}; -use datafusion_expr::expr::{Alias, Cast}; -use datafusion_expr::{Distinct, Expr, ExprSchemable, LogicalPlan, Projection, Union, Values}; -use datatypes::arrow::datatypes::DataType; -use session::context::QueryContextRef; - -use crate::optimizer::type_conversion::cast_string_to_timestamp; -use crate::plan::ExtractExpr; - -/// Rewrites string literals that feed timestamp columns at an INSERT boundary. -/// Constants are folded at the assignment to avoid changing source types; -/// `VALUES` and `UNION` inputs are rewritten per row or branch. Explicit casts -/// stay on DataFusion's existing path. -pub(super) fn rewrite_insert_assignments( - plan: LogicalPlan, - query_ctx: QueryContextRef, -) -> Result { - let LogicalPlan::Projection(assignment) = plan else { - return Ok(plan); - }; - - let converter = InsertAssignmentConverter { query_ctx }; - let mut exprs = assignment.expr.clone(); - let mut input = assignment.input.as_ref().clone(); - let mut changed = false; - for (output_idx, expr) in assignment.expr.iter().enumerate() { - let target_type = assignment.schema.field(output_idx).data_type(); - if !matches!(target_type, DataType::Timestamp(_, _)) { - continue; - } - - let Some(column) = assignment_input_column(expr) else { - continue; - }; - let Some(input_idx) = input.schema().maybe_index_of_column(column) else { - continue; - }; - - let literal = lineage_literal(&input, input_idx).cloned(); - if let Some(literal) = literal - && let Some(folded) = converter.convert_literal(&literal, target_type) - { - let (qualifier, field) = assignment.schema.qualified_field(output_idx); - exprs[output_idx] = folded.alias_qualified(qualifier.cloned(), field.name()); - changed = true; - continue; - } - - // A hand-built DML plan (e.g. from flow) may share one source column - // between targets; rewriting it in place would retype every reader. - if assignment - .expr - .iter() - .enumerate() - .any(|(idx, other)| idx != output_idx && other.column_refs().contains(column)) - { - continue; - } - if let Some(rewritten) = converter.rewrite_output_column(&input, input_idx, target_type)? { - input = rewritten; - changed = true; - } - } - - if !changed { - return Ok(LogicalPlan::Projection(assignment)); - } - Projection::try_new(exprs, Arc::new(input)).map(LogicalPlan::Projection) -} - -/// Resolves a literal when every row carries the same value at `output_idx`. -fn lineage_literal(plan: &LogicalPlan, output_idx: usize) -> Option<&ScalarValue> { - if output_idx >= plan.schema().fields().len() { - return None; - } - - match plan { - LogicalPlan::Projection(projection) => match unalias(&projection.expr[output_idx]) { - Expr::Literal(value, _) => Some(value), - Expr::Column(column) => { - let input_idx = projection.input.schema().maybe_index_of_column(column)?; - lineage_literal(projection.input.as_ref(), input_idx) - } - _ => None, - }, - // These nodes drop, reorder or deduplicate rows without touching the - // value a surviving row carries, and their schema stays positional. - LogicalPlan::Filter(_) - | LogicalPlan::Sort(_) - | LogicalPlan::Limit(_) - | LogicalPlan::SubqueryAlias(_) - | LogicalPlan::Distinct(Distinct::All(_)) => { - let inputs = plan.inputs(); - let [input] = inputs.as_slice() else { - return None; - }; - lineage_literal(input, output_idx) - } - _ => None, - } -} - -struct InsertAssignmentConverter { - query_ctx: QueryContextRef, -} - -impl InsertAssignmentConverter { - fn rewrite_output_column( - &self, - plan: &LogicalPlan, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - if output_idx >= plan.schema().fields().len() { - return Ok(None); - } - - match plan { - // DataFusion pushes INSERT assignment casts into Values, so inspect - // them even though the Values schema already has the target type. - LogicalPlan::Values(values) => self.rewrite_values(values, output_idx, target_type), - // Once the source query has produced the target type, its casts and - // coercions belong to the source query rather than INSERT assignment. - _ if plan.schema().field(output_idx).data_type() == target_type => Ok(None), - LogicalPlan::Projection(projection) => { - self.rewrite_projection(projection, output_idx, target_type) - } - LogicalPlan::Union(union) => self.rewrite_union(union, output_idx, target_type), - // Filter and Sort are excluded here: their predicates and keys read - // the retyped column, which would change the source query. Constants - // still reach the assignment through `lineage_literal`. Distinct::All - // holds no expressions; allowing it shifts dedup keys from raw - // strings to parsed instants. - LogicalPlan::Limit(_) - | LogicalPlan::SubqueryAlias(_) - | LogicalPlan::Distinct(Distinct::All(_)) => { - self.rewrite_passthrough(plan, output_idx, target_type) - } - _ => Ok(None), - } - } - - fn rewrite_projection( - &self, - projection: &Projection, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - match unalias(&projection.expr[output_idx]) { - Expr::Literal(value, _) => { - let Some(converted) = self.convert_literal(value, target_type) else { - return Ok(None); - }; - let (qualifier, field) = projection.schema.qualified_field(output_idx); - let mut exprs = projection.expr.clone(); - exprs[output_idx] = converted.alias_qualified(qualifier.cloned(), field.name()); - - Projection::try_new(exprs, projection.input.clone()) - .map(LogicalPlan::Projection) - .map(Some) - } - Expr::Column(column) => { - // Retyping the input column affects every output column reading - // it, so only follow lineage with a single consumer. - if projection - .expr - .iter() - .enumerate() - .any(|(idx, other)| idx != output_idx && other.column_refs().contains(column)) - { - return Ok(None); - } - - let input = projection.input.as_ref(); - let Some(input_idx) = input.schema().maybe_index_of_column(column) else { - return Ok(None); - }; - let Some(rewritten) = self.rewrite_output_column(input, input_idx, target_type)? - else { - return Ok(None); - }; - - Projection::try_new(projection.expr.clone(), Arc::new(rewritten)) - .map(LogicalPlan::Projection) - .map(Some) - } - _ => Ok(None), - } - } - - fn rewrite_values( - &self, - values: &Values, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - let mut rewritten = values.clone(); - let mut changed = false; - for row in &mut rewritten.values { - let Some(expr) = row.get_mut(output_idx) else { - return Ok(None); - }; - - if let Expr::Cast(Cast { - expr: inner, - data_type, - }) = expr - && data_type == target_type - && let Expr::Literal(value, _) = inner.as_ref() - && let Some(value) = self.convert_literal(value, target_type) - { - *expr = value; - changed = true; - continue; - } - - if expr.get_type(values.schema.as_ref()).ok().as_ref() != Some(target_type) { - return Ok(None); - } - } - - Ok(changed.then_some(LogicalPlan::Values(rewritten))) - } - - fn rewrite_union( - &self, - union: &Union, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - let mut inputs = Vec::with_capacity(union.inputs.len()); - // A union has one shared schema, so rewrite it only when the same output - // column is an INSERT-assigned literal in every branch. - for input in &union.inputs { - let Some(rewritten) = self.rewrite_output_column(input, output_idx, target_type)? - else { - return Ok(None); - }; - inputs.push(Arc::new(rewritten)); - } - - // Untouched columns may still have mismatched branch types before - // TypeCoercion runs, so rebuild loosely like the SQL planner does. - Union::try_new_with_loose_types(inputs) - .map(LogicalPlan::Union) - .map(Some) - } - - fn rewrite_passthrough( - &self, - plan: &LogicalPlan, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - let inputs = plan.inputs(); - let [input] = inputs.as_slice() else { - return Ok(None); - }; - let Some(rewritten_input) = self.rewrite_output_column(input, output_idx, target_type)? - else { - return Ok(None); - }; - - plan.with_new_exprs(plan.expressions_consider_join(), vec![rewritten_input]) - .map(Some) - } - - fn convert_literal(&self, value: &ScalarValue, target_type: &DataType) -> Option { - let ScalarValue::Utf8(Some(value)) = value else { - return None; - }; - cast_string_to_timestamp(value, target_type, Some(&self.query_ctx.timezone())) - .ok() - .filter(|value| !value.is_null()) - .map(|value| Expr::Literal(value, None)) - } -} - -fn assignment_input_column(expr: &Expr) -> Option<&Column> { - let expr = match unalias(expr) { - Expr::Cast(Cast { expr, .. }) => expr.as_ref(), - expr => expr, - }; - let Expr::Column(column) = expr else { - return None; - }; - Some(column) -} - -fn unalias(expr: &Expr) -> &Expr { - match expr { - Expr::Alias(Alias { expr, .. }) => unalias(expr), - expr => expr, - } -} - -#[cfg(test)] -mod tests { - use datafusion_common::ScalarValue; - use datafusion_common::arrow::datatypes::TimeUnit; - use session::context::QueryContext; - - use super::*; - - #[test] - fn test_convert_literal_falls_back_for_unsupported_literal() { - let converter = InsertAssignmentConverter { - query_ctx: QueryContext::arc(), - }; - let target_type = DataType::Timestamp(TimeUnit::Nanosecond, None); - - for literal in ["1970-01-01", "-8-01-01 00:00:01.5"] { - assert_eq!( - converter - .convert_literal(&ScalarValue::Utf8(Some(literal.to_string())), &target_type,), - None - ); - } - } -} diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index a7d4a84cb28..c3871a96afa 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -866,6 +866,30 @@ mod tests { engine.planner().plan(&stmt, query_ctx).await.unwrap() } + /// Plans `sql` and runs the DataFusion analyzer, which is where + /// `InsertAssignmentRule` sits. Planning alone stops short of it, so these + /// assertions would not see the assignment rewrite at all. + async fn analyze_insert( + engine: &QueryEngineRef, + sql: &str, + query_ctx: &QueryContextRef, + ) -> String { + let stmt = QueryLanguageParser::parse_sql(sql, query_ctx).unwrap(); + let plan = engine + .planner() + .plan(&stmt, query_ctx.clone()) + .await + .unwrap(); + let context = engine.engine_context(query_ctx.clone()); + let state = context.state(); + state + .analyzer() + .execute_and_check(plan, state.config_options(), |_, _| {}) + .unwrap() + .display_indent_schema() + .to_string() + } + #[tokio::test] async fn test_insert_timestamp_literals_use_query_timezone() { let query_ctx = Arc::new( @@ -898,20 +922,6 @@ mod tests { ) AS source", &[1_785_902_400_001_i64][..], ), - ( - "INSERT INTO timestamps (ts, st) \ - SELECT '2026-08-06 12:00:00.001', now() \ - UNION ALL \ - SELECT '2026-08-07 12:00:00.001', now()", - &[1_785_988_800_001_i64, 1_786_075_200_001_i64][..], - ), - ( - "INSERT INTO timestamps (ts, st) \ - SELECT '2026-08-16 12:00:00.001', now() \ - UNION \ - SELECT '2026-08-17 12:00:00.001', now()", - &[1_786_852_800_001_i64, 1_786_939_200_001_i64][..], - ), ( "INSERT INTO timestamps (ts, st) \ SELECT '2026-08-18 12:00:00.001', max(st) \ @@ -940,14 +950,7 @@ mod tests { &[1_786_766_400_001_i64][..], ), ] { - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx.clone()) - .await - .unwrap() - .display_indent() - .to_string(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; for expected_timestamp in expected_timestamps { assert!( @@ -968,19 +971,19 @@ mod tests { let engine = create_timestamp_test_engine().await; let sql = "INSERT INTO timestamps (ts, st) \ VALUES (CAST('2026-08-08 12:00:00.001' AS TIMESTAMP), now())"; - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx) - .await - .unwrap() - .display_indent() - .to_string(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; + // An explicit cast reaches the analyzer as an `arrow_cast` call rather + // than an `Expr::Cast`, which is how it stays out of the rewrite. assert!( plan.contains("arrow_cast(Utf8(\"2026-08-08 12:00:00.001\")"), "{plan}" ); + // 12:00:00.001 read as Shanghai local time; the source query keeps UTC. + assert!( + !plan.contains("TimestampMillisecond(1786104000001, None)"), + "{plan}" + ); } #[tokio::test] @@ -1017,14 +1020,7 @@ mod tests { ][..], ), ] { - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx.clone()) - .await - .unwrap() - .display_indent_schema() - .to_string(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; for expected in expected { assert!(plan.contains(expected), "{plan}"); @@ -1033,30 +1029,37 @@ mod tests { } #[tokio::test] - async fn test_insert_union_tolerates_uncoerced_untouched_column() { + async fn test_insert_union_converts_via_assignment_cast() { let query_ctx = Arc::new( QueryContextBuilder::default() .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()) .build(), ); let engine = create_timestamp_test_engine().await; - // `st` stays Timestamp vs Null across branches until TypeCoercion runs. - let sql = "INSERT INTO timestamps (ts, st) \ - SELECT '2026-08-06 12:00:00.001', now() \ - UNION ALL \ - SELECT '2026-08-07 12:00:00.001', NULL"; - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx) - .await - .unwrap() - .display_indent() - .to_string(); + // Branches disagree, so the conversion stays as a cast on the + // assignment instead of folding. One cast covers every branch, which is + // why a NULL branch no longer cancels the conversion for the column and + // why UNION's dedup keys stay on the original strings. + for sql in [ + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-06 12:00:00.001', now() \ + UNION ALL \ + SELECT '2026-08-07 12:00:00.001', NULL", + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-16 12:00:00.001', now() \ + UNION \ + SELECT '2026-08-17 12:00:00.001', now()", + ] { + let plan = analyze_insert(&engine, sql, &query_ctx).await; - for expected_timestamp in [1_785_988_800_001_i64, 1_786_075_200_001_i64] { assert!( - plan.contains(&format!("TimestampMillisecond({expected_timestamp}, None)")), + plan.contains("AS Timestamp(ms, \"Asia/Shanghai\")"), + "{plan}" + ); + // The branches themselves are untouched. + assert!( + plan.contains("Utf8(\"2026-08-07 12:00:00.001\")") + || plan.contains("Utf8(\"2026-08-17 12:00:00.001\")"), "{plan}" ); } @@ -1074,14 +1077,8 @@ mod tests { SELECT '2026-08-10 12:00:00.001', now() \ UNION ALL \ SELECT CAST('2026-08-11 12:00:00.001' AS TIMESTAMP), now()"; - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx) - .await - .unwrap() - .display_indent_schema() - .to_string(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; + assert!( !plan.contains("TimestampMillisecond(1786334400001, None)"), "{plan}" @@ -1090,6 +1087,11 @@ mod tests { plan.contains("arrow_cast(Utf8(\"2026-08-11 12:00:00.001\")"), "{plan}" ); + // TypeCoercion has already settled this union to timestamp, so the + // assignment has nothing left to reinterpret. Retargeting the cast here + // would leave a Timestamp(None) -> Timestamp(Some(tz)) step behind, + // which shifts the value instead of relabelling it. + assert!(!plan.contains("Asia/Shanghai"), "{plan}"); } #[tokio::test] diff --git a/tests/cases/standalone/common/insert/insert_default_timezone.result b/tests/cases/standalone/common/insert/insert_default_timezone.result index 413eb5d46ae..41d8ca6cba8 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.result +++ b/tests/cases/standalone/common/insert/insert_default_timezone.result @@ -114,6 +114,27 @@ INSERT INTO test3 (ts, st, ts_ns) VALUES ( Affected Rows: 1 +-- a NULL branch must not cancel the conversion for the whole column +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-13 12:00:00.001' AS a, '2026-08-13 12:00:00.123456789' AS b +UNION ALL +SELECT '2026-08-14 12:00:00.001', NULL; + +Affected Rows: 2 + +-- NULL in the first branch: the union's schema starts out as Null +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-15 12:00:00.001' AS a, NULL AS b +UNION ALL +SELECT '2026-08-19 12:00:00.001', '2026-08-19 12:00:00.123456789'; + +Affected Rows: 2 + +-- the assignment cast also lands on non-literal VALUES expressions +INSERT INTO test3 (ts, st) VALUES (concat('2026-08-20 ', '12:00:00.001'), now()); + +Affected Rows: 1 + SELECT ts, ts_ns FROM test3 ORDER BY ts; +-------------------------+-------------------------------+ @@ -130,10 +151,43 @@ SELECT ts, ts_ns FROM test3 ORDER BY ts; | 2026-08-10T12:00:00.001 | | | 2026-08-11T12:00:00.001 | | | 2026-08-12T04:00:00.123 | 2026-08-12T04:00:00.123456789 | +| 2026-08-13T04:00:00.001 | 2026-08-13T04:00:00.123456789 | +| 2026-08-14T04:00:00.001 | | +| 2026-08-15T04:00:00.001 | | | 2026-08-16T04:00:00.001 | | | 2026-08-17T04:00:00.001 | | +| 2026-08-19T04:00:00.001 | 2026-08-19T04:00:00.123456789 | +| 2026-08-20T04:00:00.001 | | +-------------------------+-------------------------------+ +-- UNION dedup keys must stay on the source strings: these two spell the same +-- instant differently, so the source query yields two rows and both are kept. +CREATE TABLE test4 (ts TIMESTAMP TIME INDEX) WITH ('append_mode'='true'); + +Affected Rows: 0 + +INSERT INTO test4 (ts) +SELECT '2026-08-06 04:00:00' UNION SELECT '2026-08-06 04:00:00.000'; + +Affected Rows: 2 + +SELECT count(*) FROM test4; + ++----------+ +| count(*) | ++----------+ +| 2 | ++----------+ + +SELECT ts FROM test4 ORDER BY ts; + ++---------------------+ +| ts | ++---------------------+ +| 2026-08-05T20:00:00 | +| 2026-08-05T20:00:00 | ++---------------------+ + SET time_zone = 'UTC'; Affected Rows: 0 @@ -150,3 +204,7 @@ DROP TABLE test3; Affected Rows: 0 +DROP TABLE test4; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/insert/insert_default_timezone.sql b/tests/cases/standalone/common/insert/insert_default_timezone.sql index 7cbef15d6d9..f7c196699b7 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.sql +++ b/tests/cases/standalone/common/insert/insert_default_timezone.sql @@ -56,8 +56,34 @@ INSERT INTO test3 (ts, st, ts_ns) VALUES ( '2026-08-09 12:00:00.123456789' ); +-- a NULL branch must not cancel the conversion for the whole column +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-13 12:00:00.001' AS a, '2026-08-13 12:00:00.123456789' AS b +UNION ALL +SELECT '2026-08-14 12:00:00.001', NULL; + +-- NULL in the first branch: the union's schema starts out as Null +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-15 12:00:00.001' AS a, NULL AS b +UNION ALL +SELECT '2026-08-19 12:00:00.001', '2026-08-19 12:00:00.123456789'; + +-- the assignment cast also lands on non-literal VALUES expressions +INSERT INTO test3 (ts, st) VALUES (concat('2026-08-20 ', '12:00:00.001'), now()); + SELECT ts, ts_ns FROM test3 ORDER BY ts; +-- UNION dedup keys must stay on the source strings: these two spell the same +-- instant differently, so the source query yields two rows and both are kept. +CREATE TABLE test4 (ts TIMESTAMP TIME INDEX) WITH ('append_mode'='true'); + +INSERT INTO test4 (ts) +SELECT '2026-08-06 04:00:00' UNION SELECT '2026-08-06 04:00:00.000'; + +SELECT count(*) FROM test4; + +SELECT ts FROM test4 ORDER BY ts; + SET time_zone = 'UTC'; DROP TABLE test1; @@ -65,3 +91,5 @@ DROP TABLE test1; DROP TABLE test2; DROP TABLE test3; + +DROP TABLE test4; From 6d86e6ff068fe8471c3d7a28e70f9144e7684e15 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 25 Aug 2026 13:08:44 +0000 Subject: [PATCH 02/22] feat: synthesize OTLP resource descriptor for the semantic entity graph (#8904) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(servers): compose OTLP metrics job from service.namespace/service.name The OTel Prometheus compatibility spec defines job as "/" when the namespace is present. The OTLP metrics path only used the bare service.name, so the job tag diverged from target_info produced by Prometheus-side exporters for the same resource. Compose the namespace form, and keep not fabricating a job when service.name is absent. Behavior change: resources carrying service.namespace now get "namespace/name" as their job tag value. Signed-off-by: Dennis Zhuang * feat(otlp): synthesize otel_resource_info at OTLP metrics ingestion Ordinary OTLP metrics scatter filtered resource attributes as tags over every logical metric table, so metrics-only services contribute nothing to the semantic entity graph. Each request now also projects its distinct resources into one info-metric-shaped mito table, otel_resource_info: a fixed allowlist of identity-relevant attributes under their raw OTel keys (independent of the label translation strategy and the promote/ignore headers) plus derived job/instance compatibility columns, value 1.0, and the newest data-point timestamp. The descriptor is written after the main insert is committed; a failure there (conflicting pre-existing table, auto-create disabled) degrades to an OTLP partial_success warning with rejected_data_points = 0 instead of failing the request and triggering client retries of already-accepted data. A request writing a metric named otel_resource_info suppresses synthesis. Legacy mode is unchanged. Signed-off-by: Dennis Zhuang * feat(operator): otel info-metric conventions with host/container entities Whitelist the ingestion-synthesized otel_resource_info descriptor via a new otel_info_metrics conventions map, gated on source=opentelemetry (the existing gate hardcoded source=prometheus). Its declarations use explicit descriptive lists instead of descriptive_rest so identifying attributes of other entities do not leak into service.instance. Conventions tightened per the Astronomy Shop findings: host identity is host.id with host.name descriptive only (host.name is not stable across SDKs and resource detectors), a generic container entity (new entity type) is declared only when container.id is present, and trace-v1 tables now synthesize host/container from their flattened resource attributes too. New co-declared edges: service.instance runs_on container, container runs_on host. Signed-off-by: Dennis Zhuang * test(otlp): cover the resource descriptor in integration tests Covers the descriptor's raw-key columns and info-metric options through the HTTP path, the namespace/name job composition end-to-end, column names being independent of the translation strategy, the allowlist excluding unlisted resource attributes, auto-create after a drop, the metric-name collision suppressing synthesis, and the partial-success warning (rejected_data_points = 0) when a pre-existing incompatible table fails the descriptor write while metric data is accepted. Signed-off-by: Dennis Zhuang * chore: cargo fmt Signed-off-by: Dennis Zhuang * fix(frontend): degrade descriptor permission denial to a warning A table-level permission policy denying otel_resource_info would have failed the whole OTLP metrics request because the descriptor's permission check ran before the main insert. The descriptor is derived enrichment: check its permission in the degrade path so a denial skips the write and surfaces as the partial-success warning, like any other descriptor write failure. Signed-off-by: Dennis Zhuang * fix(otlp): guard descriptor writes with semantic ownership markers A pre-existing schema-compatible table named otel_resource_info would silently receive descriptor rows while its missing semantic stamps kept it out of the entity graph. The descriptor write now requires the auto-created table's ownership markers (mito engine + signal_type + source + metric.type=info + metadata_quality=declared) and otherwise degrades to the partial-success warning; the entity-graph gate for the otel whitelist likewise requires metric.type=info, so a user table stamped with only signal/source no longer picks up implicit declarations. Also fold the descriptor write cost into the response and surface the degrade warning through the otel-arrow BatchStatus status_message. Integration tests pin the full marker set on auto-create and that an existing owned descriptor keeps accepting writes without degrading — a missing marker would otherwise silently stop every descriptor write after the first request. Signed-off-by: Dennis Zhuang * perf(otlp): build descriptor rows without the per-resource BTreeMap Projecting a resource allocated a BTreeMap and then collected it into the row key, and every attribute was matched against the allowlist by linear scan. Collect the tags into a Vec and sort once, and match the allowlist instead of scanning it. Measured on the conversion path: descriptor work drops 16-18%, from 10.6% to 8.9% of conversion CPU on the worst shape (1000 resources with 4 data points each), where the cost tracks resource count rather than data-point count. Also trims the comments and tests added with the descriptor to what carries information. Signed-off-by: Dennis Zhuang * test(otlp): pin the descriptor permission-denial degrade path A policy denying the descriptor table must not fail the metrics request, which the fix in 2401b3dd9c does but nothing covered. Verified as a regression guard by mutation: moving the permission check back before the main write makes this test fail. Signed-off-by: Dennis Zhuang * test(otlp): keep legacy mode covered after trimming the unit tests Trimming the descriptor tests dropped the only assertion that legacy mode skips the job/instance remap and the promote filter. Both alter the columns of tables already in use, so fold the check into the legacy conversion test rather than leaving it uncovered. Signed-off-by: Dennis Zhuang * fix(semantic-graph): stop encoding column names into composite entity ids A composite entity id rendered the identifying columns as sorted `col=value` pairs, so the same identity split into one entity per signal: a trace table names its columns service_name and resource_attributes.service.instance.id where a metric table names them job and instance. One service instance became two nodes with two parallel edge sets, breaking the walk from a trace to that instance's metrics. Render an id as its values in declared order instead, escaping the separator so components stay distinguishable, which is what single-column ids already did by keeping only the value. entity_id_attrs still carries the structured form. Values alone are not enough for a namespaced service: the metric side folds service.namespace into job while traces keep the bare name. Add qualified_by to the conventions so the trace declarations compose the namespace the same way, per the OTel rule that job is / or the bare name when the namespace is empty. A table without the namespace column keeps the unqualified identity rather than losing the declaration. Conventions validation now rejects one entity type declared with a different number of id columns by two sources, which would silently produce ids that can never match. Signed-off-by: Dennis Zhuang * style(otlp): import the parent module by crate path check-super-imports.py, part of the CI format gate, rejects a file-level `use super::`. Signed-off-by: Dennis Zhuang * feat(otlp): gate the resource descriptor, and fix what review found Synthesizing greptime_otel_resource_info creates and writes a table the user never sent, so it is now off unless otlp.experimental_enable_resource_info says otherwise. With it off the request costs exactly what it did before the descriptor existed: nothing is projected, no table is created, no write and no permission check happen. Tests run with it on. StandaloneOptions carried no otlp field, so the whole [otlp] section was silently dropped in standalone mode; map it through, or the new option (and trace_ingest_chunk_size before it) would do nothing there. Renamed from otel_resource_info: the greptime_ prefix marks the table as engine-managed and makes a collision with a user metric unlikely, which is what the pre-existing-table ownership check and its per-request catalog lookup were defending against. Both are gone. A request may carry data for several graph windows, but the descriptor folded every data-point time into one row at the newest of them, leaving the earlier windows with metric rows and no entities. Key the rows by window as well, and take the times from the data points the encoder actually writes: it drops exponential histograms, and a resource carrying nothing else was being described as an entity with no measurements. Projecting a resource cloned its attributes once per data point. Nest the windows under the attributes instead, so they are moved once per resource, and walk the data-point times through a visitor rather than collecting a Vec per metric. Also documents what the two maps key and hold, and lifts the projected attribute names to constants beside KEY_SERVICE_NAME. Signed-off-by: Dennis Zhuang * fix(otlp): skip the descriptor's work entirely when it is disabled The collision scan over the request's output tables ran even with the feature off. Short-circuit on the option instead, and update the config snapshot the new [otlp] section changed. Also drops the doc comment orphaned by the deleted ownership check: it had attached itself to the trait impl and described a check that no longer exists. Signed-off-by: Dennis Zhuang * refactor(semantic-graph): drop the expect and name the service identity The CASE is built through Case directly rather than the fallible when().otherwise() builder, so the non-test path no longer carries an expect (architecture-invariants $4). service_identity returned two same-typed Options that both call sites destructured positionally; a named struct makes a swap fail to compile. Also records that id-column order is part of the identity, where the option docs and the conventions authors will read it. Signed-off-by: Dennis Zhuang * chore(semantic-graph): drop comments that narrate the code Signed-off-by: Dennis Zhuang * fix(semantic-graph): cast duration_nano before the trace-table union Trace tables written before the signed-integer ingest change hold duration_nano as UInt64 and later ones as Int64. The calls derivation unions the per-table selects, and the two have no common integer type, so a deployment holding both shapes could not build the plan. The cross-table test now spans both. Signed-off-by: Dennis Zhuang * refactor(semantic-graph): drop the redundant duration_nano casts Signed-off-by: Dennis Zhuang * fix(otlp): decide exponential histogram acceptance in one place The resource descriptor mirrored only the experimental gate, so with both experimental flags on a resource whose only metric is a delta exponential histogram was described as an entity with no measurements. The encoder's whole-metric rules move into exponential_histogram_gate, which both call, and the descriptor takes its timestamps through exponential_histogram_value so per-point rejections drop out too. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- config/config.md | 2 + config/frontend.example.toml | 3 + config/standalone.example.toml | 3 + ...25-entity-relationships-and-graph-query.md | 8 +- src/common/catalog/src/consts.rs | 5 + src/frontend/src/instance.rs | 1 + src/frontend/src/instance/builder.rs | 1 + src/frontend/src/instance/entity_graph.rs | 230 ++++++++- src/frontend/src/instance/otlp.rs | 47 +- src/frontend/src/metrics.rs | 8 + src/frontend/src/service_config/otlp.rs | 7 + src/operator/src/statement/semantic_graph.rs | 202 ++++++-- .../statement/semantic_graph/conventions.rs | 74 ++- .../statement/semantic_graph/conventions.yaml | 50 +- .../statement/semantic_graph/relationships.rs | 79 ++- src/servers/src/http/otlp.rs | 5 +- src/servers/src/otlp/metrics.rs | 467 ++++++++++++++---- src/servers/src/otlp/metrics/resource_info.rs | 393 +++++++++++++++ src/servers/src/otlp/trace.rs | 8 + src/session/src/protocol_ctx.rs | 3 + src/table/src/requests/semantic.rs | 4 + tests-integration/src/instance.rs | 115 +++++ tests-integration/src/standalone.rs | 5 + tests-integration/tests/http.rs | 220 ++++++++- .../common/system/semantic_graph.result | 261 ++++++++-- .../common/system/semantic_graph.sql | 119 +++++ tests/conf/frontend-test.toml.template | 3 + tests/conf/standalone-test.toml.template | 3 + 28 files changed, 2093 insertions(+), 233 deletions(-) create mode 100644 src/servers/src/otlp/metrics/resource_info.rs diff --git a/config/config.md b/config/config.md index 91cb4fe7669..992296d4827 100644 --- a/config/config.md +++ b/config/config.md @@ -76,6 +76,7 @@ | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | | `otlp.experimental_enable_exponential_histogram` | Bool | `false` | Experimental: enable cumulative OTLP exponential histogram ingestion. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | +| `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric
resource attributes, so metrics-only services reach the semantic graph. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | | `prom_store.with_metric_engine` | Bool | `true` | Whether to store the data from Prometheus remote write in metric engine. | @@ -319,6 +320,7 @@ | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | | `otlp.experimental_enable_exponential_histogram` | Bool | `false` | Experimental: enable cumulative OTLP exponential histogram ingestion. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | +| `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric
resource attributes, so metrics-only services reach the semantic graph. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | | `prom_store.with_metric_engine` | Bool | `true` | Whether to store the data from Prometheus remote write in metric engine. | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 3635a63f063..80534855721 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -241,6 +241,9 @@ enable = true experimental_enable_exponential_histogram = false ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 +## Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric +## resource attributes, so metrics-only services reach the semantic graph. +experimental_enable_resource_info = false ## Prometheus remote storage options [prom_store] diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 72ba992cce0..70f74c7f946 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -208,6 +208,9 @@ enable = true experimental_enable_exponential_histogram = false ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 +## Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric +## resource attributes, so metrics-only services reach the semantic graph. +experimental_enable_resource_info = false ## Prometheus remote storage options [prom_store] diff --git a/docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md b/docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md index 995cf6cab7e..1c78e91f252 100644 --- a/docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md +++ b/docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md @@ -149,7 +149,7 @@ The decisions behind the shape: - **`provenance` is part of edge identity**, so a declared edge and a derived edge between the same pair coexist, and a user-declared edge survives when derivation is off. - **`confidence` expresses derivation certainty, not statistical completeness**: `1.0` for a successfully paired or declared edge, `< 1.0` for virtual-node or agent-inferred edges. It does not correct for trace sampling. - **RED metrics describe the observed span-pair population** — the `service_graph` connector semantics. Under sampling, counts understate true traffic, and ratios (error rate, mean duration) are representative only when sampling is unbiased with respect to status and latency — tail sampling that keeps errors and slow traces skews both. -- **Endpoint encoding (v1).** A single-attribute identity is the value verbatim; a composite identity is the attributes sorted by key, rendered `k1=v1,k2=v2` — a human-readable traversal key, not a collision-free canonical encoding. Nodes carry the structured `entity_id_attrs`; edges carry only the string key. Known v1 limitation: composite components containing unescaped `,` or `=` are not guaranteed to be distinguished. +- **Endpoint encoding.** An identity renders as its values in declared order (broad to narrow), joined by `,` with `,` and `\` escaped in each value. The identifying attribute *keys* are deliberately absent: the same identity reaches the graph under different column names per signal — a trace table's `service_name` against a metric table's `job` — and encoding the keys would split one entity into one entity per signal. Where a source carries a qualifier separately (`service.namespace` beside `service.name`) the convention folds it in as `/`, reproducing the identity a source that pre-composed it (Prometheus `job`) already emits. Nodes carry the structured `entity_id_attrs`; edges carry only the string key. - **Hand-declared edges** (`provenance = 'declared'`) live in a physical table of the same shape (plus a business validity window) in `greptime_private` and are unioned into the computed `semantic_relationships`. The relationship vocabulary is small, typed, and inverse-paired; stored direction is `src -> dst`, the inverse is a query concern: @@ -197,7 +197,7 @@ GROUP BY 1, src_id, dst_id; Edge endpoints are built from each trace table's `service` entity declaration (the same id expression the registry uses), so edges land on exactly the entity ids the registry emits, self-calls are judged on the full endpoint id, and a table whose declaration is unusable contributes no edges rather than silently switching identity. The join is bounded by a time-proximity window and the results of **all** trace tables combine into one edge per window with merged RED metrics. One assumption to note: a deployment can route traces to multiple tables (`x-greptime-trace-table-name`), so one distributed trace is not guaranteed to live in a single table — pairing must account for that, and how is an implementation concern. Uninstrumented peers become **virtual nodes**: a client span with no matching server span yields an edge to a synthetic node named from `peer.service`/`db.name`/`server.address`. -**The entity registry.** Per declaring table: filter to the window, project (time bin, casts, JSON assembly), `DISTINCT`. The per-table results are unioned without any cross-source merge: no descriptive-conflict or JSON-merging policy lives in the derivation. Rows whose identity columns are NULL identify nothing and are filtered out; single-column ids are used verbatim, composite ids render the sorted `k=v` form plus the `entity_id_attrs` JSON. Consumers deduplicate with `SELECT DISTINCT entity_type, entity_id`; a unique node set is the snapshot relation's job (see *Querying the graph*). +**The entity registry.** Per declaring table: filter to the window, project (time bin, casts, JSON assembly), `DISTINCT`. The per-table results are unioned without any cross-source merge: no descriptive-conflict or JSON-merging policy lives in the derivation. Rows whose identity columns are NULL identify nothing and are filtered out; ids render as their escaped values joined in declared order, with the `entity_id_attrs` JSON alongside once an id is assembled from more than one column. Consumers deduplicate with `SELECT DISTINCT entity_type, entity_id`; a unique node set is the snapshot relation's job (see *Querying the graph*). **Attribute edges.** Shared attributes provide join keys; co-declaration or a relationship template provides the relationship semantics. A shared value alone determines neither direction nor edge type, so the derivation is rule-based, not join-everything. A table declaring both a `service.instance` and a `host` identity on the same rows derives `runs_on` between them: the row itself witnesses the relationship, and the built-in vocabulary fixes the direction. Cross-table edges require a declared relationship template (endpoint mappings, `rel_type`, direction): a derivation *rule* declared on schema, not per-instance edge data. @@ -284,7 +284,7 @@ Drawbacks: - Read-time derivation costs a self-join over trace tables on every scan of the computed tables. The default one-hour window bounds it, but very large deployments will want materialisation (Future Work). - `semantic_entities` exposes per-observation rows; deduplication falls on the consumer (`DISTINCT`) until the snapshot relation lands in M2. -- The v1 `k=v` id encoding can collide on hostile values (unescaped `,`/`=`); hardening is an open question. +- The id encoding is readable rather than canonical: escaping keeps components distinguishable, but a value containing the qualifier separator (`/`) is still ambiguous against a qualified name. - The graph is only as connected as the identity columns tables actually share: the `job`/`service_name` misalignment above is the canonical example, and enrichment discipline gates connectivity. # Implementation Plan @@ -306,7 +306,7 @@ Drawbacks: # Open Questions -1. **Endpoint encoding hardening.** The v1 `k=v` rendering (with `entity_id_attrs` JSON as the source of truth) is readable but collides when values contain `,`/`=`. Open: a fixed-width `entity_id_hash` join key (New-Relic-style); typed value encoding; the criteria for introducing the `semantic_key=column` mapping form; `schema_url`-gated entity merge. Related and unresolved: whether the identifying attribute *key* should participate in entity equality — a single-attribute id keeps only the value, so `job=api` and `service_name=api` currently name the same entity even when the keys differ semantically. +1. **Endpoint encoding hardening.** The rendering (with `entity_id_attrs` JSON as the source of truth) is readable rather than canonical. Open: a fixed-width `entity_id_hash` join key (New-Relic-style); typed value encoding; the criteria for introducing the `semantic_key=column` mapping form; `schema_url`-gated entity merge. Settled: identifying attribute *keys* do not participate in entity equality, so `job=api` and `service_name=api` name the same entity — the property that lets one service reach the graph from several signals. 2. **Identity reconciliation across signals.** Trace edges key on the trace table's service declaration; host edges key on `host.id`/`k8s.pod.uid`. Linking them relies on tables carrying both columns; how aggressively to enrich (a `k8sattributes`-equivalent) gates how connected the graph is. 3. **`GRAPH_TABLE` surface scope.** How much of SQL/PGQ to implement vs. a minimal single-hop subset. Start minimal and grow with demand. 4. **Visibility beyond table grants.** The derivation contract filters by table-read permission; whether topology needs finer policies (row-level, tenant-scoped visibility of cross-schema edges) is open. diff --git a/src/common/catalog/src/consts.rs b/src/common/catalog/src/consts.rs index b976cbc256e..63aa103b83d 100644 --- a/src/common/catalog/src/consts.rs +++ b/src/common/catalog/src/consts.rs @@ -220,6 +220,11 @@ pub const SEMANTIC_RELATIONSHIPS_TABLE_NAME: &str = "semantic_relationships"; /// `semantic_relationships`; see [`is_ddl_reserved_table`] for its lifecycle. pub const SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME: &str = "semantic_relationships_declared"; +/// Width of the window the graph derivation bins observations into. Sources +/// synthesizing observations must land a row in every window they describe, +/// so this is shared rather than restated per crate. +pub const SEMANTIC_GRAPH_WINDOW_NANOS: i64 = 60 * 1_000_000_000; + // Column names of the graph tables: `catalog` exposes these schemas and the // read-time plans in `operator` must project exactly them. pub const OBSERVED_AT_COLUMN: &str = "observed_at"; diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 78b701a3442..412ff495f65 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -134,6 +134,7 @@ pub struct Instance { slow_query_options: SlowQueryOptions, influxdb_default_merge_mode: InfluxdbMergeMode, trace_ingest_chunk_size: usize, + otlp_resource_info: bool, suspend: Arc, // cache for otlp metrics diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index c49a8394beb..0f87dc424f0 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -367,6 +367,7 @@ impl FrontendBuilder { slow_query_options: self.options.slow_query.clone(), influxdb_default_merge_mode: self.options.influxdb.default_merge_mode, trace_ingest_chunk_size: self.options.otlp.trace_ingest_chunk_size, + otlp_resource_info: self.options.otlp.experimental_enable_resource_info, suspend: Arc::new(AtomicBool::new(false)), }) } diff --git a/src/frontend/src/instance/entity_graph.rs b/src/frontend/src/instance/entity_graph.rs index 14f18ed1b79..a0817cf2806 100644 --- a/src/frontend/src/instance/entity_graph.rs +++ b/src/frontend/src/instance/entity_graph.rs @@ -23,7 +23,7 @@ //! the query engine. Injected into the catalog manager after the engine is built, //! breaking the `catalog -> query` cycle. -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::sync::Weak; use async_trait::async_trait; @@ -60,8 +60,9 @@ use table::TableRef; use table::metadata::TableInfo; use table::predicate::{TimeRangeExtraction, extract_time_range_strict}; use table::requests::{ - EntityRole, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_METRIC, SOURCE_PROMETHEUS, - is_trace_v1_table, parse_entity_columns, parse_entity_option_key, + EntityRole, SEMANTIC_METRIC_TYPE, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_METRIC, + SOURCE_OPENTELEMETRY, SOURCE_PROMETHEUS, is_trace_v1_table, parse_entity_columns, + parse_entity_option_key, }; use crate::error; @@ -177,6 +178,7 @@ impl EntityGraphProviderImpl { time_index: time_index.clone(), entity_type, id_columns, + id_qualifier: None, descriptive_columns, scope_columns, }) @@ -188,10 +190,11 @@ impl EntityGraphProviderImpl { /// All entity declarations of one table: the explicit options plus the /// zero-configuration conventions (`otlp_trace_entities` for trace-v1 /// tables — including the `service` identity of tables created before the - /// ingest-side auto-stamp — and the Prometheus descriptor whitelist). An - /// explicit declaration of a type always suppresses the implicit one, even - /// when the explicit declaration is invalid and skipped: silently falling - /// back would change entity identity behind the user's back. + /// ingest-side auto-stamp — and the Prometheus/OTel descriptor + /// whitelists). An explicit declaration of a type always suppresses the + /// implicit one, even when the explicit declaration is invalid and + /// skipped: silently falling back would change entity identity behind the + /// user's back. fn declarations_for( table_info: &TableInfo, conventions: &Conventions, @@ -204,7 +207,20 @@ impl EntityGraphProviderImpl { &mut declarations, ); } - Self::extend_with_prometheus_conventions(table_info, conventions, &mut declarations); + Self::extend_with_info_metric_conventions( + table_info, + &conventions.prometheus_info_metrics, + SOURCE_PROMETHEUS, + None, + &mut declarations, + ); + Self::extend_with_info_metric_conventions( + table_info, + &conventions.otel_info_metrics, + SOURCE_OPENTELEMETRY, + Some(servers::semantic::METRIC_TYPE_INFO), + &mut declarations, + ); declarations } @@ -216,29 +232,34 @@ impl EntityGraphProviderImpl { }) } - /// Implicit declarations of the well-known Prometheus entity-descriptor - /// metrics (the `prometheus_info_metrics` whitelist of `conventions.yaml`), - /// gated on the ingest-stamped `signal_type=metric` + `source=prometheus` - /// options and keyed by table name. The metric engine's physical table + /// Implicit declarations of the well-known entity-descriptor metrics + /// (the `prometheus_info_metrics` / `otel_info_metrics` whitelists of + /// `conventions.yaml`), gated on the ingest-stamped `signal_type=metric` + /// option plus the whitelist's expected `source`; OTel descriptors also + /// require `metric.type=info`. The metric engine's physical table /// aggregates every logical table's columns and must not contribute a /// duplicate source. - fn extend_with_prometheus_conventions( + fn extend_with_info_metric_conventions( table_info: &TableInfo, - conventions: &Conventions, + whitelist: &BTreeMap>, + expected_source: &str, + expected_metric_type: Option<&str>, declarations: &mut Vec, ) { - let Some(implicit_entities) = conventions.prometheus_info_metrics.get(&table_info.name) - else { + let Some(implicit_entities) = whitelist.get(&table_info.name) else { return; }; let options = &table_info.meta.options.extra_options; if options.get(SEMANTIC_SIGNAL_TYPE).map(String::as_str) != Some(SIGNAL_TYPE_METRIC) - || options.get(SEMANTIC_SOURCE).map(String::as_str) != Some(SOURCE_PROMETHEUS) + || options.get(SEMANTIC_SOURCE).map(String::as_str) != Some(expected_source) + || expected_metric_type.is_some_and(|expected| { + options.get(SEMANTIC_METRIC_TYPE).map(String::as_str) != Some(expected) + }) || table_info.is_physical_table() { debug!( - "Table `{}` matches the info-metric whitelist but is not a prometheus logical \ - metric table; skipping its implicit declarations", + "Table `{}` matches the info-metric whitelist but is not an eligible \ + `{expected_source}` info-metric source; skipping its implicit declarations", table_info.name ); return; @@ -296,12 +317,19 @@ impl EntityGraphProviderImpl { .cloned() .collect() }; + // A table predating the qualifier column keeps the unqualified + // identity rather than losing the declaration. + let id_qualifier = implicit + .qualified_by + .clone() + .filter(|c| schema.column_schema_by_name(c).is_some()); declarations.push(EntityDeclaration { schema: table_info.schema_name.clone(), table: table_info.name.clone(), time_index: time_index.clone(), entity_type: implicit.entity.clone(), id_columns: implicit.id.clone(), + id_qualifier, descriptive_columns, scope_columns: vec![], }); @@ -832,6 +860,24 @@ mod tests { declarations[2].id_columns, vec!["service_name", "resource_attributes.service.instance.id"] ); + assert_eq!(declarations[1].id_qualifier, None); + + let namespaced = table_info( + &[ + "service_name", + "resource_attributes.service.namespace", + "resource_attributes.service.instance.id", + ], + &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)], + ); + for declaration in sorted_declarations(&namespaced) { + assert_eq!( + declaration.id_qualifier.as_deref(), + Some("resource_attributes.service.namespace"), + "{} must qualify its identity like the metric side's job", + declaration.entity_type + ); + } // Missing uid column: no pod entity synthesized, no name-based guess. let no_uid = table_info( @@ -856,6 +902,47 @@ mod tests { assert!(sorted_declarations(&invalid_explicit).is_empty()); } + #[test] + fn trace_table_host_and_container_require_stable_ids() { + let with_ids = table_info( + &[ + "service_name", + "resource_attributes.host.id", + "resource_attributes.host.name", + "resource_attributes.container.id", + ], + &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)], + ); + let declarations = sorted_declarations(&with_ids); + let types: Vec<&str> = declarations + .iter() + .map(|d| d.entity_type.as_str()) + .collect(); + assert_eq!(types, vec!["container", "host", "service"]); + assert_eq!( + declarations[1].id_columns, + vec!["resource_attributes.host.id"] + ); + assert_eq!( + declarations[1].descriptive_columns, + vec!["resource_attributes.host.name"] + ); + + let names_only = table_info( + &[ + "service_name", + "resource_attributes.host.name", + "resource_attributes.container.name", + ], + &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)], + ); + let types: Vec = sorted_declarations(&names_only) + .into_iter() + .map(|d| d.entity_type) + .collect(); + assert_eq!(types, vec!["service"]); + } + #[test] fn prometheus_implicit_declarations_are_gated() { let labels: &[&str] = &["namespace", "pod", "node"]; @@ -900,6 +987,111 @@ mod tests { assert_eq!(declarations[1].id_columns, vec!["pod"]); } + const OTEL_STAMPS: &[(&str, &str)] = &[ + (SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC), + (SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY), + (SEMANTIC_METRIC_TYPE, servers::semantic::METRIC_TYPE_INFO), + ]; + + #[test] + fn greptime_otel_resource_info_gets_implicit_declarations() { + let info = prom_table_info( + "greptime_otel_resource_info", + &[ + "job", + "instance", + "service.name", + "service.namespace", + "host.id", + "host.name", + "container.id", + "container.name", + "k8s.pod.uid", + "k8s.pod.name", + "k8s.namespace.name", + ], + OTEL_STAMPS, + ); + let declarations = sorted_declarations(&info); + let types: Vec<&str> = declarations + .iter() + .map(|d| d.entity_type.as_str()) + .collect(); + assert_eq!( + types, + vec![ + "container", + "host", + "k8s.pod", + "service", + "service.instance" + ] + ); + assert_eq!(declarations[1].id_columns, vec!["host.id"]); + assert_eq!(declarations[1].descriptive_columns, vec!["host.name"]); + assert_eq!(declarations[3].id_columns, vec!["job"]); + assert_eq!( + declarations[3].descriptive_columns, + vec!["service.name", "service.namespace"] + ); + assert_eq!(declarations[4].id_columns, vec!["job", "instance"]); + assert!(declarations[4].descriptive_columns.is_empty()); + + let partial = prom_table_info( + "greptime_otel_resource_info", + &["job", "service.name", "host.id"], + OTEL_STAMPS, + ); + let types: Vec = sorted_declarations(&partial) + .into_iter() + .map(|d| d.entity_type) + .collect(); + assert_eq!(types, vec!["host", "service"]); + } + + #[test] + fn otel_implicit_declarations_are_gated() { + let labels: &[&str] = &["job", "instance", "host.id"]; + assert!( + sorted_declarations(&prom_table_info( + "greptime_otel_resource_info", + labels, + PROM_STAMPS + )) + .is_empty() + ); + let mut stamps = OTEL_STAMPS.to_vec(); + stamps.push((PHYSICAL_TABLE_METADATA_KEY, "true")); + assert!( + sorted_declarations(&prom_table_info( + "greptime_otel_resource_info", + labels, + &stamps + )) + .is_empty() + ); + assert!( + conventions() + .unwrap() + .otel_info_metrics + .contains_key(servers::otlp::metrics::OTEL_RESOURCE_INFO_TABLE_NAME) + ); + + let wrong_type = [ + (SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC), + (SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY), + (SEMANTIC_METRIC_TYPE, servers::semantic::METRIC_TYPE_GAUGE), + ]; + assert!( + sorted_declarations(&prom_table_info( + "greptime_otel_resource_info", + labels, + &wrong_type + )) + .is_empty() + ); + } + #[test] fn target_info_descriptive_rest_covers_remaining_tags() { let info = prom_table_info( diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index 08086198536..ccc09c9978c 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -50,7 +50,7 @@ use table::requests::{ use self::trace_ingest::trace_conventions; use crate::instance::Instance; -use crate::metrics::{OTLP_LOGS_ROWS, OTLP_METRICS_ROWS}; +use crate::metrics::{OTLP_LOGS_ROWS, OTLP_METRICS_ROWS, OTLP_RESOURCE_INFO_WRITE_ERRORS}; fn trace_permission_targets( table_name: &str, @@ -119,9 +119,15 @@ impl OpenTelemetryProtocolHandler for Instance { .cloned() .unwrap_or_default(); metric_ctx.is_legacy = is_legacy; + metric_ctx.resource_info = self.otlp_resource_info; - let (requests, rows, semantic_index, mut outcome) = - otlp::metrics::to_grpc_insert_requests(request, &mut metric_ctx)?; + let otlp::metrics::MetricsConversion { + requests, + rows, + semantic_index, + resource_info, + mut outcome, + } = otlp::metrics::to_grpc_insert_requests(request, &mut metric_ctx)?; if outcome.rejected_data_points > 0 { warn!( "Rejected {} OTLP exponential histogram data points: {}", @@ -156,7 +162,7 @@ impl OpenTelemetryProtocolHandler for Instance { // OTLP tables have one sample field in both the legacy and physical paths. let output = if metric_ctx.is_legacy || !metric_ctx.with_metric_engine { - self.handle_row_inserts(requests, ctx, false, true) + self.handle_row_inserts(requests, ctx.clone(), false, true) .await .map_err(BoxedError::new) .context(error::ExecuteGrpcQuerySnafu) @@ -165,12 +171,43 @@ impl OpenTelemetryProtocolHandler for Instance { .extension(PHYSICAL_TABLE_PARAM) .unwrap_or(GREPTIME_PHYSICAL_TABLE) .to_string(); - self.handle_metric_row_inserts(requests, ctx, physical_table) + self.handle_metric_row_inserts(requests, ctx.clone(), physical_table) .await .map_err(BoxedError::new) .context(error::ExecuteGrpcQuerySnafu) }?; outcome.write_cost = output.meta.cost; + + // Derived enrichment, written after the metric data is committed: + // failing here would make the client retry data the server already + // accepted, so every failure degrades to a warning instead. + if let Some(resource_info) = resource_info { + let written = match self.check_row_insert_permission( + &resource_info, + &ctx, + PermissionReq::Action(OTLP_WRITE), + ) { + Ok(_) => self + .handle_row_inserts(resource_info, ctx, false, false) + .await + .map_err(BoxedError::new) + .map_err(|e| e.to_string()), + Err(e) => Err(e.to_string()), + }; + match written { + Ok(descriptor_output) => outcome.write_cost += descriptor_output.meta.cost, + Err(e) => { + OTLP_RESOURCE_INFO_WRITE_ERRORS.inc(); + warn!("Failed to write the OTLP resource descriptor table: {e}"); + outcome.error_message.get_or_insert(format!( + "metric data was accepted, but writing the resource \ + descriptor table `{}` failed: {e}", + otlp::metrics::OTEL_RESOURCE_INFO_TABLE_NAME + )); + } + } + } + Ok(outcome) } diff --git a/src/frontend/src/metrics.rs b/src/frontend/src/metrics.rs index aba33637cfc..f649db0b46f 100644 --- a/src/frontend/src/metrics.rs +++ b/src/frontend/src/metrics.rs @@ -45,6 +45,14 @@ lazy_static! { ) .unwrap(); + /// Failed writes of the synthesized OTLP resource descriptor table; these + /// surface as OTLP partial-success warnings, not request failures. + pub static ref OTLP_RESOURCE_INFO_WRITE_ERRORS: IntCounter = register_int_counter!( + "greptime_frontend_otlp_resource_info_write_errors", + "frontend otlp resource descriptor write errors" + ) + .unwrap(); + /// The number of OpenTelemetry traces send by frontend node. pub static ref OTLP_TRACES_ROWS: IntCounter = register_int_counter!( "greptime_frontend_otlp_traces_rows", diff --git a/src/frontend/src/service_config/otlp.rs b/src/frontend/src/service_config/otlp.rs index 5dcab2c8ce2..d6a36f702e4 100644 --- a/src/frontend/src/service_config/otlp.rs +++ b/src/frontend/src/service_config/otlp.rs @@ -23,6 +23,11 @@ pub struct OtlpOptions { pub experimental_enable_exponential_histogram: bool, /// Maximum spans per trace ingest chunk. Set to 0 to disable splitting. pub trace_ingest_chunk_size: usize, + /// Whether to synthesize the `greptime_otel_resource_info` descriptor + /// table from the resource attributes of OTLP metrics, so metrics-only + /// services reach the semantic graph. Off by default: it creates and + /// writes a table the user did not send. + pub experimental_enable_resource_info: bool, } impl Default for OtlpOptions { @@ -31,6 +36,7 @@ impl Default for OtlpOptions { enable: true, experimental_enable_exponential_histogram: false, trace_ingest_chunk_size: DEFAULT_TRACE_INGEST_CHUNK_SIZE, + experimental_enable_resource_info: false, } } } @@ -45,6 +51,7 @@ mod tests { assert!(default.enable); assert!(!default.experimental_enable_exponential_histogram); assert_eq!(default.trace_ingest_chunk_size, 512); + assert!(!default.experimental_enable_resource_info); let options: OtlpOptions = toml::from_str("enable = false").unwrap(); assert!(!options.enable); diff --git a/src/operator/src/statement/semantic_graph.rs b/src/operator/src/statement/semantic_graph.rs index 7014a94f8ac..86827cf8c7a 100644 --- a/src/operator/src/statement/semantic_graph.rs +++ b/src/operator/src/statement/semantic_graph.rs @@ -41,8 +41,9 @@ use common_catalog::consts::{ ENTITY_DESCRIPTIVE_COLUMN, ENTITY_ID_ATTRS_COLUMN, ENTITY_ID_COLUMN, ENTITY_SCOPE_COLUMN, ENTITY_TYPE_COLUMN, ERROR_COUNT_COLUMN, FRESH_UNTIL_COLUMN, GENERATION_ID_COLUMN, OBSERVED_AT_COLUMN, PROVENANCE_COLUMN, REL_TYPE_COLUMN, REQUEST_COUNT_COLUMN, - SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SOURCE_TABLES_COLUMN, SRC_ID_COLUMN, - SRC_TYPE_COLUMN, VALID_FROM_COLUMN, VALID_UNTIL_COLUMN, WINDOW_END_COLUMN, WINDOW_START_COLUMN, + SEMANTIC_GRAPH_WINDOW_NANOS, SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SOURCE_TABLES_COLUMN, + SRC_ID_COLUMN, SRC_TYPE_COLUMN, VALID_FROM_COLUMN, VALID_UNTIL_COLUMN, WINDOW_END_COLUMN, + WINDOW_START_COLUMN, }; use common_function::function::FunctionContext; use common_function::function_registry::FUNCTION_REGISTRY; @@ -60,7 +61,7 @@ use datafusion::dataframe::DataFrame; use datafusion::functions::{core as core_fns, datetime as datetime_fns, string as string_fns}; use datafusion::functions_nested::expr_fn::make_array; use datafusion_common::{Column, Result as DfResult, ScalarValue}; -use datafusion_expr::{Expr, LogicalPlan, ScalarUDF, cast, ident, lit}; +use datafusion_expr::{Case, Expr, LogicalPlan, ScalarUDF, cast, ident, lit}; pub use relationships::{ CallsSource, CoDeclaredSource, DeclaredSource, RelationshipSources, build_relationships_plan, }; @@ -122,9 +123,10 @@ pub fn declared_relationships_schema_matches(table_info: &table::metadata::Table }) } -/// Bin width for the temporal window of derived rows: 60s buckets, matching the -/// service-graph convention. -const BIN_NANOS: i64 = 60 * 1_000_000_000; +/// Bin width for the temporal window of derived rows, matching the +/// service-graph convention. Shared so ingestion-synthesized observations +/// land in the same buckets. +const BIN_NANOS: i64 = SEMANTIC_GRAPH_WINDOW_NANOS; /// Default retention for the declared-edge table; expiry slides the topology window. const DEFAULT_DECLARED_RELATIONSHIPS_TTL: &str = "90d"; @@ -284,8 +286,12 @@ pub struct EntityDeclaration { /// The table's time index column, used for the temporal window filter. pub time_index: String, pub entity_type: String, - /// Identifying columns (>= 1). One column → id verbatim; several → composite. + /// Identifying columns (>= 1), ordered broad to narrow. pub id_columns: Vec, + /// Optional column qualifying the first identity component, so a source + /// carrying the parts separately (a trace table's `service.namespace`) + /// yields the same id as one carrying them pre-composed (`job`). + pub id_qualifier: Option, /// Descriptive columns snapshotted into the `descriptive` JSON (may be empty). pub descriptive_columns: Vec, /// Scope columns (namespace/environment). One column → scope verbatim; @@ -417,28 +423,91 @@ fn cast_string_or_empty(column: &str) -> Expr { core_fns::coalesce().call(vec![cast(ident(column), DataType::Utf8), lit("")]) } -/// The canonical entity-id expression for `id_columns`: the value verbatim for -/// a single column, the sorted `k=v,k=v` rendering for a composite. `col` -/// constructs the column reference (unqualified for registry branches, -/// join-side-qualified for the calls derivation). /// A row identifies an entity only when every identity component is present /// and non-empty: kube-state-metrics descriptors emit empty-string labels (an /// unscheduled pod's `node`, an owner-less pod's `owner_*`), and an empty -/// string is never a meaningful entity id. +/// string is never a meaningful entity id. The qualifier is optional by +/// construction and so is not guarded. pub(crate) fn identifies(column: &str) -> Expr { ident(column) .is_not_null() .and(cast(ident(column), DataType::Utf8).not_eq(lit(""))) } -fn entity_id_expr(id_columns: &[String], col: &dyn Fn(&str) -> Expr) -> Expr { - if let [id] = id_columns { - cast(col(id), DataType::Utf8) - } else { - let mut cols = id_columns.to_vec(); - cols.sort(); - sorted_kv_expr_with(&cols, false, col) +const ID_SEPARATOR: &str = ","; +const ID_ESCAPE: &str = "\\"; +/// Left unescaped: `/` is how Prometheus renders `job`. +const ID_QUALIFIER_SEPARATOR: &str = "/"; + +/// Escapes so a composite id decodes back to its components. The escape +/// character goes first, or it would double the escapes the separator pass +/// introduces. +fn escaped_id_value(value: Expr) -> Expr { + let escape_escape = string_fns::replace().call(vec![ + value, + lit(ID_ESCAPE), + lit(format!("{ID_ESCAPE}{ID_ESCAPE}")), + ]); + string_fns::replace().call(vec![ + escape_escape, + lit(ID_SEPARATOR), + lit(format!("{ID_ESCAPE}{ID_SEPARATOR}")), + ]) +} + +/// The identity values in declared order (broad to narrow), escaped and +/// joined. +/// +/// The identifying *column names* are deliberately absent: the same identity +/// reaches us under different names per source — a trace table's +/// `service_name` against a metric table's `job` — and encoding them would +/// split one entity into one per signal. `entity_id_attrs` keeps the +/// structured form. +/// +/// `col` constructs the column reference (unqualified for registry branches, +/// join-side-qualified for the calls derivation). +fn entity_id_expr( + id_columns: &[String], + qualifier: Option<&str>, + col: &dyn Fn(&str) -> Expr, +) -> Expr { + let mut parts = Vec::with_capacity(id_columns.len() * 2); + for (index, column) in id_columns.iter().enumerate() { + if index > 0 { + parts.push(lit(ID_SEPARATOR)); + } + let value = escaped_id_value(cast(col(column), DataType::Utf8)); + parts.push(match qualifier { + Some(qualifier) if index == 0 => qualified_id_expr(qualifier, value, col), + _ => value, + }); } + if let [single] = parts.as_slice() { + single.clone() + } else { + concat_expr(parts) + } +} + +/// `/`, or bare `value` when the qualifier is empty on this +/// row — the spec's rule composing `job` from `service.namespace` and +/// `service.name`. +fn qualified_id_expr(qualifier: &str, value: Expr, col: &dyn Fn(&str) -> Expr) -> Expr { + let qualifier = escaped_id_value( + core_fns::coalesce().call(vec![cast(col(qualifier), DataType::Utf8), lit("")]), + ); + Expr::Case(Case::new( + None, + vec![( + Box::new(qualifier.clone().eq(lit(""))), + Box::new(value.clone()), + )], + Some(Box::new(concat_expr(vec![ + qualifier, + lit(ID_QUALIFIER_SEPARATOR), + value, + ]))), + )) } /// The `parse_json` UDF, shared by all derivation plans. Resolved from the @@ -583,14 +652,22 @@ fn registry_source( let mut rows = Vec::with_capacity(1 + rest.len()); for decl in std::iter::once(first).chain(rest) { // CAST even a single-column id: id columns need not be strings, and - // the computed table declares entity_id STRING. Composite ids - // additionally carry a JSON object of the id columns in - // entity_id_attrs. - let entity_id = entity_id_expr(&decl.id_columns, &|c| ident(c)); - let entity_id_attrs = if decl.id_columns.len() == 1 { + // the computed table declares entity_id STRING. An id assembled from + // more than one column additionally carries its parts as a JSON + // object in entity_id_attrs, the structured form of the id string. + let entity_id = entity_id_expr(&decl.id_columns, decl.id_qualifier.as_deref(), &|c| { + ident(c) + }); + let id_parts = decl + .id_qualifier + .iter() + .chain(&decl.id_columns) + .cloned() + .collect::>(); + let entity_id_attrs = if id_parts.len() == 1 { null_json() } else { - json_object_expr(&decl.id_columns) + json_object_expr(&id_parts) }; let scope = match decl.scope_columns.as_slice() { @@ -811,11 +888,74 @@ mod tests { time_index: "ts".to_string(), entity_type: entity_type.to_string(), id_columns: id_columns.iter().map(|s| s.to_string()).collect(), + id_qualifier: None, descriptive_columns: vec![], scope_columns: vec![], } } + /// The id string must decode back to its components, and a qualifier must + /// reproduce the identity a source that pre-composed it already emits. + #[tokio::test] + async fn registry_id_escaping_and_qualifier() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("service_name", DataType::Utf8, false), + Field::new("instance", DataType::Utf8, false), + Field::new("namespace", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_000, 2_000, 3_000])) as ArrayRef, + Arc::new(StringArray::from(vec!["cart", "a,b", "we\\ird"])), + Arc::new(StringArray::from(vec!["i-1", "c", "i-3"])), + Arc::new(StringArray::from(vec![Some("shop"), None, Some("")])), + ], + ) + .unwrap(); + let ctx = SessionContext::new(); + ctx.register_table( + "svc", + Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), + ) + .unwrap(); + + let mut declaration = decl("service.instance", &["service_name", "instance"]); + declaration.id_qualifier = Some("namespace".to_string()); + declaration.time_index = "ts".to_string(); + let plan = build_registry_plan( + vec![RegistrySource { + declarations: vec![declaration], + scan: ctx.table("svc").await.unwrap(), + }], + &test_window(), + ) + .unwrap() + .unwrap(); + + let mut ids: Vec = collect(&ctx, plan) + .await + .iter() + .flat_map(|batch| strings(batch, 5)) + .collect(); + ids.sort(); + assert_eq!( + ids, + vec![ + // an absent qualifier leaves the identity bare, and a value + // holding the separator stays distinguishable from two values + "a\\,b,c".to_string(), + "shop/cart,i-1".to_string(), + "we\\\\ird,i-3".to_string(), + ] + ); + } + #[tokio::test] async fn registry_single_column_identity() { let ctx = metric_table_ctx(); @@ -900,24 +1040,24 @@ mod tests { .collect(); rows.sort(); - // Composite id -> sorted `k=v,k=v` plus a JSON object of the id columns; - // descriptive JSON keeps `\`, `"` and control characters intact in - // runtime values, NULL -> "". + // Composite id -> values in declared order plus a JSON object of the + // id columns; descriptive JSON keeps `\`, `"` and control characters + // intact in runtime values, NULL -> "". assert_eq!( rows, vec![ ( - "pid=42,service_name=cart".to_string(), + "cart,42".to_string(), Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()), Some(r#"{"host":""}"#.to_string()), ), ( - "pid=42,service_name=cart".to_string(), + "cart,42".to_string(), Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()), Some(r#"{"host":"h2"}"#.to_string()), ), ( - "pid=42,service_name=cart".to_string(), + "cart,42".to_string(), Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()), Some(r#"{"host":"we\"ird\\\nhost"}"#.to_string()), ), diff --git a/src/operator/src/statement/semantic_graph/conventions.rs b/src/operator/src/statement/semantic_graph/conventions.rs index 36d058921b5..9ec6f52e158 100644 --- a/src/operator/src/statement/semantic_graph/conventions.rs +++ b/src/operator/src/statement/semantic_graph/conventions.rs @@ -19,7 +19,7 @@ //! with the binary, not an operator-editable configuration surface; explicit //! `greptime.semantic.entity.*` declarations always override it. -use std::collections::{BTreeMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::LazyLock; use serde::Deserialize; @@ -43,15 +43,20 @@ pub struct VirtualDstCandidate { pub connection_type: String, } -/// One implicit entity declaration: of a whitelisted Prometheus info metric, -/// or of a trace-v1 table's flattened resource attributes. +/// One implicit entity declaration: of a whitelisted Prometheus or OTel info +/// metric, or of a trace-v1 table's flattened resource attributes. #[derive(Debug, Deserialize)] #[serde(deny_unknown_fields)] pub struct ImplicitEntity { pub entity: String, - /// Identifying label columns; every one must exist on the table for the - /// declaration to apply. + /// Identifying label columns, ordered broad to narrow; every one must + /// exist on the table for the declaration to apply. pub id: Vec, + /// Column qualifying `id[0]` as `/`, so a source + /// carrying the parts separately matches one carrying them pre-composed. + /// Skipped when the column is absent or empty. + #[serde(default)] + pub qualified_by: Option, /// Descriptive label columns, filtered to those present (kube-state-metrics /// label sets vary across versions). #[serde(default)] @@ -70,7 +75,12 @@ pub struct Conventions { pub trace_co_declared_edges: Vec, pub virtual_dst_candidates: Vec, pub otlp_trace_entities: Vec, + /// Table name -> the entities that table declares, for Prometheus-sourced + /// descriptor metrics (`source = prometheus`). pub prometheus_info_metrics: BTreeMap>, + /// Table name -> the entities that table declares, for OTLP-sourced + /// descriptor tables (`source = opentelemetry`). + pub otel_info_metrics: BTreeMap>, } /// The built-in entity-type vocabulary. User-declared types are open-ended; @@ -78,6 +88,7 @@ pub struct Conventions { pub const ENTITY_TYPE_SERVICE: &str = "service"; pub const ENTITY_TYPE_SERVICE_INSTANCE: &str = "service.instance"; pub const ENTITY_TYPE_HOST: &str = "host"; +pub const ENTITY_TYPE_CONTAINER: &str = "container"; pub const ENTITY_TYPE_PROCESS: &str = "process"; pub const ENTITY_TYPE_K8S_POD: &str = "k8s.pod"; pub const ENTITY_TYPE_K8S_NODE: &str = "k8s.node"; @@ -88,10 +99,11 @@ pub const ENTITY_TYPE_GEN_AI_AGENT: &str = "gen_ai.agent"; pub const ENTITY_TYPE_GEN_AI_MODEL: &str = "gen_ai.model"; pub const ENTITY_TYPE_GEN_AI_TOOL: &str = "gen_ai.tool"; -const ENTITY_TYPES: [&str; 12] = [ +const ENTITY_TYPES: [&str; 13] = [ ENTITY_TYPE_SERVICE, ENTITY_TYPE_SERVICE_INSTANCE, ENTITY_TYPE_HOST, + ENTITY_TYPE_CONTAINER, ENTITY_TYPE_PROCESS, ENTITY_TYPE_K8S_POD, ENTITY_TYPE_K8S_NODE, @@ -200,11 +212,16 @@ fn validate(conventions: &Conventions) -> Result<(), String> { let per_table = conventions .prometheus_info_metrics .iter() + .chain(&conventions.otel_info_metrics) .map(|(table, entities)| (table.as_str(), entities)) .chain(std::iter::once(( "otlp traces", &conventions.otlp_trace_entities, ))); + // An entity type declared with a different number of id columns by two + // sources yields ids that can never match, silently splitting one entity + // in two. + let mut arity = HashMap::new(); for (table, entities) in per_table { let mut seen_types = HashSet::new(); for implicit in entities { @@ -226,6 +243,12 @@ fn validate(conventions: &Conventions) -> Result<(), String> { implicit.entity )); } + if implicit.qualified_by.as_ref().is_some_and(String::is_empty) { + return Err(format!( + "entity `{}` of info metric `{table}` has an empty qualified_by", + implicit.entity + )); + } if implicit.descriptive_rest && !implicit.descriptive.is_empty() { return Err(format!( "entity `{}` of info metric `{table}` sets both descriptive and \ @@ -233,6 +256,17 @@ fn validate(conventions: &Conventions) -> Result<(), String> { implicit.entity )); } + match arity.insert(implicit.entity.as_str(), implicit.id.len()) { + Some(previous) if previous != implicit.id.len() => { + return Err(format!( + "entity `{}` is declared with {previous} id columns elsewhere but \ + {} for `{table}`", + implicit.entity, + implicit.id.len() + )); + } + _ => {} + } } } Ok(()) @@ -249,25 +283,28 @@ mod tests { /// Each case must fail on exactly the rule it names, so the shared /// boilerplate is valid and the mutated part is inside the vocabulary. - fn broken(edges: &str, info_metrics: &str) -> String { + fn broken(edges: &str, info_metrics: &str, otel_metrics: &str) -> String { format!( "co_declared_edges: [{edges}]\ntrace_co_declared_edges: []\n\ virtual_dst_candidates: []\notlp_trace_entities: []\n\ - prometheus_info_metrics: {{{info_metrics}}}" + prometheus_info_metrics: {{{info_metrics}}}\n\ + otel_info_metrics: {{{otel_metrics}}}" ) } #[test] fn validation_rejects_broken_conventions() { - let err = |edges, info| parse(&broken(edges, info)).unwrap_err(); + let err = |edges, info, otel| parse(&broken(edges, info, otel)).unwrap_err(); - assert!(err("{src: host, dst: service, rel: pets}", "").contains("unknown rel_type")); + assert!(err("{src: host, dst: service, rel: pets}", "", "").contains("unknown rel_type")); assert!( - err("{src: k8s.pods, dst: k8s.node, rel: runs_on}", "").contains("unknown entity type") + err("{src: k8s.pods, dst: k8s.node, rel: runs_on}", "", "") + .contains("unknown entity type") ); assert!( err( "{src: host, dst: service, rel: uses}, {src: host, dst: service, rel: uses}", + "", "" ) .contains("duplicate edge rule") @@ -275,14 +312,27 @@ mod tests { assert!( err( "", - "t: [{entity: host, id: [x], descriptive: [y], descriptive_rest: true}]" + "t: [{entity: host, id: [x], descriptive: [y], descriptive_rest: true}]", + "" ) .contains("descriptive_rest") ); + // the otel map runs through the same per-table validation + assert!(err("", "", "t: [{entity: hosts, id: [x]}]").contains("unknown entity type")); + // ids of a different arity for one type can never match each other + assert!( + err( + "", + "a: [{entity: host, id: [x]}]", + "b: [{entity: host, id: [x, y]}]" + ) + .contains("id columns") + ); // Unknown YAML keys are rejected, catching typos in the embedded file. assert!( parse(&broken( "{src: host, dst: service, rel: uses, direction: down}", + "", "" )) .is_err() diff --git a/src/operator/src/statement/semantic_graph/conventions.yaml b/src/operator/src/statement/semantic_graph/conventions.yaml index 9fc37fe2a02..808fe0f2658 100644 --- a/src/operator/src/statement/semantic_graph/conventions.yaml +++ b/src/operator/src/statement/semantic_graph/conventions.yaml @@ -2,6 +2,10 @@ # versioned with the binary (include_str!), not an operator-editable surface; # non-standard deployments override per table with explicit # greptime.semantic.entity.* declarations, which always win. +# +# Every `id` list is ordered broad to narrow, and that order is part of the +# identity: the id renders as the values joined in this order, so all +# declarations of one entity type must agree on it or they name two entities. co_declared_edges: - { src: service.instance, dst: host, rel: runs_on } @@ -11,6 +15,8 @@ co_declared_edges: - { src: service.instance, dst: service, rel: part_of } - { src: k8s.pod, dst: k8s.workload, rel: part_of } - { src: service.instance, dst: k8s.pod, rel: runs_on } + - { src: service.instance, dst: container, rel: runs_on } + - { src: container, dst: host, rel: runs_on } # Applied only to trace sources; the endpoints are Greptime entity types # derived from GenAI semantic-convention attributes. @@ -30,12 +36,25 @@ virtual_dst_candidates: - { column: span_attributes.server.address, connection_type: virtual_node } # Implicit declarations for greptime_trace_v1 tables; the pod UID bridges -# trace-side pods onto the kube-state-metrics entities below. +# trace-side pods onto the kube-state-metrics entities below. host.name and +# container.name vary by SDK and resource detector, so identity is the id. +# qualified_by folds the namespace into the service identity the same way the +# metric sources' `job` already carries it, so both signals name one entity. otlp_trace_entities: - entity: service id: [service_name] + qualified_by: resource_attributes.service.namespace - entity: service.instance id: [service_name, resource_attributes.service.instance.id] + qualified_by: resource_attributes.service.namespace + - entity: host + id: [resource_attributes.host.id] + descriptive: + - resource_attributes.host.name + - entity: container + id: [resource_attributes.container.id] + descriptive: + - resource_attributes.container.name - entity: k8s.pod id: [resource_attributes.k8s.pod.uid] descriptive: @@ -113,3 +132,32 @@ prometheus_info_metrics: - entity: service.instance id: [job, instance] descriptive_rest: true + +# Implicit declarations for the OTLP resource descriptor synthesized at +# metrics ingestion, gated on source=opentelemetry and metric.type=info. +# job/instance are the derived Prometheus-compatible identity, so services +# line up with target_info. Unlike target_info the attribute set is known, so +# descriptive columns are listed instead of swept up by descriptive_rest, +# which would misfile other entities' ids under service.instance. +otel_info_metrics: + greptime_otel_resource_info: + - entity: service + id: [job] + descriptive: + - service.name + - service.namespace + - entity: service.instance + id: [job, instance] + - entity: host + id: [host.id] + descriptive: + - host.name + - entity: container + id: [container.id] + descriptive: + - container.name + - entity: k8s.pod + id: [k8s.pod.uid] + descriptive: + - k8s.pod.name + - k8s.namespace.name diff --git a/src/operator/src/statement/semantic_graph/relationships.rs b/src/operator/src/statement/semantic_graph/relationships.rs index 3b5c6d37b78..257c844a9a8 100644 --- a/src/operator/src/statement/semantic_graph/relationships.rs +++ b/src/operator/src/statement/semantic_graph/relationships.rs @@ -200,9 +200,9 @@ fn co_declared_branch( bin.clone() + bin_interval(), bin.clone() + bin_interval(), lit(src.entity_type.as_str()), - entity_id_expr(&src.id_columns, &|c| ident(c)), + entity_id_expr(&src.id_columns, src.id_qualifier.as_deref(), &|c| ident(c)), lit(dst.entity_type.as_str()), - entity_id_expr(&dst.id_columns, &|c| ident(c)), + entity_id_expr(&dst.id_columns, dst.id_qualifier.as_deref(), &|c| ident(c)), lit(*rel_type), lit(*provenance), lit(1.0_f64), @@ -501,6 +501,13 @@ fn span_predicate(service: &EntityDeclaration, window: &GraphQueryWindow, strict } /// One trace table's client spans, normalized to +/// A trace table's `duration_nano` is UInt64 on tables created before the +/// signed-integer ingest change and Int64 after it. The per-table selects are +/// unioned, and those two have no common integer type. +fn duration_nano_expr() -> Expr { + cast(ident(DURATION_NANO_COLUMN), DataType::Int64).alias(DURATION_NANO_COLUMN) +} + /// `(timestamp, trace_id, span_id, src_id, status_code, duration_nano, /// virtual_dst, virtual_conn)`. `src_id` is built from the table's `service` /// declaration, so edges land on exactly the entity ids the registry emits (a @@ -553,9 +560,12 @@ fn client_spans( ident(SPAN_ID_COLUMN), // The cast inside entity_id_expr also normalizes tag columns, which // come out of the storage engine dictionary-encoded. - entity_id_expr(&service.id_columns, &|c| ident(c)).alias(SRC_ID_COLUMN), + entity_id_expr(&service.id_columns, service.id_qualifier.as_deref(), &|c| { + ident(c) + }) + .alias(SRC_ID_COLUMN), ident(SPAN_STATUS_CODE_COLUMN), - ident(DURATION_NANO_COLUMN), + duration_nano_expr(), virtual_dst.alias("virtual_dst"), virtual_conn.alias("virtual_conn"), ]) @@ -578,9 +588,12 @@ fn server_spans( ident(TRACE_TIMESTAMP_COLUMN), ident(TRACE_ID_COLUMN), ident(PARENT_SPAN_ID_COLUMN), - entity_id_expr(&service.id_columns, &|c| ident(c)).alias(DST_ID_COLUMN), + entity_id_expr(&service.id_columns, service.id_qualifier.as_deref(), &|c| { + ident(c) + }) + .alias(DST_ID_COLUMN), ident(SPAN_STATUS_CODE_COLUMN), - ident(DURATION_NANO_COLUMN), + duration_nano_expr(), ]) } @@ -609,7 +622,10 @@ fn agent_calls_branch( ident(TRACE_TIMESTAMP_COLUMN), ident(TRACE_ID_COLUMN), ident(SPAN_ID_COLUMN), - entity_id_expr(&agent.id_columns, &|c| ident(c)).alias(SRC_ID_COLUMN), + entity_id_expr(&agent.id_columns, agent.id_qualifier.as_deref(), &|c| { + ident(c) + }) + .alias(SRC_ID_COLUMN), ])?; let child = trace .scan @@ -619,9 +635,12 @@ fn agent_calls_branch( ident(TRACE_TIMESTAMP_COLUMN), ident(TRACE_ID_COLUMN), ident(PARENT_SPAN_ID_COLUMN), - entity_id_expr(&agent.id_columns, &|c| ident(c)).alias(DST_ID_COLUMN), + entity_id_expr(&agent.id_columns, agent.id_qualifier.as_deref(), &|c| { + ident(c) + }) + .alias(DST_ID_COLUMN), ident(SPAN_STATUS_CODE_COLUMN), - ident(DURATION_NANO_COLUMN), + duration_nano_expr(), ])?; parents = union_all(parents, parent)?; children = union_all(children, child)?; @@ -732,6 +751,16 @@ mod tests { name: &str, extra: &[(&str, &[Option<&str>])], spans: &[Span<'_>], + ) { + register_typed_trace_table(ctx, name, extra, spans, DataType::UInt64) + } + + fn register_typed_trace_table( + ctx: &SessionContext, + name: &str, + extra: &[(&str, &[Option<&str>])], + spans: &[Span<'_>], + duration_type: DataType, ) { let mut fields = vec![ Field::new( @@ -745,7 +774,7 @@ mod tests { Field::new(SPAN_KIND_COLUMN, DataType::Utf8, false), Field::new(SPAN_STATUS_CODE_COLUMN, DataType::Utf8, false), Field::new(SERVICE_NAME_COLUMN, DataType::Utf8, false), - Field::new(DURATION_NANO_COLUMN, DataType::UInt64, false), + Field::new(DURATION_NANO_COLUMN, duration_type.clone(), false), ]; for (column, _) in extra { fields.push(Field::new(*column, DataType::Utf8, true)); @@ -774,9 +803,14 @@ mod tests { Arc::new(StringArray::from( spans.iter().map(|s| s.6).collect::>(), )), - Arc::new(UInt64Array::from( - spans.iter().map(|s| s.7).collect::>(), - )), + match duration_type { + DataType::Int64 => Arc::new(Int64Array::from( + spans.iter().map(|s| s.7 as i64).collect::>(), + )) as ArrayRef, + _ => Arc::new(UInt64Array::from( + spans.iter().map(|s| s.7).collect::>(), + )), + }, ]; for (_, values) in extra { columns.push(Arc::new(StringArray::from(values.to_vec()))); @@ -843,6 +877,7 @@ mod tests { time_index: TRACE_TIMESTAMP_COLUMN.to_string(), entity_type: "service".to_string(), id_columns: id_columns.iter().map(|s| s.to_string()).collect(), + id_qualifier: None, descriptive_columns: vec![], scope_columns: vec![], } @@ -972,7 +1007,7 @@ mod tests { ), ], ); - register_trace_table( + register_typed_trace_table( &ctx, "trace_b", &[], @@ -989,6 +1024,7 @@ mod tests { 1_500_000_000, ), ], + DataType::Int64, ); let a = ctx.table("trace_a").await.unwrap(); let b = ctx.table("trace_b").await.unwrap(); @@ -1294,6 +1330,7 @@ mod tests { time_index: "ts".to_string(), entity_type: entity_type.to_string(), id_columns: id_columns.iter().map(|s| s.to_string()).collect(), + id_qualifier: None, descriptive_columns: vec![], scope_columns: vec![], } @@ -1544,16 +1581,10 @@ mod tests { let total: usize = batches.iter().map(|b| b.num_rows()).sum(); assert_eq!(total, 1); let batch = &batches[0]; - // Composite service identity renders the same sorted `k=v` form the - // registry emits, so edges land on registry entity ids. - assert_eq!( - strings(batch, 5), - vec!["service_name=frontend,service_namespace=ns1"] - ); - assert_eq!( - strings(batch, 7), - vec!["service_name=cart,service_namespace=ns1"] - ); + // Composite service identity renders exactly as the registry emits it, + // so edges land on registry entity ids. + assert_eq!(strings(batch, 5), vec!["frontend,ns1"]); + assert_eq!(strings(batch, 7), vec!["cart,ns1"]); } fn build_relationships_plan_for_test( diff --git a/src/servers/src/http/otlp.rs b/src/servers/src/http/otlp.rs index d04022cdf88..3ecaba74a4e 100644 --- a/src/servers/src/http/otlp.rs +++ b/src/servers/src/http/otlp.rs @@ -118,8 +118,9 @@ pub async fn metrics( promote_scope_attrs: http_opts.promote_scope_attrs, with_metric_engine, experimental_enable_exponential_histogram, - // set is_legacy later + // set by the frontend from its config is_legacy: false, + resource_info: false, metric_type: MetricType::Init, metric_translation_strategy: http_opts.metric_translation_strategy, })); @@ -128,7 +129,7 @@ pub async fn metrics( handler.metrics(request, query_ctx).await.map(|outcome| { if outcome.accepted_data_points == 0 && outcome.rejected_data_points > 0 { OtlpMetricsResponse::Failure(outcome) - } else if outcome.rejected_data_points > 0 { + } else if outcome.rejected_data_points > 0 || outcome.error_message.is_some() { OtlpMetricsResponse::PartialSuccess(outcome) } else { OtlpMetricsResponse::FullSuccess(outcome) diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index bd0bfb2ce77..d65e117e07c 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -25,6 +25,7 @@ use common_query::native_histogram::{ }; use common_query::prelude::{GREPTIME_COUNT, greptime_timestamp, greptime_value}; use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; +use common_telemetry::warn; use lazy_static::lazy_static; use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, KeyValue, any_value}; @@ -36,7 +37,7 @@ use table::requests::{ }; use crate::error::{self, Result}; -use crate::otlp::trace::{KEY_SERVICE_INSTANCE_ID, KEY_SERVICE_NAME}; +use crate::otlp::trace::{KEY_SERVICE_INSTANCE_ID, KEY_SERVICE_NAME, KEY_SERVICE_NAMESPACE}; use crate::query_handler::MetricsIngestOutcome; use crate::row_writer::{self, MultiTableData, TableData}; pub use crate::semantic::SemanticIndex; @@ -45,8 +46,11 @@ use crate::semantic::{ METRIC_TYPE_UPDOWN_COUNTER, }; +mod resource_info; mod translator; +pub use resource_info::OTEL_RESOURCE_INFO_TABLE_NAME; +use resource_info::ResourceInfoData; pub use translator::legacy_normalize_otlp_name; pub(crate) use translator::ucum_to_openmetrics_unit; use translator::{translate_label_name, translate_metric_name}; @@ -96,28 +100,43 @@ const MIN_EXPONENTIAL_HISTOGRAM_SCALE: i32 = -4; const MAX_EXPONENTIAL_HISTOGRAM_SCALE: i32 = 8; const MAX_REJECTION_MESSAGE_BYTES: usize = 512; +/// Result of converting one OTLP metrics request. +#[derive(Debug)] +pub struct MetricsConversion { + pub requests: RowInsertRequests, + /// Row count of `requests`; the resource descriptor is not counted. + pub rows: usize, + /// Per-table semantic index for the auto-create path to stamp as table + /// options; covers the descriptor table too. + pub semantic_index: SemanticIndex, + /// The synthesized resource descriptor, written separately from the + /// metric data. See [`resource_info`]. + pub resource_info: Option, + pub outcome: MetricsIngestOutcome, +} + /// Convert OpenTelemetry metrics to GreptimeDB insert requests /// /// See /// /// for data structure of OTLP metrics. -/// -/// Returns `InsertRequests`, total number of rows to ingest, and the per-table -/// semantic index for the auto-create path to stamp as table options. pub fn to_grpc_insert_requests( request: ExportMetricsServiceRequest, metric_ctx: &mut OtlpMetricCtx, -) -> Result<( - RowInsertRequests, - usize, - SemanticIndex, - MetricsIngestOutcome, -)> { +) -> Result { let mut table_writer = MultiTableData::default(); let mut semantic_index = SemanticIndex::default(); let mut outcome = MetricsIngestOutcome::default(); + let mut resource_info = ResourceInfoData::default(); for resource in &request.resource_metrics { + if metric_ctx.resource_info + && !metric_ctx.is_legacy + && let Some(r) = resource.resource.as_ref() + { + resource_info.observe(&r.attributes, resource, metric_ctx); + } + let resource_attrs = resource.resource.as_ref().map(|r| { let mut attrs = r.attributes.clone(); process_resource_attrs(&mut attrs, metric_ctx); @@ -149,8 +168,46 @@ pub fn to_grpc_insert_requests( } let (requests, rows) = table_writer.into_row_insert_requests(); + validate_sample_kinds(&requests)?; - Ok((requests, rows, semantic_index, outcome)) + + // The metric and the descriptor would fight over the same table across + // two engines. The user's metric wins. + let resource_info = if !metric_ctx.resource_info { + None + } else if requests + .inserts + .iter() + .any(|r| r.table_name == OTEL_RESOURCE_INFO_TABLE_NAME) + { + warn!( + "Skipping OTLP resource descriptor synthesis: the request writes \ + a metric named `{OTEL_RESOURCE_INFO_TABLE_NAME}`" + ); + None + } else { + resource_info.into_row_insert_requests()? + }; + if resource_info.is_some() { + semantic_index.record_scalar( + OTEL_RESOURCE_INFO_TABLE_NAME, + SEMANTIC_METRIC_TYPE, + crate::semantic::METRIC_TYPE_INFO, + ); + semantic_index.record_scalar( + OTEL_RESOURCE_INFO_TABLE_NAME, + SEMANTIC_METRIC_METADATA_QUALITY, + METADATA_QUALITY_DECLARED, + ); + } + + Ok(MetricsConversion { + requests, + rows, + semantic_index, + outcome, + resource_info, + }) } fn validate_sample_kinds(requests: &RowInsertRequests) -> Result<()> { @@ -287,30 +344,61 @@ fn from_metric_type(data: &metric::Data) -> MetricType { } } +/// Non-scalar values (bool, arrays, maps, bytes) are not representable as tags. +fn scalar_value_string(value: Option<&AnyValue>) -> Option { + match value.and_then(|v| v.value.as_ref())? { + any_value::Value::StringValue(s) => Some(s.clone()), + any_value::Value::IntValue(v) => Some(v.to_string()), + any_value::Value::DoubleValue(v) => Some(v.to_string()), + _ => None, + } +} + +/// Prometheus-style `(job, instance)` identity. Per the OTel Prometheus +/// compatibility spec `job` folds in `service.namespace` when present; a +/// resource without `service.name` gets no job rather than a fabricated one. +/// Both fields are optional strings, so a named type keeps a swapped +/// destructuring from type-checking at the call sites. +pub(crate) struct ServiceIdentity { + pub job: Option, + pub instance: Option, +} + +pub(crate) fn service_identity(attrs: &[KeyValue]) -> ServiceIdentity { + let mut name = None; + let mut namespace = None; + let mut instance = None; + for kv in attrs { + match kv.key.as_str() { + KEY_SERVICE_NAME => name = scalar_value_string(kv.value.as_ref()), + KEY_SERVICE_NAMESPACE => namespace = scalar_value_string(kv.value.as_ref()), + KEY_SERVICE_INSTANCE_ID => instance = scalar_value_string(kv.value.as_ref()), + _ => {} + } + } + let job = name.map(|name| match namespace { + Some(ns) if !ns.is_empty() => format!("{ns}/{name}"), + _ => name, + }); + ServiceIdentity { job, instance } +} + +fn string_key_value(key: &str, value: String) -> KeyValue { + KeyValue { + key: key.to_string(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(value)), + }), + } +} + fn process_resource_attrs(attrs: &mut Vec, metric_ctx: &OtlpMetricCtx) { if metric_ctx.is_legacy { return; } - // remap service.name and service.instance.id to job and instance - let mut tmp = Vec::with_capacity(2); - for kv in attrs.iter() { - match &kv.key as &str { - KEY_SERVICE_NAME => { - tmp.push(KeyValue { - key: JOB_KEY.to_string(), - value: kv.value.clone(), - }); - } - KEY_SERVICE_INSTANCE_ID => { - tmp.push(KeyValue { - key: INSTANCE_KEY.to_string(), - value: kv.value.clone(), - }); - } - _ => {} - } - } + // remap the service identity attributes to job and instance + let ServiceIdentity { job, instance } = service_identity(attrs); // if promote all, then exclude the list, else, include the list if metric_ctx.promote_all_resource_attrs { @@ -322,7 +410,12 @@ fn process_resource_attrs(attrs: &mut Vec, metric_ctx: &OtlpMetricCtx) }); } - attrs.extend(tmp); + if let Some(job) = job { + attrs.push(string_key_value(JOB_KEY, job)); + } + if let Some(instance) = instance { + attrs.push(string_key_value(INSTANCE_KEY, instance)); + } } fn process_scope_attrs(scope: &ScopeMetrics, metric_ctx: &OtlpMetricCtx) -> Option> { @@ -536,35 +629,13 @@ fn encode_exponential_histogram( metric_ctx: &OtlpMetricCtx, outcome: &mut MetricsIngestOutcome, ) -> Result { - if !metric_ctx.experimental_enable_exponential_histogram { + if let Err(rejection) = exponential_histogram_gate(histogram, metric_ctx) { reject_data_points(outcome, histogram.data_points.len(), || { - format!( - "metric `{name}` uses OTLP exponential histograms; set otlp.experimental_enable_exponential_histogram = true to enable ingestion" - ) + rejection.message(name) })?; return Ok(false); } - match AggregationTemporality::try_from(histogram.aggregation_temporality) { - Ok(AggregationTemporality::Cumulative) => {} - Ok(AggregationTemporality::Delta) => { - reject_data_points(outcome, histogram.data_points.len(), || { - format!( - "metric `{name}` uses delta OTLP exponential histograms; only cumulative temporality is supported" - ) - })?; - return Ok(false); - } - _ => { - reject_data_points(outcome, histogram.data_points.len(), || { - format!( - "metric `{name}` has unspecified OTLP exponential histogram temporality; cumulative temporality is required" - ) - })?; - return Ok(false); - } - } - let column_schema = native_histogram_column_schema().map_err(|error| { error::InternalSnafu { err_msg: error.to_string(), @@ -611,7 +682,46 @@ fn encode_exponential_histogram( Ok(emitted) } -fn exponential_histogram_value( +pub(crate) enum ExponentialHistogramRejection { + Disabled, + DeltaTemporality, + UnspecifiedTemporality, +} + +impl ExponentialHistogramRejection { + fn message(&self, name: &str) -> String { + match self { + Self::Disabled => format!( + "metric `{name}` uses OTLP exponential histograms; set otlp.experimental_enable_exponential_histogram = true to enable ingestion" + ), + Self::DeltaTemporality => format!( + "metric `{name}` uses delta OTLP exponential histograms; only cumulative temporality is supported" + ), + Self::UnspecifiedTemporality => format!( + "metric `{name}` has unspecified OTLP exponential histogram temporality; cumulative temporality is required" + ), + } + } +} + +/// Whole-metric acceptance, decided once for the encoder and for the resource +/// descriptor, which must not describe a resource whose data was rejected. +/// Individual points can still fail [`exponential_histogram_value`]. +pub(crate) fn exponential_histogram_gate( + histogram: &ExponentialHistogram, + metric_ctx: &OtlpMetricCtx, +) -> std::result::Result<(), ExponentialHistogramRejection> { + if !metric_ctx.experimental_enable_exponential_histogram { + return Err(ExponentialHistogramRejection::Disabled); + } + match AggregationTemporality::try_from(histogram.aggregation_temporality) { + Ok(AggregationTemporality::Cumulative) => Ok(()), + Ok(AggregationTemporality::Delta) => Err(ExponentialHistogramRejection::DeltaTemporality), + _ => Err(ExponentialHistogramRejection::UnspecifiedTemporality), + } +} + +pub(crate) fn exponential_histogram_value( data_point: &ExponentialHistogramDataPoint, ) -> std::result::Result<(ValueData, i64), String> { if data_point.start_time_unix_nano > data_point.time_unix_nano { @@ -823,29 +933,21 @@ fn write_attributes( }; let tags = attrs.iter().filter_map(|attr| { - attr.value - .as_ref() - .and_then(|v| v.value.as_ref()) - .and_then(|val| { - let key = match attribute_type { - AttributeType::Resource | AttributeType::DataPoint => { - translate_label_name(&attr.key, metric_ctx.metric_translation_strategy) - } - AttributeType::Scope => { - format!( - "otel_scope_{}", - translate_label_name(&attr.key, metric_ctx.metric_translation_strategy) - ) - } - AttributeType::Legacy => legacy_normalize_otlp_name(&attr.key), - }; - match val { - any_value::Value::StringValue(s) => Some((key, s.clone())), - any_value::Value::IntValue(v) => Some((key, v.to_string())), - any_value::Value::DoubleValue(v) => Some((key, v.to_string())), - _ => None, // TODO(sunng87): allow different type of values - } - }) + // TODO(sunng87): allow different type of values + let value = scalar_value_string(attr.value.as_ref())?; + let key = match attribute_type { + AttributeType::Resource | AttributeType::DataPoint => { + translate_label_name(&attr.key, metric_ctx.metric_translation_strategy) + } + AttributeType::Scope => { + format!( + "otel_scope_{}", + translate_label_name(&attr.key, metric_ctx.metric_translation_strategy) + ) + } + AttributeType::Legacy => legacy_normalize_otlp_name(&attr.key), + }; + Some((key, value)) }); row_writer::write_tags(writer, tags, row)?; @@ -1295,6 +1397,169 @@ mod tests { } } + fn descriptor_ctx() -> OtlpMetricCtx { + OtlpMetricCtx { + resource_info: true, + ..Default::default() + } + } + + fn attr_value(attrs: &[KeyValue], key: &str) -> Option { + attrs + .iter() + .find(|kv| kv.key == key) + .and_then(|kv| scalar_value_string(kv.value.as_ref())) + } + + fn gauge_request( + resource_attrs: Vec, + metric_name: &str, + ) -> ExportMetricsServiceRequest { + use otel_arrow_rust::proto::opentelemetry::resource::v1::Resource; + ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + resource: Some(Resource { + attributes: resource_attrs, + ..Default::default() + }), + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + name: metric_name.to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint { + time_unix_nano: 1_000_000, + value: Some(Value::AsInt(1)), + ..Default::default() + }], + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }], + } + } + + fn column_names(request: &RowInsertRequests, table: &str) -> Vec { + request + .inserts + .iter() + .find(|r| r.table_name == table) + .unwrap_or_else(|| panic!("missing table {table}")) + .rows + .as_ref() + .unwrap() + .schema + .iter() + .map(|c| c.column_name.clone()) + .collect() + } + + #[test] + fn test_conversion_synthesizes_resource_descriptor() { + set_default_prefix(None).unwrap(); + let request = gauge_request( + vec![keyvalue("service.name", "api"), keyvalue("host.id", "h-1")], + "my_gauge", + ); + let conversion = to_grpc_insert_requests(request, &mut descriptor_ctx()).unwrap(); + + // descriptor keeps raw OTel keys while the metric table's labels went + // through the (default underscore-escaping) translation strategy + let resource_info = conversion.resource_info.expect("descriptor synthesized"); + let descriptor_cols = column_names(&resource_info, OTEL_RESOURCE_INFO_TABLE_NAME); + assert!(descriptor_cols.contains(&"host.id".to_string())); + assert!(descriptor_cols.contains(&"service.name".to_string())); + assert!(descriptor_cols.contains(&"job".to_string())); + let metric_cols = column_names(&conversion.requests, "my_gauge"); + assert!(metric_cols.contains(&"service_name".to_string())); + assert!(!metric_cols.contains(&"service.name".to_string())); + // host.id is not in the promote list: only the descriptor keeps it + assert!(!metric_cols.contains(&"host_id".to_string())); + + let decoded = decode(&conversion.semantic_index); + let t = &decoded[OTEL_RESOURCE_INFO_TABLE_NAME]; + assert_eq!( + t.get(SEMANTIC_METRIC_TYPE).map(String::as_str), + Some("info") + ); + assert_eq!( + t.get(SEMANTIC_METRIC_METADATA_QUALITY).map(String::as_str), + Some("declared") + ); + } + + #[test] + fn test_conversion_skips_descriptor_for_legacy_mode() { + set_default_prefix(None).unwrap(); + let request = gauge_request( + vec![keyvalue("service.name", "api"), keyvalue("host.id", "h-1")], + "my_gauge", + ); + let mut ctx = OtlpMetricCtx { + is_legacy: true, + ..descriptor_ctx() + }; + let conversion = to_grpc_insert_requests(request, &mut ctx).unwrap(); + assert!(conversion.resource_info.is_none()); + + // legacy tables predate job/instance and the promote filter; adding + // either would alter the schema of tables already in use + let cols = column_names(&conversion.requests, "my_gauge"); + assert!(!cols.contains(&"job".to_string())); + assert!(cols.contains(&"host_id".to_string())); + } + + #[test] + fn test_conversion_skips_descriptor_on_metric_name_collision() { + set_default_prefix(None).unwrap(); + let request = gauge_request( + vec![keyvalue("service.name", "api")], + OTEL_RESOURCE_INFO_TABLE_NAME, + ); + let conversion = to_grpc_insert_requests(request, &mut descriptor_ctx()).unwrap(); + assert!(conversion.resource_info.is_none()); + // the metric itself still goes through the main path + assert!( + conversion + .requests + .inserts + .iter() + .any(|r| r.table_name == OTEL_RESOURCE_INFO_TABLE_NAME) + ); + } + + #[test] + fn test_job_composition_follows_service_namespace() { + let mut attrs = vec![ + keyvalue("service.name", "api"), + keyvalue("service.namespace", "shop"), + ]; + process_resource_attrs(&mut attrs, &OtlpMetricCtx::default()); + assert_eq!(attr_value(&attrs, "job").as_deref(), Some("shop/api")); + + let mut attrs = vec![keyvalue("service.name", "api")]; + process_resource_attrs(&mut attrs, &OtlpMetricCtx::default()); + assert_eq!(attr_value(&attrs, "job").as_deref(), Some("api")); + + let mut attrs = vec![ + keyvalue("service.name", "api"), + keyvalue("service.namespace", ""), + ]; + process_resource_attrs(&mut attrs, &OtlpMetricCtx::default()); + assert_eq!(attr_value(&attrs, "job").as_deref(), Some("api")); + + // no service.name: no job is fabricated + let mut attrs = vec![ + keyvalue("service.namespace", "shop"), + keyvalue("service.instance.id", "inst-1"), + ]; + process_resource_attrs(&mut attrs, &OtlpMetricCtx::default()); + assert_eq!(attr_value(&attrs, "job"), None); + assert_eq!(attr_value(&attrs, "instance").as_deref(), Some("inst-1")); + } + #[test] fn test_encode_gauge() { let mut tables = MultiTableData::default(); @@ -2104,8 +2369,12 @@ mod tests { AggregationTemporality::Cumulative, ), ]); - let (requests, _, semantic_index, outcome) = - to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + let MetricsConversion { + requests, + semantic_index, + outcome, + .. + } = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); assert_eq!(outcome.accepted_data_points, 1); assert_eq!(outcome.rejected_data_points, 1); @@ -2127,8 +2396,9 @@ mod tests { vec![], AggregationTemporality::Cumulative, )]); - let (_, _, _, outcome) = - to_grpc_insert_requests(empty, &mut OtlpMetricCtx::default()).unwrap(); + let outcome = to_grpc_insert_requests(empty, &mut OtlpMetricCtx::default()) + .unwrap() + .outcome; assert_eq!(outcome.rejected_data_points, 0); assert_eq!(outcome.error_message, None); } @@ -2208,8 +2478,12 @@ mod tests { .collect(), }; - let (requests, rows, _, outcome) = - to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + let MetricsConversion { + requests, + rows, + outcome, + .. + } = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); assert_eq!(outcome.accepted_data_points, 2); assert_eq!(rows, 6); @@ -2243,8 +2517,13 @@ mod tests { experimental_enable_exponential_histogram: true, ..Default::default() }; - let (requests, rows, semantic_index, outcome) = - to_grpc_insert_requests(request, &mut ctx).unwrap(); + let MetricsConversion { + requests, + rows, + semantic_index, + outcome, + .. + } = to_grpc_insert_requests(request, &mut ctx).unwrap(); assert_eq!(outcome.accepted_data_points, 0); assert_eq!(outcome.rejected_data_points, 1); @@ -2269,14 +2548,17 @@ mod tests { experimental_enable_exponential_histogram: true, ..Default::default() }; - let (new_requests, _, _, _) = - to_grpc_insert_requests(request.clone(), &mut new_ctx).unwrap(); + let new_requests = to_grpc_insert_requests(request.clone(), &mut new_ctx) + .unwrap() + .requests; let mut legacy_ctx = OtlpMetricCtx { experimental_enable_exponential_histogram: true, is_legacy: true, ..Default::default() }; - let (legacy_requests, _, _, _) = to_grpc_insert_requests(request, &mut legacy_ctx).unwrap(); + let legacy_requests = to_grpc_insert_requests(request, &mut legacy_ctx) + .unwrap() + .requests; let new_insert = &new_requests.inserts[0]; let legacy_insert = &legacy_requests.inserts[0]; @@ -2325,8 +2607,9 @@ mod tests { vec![ExponentialHistogramDataPoint::default()], AggregationTemporality::Cumulative, )]); - let (_, _, _, outcome) = - to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + let outcome = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()) + .unwrap() + .outcome; assert_eq!(outcome.rejected_data_points, 1); assert!(outcome.error_message.unwrap().len() <= MAX_REJECTION_MESSAGE_BYTES); diff --git a/src/servers/src/otlp/metrics/resource_info.rs b/src/servers/src/otlp/metrics/resource_info.rs new file mode 100644 index 00000000000..cd3427f7163 --- /dev/null +++ b/src/servers/src/otlp/metrics/resource_info.rs @@ -0,0 +1,393 @@ +// 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. + +//! Resource descriptor synthesized from OTLP metrics requests, so the entity +//! graph reads one table instead of scanning every logical metric table for +//! attributes the promote filter may have dropped. +//! +//! Columns are a fixed allowlist under the raw OTel attribute names: the +//! conventions whitelist matches fixed names, so they must not follow the +//! per-request label translation strategy or the promote/ignore headers. + +use std::collections::BTreeMap; + +use api::v1::RowInsertRequests; +use common_catalog::consts::SEMANTIC_GRAPH_WINDOW_NANOS; +use common_grpc::precision::Precision; +use common_query::prelude::{greptime_timestamp, greptime_value}; +use otel_arrow_rust::proto::opentelemetry::common::v1::KeyValue; +use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ResourceMetrics, metric}; +use session::protocol_ctx::OtlpMetricCtx; + +use crate::error::Result; +use crate::otlp::metrics::{ + INSTANCE_KEY, JOB_KEY, ServiceIdentity, exponential_histogram_gate, + exponential_histogram_value, scalar_value_string, service_identity, +}; +use crate::otlp::trace::{ + KEY_CONTAINER_ID, KEY_CONTAINER_NAME, KEY_HOST_ID, KEY_HOST_NAME, KEY_K8S_NAMESPACE_NAME, + KEY_K8S_POD_NAME, KEY_K8S_POD_UID, KEY_SERVICE_NAME, KEY_SERVICE_NAMESPACE, +}; +use crate::row_writer::{self, MultiTableData}; + +/// Prefixed like the other engine-managed tables, so a user metric is +/// unlikely to claim the name. +pub const OTEL_RESOURCE_INFO_TABLE_NAME: &str = "greptime_otel_resource_info"; + +/// Attributes projected under their raw OTel keys. `service.instance.id` is +/// absent on purpose: it lands in `instance`. +/// +/// Matched instead of scanned: this runs for every attribute of every +/// resource, and the compiler turns it into a length-and-prefix dispatch. +fn is_projected_attr(key: &str) -> bool { + matches!( + key, + KEY_SERVICE_NAME + | KEY_SERVICE_NAMESPACE + | KEY_HOST_ID + | KEY_HOST_NAME + | KEY_CONTAINER_ID + | KEY_CONTAINER_NAME + | KEY_K8S_POD_UID + | KEY_K8S_POD_NAME + | KEY_K8S_NAMESPACE_NAME + ) +} + +/// Upper bound of [`is_projected_attr`] plus the derived `job`/`instance`, +/// used to size the per-row buffers. +const MAX_PROJECTED_TAGS: usize = 11; + +/// Projected attributes (sorted `(name, value)` pairs) -> graph window -> +/// the newest data-point time seen in that window, which is what the row for +/// that window is stamped with. +/// +/// Windows are an inner map so the attributes are stored, and moved, once per +/// resource. They are keyed separately because one request may carry data for +/// several of them: a single row per resource would describe only the newest +/// window, leaving the earlier ones with metric rows but no entities. +#[derive(Debug, Default)] +pub struct ResourceInfoData { + rows: BTreeMap, BTreeMap>, +} + +impl ResourceInfoData { + /// Takes the raw attributes, before the promote filter runs on them. + pub fn observe( + &mut self, + raw_attrs: &[KeyValue], + resource: &ResourceMetrics, + metric_ctx: &OtlpMetricCtx, + ) { + let mut tags = Vec::with_capacity(MAX_PROJECTED_TAGS); + let ServiceIdentity { job, instance } = service_identity(raw_attrs); + if let Some(job) = job { + tags.push((JOB_KEY.to_string(), job)); + } + if let Some(instance) = instance { + tags.push((INSTANCE_KEY.to_string(), instance)); + } + for kv in raw_attrs { + if is_projected_attr(&kv.key) + && let Some(value) = scalar_value_string(kv.value.as_ref()) + { + tags.push((kv.key.clone(), value)); + } + } + if tags.is_empty() { + return; + } + // Sorted so equal attribute sets share a key, and so the emitted + // columns keep a stable order. + tags.sort_unstable(); + + let mut observed: BTreeMap = BTreeMap::new(); + for_each_encoded_time(resource, metric_ctx, |ts| { + let window = ts - ts.rem_euclid(SEMANTIC_GRAPH_WINDOW_NANOS); + observed + .entry(window) + .and_modify(|newest| *newest = (*newest).max(ts)) + .or_insert(ts); + }); + if observed.is_empty() { + return; + } + + let windows = self.rows.entry(tags).or_default(); + for (window, newest) in observed { + windows + .entry(window) + .and_modify(|seen| *seen = (*seen).max(newest)) + .or_insert(newest); + } + } + + /// Every projected attribute becomes a tag, so auto-create puts it in the + /// primary key where the conventions expect it. + pub fn into_row_insert_requests(self) -> Result> { + if self.rows.is_empty() { + return Ok(None); + } + + let mut writer = MultiTableData::default(); + let table = writer.get_or_default_table_data( + OTEL_RESOURCE_INFO_TABLE_NAME, + MAX_PROJECTED_TAGS + 2, + self.rows.values().map(BTreeMap::len).sum(), + ); + for (tags, windows) in &self.rows { + for ts_nanos in windows.values().copied() { + let mut row = table.alloc_one_row(); + row_writer::write_tags(table, tags.iter().cloned(), &mut row)?; + row_writer::write_f64(table, greptime_value(), 1.0, &mut row)?; + row_writer::write_ts_to_millis( + table, + greptime_timestamp(), + Some(ts_nanos), + Precision::Nanosecond, + &mut row, + )?; + table.add_row(row); + } + } + + let (requests, _) = writer.into_row_insert_requests(); + Ok(Some(requests)) + } +} + +/// Visits the times of the data points the encoder writes rows for, so a +/// resource is described exactly where it is measured rather than wherever +/// its request happens to reach. +fn for_each_encoded_time( + resource: &ResourceMetrics, + metric_ctx: &OtlpMetricCtx, + mut visit: impl FnMut(i64), +) { + fn visit_all(points: impl Iterator, visit: &mut impl FnMut(i64)) { + for ts in points { + visit(ts as i64); + } + } + for scope in &resource.scope_metrics { + for m in &scope.metrics { + match &m.data { + Some(metric::Data::Gauge(g)) => { + visit_all(g.data_points.iter().map(|p| p.time_unix_nano), &mut visit) + } + Some(metric::Data::Sum(s)) => { + visit_all(s.data_points.iter().map(|p| p.time_unix_nano), &mut visit) + } + Some(metric::Data::Histogram(h)) => { + visit_all(h.data_points.iter().map(|p| p.time_unix_nano), &mut visit) + } + Some(metric::Data::Summary(s)) => { + visit_all(s.data_points.iter().map(|p| p.time_unix_nano), &mut visit) + } + Some(metric::Data::ExponentialHistogram(h)) + if exponential_histogram_gate(h, metric_ctx).is_ok() => + { + for point in &h.data_points { + if let Ok((_, ts)) = exponential_histogram_value(point) { + visit(ts); + } + } + } + Some(metric::Data::ExponentialHistogram(_)) => {} + None => {} + } + } + } +} + +#[cfg(test)] +mod tests { + use api::v1::SemanticType; + use api::v1::value::ValueData; + use common_query::prelude::set_default_prefix; + use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, any_value}; + use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ + AggregationTemporality, ExponentialHistogram, ExponentialHistogramDataPoint, Gauge, Metric, + NumberDataPoint, ScopeMetrics, + }; + + use super::*; + + fn kv(key: &str, value: &str) -> KeyValue { + KeyValue { + key: key.into(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(value.into())), + }), + } + } + + fn gauge_at(times: &[i64]) -> ResourceMetrics { + ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + data: Some(metric::Data::Gauge(Gauge { + data_points: times + .iter() + .map(|ts| NumberDataPoint { + time_unix_nano: *ts as u64, + ..Default::default() + }) + .collect(), + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + } + } + + #[test] + fn observe_projects_allowlist_and_dedups_per_request() { + let mut data = ResourceInfoData::default(); + let attrs = vec![ + kv("service.name", "api"), + kv("service.namespace", "shop"), + kv("service.instance.id", "inst-1"), + kv("host.id", "h-1"), + kv("os.type", "linux"), + ]; + data.observe(&attrs, &gauge_at(&[100, 50]), &OtlpMetricCtx::default()); + assert_eq!(data.rows.len(), 1); + let (tags, windows) = data.rows.iter().next().unwrap(); + assert_eq!(windows.values().copied().collect::>(), vec![100]); + assert!(tags.contains(&("job".to_string(), "shop/api".to_string()))); + assert!(tags.contains(&("instance".to_string(), "inst-1".to_string()))); + assert!(tags.contains(&("service.name".to_string(), "api".to_string()))); + assert!( + tags.iter() + .all(|(k, _)| k != "os.type" && k != "service.instance.id") + ); + + data.observe( + &[kv("host.id", "h-2")], + &gauge_at(&[10]), + &OtlpMetricCtx::default(), + ); + assert_eq!(data.rows.len(), 2); + + let mut empty = ResourceInfoData::default(); + empty.observe( + &[kv("os.type", "linux")], + &gauge_at(&[100]), + &OtlpMetricCtx::default(), + ); + assert!(empty.into_row_insert_requests().unwrap().is_none()); + } + + /// Earlier windows would keep their metric rows but lose their entities. + #[test] + fn observe_keeps_one_row_per_graph_window() { + let window = SEMANTIC_GRAPH_WINDOW_NANOS; + let mut data = ResourceInfoData::default(); + data.observe( + &[kv("service.name", "api")], + &gauge_at(&[window + 1, window + 2, 3 * window + 7]), + &OtlpMetricCtx::default(), + ); + + let windows = data.rows.values().next().unwrap(); + assert_eq!( + windows.iter().collect::>(), + vec![(&window, &(window + 2)), (&(3 * window), &(3 * window + 7))] + ); + } + + /// Describing a resource whose only data the encoder drops invents an + /// entity with no measurements. + #[test] + fn observe_ignores_data_the_encoder_drops() { + let exponential = |temporality: AggregationTemporality| ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + data: Some(metric::Data::ExponentialHistogram(ExponentialHistogram { + data_points: vec![ExponentialHistogramDataPoint { + time_unix_nano: 100, + ..Default::default() + }], + aggregation_temporality: temporality as i32, + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }; + let enabled = OtlpMetricCtx { + experimental_enable_exponential_histogram: true, + ..Default::default() + }; + + for (resource, ctx) in [ + ( + exponential(AggregationTemporality::Cumulative), + OtlpMetricCtx::default(), + ), + (exponential(AggregationTemporality::Delta), enabled), + ] { + let mut data = ResourceInfoData::default(); + data.observe(&[kv("service.name", "api")], &resource, &ctx); + assert!(data.into_row_insert_requests().unwrap().is_none()); + } + } + + #[test] + fn rows_carry_raw_key_tags_value_and_millis_timestamp() { + set_default_prefix(None).unwrap(); + let mut data = ResourceInfoData::default(); + data.observe( + &[kv("service.name", "api"), kv("host.id", "h-1")], + &gauge_at(&[1_700_000_000_123_456_789]), + &OtlpMetricCtx::default(), + ); + let requests = data.into_row_insert_requests().unwrap().unwrap(); + assert_eq!(requests.inserts.len(), 1); + let insert = &requests.inserts[0]; + assert_eq!(insert.table_name, OTEL_RESOURCE_INFO_TABLE_NAME); + + let rows = insert.rows.as_ref().unwrap(); + let names = rows + .schema + .iter() + .map(|c| c.column_name.as_str()) + .collect::>(); + assert_eq!( + names, + vec![ + "host.id", + "job", + "service.name", + greptime_value(), + greptime_timestamp() + ] + ); + for column in &rows.schema[..3] { + assert_eq!(column.semantic_type, SemanticType::Tag as i32); + } + + assert_eq!(rows.rows.len(), 1); + let values = &rows.rows[0].values; + assert_eq!(values[3].value_data, Some(ValueData::F64Value(1.0))); + assert_eq!( + values[4].value_data, + Some(ValueData::TimestampMillisecondValue(1_700_000_000_123)) + ); + } +} diff --git a/src/servers/src/otlp/trace.rs b/src/servers/src/otlp/trace.rs index a8439cdb2c4..9ee16507bda 100644 --- a/src/servers/src/otlp/trace.rs +++ b/src/servers/src/otlp/trace.rs @@ -46,7 +46,15 @@ pub const TRACE_STATE_COLUMN: &str = "trace_state"; // const keys pub const KEY_SERVICE_NAME: &str = "service.name"; +pub const KEY_SERVICE_NAMESPACE: &str = "service.namespace"; pub const KEY_SERVICE_INSTANCE_ID: &str = "service.instance.id"; +pub const KEY_HOST_ID: &str = "host.id"; +pub const KEY_HOST_NAME: &str = "host.name"; +pub const KEY_CONTAINER_ID: &str = "container.id"; +pub const KEY_CONTAINER_NAME: &str = "container.name"; +pub const KEY_K8S_POD_UID: &str = "k8s.pod.uid"; +pub const KEY_K8S_POD_NAME: &str = "k8s.pod.name"; +pub const KEY_K8S_NAMESPACE_NAME: &str = "k8s.namespace.name"; pub const KEY_SPAN_KIND: &str = "span.kind"; // jaeger const keys, not sure if they are general diff --git a/src/session/src/protocol_ctx.rs b/src/session/src/protocol_ctx.rs index 537dec5c414..d9a4a386f91 100644 --- a/src/session/src/protocol_ctx.rs +++ b/src/session/src/protocol_ctx.rs @@ -56,6 +56,9 @@ pub struct OtlpMetricCtx { pub with_metric_engine: bool, pub experimental_enable_exponential_histogram: bool, pub is_legacy: bool, + /// Set from the server's `otlp.experimental_enable_resource_info`; off + /// means the resource descriptor is not synthesized at all. + pub resource_info: bool, pub metric_type: MetricType, pub metric_translation_strategy: OtlpMetricTranslationStrategy, } diff --git a/src/table/src/requests/semantic.rs b/src/table/src/requests/semantic.rs index 2daa440c11d..471be8b891c 100644 --- a/src/table/src/requests/semantic.rs +++ b/src/table/src/requests/semantic.rs @@ -98,6 +98,10 @@ pub const SEMANTIC_ENTITY_SERVICE_ID: &str = "greptime.semantic.entity.service.i /// ([`has_stable_string_form`]). #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum EntityRole { + /// Identifying columns. Their **order is part of the identity**: the + /// entity id is their values joined in the declared order, so two tables + /// naming the same entity must list them the same way (broad to narrow) + /// or they name two entities. Id, Descriptive, Scope, diff --git a/tests-integration/src/instance.rs b/tests-integration/src/instance.rs index 08d2a1aff15..9f48d01a467 100644 --- a/tests-integration/src/instance.rs +++ b/tests-integration/src/instance.rs @@ -691,4 +691,119 @@ mod tests { "denied source leaked into entities:\n{pretty_print}" ); } + + /// Denying the descriptor's table must not fail the request: the metric + /// data is already committed, so the client would retry what the server + /// accepted. + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_descriptor_permission_denial_degrades() { + use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; + use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, KeyValue, any_value}; + use otel_arrow_rust::proto::opentelemetry::metrics::v1::number_data_point::Value; + use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ + Gauge, Metric, NumberDataPoint, ResourceMetrics, ScopeMetrics, metric, + }; + use otel_arrow_rust::proto::opentelemetry::resource::v1::Resource; + use servers::query_handler::OpenTelemetryProtocolHandler; + + common_telemetry::init_default_ut_logging(); + + /// Denies writes to the descriptor table only. + struct DenyResourceInfo; + + impl auth::PermissionChecker for DenyResourceInfo { + fn check_permission( + &self, + _user_info: auth::UserInfoRef, + _req: auth::PermissionReq, + ) -> auth::error::Result { + Ok(auth::PermissionResp::Allow) + } + + fn check_permission_with_table_targets( + &self, + _user_info: auth::UserInfoRef, + _req: auth::PermissionReq, + targets: auth::PermissionTableTargets, + ) -> auth::error::Result { + if let auth::PermissionTableTargets::Resolved(targets) = targets + && targets + .iter() + .any(|target| target.table == "greptime_otel_resource_info") + { + return Ok(auth::PermissionResp::Reject); + } + Ok(auth::PermissionResp::Allow) + } + } + + let plugins = Plugins::new(); + plugins.insert::(Arc::new(DenyResourceInfo)); + + let standalone = GreptimeDbStandaloneBuilder::new("test_otlp_descriptor_denied") + .with_plugin(plugins) + .build() + .await; + let instance = standalone.fe_instance().clone(); + + let attr = |key: &str, value: &str| KeyValue { + key: key.to_string(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(value.to_string())), + }), + }; + let request = ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + resource: Some(Resource { + attributes: vec![attr("service.name", "api"), attr("host.id", "h-1")], + ..Default::default() + }), + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + name: "denied_descriptor_gauge".to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint { + time_unix_nano: 1_700_000_000_000_000_000, + value: Some(Value::AsDouble(1.0)), + ..Default::default() + }], + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }], + }; + + let outcome = + OpenTelemetryProtocolHandler::metrics(&*instance, request, QueryContext::arc()) + .await + .unwrap(); + let warning = outcome + .error_message + .expect("denial must surface as a warning"); + assert!( + warning.contains("greptime_otel_resource_info"), + "unexpected warning: {warning}" + ); + + // the metric data is committed, the descriptor table was never created + let sql = "select table_name from information_schema.tables \ + where table_schema = 'public' order by table_name"; + let output = query(&instance, sql).await; + let OutputData::Stream(s) = output.data else { + unreachable!() + }; + let batches = common_recordbatch::util::collect_batches(s).await.unwrap(); + let pretty_print = batches.pretty_print().unwrap(); + assert!( + pretty_print.contains("denied_descriptor_gauge"), + "metric data was not committed:\n{pretty_print}" + ); + assert!( + !pretty_print.contains("greptime_otel_resource_info"), + "denied descriptor table was created:\n{pretty_print}" + ); + } } diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 1085c21983e..16fc2e65634 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -369,6 +369,11 @@ impl GreptimeDbStandaloneBuilder { grpc: GrpcOptions::default().with_server_addr("127.0.0.1:4001"), slow_query: self.slow_query_options.clone(), auto_create_table: self.auto_create_table, + // Tests cover the descriptor, so they run with it enabled. + otlp: frontend::service_config::OtlpOptions { + experimental_enable_resource_info: true, + ..Default::default() + }, ..StandaloneOptions::default() }; diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 7a787833376..4a8f96260e8 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -164,6 +164,7 @@ macro_rules! http_tests { test_otlp_metrics_new, test_otlp_exponential_histogram, test_otlp_metric_translation_strategies, + test_otlp_metrics_resource_info_conflicts, test_otlp_traces_v0, test_otlp_traces_v1, test_otlp_traces_v1_entity_graph, @@ -2208,6 +2209,7 @@ enable = true enable = true experimental_enable_exponential_histogram = false trace_ingest_chunk_size = 512 +experimental_enable_resource_info = true [prom_store] enable = true @@ -6500,7 +6502,7 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { .await; assert_eq!(StatusCode::OK, res.status()); - let expected = "[[\"claude_code_cost_usage_USD_total\"],[\"claude_code_token_usage_tokens_total\"],[\"demo\"],[\"greptime_physical_table\"],[\"numbers\"]]"; + let expected = "[[\"claude_code_cost_usage_USD_total\"],[\"claude_code_token_usage_tokens_total\"],[\"demo\"],[\"greptime_otel_resource_info\"],[\"greptime_physical_table\"],[\"numbers\"]]"; validate_data("otlp_metrics_all_tables", &client, "show tables;", expected).await; // Metric-engine logical table carries the semantic identity. Match substrings @@ -6567,6 +6569,33 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { ) .await; + // The synthesized resource descriptor: a plain mito info table keyed by + // the raw OTel attribute keys (service.name is the only allowlisted + // resource attribute in this payload), independent of the promote/ignore + // headers. The timestamp is the newest data-point time. + validate_data( + "otlp_metrics_resource_info", + &client, + "select job, \"service.name\", greptime_value, greptime_timestamp from greptime_otel_resource_info;", + "[[\"claude-code\",\"claude-code\",1.0,1753780559836]]", + ) + .await; + // All four ownership markers must land on auto-create: the descriptor + // write path refuses tables that miss any of them, so a missing stamp + // here would silently stop every follow-up descriptor write. + validate_data( + "otlp_metrics_resource_info_options", + &client, + "select count(*) from information_schema.tables where table_name = 'greptime_otel_resource_info' \ + and engine = 'mito' \ + and create_options like '%greptime.semantic.metric.type=info%' \ + and create_options like '%greptime.semantic.metric.metadata_quality=declared%' \ + and create_options like '%greptime.semantic.signal_type=metric%' \ + and create_options like '%greptime.semantic.source=opentelemetry%';", + "[[1]]", + ) + .await; + // drop table let res = client .get("/v1/sql?sql=drop table `claude_code_cost_usage_USD_total`;") @@ -6578,6 +6607,11 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { .send() .await; assert_eq!(res.status(), StatusCode::OK); + let res = client + .get("/v1/sql?sql=drop table greptime_otel_resource_info;") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); // write metrics data // with scope attrs @@ -6704,6 +6738,16 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { ) .await; + // the descriptor was auto-created again after the earlier drop + validate_data( + "otlp_metrics_resource_info_recreated", + &client, + "select count(*) from information_schema.tables \ + where table_name = 'greptime_otel_resource_info' and engine = 'mito';", + "[[1]]", + ) + .await; + // drop table let res = client .get("/v1/sql?sql=drop table `claude_code_cost_usage_USD_total`;") @@ -6715,6 +6759,44 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { .send() .await; assert_eq!(res.status(), StatusCode::OK); + let res = client + .get("/v1/sql?sql=drop table greptime_otel_resource_info;") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + + // job composition end-to-end: service.namespace folds into "/" + // on both the descriptor and the metric tables + let content = r#" +{"resourceMetrics":[{"resource":{"attributes":[{"key":"service.name","value":{"stringValue":"api"}},{"key":"service.namespace","value":{"stringValue":"shop"}},{"key":"host.id","value":{"stringValue":"h-1"}}]},"scopeMetrics":[{"scope":{"name":"s"},"metrics":[{"name":"ns_gauge","gauge":{"dataPoints":[{"timeUnixNano":"1753780559836000000","asDouble":1.0}]}}]}]}]} + "#; + let req: ExportMetricsServiceRequest = serde_json::from_str(content).unwrap(); + let res = send_req( + &client, + vec![( + HeaderName::from_static("content-type"), + HeaderValue::from_static("application/x-protobuf"), + )], + "/v1/otlp/v1/metrics", + req.encode_to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + validate_data( + "otlp_metrics_namespace_job_descriptor", + &client, + "select job, \"service.name\", \"service.namespace\", \"host.id\" from greptime_otel_resource_info;", + "[[\"shop/api\",\"api\",\"shop\",\"h-1\"]]", + ) + .await; + validate_data( + "otlp_metrics_namespace_job_metric_table", + &client, + "select job from ns_gauge;", + "[[\"shop/api\"]]", + ) + .await; guard.remove_all().await; } @@ -6912,6 +6994,24 @@ pub async fn test_otlp_metric_translation_strategies(store_type: StorageType) { ) .await; + // the descriptor's columns are the fixed raw keys regardless of the + // translation strategy, and only allowlisted attributes are projected + validate_data( + "otlp_metric_translation_strategy_resource_info", + &client, + "select job, \"service.name\" from greptime_otel_resource_info;", + "[[\"strategy-service\",\"strategy-service\"]]", + ) + .await; + validate_data( + "otlp_metric_translation_strategy_resource_info_allowlist", + &client, + "select count(*) from information_schema.columns \ + where table_name = 'greptime_otel_resource_info' and column_name = 'resource.attr';", + "[[0]]", + ) + .await; + let res = send_req( &client, vec![ @@ -6934,6 +7034,122 @@ pub async fn test_otlp_metric_translation_strategies(store_type: StorageType) { guard.remove_all().await; } +pub async fn test_otlp_metrics_resource_info_conflicts(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = + setup_test_http_app_with_frontend(store_type, "test_otlp_metrics_resource_info_conflicts") + .await; + let client = TestClient::new(app).await; + + let content_type = || { + ( + HeaderName::from_static("content-type"), + HeaderValue::from_static("application/x-protobuf"), + ) + }; + + // A metric named like the descriptor wins the table name: it goes through + // the normal metric-engine path and no descriptor is synthesized. + let content = r#" +{"resourceMetrics":[{"resource":{"attributes":[{"key":"service.name","value":{"stringValue":"api"}},{"key":"host.id","value":{"stringValue":"h-1"}}]},"scopeMetrics":[{"scope":{"name":"s"},"metrics":[{"name":"greptime_otel_resource_info","gauge":{"dataPoints":[{"timeUnixNano":"1753780559836000000","asDouble":1.0}]}}]}]}]} + "#; + let req: ExportMetricsServiceRequest = serde_json::from_str(content).unwrap(); + let res = send_req( + &client, + vec![content_type()], + "/v1/otlp/v1/metrics", + req.encode_to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + let body = ExportMetricsServiceResponse::decode(res.bytes().await).unwrap(); + assert!(body.partial_success.is_none()); + validate_data( + "otlp_metrics_name_collision_engine", + &client, + "select count(*) from information_schema.tables \ + where table_name = 'greptime_otel_resource_info' and engine = 'metric';", + "[[1]]", + ) + .await; + let res = execute_sql(&client, "drop table greptime_otel_resource_info;").await; + assert_eq!(res.status(), StatusCode::OK); + + // A pre-existing table with an incompatible schema fails the descriptor + // write, which degrades to an OTLP partial-success warning while the + // metric data itself is accepted. + let res = execute_sql( + &client, + "create table greptime_otel_resource_info (ts timestamp time index, job double, greptime_value double);", + ) + .await; + assert_eq!(res.status(), StatusCode::OK); + + let content = r#" +{"resourceMetrics":[{"resource":{"attributes":[{"key":"service.name","value":{"stringValue":"api"}}]},"scopeMetrics":[{"scope":{"name":"s"},"metrics":[{"name":"conflict_gauge","gauge":{"dataPoints":[{"timeUnixNano":"1753780559836000000","asDouble":7.0}]}}]}]}]} + "#; + let req: ExportMetricsServiceRequest = serde_json::from_str(content).unwrap(); + let res = send_req( + &client, + vec![content_type()], + "/v1/otlp/v1/metrics", + req.encode_to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + let body = ExportMetricsServiceResponse::decode(res.bytes().await).unwrap(); + let partial_success = body.partial_success.as_ref().unwrap(); + assert_eq!(partial_success.rejected_data_points, 0); + assert!( + partial_success + .error_message + .contains("greptime_otel_resource_info"), + "unexpected warning: {}", + partial_success.error_message + ); + validate_data( + "otlp_metrics_conflict_main_data_accepted", + &client, + "select greptime_value from conflict_gauge;", + "[[7.0]]", + ) + .await; + + // The descriptor is auto-created again once the conflicting table is + // gone, and a follow-up write into it does not degrade. + let res = execute_sql(&client, "drop table greptime_otel_resource_info;").await; + assert_eq!(res.status(), StatusCode::OK); + for round in 0..2 { + let req: ExportMetricsServiceRequest = serde_json::from_str(content).unwrap(); + let res = send_req( + &client, + vec![content_type()], + "/v1/otlp/v1/metrics", + req.encode_to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + let body = ExportMetricsServiceResponse::decode(res.bytes().await).unwrap(); + assert!( + body.partial_success.is_none(), + "round {round} unexpectedly degraded: {:?}", + body.partial_success + ); + } + validate_data( + "otlp_metrics_owned_descriptor_keeps_accepting", + &client, + "select count(*) from greptime_otel_resource_info;", + "[[1]]", + ) + .await; + + guard.remove_all().await; +} + pub async fn test_otlp_traces_v0(store_type: StorageType) { // init common_telemetry::init_default_ut_logging(); @@ -7136,7 +7352,7 @@ pub async fn test_otlp_traces_v1_entity_graph(store_type: StorageType) { let res = send_trace_v1_req(&client, "graph_traces", req, false).await; assert_eq!(StatusCode::OK, res.status()); - let expected = r#"[["service","frontend","service","cart","calls"],["service.instance","resource_attributes.service.instance.id=binst-1,service_name=batch","service","batch","part_of"],["service.instance","resource_attributes.service.instance.id=inst-1,service_name=frontend","service","frontend","part_of"],["service.instance","resource_attributes.service.instance.id=inst-1,service_name=frontend","k8s.pod","uid-9","runs_on"]]"#; + let expected = r#"[["service","frontend","service","cart","calls"],["service.instance","batch,binst-1","service","batch","part_of"],["service.instance","frontend,inst-1","service","frontend","part_of"],["service.instance","frontend,inst-1","k8s.pod","uid-9","runs_on"]]"#; validate_data( "otlp_traces_entity_graph", &client, diff --git a/tests/cases/standalone/common/system/semantic_graph.result b/tests/cases/standalone/common/system/semantic_graph.result index 6af28cfa2b5..f8894319842 100644 --- a/tests/cases/standalone/common/system/semantic_graph.result +++ b/tests/cases/standalone/common/system/semantic_graph.result @@ -94,26 +94,26 @@ select entity_type, entity_id, entity_id_attrs, scope, descriptive, source_table from greptime_private.semantic_entities order by entity_type, entity_id; -+------------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ -| entity_type | entity_id | entity_id_attrs | scope | descriptive | source_tables | -+------------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ -| host | h1 | | | | ["public.graph_app_metrics"] | -| process | host=h1,service_name=cart | {"host":"h1","service_name":"cart"} | | {"env":"us-east"} | ["public.graph_app_metrics"] | -| service | cart | | us-east | | ["public.graph_app_metrics"] | -| service.instance | cart-0 | | | | ["public.graph_app_metrics"] | -+------------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ ++------------------+-----------+-------------------------------------+---------+-------------------+------------------------------+ +| entity_type | entity_id | entity_id_attrs | scope | descriptive | source_tables | ++------------------+-----------+-------------------------------------+---------+-------------------+------------------------------+ +| host | h1 | | | | ["public.graph_app_metrics"] | +| process | cart,h1 | {"host":"h1","service_name":"cart"} | | {"env":"us-east"} | ["public.graph_app_metrics"] | +| service | cart | | us-east | | ["public.graph_app_metrics"] | +| service.instance | cart-0 | | | | ["public.graph_app_metrics"] | ++------------------+-----------+-------------------------------------+---------+-------------------+------------------------------+ select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence from greptime_private.semantic_relationships order by rel_type, src_id; -+------------------+---------------------------+----------+--------+----------+------------+------------+ -| src_type | src_id | dst_type | dst_id | rel_type | provenance | confidence | -+------------------+---------------------------+----------+--------+----------+------------+------------+ -| service.instance | cart-0 | service | cart | part_of | attribute | 1.0 | -| service.instance | cart-0 | host | h1 | runs_on | attribute | 1.0 | -| process | host=h1,service_name=cart | host | h1 | runs_on | attribute | 1.0 | -+------------------+---------------------------+----------+--------+----------+------------+------------+ ++------------------+---------+----------+--------+----------+------------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | confidence | ++------------------+---------+----------+--------+----------+------------+------------+ +| service.instance | cart-0 | service | cart | part_of | attribute | 1.0 | +| process | cart,h1 | host | h1 | runs_on | attribute | 1.0 | +| service.instance | cart-0 | host | h1 | runs_on | attribute | 1.0 | ++------------------+---------+----------+--------+----------+------------+------------+ drop table graph_app_metrics; @@ -570,38 +570,38 @@ select entity_type, entity_id, source_tables from greptime_private.semantic_entities order by entity_type, entity_id, source_tables; -+------------------+-----------------------------------------------------------+------------------------------------+ -| entity_type | entity_id | source_tables | -+------------------+-----------------------------------------------------------+------------------------------------+ -| k8s.container | container=main,uid=uid-1 | ["public.kube_pod_container_info"] | -| k8s.node | node-a | ["public.kube_pod_info"] | -| k8s.pod | uid-1 | ["public.kube_pod_info"] | -| k8s.pod | uid-1 | ["public.kube_pod_owner"] | -| k8s.pod | uid-1 | ["public.kube_pod_container_info"] | -| k8s.pod | uid-2 | ["public.kube_pod_info"] | -| k8s.pod | uid-2 | ["public.kube_pod_owner"] | -| k8s.pod | uid-3 | ["public.kube_pod_info"] | -| k8s.service | svc-uid-1 | ["public.kube_service_info"] | -| k8s.workload | namespace=default,owner_kind=ReplicaSet,owner_name=api-rs | ["public.kube_pod_owner"] | -| service | shop/api | ["public.target_info"] | -| service.instance | instance=inst-1,job=shop/api | ["public.target_info"] | -+------------------+-----------------------------------------------------------+------------------------------------+ ++------------------+---------------------------+------------------------------------+ +| entity_type | entity_id | source_tables | ++------------------+---------------------------+------------------------------------+ +| k8s.container | uid-1,main | ["public.kube_pod_container_info"] | +| k8s.node | node-a | ["public.kube_pod_info"] | +| k8s.pod | uid-1 | ["public.kube_pod_info"] | +| k8s.pod | uid-1 | ["public.kube_pod_owner"] | +| k8s.pod | uid-1 | ["public.kube_pod_container_info"] | +| k8s.pod | uid-2 | ["public.kube_pod_info"] | +| k8s.pod | uid-2 | ["public.kube_pod_owner"] | +| k8s.pod | uid-3 | ["public.kube_pod_info"] | +| k8s.service | svc-uid-1 | ["public.kube_service_info"] | +| k8s.workload | default,ReplicaSet,api-rs | ["public.kube_pod_owner"] | +| service | shop/api | ["public.target_info"] | +| service.instance | shop/api,inst-1 | ["public.target_info"] | ++------------------+---------------------------+------------------------------------+ -- SQLNESS PROTOCOL MYSQL select src_type, src_id, dst_type, dst_id, rel_type, provenance from greptime_private.semantic_relationships order by rel_type, src_id, dst_id; -+------------------+------------------------------+---------------+-----------------------------------------------------------+----------+------------+ -| src_type | src_id | dst_type | dst_id | rel_type | provenance | -+------------------+------------------------------+---------------+-----------------------------------------------------------+----------+------------+ -| k8s.pod | uid-1 | k8s.container | container=main,uid=uid-1 | contains | attribute | -| service.instance | instance=inst-1,job=shop/api | service | shop/api | part_of | attribute | -| k8s.pod | uid-1 | k8s.workload | namespace=default,owner_kind=ReplicaSet,owner_name=api-rs | part_of | attribute | -| k8s.pod | uid-2 | k8s.workload | namespace=default,owner_kind=ReplicaSet,owner_name=api-rs | part_of | attribute | -| k8s.pod | uid-1 | k8s.node | node-a | runs_on | attribute | -| k8s.pod | uid-2 | k8s.node | node-a | runs_on | attribute | -+------------------+------------------------------+---------------+-----------------------------------------------------------+----------+------------+ ++------------------+-----------------+---------------+---------------------------+----------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | ++------------------+-----------------+---------------+---------------------------+----------+------------+ +| k8s.pod | uid-1 | k8s.container | uid-1,main | contains | attribute | +| service.instance | shop/api,inst-1 | service | shop/api | part_of | attribute | +| k8s.pod | uid-1 | k8s.workload | default,ReplicaSet,api-rs | part_of | attribute | +| k8s.pod | uid-2 | k8s.workload | default,ReplicaSet,api-rs | part_of | attribute | +| k8s.pod | uid-1 | k8s.node | node-a | runs_on | attribute | +| k8s.pod | uid-2 | k8s.node | node-a | runs_on | attribute | ++------------------+-----------------+---------------+---------------------------+----------+------------+ drop table kube_pod_info; @@ -627,3 +627,182 @@ drop table http_requests_total; Affected Rows: 0 +-- OTel conventions: the ingestion-synthesized greptime_otel_resource_info descriptor +-- (stamped signal_type=metric + source=opentelemetry + metric.type=info) gets implicit +-- declarations under the raw OTel column names. host/container identities are +-- the stable ids only, so rows with an empty host.id / container.id link no +-- infrastructure, and a row without an instance value declares no +-- service.instance. The same table name stamped source=prometheus is not +-- whitelisted. +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "service.name" string, + "service.namespace" string, + "host.id" string, + "host.name" string, + "container.id" string, + "container.name" string, + greptime_value double, + primary key (job, instance, "service.name", "service.namespace", + "host.id", "host.name", "container.id", "container.name") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'opentelemetry', + 'greptime.semantic.metric.type' = 'info' +); + +Affected Rows: 0 + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'api', 'shop', 'h-1', 'node-a', 'c-1', 'api-ctr', 1), + (now(), 'shop/api', 'inst-2', 'api', 'shop', '', 'laptop', '', '', 1), + (now(), 'worker', '', 'worker', '', 'h-1', 'node-a', '', '', 1); + +Affected Rows: 3 + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +order by entity_type, entity_id, source_tables; + ++------------------+-----------------+----------------------------------------+ +| entity_type | entity_id | source_tables | ++------------------+-----------------+----------------------------------------+ +| container | c-1 | ["public.greptime_otel_resource_info"] | +| host | h-1 | ["public.greptime_otel_resource_info"] | +| service | shop/api | ["public.greptime_otel_resource_info"] | +| service | worker | ["public.greptime_otel_resource_info"] | +| service.instance | shop/api,inst-1 | ["public.greptime_otel_resource_info"] | +| service.instance | shop/api,inst-2 | ["public.greptime_otel_resource_info"] | ++------------------+-----------------+----------------------------------------+ + +-- SQLNESS PROTOCOL MYSQL +select src_type, src_id, dst_type, dst_id, rel_type, provenance +from greptime_private.semantic_relationships +order by rel_type, src_id, dst_id; + ++------------------+-----------------+-----------+----------+----------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | ++------------------+-----------------+-----------+----------+----------+------------+ +| service.instance | shop/api,inst-1 | service | shop/api | part_of | attribute | +| service.instance | shop/api,inst-2 | service | shop/api | part_of | attribute | +| container | c-1 | host | h-1 | runs_on | attribute | +| service.instance | shop/api,inst-1 | container | c-1 | runs_on | attribute | +| service.instance | shop/api,inst-1 | host | h-1 | runs_on | attribute | ++------------------+-----------------+-----------+----------+----------+------------+ + +drop table greptime_otel_resource_info; + +Affected Rows: 0 + +-- The descriptor whitelist is gated on source=opentelemetry: the same table +-- shape stamped as a prometheus source contributes nothing. +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "host.id" string, + greptime_value double, + primary key (job, instance, "host.id") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +Affected Rows: 0 + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'h-1', 1); + +Affected Rows: 1 + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id +from greptime_private.semantic_entities +order by entity_type, entity_id; + +affected_rows: 0 + +drop table greptime_otel_resource_info; + +Affected Rows: 0 + +-- Cross-signal identity: the same service reaches the graph as trace spans and +-- as an ingestion-synthesized descriptor. The two sources name their identity +-- columns differently and only the metric side pre-composes the namespace into +-- job, so this asserts one entity per service and per instance, sourced from +-- both tables. +create table graph_xsignal_traces ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + "resource_attributes.service.namespace" string, + "resource_attributes.service.instance.id" string, + primary key (service_name, "resource_attributes.service.namespace", + "resource_attributes.service.instance.id") +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +Affected Rows: 0 + +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "service.name" string, + "service.namespace" string, + greptime_value double, + primary key (job, instance, "service.name", "service.namespace") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'opentelemetry', + 'greptime.semantic.metric.type' = 'info' +); + +Affected Rows: 0 + +insert into graph_xsignal_traces values + (now(), 't1', 's1', NULL, 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'api', 100, 'shop', 'inst-1'), + (now(), 't2', 's2', NULL, 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'worker', 100, '', 'inst-2'); + +Affected Rows: 2 + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'api', 'shop', 1), + (now(), 'worker', 'inst-2', 'worker', '', 1); + +Affected Rows: 2 + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +where entity_type in ('service', 'service.instance') +order by entity_type, entity_id, source_tables; + ++------------------+-----------------+----------------------------------------+ +| entity_type | entity_id | source_tables | ++------------------+-----------------+----------------------------------------+ +| service | shop/api | ["public.graph_xsignal_traces"] | +| service | shop/api | ["public.greptime_otel_resource_info"] | +| service | worker | ["public.graph_xsignal_traces"] | +| service | worker | ["public.greptime_otel_resource_info"] | +| service.instance | shop/api,inst-1 | ["public.graph_xsignal_traces"] | +| service.instance | shop/api,inst-1 | ["public.greptime_otel_resource_info"] | +| service.instance | worker,inst-2 | ["public.graph_xsignal_traces"] | +| service.instance | worker,inst-2 | ["public.greptime_otel_resource_info"] | ++------------------+-----------------+----------------------------------------+ + +drop table graph_xsignal_traces; + +Affected Rows: 0 + +drop table greptime_otel_resource_info; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/system/semantic_graph.sql b/tests/cases/standalone/common/system/semantic_graph.sql index bfc9ea8fd7b..60d95630d91 100644 --- a/tests/cases/standalone/common/system/semantic_graph.sql +++ b/tests/cases/standalone/common/system/semantic_graph.sql @@ -387,3 +387,122 @@ drop table kube_service_info; drop table target_info; drop table http_requests_total; + +-- OTel conventions: the ingestion-synthesized greptime_otel_resource_info descriptor +-- (stamped signal_type=metric + source=opentelemetry + metric.type=info) gets implicit +-- declarations under the raw OTel column names. host/container identities are +-- the stable ids only, so rows with an empty host.id / container.id link no +-- infrastructure, and a row without an instance value declares no +-- service.instance. The same table name stamped source=prometheus is not +-- whitelisted. +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "service.name" string, + "service.namespace" string, + "host.id" string, + "host.name" string, + "container.id" string, + "container.name" string, + greptime_value double, + primary key (job, instance, "service.name", "service.namespace", + "host.id", "host.name", "container.id", "container.name") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'opentelemetry', + 'greptime.semantic.metric.type' = 'info' +); + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'api', 'shop', 'h-1', 'node-a', 'c-1', 'api-ctr', 1), + (now(), 'shop/api', 'inst-2', 'api', 'shop', '', 'laptop', '', '', 1), + (now(), 'worker', '', 'worker', '', 'h-1', 'node-a', '', '', 1); + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +order by entity_type, entity_id, source_tables; + +-- SQLNESS PROTOCOL MYSQL +select src_type, src_id, dst_type, dst_id, rel_type, provenance +from greptime_private.semantic_relationships +order by rel_type, src_id, dst_id; + +drop table greptime_otel_resource_info; + +-- The descriptor whitelist is gated on source=opentelemetry: the same table +-- shape stamped as a prometheus source contributes nothing. +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "host.id" string, + greptime_value double, + primary key (job, instance, "host.id") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'h-1', 1); + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id +from greptime_private.semantic_entities +order by entity_type, entity_id; + +drop table greptime_otel_resource_info; + +-- Cross-signal identity: the same service reaches the graph as trace spans and +-- as an ingestion-synthesized descriptor. The two sources name their identity +-- columns differently and only the metric side pre-composes the namespace into +-- job, so this asserts one entity per service and per instance, sourced from +-- both tables. +create table graph_xsignal_traces ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + "resource_attributes.service.namespace" string, + "resource_attributes.service.instance.id" string, + primary key (service_name, "resource_attributes.service.namespace", + "resource_attributes.service.instance.id") +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "service.name" string, + "service.namespace" string, + greptime_value double, + primary key (job, instance, "service.name", "service.namespace") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'opentelemetry', + 'greptime.semantic.metric.type' = 'info' +); + +insert into graph_xsignal_traces values + (now(), 't1', 's1', NULL, 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'api', 100, 'shop', 'inst-1'), + (now(), 't2', 's2', NULL, 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'worker', 100, '', 'inst-2'); + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'api', 'shop', 1), + (now(), 'worker', 'inst-2', 'worker', '', 1); + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +where entity_type in ('service', 'service.instance') +order by entity_type, entity_id, source_tables; + +drop table graph_xsignal_traces; + +drop table greptime_otel_resource_info; diff --git a/tests/conf/frontend-test.toml.template b/tests/conf/frontend-test.toml.template index 632947a3cde..5a879e04882 100644 --- a/tests/conf/frontend-test.toml.template +++ b/tests/conf/frontend-test.toml.template @@ -7,3 +7,6 @@ parallelism = 4 [event_recorder] event_types = [] + +[otlp] +experimental_enable_resource_info = true diff --git a/tests/conf/standalone-test.toml.template b/tests/conf/standalone-test.toml.template index ebe25f82487..8d7e307d511 100644 --- a/tests/conf/standalone-test.toml.template +++ b/tests/conf/standalone-test.toml.template @@ -53,3 +53,6 @@ retry_delay = "500ms" [event_recorder] event_types = [] + +[otlp] +experimental_enable_resource_info = true From 144f83528db8648bc5607fdc65c249bbfa3c7ffc Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 26 Aug 2026 06:15:37 +0000 Subject: [PATCH 03/22] test: renew etcd TLS certificates (#8956) Signed-off-by: WenyXu --- .../fixtures/etcd-tls-certs/ca-key.pem | 52 +++++++++---------- .../fixtures/etcd-tls-certs/ca.crt | 34 ++++++------ .../fixtures/etcd-tls-certs/ca.srl | 2 +- .../fixtures/etcd-tls-certs/client-key.pem | 52 +++++++++---------- .../fixtures/etcd-tls-certs/client.crt | 34 ++++++------ .../fixtures/etcd-tls-certs/server-key.pem | 52 +++++++++---------- .../fixtures/etcd-tls-certs/server.crt | 32 ++++++------ 7 files changed, 129 insertions(+), 129 deletions(-) diff --git a/tests-integration/fixtures/etcd-tls-certs/ca-key.pem b/tests-integration/fixtures/etcd-tls-certs/ca-key.pem index 2066dfb7b47..07a0a86ee39 100644 --- a/tests-integration/fixtures/etcd-tls-certs/ca-key.pem +++ b/tests-integration/fixtures/etcd-tls-certs/ca-key.pem @@ -1,28 +1,28 @@ -----BEGIN PRIVATE KEY----- -MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCfIi0iTllUQyNY -ka+d9Zt30ILA0QBzLKHwez9QJj8bQLpfKR26d794CQn/OO6XO7MWsJ6nvqwRNYM0 -WX+UwUtuT9MrttYWpS5yEAVO79dMtF67gCNafXKjQyK8L1vz6Avq4XJN4VqZEifT -Fdk3s2MmS0twv4sgrK//jh7aZzBQJtw+ROe9RN1dqlKoa3TtpB7iwNTDxGiR0EYj -wHj+GdiNHtDMWvrnTNNX/IKeEwLP/Mjirzr7GnaOjR25ruqcLuwWrufu+EMhqFTU -r+LIQpZD0TFuPX6DPlQZOs/4UpjteoC84m5AnQXNlh0vzjbW1efEYWYQoKDesx9e -oE5jcABpAgMBAAECggEAA3Rw/r6CYaTG1tdBilA47LF/XTihu15mh6Y4BNadESDn -IEYav0p2mDW4sgH7lWyg4ecPaBHo124jfUFM1p8Zsvk9sERwbR5vSIpWKyq5hbtM -Fo2yIXakGgJM9ev19vDR7WqxRMVAkR6GnyZpcwmh0e+ENvGZpNR1n7qSAJOJRibW -Ye8qMF/zMLF9VekBRPFJFIZBvKgHxSGCLbXRzsi5LKMbwfBoJmxEZSOC9zODl43o -Z2i+LvlS0R/Rsf2KG+Hwgj1P7HH1znJeA5bVJFDewljjSwu1mk0JpQL7Jgp2xpqs -F/mRZmtoQSH3EsTm2Kjz1sLGXwUb9JaQ274aV4t7pwKBgQDXYF+qFYip7k+b9FTF -sFbfp3CjOqxXiAxnoE5hlfV/h0w4V1Su42AdejnbIVYFX8bvmezkNWPXhB6P5oHU -4nO+oAuhsRT4317mfRM5NWO6mv2eglSBnS04GLsX142A7jMPr9e+Cqh4dYfw/A1L -bSdsdILemTFQrak27pC11mHZOwKBgQC9JhDhQf5uRh6k8GvgKGu9ZEZmkYCVu2sr -7kpD60UhklnyQo6DKYsa6OcChRQMv9ss3Hqu3Y1G0/sTe1Tjz5zwrWWAZoLlbdOt -vipesvwNTxvJoUWPtrjOoFZtmyg0Q0TqxaxfEWAssDU9IfnIy+zx7sAg4zZIPUyy -6kudaY1SqwKBgQCAQ/71Fjn7qddzc4GA8lHqhJeKPokg3/8zP78uUtaQCo2UCD6A -oR0+sOn/3MyUCsQ5MZxpFHrPgPmKjabIl8yCvGHw+7sXtD+aWOa37VnlaiSc39Vg -E7E4dVIHEvJM1I9ISlrb7REEHErHc/Se9PTDnGfMFcPO3n2mH1HDWVeQvQKBgGGn -BHH3a08tXmbTRS5uT+lwmrQbjKJBJ3x/wtG75m4Fq/BaEk9/JDUZZyKy5/4JEzPf -BGvBME4P5QFS3CndJu5O5ydaRVwDzpRVqHRJvb11SShY3Zvrvw/WUai2wRPyYuM+ -eNaAFwIbWvEb2GSle8gP9htEkuLK2w1HzxAOzYqPAoGAKl5l/WsoLfQP5UIUjSQe -JtVTDVmoz8m8Mg2nn56UdUetG246EVhDx5ngj84kZDM0GDXhfVWV6MWnIZr5NQsr -bORIDNP7Czat8FnlXeAQ+R6A/87o5g5p6ydH9jYqrrLxYjaZyXJptE+ja4zW7Lbl -ogWAc76iZvXLXqKRfepG0zQ= +MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQC05De2Tk8TMZ0J +ME+7Rp+zvvLkEBPPXCCFcXfo91gYVHMMr1TGQJuVtbx7rIqC1tFhc2luAGfsBmW3 +IfTXZjzY6834ZUzmPL7uuBfoRpd0Paj5X4owpn0eHIgcsqx4juHSTYc3fdnrTLpq +lNMulZEemkNULLCE7pyxg3hhYIx9JEe9nv26fkyr1vW5z1cPlEpyOK+SUIeqHgXi +VtWPU9WGv5sdQeXnD4E1JTX2URihb7C8xali3YVbI3SUxy/gAhM6swQgZ4SLgW8k +f87a+PcP9C1N5Gujfw4fT90dJoutyiTSbcDrjZLjF03TgjcDaACvntb0qN85XD96 +dIvhg5jfAgMBAAECggEAAd3vHe6d+KTMDZ1HwPAAKS151ipnpX3W3aRgi/FC6OMm +r0VUpWC4jCJlif3kYvaID66yFeucJveD75CJw6VaF4bSnTa9GmKM5F2BQhsrhKj+ +69jY+MU8cztGfJMdgczNfi79qDc8zSW6TUFpoMhmD6Og8icRCODbkbRWDDnnQi1d +1CVe/lZ4yIapCGgzHfQxdl9I8LEZWyeAfqMCvOTI2eE4Z4ZYto/EmuRoJN3oqOx1 +MAtpS5mw6RTFVitO4b0j5gwAN2u6Cgo8jDc5UTzHGAiGc/dAn7S+mNGv+IXP1/aL +UA/ZsBwQhZbqtGCkSmCo5gGDL+anDTiIy6nRCwIWWQKBgQDtE7V4pNj5J2MciK+h +KAR/kT36NuGoZSnf585HJVB/xAs/d8F4kU5pCofmrtm2Bl5Lx4sCE3wPmv/HMYcI +QH2zt/YxPA3F/9zoBowfTNB8/iCDlzNb/pUry+qkBhKG2HgWaUkkGEtiO16PeCF2 +bidYA+h+XZBnHKexAKL74ee71wKBgQDDVHJ6HRo1uTAT/8Uz3AF5NRj+qj9QC3UX +On6dtzbtDGsJFkcqW5CV+LwUPN6zQ7MW7U6AEeH99q1UAUDkd5Q8/4bqs+0y2Ss7 +0UUWr1LDJVYo6FIs/f2ON0e9LzcJcG9TgZ6T95TL1kGEG1Ng96Je3hAZWolLmr+x +yfjTSYKqOQKBgEwRWtTO799hx+dL5C5tTKQx0hUKrvT9IKZ7FjC1xFJ6cLF4l1c8 +KFCD1H8r8yb4fCEMcYnE/rVzIkajmZQIPU0A5bl+b1zsb9Dy6NrSJsM0NvKB/TSz +RuG6mBrw59jkdAOc3J78PJPUQM7/2JzLU0xmVJ7XHpI3G4crkSAIp/YZAoGATX6m +gF4ldOUI9xZFleKWTxFK3laLEeXJybJyY34582g22v8UsvBq96UccWcI79RPLCxw +NY1ivNBuSeLJbRsoG99BFsLVu5O/fFb1cx+R1Uxt14L8f08xlofGFX+y9TK/aEEH +uirCxPA3RANXXCRDLiIp/vUVfYJixVWdO65xgbkCgYAVQhvJ3jPXZLXPl9fVC3Ro +qweEZHs0knPc4mqL/7bE+BCE1HcccU7O7+ZKbhW2q11EMrPJ7cHUqweHE3GsTgRC +C5FIvRxInZUgtPNx2UIZAVNjxRdCZJHBc/59eu8wVp9f5ETiSFmdz5xlsJECSxdO +mBakpZMqnESl1/6FMtcqeQ== -----END PRIVATE KEY----- diff --git a/tests-integration/fixtures/etcd-tls-certs/ca.crt b/tests-integration/fixtures/etcd-tls-certs/ca.crt index 8723f7ffedd..ecd404c8918 100644 --- a/tests-integration/fixtures/etcd-tls-certs/ca.crt +++ b/tests-integration/fixtures/etcd-tls-certs/ca.crt @@ -1,21 +1,21 @@ -----BEGIN CERTIFICATE----- -MIIDeTCCAmGgAwIBAgIUeRD4bMVd0XEzS3UHo2ptnknsWzEwDQYJKoZIhvcNAQEL +MIIDeTCCAmGgAwIBAgIUM4xko7n+KubCgGJnX4EaUzP7ZicwDQYJKoZIhvcNAQEL BQAwTDELMAkGA1UEBhMCVVMxCzAJBgNVBAgMAkNBMQswCQYDVQQHDAJTRjERMA8G -A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjUwODI1MTgwNjU1 -WhcNMjYwODI1MTgwNjU1WjBMMQswCQYDVQQGEwJVUzELMAkGA1UECAwCQ0ExCzAJ +A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjYwODI2MDQxMzQy +WhcNMjcwODI2MDQxMzQyWjBMMQswCQYDVQQGEwJVUzELMAkGA1UECAwCQ0ExCzAJ BgNVBAcMAlNGMREwDwYDVQQKDAhHcmVwdGltZTEQMA4GA1UEAwwHZXRjZC1jYTCC -ASIwDQYJKoZIhvcNAQEBBQADggEPADCCAQoCggEBAJ8iLSJOWVRDI1iRr531m3fQ -gsDRAHMsofB7P1AmPxtAul8pHbp3v3gJCf847pc7sxawnqe+rBE1gzRZf5TBS25P -0yu21halLnIQBU7v10y0XruAI1p9cqNDIrwvW/PoC+rhck3hWpkSJ9MV2TezYyZL -S3C/iyCsr/+OHtpnMFAm3D5E571E3V2qUqhrdO2kHuLA1MPEaJHQRiPAeP4Z2I0e -0Mxa+udM01f8gp4TAs/8yOKvOvsado6NHbmu6pwu7Bau5+74QyGoVNSv4shClkPR -MW49foM+VBk6z/hSmO16gLzibkCdBc2WHS/ONtbV58RhZhCgoN6zH16gTmNwAGkC -AwEAAaNTMFEwHQYDVR0OBBYEFOCkI3Uyx7F38LtXtSlg3ORE9ur6MB8GA1UdIwQY -MBaAFOCkI3Uyx7F38LtXtSlg3ORE9ur6MA8GA1UdEwEB/wQFMAMBAf8wDQYJKoZI -hvcNAQELBQADggEBAJXbWJJ9b+5nXjRiFXOg+wQnzn7kzMf6sWKIFk+AuKJWWt78 -O00t6vyAz6zel5Cj3ho9yAaMFNy8vYEnJCYngy5pT/2hOncnz/w7IKTeoEhzqAnf -MCEmCgHbTKDoFfMfrrRwtyoePVfx4xbiGVWMQFTPG+WNlE/ivFMRvFwsgPJ+7SUK -mR2FscH6DVo4sqF6s729lTmr6/U7bOD5l2HYGpKJ2cjCI8+HDv55aAT43tsTwBoJ -BgX+RT4w8ryGhYV+hrRTjPlMHWiHsNbzJeAi5bNA0f6vUP6k6zHn/Ur3MNb1avcf -cSsdYU2dhs7PeB6IpoJ0QCeJw9MIFK3XTnmbeVY= +ASIwDQYJKoZIhvcNAQEBBQADggEPADCCAQoCggEBALTkN7ZOTxMxnQkwT7tGn7O+ +8uQQE89cIIVxd+j3WBhUcwyvVMZAm5W1vHusioLW0WFzaW4AZ+wGZbch9NdmPNjr +zfhlTOY8vu64F+hGl3Q9qPlfijCmfR4ciByyrHiO4dJNhzd92etMumqU0y6VkR6a +Q1QssITunLGDeGFgjH0kR72e/bp+TKvW9bnPVw+USnI4r5JQh6oeBeJW1Y9T1Ya/ +mx1B5ecPgTUlNfZRGKFvsLzFqWLdhVsjdJTHL+ACEzqzBCBnhIuBbyR/ztr49w/0 +LU3ka6N/Dh9P3R0mi63KJNJtwOuNkuMXTdOCNwNoAK+e1vSo3zlcP3p0i+GDmN8C +AwEAAaNTMFEwHQYDVR0OBBYEFDjvuYaWEoViLOKyCse6mAwH66mIMB8GA1UdIwQY +MBaAFDjvuYaWEoViLOKyCse6mAwH66mIMA8GA1UdEwEB/wQFMAMBAf8wDQYJKoZI +hvcNAQELBQADggEBAJuGxtprq/GF/fwUJ5hk0cG80VY24G+iKKi7Jw7PWPEMvnlQ +/ZIKgNDupE8mSSyokicKZMrtwT8QF6wlDLVcz4LTPCtKUmYU3FiuNE4V57+yw7zt +XZn6rwiWGxeyfY0ZktLXiSYA5FydOJ9GQQWgxkpbd5QG1L9aoFW49cbz8ob6efxG +p+YRgr0Fr4hXn7ImJWWqeawf2tg4iLwpV4sem5atpI7y9qamxrGGUp2h/CnADqUz +YWYIYP1RCPp2bXzxO40Ni1NYmFsGUU/SL1vtJw5n2Zmx/Phe7u/CLe05dH1O54Qf +vCbk6X0CthFI0HzSepbEVAuniVHkrpWwDLHcQ6s= -----END CERTIFICATE----- diff --git a/tests-integration/fixtures/etcd-tls-certs/ca.srl b/tests-integration/fixtures/etcd-tls-certs/ca.srl index 69ac026da25..ad32d06e16e 100644 --- a/tests-integration/fixtures/etcd-tls-certs/ca.srl +++ b/tests-integration/fixtures/etcd-tls-certs/ca.srl @@ -1 +1 @@ -4484A9BBD4F25F32994F3C03D80D294105B119E6 +4484A9BBD4F25F32994F3C03D80D294105B119E8 diff --git a/tests-integration/fixtures/etcd-tls-certs/client-key.pem b/tests-integration/fixtures/etcd-tls-certs/client-key.pem index 13a23c53478..1d64fc66a3d 100644 --- a/tests-integration/fixtures/etcd-tls-certs/client-key.pem +++ b/tests-integration/fixtures/etcd-tls-certs/client-key.pem @@ -1,28 +1,28 @@ -----BEGIN PRIVATE KEY----- -MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQCOcrzRmE3sEGYo -pazDrpkkzbiaBhHClfmBiZAAGV2cvSXY8e9pc21GOgrGBNeyPPieS305rAE02EZl -oIAt97XRFxkZniUy8C5+PT3eicAKPRcrgGhIrYxUs39rdRocqVVpZsiZGzhHhboB -bXNW1l37SDQpuxAQJuhHu2wu9J7c5/TBXWN9ZOxzWTlrIdEQOA46VclE/N1TlaKs -CSWw8qH87YLwkLEXcFLOJfmWNZxOOhdL5MSzKKNcQKbpmkd65B/PHNwBIBZHpI+K -xphK1SybA+ESkN09qZaa1Kj3e0I62UjHNXUCuAod6RhRRmpiFDFgggXhD9T9Ttuf -1Lq38XxDAgMBAAECggEAANNhqQldwHzi9WCxXObLoufBe9lM73+1Df1IQaDHkZR3 -jHliixxF18XQW7qZ5tAJqjsDWLPm4BxSKXozsjsvBdgIUjbqkyx1dsGlO9TjeGu5 -H2/8jlH14boIMWJ9Y0JBWwvU7EnVU5IkUOEiwsxvFlzM5Q0H+841CrSD+Xd7iSZE -ZrpyZgpHNmvc/u7AmseZVccjX21dxwWtcd7cGUmFWC9o2a/hkgG6+baA8zH2ze3J -Smeowa/oOyOVkTM1EHLCBvddE+/W+BeMGPY7A6SztVw0gLqlSOL3cr/g2llUTmFg -T6MOS6FZZ7MD5c2HBL0V3OXIhxNABzD+c0O3CaRbSQKBgQC/+QuhMpDCzxBPaL/K -sqRp4sQgURDCXPA1nAfI+NoHos/CbjwcMGp6LUSiBW0tLTJ+LZoVgZXj6j/uxji9 -DWxQfULvQ7oanm/q/D+40NnupQQTzmMT04z/mV3XjFtzEhbO30YZfqwXy+tqB8vR -g1KajXVqLmCxKDtsFedAlIVxRQKBgQC99TKGcBwsKlYJJThU2eIIy8G0GkjjXleK -ROMO5i0keJtLac1oRvXo4JN5UBLJM6ZpSy0TEa2bxfzy6Uea+41aZmdCpPg0Kcsq -Pb1S3rsPpVdsfxSxsdKvwZmHXOzAdF+fom1+ZRAdo/8ALIPaNQSSaw+KkmmpHPEV -qdzzWJ4b5wKBgFubEtKcF3nuZwENoh+ueUhRvncRV+b3hGSAjTJ4lUn5hhxoj+R/ -sf+VJGAQKNXa8HJHfnRuvsDgYhulmSOViS8rZspXzjGvkwZV0m51stjvA3AUFzE5 -zNmXLLGTt3vEkP+siX3W9XXxh+ezyq2ydbNsdy/w65D9+sUL+qrVdIvlAoGAOOQy -2ajCB0g2tE59bIxE8jV0MiidI9uhhDvVdSTi6EVm3VM2vcBi7fg0suSUe8YIVQi6 -2zc0M688btQHKhek4ipBSuh1ncnWmzQae7NRewIeCNSWshF79D+bZ7sg/RLdgMX4 -3R4PkZEIUlkCtFuknuWJpgrrskaEveQ91HP6BokCgYAX42nKPANIMrTXuJ2vACWu -qKQ9atAC1ngVzYe2PI8rHjA4gCDraGfCaILewrd89rK9ukBeSDIm5Mn448q3ystw -p/QMxmMdHADCD6UwCjlt9N0Xigv+IKbDzDzhartRVIPpZdTViQXW3cSHt9id8zvO -HTu5C6W+N9Vb7j8Ak/mvzw== +MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQC5F3hVsri2FBzP +dtRe+aJFh9yqhXnJ7zN+xbUmNWsTZVnPwrbVo9H6PtGWzkbZ6YMBzMlQSYP2txaS +ROjE8wwEs38eKbyJrHO2Ccq3iR9NfDxSbHTClngsOjf9fEoYExif84VKWLO6+EKL +KqqOfj3+iuI3hbabr0MqGu4e2Nr6MlQtViBxTs5aSnahc8hmZbujBcWRTsw/GNx5 +Am0a8BrrJRVmVR+BzlYtzbf1ltfLxIhnc8dAGyT9B6GTwcMNST6m4WtTzR/uOr6S +2YEbIqwCHB1VQ+oJUGJvExqaNLoM/x5a8sUNIWx3LBN0/XC4D+MMmh/bbyb2cAOE +gyUYu+hLAgMBAAECggEAEkt+wksVP9toEkecFpEqgss65VimdkPK/WkHDpzxawGL +FDfvq4Py9tRrvLIXRbseObjSCI9jrai2Ne9Sv1jeDhQ8kzDWk8MfItX/8A6Whf8A +KrjRUlneIfz/2Mcxaa9wMSs4YxC2dIYidFqwxhSLWyluJ8Ud4hs7kndb+beUmW3i +DgkLsQ9ytQ0HDQrAERnrbRf19fmVaqQapfxpYT2gg80BMIq+zlzoiXareDz3r/1g +aoponxArWV6yCvQOYYjWcDr4WZPvIA2APEX5isxyvMBvb72Fpd2z+0JX+4LVe6N1 +FqZLNpSeJn6O+Q1vJSScSJcoJYkLWCfZbFLHPttZGQKBgQDdVr5PtgPIVM+67cvv +p3x7NNC10Gp387PgQS/FlG+QtbP87RTBnANT1uFd/FMFqL35or7Y46YKDUdk00OZ +LinpM4ARFRYFqkthM6bcUWatpF6Fg2dfsKuD4H6ARp8YHhlmCnDCQ0swZp1HCmAB +6vltUNTZjQjRdd44GlQ0icX0mQKBgQDWE50TXRszea75S+zUfXeyvxMyptIOS/rY +L50XVuhGVknH4OuBQbGF6UZulIGSpN7gmmbO29Yx1HKmqiJVaWlvB71baK4hWLDd +2grIUy0zuq1B7I1IOLmOpMcL3xau76g7ZWRUprX26hxAuxYyPb9V8gNo1P7HN8si +VLJVE8ZugwKBgQCaxoGmM91JRSVNzeOB3ljJvxEDUo5g+uWZt3u0aivpwWXvQ8nz +6Sjag7RsiHl1x52w5wEVoXsGJGr8Mk9e2k0saXrwdxJDO+YiPoA8KB/o5LvEGTM8 +UspdGarcAIZX0xRnqn1XGr+FRPxOJQ8lyC5LJu7wghLchdOy35Zqdr0aYQKBgQC/ +38yFsonS1VnS8A5RVjOW7lPSrlrPnaIzalmutaJyiJyQnjP3Il5u2+rY6hpIyaVK +QpmrBrcw6m3om80yKMzrS1CZQXXxRYEhF3Fao9J77vGjiNYIyW7nPyF4rneyS/PJ +aNNIXDP0H1k7W3RFi7qW2dfceivxezyChM9iGdtc6QKBgAIfb63+1q98MNlAB7eD +x7K8vuSxHrJ9gX9KcLrG1ScLOxnAC4Op+8xf06ZHt828CajchAOqee2nh52MsFOr +AdPNdhS/8mTq8KosTfgVuawv4yTHpZViMR2KSOcTKhKaBmhufY+k8mJ0YAAOTnzO +JJqjDSJdjWsDG+7UlNgFcvuU -----END PRIVATE KEY----- diff --git a/tests-integration/fixtures/etcd-tls-certs/client.crt b/tests-integration/fixtures/etcd-tls-certs/client.crt index 935351a48b1..0dee5915ffa 100644 --- a/tests-integration/fixtures/etcd-tls-certs/client.crt +++ b/tests-integration/fixtures/etcd-tls-certs/client.crt @@ -1,20 +1,20 @@ -----BEGIN CERTIFICATE----- -MIIDMjCCAhqgAwIBAgIURISpu9TyXzKZTzwD2A0pQQWxGeYwDQYJKoZIhvcNAQEL +MIIDMjCCAhqgAwIBAgIURISpu9TyXzKZTzwD2A0pQQWxGegwDQYJKoZIhvcNAQEL BQAwTDELMAkGA1UEBhMCVVMxCzAJBgNVBAgMAkNBMQswCQYDVQQHDAJTRjERMA8G -A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjUwODI1MTgwNjU1 -WhcNMjYwODI1MTgwNjU1WjAWMRQwEgYDVQQDDAtldGNkLWNsaWVudDCCASIwDQYJ -KoZIhvcNAQEBBQADggEPADCCAQoCggEBAI5yvNGYTewQZiilrMOumSTNuJoGEcKV -+YGJkAAZXZy9Jdjx72lzbUY6CsYE17I8+J5LfTmsATTYRmWggC33tdEXGRmeJTLw -Ln49Pd6JwAo9FyuAaEitjFSzf2t1GhypVWlmyJkbOEeFugFtc1bWXftINCm7EBAm -6Ee7bC70ntzn9MFdY31k7HNZOWsh0RA4DjpVyUT83VOVoqwJJbDyofztgvCQsRdw -Us4l+ZY1nE46F0vkxLMoo1xApumaR3rkH88c3AEgFkekj4rGmErVLJsD4RKQ3T2p -lprUqPd7QjrZSMc1dQK4Ch3pGFFGamIUMWCCBeEP1P1O25/UurfxfEMCAwEAAaNC -MEAwHQYDVR0OBBYEFINcPBk3iqCw1B2sZzcDpSF6ZNGnMB8GA1UdIwQYMBaAFOCk -I3Uyx7F38LtXtSlg3ORE9ur6MA0GCSqGSIb3DQEBCwUAA4IBAQBgyF1XSpAyiQNA -sFCqXrBfwPnv6V+R5jhO6Glmn+Nhj0gbzNBlU8EKsK/WZllDrJjmxubPBqb563Zt -J+QcNLqugZdioXPbxmmi9xj3oK25pUCRBj03nlHpGqIwuCJyn/ZxlfXMpNYE1KjM -nprFNUKgnYPk2AGffAngQdm0yrnmSpxXhNTuhQE8avgRhd5nEUEDyPIvHWwo4iU6 -pw2aP2FpWorESv9LVna+JortThYDQrIFgCtPM6BP1N7Aeqn0H846+JsXlq2ch1Iy -aBRqKNQRc1Q2Qb9/QCDEJ3dEixgFjxztIaqrIcr18ZJ797MhTFS3aVEwE8clSRR5 -9ow9jK8B +A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjYwODI2MDQxMzQy +WhcNMjcwODI2MDQxMzQyWjAWMRQwEgYDVQQDDAtldGNkLWNsaWVudDCCASIwDQYJ +KoZIhvcNAQEBBQADggEPADCCAQoCggEBALkXeFWyuLYUHM921F75okWH3KqFecnv +M37FtSY1axNlWc/CttWj0fo+0ZbORtnpgwHMyVBJg/a3FpJE6MTzDASzfx4pvIms +c7YJyreJH018PFJsdMKWeCw6N/18ShgTGJ/zhUpYs7r4Qosqqo5+Pf6K4jeFtpuv +Qyoa7h7Y2voyVC1WIHFOzlpKdqFzyGZlu6MFxZFOzD8Y3HkCbRrwGuslFWZVH4HO +Vi3Nt/WW18vEiGdzx0AbJP0HoZPBww1JPqbha1PNH+46vpLZgRsirAIcHVVD6glQ +Ym8TGpo0ugz/HlryxQ0hbHcsE3T9cLgP4wyaH9tvJvZwA4SDJRi76EsCAwEAAaNC +MEAwHQYDVR0OBBYEFD1aTqJO87ovA2usg/v8UckSsnwhMB8GA1UdIwQYMBaAFDjv +uYaWEoViLOKyCse6mAwH66mIMA0GCSqGSIb3DQEBCwUAA4IBAQArP3mtjdj5a9uT +XTys4DUjp+CJaSmsZ/4ICf/KeddaEyijx6uoiqbNsAtfwlD5A5BtyjJrxg3SpjqD +Fm5mTV0HlwvYHxcFPTUlR56//sOyeRhm+kVPp4xdVq2p6IZXFF/hDkeT/AETWOjj +l4+90gAsJsUHrKUZmzlqHBeOXoAhzo2PwSh8KLcGK0705FGql4QnS731PHo9a9qp +e8uxk1KKoc+RHm2OjJvZ7hvkCXuCm3TJeOHQEtjZ5buQg40BRrt11GVnsPePpnQR +8OtV0DvPT3UcOss88kLDQuRuZeJf8UyaevGGvItv3+Yk5cwUBSaJ0mY6sgrmkXn0 +TstUzYIT -----END CERTIFICATE----- diff --git a/tests-integration/fixtures/etcd-tls-certs/server-key.pem b/tests-integration/fixtures/etcd-tls-certs/server-key.pem index fca8a0122f9..a73ac974434 100644 --- a/tests-integration/fixtures/etcd-tls-certs/server-key.pem +++ b/tests-integration/fixtures/etcd-tls-certs/server-key.pem @@ -1,28 +1,28 @@ -----BEGIN PRIVATE KEY----- -MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCtYFxH+oet6jTD -0eF55ISO+bwm1NPTluP9zLJrQFNDkqlddyiIasK5q7qN1pPXOw9mJF/EN8Hf1lWj -42an6tr7PgE0NHLcOrqDv4dRjegRi87IWiPkfep8peUE3glI4IOMXfZmAUNp+nRO -JP/ymmR0pzHYJlx+D082OYmeBny7uFQxo+vxgVNo2MVVwy6oTv+jtjLgZ4iOrIvh -iILoXtXcdg/tjB1rS2hSABGPJMNHwKswcrt6xXv+7Llsxroq0cx0E2PrAbc5qB40 -v1Qp7PjFTJ0drgf0g51LkOb7znpB1viPD2BLMxpELTZqa4OruarcDNxShrK8WWMP -1mygLHLxAgMBAAECggEAC4TgtLaQP62VImKGLs1QQliU1+ahiUgX70Ojog0UyyNK -JefWFVQ0ilX+z9AvI+hsYj6t7zE+LBNHPtuLtUHdGT66IUAP1pJ/VGQMB07cmZfW -ngihJFv6UZxLDkr7RnCGRPP0PDw+wKKPiiaaq8F2xapbHS+VSxnUyzdA7bMkI+uj -+gLAFj+kEU0FNbIp/mBzz8sksDzcB29oy3VnZkR1211KkL1BJDvr9oKiUg8NFKaN -NF2WXRYI/m0CQMlX6zlTyv7+h/Ay7y8oNG7cfHavmCrWvYIMfaDggYxbIpNdGxb5 -MAdbrtXJ6daUBa2sezxgQUyJX8+8GGF2zpGvCfhkfQKBgQDiAwA8rlxrBnXclwRK -9s/apILKCxPEGwU6XZi/8na3ugIO14y2WD0zZTjUK/3xwoKN6j6+ObLkUBX2RRYZ -ewwxpqos2cyxi5EBKPhsydqbJQlat68I5+NOG4KOZQkqO92tBtgY3gGFWVntaqtE -UVE/Clh2gFcc/5nR/5C5Mil8LQKBgQDEYXtrl2kl4qcyXuuXKDA/bLSnDuIEMhLO -4oIiqmywXQgjPjGAWy0eztjMq2a353pqTFCqto/cfOif+rjkB7oFfi9gMm0plIKI -+F8QsBB7HpGqWos3zb1yt8X9hbxKuFFwZS/un3Li4lwPlk8EZ1YIvhC44BHAvNkD -VzxzBwoYVQKBgArRe/RroC7bS074x5LTB5X+o+gJ6bNMW860Zjhh4b7fn3OYa7ra -tGs+YB7/0BL/bYJfgQtX9bEqCDMWkX08v5Os1553+m1RMeqtTF7gtp8Qgcce3bj+ -aIn3lSM9wNeNsAm1NyjRj58TbNOJdJM7lTkARMW/VOwla/Z6VjIXLZctAoGBAJKf -xira7eMfi36MaJJ/qyZv36Ir9ozzZh+Z91gyrtwvWfgWY5dWfCXYgv6tqw/8gOYE -/OW5UUhq6rUn2gxHyJh5Up4ciGzXOW9TIoevLV7/v/rVh8SulJimpelYhPG1FPk6 -U8NywbCtGdd5fp3nGdGFN68Rfa/OUKmx5KxtwRfRAoGAWpfzmoI8EmCBZNXYk03k -hnDKmICsKi9seDclcQq2PbyNvE7LDEd8O4pDnwV6RwKVEDTer9Y4IGcjkEBZcJVm -za3l4En9+M27Xlf1I8XeB8o8RYbFjoDQyScFFIgQpvTDsOezXeDR55FGPZN7V6Si -lKmd+8BuzJ7d31KJtRjoS1U= +MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQDQLtVfdk9l1rdH +/ssuy/BK9LpmNPx7gatp/E3EqpgekBjizMXIxnCHDFNM4Dtq1v3afaxeijXnnW/b +0CVKSeGkOBmb0dO0K21xe86Y5VfNHkGotQyCZ8mPOzep/EQOGJY5+dprErL0YKM+ +2drgiRNGwidmNk/jD3OtzXSMR+8UvDWJSxUoZ1ncvgmqf3zTTiaTiABHrfLdXaGI +23+JftyLgoeU4vEI+x+RQJMAa6Gzjo9vALxX3kTtR/uH34hfm8vb3yo68UHGKVMp +i3yi6YDP4YEao5lOpm7fIyEes10fPkD224++S34KHykAk4q+eTPetWkDYFCYzdSp +AR2Ob3cxAgMBAAECggEAMDgAxO0mw7xBVGQgFJU48WuQtvaj2kl09gXxz1ECDeYr +VXC/iNrpmmYQ7zfqmzrzrkU4hOc3SA/PplamJHhLUpmJ2Oz3P35liYj3F7PbK8/L +vnM81AGNDmdVY8Jh0u//76q+29kHaRHvDbIw/5vQQq3aqVKAG2PrU8DIM2u/5Qmt +0GHJTMmyVADAOLe7j2agN9hrOrjeoWobcELwtwgKXxmU0l+g0XX1OdJpbjA3qDPa +S5a4MzcucGINeB3bohcqbcx2f1HwVUoC19mti3Cz+KvxF2YZV9bS9izW4hSEsAwU +0upNHF3MwTUQg7ftLPAZcfNoXIZMFXwYMS1HW2v7HwKBgQDuHHkFRN2lvor/Ycjr +gJD6LGb9RlQ7NMYBQHB+5g49Y/q8oDZJgyl4CfgYWdVjjOfEX8oMst1WU/Stbva5 +HlHR7yYFyebScQYOPHZaPyVKGIbcl2p7Qy9H4pxJc1ZMNxO8SxoSC68LZEan08FD +/Qsish51fZ2Sqf37b+8Uz/CGdwKBgQDf0sIp88MPGbhtzdYtF9gA4AoIOPo9UxeI +F6wTb05uCqZyK+hcdK37av1xcNpczd/yvlALxRdM9xGedJkcS2cgeExtMRcfhAzL +as0dYPzuhRkpEIfN/y9cqtHj2wigrbu35eDgRRnJDgXgL9dYl6r+Swqn3jBaasP7 +KtQdQb/RlwKBgQC7qe0n3fLi4p4iUStNkPKyebRiAb/5Ocqkyejf2ul2MQo5B/xB +TAKu/QxwBL1NzIwOFYDlKUOQ+nJpDn+dvuu1jcpl3Y7yZOnk5npQ/luhXltMGHpv +06+79DpBGYn2X6JKUNanSlYXoFyfgSFdOF5CZifjabF7Gkd2l+3SdWCYWQKBgQDB +3k4sBGZKaB7ljUscl/CTIXvPD3tBLv3M9aQo2Vp32mW9suZ7Xt1sTonkfrnFdNWr +7shqyXabRc5PD/OnHHDhIRIh6kl7FOf4MjQkZGPxPfxDI3xeI9EkVRmkYY6hjppw +eX9FAtWI3sqcGxROOmD0Do/WQ5BiYOQMZFaCWPcLVQKBgGVT+iUYkYTs+UJeD4RC +QeL6QEJDszvAI/gavUp7sD/LWn7Ig+hJiAl2O+HFYMurh7eda0JnE2KF40UyPiYA +DDl19xRAXrLPOlOj9YKDA9HkROcaVZCQCqyDalKsbw1IqMrAkRmLCIZW3EHaMrUf +UNjg+GaxwvrRj11gbd4IPgai -----END PRIVATE KEY----- diff --git a/tests-integration/fixtures/etcd-tls-certs/server.crt b/tests-integration/fixtures/etcd-tls-certs/server.crt index 6ed8b1054e7..8dcb7b13b6d 100644 --- a/tests-integration/fixtures/etcd-tls-certs/server.crt +++ b/tests-integration/fixtures/etcd-tls-certs/server.crt @@ -1,21 +1,21 @@ -----BEGIN CERTIFICATE----- -MIIDjDCCAnSgAwIBAgIURISpu9TyXzKZTzwD2A0pQQWxGeUwDQYJKoZIhvcNAQEL +MIIDjDCCAnSgAwIBAgIURISpu9TyXzKZTzwD2A0pQQWxGecwDQYJKoZIhvcNAQEL BQAwTDELMAkGA1UEBhMCVVMxCzAJBgNVBAgMAkNBMQswCQYDVQQHDAJTRjERMA8G -A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjUwODI1MTgwNjU1 -WhcNMjYwODI1MTgwNjU1WjATMREwDwYDVQQDDAhldGNkLXRsczCCASIwDQYJKoZI -hvcNAQEBBQADggEPADCCAQoCggEBAK1gXEf6h63qNMPR4XnkhI75vCbU09OW4/3M -smtAU0OSqV13KIhqwrmruo3Wk9c7D2YkX8Q3wd/WVaPjZqfq2vs+ATQ0ctw6uoO/ -h1GN6BGLzshaI+R96nyl5QTeCUjgg4xd9mYBQ2n6dE4k//KaZHSnMdgmXH4PTzY5 -iZ4GfLu4VDGj6/GBU2jYxVXDLqhO/6O2MuBniI6si+GIguhe1dx2D+2MHWtLaFIA -EY8kw0fAqzByu3rFe/7suWzGuirRzHQTY+sBtzmoHjS/VCns+MVMnR2uB/SDnUuQ -5vvOekHW+I8PYEszGkQtNmprg6u5qtwM3FKGsrxZYw/WbKAscvECAwEAAaOBnjCB +A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjYwODI2MDQxMzQy +WhcNMjcwODI2MDQxMzQyWjATMREwDwYDVQQDDAhldGNkLXRsczCCASIwDQYJKoZI +hvcNAQEBBQADggEPADCCAQoCggEBANAu1V92T2XWt0f+yy7L8Er0umY0/HuBq2n8 +TcSqmB6QGOLMxcjGcIcMU0zgO2rW/dp9rF6KNeedb9vQJUpJ4aQ4GZvR07QrbXF7 +zpjlV80eQai1DIJnyY87N6n8RA4Yljn52msSsvRgoz7Z2uCJE0bCJ2Y2T+MPc63N +dIxH7xS8NYlLFShnWdy+Cap/fNNOJpOIAEet8t1doYjbf4l+3IuCh5Ti8Qj7H5FA +kwBrobOOj28AvFfeRO1H+4ffiF+by9vfKjrxQcYpUymLfKLpgM/hgRqjmU6mbt8j +IR6zXR8+QPbbj75LfgofKQCTir55M961aQNgUJjN1KkBHY5vdzECAwEAAaOBnjCB mzAJBgNVHRMEAjAAMAsGA1UdDwQEAwIEMDBBBgNVHREEOjA4gglsb2NhbGhvc3SC CGV0Y2QtdGxzggkxMjcuMC4wLjGHBH8AAAGHEAAAAAAAAAAAAAAAAAAAAAEwHQYD -VR0OBBYEFK8y9gWv8Si9AusP280BrgDF6y3jMB8GA1UdIwQYMBaAFOCkI3Uyx7F3 -8LtXtSlg3ORE9ur6MA0GCSqGSIb3DQEBCwUAA4IBAQBMBsAPIlk17I2ioN9xwxwl -yYbSjZlTs+18wSZAoCNLfxwWIYAQ08dPoEUdALfMUPEEe1Ol2IAp/qx2JoDrZWje -maA43hppBUpFFCkTKWUMYxsetN7d7BVxWL43GDoMwoD/k36nhoUKBUjlbF0+nkem -dOMr8SA4GZbEg7qk2cL93g0UHv/Z/dZgKf3epZkR9hEN4/R2jSP6OPmY9XIvQCoZ -8D9jgbQKavAAXmUcP5a81alQMRroGaBCzI1f5OlS3EuVE7ZTEBgxKK1idbouSrPt -4UHEbOGz9zXDjmfut6CTm247+lJm9jzYe2Xx+XGw29l0pzd/8tGFN9zIvW8mCv7a +VR0OBBYEFNqGaSuuVHwEIvQvYxt9H0kdWFPSMB8GA1UdIwQYMBaAFDjvuYaWEoVi +LOKyCse6mAwH66mIMA0GCSqGSIb3DQEBCwUAA4IBAQAHttiuyEZieeTsggAKzAIC +UXsjWCQMAmjRxery6mzM8P33CdCKA4cwGyy6tzoFPiZC/wh+gPq+LiUO9G+vkqNA +NQ/DSJS+fsqsOBWUt5DSYekD25o80mZVE/4H/ynFOgENBqgkVl1ikpkXu9BK8wTm +Wij+p9RWyNup8lr9OFYZRpy7wrxPp7opwp3qG2Bfu/BGY1YEr04CIrLMKnk1zRMm +73X1lcrlrM9ycCmYcfXxyU69G0iFKaPcbS7T0fngm1NciA1qxYYLuMyzDtF7/SyC +J2ynvDBBnc9+/X8piu8p9D0I9mbcbEySv6zZfYMT0kqXbsn2GijZTH5mCJ3P8cwc -----END CERTIFICATE----- From 35ea88a4ef4b06c8c5d10e57cc8bad44f4f38569 Mon Sep 17 00:00:00 2001 From: localhost Date: Wed, 26 Aug 2026 12:11:14 +0000 Subject: [PATCH 04/22] feat(ci): run query regression on ephemeral Aliyun ECS runners (#8937) * feat(ci): add aliyun ecs ephemeral runner path for query regression Signed-off-by: paomian * fix: improve condition for query-regression job execution in workflow * feat: update Docker installation to use official repository and add GPG key handling * Refactor query regression runner setup and configuration - Removed deprecated PersistentVolumeClaim for build cache. - Introduced a new bootstrap script for setting up the ECS runner host. - Deleted obsolete Helm values files for runner configuration. - Updated the Aliyun ECS runner provisioning script to reflect new cache paths. - Modified GitHub workflows to use the new Aliyun ECS runner setup. - Adjusted documentation to clarify the new runner lifecycle and provisioning process. * fix: enhance runner service management during bootstrap process * fix: update alibabacloud_tea_openapi dependency version in metadata * feat: enhance ECS runner scripts with region_id and resource_group_id support * fix: move containerd content store to data root for improved storage management * feat: rename query-regression runner to ephemeral-github runner and update related scripts * fix: update sentinel polling method to use serial console output for improved reliability * fix: add environment variable checks for Alibaba Cloud access keys in ECS client * fix: improve error handling in GitHub API requests for better diagnostics * fix: improve cache disk detection logic for Aliyun ECS instances * fix: enhance cache disk waiting logic with detailed output and error handling * fix: update dependency version for alibabacloud_tea_openapi in teardown script * fix: enhance cache disk waiting logic for better compatibility and clarity * fix: enhance console output handling and add incremental logging during instance provisioning * fix: add PATH environment variable for runner jobs in service and provision script * fix: add machine telemetry sampling and logging during query regression jobs * fix: update query regression documentation and provision script for cache disk handling * fix: update SCCACHE_CACHE_SIZE validation to 10G for improved caching efficiency * fix: remove outdated cache size checks and cleanup logic for fresh system disk runs * fix: enhance instance deletion logic with region handling and console output export * fix: add swap file setup and OOM handling for ECS runner to improve stability * fix: update OOM handling and service restart logic for ECS runner to enhance stability * fix: increase system disk size to 100 GiB for cold double nightly builds to prevent ENOSPC errors * fix: increase system disk size to 150 GiB for ECS runner to prevent ENOSPC errors * fix: add keep_instance option to preserve ECS instance for post-mortem debugging * fix: disable unattended upgrades to prevent job cancellations during library updates * fix: reduce system disk size to 40 GiB for ECS runner to prevent ENOSPC errors * feat: Refactor Aliyun ECS runner provisioning and introduce nightly regression comparison - Update `aliyun-ecs-runner-provision.py` to remove cache disk handling, simplifying the provisioning process. - Introduce `query-regression-nightly-refs.py` to resolve and compare SHAs from successful nightly builds. - Create `query-regression-nightly.yml` workflow to trigger nightly comparisons based on successful builds. - Enhance `query-regression.yml` to include a `test-tooling` job for validating Python scripts before provisioning. - Update tests for the new nightly reference selection logic and refactor existing tests to align with the new caching strategy. - Modify documentation to reflect changes in caching and nightly comparison workflows. * fix: enhance runner image tool verification with detailed checks * fix: improve error handling in runner image tool verification * fix: update tool versions in ECS image and workflow for consistency * fix: correct typo in error message for unparseable ECS creation time * fix: update README and workflow files for query regression tests and image hygiene --------- Signed-off-by: paomian --- .../query-regression/README.md | 453 ++++++---------- .../query-regression/cache-pvc.yaml | 15 - .../ecs-image/bootstrap-runner-host.sh | 195 +++++++ .../ecs-image/build-ecs-image.py | 392 ++++++++++++++ .../ecs-image/ephemeral-github-runner.service | 45 ++ .../ecs-image/start-runner.sh | 37 ++ .../query-regression/values-8-cores.yaml | 146 ------ .../query-regression/values-paused.yaml | 2 - .../scripts/aliyun-ecs-runner-provision.py | 485 ++++++++++++++++++ .github/scripts/aliyun-ecs-runner-teardown.py | 237 +++++++++ .../scripts/query-regression-nightly-refs.py | 270 ++++++++++ .github/workflows/checks.yml | 8 + .../workflows/query-regression-janitor.yml | 36 ++ .../workflows/query-regression-nightly.yml | 104 ++++ .github/workflows/query-regression.yml | 433 +++++++++++----- .github/workflows/release.yml | 2 +- tests/perf/AGENTS.md | 20 +- tests/perf/README.md | 7 +- tests/perf/test_aliyun_ecs_runner_scripts.py | 122 +++++ .../test_query_regression_nightly_refs.py | 120 +++++ 20 files changed, 2538 insertions(+), 591 deletions(-) delete mode 100644 .github/runner-scale-sets/query-regression/cache-pvc.yaml create mode 100644 .github/runner-scale-sets/query-regression/ecs-image/bootstrap-runner-host.sh create mode 100644 .github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py create mode 100644 .github/runner-scale-sets/query-regression/ecs-image/ephemeral-github-runner.service create mode 100644 .github/runner-scale-sets/query-regression/ecs-image/start-runner.sh delete mode 100644 .github/runner-scale-sets/query-regression/values-8-cores.yaml delete mode 100644 .github/runner-scale-sets/query-regression/values-paused.yaml create mode 100644 .github/scripts/aliyun-ecs-runner-provision.py create mode 100644 .github/scripts/aliyun-ecs-runner-teardown.py create mode 100644 .github/scripts/query-regression-nightly-refs.py create mode 100644 .github/workflows/query-regression-janitor.yml create mode 100644 .github/workflows/query-regression-nightly.yml create mode 100644 tests/perf/test_aliyun_ecs_runner_scripts.py create mode 100644 tests/perf/test_query_regression_nightly_refs.py diff --git a/.github/runner-scale-sets/query-regression/README.md b/.github/runner-scale-sets/query-regression/README.md index a1d58de7166..ce447dbdd0a 100644 --- a/.github/runner-scale-sets/query-regression/README.md +++ b/.github/runner-scale-sets/query-regression/README.md @@ -1,321 +1,184 @@ # Query regression self-hosted runners -The `Query Regression` workflow uses the dedicated ARC runner scale set -`perf-regression-8-cores`. ARC runner Pods run in the target Kubernetes cluster -and connect outbound to GitHub. The live scale set is currently **paused**: -`minRunners=0`, `maxRunners=0`, and no runner Pods. Do not resume it without -explicit approval. +The `Query Regression` workflow runs on **Aliyun ECS ephemeral runners** +(`aliyun-ecs`, the default and only automated path): a `provision` job on +`ubuntu-latest` creates one pay-as-you-go ECS instance per run, and the +instance registers itself as an ephemeral GitHub runner with a per-run +label. A `teardown` job (`if: always()`) deletes the instance; +`query-regression-janitor.yml` sweeps tagged leftovers older than 4 hours +daily. -## Prerequisites and trust admission +Build caches live on the instance's system disk, so **every run compiles +cold**; only the within-run reuse (base build warms the candidate build +through the shared target dir and sccache) applies. There is no retained +data disk. -Install the ARC scale set controller if it is not already installed: +Dispatching with any other `runner` value uses it as a literal self-hosted +runner label, which is how a manually prepared host (see +`ecs-image/bootstrap-runner-host.sh`) runs the workflow. For PR labels, set +the repository variable `QUERY_REGRESSION_PR_RUNNER` to such a label to +redirect PR runs away from ECS. + +The office ARC scale set `perf-regression-8-cores` that previously ran this +workflow is retired; see git history for its values files and pause/deploy +procedures. Tearing down the office cluster (helm release, runner namespace, +and the `query-regression-build-cache` PVC) is a manual operator action +outside this repository. This directory keeps its historical +`runner-scale-sets` name for path stability. + +## Nightly vs previous nightly + +`query-regression-nightly.yml` runs after a successful `GreptimeDB Nightly +Build` (`workflow_run`). It resolves that run's `head_sha` as the candidate +and the previous successful nightly (same branch, typically Friday when +Monday's nightly fires) as the base, then calls `query-regression.yml` with +those immutable SHAs. Both binaries are still compiled on the ECS runner; +this is not an artifact-download path. `workflow_dispatch` can pass explicit +`base_ref` / `candidate_ref` or a nightly run id. If there is no previous +nightly, or both nightlies built the same commit, the comparison is skipped. + +## Aliyun ECS path + +Configuration lives in repository variables/secrets: + +| Kind | Name | Purpose | +| --- | --- | --- | +| secret | `ALICLOUD_ECS_ACCESS_KEY_ID` / `ALICLOUD_ECS_ACCESS_KEY_SECRET` | RAM user scoped to ECS RunInstances/DeleteInstances/Describe*/CreateImage/RunCommand. Used only by provision/teardown jobs on `ubuntu-latest`; never reaches the ECS instance. | +| secret | `GH_PERSONAL_ACCESS_TOKEN` | Creates the short-lived runner registration token (shared with the jsonbench EC2 path). | +| vars | `ALIYUN_ECS_REGION_ID` / `ALIYUN_ECS_VSWITCH_ID` / `ALIYUN_ECS_SECURITY_GROUP_ID` | Network placement. The security group should allow egress only; no inbound rules are needed. The vSwitch pins the zone. | +| vars | `ALIYUN_ECS_INSTANCE_TYPE` | Dedicated (non-burstable, non-shared) instance family. Prefer 32 GiB (e.g. `ecs.g8i.2xlarge`); `ecs.c9i.2xlarge` is 8c16g and nightly thin-LTO of greptime peaks above that. Both base and candidate clusters run on the same machine, so noisy neighbors break thresholds. | +| vars | `QUERY_REGRESSION_ECS_IMAGE_ID` | Custom image built by `ecs-image/build-ecs-image.py`. | + +The system disk is 40 GiB, which covers the image, a 16 GiB swapfile, the +checkout, and cold build caches (target dir, cargo registry, sccache). ENOSPC +stops the runner itself from writing logs, which GitHub reports as `The +operation was canceled` with no telemetry, indistinguishable from a +platform-side cancellation. +cloud-init masks `systemd-oomd`, disables `unattended-upgrades` / +`apt-daily-upgrade`, creates `/swapfile`, and sets `OOMPolicy=continue` +on the runner unit. Ubuntu 24.04 defaults to `DefaultOOMPolicy=stop`, +which SIGTERM-s the whole unit when rustc is OOM-killed and GitHub +reports `The operation was canceled` with no telemetry. Unattended +upgrades can do the same via `systemctl restart` of the runner after a +library update; GitHub then records `UserCancelled` even though the job +was still valid. Swap is a safety net for 16 GiB types, not a substitute +for 32 GiB; linking on swap is slow. + +When a run dies with an unexplained `The operation was canceled` (no +"Canceled by" banner, healthy machine), re-dispatch with `keep_instance` +checked: the teardown job is skipped and the instance survives for +post-mortem inspection. The security group is egress-only, so inspect via +Cloud Assistant (`RunCommand`) or VNC: the runner's `_diag` logs under the +runner home record reconnects and worker crashes, `journalctl -u +ephemeral-github-runner.service` mirrors the console stream, `dmesg -T` +shows kernel OOM kills, and `machine-telemetry.log` in the job workspace +has the 30 s sampler history. The janitor sweep still deletes the instance +after its TTL, so finish the inspection within that window. + +Trust model on the ECS path: the instance receives only the one-hour runner +registration token via user data and holds no cloud credentials; the Aliyun +AK/SK exist only in the control-plane jobs. Instance tags +(`managed-by=query-regression-ci`, `query-regression-run-id`) feed the janitor +sweep. + +### Building and updating the ECS image + +The runner `Dockerfile` in the parent directory stays the single source of the +tool contract. Build a new ECS image from it: ```bash -helm upgrade --install arc \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set-controller \ - --namespace arc-systems \ - --create-namespace \ - --version 0.14.2 +ALIBABA_CLOUD_ACCESS_KEY_ID=... ALIBABA_CLOUD_ACCESS_KEY_SECRET=... \ +uv run .github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py \ + --region-id --vswitch-id --security-group-id \ + --base-image-id ``` -Create the GitHub App secret in the runner namespace. Prefer an App limited to -`GreptimeTeam/greptimedb`: +The script boots a temporary builder instance, `docker build`s the runner +image, materializes `/opt/rustup`, `/opt/cargo`, `/usr/local/bin` tools, and +`/home/runner` (actions-runner) onto the host, installs the ephemeral-runner +systemd unit from `ecs-image/`, and snapshots a custom image. It prints the +image id; set it as `QUERY_REGRESSION_ECS_IMAGE_ID`, and bump +`RUNNER_IMAGE_EPOCH` in `query-regression.yml` at the same time so the target +cache invalidates. Builder sentinel polling requires the Cloud Assistant +agent, which Aliyun public Ubuntu images include. -```bash -kubectl -n arc-runners create secret generic greptimedb-arc-github-app \ - --from-literal=github_app_id= \ - --from-literal=github_app_installation_id= \ - --from-file=github_app_private_key= -``` - -The values files here reference that secret by name. +## Trust admission for PR runs A maintainer applying the `query-regression` or `heavy-regression` label is -**trust admission for that exact PR revision**. `query-regression` runs the six -routine default cases; `heavy-regression` runs only the high-cardinality -`prom_remote_write_7913` remote-write case. The admitted job may use this scale -set's dedicated, writable persistent cache. `pull_request: labeled` is the only -PR trigger: the label event snapshots its merge, head, and base SHAs. A queued -job fetches that immutable event merge SHA directly, verifies it is a two-parent -merge whose parents include the snapshotted head exactly once, and uses its -other parent as the actual base build revision. The snapshotted event base is -retained for audit only, so a difference from the merge's non-head parent is -not a failure. The job never follows a newer mutable PR merge ref. An -unavailable event merge, or one that does not contain exactly one snapshotted -head parent, fails closed. A later PR head change does not retarget an already -queued run: it may execute only its previously trusted event revision if that -revision remains fetchable. To run the new revision, the maintainer must review -it, remove the label, and re-add the desired regression label; cancel the old -run if it is no longer wanted. An existing label does not automatically rerun -the benchmark. +**trust admission for that exact PR revision**. `query-regression` runs the +six routine default cases; `heavy-regression` runs only the high-cardinality +`prom_remote_write_7913` remote-write case. `pull_request: labeled` is the only PR +trigger: the label event snapshots its merge, head, and base SHAs. A queued +job fetches that immutable event merge SHA directly, verifies it is a +two-parent merge whose parents include the snapshotted head exactly once, and +uses its other parent as the actual base build revision. The snapshotted +event base is retained for audit only, so a difference from the merge's +non-head parent is not a failure. The job never follows a newer mutable PR +merge ref. An unavailable event merge, or one that does not contain exactly +one snapshotted head parent, fails closed. A later PR head change does not +retarget an already queued run: it may execute only its previously trusted +event revision if that revision remains fetchable. To run the new revision, +the maintainer must review it, remove the label, and re-add the desired +regression label; cancel the old run if it is no longer wanted. An existing +label does not automatically rerun the benchmark. -Admission does not relax runner hardening or GitHub permissions. Keep -service-account token mounting disabled; do not mount host paths, the Docker -socket, kubeconfig, or long-lived credentials. The runner and cache initializer -use UID/GID 1001, disallow privilege escalation, drop all capabilities, and use -the RuntimeDefault seccomp profile. Keep GitHub tokens least-privilege and -review workflow changes before admission. Where the CNI supports it, restrict -egress to required GitHub Actions, artifact/cache, Rust/crate/toolchain, DNS, -and image-registry endpoints; block unrelated cluster services, private ranges, -and metadata endpoints unless a case requires them. - -### Network routing prerequisite - -Required split routing is an **external environment-specific prerequisite**. The -responsible network operator must route GitHub Actions, GitHub content, -artifact/cache, crates.io, Rust toolchain, and image-registry traffic through -the approved path rather than the VPN where applicable. Neither this repository -nor Kubernetes configures that route. Verify it with the responsible network -operator before any canary. +Admission does not relax runner hardening or GitHub permissions. Keep the ECS +instance free of cloud credentials and long-lived tokens, keep the security +group egress-only, keep GitHub tokens least-privilege, and review workflow +changes before admission. ## Runner image and workflow tools -Build and push the derived runner image; it preserves the official -`/home/runner/run.sh` entrypoint and supplies CI tools needed at runtime. The -image builds `otelgen` from +The runner `Dockerfile` builds `otelgen` from [`WenyXu/otelgen`](https://github.com/WenyXu/otelgen) commit [`863a3f395d062c7322cc1de08a38774b7fdaa6c8`](https://github.com/WenyXu/otelgen/commit/863a3f395d062c7322cc1de08a38774b7fdaa6c8) -so trace cases do not download or compile tools during a benchmark run: +so trace cases do not download or compile tools during a benchmark run. It is +built only as a toolchain factory: `build-ecs-image.py` and +`bootstrap-runner-host.sh` both materialize its contents onto a host, and the +benchmark itself runs host-native. -```bash -docker build \ - --platform linux/amd64 \ - -f .github/runner-scale-sets/query-regression/Dockerfile \ - -t greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner:latest \ - .github/runner-scale-sets/query-regression - -docker push greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner:latest -``` - -Deploy by digest, not mutable tag, by updating both image references in -`values-8-cores.yaml` after a rebuild. Update `RUNNER_IMAGE_DIGEST` and bump -`RUNNER_IMAGE_EPOCH` in `query-regression.yml` at the same time. If the registry -is private, use a dedicated read-only pull secret only as `imagePullSecrets`; -never expose registry credentials to runner containers. -Both digest-pinned init and runner containers use `IfNotPresent`: the immutable -digest makes a cached image safe and avoids adding a registry dependency to every -runner startup. - -The runner optionally imports only the non-sensitive `HTTP_PROXY`, -`HTTPS_PROXY`, and `NO_PROXY` variables from the -`query-regression-runner-local-env` ConfigMap. Manage that ConfigMap locally in -the target namespace; private endpoint configuration must not be committed, and -credentials or secrets must never be placed in a ConfigMap. - -Before builds, the workflow asserts UID/GID 1001 and exact image tool versions: -`libprotoc 3.21.12`, `uv 0.11.26`, `mold 2.30.0`, `Python 3.12.3`, `sccache -0.16.0`, `otelgen` commit `863a3f395d062c7322cc1de08a38774b7fdaa6c8`, -root-owned `rustup 1.29.0`, and the image-baked -`nightly-2026-03-21` Rust toolchain. Rustup, Cargo, and Rustc must resolve from -`/opt/cargo/bin`; the runner cannot write `/opt/rustup` or `/opt/cargo/bin`. -Protobuf well-known includes, including `google/protobuf/any.proto` and -`google/protobuf/empty.proto`, are an image contract and must compile with -`protoc`. +Before builds, the workflow asserts the runner UID/GID (1001 in the ECS image; +overridable via `QUERY_REGRESSION_RUNNER_UID`/`QUERY_REGRESSION_RUNNER_GID`) +and exact tool versions: `libprotoc 3.21.12`, `uv 0.11.26`, `mold 2.40.4`, +`Python 3.14.4`, `sccache 0.16.0`, `otelgen` commit +`863a3f395d062c7322cc1de08a38774b7fdaa6c8`, root-owned `rustup 1.29.0`, and +the image-baked `nightly-2026-03-21` Rust toolchain. `mold` and `python3` +come from apt at image-build time (not Ubuntu 24.04's default 3.12); bump +the Verify pins together with `QUERY_REGRESSION_ECS_IMAGE_ID` when the +image is rebuilt. Rustup, Cargo, and Rustc +must resolve from `/opt/cargo/bin`; the runner cannot write `/opt/rustup` or +`/opt/cargo/bin`. Protobuf well-known includes, including +`google/protobuf/any.proto` and `google/protobuf/empty.proto`, are an image +contract and must compile with `protoc`. `actions-rust-lang/setup-rust-toolchain@v1` is intentionally removed. The -workflow sets its warning-denying mold `RUSTFLAGS` directly, disables automatic -Rustup installation, and performs no runtime toolchain downloads. +workflow sets its warning-denying mold `RUSTFLAGS` directly, disables +automatic Rustup installation, and performs no runtime toolchain downloads. -The workflow no longer uses GitHub `rust-cache`, `setup-protoc`, `setup-uv`, or -runtime Rust setup: the image establishes immutable executable state and the PVC -supplies only reusable Cargo data. Do not reintroduce those actions unless the -corresponding cache or image contract changes. +The workflow no longer uses GitHub `rust-cache`, `setup-protoc`, `setup-uv`, +or runtime Rust setup: the image establishes immutable executable state and +each instance's system disk holds only that run's Cargo data. Do not +reintroduce those actions unless the image contract changes. -## Capacity and persistent cache +## Capacity -`values-8-cores.yaml` is normal operation: `minRunners=0`, `maxRunners=1`. -`values-paused.yaml` is the mandatory pause overlay: `minRunners=0`, -`maxRunners=0`. The job uses group `query-regression-persistent-cache-v1`, -`queue: max`, and `cancel-in-progress: false`; admitted jobs queue rather than -replacing older pending jobs. During maintenance, cancel admitted queued runs as -well as pausing ARC. Runner Pods have `activeDeadlineSeconds=12600`. -The runner requests 6 CPU and limits at 8 CPU to preserve `minipc-3` -allocatable-capacity scheduling headroom; do not reset the request to 8 CPU -without revalidating scheduling capacity. +Each run provisions its own ECS instance, so overlapping dispatches proceed +in parallel. There is no workflow `concurrency` group. All Cargo state +(`CARGO_HOME` including registry/git, `CARGO_TARGET_DIR`, sccache, cache +metadata) lives on the 40 GiB system disk and is discarded with the VM. +`RUSTUP_HOME=/opt/rustup` and `/opt/cargo/bin` are image-owned. The runner +sets `RUSTC_WRAPPER=/usr/local/bin/sccache`, +`SCCACHE_DIR=/home/runner/.cache/sccache`, `SCCACHE_CACHE_SIZE=10G`, and +`CARGO_INCREMENTAL=0`. sccache uses its local disk backend and self-evicts +at 10G. Base and candidate builds share the target on that disk; Cargo +fingerprints invalidate source and dependency changes. -The cache claim `query-regression-build-cache` is a nominal 600Gi `local-path` -PVC in `arc-runners`. It is `ReadWriteOnce`; `local-path` uses -WaitForFirstConsumer binding and Delete reclaim behavior, produces a -node-affine local PV, is non-expandable, and the 600Gi request is not a hard -storage quota. The runner's `minipc-3` selector is its only consumer candidate. - -The initializer mounts the PVC root at `/cache`, creates and write-tests these -versioned subpaths as non-root UID/GID 1001, and the runner mounts them as: - -| Persistent state | PVC subpath | Runner mount | -| --- | --- | --- | -| Ephemeral Cargo home | `emptyDir` | `/home/runner/.cargo` | -| Cargo registry data | `cargo-registry-v1` | `/home/runner/.cargo/registry` | -| Cargo Git data | `cargo-git-v1` | `/home/runner/.cargo/git` | -| Cargo target | `query-regression-target-v1` | `/home/runner/query-regression-target` | -| Cache metadata | `meta-v1` | `/home/runner/query-regression-cache-meta` | -| sccache local disk cache | `sccache-v1` | `/home/runner/.cache/sccache` | -| Immutable Rust toolchain | image-owned | `/opt/rustup`, `/opt/cargo/bin` | - -The Pod security context uses UID/GID and `fsGroup` 1001 with -`fsGroupChangePolicy: OnRootMismatch`; no privileged `chown` or raw `hostPath` -is used. `CARGO_HOME` is a per-Pod `emptyDir`; only its nested `registry` and -`git` mounts are persistent. `RUSTUP_HOME=/opt/rustup` and `/opt/cargo/bin` are -image-owned immutable paths, while `CARGO_TARGET_DIR`, cache metadata, and -`SCCACHE_DIR` are persistent absolute paths. The runner sets -`RUSTC_WRAPPER=/usr/local/bin/sccache`, -`SCCACHE_DIR=/home/runner/.cache/sccache`, `SCCACHE_CACHE_SIZE=40G`, and -`CARGO_INCREMENTAL=0`. sccache uses its local PVC disk backend and self-evicts -at 40G; do not add runtime downloads, object storage, or a shared backend. - -The repository's `.cargo/config.toml` remains a trusted per-revision build input. -In contrast, `$CARGO_HOME/config*`, credentials, installed bins, and Cargo -metadata outside the persistent `registry` and `git` data mounts are ephemeral -and cannot survive to another Pod. - -The local disk backend has a one-server constraint. `maxRunners=1` and the -unchanged `query-regression-persistent-cache-v1` workflow concurrency group -serialize runs; do not increase runner capacity or relax that serialization -while this backend is in use. Base and candidate builds share the target; Cargo -fingerprints invalidate source and dependency changes. The workflow records the -sccache version and relevant environment in the target ABI marker, starts and -zeros sccache after cache and toolchain checks, shows initial/base/candidate -statistics, and resets statistics between base and candidate builds. - -### Disk preflight and cleanup contract - -Before applying or unpausing, verify the backing filesystem on `minipc-3` has -at least 900GiB free. The current local-path provisioner source is -`/opt/local-path-provisioner`; measure the filesystem containing it: - -```bash -df -PB1G /opt/local-path-provisioner -``` - -The workflow reports `du`, `df -P`, human-readable free space, and inode -availability before builds and in an always-run report. Its cleanup is narrow -and non-destructive: - -- warn at target size 400GiB; at 450GiB clear only the complete target root; -- warn at Cargo registry-plus-Git data size 60GiB; at 80GiB remove only - `registry/src` and `git/checkouts`, then abort if that persistent data remains - at least 80GiB; -- below 300GiB backing free space, clear the complete target root first, - remeasure, then remove only those Cargo extracted trees and checkouts; abort - if free space is still below 300GiB; -- never automatically remove Cargo registry cache/index, Git database, the - image-owned Cargo bin or Rustup toolchain, cache metadata, the self-evicting - sccache directory, or the PVC. - -The target clear uses fixed absolute roots and removes all entries, including -dotfiles. After migration validation, remove obsolete `cargo-home-v1` and -`rustup-home-v1` only in explicit maintenance while ARC is 0/0 and no runner Pod -exists; they are not mounted by the current configuration. - -## Deploy and pause safely - -First verify the configured external network route and the disk preflight. Apply the PVC; while -the scale set is paused, it is expected to remain `Pending` because -WaitForFirstConsumer has no scheduled runner: - -```bash -kubectl apply --dry-run=server \ - -f .github/runner-scale-sets/query-regression/cache-pvc.yaml -kubectl apply -f .github/runner-scale-sets/query-regression/cache-pvc.yaml -``` - -Render normal and paused configurations. Normal values are always first; the -pause overlay is always last: - -```bash -helm template perf-regression-8-cores \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ - --namespace arc-runners --version 0.14.2 \ - --set controllerServiceAccount.name=arc-gha-rs-controller \ - --set controllerServiceAccount.namespace=arc-systems \ - -f .github/runner-scale-sets/query-regression/values-8-cores.yaml - -helm template perf-regression-8-cores \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ - --namespace arc-runners --version 0.14.2 \ - --set controllerServiceAccount.name=arc-gha-rs-controller \ - --set controllerServiceAccount.namespace=arc-systems \ - -f .github/runner-scale-sets/query-regression/values-8-cores.yaml \ - -f .github/runner-scale-sets/query-regression/values-paused.yaml -``` - -The **first post-merge Helm deployment must reconcile the release in paused -mode**. Keep the pause overlay last: - -```bash -# First post-merge deployment and every return to paused mode: 0/0. -helm upgrade --install perf-regression-8-cores \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ - --namespace arc-runners --create-namespace --version 0.14.2 \ - --reset-values --wait \ - -f .github/runner-scale-sets/query-regression/values-8-cores.yaml \ - -f .github/runner-scale-sets/query-regression/values-paused.yaml - -# Expect 0/0 and no runner resources before considering normal mode. -kubectl -n arc-runners get autoscalingrunnerset perf-regression-8-cores \ - -o jsonpath='{.spec.minRunners}{"/"}{.spec.maxRunners}{"\n"}' -kubectl -n arc-runners get ephemeralrunners,pods \ - -l actions.github.com/scale-set-name=perf-regression-8-cores -``` - -Only after that verification and separate explicit approval, apply normal 0/1 -operation without the pause overlay: - -```bash -helm upgrade --install perf-regression-8-cores \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ - --namespace arc-runners --create-namespace --version 0.14.2 \ - --reset-values --wait \ - -f .github/runner-scale-sets/query-regression/values-8-cores.yaml -``` - -Do not use bare `helm rollback`, `--atomic`, or `--reuse-values`: a stored -revision can restore nonzero runner capacity. Inspect rendered manifests for -capacity, the `minipc-3` selector, cache claim and mounts, initializer, -security context, and resources. After approved normal mode receives its first -canary, the PVC binds to `minipc-3`. - -For local-PV node loss, cache recovery is intentionally disposable: return to -0/0, recreate the PVC on a healthy node, cold-fill it, and run a new canary. - -## Canary and rollback - -With explicit approval, run two identical `workflow_dispatch` canaries on -`perf-regression-8-cores`, using immutable full base and candidate commit SHAs -and `cargo_profile=nightly`. The first is the cold fill; the second verifies warm -reuse. Record the workflow's base/candidate build elapsed logs and cache -ABI-marker output, initial/base/candidate sccache statistics, and cache report. -Confirm the image tool contract (root-owned Rustup/Cargo paths, baked nightly, -and non-writable `/opt` roots) and that the ephemeral Cargo home contains only -the mounted registry/Git data before Cargo creates per-Pod state. -Obtain dependency and tool network byte counters from the configured -environment counter source, filtered to `minipc-3` and the dependency/tool -destinations. - -Accept the canary only when all of the following hold: - -- exactly one runner Pod runs on `minipc-3`, and both jobs use the same bound PV; -- UID/GID 1001 cache mounts are writable; the warm run does not invalidate the - target or bulk-redownload crates or toolchains; sccache reports separate base - and candidate build statistics without server or cache-path errors; immutable - Rustup/Cargo roots remain non-writable and only registry/Git data persists; -- warm base build time is at most 50% of cold base build time; -- warm dependency/tool network bytes are at most 10% of cold fill bytes; -- cache sizes remain below soft watermarks, node free space remains at least - 300GiB, and the benchmark is correct without TLS EOFs or timeouts; -- the configured environment counter source confirms this traffic is outside VPN - accounting. - -Immediately return to 0/0 after either canary unless ongoing normal operation -has been explicitly approved; return immediately on any traffic, cache, disk, -TLS, or correctness failure. To roll back, use the paused Helm upgrade above, -or another explicit `helm upgrade --install` with known-good values followed by -`values-paused.yaml`, `--reset-values`, and `--wait`. Do not delete the PVC -automatically; preserve it for diagnosis unless intentionally discarding cache. +The workflow reports `du` / `df` in telemetry. It does not try to reclaim +space across runs: a cold 40 GiB disk that fills up fails the build. ## Future optional phases -The current phase uses a digest-pinned image with sccache 0.16.0 and no shared -cache service. Optional follow-ups are an image additionally seeded with the -exact Rust toolchain and `cargo fetch --locked`; or an internal read/write -sccache backend or Cargo/Git mirror. Evaluate them only if persistent PVC reuse -is insufficient. +The current phase uses a materialized runner toolchain with sccache 0.16.0 and +no shared cache service. Optional follow-ups are an image additionally seeded +with `cargo fetch --locked` results, or an internal read/write sccache +backend or Cargo/Git mirror. Evaluate them only if cold compile time on the +system disk becomes the bottleneck. diff --git a/.github/runner-scale-sets/query-regression/cache-pvc.yaml b/.github/runner-scale-sets/query-regression/cache-pvc.yaml deleted file mode 100644 index 692f3a2d6e2..00000000000 --- a/.github/runner-scale-sets/query-regression/cache-pvc.yaml +++ /dev/null @@ -1,15 +0,0 @@ -apiVersion: v1 -kind: PersistentVolumeClaim -metadata: - name: query-regression-build-cache - namespace: arc-runners - labels: - app.kubernetes.io/part-of: perf-regression-8-cores -spec: - volumeMode: Filesystem - accessModes: - - ReadWriteOnce - storageClassName: local-path - resources: - requests: - storage: 600Gi diff --git a/.github/runner-scale-sets/query-regression/ecs-image/bootstrap-runner-host.sh b/.github/runner-scale-sets/query-regression/ecs-image/bootstrap-runner-host.sh new file mode 100644 index 00000000000..883bbede092 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/ecs-image/bootstrap-runner-host.sh @@ -0,0 +1,195 @@ +#!/usr/bin/env bash +# 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. + +# Bootstrap an Ubuntu 24.04 host as a query-regression self-hosted runner. +# +# This is the host-native equivalent of the ECS image contract: it +# installs docker-ce from Docker's official repository, builds the runner +# image from the sibling Dockerfile (the single source of the tool contract), +# materializes the toolchain onto the host, and optionally registers the +# runner and installs its systemd service. +# +# Large state lives under --data-root (a data disk): docker's data-root, the +# runner home, and the rustup/cargo toolchain roots, symlinked back to the +# hard-coded contract paths (/home/runner, /opt/rustup, /opt/cargo). +# +# The script is idempotent: re-running refreshes the toolchain in place, and +# the runner registration (.runner/.credentials) and _work survive. A live +# runner service is stopped first and restarted at the end. +# +# Usage: +# sudo bash bootstrap-runner-host.sh --repo-dir /path/to/greptimedb +# sudo RUNNER_TOKEN= bash bootstrap-runner-host.sh --repo-dir . \ +# --register --repo-url https://github.com// +# +# Options: +# --repo-dir PATH greptimedb checkout containing the runner Dockerfile (required) +# --data-root PATH data disk mount point (default: /data) +# --uid / --gid N runner user id (default: 3141; the ECS image contract is 1001) +# --register also register the runner and install the systemd service +# --repo-url URL repository URL for registration (required with --register) +# --runner-name N runner name (default: qreg-host) +# --labels L runner labels (default: perf-regression-8-cores) +# +# With --register, provide a fresh registration token via RUNNER_TOKEN. + +set -euo pipefail + +REPO_DIR="" +DATA_ROOT="/data" +RUNNER_UID="3141" +RUNNER_GID="3141" +REGISTER="false" +REPO_URL="" +RUNNER_NAME="qreg-host" +RUNNER_LABELS="perf-regression-8-cores" +IMAGE_TAG="qreg-runner:manual" + +while [[ $# -gt 0 ]]; do + case "$1" in + --repo-dir) REPO_DIR="$2"; shift 2 ;; + --data-root) DATA_ROOT="$2"; shift 2 ;; + --uid) RUNNER_UID="$2"; shift 2 ;; + --gid) RUNNER_GID="$2"; shift 2 ;; + --register) REGISTER="true"; shift ;; + --repo-url) REPO_URL="$2"; shift 2 ;; + --runner-name) RUNNER_NAME="$2"; shift 2 ;; + --labels) RUNNER_LABELS="$2"; shift 2 ;; + -h | --help) sed -n '17,40p' "$0"; exit 0 ;; + *) echo "Unknown argument: $1" >&2; exit 2 ;; + esac +done + +[[ $EUID -eq 0 ]] || { echo "Run as root, e.g. sudo bash $0 ..." >&2; exit 1; } +[[ -n "${REPO_DIR}" ]] || { echo "--repo-dir is required" >&2; exit 2; } +DOCKERFILE="${REPO_DIR}/.github/runner-scale-sets/query-regression/Dockerfile" +[[ -f "${DOCKERFILE}" ]] || { echo "Dockerfile not found at ${DOCKERFILE}" >&2; exit 1; } +if [[ "${REGISTER}" == "true" ]]; then + [[ -n "${REPO_URL}" ]] || { echo "--repo-url is required with --register" >&2; exit 2; } + [[ -n "${RUNNER_TOKEN:-}" ]] || { echo "RUNNER_TOKEN env is required with --register" >&2; exit 2; } +fi + +step() { printf '\n==> %s\n' "$*"; } + +# A re-run refreshes files under the (possibly live) runner service, so stop +# it first and restart it at the end (unless --register re-creates it). +RUNNER_SERVICE="$(systemctl list-units --type=service --all --no-legend 'actions.runner.*' 2>/dev/null | awk '{print $1}' | head -n 1 || true)" +RUNNER_WAS_ACTIVE="false" +if [[ -n "${RUNNER_SERVICE}" ]] && systemctl is-active --quiet "${RUNNER_SERVICE}"; then + step "Stop running runner service ${RUNNER_SERVICE}" + systemctl stop "${RUNNER_SERVICE}" + RUNNER_WAS_ACTIVE="true" +fi + +step "Install docker-ce from Docker's official repository" +for pkg in docker.io docker-doc docker-compose docker-compose-v2 podman-docker containerd runc; do + apt-get remove -y "${pkg}" 2>/dev/null || true +done +apt-get update +apt-get install -y ca-certificates curl +install -m 0755 -d /etc/apt/keyrings +curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc +chmod a+r /etc/apt/keyrings/docker.asc +# shellcheck disable=SC1091 +. /etc/os-release +tee /etc/apt/sources.list.d/docker.sources > /dev/null <&2 + exit 1 +fi +echo "{ \"data-root\": \"${DATA_ROOT}/docker\" }" > /etc/docker/daemon.json +# containerd keeps its own content store (where build-time image data lands); +# move it off the system disk too. +mkdir -p /etc/containerd +containerd config default | sed "s|root = \"/var/lib/containerd\"|root = \"${DATA_ROOT}/containerd\"|" > /etc/containerd/config.toml +systemctl start containerd docker +docker info | grep "Docker Root Dir" + +step "Build the runner image" +docker build --platform linux/amd64 -f "${DOCKERFILE}" -t "${IMAGE_TAG}" "${REPO_DIR}" + +step "Materialize the toolchain onto the host" +mkdir -p "${DATA_ROOT}/opt" "${DATA_ROOT}/runner" /opt /home +container="$(docker create "${IMAGE_TAG}")" +trap 'docker rm -f "${container}" >/dev/null 2>&1 || true' EXIT +# The `/.` suffix copies directory *contents*, so re-running over an existing +# materialization refreshes it in place instead of nesting runner/runner. +# Files absent from the image (runner registration, _work) are preserved. +docker cp "${container}:/home/runner/." "${DATA_ROOT}/runner" +docker cp "${container}:/opt/rustup/." "${DATA_ROOT}/opt/rustup" +docker cp "${container}:/opt/cargo/." "${DATA_ROOT}/opt/cargo" +for tool in uv uvx otelgen sccache; do + docker cp "${container}:/usr/local/bin/${tool}" "/usr/local/bin/${tool}" +done +docker rm "${container}" > /dev/null +trap - EXIT +ln -sfn "${DATA_ROOT}/runner" /home/runner +ln -sfn "${DATA_ROOT}/opt/rustup" /opt/rustup +ln -sfn "${DATA_ROOT}/opt/cargo" /opt/cargo +docker image rm "${IMAGE_TAG}" > /dev/null + +step "Install system packages" +apt-get install -y --no-install-recommends \ + build-essential clang cmake git gzip jq libprotobuf-dev libssl-dev mold \ + openssh-client pkg-config protobuf-compiler python3 sudo tar unzip wget xz-utils zip zstd + +step "Create runner user (${RUNNER_UID}:${RUNNER_GID}) and environment" +getent group "${RUNNER_GID}" > /dev/null || groupadd -g "${RUNNER_GID}" runner +id -u runner > /dev/null 2>&1 || useradd -u "${RUNNER_UID}" -g "${RUNNER_GID}" -d /home/runner -s /bin/bash runner +chown -R "${RUNNER_UID}:${RUNNER_GID}" "${DATA_ROOT}/runner" +echo 'PATH=/opt/cargo/bin:/usr/local/bin:/usr/bin:/bin' > "${DATA_ROOT}/runner/.env" +chown "${RUNNER_UID}:${RUNNER_GID}" "${DATA_ROOT}/runner/.env" + +if [[ "${REGISTER}" == "true" ]]; then + step "Register runner ${RUNNER_NAME} (labels: ${RUNNER_LABELS})" + if [[ -n "${RUNNER_SERVICE}" ]]; then + (cd /home/runner && ./svc.sh uninstall) || true + fi + runuser -u runner -- bash -c "cd /home/runner && HOME=/home/runner ./config.sh \ + --url '${REPO_URL}' --token '${RUNNER_TOKEN}' --name '${RUNNER_NAME}' \ + --labels '${RUNNER_LABELS}' --unattended --replace --disableupdate" + + step "Install and start the runner service" + cd /home/runner + ./svc.sh install runner + ./svc.sh start + ./svc.sh status +else + step "Skipping registration (pass --register --repo-url ... with RUNNER_TOKEN to enable)" + if [[ "${RUNNER_WAS_ACTIVE}" == "true" ]]; then + step "Restart runner service ${RUNNER_SERVICE}" + systemctl start "${RUNNER_SERVICE}" + systemctl --no-pager status "${RUNNER_SERVICE}" || true + fi +fi + +step "Done" +echo "Runner home: /home/runner -> ${DATA_ROOT}/runner" +echo "Before each workflow run on this persistent host, clean transient cargo state:" +echo " sudo rm -f /home/runner/.cargo/.package-cache" +echo " sudo rm -rf /home/runner/.cargo/.global-cache" diff --git a/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py b/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py new file mode 100644 index 00000000000..7031917f910 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py @@ -0,0 +1,392 @@ +#!/usr/bin/env python3 +# 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. + +# PEP 723 inline metadata (see .github/scripts/aliyun-ecs-runner-provision.py +# for the convention). +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "alibabacloud_ecs20140526>=4.1.0,<6", +# "alibabacloud_tea_openapi>=0.3.12,<1", +# ] +# /// + +"""Build the query-regression ECS custom image (manual ops tool). + +Boots a temporary pay-as-you-go ECS instance from a public Ubuntu 24.04 image, +builds the existing runner container image (the Dockerfile in the parent +directory remains the single source of the tool contract), materializes the +tool directories onto the host filesystem so the workflow's "Verify runner +image tools" step holds unchanged, installs the ephemeral-runner systemd unit, +and snapshots the result as a custom image. The temporary instance is deleted +afterwards. + +Sentinel polling reads the instance's serial console output +(GetInstanceConsoleOutput) and looks for marker lines the user data writes to +/dev/console. This has no in-guest agent dependency. + +Usage: + uv run .github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py \ + --region-id cn-hangzhou --vswitch-id vsw-... --security-group-id sg-... \ + --base-image-id ubuntu_24_04_x64_20G_alibase_*.vhd +""" + +from __future__ import annotations + +import argparse +import base64 +import os +import time +from pathlib import Path + +ASSETS_DIR = Path(__file__).resolve().parent +DONE_MARKER = "QREG_IMAGE_BUILD_DONE" +FAILED_MARKER = "QREG_IMAGE_BUILD_FAILED" +POLL_INTERVAL_SECONDS = 15 +CONSOLE_POLL_INTERVAL_SECONDS = 30 +BUILD_TIMEOUT_SECONDS = 60 * 60 + +# Same apt package contract as the runner Dockerfile; the base +# actions-runner image is Ubuntu 24.04, so an Ubuntu 24.04 host resolves the +# same tool versions (protoc 3.21.12, mold 2.40.4, Python 3.14.4). +# Docker itself comes from Docker's official repository (docker-ce), not the +# distribution-packaged docker.io. +APT_PACKAGES = [ + "build-essential", + "ca-certificates", + "clang", + "cmake", + "curl", + "git", + "gpg", + "gzip", + "jq", + "libprotobuf-dev", + "libssl-dev", + "mold", + "openssh-client", + "pkg-config", + "protobuf-compiler", + "python3", + "sudo", + "tar", + "unzip", + "wget", + "xz-utils", + "zip", + "zstd", +] + +DOCKER_CE_PACKAGES = "docker-ce docker-ce-cli containerd.io docker-buildx-plugin" + + +def render_user_data(dockerfile: str, start_runner: str, unit: str) -> str: + dockerfile_b64 = base64.b64encode(dockerfile.encode()).decode() + start_runner_b64 = base64.b64encode(start_runner.encode()).decode() + unit_b64 = base64.b64encode(unit.encode()).decode() + packages = " ".join(APT_PACKAGES) + return f"""#!/bin/bash +set -euo pipefail +trap 'echo "{FAILED_MARKER} at line $LINENO" > /dev/console' ERR +# Stream the full build log to the serial console so the poller (and anyone +# watching GetInstanceConsoleOutput) sees real progress, not a silent login +# prompt for the whole build. +exec > >(tee -a /dev/console) 2>&1 + +apt-get update +DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends {packages} + +# Docker from the official repository, not the distribution-packaged docker.io. +install -m 0755 -d /etc/apt/keyrings +curl -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /etc/apt/keyrings/docker.gpg +chmod a+r /etc/apt/keyrings/docker.gpg +. /etc/os-release +echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu ${{VERSION_CODENAME}} stable" \ + > /etc/apt/sources.list.d/docker.list +apt-get update +DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends {DOCKER_CE_PACKAGES} + +base64 -d > /tmp/Dockerfile <<'EOF' +{dockerfile_b64} +EOF +mkdir -p /tmp/image-context +docker build --platform linux/amd64 -f /tmp/Dockerfile -t qreg-runner:local /tmp/image-context + +# Materialize the tool contract onto the host filesystem. +container="$(docker create qreg-runner:local)" +trap 'docker rm -f "${{container}}" >/dev/null 2>&1 || true' EXIT +docker cp "${{container}}:/opt/rustup" /opt/rustup +docker cp "${{container}}:/opt/cargo" /opt/cargo +for tool in uv uvx otelgen sccache; do + docker cp "${{container}}:/usr/local/bin/${{tool}}" "/usr/local/bin/${{tool}}" +done +docker cp "${{container}}:/home/runner" /home/runner + +# Runner identity mirrors the container image (UID/GID 1001). +groupadd --gid 1001 runner 2>/dev/null || true +useradd --uid 1001 --gid 1001 --home-dir /home/runner --shell /bin/bash runner 2>/dev/null || true +chown -R 1001:1001 /home/runner + +mkdir -p /opt/ephemeral-github-runner +base64 -d > /opt/ephemeral-github-runner/start-runner.sh <<'EOF' +{start_runner_b64} +EOF +chmod 0755 /opt/ephemeral-github-runner/start-runner.sh +base64 -d > /etc/systemd/system/ephemeral-github-runner.service <<'EOF' +{unit_b64} +EOF +systemctl daemon-reload +systemctl enable ephemeral-github-runner.service + +# Keep the image free of the build-time docker state (also shrinks the +# snapshot: buildkit cache is several GB). +docker rm -f "${{container}}" >/dev/null +docker system prune -af >/dev/null +trap - EXIT + +echo "{DONE_MARKER}" > /dev/console +""" + + +def make_ecs_client(region_id: str): + from alibabacloud_ecs20140526.client import Client as EcsClient + from alibabacloud_tea_openapi.models import Config as OpenApiConfig + + return EcsClient( + OpenApiConfig( + access_key_id=os.environ["ALIBABA_CLOUD_ACCESS_KEY_ID"], + access_key_secret=os.environ["ALIBABA_CLOUD_ACCESS_KEY_SECRET"], + region_id=region_id, + endpoint=f"ecs.{region_id}.aliyuncs.com", + ) + ) + + +def wait_for_instance_status(client, region_id: str, instance_id: str, wanted: str, deadline: float) -> None: + from alibabacloud_ecs20140526 import models as ecs_models + + import json + + while time.monotonic() < deadline: + response = call_api_with_retry( + lambda: client.describe_instances( + ecs_models.DescribeInstancesRequest( + region_id=region_id, instance_ids=json.dumps([instance_id]) + ) + ), + "DescribeInstances", + ) + instances = response.body.instances.instance + if instances and instances[0].status == wanted: + return + time.sleep(POLL_INTERVAL_SECONDS) + raise TimeoutError(f"Instance {instance_id} did not reach status {wanted} in time") + + +# Aliyun error codes worth retrying: throttling and server-side faults. Client +# errors (4xx: permissions, bad parameters) are configuration problems and must +# fail fast instead of being retried. +TRANSIENT_ERROR_CODES = {"Throttling", "Throttling.User", "InternalError", "ServiceUnavailable"} + + +def is_transient(error: Exception) -> bool: + # UnretryableException comes from the darabonba network layer: connection + # reset, read timeout, DNS blip. + if type(error).__name__ == "UnretryableException": + return True + code = getattr(error, "code", "") or "" + status = getattr(error, "statusCode", None) + return code in TRANSIENT_ERROR_CODES or (isinstance(status, int) and status >= 500) + + +def call_api_with_retry(fn, description: str, attempts: int = 5): + """Retry transient API/network failures; the build is too long to die on a blip. + + Only call this with idempotent or safely-repeatable requests (reads, + StopInstance, CreateImage). Never with RunInstances: if the request + succeeded but the response was lost, a retry double-creates instances. + """ + for attempt in range(1, attempts + 1): + try: + return fn() + except Exception as error: # noqa: BLE001 + if attempt == attempts or not is_transient(error): + raise + print(f"{description} failed (attempt {attempt}/{attempts}): {error}", flush=True) + time.sleep(POLL_INTERVAL_SECONDS) + raise RuntimeError("unreachable: retry loop exited without returning") + + +def read_console_output(client, region_id: str, instance_id: str) -> str: + """Fetch the instance's serial console output; no in-guest agent needed.""" + from alibabacloud_ecs20140526 import models as ecs_models + + response = call_api_with_retry( + lambda: client.get_instance_console_output( + ecs_models.GetInstanceConsoleOutputRequest(region_id=region_id, instance_id=instance_id) + ), + "GetInstanceConsoleOutput", + ) + return base64.b64decode(response.body.console_output or "").decode("utf-8", "replace") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--region-id", default=os.environ.get("ALIYUN_ECS_REGION_ID")) + parser.add_argument("--vswitch-id", default=os.environ.get("ALIYUN_ECS_VSWITCH_ID")) + parser.add_argument("--security-group-id", default=os.environ.get("ALIYUN_ECS_SECURITY_GROUP_ID")) + parser.add_argument("--base-image-id", default=os.environ.get("ALIYUN_ECS_BASE_IMAGE_ID")) + parser.add_argument("--resource-group-id", default=os.environ.get("ALIYUN_ECS_RESOURCE_GROUP_ID")) + parser.add_argument("--instance-type", default="ecs.g7.xlarge") + parser.add_argument("--image-name", default=None, help="Defaults to a timestamped name.") + args = parser.parse_args() + + for name in ("region_id", "vswitch_id", "security_group_id", "base_image_id"): + if not getattr(args, name): + raise SystemExit(f"Missing required configuration: --{name.replace('_', '-')}") + + from alibabacloud_ecs20140526 import models as ecs_models + + client = make_ecs_client(args.region_id) + image_name = args.image_name or time.strftime( + "greptimedb-query-regression-runner-%Y%m%d%H%M%S", time.gmtime() + ) + + user_data = base64.b64encode( + render_user_data( + (ASSETS_DIR.parent / "Dockerfile").read_text(), + (ASSETS_DIR / "start-runner.sh").read_text(), + (ASSETS_DIR / "ephemeral-github-runner.service").read_text(), + ).encode() + ).decode() + + instance_id = None + try: + response = client.run_instances( + ecs_models.RunInstancesRequest( + region_id=args.region_id, + image_id=args.base_image_id, + resource_group_id=args.resource_group_id, + instance_type=args.instance_type, + v_switch_id=args.vswitch_id, + security_group_id=args.security_group_id, + instance_name=f"build-{image_name}", + description="Temporary builder for the query-regression ECS image", + amount=1, + instance_charge_type="PostPaid", + spot_strategy="NoSpot", + internet_charge_type="PayByTraffic", + internet_max_bandwidth_out=100, + # Peak usage is ~18G (OS+apt, docker image, and the materialized + # toolchain coexist briefly): deliberately tight, and a smaller + # disk makes the image snapshot faster. Instances created from + # the image get a larger system disk from the provision side. + system_disk=ecs_models.RunInstancesRequestSystemDisk( + category="cloud_essd", size="20" + ), + user_data=user_data, + tag=[ + ecs_models.RunInstancesRequestTag(key="managed-by", value="query-regression-ci"), + ecs_models.RunInstancesRequestTag(key="role", value="image-builder"), + ], + ) + ) + instance_id = response.body.instance_id_sets.instance_id_set[0] + print(f"Builder instance: {instance_id}", flush=True) + wait_for_instance_status(client, args.region_id, instance_id, "Running", time.monotonic() + 10 * 60) + + deadline = time.monotonic() + BUILD_TIMEOUT_SECONDS + while time.monotonic() < deadline: + console = read_console_output(client, args.region_id, instance_id) + if DONE_MARKER in console: + break + if FAILED_MARKER in console: + tail = "\n".join(console.splitlines()[-20:]) + raise RuntimeError(f"Image build failed on the builder; console tail:\n{tail}") + lines = console.splitlines() + print(f"Build in progress; last console line: {lines[-1] if lines else '(none yet)'}", flush=True) + time.sleep(CONSOLE_POLL_INTERVAL_SECONDS) + else: + raise TimeoutError("Image build did not finish in time") + + print("Stopping builder before image creation", flush=True) + try: + call_api_with_retry( + lambda: client.stop_instance( + ecs_models.StopInstanceRequest( + instance_id=instance_id, + # The instance is deleted right after the snapshot, but + # there is no reason to keep billing vCPU during the stop. + stopped_mode="StopCharging", + ) + ), + "StopInstance", + ) + except Exception as error: # noqa: BLE001 + # Instance families with local disks do not support StopCharging. + print(f"StopCharging unavailable ({error}); stopping with default mode", flush=True) + call_api_with_retry( + lambda: client.stop_instance(ecs_models.StopInstanceRequest(instance_id=instance_id)), + "StopInstance", + ) + wait_for_instance_status(client, args.region_id, instance_id, "Stopped", time.monotonic() + 10 * 60) + + image = call_api_with_retry( + lambda: client.create_image( + ecs_models.CreateImageRequest( + region_id=args.region_id, instance_id=instance_id, image_name=image_name + ) + ), + "CreateImage", + ) + image_id = image.body.image_id + deadline = time.monotonic() + 60 * 60 + while time.monotonic() < deadline: + description = call_api_with_retry( + lambda: client.describe_images( + ecs_models.DescribeImagesRequest(region_id=args.region_id, image_id=image_id) + ), + "DescribeImages", + ) + images = description.body.images.image + if images: + status = images[0].status + print(f"Image {image_id} status: {status} (progress {images[0].progress})", flush=True) + if status == "Available": + break + time.sleep(POLL_INTERVAL_SECONDS) + else: + raise TimeoutError( + f"Image {image_id} did not become Available in time. The snapshot " + f"continues server-side: check its status in the console and reuse it " + f"once Available instead of rebuilding." + ) + + print(f"Custom image ready: {image_id} ({image_name})", flush=True) + print(f"Set the repo variable QUERY_REGRESSION_ECS_IMAGE_ID={image_id}", flush=True) + return 0 + finally: + if instance_id: + try: + client.delete_instance( + ecs_models.DeleteInstanceRequest(instance_id=instance_id, force=True) + ) + print(f"Deleted builder instance {instance_id}", flush=True) + except Exception as error: # noqa: BLE001 + print(f"Failed to delete builder instance {instance_id}: {error}", flush=True) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/runner-scale-sets/query-regression/ecs-image/ephemeral-github-runner.service b/.github/runner-scale-sets/query-regression/ecs-image/ephemeral-github-runner.service new file mode 100644 index 00000000000..a1cbe6a4503 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/ecs-image/ephemeral-github-runner.service @@ -0,0 +1,45 @@ +# 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. + +[Unit] +Description=Ephemeral GitHub Actions runner +After=network-online.target +Wants=network-online.target +# cloud-final runs the provision user data, which writes +# /etc/ephemeral-github-runner.env and the drop-ins below. Ordering after it +# makes the runner come up configured on first boot, no restart needed. +# (cloud-init's `systemctl restart` in user data remains as the bridge for +# images built before this ordering existed.) +After=cloud-final.service + +[Service] +Type=simple +# The toolchain lives in the image at /opt/cargo/bin (the Dockerfile's ENV +# PATH does not survive materialization onto the host). Publish it here so +# every runner job inherits it. /etc/ephemeral-github-runner.env currently +# carries the same PATH line as a bridge for images built before this +# directive existed; EnvironmentFile is applied after Environment=, and the +# two values are kept identical, so there is no conflict. +Environment=PATH=/opt/cargo/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin +EnvironmentFile=/etc/ephemeral-github-runner.env +ExecStart=/opt/ephemeral-github-runner/start-runner.sh +Restart=no +# If the kernel OOM killer kills a job child (rustc/cargo), keep the runner +# service alive so the job is reported as a normal failure with logs instead +# of the runner vanishing and GitHub reporting a cancellation. cloud-init's +# oom.conf drop-in carries the same directive for images built before this. +OOMPolicy=continue + +[Install] +WantedBy=multi-user.target diff --git a/.github/runner-scale-sets/query-regression/ecs-image/start-runner.sh b/.github/runner-scale-sets/query-regression/ecs-image/start-runner.sh new file mode 100644 index 00000000000..a6aff87ca22 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/ecs-image/start-runner.sh @@ -0,0 +1,37 @@ +#!/bin/bash +# 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. + +# Configure and start the ephemeral GitHub Actions runner on an ECS instance. +# Baked into the custom image at /opt/ephemeral-github-runner/ by +# build-ecs-image.py; the per-run environment is written by cloud-init to +# /etc/ephemeral-github-runner.env. +set -euo pipefail + +# shellcheck disable=SC1091 +source /etc/ephemeral-github-runner.env + +export HOME=/home/runner +cd /home/runner + +if [[ ! -f /home/runner/.runner ]]; then + runuser -u runner -- ./config.sh \ + --url "${REPO_URL}" \ + --token "${RUNNER_TOKEN}" \ + --name "${RUNNER_NAME}" \ + --labels "${RUNNER_LABELS}" \ + --ephemeral --unattended --replace --disableupdate +fi + +exec runuser -u runner -- ./run.sh diff --git a/.github/runner-scale-sets/query-regression/values-8-cores.yaml b/.github/runner-scale-sets/query-regression/values-8-cores.yaml deleted file mode 100644 index cff1322680b..00000000000 --- a/.github/runner-scale-sets/query-regression/values-8-cores.yaml +++ /dev/null @@ -1,146 +0,0 @@ -githubConfigUrl: "https://github.com/GreptimeTeam/greptimedb" -githubConfigSecret: greptimedb-arc-github-app - -runnerScaleSetName: "perf-regression-8-cores" -minRunners: 0 -maxRunners: 1 - -template: - spec: - automountServiceAccountToken: false - activeDeadlineSeconds: 12600 - nodeSelector: - kubernetes.io/hostname: minipc-3 - securityContext: - runAsNonRoot: true - runAsUser: 1001 - runAsGroup: 1001 - fsGroup: 1001 - fsGroupChangePolicy: OnRootMismatch - seccompProfile: - type: RuntimeDefault - volumes: - - name: cargo-home - emptyDir: {} - - name: build-cache - persistentVolumeClaim: - claimName: query-regression-build-cache - initContainers: - - name: initialize-build-cache - image: greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner@sha256:e713b294e23b7e15184e558866c90025e59930033e72c97650dbc7f1ca022d11 - imagePullPolicy: IfNotPresent - command: - - /bin/sh - - -ec - - | - umask 0002 - for directory in cargo-registry-v1 cargo-git-v1 query-regression-target-v1 meta-v1 sccache-v1; do - cache_directory="/cache/${directory}" - mkdir -p "${cache_directory}" - test -w "${cache_directory}" - probe_file="${cache_directory}/.write-probe" - : > "${probe_file}" - test -f "${probe_file}" - rm "${probe_file}" - done - securityContext: - runAsNonRoot: true - runAsUser: 1001 - runAsGroup: 1001 - allowPrivilegeEscalation: false - capabilities: - drop: - - ALL - resources: - requests: - cpu: 100m - memory: 128Mi - ephemeral-storage: 1Gi - limits: - cpu: 500m - memory: 512Mi - ephemeral-storage: 1Gi - volumeMounts: - - name: build-cache - mountPath: /cache - containers: - - name: runner - image: greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner@sha256:e713b294e23b7e15184e558866c90025e59930033e72c97650dbc7f1ca022d11 - imagePullPolicy: IfNotPresent - command: ["/home/runner/run.sh"] - env: - - name: CARGO_HOME - value: /home/runner/.cargo - - name: HTTP_PROXY - valueFrom: - configMapKeyRef: - name: query-regression-runner-local-env - key: HTTP_PROXY - optional: true - - name: HTTPS_PROXY - valueFrom: - configMapKeyRef: - name: query-regression-runner-local-env - key: HTTPS_PROXY - optional: true - - name: NO_PROXY - valueFrom: - configMapKeyRef: - name: query-regression-runner-local-env - key: NO_PROXY - optional: true - - name: UV_CACHE_DIR - value: /home/runner/.cargo/uv-cache - - name: RUSTUP_HOME - value: /opt/rustup - - name: RUSTUP_TOOLCHAIN - value: nightly-2026-03-21 - - name: RUSTUP_AUTO_INSTALL - value: "0" - - name: CARGO_TARGET_DIR - value: /home/runner/query-regression-target - - name: QUERY_REGRESSION_CACHE_META - value: /home/runner/query-regression-cache-meta - - name: RUSTC_WRAPPER - value: /usr/local/bin/sccache - - name: SCCACHE_DIR - value: /home/runner/.cache/sccache - - name: SCCACHE_CACHE_SIZE - value: 40G - - name: CARGO_INCREMENTAL - value: "0" - securityContext: - runAsNonRoot: true - runAsUser: 1001 - runAsGroup: 1001 - allowPrivilegeEscalation: false - capabilities: - drop: - - ALL - resources: - requests: - cpu: "6" - memory: 24Gi - ephemeral-storage: 80Gi - limits: - cpu: "8" - memory: 24Gi - ephemeral-storage: 80Gi - volumeMounts: - - name: cargo-home - mountPath: /home/runner/.cargo - - name: build-cache - mountPath: /home/runner/.cargo/registry - subPath: cargo-registry-v1 - - name: build-cache - mountPath: /home/runner/.cargo/git - subPath: cargo-git-v1 - - name: build-cache - mountPath: /home/runner/query-regression-target - subPath: query-regression-target-v1 - - name: build-cache - mountPath: /home/runner/query-regression-cache-meta - subPath: meta-v1 - - name: build-cache - mountPath: /home/runner/.cache/sccache - subPath: sccache-v1 diff --git a/.github/runner-scale-sets/query-regression/values-paused.yaml b/.github/runner-scale-sets/query-regression/values-paused.yaml deleted file mode 100644 index fe734ef7615..00000000000 --- a/.github/runner-scale-sets/query-regression/values-paused.yaml +++ /dev/null @@ -1,2 +0,0 @@ -minRunners: 0 -maxRunners: 0 diff --git a/.github/scripts/aliyun-ecs-runner-provision.py b/.github/scripts/aliyun-ecs-runner-provision.py new file mode 100644 index 00000000000..74a69681f3d --- /dev/null +++ b/.github/scripts/aliyun-ecs-runner-provision.py @@ -0,0 +1,485 @@ +#!/usr/bin/env python3 +# 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. + +# This is the first `.github/scripts/` user of PEP 723 inline script metadata: +# `uv run` installs the pinned Aliyun SDK before executing. The SDK import stays +# lazy (inside `make_ecs_client`) so unit tests can import this module with a +# plain stdlib interpreter. +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# # Pin exactly after the first live run against the Aliyun API. +# "alibabacloud_ecs20140526>=4.1.0,<6", +# "alibabacloud_tea_openapi>=0.3.12,<1", +# ] +# /// + +"""Provision an ephemeral Aliyun ECS query-regression runner. + +Creates one pay-as-you-go ECS instance from the query-regression custom image +and waits until the instance registers itself as an ephemeral GitHub Actions +runner with a per-run label. Build caches live on the instance system disk +and disappear with the VM. The instance id is written to job outputs as soon +as the VM exists so the teardown job is the single DeleteInstance caller (a +timeout here used to delete and then teardown deleted again). On runner +online timeout the console is dumped and the job fails; teardown releases +the instance. + +The Aliyun credentials are used only in this control-plane job (ubuntu-latest) +and are never passed to the ECS instance; the instance receives only a +short-lived GitHub runner registration token via user data. +""" + +from __future__ import annotations + +import argparse +import base64 +import json +import os +import time +import urllib.error +import urllib.request +from dataclasses import dataclass + +RUNNER_NAME_PREFIX = "qreg-ecs" +RUNNER_LABEL_PREFIX = "query-regression-ecs" +MANAGED_BY_TAG_KEY = "managed-by" +MANAGED_BY_TAG_VALUE = "query-regression-ci" +RUN_TAG_KEY = "query-regression-run-id" + +# Runner cache paths on the instance system disk. They are created empty +# every provision and deleted with the VM. +CACHE_PATHS = ( + "/home/runner/.cargo/registry", + "/home/runner/.cargo/git", + "/home/runner/query-regression-target", + "/home/runner/query-regression-cache-meta", + "/home/runner/.cache/sccache", +) + +RUNNER_ONLINE_TIMEOUT_SECONDS = 10 * 60 +POLL_INTERVAL_SECONDS = 5 + +# Nightly thin-LTO of greptime peaks above ecs.c9i.2xlarge's 16 GiB. A swap +# file plus masking systemd-oomd lets the kernel reclaim rustc pages instead +# of SIGTERM-ing the runner service cgroup (GitHub then reports the step as +# "The operation was canceled"). Sized to match that 16 GiB instance; a +# 32 GiB type still benefits as a safety net. Lives on the system disk next +# to the cold build caches. +SWAP_FILE = "/swapfile" +SWAP_SIZE_GIB = 16 + + +@dataclass(frozen=True) +class ProvisionConfig: + region_id: str + vswitch_id: str + security_group_id: str + image_id: str + instance_type: str + repo: str + run_id: str + github_token: str + # Optional resource group; required when the RAM grant is scoped to one. + resource_group_id: str | None = None + # Runner identity inside the image; the workflow asserts the same values. + runner_uid: str = "1001" + runner_gid: str = "1001" + + +def runner_name_for_run(run_id: str) -> str: + return f"{RUNNER_NAME_PREFIX}-{run_id}" + + +def runner_label_for_run(run_id: str) -> str: + return f"{RUNNER_LABEL_PREFIX}-{run_id}" + + +def render_user_data( + runner_name: str, + runner_label: str, + runner_token: str, + repo: str, + runner_uid: str = "1001", + runner_gid: str = "1001", +) -> str: + """Render the cloud-init shell script for the runner instance.""" + destinations = " ".join(f'"{dst}"' for dst in CACHE_PATHS) + cache_setup = f"""# Caches live on the system disk, are deleted with the instance, and every +# run compiles cold. Within-run reuse (base warming candidate via the shared +# target dir and sccache) still applies. +mkdir -p {destinations} +chown -R {runner_uid}:{runner_gid} /home/runner""" + swap_setup = f"""# Mask systemd-oomd before enabling swap: Ubuntu 24.04 kills the whole +# service cgroup on PSI pressure, and swap thrashing looks like pressure. +systemctl disable --now systemd-oomd.socket systemd-oomd.service || true +systemctl mask systemd-oomd.socket systemd-oomd.service || true +# apt-daily-upgrade / unattended-upgrades can `systemctl restart` services +# whose libraries were updated. That SIGTERM-s the runner mid-job; GitHub +# records UserCancelled (job still valid, child exit 143). These VMs live +# ~1h; security updates belong in the image, not at job time. +systemctl disable --now unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service || true +systemctl mask unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service || true +systemctl stop unattended-upgrades.service || true +killall -9 unattended-upgr || true +cat > /etc/apt/apt.conf.d/99disable-auto-updates <<'APTEOF' +APT::Periodic::Update-Package-Lists "0"; +APT::Periodic::Unattended-Upgrade "0"; +APT::Periodic::Download-Upgradeable-Packages "0"; +APTEOF +if [[ ! -f "{SWAP_FILE}" ]]; then + fallocate --length {SWAP_SIZE_GIB}G "{SWAP_FILE}" + chmod 600 "{SWAP_FILE}" + mkswap "{SWAP_FILE}" +fi +swapon "{SWAP_FILE}" +sysctl --write vm.swappiness=10 +swapon --show +free --human""" + return f"""#!/bin/bash +set -euo pipefail + +{cache_setup} + +{swap_setup} + +cat > /etc/ephemeral-github-runner.env <<'ENVEOF' +RUNNER_NAME={runner_name} +RUNNER_LABELS={runner_label} +RUNNER_TOKEN={runner_token} +REPO_URL=https://github.com/{repo} +# The container image baked /opt/cargo/bin into PATH via Dockerfile ENV; on +# the host nothing inherits that, so publish it through the unit's +# EnvironmentFile: runner job processes inherit the runner's environment. +PATH=/opt/cargo/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin +ENVEOF + +# Stream the runner service's own logs (config output, job lifecycle, crash +# messages) to the serial console, next to the kernel's OOM-killer records. +# Teardown dumps the console tail before deleting the instance, so the +# witness of what killed the runner survives the machine. +mkdir -p /etc/systemd/system/ephemeral-github-runner.service.d +cat > /etc/systemd/system/ephemeral-github-runner.service.d/console.conf <<'CONFEOF' +[Service] +StandardOutput=journal+console +StandardError=journal+console +CONFEOF +# systemd's DefaultOOMPolicy=stop SIGTERM-s the whole unit when *any* +# process in the cgroup is OOM-killed (typically rustc during thin LTO). +# That is the "Session terminated, killing shell..." / job Canceled path: +# runuser dies, GitHub deletes the ephemeral registration, always() steps +# never run. continue keeps the runner up so cargo can report signal 9. +cat > /etc/systemd/system/ephemeral-github-runner.service.d/oom.conf <<'CONFEOF' +[Service] +OOMPolicy=continue +CONFEOF +systemctl daemon-reload +# restart (not start): the image enables this unit, so it may already +# have come up without the drop-ins if it raced cloud-init. restart +# applies OOMPolicy=continue to the running cgroup. --no-block matters: +# newer image units carry After=cloud-final.service, and this script IS +# cloud-final — a blocking restart would deadlock until TimeoutStartSec. +systemctl restart --no-block ephemeral-github-runner.service +""" + + +def encode_user_data(script: str) -> str: + return base64.b64encode(script.encode("utf-8")).decode("ascii") + + +def github_api(token: str, method: str, path: str, body: dict | None = None) -> dict: + data = json.dumps(body).encode("utf-8") if body is not None else None + request = urllib.request.Request( + f"https://api.github.com{path}", + data=data, + method=method, + headers={ + "Accept": "application/vnd.github+json", + "Authorization": f"Bearer {token}", + "X-GitHub-Api-Version": "2022-11-28", + }, + ) + try: + with urllib.request.urlopen(request, timeout=30) as response: + return json.loads(response.read().decode("utf-8")) + except urllib.error.HTTPError as error: + # GitHub's error body says exactly why (e.g. "Must have admin rights to + # Repository" for a PAT without the required scope); surface it instead + # of a bare "HTTP Error 403". + body = error.read().decode("utf-8", "replace") + raise SystemExit( + f"GitHub API {method} {path} failed: HTTP {error.code}: {body}\n" + "The token comes from the GH_PERSONAL_ACCESS_TOKEN secret; it needs " + "'repo' scope (classic PAT) or 'Administration: write' on the " + "repository (fine-grained PAT)." + ) from error + + +def create_registration_token(github_token: str, repo: str) -> str: + response = github_api( + github_token, + "POST", + f"/repos/{repo}/actions/runners/registration-token", + body={}, + ) + return response["token"] + + +def find_runner_by_name(github_token: str, repo: str, name: str) -> dict | None: + page = 1 + while True: + response = github_api( + github_token, + "GET", + f"/repos/{repo}/actions/runners?per_page=100&page={page}", + ) + runners = response.get("runners", []) + for runner in runners: + if runner.get("name") == name: + return runner + if len(runners) < 100: + return None + page += 1 + + +def make_ecs_client(config: ProvisionConfig): + from alibabacloud_ecs20140526.client import Client as EcsClient + from alibabacloud_tea_openapi.models import Config as OpenApiConfig + + access_key_id = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_ID", "") + access_key_secret = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_SECRET", "") + if not access_key_id or not access_key_secret: + raise SystemExit( + "ALIBABA_CLOUD_ACCESS_KEY_ID/SECRET are empty; in CI they come from the " + "ALICLOUD_ECS_ACCESS_KEY_ID/SECRET repository secrets (a missing or " + "misnamed secret expands to an empty string)." + ) + return EcsClient( + OpenApiConfig( + access_key_id=access_key_id, + access_key_secret=access_key_secret, + region_id=config.region_id, + endpoint=f"ecs.{config.region_id}.aliyuncs.com", + ) + ) + + +def wait_for_instance_status(client, region_id: str, instance_id: str, wanted: str, deadline: float) -> None: + from alibabacloud_ecs20140526 import models as ecs_models + + while time.monotonic() < deadline: + response = client.describe_instances( + ecs_models.DescribeInstancesRequest( + region_id=region_id, instance_ids=json.dumps([instance_id]) + ) + ) + instances = response.body.instances.instance + if instances and instances[0].status == wanted: + return + time.sleep(POLL_INTERVAL_SECONDS) + raise TimeoutError(f"Instance {instance_id} did not reach status {wanted} in time") + + +def fetch_console_output(client, region_id: str, instance_id: str) -> str: + from alibabacloud_ecs20140526 import models as ecs_models + + response = client.get_instance_console_output( + ecs_models.GetInstanceConsoleOutputRequest(region_id=region_id, instance_id=instance_id) + ) + return base64.b64decode(response.body.console_output or "").decode("utf-8", "replace") + + +class ConsoleTailer: + """Incrementally prints new serial-console lines so cloud-init progress is + visible in the CI log while waiting, not only after a failure dump.""" + + def __init__(self, client, region_id: str, instance_id: str) -> None: + self.client = client + self.region_id = region_id + self.instance_id = instance_id + self.printed_lines = 0 + + def poll(self) -> None: + try: + lines = fetch_console_output(self.client, self.region_id, self.instance_id).splitlines() + except Exception as error: # noqa: BLE001 + print(f"[console] unable to fetch console output: {error}", flush=True) + return + # The API returns a bounded tail; if it ever shrinks, restart from the + # beginning of the new buffer rather than skipping lines. + if len(lines) < self.printed_lines: + self.printed_lines = 0 + for line in lines[self.printed_lines :]: + print(f"[console] {line}", flush=True) + self.printed_lines = len(lines) + + +def dump_console_output(client, region_id: str, instance_id: str) -> None: + try: + output = fetch_console_output(client, region_id, instance_id) + except Exception as error: # noqa: BLE001 + print(f"Unable to fetch console output for {instance_id}: {error}", flush=True) + return + tail = "\n".join(output.splitlines()[-80:]) + print(f"::group::Console output tail for {instance_id}\n{tail}\n::endgroup::", flush=True) + summary_path = os.environ.get("GITHUB_STEP_SUMMARY") + if summary_path: + with open(summary_path, "a", encoding="utf-8") as summary: + summary.write(f"
ECS console output ({instance_id})\n\n```\n") + summary.write(tail) + summary.write("\n```\n
\n") + + +def append_github_output(name: str, value: str) -> None: + output_path = os.environ.get("GITHUB_OUTPUT") + if output_path: + with open(output_path, "a", encoding="utf-8") as output: + output.write(f"{name}={value}\n") + + +def run_instance(client, config: ProvisionConfig, user_data: str) -> str: + from alibabacloud_ecs20140526 import models as ecs_models + + request = ecs_models.RunInstancesRequest( + region_id=config.region_id, + image_id=config.image_id, + resource_group_id=config.resource_group_id, + instance_type=config.instance_type, + v_switch_id=config.vswitch_id, + security_group_id=config.security_group_id, + instance_name=runner_name_for_run(config.run_id), + host_name=runner_name_for_run(config.run_id), + description=f"Ephemeral query-regression runner for run {config.run_id}", + amount=1, + instance_charge_type="PostPaid", + spot_strategy="NoSpot", + internet_charge_type="PayByTraffic", + internet_max_bandwidth_out=100, + # Holds the image (~12G), the 16G swapfile, the source checkout, the + # base+candidate cluster data homes, and cold build caches. Deleted + # with the instance. + system_disk=ecs_models.RunInstancesRequestSystemDisk(category="cloud_essd", size="40"), + user_data=user_data, + tag=[ + ecs_models.RunInstancesRequestTag(key=MANAGED_BY_TAG_KEY, value=MANAGED_BY_TAG_VALUE), + ecs_models.RunInstancesRequestTag(key=RUN_TAG_KEY, value=config.run_id), + ], + ) + response = client.run_instances(request) + instance_id = response.body.instance_id_sets.instance_id_set[0] + return instance_id + + +def provision(config: ProvisionConfig) -> int: + client = make_ecs_client(config) + runner_name = runner_name_for_run(config.run_id) + runner_label = runner_label_for_run(config.run_id) + + registration_token = create_registration_token(config.github_token, config.repo) + user_data = encode_user_data( + render_user_data( + runner_name, + runner_label, + registration_token, + config.repo, + config.runner_uid, + config.runner_gid, + ) + ) + + print(f"::group::Run instance {runner_name}", flush=True) + instance_id = run_instance(client, config, user_data) + print(f"Instance id: {instance_id}", flush=True) + # Teardown is the only DeleteInstance caller. Write outputs immediately + # so a later status/online failure still releases this VM. + append_github_output("label", runner_label) + append_github_output("instance_id", instance_id) + append_github_output("runner_name", runner_name) + wait_for_instance_status(client, config.region_id, instance_id, "Running", time.monotonic() + 5 * 60) + print("::endgroup::", flush=True) + + print("::group::Wait for runner online", flush=True) + deadline = time.monotonic() + RUNNER_ONLINE_TIMEOUT_SECONDS + tailer = ConsoleTailer(client, config.region_id, instance_id) + last_console_poll = 0.0 + while time.monotonic() < deadline: + runner = find_runner_by_name(config.github_token, config.repo, runner_name) + if runner is not None: + status = runner.get("status") + print(f"Runner {runner_name} status: {status}", flush=True) + if status == "online": + print("::endgroup::", flush=True) + print(f"Runner {runner_name} is online with label {runner_label}", flush=True) + return 0 + # The serial console lags the guest by a minute or so; 30s polling is + # enough to follow cloud-init without tripping API throttling. + if time.monotonic() - last_console_poll >= 30: + tailer.poll() + last_console_poll = time.monotonic() + time.sleep(POLL_INTERVAL_SECONDS) + + print("::endgroup::", flush=True) + print( + f"Runner {runner_name} did not come online in time; " + "leaving instance for the teardown job to delete", + flush=True, + ) + dump_console_output(client, config.region_id, instance_id) + return 1 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--region-id", default=os.environ.get("ALIYUN_ECS_REGION_ID")) + parser.add_argument("--vswitch-id", default=os.environ.get("ALIYUN_ECS_VSWITCH_ID")) + parser.add_argument("--security-group-id", default=os.environ.get("ALIYUN_ECS_SECURITY_GROUP_ID")) + parser.add_argument("--image-id", default=os.environ.get("QUERY_REGRESSION_ECS_IMAGE_ID")) + parser.add_argument("--instance-type", default=os.environ.get("ALIYUN_ECS_INSTANCE_TYPE")) + parser.add_argument("--repo", default=os.environ.get("GITHUB_REPOSITORY")) + parser.add_argument("--run-id", default=os.environ.get("GITHUB_RUN_ID")) + parser.add_argument("--github-token", default=os.environ.get("GH_PERSONAL_ACCESS_TOKEN")) + parser.add_argument("--runner-uid", default=os.environ.get("QUERY_REGRESSION_RUNNER_UID", "1001")) + parser.add_argument("--runner-gid", default=os.environ.get("QUERY_REGRESSION_RUNNER_GID", "1001")) + parser.add_argument("--resource-group-id", default=os.environ.get("ALIYUN_ECS_RESOURCE_GROUP_ID")) + args = parser.parse_args() + + missing = [ + name + for name, value in vars(args).items() + if name not in ("resource_group_id",) + and (value is None or (isinstance(value, str) and not value)) + ] + if missing: + flags = ", ".join(f"--{name.replace('_', '-')}" for name in missing) + raise SystemExit(f"Missing required configuration: {flags} (or their env defaults)") + + config = ProvisionConfig( + region_id=args.region_id, + vswitch_id=args.vswitch_id, + security_group_id=args.security_group_id, + image_id=args.image_id, + instance_type=args.instance_type, + repo=args.repo, + run_id=args.run_id, + github_token=args.github_token, + resource_group_id=args.resource_group_id or None, + runner_uid=args.runner_uid, + runner_gid=args.runner_gid, + ) + return provision(config) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/scripts/aliyun-ecs-runner-teardown.py b/.github/scripts/aliyun-ecs-runner-teardown.py new file mode 100644 index 00000000000..36c80253b95 --- /dev/null +++ b/.github/scripts/aliyun-ecs-runner-teardown.py @@ -0,0 +1,237 @@ +#!/usr/bin/env python3 +# 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. + +# PEP 723 inline metadata (see aliyun-ecs-runner-provision.py for the +# convention). The SDK import stays lazy so unit tests run on a plain +# stdlib interpreter. +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "alibabacloud_ecs20140526>=4.1.0,<6", +# "alibabacloud_tea_openapi>=0.3.12,<1", +# ] +# /// + +"""Tear down ephemeral Aliyun ECS query-regression runners. + +Two modes: + +- Targeted (per workflow run): delete one instance by id and deregister its + runner by name. Both steps are idempotent and best-effort; a missing + instance or runner is not an error. +- Sweep (scheduled janitor): delete every instance tagged as managed by + query-regression CI whose creation time is older than the given TTL, and + deregister the matching runners. This is the safety net for runs whose + teardown job never executed. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import os +import sys +from datetime import datetime, timedelta, timezone +from pathlib import Path + +# Reuse the GitHub API client, runner lookup, and tag constants from the +# provision script. +_PROVISION_SPEC = importlib.util.spec_from_file_location( + "aliyun_ecs_runner_provision", + Path(__file__).resolve().parent / "aliyun-ecs-runner-provision.py", +) +assert _PROVISION_SPEC is not None and _PROVISION_SPEC.loader is not None +provision = importlib.util.module_from_spec(_PROVISION_SPEC) +sys.modules[_PROVISION_SPEC.name] = provision +_PROVISION_SPEC.loader.exec_module(provision) + +SWEEP_TTL = timedelta(hours=4) +# ECS DescribeInstances creation_time is ISO8601 UTC, e.g. "2026-08-17T02:13Z". +CREATION_TIME_FORMATS = ("%Y-%m-%dT%H:%M:%SZ", "%Y-%m-%dT%H:%MZ") + + +def parse_creation_time(value: str) -> datetime: + for fmt in CREATION_TIME_FORMATS: + try: + return datetime.strptime(value, fmt).replace(tzinfo=timezone.utc) + except ValueError: + continue + raise ValueError(f"Unparsable ECS creation time: {value}") + + +def expired_instance_names( + instances: list[tuple[str, str, str]], now: datetime, ttl: timedelta +) -> list[tuple[str, str]]: + """Pick (instance_id, instance_name) pairs whose creation is older than ttl. + + `instances` items are (instance_id, instance_name, creation_time). + """ + expired = [] + for instance_id, instance_name, creation_time in instances: + age = now - parse_creation_time(creation_time) + if age >= ttl: + expired.append((instance_id, instance_name)) + return expired + + +def make_ecs_client(region_id: str): + from alibabacloud_ecs20140526.client import Client as EcsClient + from alibabacloud_tea_openapi.models import Config as OpenApiConfig + + access_key_id = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_ID", "") + access_key_secret = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_SECRET", "") + if not access_key_id or not access_key_secret: + raise SystemExit( + "ALIBABA_CLOUD_ACCESS_KEY_ID/SECRET are empty; in CI they come from the " + "ALICLOUD_ECS_ACCESS_KEY_ID/SECRET repository secrets (a missing or " + "misnamed secret expands to an empty string)." + ) + return EcsClient( + OpenApiConfig( + access_key_id=access_key_id, + access_key_secret=access_key_secret, + region_id=region_id, + endpoint=f"ecs.{region_id}.aliyuncs.com", + ) + ) + + +def delete_instance(client, instance_id: str, region_id: str | None = None) -> bool: + import time + + from alibabacloud_ecs20140526 import models as ecs_models + + # Export the serial console before destroying the evidence: a cancelled + # run still runs this teardown job, and the console tail (cloud-init, + # runner service logs, kernel OOM records) is the only witness of what + # happened on the machine. + if region_id is not None: + provision.dump_console_output(client, region_id, instance_id) + + # DeleteInstance rejects instances that are still Initializing (or in + # another transitional status). Teardown runs right after a cancelled + # run, when the instance may be only seconds old, so retry the transient + # status errors for a few minutes instead of leaking the instance. + deadline = time.monotonic() + 5 * 60 + while True: + try: + client.delete_instance(ecs_models.DeleteInstanceRequest(instance_id=instance_id, force=True)) + print(f"Deleted instance {instance_id}", flush=True) + return True + except Exception as error: # noqa: BLE001 + message = str(error) + if "InvalidInstanceId.NotFound" in message: + print(f"Instance {instance_id} already gone", flush=True) + return True + if "IncorrectInstanceStatus" in message and time.monotonic() < deadline: + print(f"Instance {instance_id} is in a transitional status; retrying delete", flush=True) + time.sleep(15) + continue + print(f"Failed to delete instance {instance_id}: {error}", flush=True) + return False + + +def deregister_runner(token: str, repo: str, runner_name: str) -> bool: + runner = provision.find_runner_by_name(token, repo, runner_name) + if runner is None: + print(f"Runner {runner_name} is not registered", flush=True) + return True + try: + provision.github_api(token, "DELETE", f"/repos/{repo}/actions/runners/{runner['id']}") + print(f"Deregistered runner {runner_name} (id {runner['id']})", flush=True) + return True + except Exception as error: # noqa: BLE001 + print(f"Failed to deregister runner {runner_name}: {error}", flush=True) + return False + + +def list_managed_instances(client, region_id: str) -> list[tuple[str, str, str]]: + from alibabacloud_ecs20140526 import models as ecs_models + + result: list[tuple[str, str, str]] = [] + next_token = None + while True: + request = ecs_models.DescribeInstancesRequest( + region_id=region_id, + tag=[ + ecs_models.DescribeInstancesRequestTag( + key=provision.MANAGED_BY_TAG_KEY, value=provision.MANAGED_BY_TAG_VALUE + ) + ], + max_results=100, + next_token=next_token, + ) + response = client.describe_instances(request) + for instance in response.body.instances.instance: + result.append((instance.instance_id, instance.instance_name, instance.creation_time)) + next_token = response.body.next_token + if not next_token: + return result + + +def sweep(client, region_id: str, repo: str, github_token: str, ttl: timedelta) -> int: + instances = list_managed_instances(client, region_id) + print(f"Found {len(instances)} managed instance(s) in {region_id}", flush=True) + expired = expired_instance_names(instances, datetime.now(timezone.utc), ttl) + ok = True + for instance_id, instance_name in expired: + print(f"Instance {instance_id} ({instance_name}) exceeds TTL {ttl}; deleting", flush=True) + # Runner names mirror instance names by construction in the provision + # script (both are qreg-ecs-). + ok &= delete_instance(client, instance_id, region_id) + ok &= deregister_runner(github_token, repo, instance_name) + return 0 if ok else 1 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--region-id", default=os.environ.get("ALIYUN_ECS_REGION_ID")) + parser.add_argument("--repo", default=os.environ.get("GITHUB_REPOSITORY")) + parser.add_argument("--github-token", default=os.environ.get("GH_PERSONAL_ACCESS_TOKEN")) + parser.add_argument("--instance-id", default=os.environ.get("QUERY_REGRESSION_ECS_INSTANCE_ID")) + parser.add_argument("--runner-name", default=os.environ.get("QUERY_REGRESSION_ECS_RUNNER_NAME")) + parser.add_argument( + "--sweep", + action="store_true", + help="Janitor mode: delete all managed instances older than --sweep-ttl-hours.", + ) + parser.add_argument("--sweep-ttl-hours", type=float, default=SWEEP_TTL.total_seconds() / 3600) + args = parser.parse_args() + + for name in ("region_id", "repo", "github_token"): + if not getattr(args, name): + raise SystemExit(f"Missing required configuration: --{name.replace('_', '-')}") + + client = make_ecs_client(args.region_id) + + if args.sweep: + return sweep( + client, args.region_id, args.repo, args.github_token, timedelta(hours=args.sweep_ttl_hours) + ) + + ok = True + if args.instance_id: + ok &= delete_instance(client, args.instance_id, args.region_id) + else: + print("No instance id given; skipping instance deletion", flush=True) + if args.runner_name: + ok &= deregister_runner(args.github_token, args.repo, args.runner_name) + else: + print("No runner name given; skipping runner deregistration", flush=True) + return 0 if ok else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/scripts/query-regression-nightly-refs.py b/.github/scripts/query-regression-nightly-refs.py new file mode 100644 index 00000000000..6e9f54682d4 --- /dev/null +++ b/.github/scripts/query-regression-nightly-refs.py @@ -0,0 +1,270 @@ +#!/usr/bin/env python3 +# 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. + +"""Resolve consecutive successful Nightly Build SHAs for query-regression. + +Query regression builds both binaries from git; this script only picks the +head SHAs of two Nightly Build workflow runs (today vs the previous success). +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import urllib.error +import urllib.parse +import urllib.request +from dataclasses import dataclass +from typing import Any + + +NIGHTLY_WORKFLOW = "nightly-build.yml" +NIGHTLY_EVENTS = frozenset({"schedule", "workflow_dispatch"}) + + +@dataclass(frozen=True) +class WorkflowRun: + id: int + head_sha: str + head_branch: str + html_url: str + created_at: str + conclusion: str + event: str + + +@dataclass(frozen=True) +class RefPair: + base: WorkflowRun | None + candidate: WorkflowRun | None + skip: bool + reason: str + + +def parse_run(payload: dict[str, Any]) -> WorkflowRun: + return WorkflowRun( + id=int(payload["id"]), + head_sha=str(payload["head_sha"]), + head_branch=str(payload.get("head_branch") or ""), + html_url=str(payload.get("html_url") or ""), + created_at=str(payload.get("created_at") or ""), + conclusion=str(payload.get("conclusion") or ""), + event=str(payload.get("event") or ""), + ) + + +def github_get(token: str, path: str, query: dict[str, str] | None = None) -> dict[str, Any]: + encoded = f"?{urllib.parse.urlencode(query)}" if query else "" + request = urllib.request.Request( + f"https://api.github.com{path}{encoded}", + method="GET", + headers={ + "Accept": "application/vnd.github+json", + "Authorization": f"Bearer {token}", + "X-GitHub-Api-Version": "2022-11-28", + }, + ) + try: + with urllib.request.urlopen(request, timeout=30) as response: + return json.loads(response.read().decode("utf-8")) + except urllib.error.HTTPError as error: + body = error.read().decode("utf-8", "replace") + raise SystemExit( + f"GitHub API GET {path} failed: HTTP {error.code}: {body}" + ) from error + + +def list_successful_nightly_runs( + token: str, + repo: str, + *, + branch: str | None = None, + per_page: int = 30, +) -> list[WorkflowRun]: + query = { + "status": "success", + "per_page": str(per_page), + } + if branch: + query["branch"] = branch + payload = github_get( + token, + f"/repos/{repo}/actions/workflows/{NIGHTLY_WORKFLOW}/runs", + query, + ) + runs = [parse_run(item) for item in payload.get("workflow_runs") or []] + return [ + run + for run in runs + if run.conclusion == "success" and run.event in NIGHTLY_EVENTS and run.head_sha + ] + + +def fetch_run(token: str, repo: str, run_id: int) -> WorkflowRun: + payload = github_get(token, f"/repos/{repo}/actions/runs/{run_id}") + return parse_run(payload) + + +def select_base_and_candidate( + runs: list[WorkflowRun], + *, + candidate_run_id: int | None = None, + candidate: WorkflowRun | None = None, +) -> RefPair: + """Pick candidate (newest / requested) and the previous successful run. + + `runs` must be newest-first. Same-SHA consecutive nightlies are skipped + (no commits that day). + """ + if candidate is None: + if candidate_run_id is not None: + candidate = next((run for run in runs if run.id == candidate_run_id), None) + if candidate is None: + return RefPair( + None, + None, + True, + f"candidate nightly run {candidate_run_id} was not a successful " + f"{NIGHTLY_WORKFLOW} run", + ) + elif runs: + candidate = runs[0] + else: + return RefPair(None, None, True, "no successful Nightly Build runs") + + branch = candidate.head_branch + older = [ + run + for run in runs + if run.id != candidate.id + and (not branch or not run.head_branch or run.head_branch == branch) + and (not candidate.created_at or run.created_at <= candidate.created_at) + ] + if not older: + return RefPair( + None, + candidate, + True, + "no previous successful Nightly Build to compare against", + ) + base = older[0] + if base.head_sha.lower() == candidate.head_sha.lower(): + return RefPair( + base, + candidate, + True, + f"previous nightly SHA {base.head_sha} matches candidate; nothing to compare", + ) + return RefPair(base, candidate, False, "") + + +def _override_run(sha: str) -> WorkflowRun: + return WorkflowRun( + id=0, + head_sha=sha, + head_branch="", + html_url="", + created_at="", + conclusion="success", + event="workflow_dispatch", + ) + + +def override_pair(base_ref: str, candidate_ref: str) -> RefPair: + return RefPair(_override_run(base_ref), _override_run(candidate_ref), False, "explicit refs") + + +def write_outputs(pair: RefPair) -> None: + values = { + "skip": "true" if pair.skip else "false", + "reason": pair.reason, + "base_sha": pair.base.head_sha if pair.base else "", + "candidate_sha": pair.candidate.head_sha if pair.candidate else "", + "base_run_url": pair.base.html_url if pair.base else "", + "candidate_run_url": pair.candidate.html_url if pair.candidate else "", + "base_run_id": str(pair.base.id) if pair.base and pair.base.id else "", + "candidate_run_id": str(pair.candidate.id) if pair.candidate and pair.candidate.id else "", + } + output_path = os.environ.get("GITHUB_OUTPUT") + if output_path: + with open(output_path, "a", encoding="utf-8") as handle: + for key, value in values.items(): + handle.write(f"{key}={value}\n") + for key, value in values.items(): + print(f"{key}={value}") + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--repo", default=os.environ.get("GITHUB_REPOSITORY")) + parser.add_argument( + "--token", + default=os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN"), + ) + parser.add_argument("--branch", default=os.environ.get("NIGHTLY_BRANCH") or "") + parser.add_argument( + "--candidate-run-id", + default=os.environ.get("CANDIDATE_RUN_ID") or "", + ) + parser.add_argument("--base-ref", default=os.environ.get("BASE_REF") or "") + parser.add_argument("--candidate-ref", default=os.environ.get("CANDIDATE_REF") or "") + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + base_ref = args.base_ref.strip() + candidate_ref = args.candidate_ref.strip() + if base_ref and candidate_ref: + write_outputs(override_pair(base_ref, candidate_ref)) + return 0 + if bool(base_ref) != bool(candidate_ref): + print("Provide both --base-ref and --candidate-ref, or neither.", file=sys.stderr) + return 2 + if not args.repo or not args.token: + print("--repo and --token (or GITHUB_REPOSITORY / GITHUB_TOKEN) are required.", file=sys.stderr) + return 2 + + candidate_run_id = int(args.candidate_run_id) if str(args.candidate_run_id).strip() else None + candidate: WorkflowRun | None = None + branch = args.branch.strip() or None + if candidate_run_id is not None: + candidate = fetch_run(args.token, args.repo, candidate_run_id) + if candidate.conclusion != "success": + write_outputs( + RefPair( + None, + candidate, + True, + f"nightly run {candidate_run_id} conclusion is {candidate.conclusion}", + ) + ) + return 0 + branch = branch or candidate.head_branch or None + + runs = list_successful_nightly_runs(args.token, args.repo, branch=branch) + pair = select_base_and_candidate( + runs, + candidate_run_id=candidate_run_id if candidate is None else None, + candidate=candidate, + ) + write_outputs(pair) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 7f064f52eb5..8d7a37080b3 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -80,6 +80,14 @@ jobs: run: python3 scripts/check-enterprise-license-test.py - name: Run fuzz orchestration script tests run: .github/scripts/run-fuzz-targets-test.sh + - name: Run query-regression Python tooling tests + run: | + python3 tests/perf/test_query_regression_runner_compaction_toctou.py + python3 tests/perf/test_query_regression_runner_otlp_trace_load.py + python3 tests/perf/test_query_regression_summary_otlp.py + python3 tests/perf/test_query_regression_case_selection.py + python3 tests/perf/test_query_regression_nightly_refs.py + python3 tests/perf/test_aliyun_ecs_runner_scripts.py check: if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} diff --git a/.github/workflows/query-regression-janitor.yml b/.github/workflows/query-regression-janitor.yml new file mode 100644 index 00000000000..ce277ce2f02 --- /dev/null +++ b/.github/workflows/query-regression-janitor.yml @@ -0,0 +1,36 @@ +name: Query Regression ECS Janitor + +# Safety net for the aliyun-ecs query-regression path: deletes managed ECS +# instances (and their runner registrations) that outlived their workflow run, +# so a failed teardown never leaks pay-as-you-go spend. + +on: + schedule: + - cron: "23 3 * * *" + workflow_dispatch: + +permissions: + contents: read + +jobs: + sweep: + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - name: Checkout teardown script + uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Install uv + uses: astral-sh/setup-uv@v6 + + - name: Sweep expired ECS runners + env: + ALIBABA_CLOUD_ACCESS_KEY_ID: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + ALIBABA_CLOUD_ACCESS_KEY_SECRET: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + ALIYUN_ECS_REGION_ID: ${{ vars.ALIYUN_ECS_REGION_ID }} + run: >- + uv run .github/scripts/aliyun-ecs-runner-teardown.py + --sweep --sweep-ttl-hours 4 diff --git a/.github/workflows/query-regression-nightly.yml b/.github/workflows/query-regression-nightly.yml new file mode 100644 index 00000000000..027fcd16987 --- /dev/null +++ b/.github/workflows/query-regression-nightly.yml @@ -0,0 +1,104 @@ +name: Query Regression Nightly + +# After a successful GreptimeDB Nightly Build, compare that commit against +# the previous successful nightly. The reusable Query Regression workflow +# still compiles both SHAs; this wrapper only resolves which two SHAs. + +on: + workflow_run: + workflows: + - GreptimeDB Nightly Build + types: + - completed + workflow_dispatch: + inputs: + base_ref: + description: Base ref/SHA (empty = previous successful nightly) + required: false + type: string + default: "" + candidate_ref: + description: Candidate ref/SHA (empty = latest successful nightly) + required: false + type: string + default: "" + candidate_run_id: + description: Nightly Build run id to treat as candidate (empty = latest) + required: false + type: string + default: "" + case: + description: Query perf case path(s), all, or heavy + required: false + type: string + default: all + +permissions: + contents: read + actions: read + +jobs: + resolve-refs: + name: Resolve previous vs current nightly SHAs + if: ${{ github.event_name == 'workflow_dispatch' || github.event.workflow_run.conclusion == 'success' }} + runs-on: ubuntu-latest + timeout-minutes: 5 + outputs: + skip: ${{ steps.refs.outputs.skip }} + reason: ${{ steps.refs.outputs.reason }} + base_sha: ${{ steps.refs.outputs.base_sha }} + candidate_sha: ${{ steps.refs.outputs.candidate_sha }} + steps: + - name: Checkout ref resolver + uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Resolve nightly SHAs + id: refs + env: + GITHUB_TOKEN: ${{ github.token }} + GITHUB_REPOSITORY: ${{ github.repository }} + CANDIDATE_RUN_ID: ${{ github.event_name == 'workflow_run' && github.event.workflow_run.id || inputs.candidate_run_id || '' }} + BASE_REF: ${{ github.event_name == 'workflow_dispatch' && inputs.base_ref || '' }} + CANDIDATE_REF: ${{ github.event_name == 'workflow_dispatch' && inputs.candidate_ref || '' }} + run: python3 .github/scripts/query-regression-nightly-refs.py + + - name: Summarize comparison + if: always() + env: + SKIP: ${{ steps.refs.outputs.skip }} + REASON: ${{ steps.refs.outputs.reason }} + BASE_SHA: ${{ steps.refs.outputs.base_sha }} + CANDIDATE_SHA: ${{ steps.refs.outputs.candidate_sha }} + BASE_RUN_URL: ${{ steps.refs.outputs.base_run_url }} + CANDIDATE_RUN_URL: ${{ steps.refs.outputs.candidate_run_url }} + run: | + set -euo pipefail + { + if [[ "${SKIP}" == "true" ]]; then + printf 'Skipping query-regression nightly: %s\n' "${REASON}" + else + printf 'Comparing previous nightly `%s` -> current nightly `%s`\n' \ + "${BASE_SHA}" "${CANDIDATE_SHA}" + if [[ -n "${BASE_RUN_URL}" ]]; then + printf -- '- Previous Nightly Build: %s\n' "${BASE_RUN_URL}" + fi + if [[ -n "${CANDIDATE_RUN_URL}" ]]; then + printf -- '- Current Nightly Build: %s\n' "${CANDIDATE_RUN_URL}" + fi + fi + } | tee -a "${GITHUB_STEP_SUMMARY}" + + query-regression: + name: Query regression nightly + needs: [resolve-refs] + if: ${{ needs.resolve-refs.outputs.skip != 'true' && needs.resolve-refs.outputs.base_sha != '' && needs.resolve-refs.outputs.candidate_sha != '' }} + uses: ./.github/workflows/query-regression.yml + secrets: inherit + with: + case: ${{ github.event_name == 'workflow_dispatch' && inputs.case || 'all' }} + base_ref: ${{ needs.resolve-refs.outputs.base_sha }} + candidate_ref: ${{ needs.resolve-refs.outputs.candidate_sha }} + cargo_profile: nightly + runner: aliyun-ecs diff --git a/.github/workflows/query-regression.yml b/.github/workflows/query-regression.yml index 0806d4d3e38..209c8d76fff 100644 --- a/.github/workflows/query-regression.yml +++ b/.github/workflows/query-regression.yml @@ -32,10 +32,10 @@ on: type: string default: nightly runner: - description: Self-hosted runner label or ARC runner scale set for this query regression run + description: Self-hosted runner label; aliyun-ecs provisions a fresh ECS instance per run required: false type: string - default: perf-regression-8-cores + default: aliyun-ecs workflow_dispatch: inputs: case: @@ -68,12 +68,20 @@ on: - release - dev runner: - description: Self-hosted runner label or ARC runner scale set for this query regression run + description: >- + Self-hosted runner label; aliyun-ecs provisions a fresh ECS instance + per run, any other value is used as a literal runner label required: true - type: choice - default: perf-regression-8-cores - options: - - perf-regression-8-cores + type: string + default: aliyun-ecs + keep_instance: + description: >- + Debug: keep the ECS instance after the run (skip teardown) so its + runner _diag logs, journal, and telemetry can be inspected; the + janitor still sweeps it after the TTL + required: false + type: boolean + default: false pull_request: types: [labeled] @@ -81,19 +89,99 @@ permissions: contents: read jobs: - query-regression: + test-tooling: + # Stdlib unittests for the workflow Python (case selection, report + # helpers, nightly SHA picking, and rendered ECS user-data). They do + # not talk to Aliyun or the Actions runner process; running them on + # ubuntu-latest fails fast before any ECS spend. Not gated on the + # regression labels: those labels boot a VM, these tests should not. + # Ordinary PRs also run the same tests from checks.yml, because this + # workflow only starts on `labeled` (or dispatch / workflow_call). + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Test query regression tooling + run: | + python3 tests/perf/test_query_regression_runner_compaction_toctou.py + python3 tests/perf/test_query_regression_runner_otlp_trace_load.py + python3 tests/perf/test_query_regression_summary_otlp.py + python3 tests/perf/test_query_regression_case_selection.py + python3 tests/perf/test_query_regression_nightly_refs.py + python3 tests/perf/test_aliyun_ecs_runner_scripts.py + + provision: + # Runs when the aliyun-ecs path is selected: explicitly via the runner + # input, or for PR labels by default (a QUERY_REGRESSION_PR_RUNNER + # repository variable set to another value redirects PRs to that literal + # runner label instead). Uses trusted scripts from the PR base (or the + # dispatched ref), never from candidate code. Waits for test-tooling so + # a broken user-data template does not still create a VM. + needs: [test-tooling] if: >- - ${{ github.event_name != 'pull_request' || - (github.event_name == 'pull_request' && - !github.event.pull_request.draft && - (github.event.label.name == 'query-regression' || - github.event.label.name == 'heavy-regression')) }} - runs-on: ${{ github.event_name != 'pull_request' && inputs.runner || 'perf-regression-8-cores' }} + ${{ !failure() && !cancelled() && + ((github.event_name != 'pull_request' && inputs.runner == 'aliyun-ecs') || + (github.event_name == 'pull_request' && + !github.event.pull_request.draft && + (github.event.label.name == 'query-regression' || + github.event.label.name == 'heavy-regression') && + (vars.QUERY_REGRESSION_PR_RUNNER || 'aliyun-ecs') == 'aliyun-ecs')) }} + runs-on: ubuntu-latest + timeout-minutes: 45 + outputs: + label: ${{ steps.provision.outputs.label }} + instance_id: ${{ steps.provision.outputs.instance_id }} + runner_name: ${{ steps.provision.outputs.runner_name }} + steps: + - name: Checkout trusted provisioning scripts + uses: actions/checkout@v4 + with: + repository: ${{ github.repository }} + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || github.sha }} + persist-credentials: false + + - name: Install uv + uses: astral-sh/setup-uv@v6 + + - name: Provision ECS runner + id: provision + env: + ALIBABA_CLOUD_ACCESS_KEY_ID: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + ALIBABA_CLOUD_ACCESS_KEY_SECRET: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + ALIYUN_ECS_REGION_ID: ${{ vars.ALIYUN_ECS_REGION_ID }} + ALIYUN_ECS_VSWITCH_ID: ${{ vars.ALIYUN_ECS_VSWITCH_ID }} + ALIYUN_ECS_SECURITY_GROUP_ID: ${{ vars.ALIYUN_ECS_SECURITY_GROUP_ID }} + ALIYUN_ECS_INSTANCE_TYPE: ${{ vars.ALIYUN_ECS_INSTANCE_TYPE }} + ALIYUN_ECS_RESOURCE_GROUP_ID: ${{ vars.ALIYUN_ECS_RESOURCE_GROUP_ID }} + QUERY_REGRESSION_ECS_IMAGE_ID: ${{ vars.QUERY_REGRESSION_ECS_IMAGE_ID }} + QUERY_REGRESSION_RUNNER_UID: ${{ vars.QUERY_REGRESSION_RUNNER_UID || '1001' }} + QUERY_REGRESSION_RUNNER_GID: ${{ vars.QUERY_REGRESSION_RUNNER_GID || '1001' }} + run: >- + uv run .github/scripts/aliyun-ecs-runner-provision.py + + query-regression: + needs: [provision, test-tooling] + # `!failure() && !cancelled()` both suppresses the implicit success() and + # lets the job run when provision was intentionally skipped because a + # literal runner label was selected; a failed or cancelled provision still + # blocks the run because no ECS runner would be waiting. + if: >- + ${{ !failure() && !cancelled() && + (github.event_name != 'pull_request' || + (github.event_name == 'pull_request' && + !github.event.pull_request.draft && + (github.event.label.name == 'query-regression' || + github.event.label.name == 'heavy-regression'))) }} + runs-on: >- + ${{ needs.provision.outputs.label || + (github.event_name != 'pull_request' && inputs.runner || + (vars.QUERY_REGRESSION_PR_RUNNER || 'aliyun-ecs')) }} timeout-minutes: 180 - concurrency: - group: query-regression-persistent-cache-v1 - queue: max - cancel-in-progress: false env: CARGO_PROFILE: ${{ github.event_name == 'pull_request' && 'nightly' || inputs.cargo_profile }} CARGO_HOME: /home/runner/.cargo @@ -105,7 +193,8 @@ jobs: QUERY_REGRESSION_CACHE_META: /home/runner/query-regression-cache-meta RUSTC_WRAPPER: /usr/local/bin/sccache SCCACHE_DIR: /home/runner/.cache/sccache - SCCACHE_CACHE_SIZE: 40G + # Caps the local sccache on the system disk next to target and cargo. + SCCACHE_CACHE_SIZE: 10G CARGO_INCREMENTAL: "0" RUSTFLAGS: -D warnings -C link-arg=-fuse-ld=mold QUERY_REGRESSION_CACHE_EPOCH: "1" @@ -114,7 +203,55 @@ jobs: EVENT_MERGE_SHA: ${{ github.event_name == 'pull_request' && github.sha || '' }} EVENT_HEAD_SHA: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || '' }} EVENT_BASE_SHA: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || '' }} + # Runner identity contract. The ECS image uses 1001; a manually prepared + # host may override via repo variables when 1001 is already taken. + EXPECTED_RUNNER_UID: ${{ vars.QUERY_REGRESSION_RUNNER_UID || '1001' }} + EXPECTED_RUNNER_GID: ${{ vars.QUERY_REGRESSION_RUNNER_GID || '1001' }} steps: + - name: Report provisioned ECS runner + if: ${{ needs.provision.outputs.instance_id != '' }} + shell: bash + run: | + set -euo pipefail + { + printf -- '- ECS instance: `%s` (type `%s`, image `%s`, region `%s`)\n' \ + "${{ needs.provision.outputs.instance_id }}" \ + "${{ vars.ALIYUN_ECS_INSTANCE_TYPE }}" \ + "${{ vars.QUERY_REGRESSION_ECS_IMAGE_ID }}" \ + "${{ vars.ALIYUN_ECS_REGION_ID }}" + } >> "${GITHUB_STEP_SUMMARY}" + + - name: Start machine telemetry sampler + shell: bash + # The ECS instance is deleted when the run ends, so machine state must + # be captured while the job runs. A background sampler appends load, + # memory, disk, and top-memory processes to a log every 30s; the + # "Dump machine telemetry" step (if: always()) prints and uploads it. + # Note: when the job is *cancelled*, even always() steps are killed — + # telemetry covers failures, not cancellations. + run: | + set -euo pipefail + log="${GITHUB_WORKSPACE}/machine-telemetry.log" + { + echo "== baseline $(date -u +%Y-%m-%dT%H:%M:%SZ) ==" + nproc + free -h + df -h / /home/runner + } >> "${log}" + nohup bash -c " + while true; do + { + date -u '+== %Y-%m-%dT%H:%M:%SZ ==' + uptime + free -m + df -h / /home/runner + ps -eo pid,comm,%mem,%cpu --sort=-%mem | head -6 + } >> '${log}' 2>&1 + sleep 30 + done + " >/dev/null 2>&1 & + echo "Telemetry sampler started (pid $!), logging to ${log}" + - name: Checkout base source uses: actions/checkout@v4 with: @@ -263,14 +400,64 @@ jobs: shell: bash run: | set -euo pipefail - [[ "$(id -u)" == "1001" ]] - [[ "$(id -g)" == "1001" ]] - [[ "${UV_CACHE_DIR}" == "/home/runner/.cargo/uv-cache" ]] - [[ "$(protoc --version)" == "libprotoc 3.21.12" ]] + errors=() + record() { + printf 'FAIL %s\n' "$*" + errors+=("$*") + } + require_eq() { + local name="$1" actual="$2" expected="$3" + printf 'check %s: %s\n' "${name}" "${actual}" + [[ "${actual}" == "${expected}" ]] || record "${name}: expected '${expected}', got '${actual}'" + } + require_match() { + local name="$1" actual="$2" pattern="$3" + printf 'check %s: %s\n' "${name}" "${actual}" + [[ "${actual}" =~ $pattern ]] || record "${name}: expected to match ${pattern}, got '${actual}'" + } + require() { + local name="$1" + shift + printf 'check %s\n' "${name}" + "$@" || record "${name}" + } + capture() { + local out + if out="$("$@" 2>&1)"; then + printf '%s' "${out}" + else + printf '' "$*" + fi + } + + printf 'uid=%s gid=%s PATH=%s CARGO_HOME=%s\n' "$(id -u)" "$(id -g)" "${PATH}" "${CARGO_HOME}" + ls -la "${CARGO_HOME}" 2>&1 || printf '(CARGO_HOME missing)\n' + + # Image/host hygiene: CARGO_HOME must be empty of config and + # install-state before cargo/rustup run. Those tools create + # .package-cache / bin on first use, so this cannot come after. + cargo_home_dirty=() + for entry in config config.toml credentials credentials.toml bin .crates.toml .crates2.json .global-cache .package-cache; do + if [[ -e "${CARGO_HOME}/${entry}" ]]; then + cargo_home_dirty+=("${entry}") + record "CARGO_HOME must not contain ${entry}" + fi + done + if (( ${#cargo_home_dirty[@]} > 0 )); then + printf 'CARGO_HOME contents:\n' >&2 + ls -la "${CARGO_HOME}" >&2 || true + fi + + require_eq uid "$(id -u)" "${EXPECTED_RUNNER_UID}" + require_eq gid "$(id -g)" "${EXPECTED_RUNNER_GID}" + require_eq UV_CACHE_DIR "${UV_CACHE_DIR}" "/home/runner/.cargo/uv-cache" + require_eq protoc "$(capture protoc --version)" "libprotoc 3.21.12" temporary_proto_dir="$(mktemp --directory)" trap 'rm -rf "${temporary_proto_dir}"' EXIT - test -r /usr/include/google/protobuf/any.proto - test -r /usr/include/google/protobuf/empty.proto + require "readable /usr/include/google/protobuf/any.proto" \ + test -r /usr/include/google/protobuf/any.proto + require "readable /usr/include/google/protobuf/empty.proto" \ + test -r /usr/include/google/protobuf/empty.proto printf '%s\n' \ 'syntax = "proto3";' \ 'package smoke;' \ @@ -278,47 +465,51 @@ jobs: 'import "google/protobuf/empty.proto";' \ 'message Smoke { google.protobuf.Any any = 1; google.protobuf.Empty empty = 2; }' \ > "${temporary_proto_dir}/smoke.proto" - protoc --proto_path="${temporary_proto_dir}" --proto_path=/usr/include \ + if protoc --proto_path="${temporary_proto_dir}" --proto_path=/usr/include \ --descriptor_set_out="${temporary_proto_dir}/smoke.pb" \ - "${temporary_proto_dir}/smoke.proto" - test -s "${temporary_proto_dir}/smoke.pb" - [[ "$(uv --version)" =~ ^uv[[:space:]]0\.11\.26([[:space:]]|$) ]] - mold_version="$(mold --version)" - [[ "${mold_version}" =~ ^mold[[:space:]]2\.30\.0([[:space:]]|$) ]] - [[ "$(python3 --version)" == "Python 3.12.3" ]] + "${temporary_proto_dir}/smoke.proto"; then + require "protoc smoke descriptor is non-empty" test -s "${temporary_proto_dir}/smoke.pb" + else + record "protoc smoke compile failed" + fi + require_match uv "$(capture uv --version)" '^uv[[:space:]]0\.11\.26([[:space:]]|$)' + require_match mold "$(capture mold --version)" '^mold[[:space:]]2\.40\.4([[:space:]]|$)' + require_eq python3 "$(capture python3 --version)" "Python 3.14.4" otelgen_path="$(command -v otelgen || true)" otelgen_version="" if [[ -n "${otelgen_path}" ]]; then - otelgen_version="$("${otelgen_path}" --version 2>&1 || true)" + otelgen_version="$(capture "${otelgen_path}" --version)" fi - printf 'otelgen path: %s\n' "${otelgen_path:-}" - printf 'otelgen version: %s\n' "${otelgen_version}" - [[ "${otelgen_path}" == "/usr/local/bin/otelgen" ]] - [[ "${otelgen_version}" == *"863a3f395d062c7322cc1de08a38774b7fdaa6c8"* ]] - sccache_version="$(sccache --version)" - [[ "${sccache_version}" =~ ^sccache[[:space:]]0\.16\.0([[:space:]]|$) ]] - [[ "$(command -v rustup)" == "/opt/cargo/bin/rustup" ]] - [[ "$(command -v cargo)" == "/opt/cargo/bin/cargo" ]] - [[ "$(command -v rustc)" == "/opt/cargo/bin/rustc" ]] - [[ "$(rustup --version)" =~ ^rustup[[:space:]]1\.29\.0([[:space:]]|$) ]] - cargo_version="$(cargo --version)" - [[ "${cargo_version}" =~ ^cargo[[:space:]]1\.96\.0-nightly[[:space:]]\(cbb9bb8bd[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$ ]] - rustc_version="$(rustc --version)" - [[ "${rustc_version}" =~ ^rustc[[:space:]]1\.96\.0-nightly[[:space:]]\(ac7f9ec7d[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$ ]] - active_toolchain="$(rustup show active-toolchain)" - [[ "${active_toolchain}" =~ ^nightly-2026-03-21-x86_64-unknown-linux-gnu([[:space:]]|$) ]] - [[ "${RUSTUP_HOME}" == "/opt/rustup" ]] - [[ "${RUSTUP_TOOLCHAIN}" == "nightly-2026-03-21" ]] - [[ "${RUSTUP_AUTO_INSTALL}" == "0" ]] - test -r /opt/rustup && test -x /opt/rustup - test ! -w /opt/rustup - test ! -w /opt/cargo/bin - for entry in config config.toml credentials credentials.toml bin .crates.toml .crates2.json .global-cache .package-cache; do - test ! -e "${CARGO_HOME}/${entry}" - done + require_eq otelgen_path "${otelgen_path}" "/usr/local/bin/otelgen" + printf 'check otelgen_version: %s\n' "${otelgen_version}" + [[ "${otelgen_version}" == *"863a3f395d062c7322cc1de08a38774b7fdaa6c8"* ]] \ + || record "otelgen_version: expected commit 863a3f395d062c7322cc1de08a38774b7fdaa6c8, got '${otelgen_version}'" + require_match sccache "$(capture sccache --version)" '^sccache[[:space:]]0\.16\.0([[:space:]]|$)' + require_eq rustup_path "$(command -v rustup || true)" "/opt/cargo/bin/rustup" + require_eq cargo_path "$(command -v cargo || true)" "/opt/cargo/bin/cargo" + require_eq rustc_path "$(command -v rustc || true)" "/opt/cargo/bin/rustc" + require_match rustup "$(capture rustup --version)" '^rustup[[:space:]]1\.29\.0([[:space:]]|$)' + require_match cargo "$(capture cargo --version)" \ + '^cargo[[:space:]]1\.96\.0-nightly[[:space:]]\(cbb9bb8bd[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$' + require_match rustc "$(capture rustc --version)" \ + '^rustc[[:space:]]1\.96\.0-nightly[[:space:]]\(ac7f9ec7d[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$' + require_match active_toolchain "$(capture rustup show active-toolchain)" \ + '^nightly-2026-03-21-x86_64-unknown-linux-gnu([[:space:]]|$)' + require_eq RUSTUP_HOME "${RUSTUP_HOME}" "/opt/rustup" + require_eq RUSTUP_TOOLCHAIN "${RUSTUP_TOOLCHAIN}" "nightly-2026-03-21" + require_eq RUSTUP_AUTO_INSTALL "${RUSTUP_AUTO_INSTALL}" "0" + require "readable /opt/rustup" test -r /opt/rustup + require "executable /opt/rustup" test -x /opt/rustup + require "runner cannot write /opt/rustup" test ! -w /opt/rustup + require "runner cannot write /opt/cargo/bin" test ! -w /opt/cargo/bin mkdir -p "${CARGO_HOME}/registry" "${CARGO_HOME}/git" + if (( ${#errors[@]} > 0 )); then + printf '\nVerify runner image tools failed (%d checks):\n' "${#errors[@]}" >&2 + printf ' - %s\n' "${errors[@]}" >&2 + exit 1 + fi - - name: Prepare persistent query regression cache + - name: Prepare query regression cache shell: bash working-directory: src run: | @@ -362,24 +553,6 @@ jobs: find "${root}" -mindepth 1 -maxdepth 1 -exec rm -rf -- {} + } - clear_cargo_extracted_trees() { - require_expected_root CARGO_REGISTRY "${CARGO_HOME}/registry" "${EXPECTED_CARGO_REGISTRY}" - require_expected_root CARGO_GIT "${CARGO_HOME}/git" "${EXPECTED_CARGO_GIT}" - rm -rf -- "${CARGO_HOME}/registry/src" "${CARGO_HOME}/git/checkouts" - } - - cargo_size_kib() { - du -sk -- "${CARGO_HOME}/registry" "${CARGO_HOME}/git" | awk '{ total += $1 } END { print total }' - } - - target_size_kib() { - du -sk -- "${CARGO_TARGET_DIR}" | cut -f1 - } - - free_kib() { - df -Pk "${CARGO_TARGET_DIR}" | awk 'NR == 2 { print $4 }' - } - report_cache_usage() { du -sh -- "${CARGO_HOME}" "${CARGO_HOME}/registry" "${CARGO_HOME}/git" \ "${RUSTUP_HOME}" "${CARGO_TARGET_DIR}" "${QUERY_REGRESSION_CACHE_META}" "${SCCACHE_DIR}" @@ -400,7 +573,7 @@ jobs: printf 'Refusing unexpected RUSTC_WRAPPER: %s\n' "${RUSTC_WRAPPER}" >&2 exit 1 } - [[ "${SCCACHE_CACHE_SIZE}" == "40G" ]] || { + [[ "${SCCACHE_CACHE_SIZE}" == "10G" ]] || { printf 'Refusing unexpected SCCACHE_CACHE_SIZE: %s\n' "${SCCACHE_CACHE_SIZE}" >&2 exit 1 } @@ -487,46 +660,10 @@ jobs: report_cache_usage - target_size="$(target_size_kib)" - if (( target_size >= 400 * 1024 * 1024 )); then - printf 'Warning: target cache is at least 400 GiB (%s KiB)\n' "${target_size}" >&2 - fi - if (( target_size >= 450 * 1024 * 1024 )); then - printf 'Target cache is at least 450 GiB; clearing complete target root\n' >&2 - clear_directory "${CARGO_TARGET_DIR}" - fi - - cargo_size="$(cargo_size_kib)" - if (( cargo_size >= 60 * 1024 * 1024 )); then - printf 'Warning: Cargo cache is at least 60 GiB (%s KiB)\n' "${cargo_size}" >&2 - fi - if (( cargo_size >= 80 * 1024 * 1024 )); then - printf 'Cargo cache is at least 80 GiB; removing extracted sources and checkouts\n' >&2 - clear_cargo_extracted_trees - cargo_size="$(cargo_size_kib)" - if (( cargo_size >= 80 * 1024 * 1024 )); then - printf 'Cargo cache remains at least 80 GiB after cleanup (%s KiB)\n' "${cargo_size}" >&2 - exit 1 - fi - fi - - free_space="$(free_kib)" - if (( free_space < 300 * 1024 * 1024 )); then - printf 'Free space is below 300 GiB; clearing complete target root\n' >&2 - clear_directory "${CARGO_TARGET_DIR}" - free_space="$(free_kib)" - if (( free_space < 300 * 1024 * 1024 )); then - printf 'Free space remains below 300 GiB; removing Cargo extracted sources and checkouts\n' >&2 - clear_cargo_extracted_trees - free_space="$(free_kib)" - if (( free_space < 300 * 1024 * 1024 )); then - printf 'Free space remains below 300 GiB after cleanup (%s KiB)\n' "${free_space}" >&2 - exit 1 - fi - fi - fi - - report_cache_usage + # Every run starts from a fresh system disk, so size/free-space + # watermarks from the retired retained-disk era are not restored. + # A run that overflows the disk fails the build outright, which the + # telemetry step makes diagnosable. sccache --start-server sccache --zero-stats @@ -567,14 +704,6 @@ jobs: git reset --hard "${VERIFIED_CANDIDATE_SHA}" git clean -ffdx - - name: Test query regression tooling - working-directory: src - run: | - uv run --no-project python tests/perf/test_query_regression_runner_compaction_toctou.py - uv run --no-project python tests/perf/test_query_regression_runner_otlp_trace_load.py - uv run --no-project python tests/perf/test_query_regression_summary_otlp.py - uv run --no-project python tests/perf/test_query_regression_case_selection.py - - name: Build candidate greptime and query regression helpers working-directory: src run: | @@ -643,6 +772,7 @@ jobs: query-regression-work/**/logs/** query-regression-work/**/otelgen/** query-regression-summary.md + machine-telemetry.log if-no-files-found: warn retention-days: 7 @@ -657,7 +787,7 @@ jobs: if-no-files-found: warn retention-days: 7 - - name: Report persistent cache usage + - name: Report cache usage if: ${{ always() }} shell: bash run: | @@ -697,6 +827,55 @@ jobs: printf 'sccache is unavailable (report only)\n' >&2 fi + - name: Dump machine telemetry + if: ${{ always() }} + shell: bash + run: | + log="${GITHUB_WORKSPACE}/machine-telemetry.log" + if [[ -f "${log}" ]]; then + echo "::group::Machine telemetry (last 200 lines)" + tail -n 200 "${log}" + echo "::endgroup::" + else + echo "No telemetry log found (sampler never started?)" + fi + echo "::group::dmesg tail (OOM killer records)" + sudo dmesg -T 2>/dev/null | tail -n 50 || dmesg -T 2>/dev/null | tail -n 50 || \ + echo "dmesg unavailable without root" + echo "::endgroup::" + - name: Fail on regression failure if: ${{ steps.run.outputs.status != '0' }} run: exit 1 + + teardown: + # Releases the dynamically provisioned ECS runner. Runs even when the + # benchmark job fails or is cancelled; skipped when a literal runner label + # was selected because the provision outputs are empty, or when the + # dispatch set keep_instance to preserve the machine for post-mortem + # debugging (the janitor sweep still reclaims it after the TTL). + if: ${{ always() && needs.provision.outputs.instance_id != '' && !inputs.keep_instance }} + needs: [provision, query-regression] + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - name: Checkout trusted teardown scripts + uses: actions/checkout@v4 + with: + repository: ${{ github.repository }} + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || github.sha }} + persist-credentials: false + + - name: Install uv + uses: astral-sh/setup-uv@v6 + + - name: Teardown ECS runner + env: + ALIBABA_CLOUD_ACCESS_KEY_ID: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + ALIBABA_CLOUD_ACCESS_KEY_SECRET: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + ALIYUN_ECS_REGION_ID: ${{ vars.ALIYUN_ECS_REGION_ID }} + QUERY_REGRESSION_ECS_INSTANCE_ID: ${{ needs.provision.outputs.instance_id }} + QUERY_REGRESSION_ECS_RUNNER_NAME: ${{ needs.provision.outputs.runner_name }} + run: >- + uv run .github/scripts/aliyun-ecs-runner-teardown.py diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 52d62299a07..0a9c39c9034 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -510,7 +510,7 @@ jobs: base_ref: ${{ needs.prepare-release-validation.outputs.previous-release-tag }} candidate_ref: ${{ needs.prepare-release-validation.outputs.candidate-ref }} cargo_profile: nightly - runner: perf-regression-8-cores + runner: aliyun-ecs release-images-to-dockerhub: name: Build and push images to DockerHub diff --git a/tests/perf/AGENTS.md b/tests/perf/AGENTS.md index 1f13dea58dc..822b99080f9 100644 --- a/tests/perf/AGENTS.md +++ b/tests/perf/AGENTS.md @@ -3,9 +3,23 @@ - Keep GitHub Actions YAML thin. Put non-trivial control flow, case expansion, report generation, and metadata writing in scripts under `.github/scripts/`; workflow steps should mostly invoke those scripts. +- Runner lifecycle: the default path provisions one ephemeral Aliyun ECS + instance per run via `.github/scripts/aliyun-ecs-runner-provision.py` and + always releases it via `aliyun-ecs-runner-teardown.py`; a scheduled janitor + workflow sweeps leftovers. Build caches live on that instance's system disk + and are discarded with the VM. Runs do not share a workflow concurrency + group. The ECS custom image is built from the runner Dockerfile by + `.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py`; + keep the Dockerfile the single source of the tool contract. Dispatching with + any other `runner` value treats it as a literal self-hosted runner label + (see `ecs-image/bootstrap-runner-host.sh` for preparing such a host). - Query regression PR runs should build base/candidate binaries once, then run the default case set. Do not hard-code a single case such as `promql_pushdown_7913` into the workflow path. +- Scheduled nightly comparison lives in `query-regression-nightly.yml`: it + waits for a successful Nightly Build, then calls `query-regression.yml` + with the previous vs current nightly SHAs. Keep SHA selection in + `.github/scripts/query-regression-nightly-refs.py`. - The case DSL is not required to keep compatibility inside this PR. When the DSL changes, update TOML cases, the outer lifecycle script, Rust helpers, and docs together. @@ -19,8 +33,10 @@ - Keep the direct-SST generator generic. Issue-specific behavior belongs in case files and thresholds, not in Rust generator logic. - Before pushing perf harness changes, run at least: - - the Python tests in the `Test query regression tooling` step of - `.github/workflows/query-regression.yml` + - the Python tests in the `test-tooling` job of + `.github/workflows/query-regression.yml` (ubuntu-latest, not the ECS runner). + The Checks workflow runs the same tests on ordinary PRs so they are not + gated on the `query-regression` / `heavy-regression` labels. - `cargo fmt --all -- --check` - `cargo build -p cmd --bin query_perf_fixture --features dev-tools` - `cargo build -p cmd --bin query_regression_runner --features dev-tools` diff --git a/tests/perf/README.md b/tests/perf/README.md index cbe03b51794..47495ca531f 100644 --- a/tests/perf/README.md +++ b/tests/perf/README.md @@ -255,11 +255,12 @@ gh workflow run query-regression.yml \ -f candidate_ref= \ -f cargo_profile=nightly \ -f http_timeout=300 \ - -f runner=perf-regression-8-cores + -f runner=aliyun-ecs ``` -The selected ARC scale set must already be deployed with the runner-image -digest built from the current query-regression Dockerfile. +The `aliyun-ecs` path provisions a fresh ECS instance per run from the custom +image built from the current query-regression Dockerfile; see +`.github/runner-scale-sets/query-regression/README.md` for its configuration. ## Generator contract diff --git a/tests/perf/test_aliyun_ecs_runner_scripts.py b/tests/perf/test_aliyun_ecs_runner_scripts.py new file mode 100644 index 00000000000..581fc88abd0 --- /dev/null +++ b/tests/perf/test_aliyun_ecs_runner_scripts.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python3 +# 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. + +"""Coverage for the pure parts of the Aliyun ECS runner provision/teardown scripts.""" + +import base64 +import importlib.util +import sys +import unittest +from datetime import datetime, timedelta, timezone +from pathlib import Path + + +SCRIPTS_DIR = Path(__file__).parents[2] / ".github/scripts" + + +def load_module(name: str, filename: str): + spec = importlib.util.spec_from_file_location(name, SCRIPTS_DIR / filename) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +provision = load_module("aliyun_ecs_runner_provision_under_test", "aliyun-ecs-runner-provision.py") +teardown = load_module("aliyun_ecs_runner_teardown_under_test", "aliyun-ecs-runner-teardown.py") + + +class ProvisionNamingTest(unittest.TestCase): + def test_runner_name_and_label_derive_from_run_id(self) -> None: + self.assertEqual(provision.runner_name_for_run("12345"), "qreg-ecs-12345") + self.assertEqual(provision.runner_label_for_run("12345"), "query-regression-ecs-12345") + + +class ProvisionUserDataTest(unittest.TestCase): + def render(self) -> str: + return provision.render_user_data( + runner_name="qreg-ecs-12345", + runner_label="query-regression-ecs-12345", + runner_token="TOKEN", + repo="GreptimeTeam/greptimedb", + ) + + def test_user_data_creates_cache_paths_on_the_system_disk(self) -> None: + script = self.render() + self.assertNotIn("DISK_SERIAL", script) + self.assertNotIn("mount --bind", script) + self.assertNotIn("mkfs.ext4", script) + for destination in provision.CACHE_PATHS: + self.assertIn(f'"{destination}"', script) + + def test_user_data_wires_runner_registration(self) -> None: + script = self.render() + self.assertIn("RUNNER_NAME=qreg-ecs-12345", script) + self.assertIn("RUNNER_LABELS=query-regression-ecs-12345", script) + self.assertIn("RUNNER_TOKEN=TOKEN", script) + self.assertIn("REPO_URL=https://github.com/GreptimeTeam/greptimedb", script) + self.assertIn("PATH=/opt/cargo/bin:", script) + self.assertIn("systemctl restart --no-block ephemeral-github-runner.service", script) + + def test_encode_user_data_round_trips(self) -> None: + script = self.render() + self.assertEqual( + base64.b64decode(provision.encode_user_data(script)).decode("utf-8"), script + ) + + def test_user_data_enables_swap_and_masks_oomd(self) -> None: + script = self.render() + self.assertIn("systemctl mask systemd-oomd.socket systemd-oomd.service", script) + self.assertIn("systemctl mask unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer", script) + self.assertIn('APT::Periodic::Unattended-Upgrade "0"', script) + self.assertIn(f'fallocate --length {provision.SWAP_SIZE_GIB}G "{provision.SWAP_FILE}"', script) + self.assertIn(f'swapon "{provision.SWAP_FILE}"', script) + self.assertIn("sysctl --write vm.swappiness=10", script) + self.assertIn("OOMPolicy=continue", script) + self.assertNotIn("OOMScoreAdjust", script) + self.assertLess(script.index("swapon"), script.index("systemctl restart --no-block ephemeral-github-runner.service")) + + +class TeardownExpiryTest(unittest.TestCase): + NOW = datetime(2026, 8, 17, 6, 0, tzinfo=timezone.utc) + TTL = timedelta(hours=4) + + def test_parse_creation_time_formats(self) -> None: + self.assertEqual( + teardown.parse_creation_time("2026-08-17T01:02:03Z"), + datetime(2026, 8, 17, 1, 2, 3, tzinfo=timezone.utc), + ) + self.assertEqual( + teardown.parse_creation_time("2026-08-17T01:02Z"), + datetime(2026, 8, 17, 1, 2, tzinfo=timezone.utc), + ) + with self.assertRaises(ValueError): + teardown.parse_creation_time("not-a-time") + + def test_expired_instance_names_selects_only_old_instances(self) -> None: + instances = [ + ("i-old", "qreg-ecs-1", "2026-08-17T01:00Z"), # 5h old: expired + ("i-edge", "qreg-ecs-2", "2026-08-17T02:00Z"), # exactly TTL: expired + ("i-fresh", "qreg-ecs-3", "2026-08-17T05:30Z"), # 30m old: kept + ] + self.assertEqual( + teardown.expired_instance_names(instances, self.NOW, self.TTL), + [("i-old", "qreg-ecs-1"), ("i-edge", "qreg-ecs-2")], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/perf/test_query_regression_nightly_refs.py b/tests/perf/test_query_regression_nightly_refs.py new file mode 100644 index 00000000000..74b168bb464 --- /dev/null +++ b/tests/perf/test_query_regression_nightly_refs.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python3 +# 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. + +"""Unit tests for consecutive Nightly Build SHA selection.""" + +import importlib.util +import sys +import unittest +from pathlib import Path + + +SCRIPTS_DIR = Path(__file__).parents[2] / ".github/scripts" + + +def load_module(): + spec = importlib.util.spec_from_file_location( + "query_regression_nightly_refs_under_test", + SCRIPTS_DIR / "query-regression-nightly-refs.py", + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +refs = load_module() + + +def run(*, run_id: int, sha: str, branch: str = "main", created_at: str = "2026-08-24T00:00:00Z"): + return refs.WorkflowRun( + id=run_id, + head_sha=sha, + head_branch=branch, + html_url=f"https://github.com/example/run/{run_id}", + created_at=created_at, + conclusion="success", + event="schedule", + ) + + +class SelectNightlyRefsTest(unittest.TestCase): + def test_picks_newest_and_previous_on_same_branch(self) -> None: + pair = refs.select_base_and_candidate( + [ + run(run_id=3, sha="ccc", created_at="2026-08-22T00:00:00Z"), + run(run_id=2, sha="bbb", created_at="2026-08-21T00:00:00Z"), + run(run_id=1, sha="aaa", created_at="2026-08-20T00:00:00Z"), + ] + ) + self.assertFalse(pair.skip) + assert pair.candidate is not None and pair.base is not None + self.assertEqual(pair.candidate.head_sha, "ccc") + self.assertEqual(pair.base.head_sha, "bbb") + + def test_candidate_run_id_uses_that_run_and_the_next_older(self) -> None: + pair = refs.select_base_and_candidate( + [ + run(run_id=3, sha="ccc", created_at="2026-08-22T00:00:00Z"), + run(run_id=2, sha="bbb", created_at="2026-08-21T00:00:00Z"), + run(run_id=1, sha="aaa", created_at="2026-08-20T00:00:00Z"), + ], + candidate_run_id=2, + ) + self.assertFalse(pair.skip) + assert pair.candidate is not None and pair.base is not None + self.assertEqual(pair.candidate.id, 2) + self.assertEqual(pair.base.id, 1) + + def test_skips_when_consecutive_nightlies_share_a_sha(self) -> None: + pair = refs.select_base_and_candidate( + [ + run(run_id=2, sha="same", created_at="2026-08-22T00:00:00Z"), + run(run_id=1, sha="same", created_at="2026-08-21T00:00:00Z"), + ] + ) + self.assertTrue(pair.skip) + self.assertIn("matches candidate", pair.reason) + + def test_skips_other_branches_when_picking_previous(self) -> None: + pair = refs.select_base_and_candidate( + [ + run(run_id=3, sha="ccc", branch="main", created_at="2026-08-22T00:00:00Z"), + run(run_id=2, sha="other", branch="feat", created_at="2026-08-21T12:00:00Z"), + run(run_id=1, sha="aaa", branch="main", created_at="2026-08-21T00:00:00Z"), + ] + ) + self.assertFalse(pair.skip) + assert pair.base is not None + self.assertEqual(pair.base.head_sha, "aaa") + + def test_skips_when_only_one_nightly_exists(self) -> None: + pair = refs.select_base_and_candidate( + [run(run_id=1, sha="aaa")], + ) + self.assertTrue(pair.skip) + self.assertIn("no previous", pair.reason) + + def test_explicit_refs_bypass_github(self) -> None: + pair = refs.override_pair("base-sha", "cand-sha") + self.assertFalse(pair.skip) + assert pair.base is not None and pair.candidate is not None + self.assertEqual(pair.base.head_sha, "base-sha") + self.assertEqual(pair.candidate.head_sha, "cand-sha") + + +if __name__ == "__main__": + unittest.main() From 1409e66837d89ff7f9d7d0604dfdbba3ca9fc7a7 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 26 Aug 2026 14:00:39 +0000 Subject: [PATCH 05/22] refactor(flight): add request builder and defer DoGet execution (#8953) * refactor: add flight request builder Signed-off-by: WenyXu * refactor: use flight request builder in flow Signed-off-by: WenyXu * refactor: defer frontend flight query execution Signed-off-by: WenyXu * fix(grpc): avoid cloning requests during auth Signed-off-by: WenyXu * test(grpc): cover flight request timeout Signed-off-by: WenyXu * refactor(client): share Flight message reader Signed-off-by: WenyXu * feat(flow): add Flight DoGet timeout Signed-off-by: WenyXu * refactor(client): gate Flight DDL helpers for testing Signed-off-by: WenyXu * fix(client): restore Flight stream error semantics Signed-off-by: WenyXu * docs(grpc): document Flight stream input constructors Signed-off-by: WenyXu * fix(flight): preserve deferred stream context Signed-off-by: WenyXu * fix(client): use Flight stream SNAFU context Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/client/src/database.rs | 490 ++++++++++++------ src/client/src/error.rs | 18 + src/client/src/flight.rs | 74 ++- src/client/src/region.rs | 189 +++++-- src/datanode/src/region_server.rs | 9 +- src/flow/src/batching_mode.rs | 3 + src/flow/src/batching_mode/frontend_client.rs | 61 +-- src/servers/src/grpc/flight.rs | 75 +-- src/servers/src/grpc/flight/stream.rs | 150 ++++-- src/servers/src/grpc/greptime_handler.rs | 72 +-- tests-integration/src/grpc/flight.rs | 180 ++++++- tests-integration/tests/http.rs | 1 + 12 files changed, 923 insertions(+), 399 deletions(-) diff --git a/src/client/src/database.rs b/src/client/src/database.rs index 4fff3636b63..7c33794f477 100644 --- a/src/client/src/database.rs +++ b/src/client/src/database.rs @@ -12,20 +12,25 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; use std::pin::Pin; use std::str::FromStr; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, RwLock}; use std::task::{Context, Poll}; +use std::time::Duration; use api::v1::auth_header::AuthScheme; +#[cfg(feature = "testing")] use api::v1::ddl_request::Expr as DdlExpr; use api::v1::greptime_database_client::GreptimeDatabaseClient; use api::v1::greptime_request::Request; use api::v1::query_request::Query; +#[cfg(feature = "testing")] +use api::v1::{AlterTableExpr, CreateTableExpr, DdlRequest}; use api::v1::{ - AlterTableExpr, AuthHeader, Basic, CreateTableExpr, DdlRequest, GreptimeRequest, - InsertRequests, QueryRequest, RequestHeader, RowInsertRequests, + AuthHeader, Basic, GreptimeRequest, InsertRequests, QueryRequest, RequestHeader, + RowInsertRequests, }; use arc_swap::ArcSwapOption; use arrow_flight::{FlightData, Ticket}; @@ -34,7 +39,7 @@ use base64::Engine; use base64::prelude::BASE64_STANDARD; use common_catalog::build_db_string; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; -use common_error::ext::BoxedError; +use common_error::ext::{BoxedError, ErrorExt}; use common_grpc::flight::do_put::DoPutResponse; use common_grpc::flight::{ FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightMessage, SNAPSHOT_SEQS_METADATA_KEY, @@ -49,21 +54,23 @@ use common_telemetry::{error, warn}; use futures::future; use futures_util::{Stream, StreamExt, TryStreamExt}; use prost::Message; -use snafu::ResultExt; +use snafu::{IntoError, ResultExt}; use tonic::metadata::{AsciiMetadataKey, AsciiMetadataValue, MetadataMap, MetadataValue}; use tonic::transport::Channel; use crate::error::{ - ConvertFlightDataSnafu, Error, FlightGetSnafu, IllegalFlightMessagesSnafu, + ConvertFlightDataSnafu, Error, FlightGetSnafu, FlightStreamSnafu, IllegalFlightMessagesSnafu, InvalidTonicMetadataValueSnafu, }; -use crate::flight::decode_flight_data; +use crate::flight::{FlightMessageReader, decode_flight_data}; use crate::{Client, Result, error, from_grpc_response}; type FlightDataStream = Pin + Send>>; type DoPutResponseStream = Pin>>>; +const HINTS_METADATA_KEY: &str = "x-greptime-hints"; + /// Terminal metrics associated with a query output. /// /// For streaming outputs, metrics are only final after the stream is fully @@ -257,19 +264,17 @@ fn attach_terminal_metrics(output: Output, terminal_metrics: &OutputMetrics) -> } async fn output_from_flight_message_stream( - mut flight_message_stream: S, + remote_addr: String, + flight_message_stream: S, ) -> Result where S: Stream> + Send + Unpin + 'static, { - let Some(first_flight_message) = flight_message_stream.next().await else { - return IllegalFlightMessagesSnafu { - reason: "Expect the response not to be empty", - } - .fail(); - }; - - let first_flight_message = first_flight_message?; + let mut reader = FlightMessageReader::new(remote_addr, flight_message_stream); + let first_flight_message = reader + .read_first() + .await + .map_err(|error| flight_stream_error(reader.remote_addr(), error))?; match first_flight_message { FlightMessage::AffectedRows { rows, metrics } => { @@ -277,7 +282,10 @@ where if let Some(metrics) = metrics { terminal_metrics.update(Some(parse_terminal_metrics(&metrics)?)); } - let next_message = flight_message_stream.next().await.transpose()?; + let next_message = reader + .read_next() + .await + .map_err(|error| flight_stream_error(reader.remote_addr(), error))?; match next_message { None => terminal_metrics.mark_ready(), Some(FlightMessage::Metrics(s)) if terminal_metrics.get().is_none() => { @@ -316,15 +324,19 @@ where ); let schema_cloned = schema.clone(); let stream = Box::pin(stream!({ - while let Some(flight_message_item) = flight_message_stream.next().await { - let flight_message = match flight_message_item { - Ok(message) => message, - Err(e) => { - yield Err(BoxedError::new(e)).context(ExternalSnafu); + loop { + let flight_message = match reader.read_next().await { + Ok(Some(message)) => message, + Ok(None) => break, + Err(error) => { + yield Err(BoxedError::new(flight_stream_error( + reader.remote_addr(), + error, + ))) + .context(ExternalSnafu); break; } }; - match flight_message { FlightMessage::RecordBatch(arrow_batch) => { yield Ok(RecordBatch::from_df_record_batch( @@ -371,6 +383,25 @@ where } } +fn flight_stream_error(addr: &str, error: Error) -> Error { + let tonic_code = error.tonic_code().unwrap_or(tonic::Code::Unknown); + let message = error.to_string(); + if error.status_code().should_log_error() { + error!( + error; "Failed to receive Flight data, addr: {}, code: {}", + addr, + tonic_code + ); + } + + FlightStreamSnafu { + addr: addr.to_string(), + tonic_code, + message, + } + .into_error(BoxedError::new(error)) +} + #[derive(Clone, Debug, Default)] pub struct Database { // The "catalog" and "schema" to be used in processing the requests at the server side. @@ -389,6 +420,42 @@ pub struct Database { ctx: FlightContext, } +#[derive(Default)] +struct FlightRequestOptions { + hints: Option, + flow_extensions: Option, + snapshot_seqs: Option, + timeout: Option, +} + +impl FlightRequestOptions { + fn apply_to(self, request: &mut tonic::Request) -> Result<()> { + let metadata = request.metadata_mut(); + if let Some(hints) = self.hints { + Database::put_metadata_value(metadata, HINTS_METADATA_KEY, hints)?; + } + if let Some(flow_extensions) = self.flow_extensions { + Database::put_metadata_value(metadata, FLOW_EXTENSIONS_METADATA_KEY, flow_extensions)?; + } + if let Some(snapshot_seqs) = self.snapshot_seqs { + Database::put_metadata_value(metadata, SNAPSHOT_SEQS_METADATA_KEY, snapshot_seqs)?; + } + if let Some(timeout) = self.timeout { + request.set_timeout(timeout); + } + Ok(()) + } +} + +/// A single Flight DoGet request to a [`Database`]. +/// +/// The builder carries request-scoped metadata and options. It does not modify +/// the underlying [`Database`], so its configuration cannot affect later RPCs. +pub struct DatabaseFlightRequest<'a> { + database: &'a Database, + options: FlightRequestOptions, +} + pub struct DatabaseClient { pub addr: String, pub inner: GreptimeDatabaseClient, @@ -483,6 +550,14 @@ impl Database { }); } + /// Creates a builder for a single Flight DoGet request. + pub fn flight_request(&self) -> DatabaseFlightRequest<'_> { + DatabaseFlightRequest { + database: self, + options: FlightRequestOptions::default(), + } + } + /// Make an InsertRequests request to the database. pub async fn insert(&self, requests: InsertRequests) -> Result { self.handle(Request::Inserts(requests)).await @@ -538,44 +613,31 @@ impl Database { } fn put_hints(metadata: &mut MetadataMap, hints: &[(&str, &str)]) -> Result<()> { - let Some(value) = hints - .iter() - .map(|(k, v)| format!("{}={}", k, v)) - .reduce(|a, b| format!("{},{}", a, b)) - else { + let Some(value) = Self::encode_hints(hints) else { return Ok(()); }; - let key = AsciiMetadataKey::from_static("x-greptime-hints"); - let value = AsciiMetadataValue::from_str(&value).context(InvalidTonicMetadataValueSnafu)?; - metadata.insert(key, value); - Ok(()) + Self::put_metadata_value(metadata, HINTS_METADATA_KEY, value) } - fn put_flow_extensions( - metadata: &mut MetadataMap, - flow_extensions: &[(&str, &str)], - ) -> Result<()> { - if flow_extensions.is_empty() { - return Ok(()); - } - - let value = serde_json::to_string(&flow_extensions.to_vec()) - .expect("flow extension pairs should serialize"); - Self::put_metadata_value(metadata, FLOW_EXTENSIONS_METADATA_KEY, value) + fn encode_hints(hints: &[(&str, &str)]) -> Option { + hints + .iter() + .map(|(k, v)| format!("{}={}", k, v)) + .reduce(|a, b| format!("{},{}", a, b)) } - fn put_snapshot_seqs( - metadata: &mut MetadataMap, - snapshot_seqs: &std::collections::HashMap, - ) -> Result<()> { - if snapshot_seqs.is_empty() { - return Ok(()); - } + fn encode_flow_extensions(flow_extensions: &[(&str, &str)]) -> Option { + (!flow_extensions.is_empty()).then(|| { + serde_json::to_string(&flow_extensions.to_vec()) + .expect("flow extension pairs should serialize") + }) + } - let value = - serde_json::to_string(snapshot_seqs).expect("snapshot sequence map should serialize"); - Self::put_metadata_value(metadata, SNAPSHOT_SEQS_METADATA_KEY, value) + fn encode_snapshot_seqs(snapshot_seqs: &HashMap) -> Option { + (!snapshot_seqs.is_empty()).then(|| { + serde_json::to_string(snapshot_seqs).expect("snapshot sequence map should serialize") + }) } fn put_metadata_value( @@ -673,7 +735,7 @@ impl Database { where S: AsRef, { - self.sql_with_hint(sql, &[]).await + self.flight_request().sql(sql).await } /// Executes a SQL query with optional hints for query optimization. @@ -681,12 +743,7 @@ impl Database { where S: AsRef, { - let request = Request::Query(QueryRequest { - query: Some(Query::Sql(sql.as_ref().to_string())), - }); - self.do_get(request, hints, &[], &Default::default()) - .await - .map(OutputWithMetrics::into_output) + self.flight_request().with_hints(hints).sql(sql).await } /// Executes a SQL query and returns the output with terminal metrics. @@ -701,77 +758,33 @@ impl Database { where S: AsRef, { - self.query_with_terminal_metrics_and_flow_extensions( - QueryRequest { - query: Some(Query::Sql(sql.as_ref().to_string())), - }, - hints, - &[], - &Default::default(), - ) - .await + self.flight_request() + .with_hints(hints) + .sql_with_terminal_metrics(sql) + .await } /// Executes a logical plan directly without SQL parsing. pub async fn logical_plan(&self, logical_plan: Vec) -> Result { - self.query_with_terminal_metrics_and_flow_extensions( - QueryRequest { - query: Some(Query::LogicalPlan(logical_plan)), - }, - &[], - &[], - &Default::default(), - ) - .await - .map(OutputWithMetrics::into_output) - } - - /// Executes a query and carries flow extensions through Flight metadata. - /// - /// This is the lower-level terminal-metrics API for Flow callers that need - /// to pass JSON-bearing flow extensions without going through hint metadata. - pub async fn query_with_terminal_metrics_and_flow_extensions( - &self, - request: QueryRequest, - hints: &[(&str, &str)], - flow_extensions: &[(&str, &str)], - snapshot_seqs: &std::collections::HashMap, - ) -> Result { - self.do_get( - Request::Query(request), - hints, - flow_extensions, - snapshot_seqs, - ) - .await + self.flight_request().logical_plan(logical_plan).await } /// Creates a new table using the provided table expression. + #[cfg(feature = "testing")] pub async fn create(&self, expr: CreateTableExpr) -> Result { - let request = Request::Ddl(DdlRequest { - expr: Some(DdlExpr::CreateTable(expr)), - }); - self.do_get(request, &[], &[], &Default::default()) - .await - .map(OutputWithMetrics::into_output) + self.flight_request().create(expr).await } /// Alters an existing table using the provided alter expression. + #[cfg(feature = "testing")] pub async fn alter(&self, expr: AlterTableExpr) -> Result { - let request = Request::Ddl(DdlRequest { - expr: Some(DdlExpr::AlterTable(expr)), - }); - self.do_get(request, &[], &[], &Default::default()) - .await - .map(OutputWithMetrics::into_output) + self.flight_request().alter(expr).await } async fn do_get( &self, request: Request, - hints: &[(&str, &str)], - flow_extensions: &[(&str, &str)], - snapshot_seqs: &std::collections::HashMap, + options: FlightRequestOptions, ) -> Result { let request = self.to_rpc_request(request); let request = Ticket { @@ -779,12 +792,10 @@ impl Database { }; let mut request = tonic::Request::new(request); - let metadata = request.metadata_mut(); - Self::put_hints(metadata, hints)?; - Self::put_flow_extensions(metadata, flow_extensions)?; - Self::put_snapshot_seqs(metadata, snapshot_seqs)?; + options.apply_to(&mut request)?; let mut client = self.client.make_flight_client(false, false)?; + let remote_addr = client.addr().to_string(); let response = client.mut_inner().do_get(request).await.or_else(|e| { let tonic_code = e.code(); @@ -796,7 +807,7 @@ impl Database { e ); Err(BoxedError::new(e)).with_context(|_| FlightGetSnafu { - addr: client.addr().to_string(), + addr: remote_addr.clone(), tonic_code, }) })?; @@ -808,7 +819,7 @@ impl Database { future::ready(decode_flight_data(&mut decoder, flight_data)) }); - output_from_flight_message_stream(flight_message_stream).await + output_from_flight_message_stream(remote_addr, flight_message_stream).await } /// Ingest a stream of [RecordBatch]es that belong to a table, using Arrow Flight's "`DoPut`" @@ -847,6 +858,98 @@ impl Database { } } +impl<'a> DatabaseFlightRequest<'a> { + /// Adds query optimization hints to this Flight request. + pub fn with_hints(mut self, hints: &[(&str, &str)]) -> Self { + self.options.hints = Database::encode_hints(hints); + self + } + + /// Adds Flow extensions to this Flight request. + pub fn with_flow_extensions(mut self, flow_extensions: &[(&str, &str)]) -> Self { + self.options.flow_extensions = Database::encode_flow_extensions(flow_extensions); + self + } + + /// Adds snapshot sequence fences to this Flight request. + pub fn with_snapshot_seqs(mut self, snapshot_seqs: &HashMap) -> Self { + self.options.snapshot_seqs = Database::encode_snapshot_seqs(snapshot_seqs); + self + } + + /// Sets a timeout for this Flight request only. + pub fn with_timeout(mut self, timeout: Duration) -> Self { + self.options.timeout = Some(timeout); + self + } + + /// Executes a SQL query. + pub async fn sql(self, sql: S) -> Result + where + S: AsRef, + { + let request = Request::Query(QueryRequest { + query: Some(Query::Sql(sql.as_ref().to_string())), + }); + self.do_get(request) + .await + .map(OutputWithMetrics::into_output) + } + + /// Executes a SQL query and returns terminal metrics. + pub async fn sql_with_terminal_metrics(self, sql: S) -> Result + where + S: AsRef, + { + self.query_with_terminal_metrics(QueryRequest { + query: Some(Query::Sql(sql.as_ref().to_string())), + }) + .await + } + + /// Executes a logical plan directly without SQL parsing. + pub async fn logical_plan(self, logical_plan: Vec) -> Result { + self.query_with_terminal_metrics(QueryRequest { + query: Some(Query::LogicalPlan(logical_plan)), + }) + .await + .map(OutputWithMetrics::into_output) + } + + /// Executes a query and returns terminal metrics. + pub async fn query_with_terminal_metrics( + self, + request: QueryRequest, + ) -> Result { + self.do_get(Request::Query(request)).await + } + + /// Creates a new table using the provided table expression. + #[cfg(feature = "testing")] + pub async fn create(self, expr: CreateTableExpr) -> Result { + self.do_get(Request::Ddl(DdlRequest { + expr: Some(DdlExpr::CreateTable(expr)), + })) + .await + .map(OutputWithMetrics::into_output) + } + + /// Alters an existing table using the provided alter expression. + #[cfg(feature = "testing")] + pub async fn alter(self, expr: AlterTableExpr) -> Result { + self.do_get(Request::Ddl(DdlRequest { + expr: Some(DdlExpr::AlterTable(expr)), + })) + .await + .map(OutputWithMetrics::into_output) + } + + async fn do_get(self, request: Request) -> Result { + let Self { database, options } = self; + database.do_get(request, options).await + } +} + /// by grpc standard, only `Unavailable` is retryable, see: https://github.com/grpc/grpc/blob/master/doc/statuscodes.md#status-codes-and-their-use-in-grpc pub fn is_grpc_retryable(err: &tonic::Status) -> bool { matches!(err.code(), tonic::Code::Unavailable) @@ -934,12 +1037,14 @@ mod tests { #[test] fn test_put_flow_extensions_preserves_comma_bearing_values() { let mut metadata = MetadataMap::new(); - Database::put_flow_extensions( + Database::put_metadata_value( &mut metadata, - &[ + FLOW_EXTENSIONS_METADATA_KEY, + Database::encode_flow_extensions(&[ ("flow.return_region_seq", "true"), ("flow.incremental_after_seqs", r#"{"1":10,"2":20}"#), - ], + ]) + .unwrap(), ) .unwrap(); @@ -969,7 +1074,12 @@ mod tests { (9_007_199_254_740_993_u64, 9_007_199_254_740_995_u64), ]); - Database::put_snapshot_seqs(&mut metadata, &snapshot_seqs).unwrap(); + Database::put_metadata_value( + &mut metadata, + SNAPSHOT_SEQS_METADATA_KEY, + Database::encode_snapshot_seqs(&snapshot_seqs).unwrap(), + ) + .unwrap(); let value = metadata .get(SNAPSHOT_SEQS_METADATA_KEY) @@ -980,6 +1090,50 @@ mod tests { assert_eq!(decoded, snapshot_seqs); } + #[test] + fn test_flight_request_builder_applies_request_options() { + let database = Database::new("greptime", "public", Client::default()); + let snapshot_seqs = HashMap::from([(42, 99)]); + let request = database + .flight_request() + .with_hints(&[("query_parallelism", "1")]) + .with_flow_extensions(&[("flow.return_region_seq", "true")]) + .with_snapshot_seqs(&snapshot_seqs) + .with_timeout(Duration::from_millis(50)); + let mut tonic_request = tonic::Request::new(()); + + request.options.apply_to(&mut tonic_request).unwrap(); + + let metadata = tonic_request.metadata(); + assert_eq!( + metadata.get(HINTS_METADATA_KEY).unwrap(), + "query_parallelism=1" + ); + assert_eq!( + serde_json::from_str::>( + metadata + .get(FLOW_EXTENSIONS_METADATA_KEY) + .unwrap() + .to_str() + .unwrap(), + ) + .unwrap(), + vec![("flow.return_region_seq".to_string(), "true".to_string())] + ); + assert_eq!( + serde_json::from_str::>( + metadata + .get(SNAPSHOT_SEQS_METADATA_KEY) + .unwrap() + .to_str() + .unwrap(), + ) + .unwrap(), + snapshot_seqs + ); + assert!(metadata.get("grpc-timeout").is_some()); + } + #[test] fn test_flight_ctx() { let mut ctx = FlightContext::default(); @@ -1020,6 +1174,28 @@ mod tests { assert_eq!(expected.should_retry(), actual.should_retry()); } + #[test] + fn test_flight_stream_error_preserves_addr_and_message() { + let error = flight_stream_error( + "127.0.0.1:4001", + Status::out_of_range("message length too large").into(), + ); + + assert!(matches!( + &error, + Error::FlightStream { + addr, + tonic_code: Code::OutOfRange, + message, + .. + } if addr == "127.0.0.1:4001" && message == "message length too large" + )); + assert_eq!( + "Failed to receive Flight data from 127.0.0.1:4001, code: Operation was attempted past the valid range: message length too large", + error.to_string(), + ); + } + #[test] fn test_from_tonic_status_with_retry_hint() { let mut headers = HeaderMap::new(); @@ -1117,13 +1293,13 @@ mod tests { #[tokio::test] async fn test_affected_rows_inline_metrics_are_parsed() { - let output = output_from_flight_message_stream(futures_util::stream::iter(vec![Ok( - FlightMessage::AffectedRows { + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![Ok(FlightMessage::AffectedRows { rows: 3, metrics: Some(terminal_metrics_json()), - }, - )] - as Vec>)) + })] as Vec>), + ) .await .unwrap(); @@ -1138,14 +1314,16 @@ mod tests { #[tokio::test] async fn test_affected_rows_inline_metrics_rejects_trailing_metrics() { let metrics_json = terminal_metrics_json(); - let err = output_from_flight_message_stream(futures_util::stream::iter(vec![ - Ok(FlightMessage::AffectedRows { - rows: 3, - metrics: Some(metrics_json.clone()), - }), - Ok(FlightMessage::Metrics(metrics_json)), - ] - as Vec>)) + let err = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: Some(metrics_json.clone()), + }), + Ok(FlightMessage::Metrics(metrics_json)), + ] as Vec>), + ) .await .unwrap_err(); @@ -1167,12 +1345,14 @@ mod tests { vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], ) .unwrap(); - let output = output_from_flight_message_stream(futures_util::stream::iter(vec![ - Ok(FlightMessage::Schema(schema.arrow_schema().clone())), - Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), - Ok(FlightMessage::Metrics("{not-json}".to_string())), - ] - as Vec>)) + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + Ok(FlightMessage::Metrics("{not-json}".to_string())), + ] as Vec>), + ) .await .unwrap(); let terminal_metrics = output.metrics.clone(); @@ -1211,18 +1391,20 @@ mod tests { vec![Arc::new(Int32Vector::from_slice([2])) as VectorRef], ) .unwrap(); - let output = output_from_flight_message_stream(futures_util::stream::iter(vec![ - Ok(FlightMessage::Schema(schema.arrow_schema().clone())), - Ok(FlightMessage::RecordBatch( - first_batch.into_df_record_batch(), - )), - Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(1))), - Ok(FlightMessage::RecordBatch( - second_batch.into_df_record_batch(), - )), - Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(2))), - ] - as Vec>)) + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch( + first_batch.into_df_record_batch(), + )), + Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(1))), + Ok(FlightMessage::RecordBatch( + second_batch.into_df_record_batch(), + )), + Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(2))), + ] as Vec>), + ) .await .unwrap(); let terminal_metrics = output.metrics.clone(); diff --git a/src/client/src/error.rs b/src/client/src/error.rs index 1db0858f087..168ba422715 100644 --- a/src/client/src/error.rs +++ b/src/client/src/error.rs @@ -40,6 +40,21 @@ pub enum Error { source: BoxedError, }, + #[snafu(display( + "Failed to receive Flight data from {}, code: {}: {}", + addr, + tonic_code, + message + ))] + FlightStream { + addr: String, + tonic_code: Code, + message: String, + source: BoxedError, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to convert FlightData"))] ConvertFlightData { #[snafu(implicit)] @@ -154,6 +169,7 @@ impl ErrorExt for Error { Error::Server { code, .. } | Error::Tonic { code, .. } => *code, Error::FlightGet { source, .. } + | Error::FlightStream { source, .. } | Error::RegionServer { source, .. } | Error::FlowServer { source, .. } => source.status_code(), Error::CreateChannel { source, .. } @@ -174,6 +190,7 @@ impl ErrorExt for Error { match self { Error::Tonic { retry_hint, .. } => *retry_hint, Error::FlightGet { source, .. } + | Error::FlightStream { source, .. } | Error::RegionServer { source, .. } | Error::FlowServer { source, .. } | Error::External { source, .. } => source.retry_hint(), @@ -193,6 +210,7 @@ impl Error { pub fn tonic_code(&self) -> Option { match self { Self::FlightGet { tonic_code, .. } + | Self::FlightStream { tonic_code, .. } | Self::RegionServer { code: tonic_code, .. } diff --git a/src/client/src/flight.rs b/src/client/src/flight.rs index 4d00cba962e..1753fdd5246 100644 --- a/src/client/src/flight.rs +++ b/src/client/src/flight.rs @@ -12,12 +12,82 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::pin::Pin; + use arrow_flight::FlightData; use common_grpc::flight::{FlightDecoder, FlightMessage}; -use snafu::ResultExt; +use futures_util::stream::Peekable; +use futures_util::{Stream, StreamExt}; +use snafu::{OptionExt, ResultExt}; use crate::Result; -use crate::error::{ConvertFlightDataSnafu, Error}; +use crate::error::{ConvertFlightDataSnafu, Error, IllegalFlightMessagesSnafu}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum FlightMessageKind { + Schema, + RecordBatch, + AffectedRows, + Metrics, +} + +impl From<&FlightMessage> for FlightMessageKind { + fn from(message: &FlightMessage) -> Self { + match message { + FlightMessage::Schema(_) => Self::Schema, + FlightMessage::RecordBatch(_) => Self::RecordBatch, + FlightMessage::AffectedRows { .. } => Self::AffectedRows, + FlightMessage::Metrics(_) => Self::Metrics, + } + } +} + +pub(crate) struct FlightMessageReader { + /// Remote Flight peer associated with this response stream. + remote_addr: String, + messages: Peekable, +} + +impl FlightMessageReader +where + S: Stream> + Unpin, +{ + pub(crate) fn new(remote_addr: impl Into, messages: S) -> Self { + Self { + remote_addr: remote_addr.into(), + messages: messages.peekable(), + } + } + + pub(crate) fn remote_addr(&self) -> &str { + &self.remote_addr + } + + pub(crate) async fn read_first(&mut self) -> Result { + self.read_next().await?.context(IllegalFlightMessagesSnafu { + reason: "Expect the response not to be empty", + }) + } + + pub(crate) async fn read_next(&mut self) -> Result> { + self.messages.next().await.transpose() + } + + pub(crate) async fn peek_next_message_kind(&mut self) -> Result> { + match Pin::new(&mut self.messages).peek().await { + Some(Ok(message)) => Ok(Some(message.into())), + None => Ok(None), + Some(Err(_)) => match self.read_next().await { + // `peek` only borrows the error; consume it to preserve the source error. + Err(error) => Err(error), + Ok(_) => IllegalFlightMessagesSnafu { + reason: "Flight stream changed after peek".to_string(), + } + .fail(), + }, + } + } +} pub(crate) fn decode_flight_data( decoder: &mut FlightDecoder, diff --git a/src/client/src/region.rs b/src/client/src/region.rs index 2bb5f7dd0a6..cd155810ea9 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -47,7 +47,7 @@ use crate::error::{ self, FlightGetSnafu, IllegalDatabaseResponseSnafu, IllegalFlightMessagesSnafu, MissingFieldSnafu, Result, ServerSnafu, }; -use crate::flight::decode_flight_data; +use crate::flight::{FlightMessageKind, FlightMessageReader, decode_flight_data}; use crate::{Client, metrics}; const FLIGHT_DO_GET_TIMEOUT: Duration = Duration::from_secs(10); @@ -203,26 +203,22 @@ impl RegionRequester { async fn recordbatches_from_flight_message_stream( addr: String, - mut flight_message_stream: S, + flight_message_stream: S, ) -> Result where S: Stream> + Send + Unpin + 'static, { - let Some(first_flight_message) = flight_message_stream.next().await else { - return IllegalFlightMessagesSnafu { - reason: "Expect the response not to be empty", - } - .fail() - .map_err(|error| flight_stream_error(&addr, error)); - }; - let FlightMessage::Schema(schema) = - first_flight_message.map_err(|e| flight_stream_error(&addr, e))? + let mut reader = FlightMessageReader::new(addr.clone(), flight_message_stream); + let FlightMessage::Schema(schema) = reader + .read_first() + .await + .map_err(|error| flight_stream_error(reader.remote_addr(), error))? else { return IllegalFlightMessagesSnafu { reason: "Expect schema to be the first flight message", } .fail() - .map_err(|error| flight_stream_error(&addr, error)); + .map_err(|error| flight_stream_error(reader.remote_addr(), error)); }; let metrics = Arc::new(ArcSwapOption::from(None)); @@ -233,31 +229,23 @@ where let schema = Arc::new(datatypes::schema::Schema::try_from(schema).context(error::ConvertSchemaSnafu)?); let schema_cloned = schema.clone(); - let stream_addr = addr.clone(); + let stream_addr = addr; let stream = Box::pin(stream!({ let _span = tracing_context.attach(common_telemetry::tracing::info_span!( "poll_flight_data_stream" )); - let mut buffered_message: Option = None; let mut stream_ended = false; while !stream_ended { - // get the next message from the buffered message or read from the flight message stream - let flight_message_item = if let Some(msg) = buffered_message.take() { - Some(Ok(msg)) - } else { - flight_message_stream.next().await - }; - - let flight_message = match flight_message_item { - Some(Ok(message)) => message, - Some(Err(e)) => { - yield Err(BoxedError::new(flight_stream_error(&stream_addr, e))) + let flight_message = match reader.read_next().await { + Ok(Some(message)) => message, + Ok(None) => break, + Err(error) => { + yield Err(BoxedError::new(flight_stream_error(&stream_addr, error))) .context(ExternalSnafu); break; } - None => break, }; match flight_message { @@ -265,42 +253,54 @@ where let result_to_yield = RecordBatch::from_df_record_batch(schema_cloned.clone(), record_batch); - // get the next message from the stream. normally it should be a metrics message. - if let Some(next_flight_message_result) = flight_message_stream.next().await { - match next_flight_message_result { - Ok(FlightMessage::Metrics(s)) => { - let m = serde_json::from_str(&s).ok().map(Arc::new); - metrics_ref.swap(m); - } - Ok(FlightMessage::RecordBatch(rb)) => { - // for some reason it's not a metrics message, so we need to buffer this record batch - // and yield it in the next iteration. - buffered_message = Some(FlightMessage::RecordBatch(rb)); - } - Ok(_) => { - yield IllegalFlightMessagesSnafu { - reason: "A RecordBatch message can only be succeeded by a Metrics message or another RecordBatch message" - } - .fail() - .map_err(BoxedError::new) - .context(ExternalSnafu); - break; - } - Err(e) => { - yield Err(BoxedError::new(flight_stream_error(&stream_addr, e))) + // Metrics follow a batch so MergeScan can observe them before yielding it. + match reader.peek_next_message_kind().await { + Ok(Some(FlightMessageKind::Metrics)) => { + let metrics_message = match reader.read_next().await { + Ok(Some(FlightMessage::Metrics(metrics))) => metrics, + Ok(Some(_) | None) => { + yield IllegalFlightMessagesSnafu { + reason: "Flight stream changed after peek", + } + .fail() + .map_err(BoxedError::new) .context(ExternalSnafu); - break; + break; + } + Err(error) => { + yield Err(BoxedError::new(flight_stream_error( + &stream_addr, + error, + ))) + .context(ExternalSnafu); + break; + } + }; + let metrics = serde_json::from_str(&metrics_message).ok().map(Arc::new); + metrics_ref.swap(metrics); + } + Ok(Some(FlightMessageKind::RecordBatch)) => {} + Ok(Some(FlightMessageKind::Schema | FlightMessageKind::AffectedRows)) => { + yield IllegalFlightMessagesSnafu { + reason: "A RecordBatch message can only be succeeded by a Metrics message or another RecordBatch message" } + .fail() + .map_err(BoxedError::new) + .context(ExternalSnafu); + break; + } + Ok(None) => stream_ended = true, + Err(error) => { + yield Err(BoxedError::new(flight_stream_error(&stream_addr, error))) + .context(ExternalSnafu); + break; } - } else { - // the stream has ended - stream_ended = true; } yield Ok(result_to_yield); } FlightMessage::Metrics(s) => { - // just a branch in case of some metrics message comes after other things. + // Metrics may arrive before the next RecordBatch. let m = serde_json::from_str(&s).ok().map(Arc::new); metrics_ref.swap(m); continue; @@ -412,6 +412,7 @@ mod test { RemoteDynFilterUnregister, RemoteDynFilterUpdate, region_request, remote_dyn_filter_request, }; use common_recordbatch::adapter::RecordBatchMetrics; + use datatypes::arrow::array::Int32Array; use datatypes::prelude::{ConcreteDataType, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::Int32Vector; @@ -598,6 +599,86 @@ mod test { assert_eq!(metrics.elapsed_compute, 7); } + #[tokio::test] + async fn test_record_batch_stream_updates_following_metrics_before_yielding_batch() { + let schema = test_schema(); + let batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], + ) + .unwrap(); + + let mut recordbatches = recordbatches_from_flight_message_stream( + "test-peer".to_string(), + stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + Ok(FlightMessage::Metrics(test_metrics_json())), + ]), + ) + .await + .unwrap(); + + let batch = recordbatches.next().await.unwrap().unwrap(); + assert_eq!(batch.num_rows(), 1); + + let metrics = recordbatches.metrics().unwrap(); + assert_eq!(metrics.elapsed_compute, 7); + assert!(recordbatches.next().await.is_none()); + } + + #[tokio::test] + async fn test_record_batch_stream_preserves_peeked_record_batch() { + let schema = test_schema(); + let first_batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], + ) + .unwrap(); + let second_batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([2])) as VectorRef], + ) + .unwrap(); + + let mut recordbatches = recordbatches_from_flight_message_stream( + "test-peer".to_string(), + stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch( + first_batch.into_df_record_batch(), + )), + Ok(FlightMessage::RecordBatch( + second_batch.into_df_record_batch(), + )), + ]), + ) + .await + .unwrap(); + + let first_batch = recordbatches.next().await.unwrap().unwrap(); + let second_batch = recordbatches.next().await.unwrap().unwrap(); + assert_eq!( + first_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 1 + ); + assert_eq!( + second_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 2 + ); + assert!(recordbatches.next().await.is_none()); + } + #[tokio::test] async fn test_record_batch_stream_exposes_error_after_pre_batch_metrics() { let schema = test_schema(); diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 0d9398b0889..79d0dd19152 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -67,7 +67,8 @@ use servers::error::{ }; use servers::grpc::FlightCompression; use servers::grpc::flight::{ - FlightCraft, FlightRecordBatchSource, FlightRecordBatchStream, TonicStream, + FlightCraft, FlightRecordBatchSource, FlightRecordBatchStream, FlightRecordBatchStreamInput, + TonicStream, }; use servers::grpc::region_server::RegionServerHandler; use session::context::{ @@ -990,7 +991,11 @@ impl FlightCraft for RegionServer { }; let stream = Box::pin(FlightRecordBatchStream::new( - FlightRecordBatchSource::initializer(initializer), + FlightRecordBatchStreamInput::initializer(async move { + initializer + .await + .map(FlightRecordBatchSource::RecordBatches) + }), tracing_context, self.flight_compression, query_ctx, diff --git a/src/flow/src/batching_mode.rs b/src/flow/src/batching_mode.rs index bb946e75740..3981446721b 100644 --- a/src/flow/src/batching_mode.rs +++ b/src/flow/src/batching_mode.rs @@ -44,6 +44,8 @@ pub struct BatchingModeOptions { /// The gRPC connection timeout #[serde(with = "humantime_serde")] pub grpc_conn_timeout: Duration, + #[serde(with = "humantime_serde")] + pub experimental_flight_do_get_timeout: Duration, /// The gRPC max retry number pub experimental_grpc_max_retries: u32, /// Flow wait for available frontend timeout, @@ -72,6 +74,7 @@ impl Default for BatchingModeOptions { slow_query_threshold: Duration::from_secs(60), experimental_min_refresh_duration: Duration::new(5, 0), grpc_conn_timeout: Duration::from_secs(5), + experimental_flight_do_get_timeout: Duration::from_secs(10), experimental_grpc_max_retries: 3, experimental_frontend_scan_timeout: Duration::from_secs(30), experimental_max_filter_num_per_query: 20, diff --git a/src/flow/src/batching_mode/frontend_client.rs b/src/flow/src/batching_mode/frontend_client.rs index aa94b91b2ae..27db96748a6 100644 --- a/src/flow/src/batching_mode/frontend_client.rs +++ b/src/flow/src/batching_mode/frontend_client.rs @@ -18,7 +18,6 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex, RwLock, Weak}; use api::v1::greptime_request::Request; -use api::v1::query_request::Query; use api::v1::{CreateTableExpr, QueryRequest}; use client::{Client, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, Database, OutputWithMetrics}; use common_error::ext::BoxedError; @@ -340,53 +339,6 @@ impl FrontendClient { }) } - /// Execute a SQL statement on the frontend. - pub async fn sql(&self, catalog: &str, schema: &str, sql: &str) -> Result { - match self { - FrontendClient::Distributed { .. } => { - let db = self.get_random_active_frontend(catalog, schema).await?; - db.database - .sql(sql) - .await - .map_err(BoxedError::new) - .context(ExternalSnafu) - } - FrontendClient::Standalone { - database_client, .. - } => { - let ctx = QueryContextBuilder::default() - .current_catalog(catalog.to_string()) - .current_schema(schema.to_string()) - .build(); - let ctx = Arc::new(ctx); - { - let database_client = { - database_client - .handler - .lock() - .unwrap() - .as_ref() - .context(UnexpectedSnafu { - reason: "Standalone's frontend instance is not set", - })? - .upgrade() - .context(UnexpectedSnafu { - reason: "Failed to upgrade database client", - })? - }; - let req = Request::Query(QueryRequest { - query: Some(Query::Sql(sql.to_string())), - }); - database_client - .do_query(req, ctx) - .await - .map_err(BoxedError::new) - .context(ExternalSnafu) - } - } - } - } - /// Execute a flow query and return terminal metrics. `snapshot_seqs` are /// optional read upper bounds used only by snapshot-fenced repair chunks. pub(crate) async fn query_with_terminal_metrics( @@ -413,12 +365,12 @@ impl FrontendClient { peer: db.peer.clone(), }); db.database - .query_with_terminal_metrics_and_flow_extensions( - request, - &hints, - extensions, - snapshot_seqs, - ) + .flight_request() + .with_hints(&hints) + .with_flow_extensions(extensions) + .with_snapshot_seqs(snapshot_seqs) + .with_timeout(batch_opts.experimental_flight_do_get_timeout) + .query_with_terminal_metrics(request) .await .map_err(BoxedError::new) .context(ExternalSnafu) @@ -623,6 +575,7 @@ mod tests { use std::task::{Context, Poll}; use std::time::Duration; + use api::v1::query_request::Query; use arrow_flight::flight_service_server::FlightServiceServer; use arrow_flight::{FlightData, Ticket}; use common_query::{Output, OutputData}; diff --git a/src/servers/src/grpc/flight.rs b/src/servers/src/grpc/flight.rs index e6f35e17b29..0b7a9dd5704 100644 --- a/src/servers/src/grpc/flight.rs +++ b/src/servers/src/grpc/flight.rs @@ -30,8 +30,7 @@ use bytes::{self, Bytes}; use common_error::ext::ErrorExt; use common_grpc::flight::do_put::{DoPutMetadata, DoPutResponse}; use common_grpc::flight::{ - FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightEncoder, FlightMessage, - SNAPSHOT_SEQS_METADATA_KEY, + FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightMessage, SNAPSHOT_SEQS_METADATA_KEY, }; use common_memory_manager::MemoryGuard; use common_query::{Output, OutputData}; @@ -46,14 +45,16 @@ use prost::Message; use query::metrics::terminal_recordbatch_metrics_from_plan_if_requested; use query::options::FlowQueryExtensions; use session::context::{Channel, QueryContextRef}; -use snafu::{IntoError, ResultExt, ensure}; +use snafu::{IntoError, OptionExt, ResultExt, ensure}; use table::table_name::TableName; use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; use tonic::{Request, Response, Status, Streaming}; -use crate::error::{InvalidParameterSnafu, Result, ToJsonSnafu}; -pub use crate::grpc::flight::stream::{FlightRecordBatchSource, FlightRecordBatchStream}; +use crate::error::{InvalidParameterSnafu, InvalidQuerySnafu, Result, ToJsonSnafu}; +pub use crate::grpc::flight::stream::{ + FlightRecordBatchSource, FlightRecordBatchStream, FlightRecordBatchStreamInput, +}; use crate::grpc::greptime_handler::{ GreptimeRequestHandler, create_query_context, get_request_type, }; @@ -221,9 +222,12 @@ impl FlightCraft for GreptimeRequestHandler { ); let flight_compression = self.flight_compression; async { - let output = self - .handle_request_with_query_ctx(request, query_ctx.clone()) + let query = request.request.context(InvalidQuerySnafu { + reason: "Expecting non-empty GreptimeRequest.", + })?; + self.authenticate_request_with_query_ctx(request.header.as_ref(), &query_ctx) .await?; + let output = self.handle_request_with_query_ctx(query, query_ctx.clone()); let stream = to_flight_data_stream( output, TracingContext::from_current_span(), @@ -573,32 +577,36 @@ fn extract_json_metadata( Ok(Some(parsed)) } -fn to_flight_data_stream( - output: Output, +fn to_flight_data_stream( + output: F, tracing_context: TracingContext, flight_compression: FlightCompression, query_ctx: QueryContextRef, should_emit_terminal_metrics: bool, -) -> TonicStream { +) -> TonicStream +where + F: std::future::Future> + Send + 'static, +{ + let initializer = async move { + let output = output.await.map_err(Status::from)?; + output_to_flight_record_batch_source(output, should_emit_terminal_metrics) + }; + let stream = FlightRecordBatchStream::new( + FlightRecordBatchStreamInput::initializer(initializer), + tracing_context, + flight_compression, + query_ctx, + ); + Box::pin(stream) as _ +} + +fn output_to_flight_record_batch_source( + output: Output, + should_emit_terminal_metrics: bool, +) -> TonicResult { match output.data { - OutputData::Stream(stream) => { - let stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(stream), - tracing_context, - flight_compression, - query_ctx, - ); - Box::pin(stream) as _ - } - OutputData::RecordBatches(x) => { - let stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(x.as_stream()), - tracing_context, - flight_compression, - query_ctx, - ); - Box::pin(stream) as _ - } + OutputData::Stream(stream) => Ok(FlightRecordBatchSource::RecordBatches(stream)), + OutputData::RecordBatches(x) => Ok(FlightRecordBatchSource::RecordBatches(x.as_stream())), OutputData::AffectedRows(rows) => { let terminal_metrics = match terminal_recordbatch_metrics_from_plan_if_requested( output.meta.plan, @@ -607,20 +615,17 @@ fn to_flight_data_stream( Some(metrics) => match serde_json::to_string(&metrics) { Ok(metrics) => Some(metrics), Err(e) => { - let stream = tokio_stream::once(Err(Status::internal(format!( + return Err(Status::internal(format!( "Failed to serialize terminal metrics: {e}" - )))); - return Box::pin(stream) as _; + ))); } }, None => None, }; - let affected_rows = FlightEncoder::default().encode(FlightMessage::AffectedRows { + Ok(FlightRecordBatchSource::AffectedRows { rows, metrics: terminal_metrics, - }); - let stream = tokio_stream::iter(affected_rows.into_iter().map(Ok)); - Box::pin(stream) as _ + }) } } } diff --git a/src/servers/src/grpc/flight/stream.rs b/src/servers/src/grpc/flight/stream.rs index 319a797d9ca..b740f12d02e 100644 --- a/src/servers/src/grpc/flight/stream.rs +++ b/src/servers/src/grpc/flight/stream.rs @@ -41,27 +41,34 @@ use crate::error; use crate::grpc::FlightCompression; use crate::grpc::flight::TonicResult; -pub type FlightRecordBatchStreamInitializer = Pin< - Box< - dyn Future> - + Send - + 'static, - >, ->; - pub enum FlightRecordBatchSource { RecordBatches(SendableRecordBatchStream), - Initializer(FlightRecordBatchStreamInitializer), + AffectedRows { + rows: usize, + metrics: Option, + }, } -impl FlightRecordBatchSource { - pub fn initializer(initializer: F) -> Self - where - F: Future> - + Send - + 'static, - { - Self::Initializer(Box::pin(initializer)) +/// Determines whether a Flight result is ready now or initialized asynchronously. +pub enum FlightRecordBatchStreamInput>> +{ + Ready(FlightRecordBatchSource), + Initializer(F), +} + +impl FlightRecordBatchStreamInput { + /// Creates an input from a source that is already available. + pub fn ready(source: FlightRecordBatchSource) -> Self { + Self::Ready(source) + } +} + +impl FlightRecordBatchStreamInput { + /// Creates an input that obtains its source asynchronously. + /// + /// Errors from the initializer are returned through the Flight response stream. + pub fn initializer(initializer: F) -> Self { + Self::Initializer(initializer) } } @@ -160,31 +167,32 @@ impl FlightRecordBatchStream { Self::send_metrics(tx, metrics, metrics_str).await } - pub fn new( - source: FlightRecordBatchSource, + pub fn new( + input: FlightRecordBatchStreamInput, tracing_context: TracingContext, compression: FlightCompression, query_ctx: QueryContextRef, - ) -> Self { - let should_send_partial_metrics = query_ctx.explain_verbose(); - let can_send_metrics_before_batch = query_ctx.explain_verbose() - && query_ctx.live_analyze_metrics_enabled() - && query_ctx - .remote_query_id() - .zip(query_ctx.extension(SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY)) - .is_some_and(|(remote_query_id, capability)| capability == remote_query_id); + ) -> Self + where + F: Future> + Send + 'static, + { let (mut tx, rx) = mpsc::channel::>(1); - let source_type = match &source { - FlightRecordBatchSource::RecordBatches(_) => "record_batches", - FlightRecordBatchSource::Initializer(_) => "initializer", + let source_type = match &input { + FlightRecordBatchStreamInput::Ready(FlightRecordBatchSource::RecordBatches(_)) => { + "record_batches" + } + FlightRecordBatchStreamInput::Ready(FlightRecordBatchSource::AffectedRows { + .. + }) => "affected_rows", + FlightRecordBatchStreamInput::Initializer(_) => "initializer", }; let initializer_tracing_context = tracing_context.clone(); let join_handle = common_runtime::spawn_global( async move { - let recordbatches = async move { - match source { - FlightRecordBatchSource::RecordBatches(recordbatches) => Ok(recordbatches), - FlightRecordBatchSource::Initializer(initializer) => initializer.await, + let source = async move { + match input { + FlightRecordBatchStreamInput::Ready(source) => Ok(source), + FlightRecordBatchStreamInput::Initializer(initializer) => initializer.await, } } .trace( @@ -193,8 +201,20 @@ impl FlightRecordBatchStream { ) .await; - match recordbatches { - Ok(recordbatches) => { + match source { + Ok(FlightRecordBatchSource::RecordBatches(recordbatches)) => { + let should_send_partial_metrics = query_ctx.explain_verbose(); + let can_send_metrics_before_batch = + query_ctx.explain_verbose() + && query_ctx.live_analyze_metrics_enabled() + && query_ctx + .remote_query_id() + .zip(query_ctx.extension( + SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY, + )) + .is_some_and(|(remote_query_id, capability)| { + capability == remote_query_id + }); Self::flight_data_stream( recordbatches, tx, @@ -203,6 +223,11 @@ impl FlightRecordBatchStream { ) .await; } + Ok(FlightRecordBatchSource::AffectedRows { rows, metrics }) => { + let _ = tx + .send(Ok(FlightMessage::AffectedRows { rows, metrics })) + .await; + } Err(status) => { let _ = tx.send(Err(status)).await; } @@ -485,7 +510,9 @@ mod test { .unwrap() .as_stream(); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), QueryContext::arc(), @@ -519,10 +546,34 @@ mod test { } } + #[tokio::test] + async fn test_flight_record_batch_stream_encodes_affected_rows() { + let mut stream = FlightRecordBatchStream::new( + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::AffectedRows { + rows: 42, + metrics: Some(r#"{"region_watermarks":[]}"#.to_string()), + }), + TracingContext::default(), + FlightCompression::default(), + QueryContext::arc(), + ); + + let data = stream.next().await.unwrap().unwrap(); + let message = FlightDecoder::default().try_decode(&data).unwrap().unwrap(); + assert!(matches!( + message, + FlightMessage::AffectedRows { + rows: 42, + metrics: Some(_), + } + )); + assert!(stream.next().await.is_none()); + } + #[tokio::test] async fn test_flight_record_batch_stream_forwards_initializer_error() { let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::initializer(async { + FlightRecordBatchStreamInput::initializer(async { Err(tonic::Status::unavailable( "remote read initialization failed", )) @@ -552,9 +603,12 @@ mod test { metrics, }); let query_ctx = query_context_with_live_metrics_and_matching_capability(); - query_ctx.set_explain_verbose(true); + let initializer_query_ctx = query_ctx.clone(); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::initializer(async move { + initializer_query_ctx.set_explain_verbose(true); + Ok(FlightRecordBatchSource::RecordBatches(recordbatches)) + }), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -609,7 +663,9 @@ mod test { let query_ctx = query_context_with_live_metrics_and_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -663,7 +719,9 @@ mod test { let query_ctx = query_context_with_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -706,7 +764,9 @@ mod test { let query_ctx = Arc::new(query_ctx); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -745,7 +805,9 @@ mod test { }); let query_ctx = query_context_with_matching_capability(); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), query_ctx, diff --git a/src/servers/src/grpc/greptime_handler.rs b/src/servers/src/grpc/greptime_handler.rs index ad7779435b8..98786d30adc 100644 --- a/src/servers/src/grpc/greptime_handler.rs +++ b/src/servers/src/grpc/greptime_handler.rs @@ -15,11 +15,13 @@ //! Handler for Greptime Database service. It's implemented by frontend. use std::collections::HashMap; +use std::future::Future; use std::str::FromStr; use std::sync::{Arc, RwLock}; use std::time::Instant; use api::helper::request_type; +use api::v1::greptime_request::Request as QueryRequest; use api::v1::{GreptimeRequest, RequestHeader}; use auth::UserProviderRef; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; @@ -79,30 +81,39 @@ impl GreptimeRequestHandler { ) -> Result { let header = request.header.as_ref(); let query_ctx = create_query_context(Channel::Grpc, header, hints, HashMap::new())?; - self.handle_request_with_query_ctx(request, query_ctx).await - } - - pub(crate) async fn handle_request_with_query_ctx( - &self, - request: GreptimeRequest, - query_ctx: QueryContextRef, - ) -> Result { let query = request.request.context(InvalidQuerySnafu { reason: "Expecting non-empty GreptimeRequest.", })?; + self.authenticate_request_with_query_ctx(request.header.as_ref(), &query_ctx) + .await?; + self.handle_request_with_query_ctx(query, query_ctx).await + } - let header = request.header.as_ref(); - let user_info = context_auth::auth(self.user_provider.clone(), header, &query_ctx).await?; + pub(crate) async fn authenticate_request_with_query_ctx( + &self, + header: Option<&RequestHeader>, + query_ctx: &QueryContextRef, + ) -> Result<()> { + let user_info = context_auth::auth(self.user_provider.clone(), header, query_ctx).await?; query_ctx.set_current_user(user_info); + Ok(()) + } + pub(crate) fn handle_request_with_query_ctx( + &self, + query: QueryRequest, + query_ctx: QueryContextRef, + ) -> impl Future> + Send + 'static { let handler = self.handler.clone(); + let runtime = self.runtime.clone(); let request_type = request_type(&query).to_string(); let db = query_ctx.get_db_string(); let timer = RequestTimer::new(db.clone(), request_type); let tracing_context = TracingContext::from_current_span(); - let result_future = async move { - handler + async move { + let result_future = async move { + handler .do_query(query, query_ctx) .trace(tracing_context.attach(tracing::info_span!( "GreptimeRequestHandler::handle_request_runtime" @@ -118,26 +129,27 @@ impl GreptimeRequestHandler { } e }) - }; + }; - match &self.runtime { - Some(runtime) => { - // Executes requests in another runtime to - // 1. prevent the execution from being cancelled unexpected by Tonic runtime; - // - Refer to our blog for the rational behind it: - // https://www.greptime.com/blogs/2023-01-12-hidden-control-flow.html - // - Obtaining a `JoinHandle` to get the panic message (if there's any). - // From its docs, `JoinHandle` is cancel safe. The task keeps running even it's handle been dropped. - // 2. avoid the handler blocks the gRPC runtime incidentally. - runtime - .spawn(result_future) - .await - .context(JoinTaskSnafu) - .inspect_err(|e| { - timer.record(e.status_code()); - })? + match runtime { + Some(runtime) => { + // Executes requests in another runtime to + // 1. prevent the execution from being cancelled unexpected by Tonic runtime; + // - Refer to our blog for the rational behind it: + // https://www.greptime.com/blogs/2023-01-12-hidden-control-flow.html + // - Obtaining a `JoinHandle` to get the panic message (if there's any). + // From its docs, `JoinHandle` is cancel safe. The task keeps running even it's handle been dropped. + // 2. avoid the handler blocks the gRPC runtime incidentally. + runtime + .spawn(result_future) + .await + .context(JoinTaskSnafu) + .inspect_err(|e| { + timer.record(e.status_code()); + })? + } + None => result_future.await, } - None => result_future.await, } } diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index cf75ae7fff1..1a2d841600f 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -16,40 +16,45 @@ mod test { use std::collections::HashMap; use std::net::SocketAddr; + use std::pin::Pin; use std::sync::Arc; use std::time::{Duration, Instant}; use api::v1::auth_header::AuthScheme; + use api::v1::greptime_request::Request as GreptimeQueryRequest; use api::v1::query_request::Query; use api::v1::{Basic, ColumnDataType, ColumnDef, CreateTableExpr, QueryRequest, SemanticType}; use arrow_flight::flight_service_server::FlightServiceServer; use arrow_flight::{FlightData, FlightDescriptor, Ticket}; use auth::user_provider_from_option; + use client::region::RegionRequester; use client::{Client, Database}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; - use common_grpc::flight::do_put::DoPutMetadata; + use common_grpc::flight::do_put::{DoPutMetadata, DoPutResponse}; use common_grpc::flight::{FlightEncoder, FlightMessage}; - use common_query::OutputData; + use common_query::{Output, OutputData}; use common_recordbatch::adapter::RegionWatermarkEntry; use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; use common_telemetry::tracing_context::TracingContext; use datatypes::prelude::{ConcreteDataType, ScalarVector, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{Int32Vector, StringVector, TimestampMillisecondVector}; - use futures_util::StreamExt; + use futures_util::{Stream, StreamExt}; use hyper_util::rt::TokioIo; use itertools::Itertools; use servers::grpc::builder::GrpcServerBuilder; use servers::grpc::flight::{ FlightCraft, FlightCraftWrapper, FlightRecordBatchSource, FlightRecordBatchStream, - TonicStream, + FlightRecordBatchStreamInput, PutRecordBatchRequestStream, TonicStream, }; use servers::grpc::greptime_handler::GreptimeRequestHandler; use servers::grpc::{FlightCompression, GrpcServerConfig}; + use servers::query_handler::grpc::GrpcQueryHandler; use servers::server::Server; - use tonic::Response; + use session::context::QueryContextRef; use tonic::transport::Server as TonicServer; + use tonic::{Response, Status}; use tower::service_fn; use crate::cluster::GreptimeDbClusterBuilder; @@ -59,6 +64,12 @@ mod test { struct SlowFlightCraft; + struct ErrorFlightCraft; + + struct SlowRemoteQueryHandler { + region_requester: RegionRequester, + } + fn slow_recordbatch_stream() -> SendableRecordBatchStream { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( "value", @@ -83,9 +94,11 @@ mod test { _: tonic::Request, ) -> std::result::Result>, tonic::Status> { let stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::initializer(async { + FlightRecordBatchStreamInput::initializer(async { tokio::time::sleep(Duration::from_secs(2)).await; - Ok(slow_recordbatch_stream()) + Ok(FlightRecordBatchSource::RecordBatches( + slow_recordbatch_stream(), + )) }), TracingContext::default(), FlightCompression::default(), @@ -96,14 +109,72 @@ mod test { } } - #[tokio::test(flavor = "multi_thread")] - async fn test_do_get_timeout_does_not_cancel_slow_flight_stream() { + #[async_trait::async_trait] + impl FlightCraft for ErrorFlightCraft { + async fn do_get( + &self, + _: tonic::Request, + ) -> std::result::Result>, tonic::Status> { + let stream = FlightRecordBatchStream::new( + FlightRecordBatchStreamInput::initializer(async { + Err(Status::internal("deferred initializer detail")) + }), + TracingContext::default(), + FlightCompression::default(), + session::context::QueryContext::arc(), + ); + Ok(Response::new(Box::pin(stream))) + } + } + + #[async_trait::async_trait] + impl GrpcQueryHandler for SlowRemoteQueryHandler { + async fn do_query( + &self, + _: GreptimeQueryRequest, + _: QueryContextRef, + ) -> servers::error::Result { + let stream = self + .region_requester + .do_get_inner(Ticket::default()) + .await + .unwrap(); + Ok(Output::new_with_stream(stream)) + } + + fn handle_put_record_batch_stream( + &self, + _: PutRecordBatchRequestStream, + _: QueryContextRef, + ) -> Pin> + Send>> { + Box::pin(futures::stream::empty()) + } + } + + fn client_for_flight_craft(addr: &'static str, craft: T) -> Client + where + T: FlightCraft, + { + client_for_flight_craft_with_max_encoding(addr, craft, None) + } + + fn client_for_flight_craft_with_max_encoding( + addr: &'static str, + craft: T, + max_encoding_message_size: Option, + ) -> Client + where + T: FlightCraft, + { let (client_io, server_io) = tokio::io::duplex(1024); tokio::spawn(async move { + let flight_service = FlightServiceServer::new(FlightCraftWrapper(craft)); + let flight_service = match max_encoding_message_size { + Some(size) => flight_service.max_encoding_message_size(size), + None => flight_service, + }; TonicServer::builder() - .add_service(FlightServiceServer::new(FlightCraftWrapper( - SlowFlightCraft, - ))) + .add_service(flight_service) .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>( server_io, )])) @@ -115,7 +186,7 @@ mod test { let mut client_io = Some(client_io); channel_manager .reset_with_connector( - "slow-flight", + addr, service_fn(move |_| { let client_io = client_io.take(); @@ -127,7 +198,12 @@ mod test { }), ) .unwrap(); - let client = Client::with_manager_and_urls(channel_manager, ["slow-flight"]); + Client::with_manager_and_urls(channel_manager, [addr]) + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_do_get_timeout_does_not_cancel_slow_flight_stream() { + let client = client_for_flight_craft("slow-flight", SlowFlightCraft); let mut flight_client = client.make_flight_client(false, false).unwrap(); let start = Instant::now(); @@ -143,6 +219,64 @@ mod test { assert!(stream.message().await.unwrap().is_some()); } + + #[tokio::test(flavor = "multi_thread")] + async fn test_deferred_flight_initializer_error_preserves_message() { + let client = client_for_flight_craft("error-flight", ErrorFlightCraft); + let mut flight_client = client.make_flight_client(false, false).unwrap(); + let mut stream = flight_client + .mut_inner() + .do_get(tonic::Request::new(Ticket::default())) + .await + .unwrap() + .into_inner(); + let error = stream.message().await.unwrap_err(); + assert_eq!(tonic::Code::Internal, error.code()); + assert_eq!("deferred initializer detail", error.message()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_deferred_flight_encode_error_preserves_message() { + let client = + client_for_flight_craft_with_max_encoding("limited-flight", SlowFlightCraft, Some(1)); + let mut flight_client = client.make_flight_client(false, false).unwrap(); + let mut stream = flight_client + .mut_inner() + .do_get(tonic::Request::new(Ticket::default())) + .await + .unwrap() + .into_inner(); + let error = stream.message().await.unwrap_err(); + assert_eq!(tonic::Code::OutOfRange, error.code()); + assert!(error.message().contains("message length too large")); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_flight_request_timeout_does_not_cancel_slow_datanode_stream() { + let datanode_client = client_for_flight_craft("slow-datanode", SlowFlightCraft); + let frontend_handler = GreptimeRequestHandler::new( + Arc::new(SlowRemoteQueryHandler { + region_requester: RegionRequester::new(datanode_client, false, false), + }), + None, + None, + FlightCompression::default(), + ); + let frontend_client = client_for_flight_craft("slow-frontend", frontend_handler); + let database = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, frontend_client); + + let start = Instant::now(); + let output = database + .flight_request() + .with_timeout(Duration::from_secs(1)) + .sql("select 1") + .await + .unwrap(); + // `sql()` waits for the first Flight message. Succeeding after this delay + // proves the request timeout applied only to the response header. + assert!(start.elapsed() >= Duration::from_secs(1)); + assert!(matches!(output.data, OutputData::Stream(_))); + } #[tokio::test(flavor = "multi_thread")] async fn test_standalone_flight_do_put() { common_telemetry::init_default_ut_logging(); @@ -446,16 +580,14 @@ mod test { ); let result = client - .query_with_terminal_metrics_and_flow_extensions( - QueryRequest { - query: Some(Query::Sql( - "select ts, a, `B` from foo order by ts".to_string(), - )), - }, - &[], - &[("flow.return_region_seq", "true")], - &HashMap::from([(region_id.as_u64(), stale_snapshot_seq)]), - ) + .flight_request() + .with_flow_extensions(&[("flow.return_region_seq", "true")]) + .with_snapshot_seqs(&HashMap::from([(region_id.as_u64(), stale_snapshot_seq)])) + .query_with_terminal_metrics(QueryRequest { + query: Some(Query::Sql( + "select ts, a, `B` from foo order by ts".to_string(), + )), + }) .await .unwrap(); diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 4a8f96260e8..43ba8701d76 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2252,6 +2252,7 @@ query_timeout = "10m" slow_query_threshold = "1m" experimental_min_refresh_duration = "5s" grpc_conn_timeout = "5s" +experimental_flight_do_get_timeout = "10s" experimental_grpc_max_retries = 3 experimental_frontend_scan_timeout = "30s" experimental_max_filter_num_per_query = 20 From b31f05eb5964c17342e7606d80dd04e6b93a3e03 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Thu, 27 Aug 2026 03:37:59 +0000 Subject: [PATCH 06/22] fix: update tokio-postgres and correct explain/fetch cursor output schema (#8955) * chore(deps): update tokio-postgres * fix: describing fetch cursor and analyze --- Cargo.lock | 180 ++++++++++++++++++++------- Cargo.toml | 2 +- src/common/meta/Cargo.toml | 2 +- src/common/recordbatch/src/cursor.rs | 9 ++ src/query/src/analyze.rs | 21 +++- src/query/src/lib.rs | 2 +- src/servers/Cargo.toml | 4 +- src/servers/src/postgres/handler.rs | 29 +++++ src/session/src/lib.rs | 9 ++ 9 files changed, 202 insertions(+), 56 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ec477e5c9a3..d77ec340a8d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4628,10 +4628,23 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb" dependencies = [ "const-oid 0.9.6", + "der_derive", + "flagset", "pem-rfc7468", "zeroize", ] +[[package]] +name = "der_derive" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8034092389675178f570469e6c3b0465d3d30b4505c294a6550db47f3c17ad18" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "deranged" version = "0.5.8" @@ -5334,6 +5347,12 @@ version = "0.5.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" +[[package]] +name = "flagset" +version = "0.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7ac824320a75a52197e8f2d787f6a38b6718bb6897a35142d749af3c0e8f4fe" + [[package]] name = "flatbuffers" version = "25.2.10" @@ -5973,7 +5992,7 @@ dependencies = [ "cfg-if", "js-sys", "libc", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "wasm-bindgen", ] @@ -7866,13 +7885,14 @@ checksum = "f9fbbcab51052fe104eb5e5d351cf728d30a5be1fe14d9be8a3b097481fb97de" [[package]] name = "libredox" -version = "0.1.4" +version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1580801010e535496706ba011c15f8532df6b42297d2e471fec38ceadd8c0638" +checksum = "28d0a00925a9f930d679b6789b721e3a7f9ed110f41b86d2497caa780c3a070a" dependencies = [ "bitflags 2.12.1", "libc", - "redox_syscall 0.5.13", + "plain", + "redox_syscall 0.9.3", ] [[package]] @@ -8543,7 +8563,7 @@ checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" dependencies = [ "libc", "log", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "windows-sys 0.61.2", ] @@ -9358,6 +9378,15 @@ dependencies = [ "objc2-encode", ] +[[package]] +name = "objc2-core-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" +dependencies = [ + "bitflags 2.12.1", +] + [[package]] name = "objc2-encode" version = "4.1.0" @@ -9374,6 +9403,15 @@ dependencies = [ "objc2", ] +[[package]] +name = "objc2-system-configuration" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7216bd11cbda54ccabcab84d523dc93b858ec75ecfb3a7d89513fa22464da396" +dependencies = [ + "objc2-core-foundation", +] + [[package]] name = "object" version = "0.36.7" @@ -10552,16 +10590,7 @@ dependencies = [ "tokio", "tokio-rustls", "tokio-util", - "x509-certificate 0.25.0", -] - -[[package]] -name = "phf" -version = "0.11.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fd6780a80ae0c52cc120a26a1a42c1ae51b247a253e4e06113d23d2c2edd078" -dependencies = [ - "phf_shared 0.11.3", + "x509-certificate", ] [[package]] @@ -10814,6 +10843,12 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" +[[package]] +name = "plain" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4596b6d070b27117e987119b4dac604f3c58cfb0b191112e24771b2faeac1a6" + [[package]] name = "plist" version = "1.7.2" @@ -10924,16 +10959,16 @@ dependencies = [ [[package]] name = "postgres-types" -version = "0.2.9" +version = "0.2.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "613283563cd90e1dfc3518d548caee47e0e725455ed619881f5cf21f36de4b48" +checksum = "851ca9db4932932d69f3ea811b1abe63087a0f740a47692619dd40d4899b68be" dependencies = [ "array-init", "bytes", "chrono", "fallible-iterator", "postgres-protocol", - "serde", + "serde_core", "serde_json", ] @@ -11299,7 +11334,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4" dependencies = [ "bytes", - "heck 0.4.1", + "heck 0.5.0", "itertools 0.12.1", "log", "multimap", @@ -11319,7 +11354,7 @@ version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "itertools 0.14.0", "log", "multimap", @@ -12087,6 +12122,15 @@ dependencies = [ "bitflags 2.12.1", ] +[[package]] +name = "redox_syscall" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d678d17679829e73d371e96880897e98fee2ded7acc0a50bdf8af2affa4b2fe5" +dependencies = [ + "bitflags 2.12.1", +] + [[package]] name = "redox_users" version = "0.4.6" @@ -13785,7 +13829,7 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1961e2ef424c1424204d3a5d6975f934f56b6d50ff5732382d84ebf460e147f7" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "proc-macro2", "quote", "syn 2.0.117", @@ -14117,7 +14161,7 @@ dependencies = [ "stringprep", "thiserror 2.0.17", "tracing", - "whoami", + "whoami 1.6.0", ] [[package]] @@ -14156,7 +14200,7 @@ dependencies = [ "stringprep", "thiserror 2.0.17", "tracing", - "whoami", + "whoami 1.6.0", ] [[package]] @@ -15175,6 +15219,27 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" +[[package]] +name = "tls_codec" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0de2e01245e2bb89d6f05801c564fa27624dbd7b1846859876c7dad82e90bf6b" +dependencies = [ + "tls_codec_derive", + "zeroize", +] + +[[package]] +name = "tls_codec_derive" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d2e76690929402faae40aebdda620a2c0e25dd6d3b9afe48867dfd95991f4bd" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "tokio" version = "1.52.3" @@ -15241,9 +15306,9 @@ dependencies = [ [[package]] name = "tokio-postgres" -version = "0.7.13" +version = "0.7.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c95d533c83082bb6490e0189acaa0bbeef9084e60471b696ca6988cd0541fb0" +checksum = "a528f7d280f6d5b9cd149635c8705b0dd049754bc67d81d31fa25169a93809d3" dependencies = [ "async-trait", "byteorder", @@ -15254,29 +15319,29 @@ dependencies = [ "log", "parking_lot 0.12.4", "percent-encoding", - "phf 0.11.3", + "phf 0.13.1", "pin-project-lite", "postgres-protocol", "postgres-types", - "rand 0.9.4", - "socket2 0.5.10", + "rand 0.10.1", + "socket2 0.6.4", "tokio", "tokio-util", - "whoami", + "whoami 2.1.3", ] [[package]] name = "tokio-postgres-rustls" -version = "0.12.0" +version = "0.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04fb792ccd6bbcd4bba408eb8a292f70fc4a3589e5d793626f45190e6454b6ab" +checksum = "4c2ad44aa0ae96db89c4742212ed41645b2f597311ff6e1945542a4d9fadc2fb" dependencies = [ - "ring", "rustls", + "sha2 0.11.0", "tokio", "tokio-postgres", "tokio-rustls", - "x509-certificate 0.23.1", + "x509-cert", ] [[package]] @@ -16284,6 +16349,15 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +[[package]] +name = "wasi" +version = "0.14.7+wasi-0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "883478de20367e224c0090af9cf5f9fa85bed63a95c1abf3afc5c083ebc06e8c" +dependencies = [ + "wasip2", +] + [[package]] name = "wasip2" version = "1.0.2+wasi-0.2.9" @@ -16308,6 +16382,15 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8dad83b4f25e74f184f64c43b150b91efe7647395b42289f38e50566d82855b" +[[package]] +name = "wasite" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66fe902b4a6b8028a753d5424909b764ccf79b7a209eac9bf97e59cda9f71a42" +dependencies = [ + "wasi 0.14.7+wasi-0.2.4", +] + [[package]] name = "wasm-bindgen" version = "0.2.122" @@ -16530,7 +16613,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6994d13118ab492c3c80c1f81928718159254c53c472bf9ce36f8dae4add02a7" dependencies = [ "redox_syscall 0.5.13", - "wasite", + "wasite 0.1.0", +] + +[[package]] +name = "whoami" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "626c4bac6755d76ffc12cb01b2eac751db1996b9e0041de9aa02c8c211ddc82c" +dependencies = [ + "libc", + "libredox", + "objc2-system-configuration", + "wasite 1.0.2", "web-sys", ] @@ -17139,22 +17234,15 @@ dependencies = [ ] [[package]] -name = "x509-certificate" -version = "0.23.1" +name = "x509-cert" +version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "66534846dec7a11d7c50a74b7cdb208b9a581cad890b7866430d438455847c85" +checksum = "1301e935010a701ae5f8655edc0ad17c44bad3ac5ce8c39185f75453b720ae94" dependencies = [ - "bcder", - "bytes", - "chrono", + "const-oid 0.9.6", "der", - "hex", - "pem", - "ring", - "signature", "spki", - "thiserror 1.0.69", - "zeroize", + "tls_codec", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 7c3fccfbbf8..fab7b2aa736 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -255,7 +255,7 @@ strum = { version = "0.27", features = ["derive"] } sysinfo = "0.33" tempfile = "3" tokio = { version = "1.47", features = ["full"] } -tokio-postgres = "0.7" +tokio-postgres = "0.7.18" tokio-rustls = { version = "0.26.2", default-features = false } tokio-stream = "0.1" tokio-util = { version = "0.7", features = ["io-util", "compat"] } diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index 8e0cbbdf354..8deb3d00814 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -89,7 +89,7 @@ strum.workspace = true table = { workspace = true, features = ["testing"] } tokio.workspace = true tokio-postgres = { workspace = true, optional = true } -tokio-postgres-rustls = { version = "0.12", optional = true } +tokio-postgres-rustls = { version = "0.14", optional = true } tonic.workspace = true tracing.workspace = true typetag.workspace = true diff --git a/src/common/recordbatch/src/cursor.rs b/src/common/recordbatch/src/cursor.rs index a741953ccc2..60ae9d54fa2 100644 --- a/src/common/recordbatch/src/cursor.rs +++ b/src/common/recordbatch/src/cursor.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use datatypes::schema::SchemaRef; use futures::StreamExt; use tokio::sync::Mutex; @@ -28,12 +29,15 @@ struct Inner { /// A cursor on RecordBatchStream that fetches data batch by batch pub struct RecordBatchStreamCursor { + schema: SchemaRef, inner: Mutex, } impl RecordBatchStreamCursor { pub fn new(stream: SendableRecordBatchStream) -> RecordBatchStreamCursor { + let schema = stream.schema(); Self { + schema, inner: Mutex::new(Inner { stream, current_row_index: 0, @@ -43,6 +47,11 @@ impl RecordBatchStreamCursor { } } + /// Returns the schema of the underlying record batch stream. + pub fn schema(&self) -> SchemaRef { + self.schema.clone() + } + /// Take `size` of row from the `RecordBatchStream` and create a new /// `RecordBatch` for these rows. pub async fn take(&self, size: usize) -> Result { diff --git a/src/query/src/analyze.rs b/src/query/src/analyze.rs index ec0abf25ba6..b1522fc05d5 100644 --- a/src/query/src/analyze.rs +++ b/src/query/src/analyze.rs @@ -46,6 +46,21 @@ const STAGE: &str = "stage"; const NODE: &str = "node"; const PLAN: &str = "plan"; +/// Returns the fixed output schema of [`DistAnalyzeExec`]: +/// (`stage`: UInt32, `node`: UInt32, `plan`: Utf8). +/// +/// Exposed so protocol-level `Describe` handlers can advertise the schema +/// clients will actually receive (the physical plan is rewritten in +/// `optimize_physical_plan` and its schema differs from the DataFusion +/// `Analyze` logical plan schema). +pub fn dist_analyze_output_schema() -> SchemaRef { + SchemaRef::new(Schema::new(vec![ + Field::new(STAGE, DataType::UInt32, true), + Field::new(NODE, DataType::UInt32, true), + Field::new(PLAN, DataType::Utf8, true), + ])) +} + #[derive(Debug)] pub struct DistAnalyzeExec { input: Arc, @@ -58,11 +73,7 @@ pub struct DistAnalyzeExec { impl DistAnalyzeExec { /// Create a new DistAnalyzeExec pub fn new(input: Arc, verbose: bool, format: AnalyzeFormat) -> Self { - let schema = SchemaRef::new(Schema::new(vec![ - Field::new(STAGE, DataType::UInt32, true), - Field::new(NODE, DataType::UInt32, true), - Field::new(PLAN, DataType::Utf8, true), - ])); + let schema = dist_analyze_output_schema(); let properties = Arc::new(Self::compute_properties(&input, schema.clone())); Self { input, diff --git a/src/query/src/lib.rs b/src/query/src/lib.rs index 68d8ff3a9eb..4b2a2d5264d 100644 --- a/src/query/src/lib.rs +++ b/src/query/src/lib.rs @@ -45,7 +45,7 @@ pub(crate) mod test_util; #[cfg(test)] mod tests; -pub use crate::analyze::analyze_plan_metrics_to_json_value; +pub use crate::analyze::{analyze_plan_metrics_to_json_value, dist_analyze_output_schema}; pub use crate::datafusion::DfContextProviderAdapter; pub use crate::query_engine::{ QueryEngine, QueryEngineContext, QueryEngineFactory, QueryEngineRef, diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 7936c96fbda..bb358ea76a6 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -167,8 +167,8 @@ servers = { workspace = true, features = ["testing"] } session = { workspace = true, features = ["testing"] } table.workspace = true tempfile = "3.0.0" -tokio-postgres = "0.7" -tokio-postgres-rustls = "0.12" +tokio-postgres.workspace = true +tokio-postgres-rustls = "0.14" [target.'cfg(unix)'.dev-dependencies] pprof = { version = "0.14", features = ["criterion", "flamegraph"] } diff --git a/src/servers/src/postgres/handler.rs b/src/servers/src/postgres/handler.rs index 484bb6a1f11..433eb11cdad 100644 --- a/src/servers/src/postgres/handler.rs +++ b/src/servers/src/postgres/handler.rs @@ -40,6 +40,7 @@ use pgwire::error::{ErrorInfo, PgWireError, PgWireResult}; use pgwire::messages::PgWireBackendMessage; use pgwire::messages::copy::CopyData; use pgwire::messages::data::DataRow; +use query::dist_analyze_output_schema; use query::planner::DfLogicalPlanner; use query::query_engine::DescribeResult; use session::Session; @@ -554,6 +555,16 @@ fn describe_fields( session: &Arc, ) -> PgWireResult> { match sql_plan { + // EXPLAIN ANALYZE: at execution time the physical plan is replaced with + // GreptimeDB's `DistAnalyzeExec` (see `optimize_physical_plan`), whose + // output schema (stage/node/plan) differs from the DataFusion `Analyze` + // logical plan schema (plan_type/plan). Describe with the schema the + // client will actually receive so the DataRow field count matches. + SqlPlan::Plan(LogicalPlan::Analyze(_), _) => { + let schema: Schema = + Schema::try_from(dist_analyze_output_schema()).map_err(convert_err)?; + schema_to_pg(&schema, format, None).map_err(convert_err) + } // query SqlPlan::Plan(plan, _) if !matches!(plan, LogicalPlan::Dml(_) | LogicalPlan::Ddl(_)) => { let schema: Schema = plan.schema().clone().try_into().map_err(convert_err)?; @@ -679,6 +690,24 @@ fn describe_fields( Ok(vec![]) } } + // FETCH cursor: return the cursor's schema so the RowDescription + // matches the DataRow field count sent during Execute. + SqlPlan::Statement(Statement::FetchCursor(fetch), _) => { + let cursor_name = fetch.cursor_name.to_string(); + match session.get_cursor(&cursor_name) { + Some(cursor) => { + // `cursor.schema()` is the GreptimeDB `SchemaRef` captured + // when the cursor was declared; `schema_to_pg` accepts it + // directly (no DataFusion -> GreptimeDB conversion needed). + schema_to_pg(&cursor.schema(), format, None).map_err(convert_err) + } + None => { + // Cursor not found (e.g. DECLARE hasn't executed yet in + // an extended-protocol batch). Return NoData as a fallback. + Ok(vec![]) + } + } + } _ => { // NoData Ok(vec![]) diff --git a/src/session/src/lib.rs b/src/session/src/lib.rs index 110560b3243..44f3ee66891 100644 --- a/src/session/src/lib.rs +++ b/src/session/src/lib.rs @@ -111,6 +111,15 @@ impl Session { .into() } + /// Returns the cursor with the given name, if it exists. + /// + /// Cursors are stored in the session's mutable inner data and shared across + /// all query contexts created from this session. + pub fn get_cursor(&self, name: &str) -> Option> { + let guard = self.mutable_inner.read().unwrap(); + guard.cursors.get(name).cloned() + } + pub fn conn_info(&self) -> &ConnInfo { &self.conn_info } From a7ba8b01e1c8ced20a19e39ac68a4637dc64e266 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 27 Aug 2026 04:24:07 +0000 Subject: [PATCH 07/22] fix(flow): restore FrontendClient::sql API (#8963) fix(flow): restore batching frontend SQL Signed-off-by: WenyXu --- src/flow/src/batching_mode/frontend_client.rs | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/src/flow/src/batching_mode/frontend_client.rs b/src/flow/src/batching_mode/frontend_client.rs index 27db96748a6..48314dfa21b 100644 --- a/src/flow/src/batching_mode/frontend_client.rs +++ b/src/flow/src/batching_mode/frontend_client.rs @@ -18,6 +18,7 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex, RwLock, Weak}; use api::v1::greptime_request::Request; +use api::v1::query_request::Query; use api::v1::{CreateTableExpr, QueryRequest}; use client::{Client, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, Database, OutputWithMetrics}; use common_error::ext::BoxedError; @@ -339,6 +340,53 @@ impl FrontendClient { }) } + /// Execute a SQL statement on the frontend. + pub async fn sql(&self, catalog: &str, schema: &str, sql: &str) -> Result { + match self { + FrontendClient::Distributed { .. } => { + let db = self.get_random_active_frontend(catalog, schema).await?; + db.database + .sql(sql) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu) + } + FrontendClient::Standalone { + database_client, .. + } => { + let ctx = QueryContextBuilder::default() + .current_catalog(catalog.to_string()) + .current_schema(schema.to_string()) + .build(); + let ctx = Arc::new(ctx); + { + let database_client = { + database_client + .handler + .lock() + .unwrap() + .as_ref() + .context(UnexpectedSnafu { + reason: "Standalone's frontend instance is not set", + })? + .upgrade() + .context(UnexpectedSnafu { + reason: "Failed to upgrade database client", + })? + }; + let req = Request::Query(QueryRequest { + query: Some(Query::Sql(sql.to_string())), + }); + database_client + .do_query(req, ctx) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu) + } + } + } + } + /// Execute a flow query and return terminal metrics. `snapshot_seqs` are /// optional read upper bounds used only by snapshot-fenced repair chunks. pub(crate) async fn query_with_terminal_metrics( From beded6e232918147b93f91677e5ec8b8768d1be5 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 27 Aug 2026 04:51:39 +0000 Subject: [PATCH 08/22] fix(flow): avoid insert select HTTP/2 stalls (#8962) * fix(flow): avoid insert select HTTP/2 stalls Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): share record batch forwarding for DML Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/cmd/src/frontend.rs | 2 + src/query/src/datafusion.rs | 126 +++++++++++++++++++++++++++++------- 2 files changed, 103 insertions(+), 25 deletions(-) diff --git a/src/cmd/src/frontend.rs b/src/cmd/src/frontend.rs index e3737088e8d..a303041b5a2 100644 --- a/src/cmd/src/frontend.rs +++ b/src/cmd/src/frontend.rs @@ -434,6 +434,8 @@ impl StartCommand { // Some queries are expected to take long time. let mut channel_config = opts.datanode.client.channel_config(); channel_config.timeout = None; + // Source Flight streams and sink unary responses share pooled connections. + channel_config.http2_adaptive_window = Some(true); if opts.grpc.flight_compression.transport_compression() { channel_config.accept_compression = true; channel_config.send_compression = true; diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 754199a78fc..f9d87782a03 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -30,7 +30,7 @@ use common_function::function::FunctionContext; use common_function::function_factory::ScalarFunctionFactory; use common_query::{Output, OutputData, OutputMeta}; use common_recordbatch::adapter::{RecordBatchStreamAdapter, RegionQueryStatCounters}; -use common_recordbatch::{EmptyRecordBatchStream, SendableRecordBatchStream}; +use common_recordbatch::{EmptyRecordBatchStream, RecordBatch, SendableRecordBatchStream}; use common_telemetry::tracing; use datafusion::catalog::TableFunction; use datafusion::dataframe::DataFrame; @@ -44,6 +44,7 @@ use datafusion_expr::{ use datatypes::prelude::VectorRef; use datatypes::schema::Schema; use futures_util::StreamExt; +use futures_util::future::try_join; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt, ensure}; use sqlparser::ast::AnalyzeFormat; @@ -80,6 +81,29 @@ pub const QUERY_PARALLELISM_HINT: &str = "query_parallelism"; /// Whether to fallback to the original plan when failed to push down. pub const QUERY_FALLBACK_HINT: &str = "query_fallback"; +// An unbounded queue keeps draining source RPCs while mutation RPCs on a shared +// HTTP/2 connection are pending, trading bounded memory for request liveness. +async fn forward_record_batches( + mut stream: SendableRecordBatchStream, + batch_tx: tokio::sync::mpsc::UnboundedSender>, +) -> Result<()> { + while let Some(batch) = stream.next().await { + match batch.context(CreateRecordBatchSnafu) { + Ok(batch) => { + if batch_tx.send(Ok(batch)).is_err() { + break; + } + tokio::task::yield_now().await; + } + Err(error) => { + let _ = batch_tx.send(Err(error)); + break; + } + } + } + Ok(()) +} + fn query_load_region_id(plan: &Arc) -> Option { let mut region_id = None; let mut stack = vec![plan.clone()]; @@ -205,7 +229,7 @@ impl DatafusionQueryEngine { let Output { data, meta } = self .exec_query_plan((*dml.input).clone(), query_ctx.clone()) .await?; - let mut stream = match data { + let stream = match data { OutputData::RecordBatches(batches) => batches.as_stream(), OutputData::Stream(stream) => stream, _ => unreachable!(), @@ -214,30 +238,35 @@ impl DatafusionQueryEngine { let mut affected_rows = 0; let mut insert_cost = 0; - while let Some(batch) = stream.next().await { - let batch = batch.context(CreateRecordBatchSnafu)?; - let column_vectors = batch - .column_vectors(&table_name.to_string(), table.schema()) - .map_err(BoxedError::new) - .context(QueryExecutionSnafu)?; - - match dml.op { - WriteOp::Insert(_) => { - // We ignore the insert op. - let output = self - .insert(&table_name, column_vectors, query_ctx.clone()) - .await?; - let (rows, cost) = output.extract_rows_and_cost(); - affected_rows += rows; - insert_cost += cost; - } - WriteOp::Delete => { - affected_rows += self - .delete(&table_name, &table, column_vectors, query_ctx.clone()) - .await?; - } - _ => unreachable!("guarded by the 'ensure!' at the beginning"), + match dml.op { + WriteOp::Insert(_) => { + let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel(); + let producer = forward_record_batches(stream, batch_tx); + let consumer = self.consume_insert_record_batches( + batch_rx, + &table_name, + table.schema(), + query_ctx.clone(), + ); + let ((), (rows, cost)) = try_join(producer, consumer).await?; + affected_rows += rows; + insert_cost += cost; } + WriteOp::Delete => { + // Keep DELETE on the same producer/consumer schedule as INSERT so the source + // stream can continue draining while mutation RPCs are pending. + let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel(); + let producer = forward_record_batches(stream, batch_tx); + let consumer = self.consume_delete_record_batches( + batch_rx, + &table_name, + &table, + query_ctx.clone(), + ); + let (rows, ()) = try_join(consumer, producer).await?; + affected_rows += rows; + } + _ => unreachable!("guarded by the 'ensure!' at the beginning"), } Ok(Output::new( OutputData::AffectedRows(affected_rows), @@ -245,6 +274,53 @@ impl DatafusionQueryEngine { )) } + async fn consume_insert_record_batches( + &self, + mut batch_rx: tokio::sync::mpsc::UnboundedReceiver>, + table_name: &ResolvedTableReference, + table_schema: Arc, + query_ctx: QueryContextRef, + ) -> Result<(usize, usize)> { + let mut affected_rows = 0; + let mut insert_cost = 0; + while let Some(batch) = batch_rx.recv().await { + let batch = batch?; + let column_vectors = batch + .column_vectors(&table_name.to_string(), table_schema.clone()) + .map_err(BoxedError::new) + .context(QueryExecutionSnafu)?; + // We ignore the insert op. + let output = self + .insert(table_name, column_vectors, query_ctx.clone()) + .await?; + let (rows, cost) = output.extract_rows_and_cost(); + affected_rows += rows; + insert_cost += cost; + } + Ok((affected_rows, insert_cost)) + } + + async fn consume_delete_record_batches( + &self, + mut batch_rx: tokio::sync::mpsc::UnboundedReceiver>, + table_name: &ResolvedTableReference, + table: &TableRef, + query_ctx: QueryContextRef, + ) -> Result { + let mut affected_rows = 0; + while let Some(batch) = batch_rx.recv().await { + let batch = batch?; + let column_vectors = batch + .column_vectors(&table_name.to_string(), table.schema()) + .map_err(BoxedError::new) + .context(QueryExecutionSnafu)?; + affected_rows += self + .delete(table_name, table, column_vectors, query_ctx.clone()) + .await?; + } + Ok(affected_rows) + } + #[tracing::instrument(skip_all)] async fn delete( &self, From 9befb7e64cdccf25e96b36fa5ca8fa98f1901dd7 Mon Sep 17 00:00:00 2001 From: sun Date: Thu, 27 Aug 2026 09:57:18 +0000 Subject: [PATCH 09/22] feat: update dashboard to v0.13.14 (#8968) --- src/servers/dashboard/VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/servers/dashboard/VERSION b/src/servers/dashboard/VERSION index b68b26e6f78..0822cdb1e68 100644 --- a/src/servers/dashboard/VERSION +++ b/src/servers/dashboard/VERSION @@ -1 +1 @@ -v0.13.13 +v0.13.14 From a2f39ecf7b22da2a99502ff2838d51d90a502e66 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Thu, 27 Aug 2026 17:27:52 +0000 Subject: [PATCH 10/22] feat: harden frontend heartbeat extensions (#8803) Add a typed `Command::build_with_heartbeat_extensions` seam in `src/cmd/src/frontend.rs`. Freeze and harden `FrontendHeartbeatExtensions` in `src/frontend/src/heartbeat.rs`, with lifecycle and race coverage in `src/frontend/src/heartbeat/tests.rs`. Signed-off-by: Lei, HUANG --- src/cmd/src/frontend.rs | 209 +++++++++++++++++++++++++--- src/frontend/src/heartbeat.rs | 69 ++++++++- src/frontend/src/heartbeat/tests.rs | 167 +++++++++++++++++++++- 3 files changed, 418 insertions(+), 27 deletions(-) diff --git a/src/cmd/src/frontend.rs b/src/cmd/src/frontend.rs index a303041b5a2..1f95a3c54b4 100644 --- a/src/cmd/src/frontend.rs +++ b/src/cmd/src/frontend.rs @@ -13,7 +13,9 @@ // limitations under the License. use std::fmt::Debug; +use std::future::Future; use std::path::Path; +use std::pin::Pin; use std::sync::Arc; use std::time::Duration; @@ -61,6 +63,7 @@ use crate::options::{GlobalOptions, GreptimeOptions}; use crate::{App, create_resource_limit_metrics, log_versions, maybe_activate_heap_profile}; type FrontendOptions = GreptimeOptions; +type HeartbeatExtensionSetupFuture<'a> = Pin> + Send + 'a>>; pub struct Instance { frontend: Frontend, @@ -117,7 +120,35 @@ pub struct Command { impl Command { pub async fn build(&self, opts: FrontendOptions) -> Result { - self.subcmd.build(opts).await + self.build_with_heartbeat_extensions(opts, FrontendHeartbeatExtensions::default()) + .await + } + + /// Builds a frontend with pre-registered heartbeat extensions. + /// + /// Register extensions before calling this method. The supplied registry is installed before + /// normal plugin setup. After plugin heartbeat setup completes, the registry is frozen before + /// heartbeat handlers and the heartbeat task are built, so later registrations are rejected + /// and every heartbeat consumer observes the same extension membership. + pub async fn build_with_heartbeat_extensions( + &self, + opts: FrontendOptions, + heartbeat_extensions: FrontendHeartbeatExtensions, + ) -> Result { + let plugins = plugins_with_heartbeat_extensions(heartbeat_extensions); + self.forward_build_with_plugins(opts, plugins, |command, opts, plugins| { + command.build_with_plugins(opts, plugins) + }) + .await + } + + fn forward_build_with_plugins<'a, T>( + &'a self, + opts: FrontendOptions, + plugins: Plugins, + build: impl FnOnce(&'a StartCommand, FrontendOptions, Plugins) -> T, + ) -> T { + self.subcmd.forward_build_with_plugins(opts, plugins, build) } pub fn load_options(&self, global_options: &GlobalOptions) -> Result { @@ -131,9 +162,14 @@ pub enum SubCommand { } impl SubCommand { - async fn build(&self, opts: FrontendOptions) -> Result { + fn forward_build_with_plugins<'a, T>( + &'a self, + opts: FrontendOptions, + plugins: Plugins, + build: impl FnOnce(&'a StartCommand, FrontendOptions, Plugins) -> T, + ) -> T { match self { - SubCommand::Start(cmd) => cmd.build(opts).await, + SubCommand::Start(cmd) => build(cmd, opts, plugins), } } @@ -329,7 +365,11 @@ impl StartCommand { Ok(()) } - async fn build(&self, opts: FrontendOptions) -> Result { + async fn build_with_plugins( + &self, + opts: FrontendOptions, + plugins: Plugins, + ) -> Result { let guard = common_telemetry::init_global_logging( APP_NAME, &opts.component.logging, @@ -380,15 +420,8 @@ impl StartCommand { .await .context(error::MetaClientInitSnafu)?; - let mut plugins = Plugins::new(); - plugins::setup_frontend_plugins_pre_build( - &mut plugins, - &plugin_opts, - &opts, - Some(&meta_config), - ) - .await - .context(error::StartFrontendSnafu)?; + let mut plugins = + prepare_frontend_plugins(plugins, &plugin_opts, &opts, Some(&meta_config)).await?; // now initialize the meta_client with plugins let meta_client = meta_client::create_meta_client( @@ -499,12 +532,20 @@ impl StartCommand { let instance = Arc::new(instance); - plugins::setup_frontend_heartbeat_extensions(&mut plugins, &plugin_opts, &instance) - .await - .context(error::StartFrontendSnafu)?; - let heartbeat_extensions = plugins - .get::() - .unwrap_or_default(); + let heartbeat_instance = instance.clone(); + let heartbeat_extensions = + setup_and_freeze_frontend_heartbeat_extensions(&mut plugins, move |plugins| { + Box::pin(async move { + plugins::setup_frontend_heartbeat_extensions( + plugins, + &plugin_opts, + &heartbeat_instance, + ) + .await + .context(error::StartFrontendSnafu) + }) + }) + .await?; let heartbeat_task = Some(create_heartbeat_task_with_extensions( &opts, meta_client, @@ -526,6 +567,36 @@ impl StartCommand { } } +async fn prepare_frontend_plugins( + mut plugins: Plugins, + plugin_opts: &[PluginOptions], + opts: &frontend::frontend::FrontendOptions, + meta_config: Option<&[PluginOptions]>, +) -> Result { + plugins::setup_frontend_plugins_pre_build(&mut plugins, plugin_opts, opts, meta_config) + .await + .context(error::StartFrontendSnafu)?; + Ok(plugins) +} + +fn plugins_with_heartbeat_extensions(heartbeat_extensions: FrontendHeartbeatExtensions) -> Plugins { + let plugins = Plugins::new(); + plugins.insert(heartbeat_extensions); + plugins +} + +async fn setup_and_freeze_frontend_heartbeat_extensions( + plugins: &mut Plugins, + setup: impl for<'a> FnOnce(&'a mut Plugins) -> HeartbeatExtensionSetupFuture<'a>, +) -> Result { + setup(plugins).await?; + let extensions = plugins + .get::() + .unwrap_or_default(); + extensions.freeze(); + Ok(extensions) +} + pub fn create_heartbeat_task( options: &frontend::frontend::FrontendOptions, meta_client: MetaClientRef, @@ -565,6 +636,7 @@ fn create_heartbeat_task_with_extensions( #[cfg(test)] mod tests { use std::io::Write; + use std::sync::Arc; use std::time::Duration; use auth::{Identity, Password, UserProviderRef}; @@ -578,6 +650,105 @@ mod tests { use super::*; use crate::options::GlobalOptions; + #[derive(Debug)] + struct TestExtension(&'static str); + + #[async_trait] + impl frontend::heartbeat::FrontendHeartbeatExtension for TestExtension { + fn name(&self) -> &'static str { + self.0 + } + } + + #[test] + fn test_command_forwards_heartbeat_extensions_to_start_build() { + let command = Command { + subcmd: SubCommand::Start(StartCommand::default()), + }; + let extensions = FrontendHeartbeatExtensions::default(); + let shared_extensions = extensions.clone(); + + let forwarded_plugins = command.forward_build_with_plugins( + FrontendOptions::default(), + plugins_with_heartbeat_extensions(extensions), + |_, _, plugins| plugins, + ); + let forwarded_extensions = forwarded_plugins + .get::() + .unwrap(); + + assert_registry_identity(&shared_extensions, &forwarded_extensions); + } + + #[tokio::test] + async fn test_prefilled_heartbeat_extensions_survive_pre_build_setup() { + let extensions = FrontendHeartbeatExtensions::default(); + let shared_extensions = extensions.clone(); + + let plugins = prepare_frontend_plugins( + plugins_with_heartbeat_extensions(extensions), + &[], + &frontend::frontend::FrontendOptions::default(), + None, + ) + .await + .unwrap(); + + let setup_extensions = plugins.get_or_insert(FrontendHeartbeatExtensions::default); + assert_registry_identity(&shared_extensions, &setup_extensions); + } + + #[tokio::test] + async fn test_heartbeat_extension_setup_completes_before_freeze() { + let extensions = FrontendHeartbeatExtensions::default(); + assert_eq!( + extensions.try_register(Arc::new(TestExtension("caller"))), + Ok(()) + ); + let mut plugins = plugins_with_heartbeat_extensions(extensions.clone()); + + let finalized = setup_and_freeze_frontend_heartbeat_extensions(&mut plugins, |plugins| { + Box::pin(async move { + let setup_extensions = plugins.get_or_insert(FrontendHeartbeatExtensions::default); + assert_eq!( + setup_extensions.try_register(Arc::new(TestExtension("plugin"))), + Ok(()) + ); + Ok(()) + }) + }) + .await + .unwrap(); + + assert_eq!( + finalized + .extensions() + .iter() + .map(|extension| extension.name()) + .collect::>(), + ["caller", "plugin"] + ); + assert_eq!( + extensions.try_register(Arc::new(TestExtension("late"))), + Err(frontend::heartbeat::RegistrationError::Frozen) + ); + assert_eq!(finalized.len(), 2); + } + + fn assert_registry_identity( + expected: &FrontendHeartbeatExtensions, + actual: &FrontendHeartbeatExtensions, + ) { + let extension = Arc::new(TestExtension("cmd-test-extension")); + assert!(expected.register(extension.clone())); + let registered = actual.extensions(); + assert_eq!(registered.len(), 1); + assert!(Arc::ptr_eq( + &(extension as Arc), + ®istered[0] + )); + } + #[test] fn test_try_from_start_command() { let command = StartCommand { diff --git a/src/frontend/src/heartbeat.rs b/src/frontend/src/heartbeat.rs index 0d9cde74fbe..b7e3a7aff48 100644 --- a/src/frontend/src/heartbeat.rs +++ b/src/frontend/src/heartbeat.rs @@ -16,8 +16,9 @@ mod tests; use std::collections::{HashMap, HashSet}; +use std::fmt::{Display, Formatter}; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; -use std::sync::{Arc, Mutex as StdMutex}; +use std::sync::{Arc, Mutex as StdMutex, MutexGuard, PoisonError}; use api::v1::meta::heartbeat_request::NodeWorkloads; use api::v1::meta::{FrontendWorkloads, HeartbeatRequest, HeartbeatResponse, NodeInfo, Peer}; @@ -93,34 +94,88 @@ pub trait FrontendHeartbeatExtension: Send + Sync { } /// A shareable, ordered registry of frontend heartbeat extensions. +/// +/// Registration is available until [`FrontendHeartbeatExtensions::freeze`] is called. After that, +/// membership is immutable so every heartbeat consumer observes the same extensions. #[derive(Clone, Default)] pub struct FrontendHeartbeatExtensions { inner: Arc>, } +/// The reason a frontend heartbeat extension was not registered. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RegistrationError { + /// An extension with the same name is already registered. + Duplicate, + /// The registry is frozen and no longer accepts registrations. + Frozen, +} + +impl Display for RegistrationError { + fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { + match self { + Self::Duplicate => formatter.write_str("heartbeat extension is already registered"), + Self::Frozen => formatter.write_str("heartbeat extension registry is frozen"), + } + } +} + +impl std::error::Error for RegistrationError {} + #[derive(Default)] struct FrontendHeartbeatExtensionsInner { names: HashSet, extensions: Vec>, + frozen: bool, } impl FrontendHeartbeatExtensions { + fn lock(&self) -> MutexGuard<'_, FrontendHeartbeatExtensionsInner> { + self.inner.lock().unwrap_or_else(PoisonError::into_inner) + } + /// Registers an extension without replacing an existing extension of the same name. /// - /// Returns `true` for a new registration and `false` for an idempotent duplicate. + /// Returns `true` for a new registration and `false` for an idempotent duplicate or when the + /// registry is frozen. Use [`FrontendHeartbeatExtensions::try_register`] when the rejection + /// reason must not be discarded. pub fn register(&self, extension: Arc) -> bool { - let mut inner = self.inner.lock().unwrap(); + self.try_register(extension).is_ok() + } + + /// Registers an extension and returns the reason when registration is rejected. + /// + /// If the registry is already frozen, this method returns [`RegistrationError::Frozen`] + /// without calling [`FrontendHeartbeatExtension::name`]. If freezing races with name + /// evaluation, the registry state is checked again before insertion. + pub fn try_register( + &self, + extension: Arc, + ) -> std::result::Result<(), RegistrationError> { + if self.lock().frozen { + return Err(RegistrationError::Frozen); + } + let name = extension.name().to_string(); + let mut inner = self.lock(); + if inner.frozen { + return Err(RegistrationError::Frozen); + } if !inner.names.insert(name) { - return false; + return Err(RegistrationError::Duplicate); } inner.extensions.push(extension); - true + Ok(()) + } + + /// Prevents further registrations while preserving the current registration order. + pub fn freeze(&self) { + self.lock().frozen = true; } /// Returns the registered extensions in registration order. pub fn extensions(&self) -> Vec> { - self.inner.lock().unwrap().extensions.clone() + self.lock().extensions.clone() } /// Returns response handlers in extension registration order. @@ -133,7 +188,7 @@ impl FrontendHeartbeatExtensions { /// Returns the number of registered extensions. pub fn len(&self) -> usize { - self.inner.lock().unwrap().extensions.len() + self.lock().extensions.len() } /// Returns whether no extension is registered. diff --git a/src/frontend/src/heartbeat/tests.rs b/src/frontend/src/heartbeat/tests.rs index d3c651de326..8f5bfa614b2 100644 --- a/src/frontend/src/heartbeat/tests.rs +++ b/src/frontend/src/heartbeat/tests.rs @@ -14,8 +14,9 @@ use std::collections::{HashMap, VecDeque}; use std::future::pending; +use std::panic::{AssertUnwindSafe, catch_unwind}; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Barrier, Mutex}; use std::time::Duration; use api::v1::meta::heartbeat_request::NodeWorkloads; @@ -494,6 +495,170 @@ impl FrontendHeartbeatExtension for TestExtension { } } +struct PanickingNameExtension; + +#[async_trait] +impl FrontendHeartbeatExtension for PanickingNameExtension { + fn name(&self) -> &str { + panic!("extension name panic") + } +} + +#[test] +fn test_panicking_extension_name_does_not_corrupt_registry() { + let extensions = FrontendHeartbeatExtensions::default(); + + let result = catch_unwind(AssertUnwindSafe(|| { + extensions.register(Arc::new(PanickingNameExtension)); + })); + assert!(result.is_err()); + + let valid = TestExtension::new("valid"); + assert!(extensions.register(valid.clone())); + let registered = extensions.extensions(); + assert_eq!(registered.len(), 1); + assert!(Arc::ptr_eq( + &(valid as Arc), + ®istered[0] + )); + assert_eq!(extensions.len(), 1); +} + +#[test] +fn test_registry_recovers_from_poisoned_lock() { + let extensions = FrontendHeartbeatExtensions::default(); + let poisoned = extensions.clone(); + + let result = catch_unwind(AssertUnwindSafe(move || { + let _guard = poisoned.lock(); + panic!("poison registry lock") + })); + assert!(result.is_err()); + + let valid = TestExtension::new("valid"); + assert!(extensions.register(valid.clone())); + let registered = extensions.extensions(); + assert_eq!(registered.len(), 1); + assert!(Arc::ptr_eq( + &(valid as Arc), + ®istered[0] + )); + assert_eq!(extensions.len(), 1); +} + +#[test] +fn test_registry_freeze_rejects_late_registration_without_changing_membership() { + let extensions = FrontendHeartbeatExtensions::default(); + let calls = Arc::new(AtomicUsize::new(0)); + let handler: HeartbeatResponseHandlerRef = Arc::new(ShortCircuitHandler { + result: ShortCircuitResult::Done, + calls, + }); + let registered = TestExtension::with_handler("registered", handler.clone()); + + assert!(extensions.register(registered.clone())); + assert!(!extensions.register(TestExtension::new("registered"))); + extensions.freeze(); + + assert!(!extensions.register(TestExtension::new("late"))); + let frozen_extensions = extensions.extensions(); + assert_eq!(frozen_extensions.len(), 1); + assert!(Arc::ptr_eq( + &(registered as Arc), + &frozen_extensions[0] + )); + let frozen_handlers = extensions.response_handlers(); + assert_eq!(frozen_handlers.len(), 1); + assert!(Arc::ptr_eq(&handler, &frozen_handlers[0])); + assert_eq!(extensions.len(), 1); +} + +#[test] +fn test_try_register_distinguishes_duplicate_and_frozen() { + let extensions = FrontendHeartbeatExtensions::default(); + + assert_eq!( + extensions.try_register(TestExtension::new("registered")), + Ok(()) + ); + assert_eq!( + extensions.try_register(TestExtension::new("registered")), + Err(RegistrationError::Duplicate) + ); + + extensions.freeze(); + assert_eq!( + extensions.try_register(TestExtension::new("late")), + Err(RegistrationError::Frozen) + ); +} + +#[test] +fn test_frozen_registration_does_not_call_extension_name() { + let extensions = FrontendHeartbeatExtensions::default(); + extensions.freeze(); + + let compatible_result = catch_unwind(AssertUnwindSafe(|| { + extensions.register(Arc::new(PanickingNameExtension)) + })); + assert!(!compatible_result.unwrap()); + + let typed_result = catch_unwind(AssertUnwindSafe(|| { + extensions.try_register(Arc::new(PanickingNameExtension)) + })); + + assert_eq!(typed_result.unwrap(), Err(RegistrationError::Frozen)); + assert!(extensions.is_empty()); +} + +struct BlockingNameExtension { + entered_name: Arc, + release_name: Arc, +} + +#[async_trait] +impl FrontendHeartbeatExtension for BlockingNameExtension { + fn name(&self) -> &str { + self.entered_name.wait(); + self.release_name.wait(); + "blocking" + } +} + +#[test] +fn test_freeze_wins_while_registration_name_is_running() { + let extensions = FrontendHeartbeatExtensions::default(); + let entered_name = Arc::new(Barrier::new(2)); + let release_name = Arc::new(Barrier::new(2)); + let registering = extensions.clone(); + let extension = Arc::new(BlockingNameExtension { + entered_name: entered_name.clone(), + release_name: release_name.clone(), + }); + + let registration = std::thread::spawn(move || registering.try_register(extension)); + entered_name.wait(); + let freezing = extensions.clone(); + let (freeze_done_tx, freeze_done_rx) = std::sync::mpsc::sync_channel(1); + let freeze = std::thread::spawn(move || { + freezing.freeze(); + let _ = freeze_done_tx.send(()); + }); + if let Err(error) = freeze_done_rx.recv_timeout(Duration::from_secs(2)) { + release_name.wait(); + let registration = registration.join(); + let freeze = freeze.join(); + panic!( + "freeze did not complete while extension name was blocked: {error:?}; registration: {registration:?}; freeze: {freeze:?}" + ); + } + release_name.wait(); + + freeze.join().unwrap(); + assert_eq!(registration.join().unwrap(), Err(RegistrationError::Frozen)); + assert!(extensions.is_empty()); +} + fn test_task( connector: Arc, extensions: FrontendHeartbeatExtensions, From aaa843104bb2dba1d774b53dbff3adcf1d18eba7 Mon Sep 17 00:00:00 2001 From: wy471x Date: Fri, 28 Aug 2026 02:40:25 +0000 Subject: [PATCH 11/22] fix(mysql): interpret prepared statement datetime params in session timezone (#8923) * fix(mysql): interpret prepared statement datetime params in session timezone Binary DATETIME parameters of server-side prepared statements were converted as if UTC, ignoring the session timezone set via SET time_zone. Convert them with the session timezone and add an integration test covering prepared inserts and predicates under Asia/Shanghai. Signed-off-by: wy471x * refactor: share naive datetime timezone policy via common-time Address review feedback on the prepared-statement timezone fix: - Expose Timestamp::from_naive_datetime in common-time so the DST policy (gap -> error, ambiguous -> earlier instant) lives in one place, shared by the text protocol (Timestamp::from_str) and the MySQL binary protocol. - Route the MySQL prepared-statement datetime conversion through it. - Match the target type before converting datetime params so PreparedStmtTypeMismatch fails fast without wasted conversion. - Use the short Timezone import form for consistency with the rest of servers. Signed-off-by: wy471x --------- Signed-off-by: wy471x Co-authored-by: Ning Sun --- src/common/time/src/timestamp.rs | 75 ++++++++++++-- src/servers/src/mysql/handler.rs | 12 ++- src/servers/src/mysql/helper.rs | 171 +++++++++++++++++++++++++------ tests-integration/tests/sql.rs | 72 +++++++++++++ 4 files changed, 286 insertions(+), 44 deletions(-) diff --git a/src/common/time/src/timestamp.rs b/src/common/time/src/timestamp.rs index 7daa793bafe..88cac7f4fcc 100644 --- a/src/common/time/src/timestamp.rs +++ b/src/common/time/src/timestamp.rs @@ -483,6 +483,31 @@ impl Timestamp { ParseTimestampSnafu { raw: s }.fail() } + /// Interprets a timezone-less [`NaiveDateTime`] in the given timezone and + /// returns the corresponding timestamp. + /// + /// Datetimes that fall into a DST gap (a local time that does not exist) + /// are rejected, while ambiguous datetimes (from a repeated local time) + /// are resolved to the earlier instant. This policy is shared with + /// [`Timestamp::from_str`] so that the text and binary protocols interpret + /// datetimes consistently. + pub fn from_naive_datetime( + datetime: NaiveDateTime, + timezone: &Timezone, + ) -> crate::error::Result { + match datetime_to_utc(&datetime, timezone) { + LocalResult::Single(utc) | LocalResult::Ambiguous(utc, _) => { + Timestamp::from_chrono_datetime(utc).context(ParseTimestampSnafu { + raw: format!("{datetime} (timezone {timezone})"), + }) + } + LocalResult::None => ParseTimestampSnafu { + raw: format!("{datetime} (timezone {timezone})"), + } + .fail(), + } + } + pub fn negative(mut self) -> Self { self.value = -self.value; self @@ -531,12 +556,8 @@ fn naive_datetime_to_timestamp( .context(ParseTimestampSnafu { raw: s }); }; - match datetime_to_utc(&datetime, timezone) { - LocalResult::None => ParseTimestampSnafu { raw: s }.fail(), - LocalResult::Single(utc) | LocalResult::Ambiguous(utc, _) => { - Timestamp::from_chrono_datetime(utc).context(ParseTimestampSnafu { raw: s }) - } - } + Timestamp::from_naive_datetime(datetime, timezone) + .map_err(|_| ParseTimestampSnafu { raw: s }.build()) } impl From for Timestamp { @@ -919,6 +940,48 @@ mod tests { ); } + #[test] + fn test_from_naive_datetime() { + let datetime = NaiveDate::from_ymd_opt(2026, 8, 13) + .unwrap() + .and_hms_opt(8, 0, 0) + .unwrap(); + + // A fixed-offset timezone shifts the datetime by a constant amount. + let shanghai = Timezone::from_tz_string("Asia/Shanghai").unwrap(); + assert_eq!( + "2026-08-13 00:00:00", + Timestamp::from_naive_datetime(datetime, &shanghai) + .unwrap() + .to_chrono_datetime() + .unwrap() + .to_string() + ); + + // 2026-03-08 02:30 does not exist in America/New_York (DST gap). + let new_york = Timezone::from_tz_string("America/New_York").unwrap(); + let gap = NaiveDate::from_ymd_opt(2026, 3, 8) + .unwrap() + .and_hms_opt(2, 30, 0) + .unwrap(); + assert!(Timestamp::from_naive_datetime(gap, &new_york).is_err()); + + // 2026-11-01 01:30 is ambiguous in America/New_York; picks the first + // instant (EDT, UTC-4). + let ambiguous = NaiveDate::from_ymd_opt(2026, 11, 1) + .unwrap() + .and_hms_opt(1, 30, 0) + .unwrap(); + assert_eq!( + "2026-11-01 05:30:00", + Timestamp::from_naive_datetime(ambiguous, &new_york) + .unwrap() + .to_chrono_datetime() + .unwrap() + .to_string() + ); + } + #[test] fn test_to_iso8601_string() { set_default_timezone(Some("Asia/Shanghai")).unwrap(); diff --git a/src/servers/src/mysql/handler.rs b/src/servers/src/mysql/handler.rs index b0eb33f8fd5..30579d8846a 100644 --- a/src/servers/src/mysql/handler.rs +++ b/src/servers/src/mysql/handler.rs @@ -25,6 +25,7 @@ use common_catalog::parse_optional_catalog_and_schema_from_db_string; use common_error::ext::ErrorExt; use common_query::Output; use common_telemetry::{debug, error, tracing, warn}; +use common_time::Timezone; use datafusion_common::ParamValues; use datafusion_expr::LogicalPlan; use datatypes::prelude::ConcreteDataType; @@ -289,12 +290,13 @@ impl MysqlInstanceShim { .fail(); } + let timezone = query_ctx.timezone(); let replaced_plan = match params { Params::ProtocolParams(params) => { - replace_params_with_values(&plan, param_types, ¶ms) + replace_params_with_values(&plan, param_types, ¶ms, &timezone) } Params::CliParams(params) => { - replace_params_with_exprs(&plan, param_types, ¶ms) + replace_params_with_exprs(&plan, param_types, ¶ms, &timezone) } }?; @@ -807,6 +809,7 @@ fn replace_params_with_values( plan: &LogicalPlan, param_types: HashMap>, params: &[ParamValue], + timezone: &Timezone, ) -> Result { debug_assert_eq!(param_types.len(), params.len()); @@ -824,7 +827,7 @@ fn replace_params_with_values( for (i, param) in params.iter().enumerate() { if let Some(Some(t)) = param_types.get(&format_placeholder(i + 1)) { - let value = helper::convert_value(param, t)?; + let value = helper::convert_value(param, t, timezone)?; values.push(value.into()); } @@ -839,6 +842,7 @@ fn replace_params_with_exprs( plan: &LogicalPlan, param_types: HashMap>, params: &[sql::ast::Expr], + timezone: &Timezone, ) -> Result { debug_assert_eq!(param_types.len(), params.len()); @@ -853,7 +857,7 @@ fn replace_params_with_exprs( for (i, param) in params.iter().enumerate() { if let Some(Some(t)) = param_types.get(&format_placeholder(i + 1)) { - let value = helper::convert_expr_to_scalar_value(param, t)?; + let value = helper::convert_expr_to_scalar_value(param, t, timezone)?; values.push(value.into()); } diff --git a/src/servers/src/mysql/helper.rs b/src/servers/src/mysql/helper.rs index a4f289cdfb1..f9ab5325925 100644 --- a/src/servers/src/mysql/helper.rs +++ b/src/servers/src/mysql/helper.rs @@ -15,10 +15,11 @@ use std::ops::ControlFlow; use std::time::Duration; -use chrono::NaiveDate; +use chrono::{NaiveDate, NaiveDateTime}; use common_query::prelude::ScalarValue; use common_sql::convert::sql_value_to_value; -use common_time::{Date, Timestamp}; +use common_time::timestamp::TimeUnit; +use common_time::{Date, Timestamp, Timezone}; use datatypes::prelude::{ConcreteDataType, DataType}; use datatypes::schema::ColumnSchema; use datatypes::types::TimestampType; @@ -137,11 +138,15 @@ where /// Convert [`ParamValue`] into [`Value`] according to param type. /// It will try it's best to do type conversions if possible -pub fn convert_value(param: &ParamValue, t: &ConcreteDataType) -> Result { +pub fn convert_value( + param: &ParamValue, + t: &ConcreteDataType, + timezone: &Timezone, +) -> Result { if let ConcreteDataType::Dictionary(dictionary) = t { return Ok(ScalarValue::Dictionary( Box::new(dictionary.key_type().as_arrow_type()), - Box::new(convert_value(param, dictionary.value_type())?), + Box::new(convert_value(param, dictionary.value_type(), timezone)?), )); } @@ -221,7 +226,9 @@ pub fn convert_value(param: &ParamValue, t: &ConcreteDataType) -> Result Ok(ScalarValue::Binary(Some(b.to_vec()))), - ConcreteDataType::Timestamp(ts_type) => convert_bytes_to_timestamp(b, ts_type), + ConcreteDataType::Timestamp(ts_type) => { + convert_bytes_to_timestamp(b, ts_type, timezone) + } ConcreteDataType::Date(_) => convert_bytes_to_date(b), _ => error::PreparedStmtTypeMismatchSnafu { expected: t, @@ -233,29 +240,26 @@ pub fn convert_value(param: &ParamValue, t: &ConcreteDataType) -> Result { - let timestamp_millis = to_naive_datetime(param.value) - .map_err(|e| { + ValueInner::Datetime(_) => match t { + ConcreteDataType::Timestamp(_) => { + let datetime = to_naive_datetime(param.value).map_err(|e| { error::MysqlValueConversionSnafu { err_msg: e.to_string(), } .build() - })? - .and_utc() - .timestamp_millis(); - - match t { - ConcreteDataType::Timestamp(_) => Ok(ScalarValue::TimestampMillisecond( + })?; + let timestamp_millis = datetime_to_timestamp_millis(datetime, timezone)?; + Ok(ScalarValue::TimestampMillisecond( Some(timestamp_millis), None, - )), - _ => error::PreparedStmtTypeMismatchSnafu { - expected: t, - actual: param.coltype, - } - .fail(), + )) } - } + _ => error::PreparedStmtTypeMismatchSnafu { + expected: t, + actual: param.coltype, + } + .fail(), + }, ValueInner::Time(_) => Ok(ScalarValue::Time64Nanosecond(Some( Duration::from(param.value).as_millis() as i64, ))), @@ -264,13 +268,18 @@ pub fn convert_value(param: &ParamValue, t: &ConcreteDataType) -> Result Result { +pub fn convert_expr_to_scalar_value( + param: &Expr, + t: &ConcreteDataType, + timezone: &Timezone, +) -> Result { if let ConcreteDataType::Dictionary(dictionary) = t { return Ok(ScalarValue::Dictionary( Box::new(dictionary.key_type().as_arrow_type()), Box::new(convert_expr_to_scalar_value( param, dictionary.value_type(), + timezone, )?), )); } @@ -278,7 +287,7 @@ pub fn convert_expr_to_scalar_value(param: &Expr, t: &ConcreteDataType) -> Resul let column_schema = ColumnSchema::new("", t.clone(), true); match param { Expr::Value(v) => { - let v = sql_value_to_value(&column_schema, &v.value, None, None, true); + let v = sql_value_to_value(&column_schema, &v.value, Some(timezone), None, true); match v { Ok(v) => v .try_to_scalar_value(t) @@ -290,7 +299,7 @@ pub fn convert_expr_to_scalar_value(param: &Expr, t: &ConcreteDataType) -> Resul } } Expr::UnaryOp { op, expr } if let Expr::Value(v) = &**expr => { - let v = sql_value_to_value(&column_schema, &v.value, None, Some(*op), true); + let v = sql_value_to_value(&column_schema, &v.value, Some(timezone), Some(*op), true); match v { Ok(v) => v .try_to_scalar_value(t) @@ -308,8 +317,32 @@ pub fn convert_expr_to_scalar_value(param: &Expr, t: &ConcreteDataType) -> Resul } } -fn convert_bytes_to_timestamp(bytes: &[u8], ts_type: &TimestampType) -> Result { - let ts = Timestamp::from_str_utc(&String::from_utf8_lossy(bytes)) +/// Interprets a timezone-less datetime in the given timezone and returns the +/// corresponding epoch timestamp in milliseconds. +fn datetime_to_timestamp_millis(datetime: NaiveDateTime, timezone: &Timezone) -> Result { + let ts = Timestamp::from_naive_datetime(datetime, timezone) + .map_err(|e| { + error::MysqlValueConversionSnafu { + err_msg: e.to_string(), + } + .build() + })? + .convert_to(TimeUnit::Millisecond) + .ok_or_else(|| { + error::MysqlValueConversionSnafu { + err_msg: "Overflow when converting datetime to milliseconds".to_string(), + } + .build() + })?; + Ok(ts.value()) +} + +fn convert_bytes_to_timestamp( + bytes: &[u8], + ts_type: &TimestampType, + timezone: &Timezone, +) -> Result { + let ts = Timestamp::from_str(&String::from_utf8_lossy(bytes), Some(timezone)) .map_err(|e| { error::MysqlValueConversionSnafu { err_msg: e.to_string(), @@ -446,19 +479,20 @@ mod tests { #[test] fn test_convert_expr_to_scalar_value() { + let utc = Timezone::from_tz_string("UTC").unwrap(); let expr = Expr::Value(ValueExpr::Number("123".to_string(), false).into()); let t = ConcreteDataType::int32_datatype(); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); assert_eq!(ScalarValue::Int32(Some(123)), v); let expr = Expr::Value(ValueExpr::Number("123.456789".to_string(), false).into()); let t = ConcreteDataType::float64_datatype(); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); assert_eq!(ScalarValue::Float64(Some(123.456789)), v); let expr = Expr::Value(ValueExpr::SingleQuotedString("2001-01-02".to_string()).into()); let t = ConcreteDataType::date_datatype(); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); let scalar_v = ScalarValue::Utf8(Some("2001-01-02".to_string())) .cast_to(&arrow_schema::DataType::Date32) .unwrap(); @@ -467,7 +501,7 @@ mod tests { let expr = Expr::Value(ValueExpr::SingleQuotedString("2001-01-02 03:04:05".to_string()).into()); let t = ConcreteDataType::timestamp_microsecond_datatype(); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); let scalar_v = ScalarValue::Utf8(Some("2001-01-02 03:04:05".to_string())) .cast_to(&arrow_schema::DataType::Timestamp( arrow_schema::TimeUnit::Microsecond, @@ -478,14 +512,14 @@ mod tests { let expr = Expr::Value(ValueExpr::SingleQuotedString("hello".to_string()).into()); let t = ConcreteDataType::string_datatype(); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); assert_eq!(ScalarValue::Utf8(Some("hello".to_string())), v); let t = ConcreteDataType::dictionary_datatype( ConcreteDataType::uint32_datatype(), ConcreteDataType::string_datatype(), ); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); assert_eq!( ScalarValue::Dictionary( Box::new(arrow_schema::DataType::UInt32), @@ -496,7 +530,7 @@ mod tests { let expr = Expr::Value(ValueExpr::Null.into()); let t = ConcreteDataType::time_microsecond_datatype(); - let v = convert_expr_to_scalar_value(&expr, &t).unwrap(); + let v = convert_expr_to_scalar_value(&expr, &t, &utc).unwrap(); assert_eq!(ScalarValue::Time64Microsecond(None), v); } @@ -577,12 +611,81 @@ mod tests { ), ]; + let utc = Timezone::from_tz_string("UTC").unwrap(); for (input, ts_type, expected) in test_cases { - let result = convert_bytes_to_timestamp(input.as_bytes(), &ts_type).unwrap(); + let result = convert_bytes_to_timestamp(input.as_bytes(), &ts_type, &utc).unwrap(); assert_eq!(result, expected); } } + fn utc_millis(year: i32, month: u32, day: u32, hour: u32, minute: u32, second: u32) -> i64 { + NaiveDate::from_ymd_opt(year, month, day) + .unwrap() + .and_hms_opt(hour, minute, second) + .unwrap() + .and_utc() + .timestamp_millis() + } + + #[test] + fn test_datetime_to_timestamp_millis() { + let datetime = NaiveDate::from_ymd_opt(2026, 8, 13) + .unwrap() + .and_hms_opt(8, 0, 0) + .unwrap(); + + let utc = Timezone::from_tz_string("UTC").unwrap(); + assert_eq!( + datetime_to_timestamp_millis(datetime, &utc).unwrap(), + utc_millis(2026, 8, 13, 8, 0, 0) + ); + + let shanghai = Timezone::from_tz_string("Asia/Shanghai").unwrap(); + assert_eq!( + datetime_to_timestamp_millis(datetime, &shanghai).unwrap(), + utc_millis(2026, 8, 13, 0, 0, 0) + ); + + let minus_7 = Timezone::from_tz_string("-07:00").unwrap(); + assert_eq!( + datetime_to_timestamp_millis(datetime, &minus_7).unwrap(), + utc_millis(2026, 8, 13, 15, 0, 0) + ); + + // 2026-03-08 02:30 does not exist in America/New_York (DST gap). + let new_york = Timezone::from_tz_string("America/New_York").unwrap(); + let gap = NaiveDate::from_ymd_opt(2026, 3, 8) + .unwrap() + .and_hms_opt(2, 30, 0) + .unwrap(); + assert!(datetime_to_timestamp_millis(gap, &new_york).is_err()); + + // 2026-11-01 01:30 is ambiguous in America/New_York; picks the first (EDT, UTC-4). + let ambiguous = NaiveDate::from_ymd_opt(2026, 11, 1) + .unwrap() + .and_hms_opt(1, 30, 0) + .unwrap(); + assert_eq!( + datetime_to_timestamp_millis(ambiguous, &new_york).unwrap(), + utc_millis(2026, 11, 1, 5, 30, 0) + ); + } + + #[test] + fn test_convert_bytes_to_timestamp_with_timezone() { + let shanghai = Timezone::from_tz_string("Asia/Shanghai").unwrap(); + let result = convert_bytes_to_timestamp( + "2026-08-13 08:00:00".as_bytes(), + &TimestampType::Millisecond(TimestampMillisecondType), + &shanghai, + ) + .unwrap(); + assert_eq!( + result, + ScalarValue::TimestampMillisecond(Some(utc_millis(2026, 8, 13, 0, 0, 0)), None) + ); + } + #[test] fn test_convert_bytes_to_date() { let test_cases = vec![ diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index e46a8f58d1b..32a2881dd78 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -89,6 +89,7 @@ macro_rules! sql_tests { test_postgres_explain_bind_parameter, test_postgres_array_types, test_mysql_prepare_stmt_insert_timestamp, + test_mysql_prepare_stmt_timezone, test_mysql_federated_prepare_stmt, test_declare_fetch_close_cursor, test_alter_update_on, @@ -1779,6 +1780,77 @@ pub async fn test_mysql_prepare_stmt_insert_timestamp(store_type: StorageType) { guard.remove_all().await; } +pub async fn test_mysql_prepare_stmt_timezone(store_type: StorageType) { + let (mut guard, server) = + setup_mysql_server(store_type, "test_mysql_prepare_stmt_timezone").await; + let addr = server.bind_addr().unwrap().to_string(); + + let mut conn = MySqlConnection::connect(&format!("mysql://{addr}/public")) + .await + .unwrap(); + + conn.execute("create table demo(i bigint, ts timestamp time index)") + .await + .unwrap(); + conn.execute("SET time_zone = 'Asia/Shanghai'") + .await + .unwrap(); + + // Server-side prepared statement: the binary DATETIME parameter must be + // interpreted in the session timezone. + sqlx::query("insert into demo values(?, ?)") + .bind(1) + .bind( + NaiveDate::from_ymd_opt(2026, 8, 13) + .and_then(|x| x.and_hms_opt(8, 0, 0)) + .unwrap(), + ) + .execute(&mut conn) + .await + .unwrap(); + + // Text protocol with an equivalent timezone-less literal one hour later. + sqlx::query("insert into demo values(2, '2026-08-13 09:00:00')") + .execute(&mut conn) + .await + .unwrap(); + + // Timestamps are read back in the session timezone, and sqlx parses the + // timezone-less wire representation as UTC. + let rows = sqlx::query("select i, ts from demo order by i") + .fetch_all(&mut conn) + .await + .unwrap(); + assert_eq!(rows.len(), 2); + let ts: DateTime = rows[0].get("ts"); + assert_eq!(ts.to_string(), "2026-08-13 08:00:00 UTC"); + let ts: DateTime = rows[1].get("ts"); + assert_eq!(ts.to_string(), "2026-08-13 09:00:00 UTC"); + + // The prepared predicate compares the same instant as the text literal. + let rows = sqlx::query("select i from demo where ts = ? order by i") + .bind( + NaiveDate::from_ymd_opt(2026, 8, 13) + .and_then(|x| x.and_hms_opt(8, 0, 0)) + .unwrap(), + ) + .fetch_all(&mut conn) + .await + .unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].get::("i"), 1); + + let rows = sqlx::query("select i from demo where ts = '2026-08-13 08:00:00'") + .fetch_all(&mut conn) + .await + .unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].get::("i"), 1); + + let _ = server.shutdown().await; + guard.remove_all().await; +} + pub async fn test_mysql_federated_prepare_stmt(store_type: StorageType) { common_telemetry::init_default_ut_logging(); From bd7d2c1dfa166f579bc33c496d50d41eb29b775f Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 28 Aug 2026 03:45:41 +0000 Subject: [PATCH 12/22] fix(mito2): use target sequence for foreign SSTs (#8946) * fix(mito2): use target sequence for foreign SSTs Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * chore(mito2): address foreign SST review feedback Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/engine/apply_staging_manifest_test.rs | 215 ++++++++++++++++- src/mito2/src/sst/index.rs | 86 ++++++- src/mito2/src/sst/parquet.rs | 227 ++++++++++++------ src/mito2/src/sst/parquet/reader.rs | 21 +- src/mito2/src/worker/handle_rebuild_index.rs | 1 + 5 files changed, 466 insertions(+), 84 deletions(-) diff --git a/src/mito2/src/engine/apply_staging_manifest_test.rs b/src/mito2/src/engine/apply_staging_manifest_test.rs index e10ed0cbf00..49cc97c7d51 100644 --- a/src/mito2/src/engine/apply_staging_manifest_test.rs +++ b/src/mito2/src/engine/apply_staging_manifest_test.rs @@ -12,20 +12,24 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashSet; use std::sync::Arc; use std::{assert_matches, fs}; use api::v1::Rows; +use api::v1::region::{StrictWindow, compact_request}; use common_function::utils::partition_expr_version; use common_recordbatch::RecordBatches; +use datatypes::arrow::array::AsArray; +use datatypes::arrow::datatypes::Float64Type; use datatypes::value::Value; use partition::expr::{PartitionExpr, col}; use store_api::region_engine::{ RegionEngine, RegionRole, RemapManifestsRequest, SettableRegionRoleState, }; use store_api::region_request::{ - ApplyStagingManifestRequest, EnterStagingRequest, RegionFlushRequest, RegionPutRequest, - RegionRequest, StagingPartitionDirective, + ApplyStagingManifestRequest, EnterStagingRequest, RegionCompactRequest, RegionFlushRequest, + RegionPutRequest, RegionRequest, StagingPartitionDirective, }; use store_api::storage::{FileId, RegionId}; @@ -37,7 +41,9 @@ use crate::manifest::action::{ }; use crate::sst::FormatType; use crate::sst::file::FileMeta; -use crate::test_util::{CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema}; +use crate::test_util::{ + CreateRequestBuilder, TestEnv, build_rows, build_rows_for_key, put_rows, rows_schema, +}; fn range_expr(col_name: &str, start: i64, end: i64) -> PartitionExpr { col(col_name) @@ -45,6 +51,209 @@ fn range_expr(col_name: &str, start: i64, end: i64) -> PartitionExpr { .and(col(col_name).lt(Value::Int64(end))) } +#[tokio::test] +async fn test_apply_staging_manifest_sequence_domain() { + common_telemetry::init_default_ut_logging(); + test_apply_staging_manifest_sequence_domain_with_format(false).await; + test_apply_staging_manifest_sequence_domain_with_format(true).await; +} + +async fn test_apply_staging_manifest_sequence_domain_with_format(flat_format: bool) { + let mut env = TestEnv::with_prefix("apply-staging-sequence-domain").await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }) + .await; + let source = RegionId::new(1, 1); + let target = RegionId::new(1, 2); + let request = CreateRequestBuilder::new().build(); + let schema = rows_schema(&request); + + engine + .handle_request(source, RegionRequest::Create(request.clone())) + .await + .unwrap(); + for value in 0..3 { + put_rows( + &engine, + source, + Rows { + schema: schema.clone(), + rows: build_rows_for_key("0", 0, 1, value), + }, + ) + .await; + } + engine + .handle_request(source, RegionRequest::Flush(RegionFlushRequest::default())) + .await + .unwrap(); + let source_manifest = engine + .get_region(source) + .unwrap() + .manifest_ctx + .manifest() + .await; + assert_eq!(source_manifest.files.len(), 1); + assert_eq!( + source_manifest.files.values().next().unwrap().sequence, + Some(std::num::NonZeroU64::new(3).unwrap()) + ); + + engine + .set_region_role_state_gracefully(source, SettableRegionRoleState::StagingLeader) + .await + .unwrap(); + let partition_expr = float_range_expr("field_0", 0.1, 100.1) + .as_json_str() + .unwrap(); + let result = engine + .remap_manifests(RemapManifestsRequest { + region_id: source, + input_regions: vec![source], + region_mapping: [(source, vec![target])].into_iter().collect(), + new_partition_exprs: [(target, partition_expr.clone())].into_iter().collect(), + }) + .await + .unwrap(); + engine + .handle_request(target, RegionRequest::Create(request.clone())) + .await + .unwrap(); + engine + .handle_request( + target, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive: StagingPartitionDirective::UpdatePartitionExpr( + partition_expr.clone(), + ), + }), + ) + .await + .unwrap(); + engine + .handle_request( + target, + RegionRequest::ApplyStagingManifest(ApplyStagingManifestRequest { + partition_expr, + central_region_id: source, + manifest_path: result.manifest_paths[&target].clone(), + }), + ) + .await + .unwrap(); + + let manifest = engine + .get_region(target) + .unwrap() + .manifest_ctx + .manifest() + .await; + assert_eq!(manifest.files.len(), 1); + assert_eq!(manifest.committed_sequence, Some(1)); + let imported_file = manifest.files.values().next().unwrap(); + assert_eq!(imported_file.region_id, source); + assert_eq!( + imported_file.sequence, + Some(std::num::NonZeroU64::new(1).unwrap()) + ); + + put_rows( + &engine, + target, + Rows { + schema: schema.clone(), + rows: build_rows_for_key("0", 0, 1, 99), + }, + ) + .await; + assert_target_value(&engine, target, 99.0).await; + + engine + .handle_request(target, RegionRequest::Flush(RegionFlushRequest::default())) + .await + .unwrap(); + let input_ids = current_file_ids(&engine, target); + assert_eq!( + input_ids.len(), + 2, + "imported and target-write SSTs must both exist" + ); + + engine + .handle_request( + target, + RegionRequest::Compact(RegionCompactRequest { + options: compact_request::Options::StrictWindow(StrictWindow { + window_seconds: 60, + }), + parallelism: None, + time_range: None, + }), + ) + .await + .unwrap(); + let output_version = engine.get_region(target).unwrap().version(); + let output_files = output_version + .ssts + .levels() + .iter() + .flat_map(|level| level.files.values()) + .collect::>(); + let output_ids = output_files + .iter() + .map(|file| file.meta_ref().file_id) + .collect::>(); + assert_eq!(output_ids.len(), 1); + assert!( + output_files + .iter() + .all(|file| file.meta_ref().region_id == target) + ); + assert!( + input_ids.iter().all(|id| !output_ids.contains(id)), + "real compaction must replace both input SSTs" + ); + assert_target_value(&engine, target, 99.0).await; +} + +fn current_file_ids(engine: &crate::engine::MitoEngine, region_id: RegionId) -> HashSet { + engine + .get_region(region_id) + .unwrap() + .version() + .ssts + .levels() + .iter() + .flat_map(|level| level.files.values()) + .map(|file| file.meta_ref().file_id) + .collect() +} + +async fn assert_target_value( + engine: &crate::engine::MitoEngine, + region_id: RegionId, + expected: f64, +) { + let scan = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(scan).await.unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 1 + ); + let batch = batches.iter().next().unwrap(); + let values = batch + .column_by_name("field_0") + .unwrap() + .as_primitive::(); + assert_eq!(values.value(0), expected); +} + fn float_range_expr(col_name: &str, start: f64, end: f64) -> PartitionExpr { col(col_name) .gt_eq(Value::Float64(start.into())) diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 581d1edebdf..7954af2b920 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -706,6 +706,14 @@ pub struct IndexBuildTask { pub region_id: RegionId, /// The SST file handle to build index for. pub file: FileHandle, + /// The target region metadata used to decode rows from the SST. + /// + /// An SST may originate in another region while being visible in the target + /// manifest. This metadata defines the target schema and sequence domain; + /// applying the staging manifest only makes imported files visible. Index + /// rebuild happens later when a flush, compaction, schema change, or manual + /// index build request schedules it. + pub(crate) target_region_metadata: RegionMetadataRef, /// The manifest state this build is based on. pub(crate) source: IndexBuildSource, pub reason: IndexBuildType, @@ -846,6 +854,7 @@ impl IndexBuildTask { let mut parquet_reader = self .access_layer .read_sst(self.file.clone()) // use the latest file handle instead of creating a new one + .expected_metadata(Some(self.target_region_metadata.clone())) .build() .await?; @@ -1494,8 +1503,12 @@ mod tests { use datatypes::schema::{ ColumnSchema, FulltextOptions, SkippingIndexOptions, SkippingIndexType, }; + use datatypes::value::Value; + use index::inverted_index::format::reader::InvertedIndexReader; use object_store::ObjectStore; use object_store::services::Memory; + use partition::expr::col; + use puffin::puffin_manager::{PuffinManager, PuffinReader}; use puffin_manager::PuffinManagerFactory; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use tokio::sync::mpsc; @@ -2088,6 +2101,7 @@ mod tests { let task = IndexBuildTask { region_id, file, + target_region_metadata: version_control.current().version.metadata.clone(), source: IndexBuildSource::new( file_meta, version_control.current().version.metadata.schema_version, @@ -2120,16 +2134,34 @@ mod tests { } #[tokio::test] - async fn test_index_build_task_increments_legacy_index_version() { + async fn test_index_build_task_foreign_file_uses_target_metadata() { let env = SchedulerEnv::new().await; let mut scheduler = env.mock_index_build_scheduler(4); - let metadata = Arc::new(sst_region_metadata()); - let manifest_ctx = env.mock_manifest_context(metadata.clone()).await; - let region_id = metadata.region_id; + let source_metadata = Arc::new(sst_region_metadata()); + let mut target_metadata = (*source_metadata).clone(); + target_metadata.region_id = RegionId::new(1, 3); + let mut target_builder = RegionMetadataBuilder::new(target_metadata.region_id); + for mut column_metadata in target_metadata.column_metadatas.clone() { + if column_metadata.column_id == 2 { + column_metadata.column_schema = + column_metadata.column_schema.with_inverted_index(true); + } + target_builder.push_column_metadata(column_metadata); + } + let partition_expr = col("field_0") + .gt_eq(Value::UInt64(100)) + .and(col("field_0").lt(Value::UInt64(200))); + target_builder + .primary_key(target_metadata.primary_key.clone()) + .partition_expr_json(Some(partition_expr.as_json_str().unwrap())) + .bump_version(); + let target_metadata = Arc::new(target_builder.build().unwrap()); + let manifest_ctx = env.mock_manifest_context(target_metadata.clone()).await; + let region_id = target_metadata.region_id; let file_purger = Arc::new(NoopFilePurger {}); - let sst_info = mock_sst_file(metadata.clone(), &env, IndexBuildMode::Async).await; + let sst_info = mock_sst_file(source_metadata.clone(), &env, IndexBuildMode::Async).await; let file_meta = FileMeta { - region_id, + region_id: source_metadata.region_id, file_id: sst_info.file_id, file_size: sst_info.file_size, max_row_group_uncompressed_size: sst_info.max_row_group_uncompressed_size, @@ -2144,8 +2176,8 @@ mod tests { seed_manifest_file(&manifest_ctx, &file_meta).await; let files = HashMap::from([(file_meta.file_id, file_meta.clone())]); let version_control = - mock_version_control(metadata.clone(), file_purger.clone(), files).await; - let indexer_builder = mock_indexer_builder(metadata.clone(), &env).await; + mock_version_control(target_metadata.clone(), file_purger.clone(), files).await; + let indexer_builder = mock_indexer_builder(target_metadata.clone(), &env).await; let file = FileHandle::new(file_meta.clone(), file_purger.clone()); @@ -2155,6 +2187,7 @@ mod tests { let task = IndexBuildTask { region_id, file, + target_region_metadata: version_control.current().version.metadata.clone(), source: IndexBuildSource::new( file_meta.clone(), version_control.current().version.metadata.schema_version, @@ -2199,9 +2232,40 @@ mod tests { assert!(updated_meta.index_file_size > 0); assert_eq!(updated_meta.file_id, file_meta.file_id); assert_eq!(updated_meta.index_version, 1); + let field_0_index = updated_meta + .indexes + .iter() + .find(|index| index.column_id == 2) + .expect("field_0 should have an inverted index"); + assert_eq!( + field_0_index.created_indexes.as_slice(), + [IndexType::InvertedIndex] + ); } _ => panic!("Unexpected worker request: {:?}", worker_req), } + + let puffin_reader = env + .access_layer + .build_puffin_manager() + .reader(&RegionIndexId::new( + RegionFileId::new(source_metadata.region_id, file_meta.file_id), + 1, + )) + .await + .unwrap(); + let blob = puffin_reader + .blob(inverted_index::INDEX_BLOB_TYPE) + .await + .unwrap(); + let blob_reader = blob.reader().await.unwrap(); + let index_metadata = + index::inverted_index::format::reader::InvertedIndexBlobReader::new(blob_reader) + .metadata(None) + .await + .unwrap(); + assert!(index_metadata.metas.contains_key("2")); + assert_eq!(index_metadata.total_row_count, 100); } async fn schedule_index_build_task_with_mode(build_mode: IndexBuildMode) { @@ -2236,6 +2300,7 @@ mod tests { let task = IndexBuildTask { region_id, file, + target_region_metadata: version_control.current().version.metadata.clone(), source: IndexBuildSource::new( file_meta.clone(), version_control.current().version.metadata.schema_version, @@ -2346,6 +2411,7 @@ mod tests { let task = IndexBuildTask { region_id, file, + target_region_metadata: version_control.current().version.metadata.clone(), source: IndexBuildSource::new( file_meta.clone(), version_control.current().version.metadata.schema_version, @@ -2444,6 +2510,7 @@ mod tests { let task = IndexBuildTask { region_id, file, + target_region_metadata: version_control.current().version.metadata.clone(), source: IndexBuildSource::new( file_meta.clone(), version_control.current().version.metadata.schema_version, @@ -2505,7 +2572,7 @@ mod tests { let schema_version = metadata.schema_version; let manifest_ctx = env.mock_manifest_context(metadata.clone()).await; let file_purger = Arc::new(NoopFilePurger {}); - let indexer_builder = mock_indexer_builder(metadata, env).await; + let indexer_builder = mock_indexer_builder(metadata.clone(), env).await; let (tx, _rx) = mpsc::channel(4); let (result_tx, result_rx) = mpsc::channel::>(4); @@ -2521,6 +2588,7 @@ mod tests { let task = IndexBuildTask { region_id, file, + target_region_metadata: metadata, source: IndexBuildSource::new(file_meta, schema_version), reason, access_layer: env.access_layer.clone(), diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index e67dcf2bee5..8a05c099e98 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -118,8 +118,8 @@ mod tests { use datafusion_expr::{BinaryExpr, Expr, Literal, Operator, col, lit}; use datatypes::arrow; use datatypes::arrow::array::{ - ArrayRef, BinaryDictionaryBuilder, RecordBatch, StringArray, StringDictionaryBuilder, - TimestampMillisecondArray, UInt8Array, UInt64Array, + ArrayRef, AsArray, BinaryDictionaryBuilder, RecordBatch, StringArray, + StringDictionaryBuilder, TimestampMillisecondArray, UInt8Array, UInt64Array, }; use datatypes::arrow::datatypes::{DataType, Field, Schema, UInt32Type}; use datatypes::arrow::util::pretty::pretty_format_batches; @@ -1625,91 +1625,178 @@ mod tests { #[tokio::test] async fn test_read_with_override_sequence() { + test_read_with_override_sequence_with_format(false).await; + test_read_with_override_sequence_with_format(true).await; + } + + async fn test_read_with_override_sequence_with_format(flat_format: bool) { let mut env = TestEnv::new().await; let object_store = env.init_object_store_manager(); - let handle = sst_file_handle(0, 1000); - let file_path = FixedPathProvider { - region_file_id: handle.file_id(), - }; let metadata = Arc::new(sst_region_metadata()); - // Create batches with sequence 0 to trigger override functionality. - let source = new_flat_source_from_record_batches(vec![ - new_record_batch_with_custom_sequence(&["a", "d"], 0, 60, 0), - new_record_batch_with_custom_sequence(&["b", "f"], 0, 40, 0), - ]); + async fn read_sequences(builder: ParquetReaderBuilder) -> Vec { + let mut reader = builder.build().await.unwrap().unwrap(); + let mut sequences = Vec::new(); + while let Some(batch) = reader.next_record_batch().await.unwrap() { + let sequence = batch + .column(batch.num_columns() - 2) + .as_primitive::(); + sequences.extend((0..sequence.len()).map(|idx| sequence.value(idx))); + } + sequences + } - let write_opts = WriteOptions { - row_group_size: 50, - ..Default::default() - }; + async fn write_sst( + object_store: ObjectStore, + metadata: Arc, + handle: FileHandle, + flat_format: bool, + sequence: u64, + ) { + let file_path = FixedPathProvider { + region_file_id: handle.file_id(), + }; + let source = new_flat_source_from_record_batches(vec![ + new_record_batch_with_custom_sequence(&["a", "d"], 0, 60, sequence), + new_record_batch_with_custom_sequence(&["b", "f"], 0, 40, sequence), + ]); + let write_opts = WriteOptions { + row_group_size: 50, + ..Default::default() + }; + let mut metrics = Metrics::new(WriteType::Flush); + let mut writer = ParquetWriter::new_with_object_store( + object_store, + metadata, + IndexConfig::default(), + NoopIndexBuilder, + file_path, + &mut metrics, + ) + .await; + if flat_format { + writer + .write_all_flat(source, None, &write_opts) + .await + .unwrap(); + } else { + writer + .write_all_flat_as_primary_key(source, None, &write_opts) + .await + .unwrap(); + } + } - let mut metrics = Metrics::new(WriteType::Flush); - let mut writer = ParquetWriter::new_with_object_store( + let custom_sequence = 12345u64; + let local_zero_handle = sst_file_handle(0, 1000); + write_sst( object_store.clone(), metadata.clone(), - IndexConfig::default(), - NoopIndexBuilder, - file_path, - &mut metrics, + local_zero_handle.clone(), + flat_format, + 0, ) .await; - writer - .write_all_flat_as_primary_key(source, None, &write_opts) - .await - .unwrap() - .remove(0); + // Local all-zero SSTs retain the compatibility override. + let local_zero_none = read_sequences( + ParquetReaderBuilder::new( + FILE_DIR.to_string(), + PathType::Bare, + local_zero_handle.clone(), + object_store.clone(), + ) + .expected_metadata(Some(metadata.clone())), + ) + .await; + assert!(local_zero_none.iter().all(|sequence| *sequence == 0)); - // Read without override sequence (should read sequence 0) - let builder = ParquetReaderBuilder::new( - FILE_DIR.to_string(), - PathType::Bare, - handle.clone(), - object_store.clone(), - ); - let mut reader = builder.build().await.unwrap().unwrap(); - let mut normal_batches = Vec::new(); - while let Some(batch) = reader.next_record_batch().await.unwrap() { - normal_batches.push(batch); - } - - // Read with override sequence using FileMeta.sequence - let custom_sequence = 12345u64; - let file_meta = handle.meta_ref(); - let mut override_file_meta = file_meta.clone(); - override_file_meta.sequence = Some(std::num::NonZero::new(custom_sequence).unwrap()); - let override_handle = FileHandle::new( - override_file_meta, + let mut local_zero_meta = local_zero_handle.meta_ref().clone(); + local_zero_meta.sequence = Some(std::num::NonZeroU64::new(custom_sequence).unwrap()); + let local_zero_override_handle = FileHandle::new( + local_zero_meta, Arc::new(crate::sst::file_purger::NoopFilePurger), ); - - let builder = ParquetReaderBuilder::new( - FILE_DIR.to_string(), - PathType::Bare, - override_handle, - object_store.clone(), + let local_zero_override = read_sequences( + ParquetReaderBuilder::new( + FILE_DIR.to_string(), + PathType::Bare, + local_zero_override_handle, + object_store.clone(), + ) + .expected_metadata(Some(metadata.clone())), + ) + .await; + assert!( + local_zero_override + .iter() + .all(|sequence| *sequence == custom_sequence) ); - let mut reader = builder.build().await.unwrap().unwrap(); - let mut override_batches = Vec::new(); - while let Some(batch) = reader.next_record_batch().await.unwrap() { - override_batches.push(batch); - } - // Compare the results - assert_eq!(normal_batches.len(), override_batches.len()); - for (normal, override_batch) in normal_batches.into_iter().zip(override_batches.iter()) { - let expected_batch = { - let mut columns = normal.columns().to_vec(); - let num_cols = columns.len(); - columns[num_cols - 2] = - Arc::new(UInt64Array::from_value(custom_sequence, normal.num_rows())); - RecordBatch::try_new(normal.schema(), columns).unwrap() - }; + let local_nonzero_handle = sst_file_handle(0, 1000); + write_sst( + object_store.clone(), + metadata.clone(), + local_nonzero_handle.clone(), + flat_format, + 7, + ) + .await; - // Override batch should match expected batch - assert_eq!(*override_batch, expected_batch); - } + // Local nonzero SSTs retain physical per-row sequences, even with FileMeta.sequence. + let mut local_nonzero_meta = local_nonzero_handle.meta_ref().clone(); + local_nonzero_meta.sequence = Some(std::num::NonZeroU64::new(custom_sequence).unwrap()); + let local_nonzero_override_handle = FileHandle::new( + local_nonzero_meta, + Arc::new(crate::sst::file_purger::NoopFilePurger), + ); + let local_nonzero_override = read_sequences( + ParquetReaderBuilder::new( + FILE_DIR.to_string(), + PathType::Bare, + local_nonzero_override_handle, + object_store.clone(), + ) + .expected_metadata(Some(metadata.clone())), + ) + .await; + assert!(local_nonzero_override.iter().all(|sequence| *sequence == 7)); + + // None never overrides a local nonzero physical sequence. + let local_nonzero_none = read_sequences( + ParquetReaderBuilder::new( + FILE_DIR.to_string(), + PathType::Bare, + local_nonzero_handle.clone(), + object_store.clone(), + ) + .expected_metadata(Some(metadata.clone())), + ) + .await; + assert!(local_nonzero_none.iter().all(|sequence| *sequence == 7)); + + // A source-owned handle is foreign when read against target metadata, so the + // target-local manifest barrier is applied even for nonzero physical sequences. + let mut target_metadata = (*metadata).clone(); + target_metadata.region_id = RegionId::new(0, 1); + let target_metadata = Arc::new(target_metadata); + let mut foreign_meta = local_nonzero_handle.meta_ref().clone(); + foreign_meta.sequence = Some(std::num::NonZeroU64::new(custom_sequence).unwrap()); + let foreign_handle = FileHandle::new( + foreign_meta, + Arc::new(crate::sst::file_purger::NoopFilePurger), + ); + let foreign = read_sequences( + ParquetReaderBuilder::new( + FILE_DIR.to_string(), + PathType::Bare, + foreign_handle, + object_store, + ) + .expected_metadata(Some(target_metadata)), + ) + .await; + assert!(foreign.iter().all(|sequence| *sequence == custom_sequence)); } #[tokio::test] diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 4f890cb83e8..3adb335564b 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -22,7 +22,7 @@ use std::time::{Duration, Instant}; use api::v1::SemanticType; use common_recordbatch::filter::SimpleFilterEvaluator; -use common_telemetry::{error, tracing, warn}; +use common_telemetry::{debug, error, tracing, warn}; use datafusion::physical_plan::PhysicalExpr; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_expr::utils::expr_to_columns; @@ -471,7 +471,24 @@ impl ParquetReaderBuilder { &file_path, skip_auto_convert, )?; - if need_override_sequence(&parquet_meta) { + // `region_meta` comes from the Parquet/source file and must not be used as the + // target identity. When the caller has no current metadata, the handle is the + // only local identity available and therefore denotes a local read. + let expected_region_id = self + .expected_metadata + .as_ref() + .map(|metadata| metadata.region_id) + .unwrap_or(self.file_handle.region_id()); + let is_foreign = self.file_handle.region_id() != expected_region_id; + if is_foreign { + debug!( + "Reading foreign SST, file_id: {}, source_region_id: {}, expected_region_id: {}", + self.file_handle.file_id().file_id(), + self.file_handle.region_id(), + expected_region_id, + ); + } + if is_foreign || need_override_sequence(&parquet_meta) { read_format .set_override_sequence(self.file_handle.meta_ref().sequence.map(|x| x.get())); } diff --git a/src/mito2/src/worker/handle_rebuild_index.rs b/src/mito2/src/worker/handle_rebuild_index.rs index ad0a09e1d32..3eb45b1ab21 100644 --- a/src/mito2/src/worker/handle_rebuild_index.rs +++ b/src/mito2/src/worker/handle_rebuild_index.rs @@ -79,6 +79,7 @@ impl RegionWorkerLoop { IndexBuildTask { region_id: region.region_id, file: file.clone(), + target_region_metadata: version.metadata.clone(), source: IndexBuildSource::new(file_meta, version.metadata.schema_version), reason: build_type, access_layer: access_layer.clone(), From ba3c5a939e3319b1973bfc57db77c011f1b511f0 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Fri, 28 Aug 2026 05:56:44 +0000 Subject: [PATCH 13/22] chore(mito2): reduce default auto flush interval (#8971) Signed-off-by: evenyag --- config/config.md | 4 ++-- config/datanode.example.toml | 2 +- config/standalone.example.toml | 2 +- src/cmd/tests/load_config_test.rs | 4 ++-- src/mito2/src/config.rs | 12 ++++++++++-- tests-integration/tests/http.rs | 2 +- 6 files changed, 17 insertions(+), 9 deletions(-) diff --git a/config/config.md b/config/config.md index 992296d4827..1343f60dbdd 100644 --- a/config/config.md +++ b/config/config.md @@ -161,7 +161,7 @@ | `region_engine.mito.max_background_purges` | Integer | Auto | Max number of running background purge jobs (default: number of cpu cores). | | `region_engine.mito.experimental_compaction_memory_limit` | String | 0 | Memory budget for compaction tasks.
Supports absolute size (e.g., "2GiB", "512MB") or percentage of system memory (e.g., "50%").
Setting it to 0 or "unlimited" disables the limit. | | `region_engine.mito.experimental_compaction_on_exhausted` | String | wait | Behavior when compaction cannot acquire memory from the budget.
Options: "wait" (default, 10s), "wait()", "fail" | -| `region_engine.mito.auto_flush_interval` | String | `1h` | Interval to auto flush a region if it has not flushed yet. | +| `region_engine.mito.auto_flush_interval` | String | `10m` | Interval to auto flush a region if it has not flushed yet. | | `region_engine.mito.global_write_buffer_size` | String | Auto | Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. | | `region_engine.mito.global_write_buffer_reject_size` | String | Auto | Global write buffer size threshold to reject write requests. If not set, it's default to 2 times of `global_write_buffer_size`. | | `region_engine.mito.default_region_write_buffer_size` | String | `0` | Default write buffer size for each region. Regions stall at this size and reject writes at twice this size. Setting it to 0 disables both limits unless the table specifies `write_buffer_size`. | @@ -585,7 +585,7 @@ | `region_engine.mito.max_background_purges` | Integer | Auto | Max number of running background purge jobs (default: number of cpu cores). | | `region_engine.mito.experimental_compaction_memory_limit` | String | 0 | Memory budget for compaction tasks.
Supports absolute size (e.g., "2GiB", "512MB") or percentage of system memory (e.g., "50%").
Setting it to 0 or "unlimited" disables the limit. | | `region_engine.mito.experimental_compaction_on_exhausted` | String | wait | Behavior when compaction cannot acquire memory from the budget.
Options: "wait" (default, 10s), "wait()", "fail" | -| `region_engine.mito.auto_flush_interval` | String | `1h` | Interval to auto flush a region if it has not flushed yet. | +| `region_engine.mito.auto_flush_interval` | String | `10m` | Interval to auto flush a region if it has not flushed yet. | | `region_engine.mito.global_write_buffer_size` | String | Auto | Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. | | `region_engine.mito.global_write_buffer_reject_size` | String | Auto | Global write buffer size threshold to reject write requests. If not set, it's default to 2 times of `global_write_buffer_size` | | `region_engine.mito.default_region_write_buffer_size` | String | `0` | Default write buffer size for each region. Regions stall at this size and reject writes at twice this size. Setting it to 0 disables both limits unless the table specifies `write_buffer_size`. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 5966405452a..20c84fd9ee4 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -462,7 +462,7 @@ compress_manifest = false #+ experimental_compaction_on_exhausted = "wait" ## Interval to auto flush a region if it has not flushed yet. -auto_flush_interval = "1h" +auto_flush_interval = "10m" ## Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. ## @toml2docs:none-default="Auto" diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 70f74c7f946..6ca13566d7b 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -625,7 +625,7 @@ compress_manifest = false #+ experimental_compaction_on_exhausted = "wait" ## Interval to auto flush a region if it has not flushed yet. -auto_flush_interval = "1h" +auto_flush_interval = "10m" ## Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. ## @toml2docs:none-default="Auto" diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index 70594d0653f..8af216256f9 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -113,7 +113,7 @@ fn test_load_datanode_example_config() { }, region_engine: vec![ RegionEngineConfig::Mito(MitoConfig { - auto_flush_interval: Duration::from_secs(3600), + auto_flush_interval: Duration::from_secs(10 * 60), default_region_write_buffer_size: ReadableSize::mb(0), write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)), scan_memory_limit: MemoryLimit::Unlimited, @@ -354,7 +354,7 @@ fn test_load_standalone_example_config() { }), region_engine: vec![ RegionEngineConfig::Mito(MitoConfig { - auto_flush_interval: Duration::from_secs(3600), + auto_flush_interval: Duration::from_secs(10 * 60), default_region_write_buffer_size: ReadableSize::mb(0), write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)), scan_memory_limit: MemoryLimit::Unlimited, diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index 54445549a05..be74fa67c60 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -102,7 +102,7 @@ pub struct MitoConfig { pub experimental_compaction_on_exhausted: OnExhaustedPolicy, // Flush configs: - /// Interval to auto flush a region if it has not flushed yet (default 30 min). + /// Interval to auto flush a region if it has not flushed yet (default 10 min). #[serde(with = "humantime_serde")] pub auto_flush_interval: Duration, /// Global write buffer size threshold to trigger flush. @@ -210,7 +210,7 @@ impl Default for MitoConfig { max_background_purges: get_total_cpu_cores(), experimental_compaction_memory_limit: MemoryLimit::Unlimited, experimental_compaction_on_exhausted: OnExhaustedPolicy::default(), - auto_flush_interval: Duration::from_secs(30 * 60), + auto_flush_interval: Duration::from_secs(10 * 60), global_write_buffer_size: ReadableSize::gb(1), global_write_buffer_reject_size: ReadableSize::gb(2), default_region_write_buffer_size: ReadableSize::mb(0), @@ -381,6 +381,14 @@ impl MitoConfig { mod tests { use super::*; + #[test] + fn test_default_auto_flush_interval() { + assert_eq!( + Duration::from_secs(10 * 60), + MitoConfig::default().auto_flush_interval + ); + } + #[test] fn test_adjust_sst_metadata_and_prefilter_cache_caps_independently() { let mut config = MitoConfig::default(); diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 43ba8701d76..f4cfda24dc5 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2278,7 +2278,7 @@ experimental_manifest_keep_removed_file_ttl = "1h" compress_manifest = false experimental_compaction_memory_limit = "unlimited" experimental_compaction_on_exhausted = "wait" -auto_flush_interval = "30m" +auto_flush_interval = "10m" default_region_write_buffer_size = "0KiB" enable_write_cache = false write_cache_path = "" From 919846286964c34356487cb4cc437edc166c35c8 Mon Sep 17 00:00:00 2001 From: Dhruv Vaishnav Date: Fri, 28 Aug 2026 06:17:55 +0000 Subject: [PATCH 14/22] feat(meta): record physical table reconciliation events (#8935) * feat(meta): record physical table reconciliation events Signed-off-by: dhruvxvaishnav * docs(config): add reconciliation table event Signed-off-by: dhruvxvaishnav * fix(meta): address reconciliation event review feedback Signed-off-by: dhruvxvaishnav * fix(meta): keep reconciliation event summary volatile Signed-off-by: dhruvxvaishnav * refactor(meta): remove unused table state downcasting Signed-off-by: dhruvxvaishnav --------- Signed-off-by: dhruvxvaishnav --- config/config.md | 4 +- config/metasrv.example.toml | 2 +- config/standalone.example.toml | 2 +- src/common/meta/src/reconciliation.rs | 1 + src/common/meta/src/reconciliation/event.rs | 351 +++++++++++ .../src/reconciliation/reconcile_table.rs | 544 +++++++++++++++++- .../reconcile_table/reconcile_regions.rs | 25 +- .../reconcile_table/reconciliation_end.rs | 6 - .../reconcile_table/reconciliation_start.rs | 25 +- .../resolve_column_metadata.rs | 34 +- .../reconcile_table/update_table_info.rs | 13 +- tests-integration/tests/main.rs | 1 + .../tests/reconciliation_event.rs | 144 +++++ 13 files changed, 1102 insertions(+), 50 deletions(-) create mode 100644 src/common/meta/src/reconciliation/event.rs create mode 100644 tests-integration/tests/reconciliation_event.rs diff --git a/config/config.md b/config/config.md index 1343f60dbdd..e8a89e94e1f 100644 --- a/config/config.md +++ b/config/config.md @@ -237,7 +237,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`, `admin_function`.
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`, `reconcile_table`.
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. | @@ -456,7 +456,7 @@ | `wal.create_topic_timeout` | String | `30s` | The timeout for creating a Kafka topic.
**It's only used when the provider is `kafka`**. | | `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: `region_migration`,
`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`, `repartition`,
`repartition_group`, `wal_prune`, `batch_gc`.
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: `region_migration`,
`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`, `repartition`,
`repartition_group`, `wal_prune`, `batch_gc`, `reconcile_table`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | | `stats_persistence` | -- | -- | Configuration options for the stats persistence. | | `stats_persistence.ttl` | String | `0s` | TTL for the stats table that will be used to store the stats.
Set to `0s` to disable stats persistence.
Default is `0s`.
If you want to enable stats persistence, set the TTL to a value greater than 0.
It is recommended to set a small value, e.g., `3h`. | | `stats_persistence.interval` | String | `10m` | The interval to persist the stats. Default is `10m`.
The minimum value is `10m`, if the value is less than `10m`, it will be overridden to `10m`. | diff --git a/config/metasrv.example.toml b/config/metasrv.example.toml index 6017bfcc49a..ae373d07482 100644 --- a/config/metasrv.example.toml +++ b/config/metasrv.example.toml @@ -320,7 +320,7 @@ ttl = "90d" ## `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`, `repartition`, -## `repartition_group`, `wal_prune`, `batch_gc`. +## `repartition_group`, `wal_prune`, `batch_gc`, `reconcile_table`. ## When omitted, all current and future event types are recorded. ## Set to an empty array to disable event recording. #+ event_types = ["region_migration"] diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 6ca13566d7b..1fd8fe5a522 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -930,7 +930,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`, `admin_function`. +## `create_view`, `drop_view`, `admin_function`, `reconcile_table`. ## 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/common/meta/src/reconciliation.rs b/src/common/meta/src/reconciliation.rs index 351a415566a..06ff9e704af 100644 --- a/src/common/meta/src/reconciliation.rs +++ b/src/common/meta/src/reconciliation.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +pub(crate) mod event; pub mod manager; pub(crate) mod reconcile_catalog; pub(crate) mod reconcile_database; diff --git a/src/common/meta/src/reconciliation/event.rs b/src/common/meta/src/reconciliation/event.rs new file mode 100644 index 00000000000..9430fa575f0 --- /dev/null +++ b/src/common/meta/src/reconciliation/event.rs @@ -0,0 +1,351 @@ +// 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, SerializeEventSnafu}; +use common_event_recorder::event_table::{ + CATALOG_NAME_COLUMN, PHYSICAL_TABLE_ID_COLUMN, SCHEMA_NAME_COLUMN, TABLE_ID_COLUMN, + TABLE_NAME_COLUMN, column_schemas, nullable_string, nullable_value, +}; +use serde::Serialize; +use snafu::ResultExt; +use store_api::storage::TableId; + +use crate::reconciliation::ResolveStrategy; + +/// Stable event type stored for physical table reconciliation procedures. +pub(crate) const RECONCILE_TABLE_EVENT_TYPE: &str = "reconcile_table"; +const PAYLOAD_VERSION: u8 = 1; + +#[derive(Debug, Clone, Copy)] +enum ReconciliationEventType { + Table, +} + +impl ReconciliationEventType { + const fn as_str(self) -> &'static str { + match self { + Self::Table => RECONCILE_TABLE_EVENT_TYPE, + } + } +} + +/// Nullable object locators shared by all reconciliation event types. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(crate) struct ReconciliationLocator { + catalog_name: Option, + schema_name: Option, + table_name: Option, + table_id: Option, + physical_table_id: Option, +} + +impl ReconciliationLocator { + /// Creates a locator for a physical table with its fully qualified name and ID. + pub(crate) fn physical_table( + catalog_name: &str, + schema_name: &str, + table_name: &str, + table_id: TableId, + ) -> Self { + Self { + catalog_name: Some(catalog_name.to_string()), + schema_name: Some(schema_name.to_string()), + table_name: Some(table_name.to_string()), + table_id: Some(table_id), + ..Default::default() + } + } +} + +#[derive(Debug, Serialize)] +#[serde(untagged)] +enum ReconciliationPayload { + TableSubmitted(TableSubmittedPayload), + TableResult(TableResultPayload), +} + +#[derive(Debug, Serialize)] +struct TableSubmittedPayload { + version: u8, + resolve_strategy: &'static str, + is_subprocedure: bool, +} + +#[derive(Debug, Serialize)] +struct TableResultPayload { + version: u8, + complete: bool, + metadata_state: Option<&'static str>, + resolution_strategy_applied: Option<&'static str>, + resolved_column_count: Option, + scanned_region_count: usize, + updated_region_count: usize, + table_info_updated: bool, + last_completed_phase: Option<&'static str>, +} + +/// Event representation shared by reconciliation procedures. +#[derive(Debug)] +pub(crate) struct ReconciliationEvent { + event_type: ReconciliationEventType, + locator: ReconciliationLocator, + payload: Option, +} + +impl ReconciliationEvent { + /// Builds the bounded intent event emitted when table reconciliation is submitted. + pub(crate) fn table_submitted( + locator: ReconciliationLocator, + resolve_strategy: ResolveStrategy, + is_subprocedure: bool, + ) -> Self { + Self { + event_type: ReconciliationEventType::Table, + locator, + payload: Some(ReconciliationPayload::TableSubmitted( + TableSubmittedPayload { + version: PAYLOAD_VERSION, + resolve_strategy: resolve_strategy_name(resolve_strategy), + is_subprocedure, + }, + )), + } + } + + /// Builds a terminal event from the bounded reconciliation result summary. + #[allow(clippy::too_many_arguments)] + pub(crate) fn table_result( + locator: ReconciliationLocator, + complete: bool, + metadata_state: Option<&'static str>, + resolution_strategy_applied: Option, + resolved_column_count: Option, + scanned_region_count: usize, + updated_region_count: usize, + table_info_updated: bool, + last_completed_phase: Option<&'static str>, + ) -> Self { + Self { + event_type: ReconciliationEventType::Table, + locator, + payload: Some(ReconciliationPayload::TableResult(TableResultPayload { + version: PAYLOAD_VERSION, + complete, + metadata_state, + resolution_strategy_applied: resolution_strategy_applied.map(resolve_strategy_name), + resolved_column_count, + scanned_region_count, + updated_region_count, + table_info_updated, + last_completed_phase, + })), + } + } + + /// Builds a table lifecycle event whose reconciliation payload is null. + pub(crate) fn table_lifecycle(locator: ReconciliationLocator) -> Self { + Self::lifecycle(ReconciliationEventType::Table, locator) + } + + fn lifecycle(event_type: ReconciliationEventType, locator: ReconciliationLocator) -> Self { + Self { + event_type, + locator, + payload: None, + } + } + + fn schema() -> Vec { + column_schemas([ + &CATALOG_NAME_COLUMN, + &SCHEMA_NAME_COLUMN, + &TABLE_NAME_COLUMN, + &TABLE_ID_COLUMN, + &PHYSICAL_TABLE_ID_COLUMN, + ]) + } +} + +impl Event for ReconciliationEvent { + fn event_type(&self) -> &str { + self.event_type.as_str() + } + + fn json_payload(&self) -> Result { + match &self.payload { + Some(payload) => serde_json::to_value(payload).context(SerializeEventSnafu), + None => Ok(serde_json::Value::Null), + } + } + + fn extra_schema(&self) -> Vec { + Self::schema() + } + + fn extra_rows(&self) -> Result> { + Ok(vec![Row { + values: vec![ + nullable_string(self.locator.catalog_name.as_deref()), + nullable_string(self.locator.schema_name.as_deref()), + nullable_string(self.locator.table_name.as_deref()), + nullable_table_id(self.locator.table_id), + nullable_table_id(self.locator.physical_table_id), + ], + }]) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +fn resolve_strategy_name(strategy: ResolveStrategy) -> &'static str { + match strategy { + ResolveStrategy::UseLatest => "use_latest", + ResolveStrategy::UseMetasrv => "use_metasrv", + ResolveStrategy::AbortOnConflict => "abort_on_conflict", + } +} + +fn nullable_table_id(value: Option) -> api::v1::Value { + nullable_value(value.map(ValueData::U32Value)) +} + +#[cfg(test)] +mod tests { + use api::v1::value::ValueData; + use api::v1::{ColumnDataType, Row, SemanticType, Value}; + use common_event_recorder::Event; + use serde_json::json; + + use super::*; + + #[test] + fn reconciliation_events_use_the_shared_locator_contract() { + let table = ReconciliationEvent::table_lifecycle(ReconciliationLocator::physical_table( + "greptime", "public", "metrics", 42, + )); + assert_eq!(table.event_type(), RECONCILE_TABLE_EVENT_TYPE); + assert_eq!( + table + .extra_schema() + .into_iter() + .map(|column| { + ( + column.column_name, + ColumnDataType::try_from(column.datatype).unwrap(), + SemanticType::try_from(column.semantic_type).unwrap(), + ) + }) + .collect::>(), + vec![ + ( + "catalog_name".to_string(), + ColumnDataType::String, + SemanticType::Field, + ), + ( + "schema_name".to_string(), + ColumnDataType::String, + SemanticType::Field, + ), + ( + "table_name".to_string(), + ColumnDataType::String, + SemanticType::Field, + ), + ( + "table_id".to_string(), + ColumnDataType::Uint32, + SemanticType::Field, + ), + ( + "physical_table_id".to_string(), + ColumnDataType::Uint32, + SemanticType::Field, + ), + ] + ); + assert_eq!( + table.extra_rows().unwrap(), + vec![Row { + values: vec![ + ValueData::StringValue("greptime".to_string()).into(), + ValueData::StringValue("public".to_string()).into(), + ValueData::StringValue("metrics".to_string()).into(), + ValueData::U32Value(42).into(), + Value::default(), + ], + }] + ); + assert_eq!(table.json_payload().unwrap(), serde_json::Value::Null); + } + + #[test] + fn submitted_payloads_are_versioned_and_use_stable_strategy_names() { + for (strategy, expected) in [ + (ResolveStrategy::UseLatest, "use_latest"), + (ResolveStrategy::UseMetasrv, "use_metasrv"), + (ResolveStrategy::AbortOnConflict, "abort_on_conflict"), + ] { + let table = ReconciliationEvent::table_submitted( + ReconciliationLocator::physical_table("greptime", "public", "metrics", 42), + strategy, + true, + ); + assert_eq!( + table.json_payload().unwrap(), + json!({ + "version": 1, + "resolve_strategy": expected, + "is_subprocedure": true, + }) + ); + } + } + + #[test] + fn terminal_payloads_distinguish_complete_and_partial_results() { + let table = ReconciliationEvent::table_result( + ReconciliationLocator::physical_table("greptime", "public", "metrics", 42), + false, + Some("inconsistent"), + Some(ResolveStrategy::UseMetasrv), + Some(4), + 3, + 2, + true, + Some("update_table_info"), + ); + assert_eq!( + table.json_payload().unwrap(), + json!({ + "version": 1, + "complete": false, + "metadata_state": "inconsistent", + "resolution_strategy_applied": "use_metasrv", + "resolved_column_count": 4, + "scanned_region_count": 3, + "updated_region_count": 2, + "table_info_updated": true, + "last_completed_phase": "update_table_info", + }) + ); + } +} diff --git a/src/common/meta/src/reconciliation/reconcile_table.rs b/src/common/meta/src/reconciliation/reconcile_table.rs index 004b93d38c5..562a2061b0a 100644 --- a/src/common/meta/src/reconciliation/reconcile_table.rs +++ b/src/common/meta/src/reconciliation/reconcile_table.rs @@ -18,13 +18,12 @@ pub(crate) mod reconciliation_start; pub(crate) mod resolve_column_metadata; pub(crate) mod update_table_info; -use std::any::Any; use std::fmt::Debug; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, Error as ProcedureError, LockKey, Procedure, - Result as ProcedureResult, Status, + Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, + Procedure, Result as ProcedureResult, Status, }; use serde::{Deserialize, Serialize}; use snafu::ResultExt; @@ -42,6 +41,9 @@ use crate::key::{DeserializedValueWithBytes, TableMetadataManagerRef}; use crate::lock_key::{CatalogLock, SchemaLock, TableNameLock}; use crate::metrics; use crate::node_manager::NodeManagerRef; +use crate::reconciliation::event::{ + RECONCILE_TABLE_EVENT_TYPE, ReconciliationEvent, ReconciliationLocator, +}; use crate::reconciliation::reconcile_table::reconciliation_start::ReconciliationStart; use crate::reconciliation::reconcile_table::resolve_column_metadata::ResolveStrategy; use crate::reconciliation::utils::{ @@ -56,6 +58,105 @@ pub struct ReconcileTableContext { pub volatile_ctx: VolatileContext, } +#[derive(Debug, Clone, Copy)] +enum TablePhase { + Start, + ResolveColumnMetadata, + ReconcileRegions, + UpdateTableInfo, +} + +impl TablePhase { + const fn as_event_value(self) -> &'static str { + match self { + Self::Start => "start", + Self::ResolveColumnMetadata => "resolve_column_metadata", + Self::ReconcileRegions => "reconcile_regions", + Self::UpdateTableInfo => "update_table_info", + } + } +} + +#[derive(Debug, Clone, Copy)] +enum TableMetadataState { + Consistent, + Inconsistent, +} + +impl TableMetadataState { + const fn as_event_value(self) -> &'static str { + match self { + Self::Consistent => "consistent", + Self::Inconsistent => "inconsistent", + } + } +} + +#[derive(Debug, Default)] +struct ReconcileTableResultSummary { + metadata_state: Option, + resolution_strategy_applied: Option, + resolved_column_count: Option, + scanned_region_count: usize, + updated_region_count: usize, + table_info_updated: bool, + last_completed_phase: Option, +} + +impl ReconcileTableResultSummary { + fn record_scanned_regions(&mut self, scanned_region_count: usize) { + self.scanned_region_count = scanned_region_count; + } + + fn mark_start_completed(&mut self) { + self.last_completed_phase = Some(TablePhase::Start); + } + + fn record_metadata_state(&mut self, metadata_state: TableMetadataState) { + self.metadata_state = Some(metadata_state); + } + + fn record_resolution_strategy(&mut self, resolve_strategy: ResolveStrategy) { + self.resolution_strategy_applied = Some(resolve_strategy); + } + + fn record_resolved_columns( + &mut self, + metadata_state: TableMetadataState, + resolution_strategy_applied: Option, + resolved_column_count: Option, + ) { + self.metadata_state = Some(metadata_state); + self.resolution_strategy_applied = resolution_strategy_applied; + self.resolved_column_count = resolved_column_count; + self.last_completed_phase = Some(TablePhase::ResolveColumnMetadata); + } + + fn record_updated_regions(&mut self, updated_region_count: usize) { + self.updated_region_count = updated_region_count; + } + + fn mark_region_phase_completed(&mut self) { + self.last_completed_phase = Some(TablePhase::ReconcileRegions); + } + + fn mark_table_info_updated(&mut self) { + self.table_info_updated = true; + } + + fn mark_table_info_phase_completed(&mut self) { + self.last_completed_phase = Some(TablePhase::UpdateTableInfo); + } + + fn metadata_state(&self) -> Option<&'static str> { + self.metadata_state.map(TableMetadataState::as_event_value) + } + + fn last_completed_phase(&self) -> Option<&'static str> { + self.last_completed_phase.map(TablePhase::as_event_value) + } +} + impl ReconcileTableContext { /// Creates a new [`ReconcileTableContext`] with the given [`Context`] and [`PersistentContext`]. pub fn new(ctx: Context, persistent_ctx: PersistentContext) -> Self { @@ -147,6 +248,7 @@ impl PersistentContext { pub(crate) struct VolatileContext { pub(crate) table_meta: Option, pub(crate) metrics: ReconcileTableMetrics, + result_summary: ReconcileTableResultSummary, } pub struct ReconcileTableProcedure { @@ -256,6 +358,56 @@ impl Procedure for ReconcileTableProcedure { TableNameLock::new(table_ref.catalog, table_ref.schema, table_ref.table).into(), ]) } + + fn event(&self, ctx: &EventContext<'_>) -> Option> { + if !ctx.event_type_filter.allows(RECONCILE_TABLE_EVENT_TYPE) { + return None; + } + + let persistent_ctx = &self.context.persistent_ctx; + let table_name = &persistent_ctx.table_name; + let locator = ReconciliationLocator::physical_table( + &table_name.catalog_name, + &table_name.schema_name, + &table_name.table_name, + persistent_ctx.table_id, + ); + let event = match ctx.trigger { + EventTrigger::Submitted => ReconciliationEvent::table_submitted( + locator, + persistent_ctx.resolve_strategy, + persistent_ctx.is_subprocedure, + ), + EventTrigger::Succeeded => { + Self::result_event(locator, &self.context.volatile_ctx.result_summary, true) + } + EventTrigger::Failed | EventTrigger::Poisoned => { + Self::result_event(locator, &self.context.volatile_ctx.result_summary, false) + } + _ => ReconciliationEvent::table_lifecycle(locator), + }; + Some(Box::new(event)) + } +} + +impl ReconcileTableProcedure { + fn result_event( + locator: ReconciliationLocator, + summary: &ReconcileTableResultSummary, + complete: bool, + ) -> ReconciliationEvent { + ReconciliationEvent::table_result( + locator, + complete, + summary.metadata_state(), + summary.resolution_strategy_applied, + summary.resolved_column_count, + summary.scanned_region_count, + summary.updated_region_count, + summary.table_info_updated, + summary.last_completed_phase(), + ) + } } #[async_trait::async_trait] @@ -272,6 +424,390 @@ pub(crate) trait State: Sync + Send + Debug { ctx: &mut ReconcileTableContext, procedure_ctx: &ProcedureContext, ) -> Result<(Box, Status)>; +} - fn as_any(&self) -> &dyn Any; +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_event_recorder::{EventTypeFilter, EventTypeFilterRef}; + use common_procedure::{ + ChildSubmissionOutcome, EventContext, EventTrigger, Procedure, ProcedureId, ProcedureState, + RetryPhase, + }; + use common_procedure_test::MockContextProvider; + use serde_json::{Value, json}; + use store_api::storage::RegionId; + + use super::*; + use crate::ddl::test_util::datanode_handler::PartialSuccessDatanodeHandler; + use crate::key::DeserializedValueWithBytes; + use crate::key::table_info::TableInfoValue; + use crate::key::table_route::PhysicalTableRouteValue; + use crate::key::test_utils::new_test_table_info_with_name; + use crate::peer::Peer; + use crate::reconciliation::reconcile_table::reconcile_regions::ReconcileRegions; + use crate::reconciliation::utils::build_column_metadata_from_table_info; + use crate::rpc::router::{Region, RegionRoute}; + use crate::test_util::{MockDatanodeManager, new_ddl_context}; + + struct TableEventHarness { + procedure_id: ProcedureId, + lifecycle_state: ProcedureState, + event_type_filter: EventTypeFilterRef, + } + + impl TableEventHarness { + fn all() -> Self { + Self { + procedure_id: ProcedureId::random(), + lifecycle_state: ProcedureState::Running, + event_type_filter: Arc::new(EventTypeFilter::All), + } + } + + fn selected(event_types: impl IntoIterator) -> Self { + Self { + event_type_filter: Arc::new(EventTypeFilter::Only( + event_types.into_iter().map(str::to_string).collect(), + )), + ..Self::all() + } + } + + fn event( + &self, + procedure: &dyn Procedure, + trigger: EventTrigger, + ) -> Option> { + procedure.event(&EventContext { + procedure_id: self.procedure_id, + lifecycle_state: &self.lifecycle_state, + trigger, + event_type_filter: self.event_type_filter.clone(), + event_context: None, + }) + } + } + + #[test] + fn table_submitted_events_cover_root_and_child_intent() { + let events = TableEventHarness::all(); + let root = test_procedure(false); + let child = test_procedure(true); + for (procedure, is_subprocedure) in [(&root, false), (&child, true)] { + let submitted = events.event(procedure, EventTrigger::Submitted).unwrap(); + assert_eq!(submitted.event_type(), RECONCILE_TABLE_EVENT_TYPE); + assert_eq!( + submitted.json_payload().unwrap(), + json!({ + "version": 1, + "resolve_strategy": "use_latest", + "is_subprocedure": is_subprocedure, + }) + ); + } + } + + #[test] + fn table_non_terminal_lifecycle_events_have_null_payloads() { + let events = TableEventHarness::all(); + let mut procedure = test_procedure(true); + procedure.context.volatile_ctx.result_summary = populated_summary(); + + for trigger in [ + EventTrigger::Recovered, + EventTrigger::ChildSubmitted { + procedure_id: ProcedureId::random(), + outcome: ChildSubmissionOutcome::Accepted, + }, + EventTrigger::Retrying { + phase: RetryPhase::Execute, + attempt: 2, + }, + EventTrigger::RollingBack, + ] { + assert_eq!( + events + .event(&procedure, trigger) + .unwrap() + .json_payload() + .unwrap(), + Value::Null + ); + } + } + + #[test] + fn table_terminal_events_report_bounded_results() { + let events = TableEventHarness::all(); + let mut procedure = test_procedure(true); + procedure.context.volatile_ctx.result_summary = populated_summary(); + + for (trigger, complete) in [ + (EventTrigger::Succeeded, true), + (EventTrigger::Failed, false), + (EventTrigger::Poisoned, false), + ] { + assert_eq!( + events + .event(&procedure, trigger) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": complete, + "metadata_state": "inconsistent", + "resolution_strategy_applied": "use_latest", + "resolved_column_count": 3, + "scanned_region_count": 2, + "updated_region_count": 1, + "table_info_updated": true, + "last_completed_phase": "update_table_info", + }) + ); + } + } + + #[test] + fn table_event_filtering_uses_the_reconciliation_event_type() { + let procedure = test_procedure(false); + assert!( + TableEventHarness::selected([RECONCILE_TABLE_EVENT_TYPE]) + .event(&procedure, EventTrigger::Submitted) + .is_some() + ); + assert!( + TableEventHarness::selected(["create_table"]) + .event(&procedure, EventTrigger::Submitted) + .is_none() + ); + assert!( + TableEventHarness::selected([]) + .event(&procedure, EventTrigger::Submitted) + .is_none() + ); + } + + #[test] + fn table_result_summary_is_not_persisted() { + let events = TableEventHarness::all(); + let mut procedure = test_procedure(false); + procedure.context.volatile_ctx.result_summary = populated_summary(); + let value: Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); + assert!(value["persistent_ctx"].get("result_summary").is_none()); + + let loaded = ReconcileTableProcedure::from_json( + test_context(), + &serde_json::to_string(&value).unwrap(), + ) + .unwrap(); + assert_eq!( + events + .event(&loaded, EventTrigger::Succeeded) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": true, + "metadata_state": null, + "resolution_strategy_applied": null, + "resolved_column_count": null, + "scanned_region_count": 0, + "updated_region_count": 0, + "table_info_updated": false, + "last_completed_phase": null, + }) + ); + } + + #[test] + fn table_partial_summary_preserves_completed_work_without_advancing_phases() { + let events = TableEventHarness::all(); + let mut procedure = test_procedure(false); + let summary = &mut procedure.context.volatile_ctx.result_summary; + summary.record_scanned_regions(2); + + assert_eq!( + events + .event(&procedure, EventTrigger::Failed) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": false, + "metadata_state": null, + "resolution_strategy_applied": null, + "resolved_column_count": null, + "scanned_region_count": 2, + "updated_region_count": 0, + "table_info_updated": false, + "last_completed_phase": null, + }) + ); + + let summary = &mut procedure.context.volatile_ctx.result_summary; + summary.mark_start_completed(); + summary.record_metadata_state(TableMetadataState::Inconsistent); + assert_eq!( + events + .event(&procedure, EventTrigger::Failed) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": false, + "metadata_state": "inconsistent", + "resolution_strategy_applied": null, + "resolved_column_count": null, + "scanned_region_count": 2, + "updated_region_count": 0, + "table_info_updated": false, + "last_completed_phase": "start", + }) + ); + + let summary = &mut procedure.context.volatile_ctx.result_summary; + summary.record_resolved_columns( + TableMetadataState::Inconsistent, + Some(ResolveStrategy::UseLatest), + Some(3), + ); + summary.record_updated_regions(1); + assert_eq!( + events + .event(&procedure, EventTrigger::Failed) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": false, + "metadata_state": "inconsistent", + "resolution_strategy_applied": "use_latest", + "resolved_column_count": 3, + "scanned_region_count": 2, + "updated_region_count": 1, + "table_info_updated": false, + "last_completed_phase": "resolve_column_metadata", + }) + ); + } + + #[tokio::test] + async fn table_mixed_region_outcome_preserves_partial_result() { + let ddl_context = new_ddl_context(Arc::new(MockDatanodeManager::new( + PartialSuccessDatanodeHandler { retryable: false }, + ))); + let context = Context { + node_manager: ddl_context.node_manager, + table_metadata_manager: ddl_context.table_metadata_manager, + cache_invalidator: ddl_context.cache_invalidator, + }; + let mut procedure = ReconcileTableProcedure::new( + context, + 42, + TableName::new("greptime", "public", "metrics"), + ResolveStrategy::UseLatest, + false, + ); + + let mut table_info = new_test_table_info_with_name(42, "metrics"); + table_info.meta.column_ids = vec![0, 1, 2]; + let name_to_ids = table_info.name_to_ids().unwrap(); + let column_metadatas = build_column_metadata_from_table_info( + table_info.meta.schema.column_schemas(), + &table_info.meta.primary_key_indices, + &name_to_ids, + ) + .unwrap(); + let region_ids = [RegionId::new(42, 1), RegionId::new(42, 2)]; + procedure.context.persistent_ctx.table_info_value = Some( + DeserializedValueWithBytes::from_inner(TableInfoValue::new(table_info)), + ); + procedure.context.persistent_ctx.physical_table_route = + Some(PhysicalTableRouteValue::new(vec![ + RegionRoute { + region: Region::new_test(region_ids[0]), + leader_peer: Some(Peer::empty(1)), + ..Default::default() + }, + RegionRoute { + region: Region::new_test(region_ids[1]), + leader_peer: Some(Peer::empty(2)), + ..Default::default() + }, + ])); + + let summary = &mut procedure.context.volatile_ctx.result_summary; + summary.record_scanned_regions(region_ids.len()); + summary.mark_start_completed(); + summary.record_resolved_columns( + TableMetadataState::Inconsistent, + Some(ResolveStrategy::UseLatest), + Some(column_metadatas.len()), + ); + procedure.state = Box::new(ReconcileRegions::new(column_metadatas, region_ids.to_vec())); + + let procedure_ctx = ProcedureContext { + procedure_id: ProcedureId::random(), + provider: Arc::new(MockContextProvider::default()), + event_context: None, + }; + assert!(procedure.execute(&procedure_ctx).await.is_err()); + + assert_eq!( + TableEventHarness::all() + .event(&procedure, EventTrigger::Failed) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": false, + "metadata_state": "inconsistent", + "resolution_strategy_applied": "use_latest", + "resolved_column_count": 3, + "scanned_region_count": 2, + "updated_region_count": 1, + "table_info_updated": false, + "last_completed_phase": "resolve_column_metadata", + }) + ); + } + + fn populated_summary() -> ReconcileTableResultSummary { + ReconcileTableResultSummary { + metadata_state: Some(TableMetadataState::Inconsistent), + resolution_strategy_applied: Some(ResolveStrategy::UseLatest), + resolved_column_count: Some(3), + scanned_region_count: 2, + updated_region_count: 1, + table_info_updated: true, + last_completed_phase: Some(TablePhase::UpdateTableInfo), + } + } + + fn test_procedure(is_subprocedure: bool) -> ReconcileTableProcedure { + ReconcileTableProcedure::new( + test_context(), + 42, + TableName::new("greptime", "public", "metrics"), + ResolveStrategy::UseLatest, + is_subprocedure, + ) + } + + fn test_context() -> Context { + let ddl_context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + Context { + node_manager: ddl_context.node_manager, + table_metadata_manager: ddl_context.table_metadata_manager, + cache_invalidator: ddl_context.cache_invalidator, + } + } } diff --git a/src/common/meta/src/reconciliation/reconcile_table/reconcile_regions.rs b/src/common/meta/src/reconciliation/reconcile_table/reconcile_regions.rs index 6ee3df70f37..fd5d55fb0cb 100644 --- a/src/common/meta/src/reconciliation/reconcile_table/reconcile_regions.rs +++ b/src/common/meta/src/reconciliation/reconcile_table/reconcile_regions.rs @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::any::Any; use std::collections::{HashMap, HashSet}; use api::v1::column_def::try_as_column_def; @@ -65,7 +64,6 @@ impl State for ReconcileRegions { let table_meta = ctx.build_table_meta(&self.column_metadatas)?; ctx.volatile_ctx.table_meta = Some(table_meta); let table_id = ctx.table_id(); - let table_name = ctx.table_name(); let primary_keys = self .column_metadatas @@ -126,10 +124,12 @@ impl State for ReconcileRegions { } } - let mut results = future::join_all(sync_column_tsks) - .await - .into_iter() - .collect::>>()?; + let results = future::join_all(sync_column_tsks).await; + let updated_region_count = results.iter().filter(|result| result.is_ok()).count(); + ctx.volatile_ctx + .result_summary + .record_updated_regions(updated_region_count); + let mut results = results.into_iter().collect::>>()?; // Ensures all the column metadatas are the same. let column_metadatas = @@ -137,17 +137,22 @@ impl State for ReconcileRegions { UnexpectedSnafu { err_msg: format!( "The table column metadata schemas from datanodes are not the same, table: {}, table_id: {}", - table_name, + ctx.table_name(), table_id ), }, )?; + ctx.volatile_ctx + .result_summary + .mark_region_phase_completed(); + // Checks all column metadatas are consistent, and updates the table info if needed. if column_metadatas != self.column_metadatas { info!( "Datanode column metadatas are not consistent with metasrv, updating metasrv's column metadatas, table: {}, table_id: {}", - table_name, table_id + ctx.table_name(), + table_id ); // Safety: fetched in the above. let table_info_value = ctx.persistent_ctx.table_info_value.clone().unwrap(); @@ -159,10 +164,6 @@ impl State for ReconcileRegions { Ok((Box::new(ReconciliationEnd), Status::executing(false))) } - - fn as_any(&self) -> &dyn Any { - self - } } /// Makes an alter region request to sync columns. diff --git a/src/common/meta/src/reconciliation/reconcile_table/reconciliation_end.rs b/src/common/meta/src/reconciliation/reconcile_table/reconciliation_end.rs index 9f28d7f4be5..b6e695dca46 100644 --- a/src/common/meta/src/reconciliation/reconcile_table/reconciliation_end.rs +++ b/src/common/meta/src/reconciliation/reconcile_table/reconciliation_end.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::any::Any; - use common_procedure::{Context as ProcedureContext, Status}; use common_telemetry::info; use serde::{Deserialize, Serialize}; @@ -46,8 +44,4 @@ impl State for ReconciliationEnd { Ok((Box::new(ReconciliationEnd), Status::done())) } - - fn as_any(&self) -> &dyn Any { - self - } } diff --git a/src/common/meta/src/reconciliation/reconcile_table/reconciliation_start.rs b/src/common/meta/src/reconciliation/reconcile_table/reconciliation_start.rs index 2fcc8ded0c0..48a7fff8b96 100644 --- a/src/common/meta/src/reconciliation/reconcile_table/reconciliation_start.rs +++ b/src/common/meta/src/reconciliation/reconcile_table/reconciliation_start.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::any::Any; - use common_procedure::{Context as ProcedureContext, Status}; use common_telemetry::info; use serde::{Deserialize, Serialize}; @@ -46,7 +44,6 @@ impl State for ReconciliationStart { procedure_ctx: &ProcedureContext, ) -> Result<(Box, Status)> { let table_id = ctx.table_id(); - let table_name = ctx.table_name(); let (physical_table_id, physical_table_route) = ctx .table_metadata_manager @@ -58,14 +55,17 @@ impl State for ReconciliationStart { error::UnexpectedSnafu { err_msg: format!( "Reconcile table only works for physical table, but got logical table: {}, table_id: {}", - table_name, table_id + ctx.table_name(), + table_id ), } ); info!( "Reconciling table: {}, table_id: {}, procedure_id: {}", - table_name, table_id, procedure_ctx.procedure_id + ctx.table_name(), + table_id, + procedure_ctx.procedure_id ); // TODO(weny): Repairs the table route if needed. let region_metadata_lister = RegionMetadataLister::new(ctx.node_manager.clone()); @@ -79,6 +79,9 @@ impl State for ReconciliationStart { .list(physical_table_id, &physical_table_route.region_routes) .await? }; + ctx.volatile_ctx + .result_summary + .record_scanned_regions(region_metadatas.len()); ensure!(!region_metadatas.is_empty(), { metrics::METRIC_META_RECONCILIATION_STATS @@ -92,7 +95,8 @@ impl State for ReconciliationStart { error::UnexpectedSnafu { err_msg: format!( "No region metadata found for table: {}, table_id: {}", - table_name, table_id + ctx.table_name(), + table_id ), } }); @@ -109,11 +113,14 @@ impl State for ReconciliationStart { error::UnexpectedSnafu { err_msg: format!( "Some regions are not opened, table: {}, table_id: {}", - table_name, table_id + ctx.table_name(), + table_id ), } }); + ctx.volatile_ctx.result_summary.mark_start_completed(); + // Persist the physical table route. // TODO(weny): refetch the physical table route if repair is needed. ctx.persistent_ctx.physical_table_route = Some(physical_table_route); @@ -127,8 +134,4 @@ impl State for ReconciliationStart { Status::executing(false), )) } - - fn as_any(&self) -> &dyn Any { - self - } } diff --git a/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs b/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs index d45cfe464ea..27950f0fd5d 100644 --- a/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs +++ b/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::any::Any; - use async_trait::async_trait; use common_procedure::{Context as ProcedureContext, Status}; use common_telemetry::info; @@ -25,7 +23,7 @@ use strum::AsRefStr; use crate::error::{self, MissingColumnIdsSnafu, Result}; use crate::reconciliation::reconcile_table::reconcile_regions::ReconcileRegions; use crate::reconciliation::reconcile_table::update_table_info::UpdateTableInfo; -use crate::reconciliation::reconcile_table::{ReconcileTableContext, State}; +use crate::reconciliation::reconcile_table::{ReconcileTableContext, State, TableMetadataState}; use crate::reconciliation::utils::{ ResolveColumnMetadataResult, build_column_metadata_from_table_info, check_column_metadatas_consistent, resolve_column_metadatas_with_latest, @@ -101,6 +99,12 @@ impl State for ResolveColumnMetadata { table_name, table_id ); + ctx.volatile_ctx.result_summary.record_resolved_columns( + TableMetadataState::Consistent, + None, + Some(column_metadatas.len()), + ); + // Update metrics. ctx.mut_metrics().resolve_column_metadata_result = Some(ResolveColumnMetadataResult::Consistent); @@ -110,6 +114,10 @@ impl State for ResolveColumnMetadata { )); }; + ctx.volatile_ctx + .result_summary + .record_metadata_state(TableMetadataState::Inconsistent); + match self.strategy { ResolveStrategy::UseMetasrv => { let table_info_value = ctx.persistent_ctx.table_info_value.as_ref().unwrap(); @@ -126,6 +134,12 @@ impl State for ResolveColumnMetadata { let region_ids = resolve_column_metadatas_with_metasrv(&column_metadata, &self.region_metadata)?; + ctx.volatile_ctx.result_summary.record_resolved_columns( + TableMetadataState::Inconsistent, + Some(self.strategy), + Some(column_metadata.len()), + ); + // Update metrics. let metrics = ctx.mut_metrics(); metrics.resolve_column_metadata_result = @@ -139,6 +153,12 @@ impl State for ResolveColumnMetadata { let (column_metadatas, region_ids) = resolve_column_metadatas_with_latest(&self.region_metadata)?; + ctx.volatile_ctx.result_summary.record_resolved_columns( + TableMetadataState::Inconsistent, + Some(self.strategy), + Some(column_metadatas.len()), + ); + // Update metrics. let metrics = ctx.mut_metrics(); metrics.resolve_column_metadata_result = @@ -151,6 +171,10 @@ impl State for ResolveColumnMetadata { ResolveStrategy::AbortOnConflict => { let table_name = table_name.to_string(); + ctx.volatile_ctx + .result_summary + .record_resolution_strategy(self.strategy); + // Update metrics. let metrics = ctx.mut_metrics(); metrics.resolve_column_metadata_result = @@ -163,8 +187,4 @@ impl State for ResolveColumnMetadata { } } } - - fn as_any(&self) -> &dyn Any { - self - } } diff --git a/src/common/meta/src/reconciliation/reconcile_table/update_table_info.rs b/src/common/meta/src/reconciliation/reconcile_table/update_table_info.rs index 790184c07ab..7ef62265de3 100644 --- a/src/common/meta/src/reconciliation/reconcile_table/update_table_info.rs +++ b/src/common/meta/src/reconciliation/reconcile_table/update_table_info.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::any::Any; - use common_procedure::{Context as ProcedureContext, Status}; use common_telemetry::info; use serde::{Deserialize, Serialize}; @@ -81,6 +79,9 @@ impl State for UpdateTableInfo { ctx.table_name(), ctx.table_id() ); + ctx.volatile_ctx + .result_summary + .mark_table_info_phase_completed(); return Ok((Box::new(ReconciliationEnd), Status::executing(true))); } @@ -98,6 +99,7 @@ impl State for UpdateTableInfo { new_table_info, ) .await?; + ctx.volatile_ctx.result_summary.mark_table_info_updated(); let table_ref = ctx.table_name().table_ref(); let table_id = ctx.table_id(); @@ -119,11 +121,10 @@ impl State for UpdateTableInfo { // Update metrics. let metrics = ctx.mut_metrics(); metrics.update_table_info = true; + ctx.volatile_ctx + .result_summary + .mark_table_info_phase_completed(); Ok((Box::new(ReconciliationEnd), Status::executing(true))) } - - fn as_any(&self) -> &dyn Any { - self - } } diff --git a/tests-integration/tests/main.rs b/tests-integration/tests/main.rs index 965a3a67462..56108d4d3ff 100644 --- a/tests-integration/tests/main.rs +++ b/tests-integration/tests/main.rs @@ -37,6 +37,7 @@ mod repartition_event; #[macro_use] mod repartition_expr_version; mod mysql; +mod reconciliation_event; mod view_ddl_event; mod wal_prune_event; diff --git a/tests-integration/tests/reconciliation_event.rs b/tests-integration/tests/reconciliation_event.rs new file mode 100644 index 00000000000..c4b4b88ea2d --- /dev/null +++ b/tests-integration/tests/reconciliation_event.rs @@ -0,0 +1,144 @@ +// 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 client::OutputData; +use common_meta::reconciliation::ResolveStrategy; +use common_procedure::watcher; +use frontend::instance::Instance; +use servers::query_handler::sql::SqlQueryHandler; +use session::context::QueryContext; +use table::table_reference::TableReference; +use tests_integration::cluster::GreptimeDbClusterBuilder; + +use crate::event_recorder_test_util::{assert_eventually_eq, find_eventually_string}; + +const EVENTS_TABLE: &str = "greptime_private.events"; +const CATALOG: &str = "greptime"; +const DATABASE: &str = "reconciliation_event_database"; + +#[tokio::test(flavor = "multi_thread")] +async fn test_table_reconciliation_events() { + common_telemetry::init_default_ut_logging(); + + let cluster = GreptimeDbClusterBuilder::new("table_reconciliation_events") + .await + .with_datanodes(1) + .build(true) + .await; + let frontend = cluster.fe_instance().clone(); + run_sql(frontend.as_ref(), &format!("CREATE DATABASE {DATABASE}")).await; + run_sql( + frontend.as_ref(), + &format!("CREATE TABLE {DATABASE}.metrics (ts TIMESTAMP TIME INDEX)"), + ) + .await; + + let table_procedure_id = cluster + .metasrv + .reconciliation_manager() + .reconcile_table( + TableReference { + catalog: CATALOG, + schema: DATABASE, + table: "metrics", + }, + ResolveStrategy::UseLatest, + ) + .await + .unwrap(); + let mut procedure_watcher = cluster + .metasrv + .procedure_manager() + .procedure_watcher(table_procedure_id) + .unwrap(); + watcher::wait(&mut procedure_watcher).await.unwrap(); + let table_procedure_id = table_procedure_id.to_string(); + + assert_eventually_eq( + &frontend, + &format!( + "SELECT count(*) = 1 AS matches FROM {EVENTS_TABLE} \ + WHERE type = 'reconcile_table' AND procedure_id = '{table_procedure_id}' \ + AND json_get_string(procedure_trigger, 'type') = 'Submitted' \ + AND catalog_name = '{CATALOG}' AND schema_name = '{DATABASE}' \ + AND table_name = 'metrics' AND table_id IS NOT NULL \ + AND physical_table_id IS NULL \ + AND json_get_int(payload, 'version') = 1 \ + AND json_get_string(payload, 'resolve_strategy') = 'use_latest' \ + AND json_get_bool(payload, 'is_subprocedure') = false" + ), + "\ ++---------+ +| matches | ++---------+ +| true | ++---------+", + ) + .await; + + let table_result_payload = find_eventually_string( + &frontend, + &format!( + "SELECT json_to_string(payload) AS payload FROM {EVENTS_TABLE} \ + WHERE type = 'reconcile_table' AND procedure_id = '{table_procedure_id}' \ + AND json_get_string(procedure_trigger, 'type') = 'Succeeded' \ + AND catalog_name = '{CATALOG}' AND schema_name = '{DATABASE}' \ + AND table_name = 'metrics' AND table_id IS NOT NULL \ + AND physical_table_id IS NULL LIMIT 1" + ), + "payload", + ) + .await; + assert_eq!( + serde_json::from_str::(&table_result_payload).unwrap(), + serde_json::json!({ + "version": 1, + "complete": true, + "metadata_state": "consistent", + "resolution_strategy_applied": null, + "resolved_column_count": 1, + "scanned_region_count": 1, + "updated_region_count": 0, + "table_info_updated": true, + "last_completed_phase": "update_table_info", + }) + ); + + assert_eventually_eq( + &frontend, + &format!( + "SELECT count(*) AS terminal_event_count FROM {EVENTS_TABLE} \ + WHERE type = 'reconcile_table' AND procedure_id = '{table_procedure_id}' \ + AND json_get_string(procedure_trigger, 'type') = 'Succeeded' \ + AND catalog_name = '{CATALOG}' AND schema_name = '{DATABASE}' \ + AND table_name = 'metrics' AND table_id IS NOT NULL \ + AND physical_table_id IS NULL" + ), + "\ ++----------------------+ +| terminal_event_count | ++----------------------+ +| 1 | ++----------------------+", + ) + .await; +} + +async fn run_sql(instance: &Instance, sql: &str) { + let output = SqlQueryHandler::do_query(instance, sql, QueryContext::arc()) + .await + .remove(0) + .unwrap(); + assert!(matches!(output.data, OutputData::AffectedRows(_)), "{sql}"); +} From 28e2b5aca26fdaf49d5851e88de80ac787c4a14d Mon Sep 17 00:00:00 2001 From: Yingwen Date: Fri, 28 Aug 2026 07:55:57 +0000 Subject: [PATCH 15/22] feat: expose missing SST manifest fields (#8965) * feat: expose missing SST manifest fields Signed-off-by: evenyag * fix: preserve SST manifest column ordinals Signed-off-by: evenyag --------- Signed-off-by: evenyag --- src/metric-engine/src/engine/flush.rs | 13 +-- src/mito2/src/engine/basic_test.rs | 13 +-- src/mito2/src/region.rs | 2 + src/store-api/src/sst_entry.rs | 75 ++++++++++++++++ .../common/information_schema/ssts.result | 90 ++++++++++--------- .../common/system/information_schema.result | 2 + 6 files changed, 139 insertions(+), 56 deletions(-) diff --git a/src/metric-engine/src/engine/flush.rs b/src/metric-engine/src/engine/flush.rs index 62b2266249c..21e642de0b2 100644 --- a/src/metric-engine/src/engine/flush.rs +++ b/src/metric-engine/src/engine/flush.rs @@ -121,6 +121,7 @@ mod tests { .map(|path| path.replace(&e.file_id, "")); e.file_id = "".to_string(); e.index_version = 0; + e.max_row_group_uncompressed_size = 0; // Round down sizes to nearest 1000 to avoid exact size // comparisons that break when the SST format changes. e.file_size = e.file_size / 1000 * 1000; @@ -133,12 +134,12 @@ mod tests { assert_eq!( debug_format, r#" -ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000001/data/.parquet", file_size: 3000, index_file_path: Some("test_metric_region/11_0000000001/data/index/.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(20), origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") } -ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000002/data/.parquet", file_size: 3000, index_file_path: Some("test_metric_region/11_0000000002/data/index/.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") } -ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417473(11, 16777217), table_id: 11, region_number: 16777217, region_group: 1, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000001/metadata/.parquet", file_size: 4000, index_file_path: None, index_file_size: None, num_rows: 8, num_row_groups: 1, num_series: Some(8), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(8), origin_region_id: 47261417473(11, 16777217), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_4771708\t665867_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_4771708\t665868\0\0\x06") } -ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417474(11, 16777218), table_id: 11, region_number: 16777218, region_group: 1, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000002/metadata/.parquet", file_size: 4000, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, num_series: Some(4), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4), origin_region_id: 47261417474(11, 16777218), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_4771708\t665858_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_4771708\t665858\0\0\x06") } -ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/22_0000000042/data/.parquet", file_size: 3000, index_file_path: Some("test_metric_region/22_0000000042/data/index/.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") } -ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94506057770(22, 16777258), table_id: 22, region_number: 16777258, region_group: 1, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/22_0000000042/metadata/.parquet", file_size: 4000, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, num_series: Some(4), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4), origin_region_id: 94506057770(22, 16777258), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_9543417\t331754_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_9543417\t331754\0\0\x06") }"#, +ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000001/data/.parquet", file_size: 3000, max_row_group_uncompressed_size: 0, index_file_path: Some("test_metric_region/11_0000000001/data/index/.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(20), partition_expr: None, origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") } +ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000002/data/.parquet", file_size: 3000, max_row_group_uncompressed_size: 0, index_file_path: Some("test_metric_region/11_0000000002/data/index/.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") } +ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417473(11, 16777217), table_id: 11, region_number: 16777217, region_group: 1, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000001/metadata/.parquet", file_size: 4000, max_row_group_uncompressed_size: 0, index_file_path: None, index_file_size: None, num_rows: 8, num_row_groups: 1, num_series: Some(8), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(8), partition_expr: None, origin_region_id: 47261417473(11, 16777217), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_4771708\t665867_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_4771708\t665868\0\0\x06") } +ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417474(11, 16777218), table_id: 11, region_number: 16777218, region_group: 1, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000002/metadata/.parquet", file_size: 4000, max_row_group_uncompressed_size: 0, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, num_series: Some(4), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4), partition_expr: None, origin_region_id: 47261417474(11, 16777218), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_4771708\t665858_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_4771708\t665858\0\0\x06") } +ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/22_0000000042/data/.parquet", file_size: 3000, max_row_group_uncompressed_size: 0, index_file_path: Some("test_metric_region/22_0000000042/data/index/.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") } +ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94506057770(22, 16777258), table_id: 22, region_number: 16777258, region_group: 1, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test_metric_region/22_0000000042/metadata/.parquet", file_size: 4000, max_row_group_uncompressed_size: 0, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, num_series: Some(4), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4), partition_expr: None, origin_region_id: 94506057770(22, 16777258), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_9543417\t331754_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_9543417\t331754\0\0\x06") }"#, ); // list from storage let storage_entries = mito diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index da36c827665..6d90f90b784 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/src/mito2/src/engine/basic_test.rs @@ -882,9 +882,9 @@ async fn test_cache_null_primary_key_with_format(flat_format: bool) { #[tokio::test] async fn test_list_ssts() { test_list_ssts_with_format(false, r#" -ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000001/.parquet", file_size: 2951, index_file_path: Some("test/11_0000000001/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } -ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000002/.parquet", file_size: 2951, index_file_path: Some("test/11_0000000002/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } -ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test/22_0000000042/.parquet", file_size: 2951, index_file_path: Some("test/22_0000000042/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"# , +ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000001/.parquet", file_size: 2951, max_row_group_uncompressed_size: 0, index_file_path: Some("test/11_0000000001/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } +ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000002/.parquet", file_size: 2951, max_row_group_uncompressed_size: 0, index_file_path: Some("test/11_0000000002/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } +ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test/22_0000000042/.parquet", file_size: 2951, max_row_group_uncompressed_size: 0, index_file_path: Some("test/22_0000000042/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"# , r#" StorageSstEntry { file_path: "test/11_0000000001/.parquet", file_size: None, last_modified_ms: None, node_id: None } StorageSstEntry { file_path: "test/11_0000000001/index/.puffin", file_size: None, last_modified_ms: None, node_id: None } @@ -893,9 +893,9 @@ StorageSstEntry { file_path: "test/11_0000000002/index/.puffin", file_s StorageSstEntry { file_path: "test/22_0000000042/.parquet", file_size: None, last_modified_ms: None, node_id: None } StorageSstEntry { file_path: "test/22_0000000042/index/.puffin", file_size: None, last_modified_ms: None, node_id: None }"#).await; test_list_ssts_with_format(true, r#" -ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000001/.parquet", file_size: 3365, index_file_path: Some("test/11_0000000001/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } -ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000002/.parquet", file_size: 3365, index_file_path: Some("test/11_0000000002/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } -ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test/22_0000000042/.parquet", file_size: 3365, index_file_path: Some("test/22_0000000042/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"#, +ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000001/.parquet", file_size: 3365, max_row_group_uncompressed_size: 0, index_file_path: Some("test/11_0000000001/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } +ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "", index_version: 0, level: 0, file_path: "test/11_0000000002/.parquet", file_size: 3365, max_row_group_uncompressed_size: 0, index_file_path: Some("test/11_0000000002/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") } +ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "", index_version: 0, level: 0, file_path: "test/22_0000000042/.parquet", file_size: 3365, max_row_group_uncompressed_size: 0, index_file_path: Some("test/22_0000000042/index/.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"#, r#" StorageSstEntry { file_path: "test/11_0000000001/.parquet", file_size: None, last_modified_ms: None, node_id: None } StorageSstEntry { file_path: "test/11_0000000001/index/.puffin", file_size: None, last_modified_ms: None, node_id: None } @@ -967,6 +967,7 @@ async fn test_list_ssts_with_format( .map(|p| p.replace(&e.file_id, "")); e.file_id = "".to_string(); e.index_version = 0; + e.max_row_group_uncompressed_size = 0; format!("\n{:?}", e) }) .sorted() diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index 76204da9cf6..b68b0d1cccd 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -839,6 +839,7 @@ impl MitoRegion { level: meta.level, file_path: sst_file_path(table_dir, meta.file_id(), path_type), file_size: meta.file_size, + max_row_group_uncompressed_size: meta.max_row_group_uncompressed_size, index_file_path, index_file_size, num_rows: meta.num_rows, @@ -847,6 +848,7 @@ impl MitoRegion { min_ts: meta.time_range.0, max_ts: meta.time_range.1, sequence: meta.sequence.map(|s| s.get()), + partition_expr: meta.partition_expr.as_ref().map(ToString::to_string), origin_region_id, node_id: None, visible, diff --git a/src/store-api/src/sst_entry.rs b/src/store-api/src/sst_entry.rs index f1bb1f9c36f..c70e909ebd6 100644 --- a/src/store-api/src/sst_entry.rs +++ b/src/store-api/src/sst_entry.rs @@ -65,6 +65,8 @@ pub struct ManifestSstEntry { pub file_path: String, /// File size in bytes. pub file_size: u64, + /// Maximum uncompressed row group size in bytes. Zero means unknown. + pub max_row_group_uncompressed_size: u64, /// Full path of the index file in object store. pub index_file_path: Option, /// File size of the index file in object store. @@ -81,6 +83,8 @@ pub struct ManifestSstEntry { pub max_ts: Timestamp, /// The sequence number associated with this file. pub sequence: Option, + /// Human-readable partition expression associated with this file. + pub partition_expr: Option, /// The region id of region that creates the file. pub origin_region_id: RegionId, /// The node id fetched from the manifest. @@ -122,6 +126,12 @@ impl ManifestSstEntry { ColumnSchema::new("visible", Ty::boolean_datatype(), false), ColumnSchema::new("primary_key_min", Ty::binary_datatype(), true), ColumnSchema::new("primary_key_max", Ty::binary_datatype(), true), + ColumnSchema::new( + "max_row_group_uncompressed_size", + Ty::uint64_datatype(), + false, + ), + ColumnSchema::new("partition_expr", Ty::string_datatype(), true), ])) } @@ -139,6 +149,8 @@ impl ManifestSstEntry { let levels = entries.iter().map(|e| e.level); let file_paths = entries.iter().map(|e| e.file_path.as_str()); let file_sizes = entries.iter().map(|e| e.file_size); + let max_row_group_uncompressed_sizes = + entries.iter().map(|e| e.max_row_group_uncompressed_size); let index_file_paths = entries.iter().map(|e| e.index_file_path.as_ref()); let index_file_sizes = entries.iter().map(|e| e.index_file_size); let num_rows = entries.iter().map(|e| e.num_rows); @@ -155,6 +167,7 @@ impl ManifestSstEntry { .map(|ts| ts.value()) }); let sequences = entries.iter().map(|e| e.sequence); + let partition_exprs = entries.iter().map(|e| e.partition_expr.as_ref()); let origin_region_ids = entries.iter().map(|e| e.origin_region_id.as_u64()); let node_ids = entries.iter().map(|e| e.node_id); let visible_flags = entries.iter().map(|e| Some(e.visible)); @@ -186,6 +199,10 @@ impl ManifestSstEntry { Arc::new(BooleanArray::from_iter(visible_flags)), Arc::new(BinaryArray::from_iter(primary_key_min)), Arc::new(BinaryArray::from_iter(primary_key_max)), + Arc::new(UInt64Array::from_iter_values( + max_row_group_uncompressed_sizes, + )), + Arc::new(StringArray::from_iter(partition_exprs)), ]; DfRecordBatch::try_new(schema.arrow_schema().clone(), columns) @@ -458,6 +475,7 @@ mod tests { level: 1, file_path: "/p1".to_string(), file_size: 100, + max_row_group_uncompressed_size: 80, index_file_path: None, index_file_size: None, num_rows: 10, @@ -466,6 +484,7 @@ mod tests { min_ts: Timestamp::new_millisecond(1000), // 1s -> 1_000_000_000ns max_ts: Timestamp::new_second(2), // 2s -> 2_000_000_000ns sequence: None, + partition_expr: Some("a < 10".to_string()), origin_region_id: region_id1, node_id: Some(1), visible: false, @@ -484,6 +503,7 @@ mod tests { level: 3, file_path: "/p2".to_string(), file_size: 200, + max_row_group_uncompressed_size: 160, index_file_path: Some("idx".to_string()), index_file_size: Some(11), num_rows: 20, @@ -492,6 +512,7 @@ mod tests { min_ts: Timestamp::new_nanosecond(5), // 5ns max_ts: Timestamp::new_microsecond(2000), // 2ms -> 2_000_000ns sequence: Some(9), + partition_expr: None, origin_region_id: region_id2, node_id: None, visible: true, @@ -506,6 +527,44 @@ mod tests { // Schema checks assert_eq!(schema.arrow_schema().fields().len(), batch.num_columns()); assert_eq!(2, batch.num_rows()); + let expected_columns = [ + "table_dir", + "region_id", + "table_id", + "region_number", + "region_group", + "region_sequence", + "file_id", + "index_version", + "level", + "file_path", + "file_size", + "index_file_path", + "index_file_size", + "num_rows", + "num_row_groups", + "num_series", + "min_ts", + "max_ts", + "sequence", + "origin_region_id", + "node_id", + "visible", + "primary_key_min", + "primary_key_max", + "max_row_group_uncompressed_size", + "partition_expr", + ]; + assert_eq!( + expected_columns, + schema + .arrow_schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>() + .as_slice() + ); for (i, f) in schema.arrow_schema().fields().iter().enumerate() { assert_eq!(f.name(), batch.schema().field(i).name()); assert_eq!(f.is_nullable(), batch.schema().field(i).is_nullable()); @@ -704,6 +763,22 @@ mod tests { .unwrap(); assert_eq!(b"zzz", primary_key_max.value(0)); assert!(primary_key_max.is_null(1)); + + let max_row_group_uncompressed_sizes = batch + .column(24) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(80, max_row_group_uncompressed_sizes.value(0)); + assert_eq!(160, max_row_group_uncompressed_sizes.value(1)); + + let partition_exprs = batch + .column(25) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!("a < 10", partition_exprs.value(0)); + assert!(partition_exprs.is_null(1)); } #[test] diff --git a/tests/cases/standalone/common/information_schema/ssts.result b/tests/cases/standalone/common/information_schema/ssts.result index caf6a239607..9777e3b12a9 100644 --- a/tests/cases/standalone/common/information_schema/ssts.result +++ b/tests/cases/standalone/common/information_schema/ssts.result @@ -1,33 +1,35 @@ DESC TABLE information_schema.ssts_manifest; -+------------------+---------------------+-----+------+---------+---------------+ -| Column | Type | Key | Null | Default | Semantic Type | -+------------------+---------------------+-----+------+---------+---------------+ -| table_dir | String | | NO | | FIELD | -| region_id | UInt64 | | NO | | FIELD | -| table_id | UInt32 | | NO | | FIELD | -| region_number | UInt32 | | NO | | FIELD | -| region_group | UInt8 | | NO | | FIELD | -| region_sequence | UInt32 | | NO | | FIELD | -| file_id | String | | NO | | FIELD | -| index_version | UInt64 | | NO | | FIELD | -| level | UInt8 | | NO | | FIELD | -| file_path | String | | NO | | FIELD | -| file_size | UInt64 | | NO | | FIELD | -| index_file_path | String | | YES | | FIELD | -| index_file_size | UInt64 | | YES | | FIELD | -| num_rows | UInt64 | | NO | | FIELD | -| num_row_groups | UInt64 | | NO | | FIELD | -| num_series | UInt64 | | YES | | FIELD | -| min_ts | TimestampNanosecond | | YES | | FIELD | -| max_ts | TimestampNanosecond | | YES | | FIELD | -| sequence | UInt64 | | YES | | FIELD | -| origin_region_id | UInt64 | | NO | | FIELD | -| node_id | UInt64 | | YES | | FIELD | -| visible | Boolean | | NO | | FIELD | -| primary_key_min | Binary | | YES | | FIELD | -| primary_key_max | Binary | | YES | | FIELD | -+------------------+---------------------+-----+------+---------+---------------+ ++---------------------------------+---------------------+-----+------+---------+---------------+ +| Column | Type | Key | Null | Default | Semantic Type | ++---------------------------------+---------------------+-----+------+---------+---------------+ +| table_dir | String | | NO | | FIELD | +| region_id | UInt64 | | NO | | FIELD | +| table_id | UInt32 | | NO | | FIELD | +| region_number | UInt32 | | NO | | FIELD | +| region_group | UInt8 | | NO | | FIELD | +| region_sequence | UInt32 | | NO | | FIELD | +| file_id | String | | NO | | FIELD | +| index_version | UInt64 | | NO | | FIELD | +| level | UInt8 | | NO | | FIELD | +| file_path | String | | NO | | FIELD | +| file_size | UInt64 | | NO | | FIELD | +| index_file_path | String | | YES | | FIELD | +| index_file_size | UInt64 | | YES | | FIELD | +| num_rows | UInt64 | | NO | | FIELD | +| num_row_groups | UInt64 | | NO | | FIELD | +| num_series | UInt64 | | YES | | FIELD | +| min_ts | TimestampNanosecond | | YES | | FIELD | +| max_ts | TimestampNanosecond | | YES | | FIELD | +| sequence | UInt64 | | YES | | FIELD | +| origin_region_id | UInt64 | | NO | | FIELD | +| node_id | UInt64 | | YES | | FIELD | +| visible | Boolean | | NO | | FIELD | +| primary_key_min | Binary | | YES | | FIELD | +| primary_key_max | Binary | | YES | | FIELD | +| max_row_group_uncompressed_size | UInt64 | | NO | | FIELD | +| partition_expr | String | | YES | | FIELD | ++---------------------------------+---------------------+-----+------+---------+---------------+ DESC TABLE information_schema.ssts_storage; @@ -99,13 +101,13 @@ ADMIN FLUSH_TABLE('sst_case'); -- SQLNESS REPLACE (/public/\d+) /public/ SELECT * FROM information_schema.ssts_manifest order by file_path; -+----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+ -| table_dir | region_id | table_id | region_number | region_group | region_sequence | file_id | index_version | level | file_path | file_size | index_file_path | index_file_size | num_rows | num_row_groups | num_series | min_ts | max_ts | sequence | origin_region_id | node_id | visible | primary_key_min | primary_key_max | -+----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+ -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800001f4 | 01800001f4 | -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800005dc | 01800005dc | -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800009c4 | 01800009c4 | -+----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+ ++----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+---------------------------------+------------------------+ +| table_dir | region_id | table_id | region_number | region_group | region_sequence | file_id | index_version | level | file_path | file_size | index_file_path | index_file_size | num_rows | num_row_groups | num_series | min_ts | max_ts | sequence | origin_region_id | node_id | visible | primary_key_min | primary_key_max | max_row_group_uncompressed_size | partition_expr | ++----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+---------------------------------+------------------------+ +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800001f4 | 01800001f4 || a <| +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800005dc | 01800005dc || a >=AND a <| +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800009c4 | 01800009c4 || a >=| ++----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+---------------------------------+------------------------+ -- SQLNESS REPLACE (\s+\d+\s+) -- SQLNESS REPLACE ([0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}) @@ -168,15 +170,15 @@ ADMIN FLUSH_TABLE('sst_case'); -- SQLNESS REPLACE (/public/\d+) /public/ SELECT * FROM information_schema.ssts_manifest order by region_id, sequence; -+----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+ -| table_dir | region_id | table_id | region_number | region_group | region_sequence | file_id | index_version | level | file_path | file_size | index_file_path | index_file_size | num_rows | num_row_groups | num_series | min_ts | max_ts | sequence | origin_region_id | node_id | visible | primary_key_min | primary_key_max | -+----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+ -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800001f4 | 01800001f4 | -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true || 018000007c | -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800005dc | 01800005dc | -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true ||| -| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800009c4 | 01800009c4 | -+----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+ ++----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+---------------------------------+------------------------+ +| table_dir | region_id | table_id | region_number | region_group | region_sequence | file_id | index_version | level | file_path | file_size | index_file_path | index_file_size | num_rows | num_row_groups | num_series | min_ts | max_ts | sequence | origin_region_id | node_id | visible | primary_key_min | primary_key_max | max_row_group_uncompressed_size | partition_expr | ++----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+---------------------------------+------------------------+ +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800001f4 | 01800001f4 || a <| +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true || 018000007c || a <| +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800005dc | 01800005dc || a >=AND a <| +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true |||| a >=AND a <| +| data/greptime/public// |||||| ||| data/greptime/public//_/.parquet || data/greptime/public//_/index/.puffin ||||| | |||| true | 01800009c4 | 01800009c4 || a >=| ++----------------------------+---------------+----------+---------------+--------------+-----------------+--------------------------------------+---------------+-------+----------------------------------------------------------------------------------------+-----------+---------------------------------------------------------------------------------------------+-----------------+----------+----------------+------------+-------------------------+-------------------------+----------+------------------+---------+---------+-----------------+-----------------+---------------------------------+------------------------+ -- SQLNESS REPLACE (\s+\d+\s+) -- SQLNESS REPLACE ([0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}) diff --git a/tests/cases/standalone/common/system/information_schema.result b/tests/cases/standalone/common/system/information_schema.result index e05cd064e1d..17e1f0ef80d 100644 --- a/tests/cases/standalone/common/system/information_schema.result +++ b/tests/cases/standalone/common/system/information_schema.result @@ -436,6 +436,7 @@ order by table_schema, table_name, column_name; | greptime | information_schema | ssts_manifest | index_file_size | 13 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | YES | bigint unsigned | | | | greptime | information_schema | ssts_manifest | index_version | 8 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | NO | bigint unsigned | | | | greptime | information_schema | ssts_manifest | level | 9 | | | 3 | 0 | | | | | | select,insert | | UInt8 | tinyint unsigned | FIELD | | NO | tinyint unsigned | | | +| greptime | information_schema | ssts_manifest | max_row_group_uncompressed_size | 25 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | NO | bigint unsigned | | | | greptime | information_schema | ssts_manifest | max_ts | 18 | | | | | 9 | | | | | select,insert | | TimestampNanosecond | timestamp(9) | FIELD | | YES | timestamp(9) | | | | greptime | information_schema | ssts_manifest | min_ts | 17 | | | | | 9 | | | | | select,insert | | TimestampNanosecond | timestamp(9) | FIELD | | YES | timestamp(9) | | | | greptime | information_schema | ssts_manifest | node_id | 21 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | YES | bigint unsigned | | | @@ -443,6 +444,7 @@ order by table_schema, table_name, column_name; | greptime | information_schema | ssts_manifest | num_rows | 14 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | NO | bigint unsigned | | | | greptime | information_schema | ssts_manifest | num_series | 16 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | YES | bigint unsigned | | | | greptime | information_schema | ssts_manifest | origin_region_id | 20 | | | 20 | 0 | | | | | | select,insert | | UInt64 | bigint unsigned | FIELD | | NO | bigint unsigned | | | +| greptime | information_schema | ssts_manifest | partition_expr | 26 | 2147483647 | 2147483647 | | | | utf8 | utf8_bin | | | select,insert | | String | string | FIELD | | YES | string | | | | greptime | information_schema | ssts_manifest | primary_key_max | 24 | | | | | | | | | | select,insert | | Binary | varbinary | FIELD | | YES | varbinary | | | | greptime | information_schema | ssts_manifest | primary_key_min | 23 | | | | | | | | | | select,insert | | Binary | varbinary | FIELD | | YES | varbinary | | | | greptime | information_schema | ssts_manifest | region_group | 5 | | | 3 | 0 | | | | | | select,insert | | UInt8 | tinyint unsigned | FIELD | | NO | tinyint unsigned | | | From c01de4afdc7e2ef65edb5efdfaccebfd62b0d1c6 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 28 Aug 2026 08:49:49 +0000 Subject: [PATCH 16/22] fix(operator): whitelist private system table auto create (#8930) Signed-off-by: WenyXu --- src/common/meta/src/datanode.rs | 2 + .../src/handler/persist_stats_handler.rs | 8 +- src/operator/src/insert.rs | 68 ++++++-- tests-integration/src/cluster.rs | 14 +- tests-integration/src/standalone.rs | 14 ++ tests-integration/tests/database_ddl_event.rs | 148 ++++++++++++++++- tests-integration/tests/grpc.rs | 155 +++++++++++++++++- 7 files changed, 385 insertions(+), 24 deletions(-) diff --git a/src/common/meta/src/datanode.rs b/src/common/meta/src/datanode.rs index 8af86462129..d50c37c381b 100644 --- a/src/common/meta/src/datanode.rs +++ b/src/common/meta/src/datanode.rs @@ -31,6 +31,8 @@ use crate::heartbeat::utils::get_datanode_workloads; const DATANODE_STAT_PREFIX: &str = "__meta_datanode_stat"; pub const REGION_STATISTIC_KEY: &str = "__region_statistic"; +/// The private table name for persisted region statistics. +pub const REGION_STATS_HISTORY_TABLE_NAME: &str = "region_statistics_history"; lazy_static! { pub(crate) static ref DATANODE_LEASE_KEY_PATTERN: Regex = diff --git a/src/meta-srv/src/handler/persist_stats_handler.rs b/src/meta-srv/src/handler/persist_stats_handler.rs index b1378c39e6b..fcdb9fa04bc 100644 --- a/src/meta-srv/src/handler/persist_stats_handler.rs +++ b/src/meta-srv/src/handler/persist_stats_handler.rs @@ -22,7 +22,7 @@ use client::inserter::{Context as InserterContext, Inserter}; use common_catalog::consts::DEFAULT_PRIVATE_SCHEMA_NAME; use common_macro::{Schema, ToRow}; use common_meta::DatanodeId; -use common_meta::datanode::RegionStat; +use common_meta::datanode::{REGION_STATS_HISTORY_TABLE_NAME, RegionStat}; use common_telemetry::warn; use dashmap::DashMap; use store_api::region_engine::RegionRole; @@ -40,8 +40,6 @@ pub struct PersistStatsHandler { persist_interval: Duration, } -/// The name of the table to persist region stats. -const META_REGION_STATS_HISTORY_TABLE_NAME: &str = "region_statistics_history"; /// The default context to persist region stats. const DEFAULT_CONTEXT: InserterContext = InserterContext { catalog: DEFAULT_CATALOG_NAME, @@ -207,7 +205,7 @@ impl PersistStatsHandler { &DEFAULT_CONTEXT, RowInsertRequests { inserts: vec![RowInsertRequest { - table_name: META_REGION_STATS_HISTORY_TABLE_NAME.to_string(), + table_name: REGION_STATS_HISTORY_TABLE_NAME.to_string(), rows: Some(Rows { schema: PersistRegionStat::schema(), rows, @@ -537,7 +535,7 @@ mod tests { }; assert_eq!( request.table_name, - META_REGION_STATS_HISTORY_TABLE_NAME.to_string() + REGION_STATS_HISTORY_TABLE_NAME.to_string() ); assert_eq!(request.rows.unwrap().rows, vec![expected_row]); diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index ead517ab80f..52a73892dc0 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -29,12 +29,15 @@ use api::v1::{ use catalog::CatalogManagerRef; use client::{OutputData, OutputMeta}; use common_catalog::consts::{ - PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN, TRACE_ID_COLUMN, TRACE_TABLE_NAME, - TRACE_TABLE_NAME_SESSION_KEY, default_engine, is_ddl_reserved_table, + DEFAULT_PRIVATE_SCHEMA_NAME, PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN, TRACE_ID_COLUMN, + TRACE_TABLE_NAME, TRACE_TABLE_NAME_SESSION_KEY, default_engine, is_ddl_reserved_table, trace_operations_table_name, trace_services_table_name, }; +use common_event_recorder::DEFAULT_EVENTS_TABLE_NAME; +use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME; use common_grpc_expr::util::ColumnExpr; use common_meta::cache::TableFlownodeSetCacheRef; +use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME; use common_meta::node_manager::{AffectedRows, NodeManagerRef}; use common_meta::peer::Peer; use common_meta::rpc::ddl::TriggerReason; @@ -490,9 +493,9 @@ impl Inserter { Ok(inserts) } - /// Returns `None` if auto table creation is allowed, or `Some(reason)` if - /// disabled by either the global config or the request hint. The reason tells - /// which one, for a clearer error. + /// Returns `Some(reason)` if the config or request hint disables automatic + /// table creation. Exempt private system tables are handled by + /// [`Self::is_auto_create_exempt_private_table`]. fn auto_create_disabled_reason(&self, ctx: &QueryContextRef) -> Result> { let auto_create_table_hint = ctx .extension(AUTO_CREATE_TABLE_KEY) @@ -514,6 +517,16 @@ impl Inserter { }) } + /// Returns whether a private system table may infer and reconcile its schema + /// even when automatic table creation is disabled. + fn is_auto_create_exempt_private_table(schema: &str, table: &str) -> bool { + schema == DEFAULT_PRIVATE_SCHEMA_NAME + && matches!( + table, + DEFAULT_EVENTS_TABLE_NAME | SLOW_QUERY_TABLE_NAME | REGION_STATS_HISTORY_TABLE_NAME + ) + } + /// Ensures a trace table has the request-global schema without requiring a /// padded data row to drive on-demand creation or alteration. When /// `alter_existing` is false, a table created after planning is left for the @@ -571,12 +584,21 @@ impl Inserter { let _timer = crate::metrics::CREATE_ALTER_ON_DEMAND .with_label_values(&[auto_create_table_type.as_str()]) .start_timer(); - let catalog = ctx.current_catalog(); let schema = ctx.current_schema(); + let auto_create_disabled_reason = self.auto_create_disabled_reason(ctx)?; + // Enabled batches permit every table, so only disabled batches need a whitelist scan. + let has_auto_create_exempt_table = auto_create_disabled_reason.is_some() + && requests + .inserts + .iter() + .any(|req| Self::is_auto_create_exempt_private_table(&schema, &req.table_name)); let mut table_infos = HashMap::new(); - if let Some(disabled_reason) = self.auto_create_disabled_reason(ctx)? { + // Without exempt tables, verify existing tables and reject missing ones without inferring schemas. + if let Some(disabled_reason) = auto_create_disabled_reason + && !has_auto_create_exempt_table + { let mut instant_table_ids = HashSet::new(); for req in &requests.inserts { let table = match self.get_table(catalog, &schema, &req.table_name).await? { @@ -618,20 +640,25 @@ impl Inserter { let mut per_table_semantics: Option> = None; for req in &mut requests.inserts { + // Mixed batches need a per-table decision so an exempt table cannot authorize others. + let auto_create_allowed = auto_create_disabled_reason.is_none() + || Self::is_auto_create_exempt_private_table(&schema, &req.table_name); match self.get_table(catalog, &schema, &req.table_name).await? { Some(table) => { let table_info = table.table_info(); if table_info.is_ttl_instant_table() { instant_table_ids.insert(table_info.table_id()); } - if let Some(alter_expr) = self.get_alter_table_expr_on_demand( - req, - &table, - ctx, - accommodate_existing_schema, - is_single_value, - auto_create_table_type.alter_existing(), - )? { + if auto_create_allowed + && let Some(alter_expr) = self.get_alter_table_expr_on_demand( + req, + &table, + ctx, + accommodate_existing_schema, + is_single_value, + auto_create_table_type.alter_existing(), + )? + { alter_tables.push(alter_expr); need_refresh_table_infos.insert(( catalog.to_string(), @@ -655,6 +682,17 @@ impl Inserter { } table_infos.insert(table_info.table_id(), table_info); } + None if !auto_create_allowed + && let Some(disabled_reason) = auto_create_disabled_reason => + { + return InvalidInsertRequestSnafu { + reason: format!( + "Table `{}` does not exist, and {}", + req.table_name, disabled_reason, + ), + } + .fail(); + } None => { let semantic_index = per_table_semantics .get_or_insert_with(|| parse_per_table_semantic_index(ctx)) diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index d0cd2e584c3..e40c707fe98 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -169,6 +169,7 @@ pub struct GreptimeDbClusterBuilder { metasrv_wal_config: MetasrvWalConfig, datanode_gc_config: GcConfig, metasrv_gc_config: GcSchedulerOptions, + frontend_auto_create_table: bool, shared_home_dir: Option>, meta_selector: Option, } @@ -202,6 +203,7 @@ impl GreptimeDbClusterBuilder { metasrv_wal_config: MetasrvWalConfig::default(), datanode_gc_config: GcConfig::default(), metasrv_gc_config: GcSchedulerOptions::default(), + frontend_auto_create_table: true, shared_home_dir: None, meta_selector: None, } @@ -248,6 +250,13 @@ impl GreptimeDbClusterBuilder { self } + /// Sets whether the test frontend automatically creates tables on write. + #[must_use] + pub fn with_frontend_auto_create_table(mut self, auto_create_table: bool) -> Self { + self.frontend_auto_create_table = auto_create_table; + self + } + #[must_use] pub fn with_shared_home_dir(mut self, shared_home_dir: Arc) -> Self { self.shared_home_dir = Some(shared_home_dir); @@ -529,7 +538,10 @@ impl GreptimeDbClusterBuilder { } fn build_frontend_options(&self) -> FrontendOptions { - let mut fe_opts = FrontendOptions::default(); + let mut fe_opts = FrontendOptions { + auto_create_table: self.frontend_auto_create_table, + ..Default::default() + }; // Choose a random unused port between [14000, 24000] for local test to avoid conflicts. let port_range = 14000..=24000; diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 16fc2e65634..0704b184e89 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -27,6 +27,7 @@ use common_base::Plugins; use common_catalog::consts::{MIN_USER_FLOW_ID, MIN_USER_TABLE_ID}; use common_config::KvBackendConfig; use common_datasource::object_store::LocalFileAccess; +use common_event_recorder::EventRecorderOptions; use common_meta::cache::LayeredCacheRegistryBuilder; use common_meta::ddl::flow_meta::FlowMetadataAllocator; use common_meta::ddl::table_meta::TableMetadataAllocator; @@ -84,6 +85,7 @@ pub struct GreptimeDbStandaloneBuilder { default_store: Option, plugin: Option, slow_query_options: SlowQueryOptions, + event_recorder_options: EventRecorderOptions, auto_create_table: bool, } @@ -102,6 +104,7 @@ impl GreptimeDbStandaloneBuilder { threshold: Duration::from_secs(1), ..Default::default() }, + event_recorder_options: EventRecorderOptions::default(), auto_create_table: true, } } @@ -112,6 +115,16 @@ impl GreptimeDbStandaloneBuilder { self } + /// Sets the event recorder options for the standalone instance. + #[must_use] + pub fn with_event_recorder_options( + mut self, + event_recorder_options: EventRecorderOptions, + ) -> Self { + self.event_recorder_options = event_recorder_options; + self + } + #[must_use] pub fn with_default_store_type(self, store_type: StorageType) -> Self { Self { @@ -368,6 +381,7 @@ impl GreptimeDbStandaloneBuilder { wal: self.metasrv_wal_config.clone().into(), grpc: GrpcOptions::default().with_server_addr("127.0.0.1:4001"), slow_query: self.slow_query_options.clone(), + event_recorder: self.event_recorder_options.clone(), auto_create_table: self.auto_create_table, // Tests cover the descriptor, so they run with it enabled. otlp: frontend::service_config::OtlpOptions { diff --git a/tests-integration/tests/database_ddl_event.rs b/tests-integration/tests/database_ddl_event.rs index 4dfda4f32b9..c08b81211de 100644 --- a/tests-integration/tests/database_ddl_event.rs +++ b/tests-integration/tests/database_ddl_event.rs @@ -15,20 +15,163 @@ use std::sync::Arc; use client::{Database, OutputData}; +use common_event_recorder::{EventRecorderOptions, EventTypeFilter}; use common_test_util::temp_dir::create_temp_dir; use tests_integration::cluster::GreptimeDbClusterBuilder; use tests_integration::standalone::GreptimeDbStandaloneBuilder; use tests_integration::test_util::{ - StorageType, get_test_store_config, setup_authenticated_grpc_database, + StorageType, execute_sql, get_test_store_config, setup_authenticated_grpc_database, }; use crate::event_recorder_test_util::{ - assert_procedure_actor, assert_single_event, find_eventually_string, + assert_eventually_eq, assert_procedure_actor, assert_single_event, find_eventually_string, }; const DATABASE_NAME: &str = "database_ddl_events"; const PROCEDURE_ACTOR: &str = "procedure_actor"; const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; +const CREATION_DATABASE_NAME: &str = "event_schema_creation"; +const RECONCILIATION_DATABASE_NAME: &str = "event_schema_reconciliation"; + +#[tokio::test(flavor = "multi_thread")] +async fn test_event_table_auto_creation_with_auto_create_disabled() { + common_telemetry::init_default_ut_logging(); + let standalone = GreptimeDbStandaloneBuilder::new("test_event_table_auto_creation") + .with_auto_create_table(false) + .with_event_recorder_options(EventRecorderOptions { + event_types: Arc::new(EventTypeFilter::Only( + [String::from("create_database")].into_iter().collect(), + )), + ..Default::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + + let (database, _grpc_server) = setup_authenticated_grpc_database( + instance.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, + ) + .await; + database + .sql(format!("CREATE DATABASE {CREATION_DATABASE_NAME}")) + .await + .unwrap(); + + assert_eventually_eq( + instance, + "SELECT count(*) AS events_tables \ + FROM information_schema.tables \ + WHERE table_catalog = 'greptime' \ + AND table_schema = 'greptime_private' \ + AND table_name = 'events'", + "+---------------+\n| events_tables |\n+---------------+\n| 1 |\n+---------------+", + ) + .await; + + let procedure_id = find_eventually_string( + instance, + &format!( + "SELECT procedure_id FROM greptime_private.events \ + WHERE type = 'create_database' \ + AND schema_name = '{CREATION_DATABASE_NAME}' \ + AND json_path_match(procedure_trigger, '$.type == \"Submitted\"') \ + ORDER BY timestamp DESC LIMIT 1" + ), + "procedure_id", + ) + .await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_event_table_schema_reconciliation_with_auto_create_disabled() { + common_telemetry::init_default_ut_logging(); + let standalone = GreptimeDbStandaloneBuilder::new("test_event_table_schema_reconciliation") + .with_auto_create_table(false) + .with_event_recorder_options(EventRecorderOptions { + event_types: Arc::new(EventTypeFilter::Only( + [String::from("create_database")].into_iter().collect(), + )), + ..Default::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + + // Matches the pre-actor procedure-event schema. The recorder filter keeps + // this setup DDL from creating an event before the regression is exercised. + execute_sql( + instance, + r#" + CREATE TABLE greptime_private.events ( + "type" STRING, + payload JSON, + "timestamp" TIMESTAMP(9) NOT NULL, + procedure_id STRING, + procedure_state STRING, + procedure_error STRING, + procedure_trigger JSON, + catalog_name STRING, + schema_name STRING, + event_context JSON, + TIME INDEX ("timestamp"), + PRIMARY KEY ("type") + ) WITH (append_mode = 'true') + "#, + ) + .await; + + assert_eventually_eq( + instance, + "SELECT count(*) AS actor_columns \ + FROM information_schema.columns \ + WHERE table_catalog = 'greptime' \ + AND table_schema = 'greptime_private' \ + AND table_name = 'events' \ + AND column_name = 'actor'", + "+---------------+\n| actor_columns |\n+---------------+\n| 0 |\n+---------------+", + ) + .await; + + let (database, _grpc_server) = setup_authenticated_grpc_database( + instance.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, + ) + .await; + database + .sql(format!("CREATE DATABASE {RECONCILIATION_DATABASE_NAME}")) + .await + .unwrap(); + + assert_eventually_eq( + instance, + "SELECT count(*) AS actor_columns \ + FROM information_schema.columns \ + WHERE table_catalog = 'greptime' \ + AND table_schema = 'greptime_private' \ + AND table_name = 'events' \ + AND column_name = 'actor'", + "+---------------+\n| actor_columns |\n+---------------+\n| 1 |\n+---------------+", + ) + .await; + + let procedure_id = find_eventually_string( + instance, + &format!( + "SELECT procedure_id FROM greptime_private.events \ + WHERE type = 'create_database' \ + AND schema_name = '{RECONCILIATION_DATABASE_NAME}' \ + AND json_path_match(procedure_trigger, '$.type == \"Submitted\"') \ + ORDER BY timestamp DESC LIMIT 1" + ), + "procedure_id", + ) + .await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; +} #[tokio::test(flavor = "multi_thread")] async fn test_database_ddl_events() { @@ -43,6 +186,7 @@ async fn test_database_ddl_events() { let cluster = GreptimeDbClusterBuilder::new("test_database_ddl_events") .await .with_datanodes(1) + .with_frontend_auto_create_table(false) .with_store_config(store_config) .with_shared_home_dir(Arc::new(home_dir)) .build(true) diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index 737b4fb6c1d..f06ffa65224 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -26,9 +26,12 @@ use api::v1::{ use auth::user_provider_from_option; use base64::prelude::{BASE64_STANDARD, Engine as _}; use client::{Client, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, Database, OutputData}; -use common_catalog::consts::MITO_ENGINE; +use common_catalog::consts::{DEFAULT_PRIVATE_SCHEMA_NAME, MITO_ENGINE}; +use common_event_recorder::DEFAULT_EVENTS_TABLE_NAME; +use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME; use common_grpc::channel_manager::ClientTlsOption; use common_memory_manager::OnExhaustedPolicy; +use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME; use common_query::Output; use common_recordbatch::RecordBatches; use common_runtime::Runtime; @@ -98,6 +101,8 @@ macro_rules! grpc_tests { test_auto_create_table, test_auto_create_table_with_hints, test_auto_create_table_disabled_by_config, + test_private_system_tables_auto_create_table_with_global_disabled, + test_private_system_tables_bypass_auto_create_hint, test_otel_arrow_auth, test_otel_arrow_exponential_histogram, test_insert_and_select, @@ -679,6 +684,154 @@ pub async fn test_auto_create_table_disabled_by_config(store_type: StorageType) let _ = fe_grpc_server.shutdown().await; } +pub async fn test_private_system_tables_auto_create_table_with_global_disabled( + store_type: StorageType, +) { + let (_db, fe_grpc_server) = setup_grpc_server_with_auto_create_table_disabled( + store_type, + "test_private_system_tables_auto_create_table_with_global_disabled", + ) + .await; + let addr = fe_grpc_server.bind_addr().unwrap().to_string(); + + let grpc_client = Client::with_urls(vec![addr]); + let db = Database::new( + DEFAULT_CATALOG_NAME, + DEFAULT_PRIVATE_SCHEMA_NAME, + grpc_client, + ); + let (host, cpu, mem, ts) = expect_data(); + + for table_name in [ + DEFAULT_EVENTS_TABLE_NAME, + SLOW_QUERY_TABLE_NAME, + REGION_STATS_HISTORY_TABLE_NAME, + ] { + let result = db + .insert(InsertRequests { + inserts: vec![InsertRequest { + table_name: table_name.to_string(), + columns: vec![host.clone(), cpu.clone(), mem.clone(), ts.clone()], + row_count: 4, + }], + }) + .await; + assert_eq!(result.unwrap(), 4); + } + + let load = Column { + column_name: "load".to_string(), + values: Some(column::Values { + f64_values: vec![0.4, 0.5, 0.6, 0.7], + ..Default::default() + }), + semantic_type: SemanticType::Field as i32, + datatype: ColumnDataType::Float64 as i32, + ..Default::default() + }; + let result = db + .insert(InsertRequests { + inserts: vec![InsertRequest { + table_name: DEFAULT_EVENTS_TABLE_NAME.to_string(), + columns: vec![host, cpu, mem, ts, load], + row_count: 4, + }], + }) + .await; + assert_eq!(result.unwrap(), 4); + + let output = db + .sql(format!("SHOW CREATE TABLE {DEFAULT_EVENTS_TABLE_NAME}")) + .await + .unwrap(); + let record_batches = match output.data { + OutputData::RecordBatches(record_batches) => record_batches, + OutputData::Stream(stream) => RecordBatches::try_collect(stream).await.unwrap(), + OutputData::AffectedRows(_) => unreachable!(), + }; + assert!(record_batches.pretty_print().unwrap().contains("\"load\"")); + + let _ = fe_grpc_server.shutdown().await; +} + +pub async fn test_private_system_tables_bypass_auto_create_hint(store_type: StorageType) { + let (_db, fe_grpc_server) = setup_grpc_server( + store_type, + "test_private_system_tables_bypass_auto_create_hint", + ) + .await; + let addr = fe_grpc_server.bind_addr().unwrap().to_string(); + + let grpc_client = Client::with_urls(vec![addr]); + let db = Database::new( + DEFAULT_CATALOG_NAME, + DEFAULT_PRIVATE_SCHEMA_NAME, + grpc_client, + ); + let (host, cpu, mem, ts) = expect_data(); + + for table_name in [ + DEFAULT_EVENTS_TABLE_NAME, + SLOW_QUERY_TABLE_NAME, + REGION_STATS_HISTORY_TABLE_NAME, + ] { + let result = db + .insert_with_hints( + InsertRequests { + inserts: vec![InsertRequest { + table_name: table_name.to_string(), + columns: vec![host.clone(), cpu.clone(), mem.clone(), ts.clone()], + row_count: 4, + }], + }, + &[("auto_create_table", "false")], + ) + .await; + assert_eq!(result.unwrap(), 4); + } + + let ordinary_table = "ordinary_private_table"; + let result = db + .insert_with_hints( + InsertRequests { + inserts: vec![ + InsertRequest { + table_name: DEFAULT_EVENTS_TABLE_NAME.to_string(), + columns: vec![host.clone(), cpu.clone(), mem.clone(), ts.clone()], + row_count: 4, + }, + InsertRequest { + table_name: ordinary_table.to_string(), + columns: vec![host, cpu, mem, ts], + row_count: 4, + }, + ], + }, + &[("auto_create_table", "false")], + ) + .await; + let err = result.unwrap_err().to_string(); + assert!( + err.contains(ordinary_table) && err.contains("auto_create_table"), + "unexpected error: {err}" + ); + + let output = db.sql("SHOW TABLES").await.unwrap(); + let record_batches = match output.data { + OutputData::RecordBatches(record_batches) => record_batches, + OutputData::Stream(stream) => RecordBatches::try_collect(stream).await.unwrap(), + OutputData::AffectedRows(_) => unreachable!(), + }; + assert!( + !record_batches + .pretty_print() + .unwrap() + .contains(ordinary_table) + ); + + let _ = fe_grpc_server.shutdown().await; +} + fn expect_data() -> (Column, Column, Column, Column) { // testing data: let expected_host_col = Column { From c0b8612c41bf0b2e7e46d76c5251777c0a68057d Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Fri, 28 Aug 2026 13:45:06 +0000 Subject: [PATCH 17/22] fix(cmd): configure meta client in frontend plugin test (#8977) Provide valid distributed frontend options in `src/cmd/src/frontend.rs` so enterprise plugin setup can preserve prefilled heartbeat extensions without weakening production validation. Signed-off-by: Lei, HUANG --- src/cmd/src/frontend.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/cmd/src/frontend.rs b/src/cmd/src/frontend.rs index 1f95a3c54b4..6cae195f3e6 100644 --- a/src/cmd/src/frontend.rs +++ b/src/cmd/src/frontend.rs @@ -685,10 +685,14 @@ mod tests { let extensions = FrontendHeartbeatExtensions::default(); let shared_extensions = extensions.clone(); + let options = frontend::frontend::FrontendOptions { + meta_client: Some(MetaClientOptions::default()), + ..Default::default() + }; let plugins = prepare_frontend_plugins( plugins_with_heartbeat_extensions(extensions), &[], - &frontend::frontend::FrontendOptions::default(), + &options, None, ) .await From 9cfbc42126704c2cc8c358e85d89cec9482c9a9f Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Sat, 29 Aug 2026 06:51:48 +0000 Subject: [PATCH 18/22] refactor(mito2): prioritize file count in TWCS picker (#8765) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor(mito2): prioritize file count in TWCS picker Signed-off-by: Lei, HUANG * test(mito2): migrate TWCS tests to async picker API Signed-off-by: Lei, HUANG * refactor(mito2): remove legacy reduce_runs and merge_seq_files pickers Signed-off-by: Lei, HUANG * feat(mito2): balance TWCS picks by file group Signed-off-by: Lei, HUANG * refactor(mito2): treat multi-SST groups as barriers in TWCS pick_count_first Previously pick_count_first filtered out multi-SST file groups and could pick singleton groups across them in one interval. Now multi-SST groups split the candidates into independent segments, so a single pick never crosses a multi-SST group. Signed-off-by: Lei, HUANG * refactor(mito2): remove redundant overlaps_files_left_behind check The rebase onto main kept both the upstream fix (#8872) and the branch's selected_overlaps_unselected check for the same deletion-marker problem. The upstream check is fully subsumed: files_to_merge only differs from the window files in append mode, where filter_deleted is already false, and selected_overlaps_unselected treats partially-selected groups as unselected, making it strictly stronger. Signed-off-by: Lei, HUANG * perf(mito2): pre-filter by selection time span in selected_overlaps_unselected A window group outside the overall time span of the selected groups cannot overlap any of them, so skip the precise overlap check (and the per-group file id set lookup) for it. In typical TWCS windows the pick is clustered in one segment, so most unselected groups are rejected in O(1). Signed-off-by: Lei, HUANG * refactor(mito2): remove FileGroup abstraction from TWCS picker Signed-off-by: Lei, HUANG * feat(mito2): drain compactable backlog after successful compaction Signed-off-by: Lei, HUANG * refactor(mito2): score TWCS candidates by predicted progress Redefine what a pick candidate is worth: an interval is only eligible if compacting it makes progress on at least one axis — it reduces the physical file count given the max_output_file_size split threshold (predicted output = ceil(input bytes / threshold)), or it resolves at least one overlap between sorted runs. A pure rewrite that achieves neither (e.g. 32 large balanced files whose output would split back into just as many SSTs) is skipped instead of burning I/O. The candidate metrics are accumulated in a Candidate struct as the interval expands, and the score ranks by predicted file reduction, then overlap participants, then smaller input bytes. With no output size limit the behavior degenerates to the previous count-first rule. Signed-off-by: Lei, HUANG * fix(mito2): make TWCS input limit configurable Allow operators to tune the maximum SST inputs through a hidden environment variable while keeping 32 as the validated default. Signed-off-by: Lei, HUANG * fix(mito2): apply TWCS trigger at window level Check the total physical SST count before candidate selection so trigger values above the per-task input limit still allow bounded compactions. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/mito2/benches/bench_compaction_picker.rs | 66 -- src/mito2/src/compaction/run.rs | 481 ++--------- src/mito2/src/compaction/scheduler.rs | 27 +- src/mito2/src/compaction/scheduler_test.rs | 176 +++- src/mito2/src/compaction/twcs.rs | 863 +++++++++++++++---- src/mito2/src/engine/compaction_test.rs | 25 +- src/mito2/src/engine/merge_mode_test.rs | 3 +- src/mito2/src/gc/worker_test.rs | 10 +- src/mito2/src/worker/handle_compaction.rs | 6 + 9 files changed, 981 insertions(+), 676 deletions(-) diff --git a/src/mito2/benches/bench_compaction_picker.rs b/src/mito2/benches/bench_compaction_picker.rs index 74142127aee..439b9e6f226 100644 --- a/src/mito2/benches/bench_compaction_picker.rs +++ b/src/mito2/benches/bench_compaction_picker.rs @@ -17,7 +17,6 @@ use std::hint::black_box; use criterion::{Criterion, criterion_group, criterion_main}; use mito2::compaction::run::{ Item, Ranged, SortedRun, find_overlapping_items, find_sorted_runs, find_sorted_runs_original, - merge_seq_files, reduce_runs, }; #[derive(Clone, Debug, Eq, Hash, PartialEq)] @@ -49,20 +48,6 @@ impl Item for MockFile { } } -fn generate_test_files(n: usize) -> Vec { - let mut files = Vec::with_capacity(n); - for _ in 0..n { - // Create slightly overlapping ranges to force multiple sorted runs - files.push(MockFile { - start: 0, - end: 10, - pk: 0, - size: 10, - }); - } - files -} - fn generate_same_timestamp_files(total_files: usize, files_per_timestamp: usize) -> Vec { const TIMESTAMP_INTERVAL_SECS: i64 = 10 * 60; @@ -107,21 +92,6 @@ fn bench_find_sorted_runs(c: &mut Criterion) { group.finish(); } -fn bench_reduce_runs(c: &mut Criterion) { - let mut group = c.benchmark_group("reduce_runs"); - - for size in [10, 100, 1000].iter() { - group.bench_function(format!("size_{}", size), |b| { - let mut files = generate_test_files(*size); - let runs = find_sorted_runs(&mut files); - b.iter(|| { - reduce_runs(black_box(runs.clone())); - }); - }); - } - group.finish(); -} - fn bench_find_overlapping_items(c: &mut Criterion) { let mut group = c.benchmark_group("find_overlapping_items"); @@ -158,45 +128,9 @@ fn bench_find_overlapping_items(c: &mut Criterion) { group.finish(); } -fn bench_merge_seq_files(c: &mut Criterion) { - let mut group = c.benchmark_group("merge_seq_files"); - - for size in [10, 100, 1000].iter() { - group.bench_function(format!("size_{}", size), |b| { - // Create a set of files with varying sizes - let mut files = Vec::with_capacity(*size); - - for i in 0..*size { - // Create files with different sizes to test the scoring algorithm - let file_size = if i % 3 == 0 { - 5 - } else if i % 3 == 1 { - 10 - } else { - 15 - }; - - files.push(MockFile { - start: i as i64, - end: (i + 1) as i64, - pk: 0, - size: file_size, - }); - } - - b.iter(|| { - merge_seq_files(black_box(&files), black_box(Some(50))); - }); - }); - } - group.finish(); -} - criterion_group!( benches, bench_find_sorted_runs, - bench_reduce_runs, bench_find_overlapping_items, - bench_merge_seq_files ); criterion_main!(benches); diff --git a/src/mito2/src/compaction/run.rs b/src/mito2/src/compaction/run.rs index 1ba45690215..5698a6ce667 100644 --- a/src/mito2/src/compaction/run.rs +++ b/src/mito2/src/compaction/run.rs @@ -20,14 +20,9 @@ use std::collections::BinaryHeap; use bytes::{Buf, Bytes}; use common_base::BitVec; -use common_base::readable_size::ReadableSize; use common_time::Timestamp; -use smallvec::{SmallVec, smallvec}; -use crate::sst::file::{FileHandle, RegionFileId}; - -/// Default max compaction output file size when not specified. -const DEFAULT_MAX_OUTPUT_SIZE: u64 = ReadableSize::gb(2).as_bytes(); +use crate::sst::file::FileHandle; /// Trait for any items with specific range (both boundaries are inclusive). pub trait Ranged { @@ -152,71 +147,11 @@ pub trait Item: Ranged + Clone { fn size(&self) -> usize; } -/// A group of files that are created by the same compaction task. -#[derive(Debug, Clone)] -pub struct FileGroup { - files: SmallVec<[FileHandle; 2]>, - size: usize, - num_rows: usize, - min_timestamp: Timestamp, - max_timestamp: Timestamp, - primary_key_range: Option<(Bytes, Bytes)>, -} - -impl FileGroup { - pub(crate) fn new_with_file(file: FileHandle) -> Self { - let size = file.size() as usize; - let (min_timestamp, max_timestamp) = file.time_range(); - let num_rows = file.num_rows(); - let primary_key_range = file.primary_key_range(); - Self { - files: smallvec![file], - size, - num_rows, - min_timestamp, - max_timestamp, - primary_key_range, - } - } - - pub(crate) fn num_rows(&self) -> usize { - self.num_rows - } - - pub(crate) fn add_file(&mut self, file: FileHandle) { - self.size += file.size() as usize; - self.num_rows += file.num_rows(); - let (min_timestamp, max_timestamp) = file.time_range(); - self.min_timestamp = self.min_timestamp.min(min_timestamp); - self.max_timestamp = self.max_timestamp.max(max_timestamp); - self.primary_key_range = - merge_primary_key_ranges(self.primary_key_range.take(), file.primary_key_range()); - self.files.push(file); - } - - pub(crate) fn num_files(&self) -> usize { - self.files.len() - } - - #[cfg(test)] - pub(crate) fn files(&self) -> &[FileHandle] { - &self.files[..] - } - - pub(crate) fn file_ids(&self) -> SmallVec<[RegionFileId; 2]> { - SmallVec::from_iter(self.files.iter().map(|f| f.file_id())) - } - - pub(crate) fn into_files(self) -> impl Iterator { - self.files.into_iter() - } -} - -impl Ranged for FileGroup { +impl Ranged for FileHandle { type BoundType = Timestamp; fn range(&self) -> (Self::BoundType, Self::BoundType) { - (self.min_timestamp, self.max_timestamp) + self.time_range() } fn overlap(&self, other: &Self) -> bool { @@ -226,7 +161,7 @@ impl Ranged for FileGroup { return false; } - match (&self.primary_key_range, &other.primary_key_range) { + match (&self.primary_key_range(), &other.primary_key_range()) { (Some(lhs), Some(rhs)) => primary_key_ranges_overlap(lhs, rhs), _ => true, } @@ -239,16 +174,16 @@ impl Ranged for FileGroup { return false; } - match (&self.primary_key_range, &other.primary_key_range) { + match (&self.primary_key_range(), &other.primary_key_range()) { (Some(lhs), Some(rhs)) => primary_key_ranges_overlap(lhs, rhs), _ => true, } } } -impl Item for FileGroup { +impl Item for FileHandle { fn size(&self) -> usize { - self.size + self.size() as usize } } @@ -553,124 +488,8 @@ where runs.into_iter().map(|x| x.run).collect() } -/// Finds a set of files with minimum penalty to merge that can reduce the total num of runs. -/// The penalty of merging is defined as the size of all overlapping files between two runs. -pub fn reduce_runs(mut runs: Vec>) -> Vec { - assert!(runs.len() > 1); - // sort runs by size - runs.sort_unstable_by_key(|a| a.size); - // limit max probe runs to 100 - let probe_end = runs.len().min(100); - let mut min_penalty = usize::MAX; - let mut files = vec![]; - let mut temp_files = vec![]; - for i in 0..probe_end { - for j in i + 1..probe_end { - let (a, b) = runs.split_at_mut(j); - find_overlapping_items(&mut a[i], &mut b[0], &mut temp_files); - let penalty = temp_files.iter().map(|e| e.size()).sum(); - if penalty < min_penalty { - min_penalty = penalty; - files.clear(); - files.extend_from_slice(&temp_files); - } - } - } - files -} - -/// Finds the optimal set of adjacent files to merge based on a scoring system. -/// -/// This function evaluates all possible contiguous subsets of files to find the best -/// candidates for merging, considering: -/// -/// 1. File reduction - prioritizes merging more files to reduce the total count -/// 2. Write amplification - minimizes the ratio of largest file to total size -/// 3. Size efficiency - prefers merges that utilize available space effectively -/// -/// When multiple merge candidates have the same score, older files (those with lower indices) -/// are preferred. -/// -/// # Arguments -/// * `input_files` - Slice of files to consider for merging -/// * `max_file_size` - Optional maximum size constraint for the merged file. -/// If None, uses 1.5 times the average file size. -/// -/// # Returns -/// A vector containing the best set of adjacent files to merge. -/// Returns an empty vector if input is empty or contains only one file. -pub fn merge_seq_files(input_files: &[T], max_file_size: Option) -> Vec { - if input_files.is_empty() || input_files.len() == 1 { - return vec![]; - } - - // Limit the number of files to process to 100 to control time complexity - let files_to_process = if input_files.len() > 100 { - &input_files[0..100] - } else { - input_files - }; - - // Calculate target size based on max_file_size or average file size - let target_size = match max_file_size { - Some(size) => size as usize, - None => { - // Calculate 1.5*average_file_size if max_file_size is not provided and clamp to 2GB - let total_size: usize = files_to_process.iter().map(|f| f.size()).sum(); - ((((total_size as f64) / (files_to_process.len() as f64)) * 1.5) as usize) - .min(DEFAULT_MAX_OUTPUT_SIZE as usize) - } - }; - - // Find the best group of adjacent files to merge - let mut best_group = Vec::new(); - let mut best_score = f64::NEG_INFINITY; - - // Try different starting positions - iterate from end to start to prefer older files - for start_idx in (0..files_to_process.len()).rev() { - // Try different ending positions - also iterate from end to start - for end_idx in (start_idx + 1..files_to_process.len()).rev() { - let group = &files_to_process[start_idx..=end_idx]; - let total_size: usize = group.iter().map(|f| f.size()).sum(); - - // Skip if total size exceeds target size - if total_size > target_size { - continue; // Use continue instead of break to check smaller ranges - } - - // Calculate amplification factor (largest file size / total size) - let largest_file_size = group.iter().map(|f| f.size()).max().unwrap_or(0); - let amplification_factor = largest_file_size as f64 / total_size as f64; - - // Calculate file reduction (number of files that will be reduced) - let file_reduction = group.len() - 1; - - // Calculate score based on multiple factors: - // 1. File reduction (higher is better) - // 2. Amplification factor (lower is better) - // 3. Size efficiency (how close to target size) - let file_reduction_score = file_reduction as f64 / files_to_process.len() as f64; - let amp_factor_score = (1.0 - amplification_factor) * 1.5; // Lower amplification is better - let size_efficiency = (total_size as f64 / target_size as f64).min(1.0); // Reward using available space - - let score = file_reduction_score + amp_factor_score + size_efficiency; - - // Check if this group is better than our current best - // Use >= instead of > to prefer older files (which we encounter first due to reverse iteration) - if score >= best_score { - best_score = score; - best_group = group.to_vec(); - } - } - } - - best_group -} - #[cfg(test)] mod tests { - use std::collections::HashSet; - use bytes::Bytes; use store_api::storage::FileId; @@ -709,17 +528,6 @@ mod tests { .collect() } - fn build_items_with_size(items: &[(i64, i64, usize)]) -> Vec { - items - .iter() - .map(|(start, end, size)| MockFile { - start: *start, - end: *end, - size: *size, - }) - .collect() - } - fn pk_range(min: &'static [u8], max: &'static [u8]) -> Option<(Bytes, Bytes)> { Some((Bytes::from_static(min), Bytes::from_static(max))) } @@ -844,99 +652,6 @@ mod tests { } } - fn check_reduce_runs( - files: &[(i64, i64)], - expected_runs: &[Vec<(i64, i64)>], - expected: &[(i64, i64)], - ) { - let runs = check_sorted_runs(files, expected_runs); - if runs.len() <= 1 { - assert!(expected.is_empty()); - return; - } - let files_to_merge = reduce_runs(runs); - let file_to_merge_timestamps = files_to_merge - .into_iter() - .map(|f| (f.start, f.end)) - .collect::>(); - - let expected = expected.iter().cloned().collect::>(); - assert_eq!(&expected, &file_to_merge_timestamps); - } - - #[test] - fn test_reduce_runs() { - // [1..3] [5..6] - // [2..4] - check_reduce_runs( - &[(1, 3), (2, 4), (5, 6)], - &[vec![(1, 3), (5, 6)], vec![(2, 4)]], - &[(1, 3), (2, 4)], - ); - - // [1..2][3..5] - // [4..6] - check_reduce_runs( - &[(1, 2), (3, 5), (4, 6)], - &[vec![(1, 2), (3, 5)], vec![(4, 6)]], - &[(3, 5), (4, 6)], - ); - - // [1..2][3..4] [7..8] - // [4..6] - check_reduce_runs( - &[(1, 2), (3, 4), (4, 6), (7, 8)], - &[vec![(1, 2), (3, 4), (4, 6), (7, 8)]], - &[], - ); - - // [1..2][3........6][7..8][8..9] - // [3..4][5..6] - check_reduce_runs( - &[(1, 2), (3, 4), (5, 6), (3, 6), (7, 8), (8, 9)], - &[vec![(1, 2), (3, 6), (7, 8), (8, 9)], vec![(3, 4), (5, 6)]], - &[(5, 6), (3, 4), (3, 6)], // already satisfied - ); - - // [1..2][3........6][7..8][8..9] - // [3..4][5..6] - check_reduce_runs( - &[(1, 2), (3, 4), (5, 6), (3, 6), (7, 8), (8, 9)], - &[vec![(1, 2), (3, 6), (7, 8), (8, 9)], vec![(3, 4), (5, 6)]], - &[(3, 4), (3, 6), (5, 6)], - ); - - // [10..20] [30..40] [50........80][80...100][100..110] - // [50..60] [80..90] - // - check_reduce_runs( - &[ - (10, 20), - (30, 40), - (50, 60), - (50, 80), - (80, 90), - (80, 100), - (100, 110), - ], - &[ - vec![(10, 20), (30, 40), (50, 80), (80, 100), (100, 110)], - vec![(50, 60), (80, 90)], - ], - &[(50, 80), (80, 100), (50, 60), (80, 90)], - ); - - // [0..10] - // [0...11] - // [0....12] - // [0.....13] - check_reduce_runs( - &[(0, 10), (0, 11), (0, 12), (0, 13)], - &[vec![(0, 13)], vec![(0, 12)], vec![(0, 11)], vec![(0, 10)]], - &[(0, 10), (0, 11)], - ); - } - #[test] fn test_find_overlapping_items() { let mut result = Vec::new(); @@ -1038,27 +753,25 @@ mod tests { } #[test] - fn test_file_group_overlap_time_overlap_pk_disjoint() { - let lhs = - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( - FileId::random(), - 0, - 100, - 0, - 1, - 10, - pk_range(b"a", b"f"), - )); - let rhs = - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( - FileId::random(), - 50, - 150, - 0, - 2, - 10, - pk_range(b"x", b"z"), - )); + fn test_file_overlap_time_overlap_pk_disjoint() { + let lhs = new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + 0, + 100, + 0, + 1, + 10, + pk_range(b"a", b"f"), + ); + let rhs = new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + 50, + 150, + 0, + 2, + 10, + pk_range(b"x", b"z"), + ); assert!(!lhs.overlap(&rhs)); } @@ -1066,7 +779,7 @@ mod tests { #[test] fn test_find_sorted_runs_collapses_pk_disjoint_files_into_one_run() { let mut files = vec![ - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( + new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 0, 100, @@ -1074,8 +787,8 @@ mod tests { 1, 10, pk_range(b"a", b"f"), - )), - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( + ), + new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 50, 150, @@ -1083,7 +796,7 @@ mod tests { 2, 10, pk_range(b"x", b"z"), - )), + ), ]; let runs = find_sorted_runs(&mut files); @@ -1095,7 +808,7 @@ mod tests { #[test] fn test_find_sorted_runs_handles_2d_transitivity_break() { let mut files = vec![ - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( + new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 0, 100, @@ -1103,8 +816,8 @@ mod tests { 1, 10, pk_range(b"a", b"f"), - )), - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( + ), + new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 50, 150, @@ -1112,8 +825,8 @@ mod tests { 2, 10, pk_range(b"x", b"z"), - )), - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( + ), + new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 50, 150, @@ -1121,7 +834,7 @@ mod tests { 3, 10, pk_range(b"a", b"f"), - )), + ), ]; let runs = find_sorted_runs(&mut files); @@ -1133,7 +846,7 @@ mod tests { #[test] fn test_find_overlapping_items_skips_pk_disjoint_pairs() { - let mut left = SortedRun::from(vec![FileGroup::new_with_file( + let mut left = SortedRun::from(vec![ new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 0, @@ -1143,8 +856,8 @@ mod tests { 10, pk_range(b"a", b"f"), ), - )]); - let mut right = SortedRun::from(vec![FileGroup::new_with_file( + ]); + let mut right = SortedRun::from(vec![ new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 50, @@ -1154,7 +867,7 @@ mod tests { 10, pk_range(b"x", b"z"), ), - )]); + ]); let mut result = Vec::new(); find_overlapping_items(&mut left, &mut right, &mut result); @@ -1163,102 +876,26 @@ mod tests { } #[test] - fn test_file_group_touching_time_boundary_with_same_pk_is_not_overlap() { - let lhs = - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( - FileId::random(), - 0, - 100, - 0, - 1, - 10, - pk_range(b"a", b"f"), - )); - let rhs = - FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range( - FileId::random(), - 100, - 150, - 0, - 2, - 10, - pk_range(b"a", b"f"), - )); + fn test_file_touching_time_boundary_with_same_pk_is_not_overlap() { + let lhs = new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + 0, + 100, + 0, + 1, + 10, + pk_range(b"a", b"f"), + ); + let rhs = new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + 100, + 150, + 0, + 2, + 10, + pk_range(b"a", b"f"), + ); assert!(!lhs.overlap(&rhs)); } - - #[test] - fn test_merge_seq_files() { - // Test empty input - let files = Vec::::new(); - assert_eq!(merge_seq_files(&files, None), Vec::::new()); - - // Test single file input (should return empty vec as no merge needed) - let files = build_items(&[(1, 5)]); - assert_eq!(merge_seq_files(&files, None), Vec::::new()); - - // Test the example case: [10, 1, 1, 1] - should merge the last three files - let files = build_items_with_size(&[(1, 2, 10), (3, 4, 1), (5, 6, 1), (7, 8, 1)]); - let result = merge_seq_files(&files, None); - assert_eq!(result.len(), 3); - assert_eq!(result[0].size, 1); - assert_eq!(result[1].size, 1); - assert_eq!(result[2].size, 1); - - // Test with files of equal size - should merge as many as possible - let files = build_items_with_size(&[(1, 2, 5), (3, 4, 5), (5, 6, 5), (7, 8, 5)]); - let result = merge_seq_files(&files, Some(20)); - assert_eq!(result.len(), 4); // Should merge all 4 files as total size is 20 - - // Test with max_file_size constraint - let files = build_items_with_size(&[(1, 2, 5), (3, 4, 5), (5, 6, 5), (7, 8, 5)]); - let result = merge_seq_files(&files, Some(10)); - assert_eq!(result.len(), 2); // Should merge only 2 files as max size is 10 - - // Test with uneven file sizes - should prioritize reducing file count - let files = build_items_with_size(&[(1, 2, 2), (3, 4, 3), (5, 6, 4), (7, 8, 10)]); - let result = merge_seq_files(&files, Some(10)); - assert_eq!(result.len(), 3); // Should merge the first 3 files (total size 9) - - // Test amplification factor prioritization - // Two possible merges: [5, 5] (amp factor 0.5) vs [10, 1, 1] (amp factor 0.83) - let files = - build_items_with_size(&[(1, 2, 5), (3, 4, 5), (5, 6, 10), (7, 8, 1), (9, 10, 1)]); - let result = merge_seq_files(&files, Some(12)); - assert_eq!(result.len(), 2); - assert_eq!(result[0].size, 5); - assert_eq!(result[1].size, 5); - - // Test with large file preventing merges - let files = build_items_with_size(&[(1, 2, 100), (3, 4, 1), (5, 6, 1), (7, 8, 1)]); - let result = merge_seq_files(&files, Some(10)); - assert_eq!(result.len(), 3); // Should merge the last 3 small files - assert_eq!(result[0].size, 1); - assert_eq!(result[1].size, 1); - assert_eq!(result[2].size, 1); - - let files = build_items_with_size(&[(1, 2, 100), (3, 4, 20), (5, 6, 20), (7, 8, 20)]); - let result = merge_seq_files(&files, Some(200)); - assert_eq!(result.len(), 4); - - let files = build_items_with_size(&[(1, 2, 160), (3, 4, 20), (5, 6, 20), (7, 8, 20)]); - let result = merge_seq_files(&files, None); - assert_eq!(result.len(), 3); - assert_eq!(result[0].size, 20); - assert_eq!(result[1].size, 20); - assert_eq!(result[2].size, 20); - - let files = build_items_with_size(&[(1, 2, 100), (3, 4, 1)]); - let result = merge_seq_files(&files, Some(200)); - assert_eq!(result.len(), 2); - assert_eq!(result[0].size, 100); - assert_eq!(result[1].size, 1); - - let files = build_items_with_size(&[(1, 2, 20), (3, 4, 20), (5, 6, 20), (7, 8, 20)]); - let result = merge_seq_files(&files, Some(40)); - assert_eq!(result.len(), 2); - assert_eq!(result[0].start, 1); - assert_eq!(result[1].start, 3); - } } diff --git a/src/mito2/src/compaction/scheduler.rs b/src/mito2/src/compaction/scheduler.rs index 69a7acc5a5a..6aebb25eebc 100644 --- a/src/mito2/src/compaction/scheduler.rs +++ b/src/mito2/src/compaction/scheduler.rs @@ -264,9 +264,13 @@ impl CompactionScheduler { /// /// # Effects /// - /// Notifies waiters, schedules a pending manual or explicit automatic - /// follow-up, or removes the region status. The returned transition reports - /// a dispatched automatic follow-up or DDLs that are now safe to execute. + /// Notifies waiters, schedules a pending manual or automatic follow-up, or + /// removes the region status. A follow-up is scheduled whenever the cycle + /// latched an automatic trigger, or when `made_progress` is true (the + /// execution removed more files than it added): successful compaction keeps + /// draining the region until the picker returns no plan. The returned + /// transition reports a dispatched automatic follow-up or DDLs that are now + /// safe to execute. /// /// # Constraints /// @@ -279,12 +283,18 @@ impl CompactionScheduler { execution: &CompactionExecution, manifest_ctx: &ManifestContextRef, schema_metadata_manager: SchemaMetadataManagerRef, + made_progress: bool, ) -> CompactionTransition { if !self.is_current_execution(region_id, execution) { return CompactionTransition::NoAction; } - self.on_compaction_finished(region_id, manifest_ctx, schema_metadata_manager) - .await + self.on_compaction_finished( + region_id, + manifest_ctx, + schema_metadata_manager, + made_progress, + ) + .await } /// Completes a cooperatively canceled execution. @@ -571,6 +581,7 @@ impl CompactionScheduler { region_id: RegionId, manifest_ctx: &ManifestContextRef, schema_metadata_manager: SchemaMetadataManagerRef, + made_progress: bool, ) -> CompactionTransition { if !self.region_status.contains_key(®ion_id) { return CompactionTransition::NoAction; @@ -605,7 +616,11 @@ impl CompactionScheduler { return CompactionTransition::DdlReady(pending_ddl_requests); } - if status.active.reset_automatic_followup() + // Keep draining when the cycle latched an automatic trigger, or when the + // execution reduced the file count. A no-progress rewrite stops here so a + // split-heavy output cannot loop forever; the next flush trigger resumes. + let should_continue = status.active.reset_automatic_followup() || made_progress; + if should_continue && self.schedule_automatic_followup(region_id, manifest_ctx, schema_metadata_manager) { return CompactionTransition::AutomaticFollowupScheduled; diff --git a/src/mito2/src/compaction/scheduler_test.rs b/src/mito2/src/compaction/scheduler_test.rs index d3548ee8a5f..9545d0eea2a 100644 --- a/src/mito2/src/compaction/scheduler_test.rs +++ b/src/mito2/src/compaction/scheduler_test.rs @@ -826,6 +826,7 @@ async fn test_stale_plan_execution_does_not_affect_replacement_status() { &stale_execution, &manifest_ctx, schema_metadata_manager, + true, ) .await; assert!(pending_ddls.is_empty()); @@ -907,7 +908,7 @@ async fn test_schedule_compaction_skips_task_exceeding_memory_limit() { } #[tokio::test] -async fn test_execution_finished_without_followup_removes_status() { +async fn test_execution_finished_drains_until_no_plan() { common_telemetry::init_default_ut_logging(); let job_scheduler = Arc::new(VecScheduler::default()); let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); @@ -975,23 +976,59 @@ async fn test_execution_finished_without_followup_removes_status() { &file_metas, purger.clone(), ); - // A completed execution without an explicit follow-up removes its status. - scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + // A completed execution that reduced the file count chains another pick even + // without an explicit trigger, because the layout is still compactable. + let transition = scheduler + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) .await; + assert_matches!(transition, CompactionTransition::AutomaticFollowupScheduled); + assert!(scheduler.region_status.contains_key(®ion_id)); + let finished = recv_compaction_pick_finished(&mut rx).await; + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager.clone()) + .await; + assert_eq!(2, job_scheduler.num_jobs()); + + // Replace the layout with a single file: the drain must stop once the picker + // finds nothing to do. + let data = version_control.current(); + let file_metas: Vec<_> = data.version.ssts.levels()[0] + .files + .values() + .map(|file| file.meta_ref().clone()) + .collect(); + apply_edit(&version_control, &[(0, end)], &file_metas, purger.clone()); + let transition = scheduler + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) + .await; + assert_matches!(transition, CompactionTransition::AutomaticFollowupScheduled); + let finished = recv_compaction_pick_finished(&mut rx).await; + let transition = scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager.clone()) + .await; + assert_matches!(transition, CompactionTransition::NoAction); assert!(scheduler.region_status.is_empty()); - assert_eq!(1, job_scheduler.num_jobs()); + assert_eq!(2, job_scheduler.num_jobs()); } #[tokio::test] -async fn test_time_range_compaction_when_compaction_in_progress() { +async fn test_execution_finished_without_progress_removes_status() { common_telemetry::init_default_ut_logging(); let job_scheduler = Arc::new(VecScheduler::default()); let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); let (tx, mut rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx); let mut builder = VersionControlBuilder::new(); - let purger = builder.file_purger(); let region_id = builder.region_id(); let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); @@ -1017,6 +1054,70 @@ async fn test_time_range_compaction_when_compaction_in_progress() { .push_l0_file(90, end) .build(), ); + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .await; + let scheduled = scheduler + .schedule_automatic_compaction( + compact_request::Options::Regular(Default::default()), + &version_control, + &env.access_layer, + &manifest_ctx, + schema_metadata_manager.clone(), + ) + .unwrap(); + assert!(scheduled); + let finished = recv_compaction_pick_finished(&mut rx).await; + scheduler + .handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager.clone()) + .await; + assert_eq!(1, job_scheduler.num_jobs()); + + // The execution rewrote files without reducing the file count (e.g. its output + // was split into more files than its input). Chaining would loop without making + // progress, so the lifecycle ends here; the next flush trigger resumes compaction. + let transition = scheduler + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + false, + ) + .await; + assert_matches!(transition, CompactionTransition::NoAction); + assert!(scheduler.region_status.is_empty()); + assert_eq!(1, job_scheduler.num_jobs()); +} + +#[tokio::test] +async fn test_time_range_compaction_when_compaction_in_progress() { + common_telemetry::init_default_ut_logging(); + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + let mut builder = VersionControlBuilder::new(); + let purger = builder.file_purger(); + let region_id = builder.region_id(); + + let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); + schema_metadata_manager + .register_region_table_info( + builder.region_id().table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + kv_backend, + ) + .await; + + // 40 files to compact. The first task picks 32, leaving 8 for the pending request. + let end = 1000 * 1000; + for offset in 0..40 { + builder.push_l0_file(offset * 10, end); + } + let version_control = Arc::new(builder.build()); let manifest_ctx = env .mock_manifest_context(version_control.current().version.metadata.clone()) .await; @@ -1027,13 +1128,9 @@ async fn test_time_range_compaction_when_compaction_in_progress() { .map(|file| file.meta_ref().clone()) .collect(); - // 5 files for next compaction and removes old files. - apply_edit( - &version_control, - &[(0, end), (20, end), (40, end), (60, end), (80, end)], - &file_metas, - purger.clone(), - ); + // Replaces the files before scheduling the first compaction. + let next_files = (0..40).map(|offset| (offset * 20, end)).collect::>(); + apply_edit(&version_control, &next_files, &file_metas, purger.clone()); scheduler .schedule_automatic_compaction( @@ -1094,7 +1191,12 @@ async fn test_time_range_compaction_when_compaction_in_progress() { // On compaction finished and schedule next compaction. scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) .await; assert_eq!(1, scheduler.region_status.len()); assert!( @@ -1285,10 +1387,21 @@ async fn test_ranged_compaction_stops_without_followup() { assert_eq!(1, job_scheduler.num_jobs()); scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) .await; - assert!(!scheduler.region_status.contains_key(®ion_id)); - assert_eq!(1, job_scheduler.num_jobs()); + // The execution reduced the file count and the remaining window is still + // compactable, so the scheduler drains it with an unrestricted follow-up. + assert!(scheduler.region_status.contains_key(®ion_id)); + let followup = recv_compaction_pick_finished(&mut rx).await; + scheduler + .handle_compaction_pick_finished(followup, &manifest_ctx, schema_metadata_manager.clone()) + .await; + assert_eq!(2, job_scheduler.num_jobs()); } #[tokio::test] @@ -1369,7 +1482,12 @@ async fn test_automatic_trigger_during_execution_clears_continuation_scope() { // Finishing the manual execution must chain an unrestricted follow-up instead of // continuing with the manual request's time range. scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) .await; let continuation = recv_compaction_pick_finished(&mut rx).await; @@ -1936,7 +2054,7 @@ async fn test_on_compaction_finished_returns_pending_ddl_requests() { }); let pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager, true) .await; assert_eq!(pending_ddls.len(), 1); @@ -1985,7 +2103,12 @@ async fn test_planning_terminal_prioritizes_pending_ddl_over_automatic_followup( assert!(result.is_ok()); let pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) .await; assert!(pending_ddls.is_empty()); @@ -2034,7 +2157,7 @@ async fn test_on_compaction_finished_dispatches_pending_ddl_before_chained_regul }); let pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager, true) .await; assert_eq!(pending_ddls.len(), 1); @@ -2056,7 +2179,7 @@ async fn test_on_compaction_finished_returns_empty_when_region_absent() { let (schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager(); let pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager) + .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager, true) .await; assert!(pending_ddls.is_empty()); @@ -2111,7 +2234,12 @@ async fn test_on_compaction_finished_manual_schedule_error_cleans_status() { }); let pending_ddls = scheduler - .on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone()) + .on_compaction_finished( + region_id, + &manifest_ctx, + schema_metadata_manager.clone(), + true, + ) .await; assert!(pending_ddls.is_empty()); diff --git a/src/mito2/src/compaction/twcs.rs b/src/mito2/src/compaction/twcs.rs index 342426bbefe..cdd613b3424 100644 --- a/src/mito2/src/compaction/twcs.rs +++ b/src/mito2/src/compaction/twcs.rs @@ -15,8 +15,7 @@ use std::collections::hash_map::Entry; use std::collections::{BTreeMap, HashMap, HashSet}; use std::fmt::Debug; -use std::num::NonZeroU64; -use std::sync::Arc; +use std::sync::{Arc, LazyLock}; use common_base::readable_size::ReadableSize; use common_telemetry::{debug, info}; @@ -32,8 +31,8 @@ use crate::compaction::buckets::infer_time_bucket; use crate::compaction::compactor::CompactionRegion; use crate::compaction::picker::{Picker, PickerOutput, get_expired_ssts}; use crate::compaction::run::{ - FileGroup, Item, Ranged, find_sorted_runs, find_sorted_runs_by_time_range, - merge_primary_key_ranges, merge_seq_files, primary_key_ranges_overlap, reduce_runs, + Ranged, SortedRun, find_sorted_runs, find_sorted_runs_by_time_range, merge_primary_key_ranges, + primary_key_ranges_overlap, }; use crate::error::{JoinSnafu, Result}; use crate::sst::file::{FileHandle, Level, overlaps}; @@ -41,8 +40,24 @@ use crate::sst::version::LevelMeta; const LEVEL_COMPACTED: Level = 1; -/// Default value for max compaction input file num. -const DEFAULT_MAX_INPUT_FILE_NUM: usize = 32; +/// Default maximum number of input SST files in one compaction input. +const DEFAULT_MAX_INPUT_FILES: usize = 32; + +const MAX_INPUT_FILES_ENV: &str = "GREPTIME_TWCS_MAX_INPUT_FILES"; + +/// Maximum number of input SST files in one compaction input. +/// Configurable via [`MAX_INPUT_FILES_ENV`]. +static MAX_INPUT_FILES: LazyLock = LazyLock::new(|| { + let env_value = std::env::var(MAX_INPUT_FILES_ENV).ok(); + parse_max_input_files(env_value.as_deref()) +}); + +fn parse_max_input_files(env_value: Option<&str>) -> usize { + env_value + .and_then(|env_value| env_value.parse().ok()) + .filter(|max_input_files| *max_input_files >= 2) + .unwrap_or(DEFAULT_MAX_INPUT_FILES) +} /// `TwcsPicker` picks files of which the max timestamp are in the same time window as compaction /// candidates. @@ -52,7 +67,8 @@ pub struct TwcsPicker { pub trigger_file_num: usize, /// Compaction time window in seconds. pub time_window_seconds: Option, - /// Max allowed compaction output file size. + /// Max allowed compaction output file size. The picker also uses it to predict + /// output splits when scoring candidates. pub max_output_file_size: Option, /// Whether the target region is in append mode. pub append_mode: bool, @@ -113,7 +129,7 @@ impl TwcsPicker { output.push(CompactionOutput { output_level: LEVEL_COMPACTED, // always compact to l1 - inputs: inputs.into_iter().flat_map(|fg| fg.into_files()).collect(), + inputs, filter_deleted, output_time_range: None, // we do not enforce output time range in twcs compactions. }); @@ -138,7 +154,11 @@ impl TwcsPicker { active_window: Option, files: &Window, windows: &BTreeMap, - ) -> (Vec, bool) { + ) -> (Vec, bool) { + if files.files.len() < self.trigger_file_num { + return (vec![], false); + } + let window = &files.time_window; let mut files_to_merge: Vec<_> = files.files().cloned().collect(); @@ -148,7 +168,7 @@ impl TwcsPicker { { let (kept_files, ignored_files) = files_to_merge .into_iter() - .partition(|fg| fg.size() <= max_size as usize); + .partition(|file| file.size() <= max_size); files_to_merge = kept_files; if !ignored_files.is_empty() { info!( @@ -167,59 +187,13 @@ impl TwcsPicker { find_sorted_runs_by_time_range(&mut files_to_merge) }; let found_runs = sorted_runs.len(); - // We only remove deletion markers if we found less than 2 runs and not in append mode. - // because after compaction there will be no overlapping files. - let mut filter_deleted = - found_runs <= 2 && !self.append_mode && !window_has_overlap(files, windows); - if found_runs == 0 { - return (vec![], filter_deleted); - } - - let mut inputs = if found_runs > 1 { - reduce_runs(sorted_runs) - } else { - let run = sorted_runs.last().unwrap(); - if run.items().len() < self.trigger_file_num { - return (vec![], filter_deleted); - } - // no overlapping files, try merge small files - merge_seq_files(run.items(), self.max_output_file_size) - }; - - // Limits the number of input files. - let total_input_files: usize = inputs.iter().map(|fg| fg.num_files()).sum(); - if total_input_files > DEFAULT_MAX_INPUT_FILE_NUM { - // Sorts file groups by size first. - inputs.sort_unstable_by_key(|fg| fg.size()); - let mut num_picked_files = 0; - inputs = inputs - .into_iter() - .take_while(|fg| { - let current_group_file_num = fg.num_files(); - if current_group_file_num + num_picked_files <= DEFAULT_MAX_INPUT_FILE_NUM { - num_picked_files += current_group_file_num; - true - } else { - false - } - }) - .collect::>(); - info!( - "Compaction for region {} enforces max input file num limit: {}, current total: {}, input: {:?}", - region_id, DEFAULT_MAX_INPUT_FILE_NUM, total_input_files, inputs - ); - } - - // The inputs are only a subset of the window: `reduce_runs` and `merge_seq_files` both - // narrow the selection down and the file num limit above narrows it further. A deletion - // marker may only be dropped when every file that can hold a row it masks is compacted - // along with it, so re-check the final selection against what stays in the window. - if filter_deleted && overlaps_files_left_behind(&inputs, &files_to_merge) { - filter_deleted = false; - } + let inputs = pick_count_first(sorted_runs, self.max_output_file_size); + let filter_deleted = !self.append_mode + && !window_has_overlap(files, windows) + && !selected_overlaps_unselected(&inputs, files); if inputs.len() > 1 { - // If we have more than one group to compact. + // If we have more than one file to compact. log_pick_result( region_id, *window, @@ -235,18 +209,225 @@ impl TwcsPicker { } } -/// Returns whether any file group of `window_files` that is not part of `inputs` overlaps -/// `inputs`. Such a group keeps rows a deletion marker in `inputs` masks, so the compaction -/// must not filter deleted rows out. +#[derive(Debug)] +struct OrderedFile<'a> { + file: &'a FileHandle, + run_id: usize, + position_in_run: usize, +} + +/// Metrics of a candidate compaction input, accumulated incrementally as the +/// candidate interval expands one file at a time. +#[derive(Debug, Default)] +struct Candidate { + /// Number of files in the interval. + num_files: usize, + /// Total input bytes of the interval. + total_size: usize, + /// Size of the largest single file in the interval. + largest_file_size: usize, + /// Files in the interval that overlap another interval file from a different + /// sorted run. Resolving these overlaps is what reduces sorted runs. + overlap_participants: usize, +} + +impl Candidate { + /// Absorbs `file` into the candidate. `preceding` holds the interval files added + /// before it and `participations` tracks which of them already participate in an + /// intra-interval overlap; both exist only for overlap accounting. + fn absorb( + &mut self, + file: &OrderedFile, + preceding: &[OrderedFile], + participations: &mut Vec, + ) { + self.num_files += 1; + let file_size = file.file.size() as usize; + self.total_size += file_size; + self.largest_file_size = self.largest_file_size.max(file_size); + + let mut participates = false; + for (offset, other) in preceding.iter().enumerate() { + if file.run_id != other.run_id && file.file.overlap_inclusive(other.file) { + if !participations[offset] { + participations[offset] = true; + self.overlap_participants += 1; + } + participates = true; + } + } + participations.push(participates); + if participates { + self.overlap_participants += 1; + } + } + + /// Predicted number of output files after compacting this candidate. Compaction + /// output holds at most the input bytes, so a split threshold of + /// `max_output_file_size` yields at most `ceil(total_size / threshold)` files. + /// Without a threshold the output is a single file. + fn predicted_output_files(&self, max_output_file_size: Option) -> usize { + match max_output_file_size { + Some(max) if max > 0 => self.total_size.div_ceil(max as usize).max(1), + _ => 1, + } + } + + /// Net file count reduction compacting this candidate would achieve, i.e. how + /// many fewer physical SSTs the window holds afterwards. Zero when the output + /// would be split back into at least as many files as the input. + fn file_reduction(&self, max_output_file_size: Option) -> usize { + self.num_files + .saturating_sub(self.predicted_output_files(max_output_file_size)) + } + + /// A large historical file is only rewritten once its peers contribute at least + /// the same amount of data, so every rewrite moves it into a larger size tier. + fn is_balanced(&self) -> bool { + self.largest_file_size <= self.total_size - self.largest_file_size + } + + /// A candidate is worth compacting only if it makes progress on at least one + /// axis: it reduces the physical file count, or it resolves at least one + /// overlap (merging sorted runs and reducing read amplification). A pure + /// rewrite that achieves neither only burns I/O. + fn makes_progress(&self, max_output_file_size: Option) -> bool { + self.file_reduction(max_output_file_size) > 0 || self.overlap_participants > 0 + } + + fn score(&self, max_output_file_size: Option) -> CandidateScore { + CandidateScore { + file_reduction: self.file_reduction(max_output_file_size), + overlap_participants: self.overlap_participants, + total_size: self.total_size, + } + } +} + +/// Score of a candidate compaction input. Higher is better: first by the predicted +/// net file count reduction, then by the number of files participating in an +/// overlap, then by smaller input bytes. +#[derive(Debug)] +struct CandidateScore { + file_reduction: usize, + overlap_participants: usize, + total_size: usize, +} + +impl CandidateScore { + fn is_better_than(&self, other: &Self) -> bool { + self.file_reduction + .cmp(&other.file_reduction) + .then_with(|| self.overlap_participants.cmp(&other.overlap_participants)) + .then_with(|| other.total_size.cmp(&self.total_size)) + .is_gt() + } +} + +/// Picks a contiguous (in global time order) interval of files to compact. /// -/// Overlapping is inclusive on both boundaries here: two files that only share a boundary -/// timestamp can still hold the same row. -fn overlaps_files_left_behind(inputs: &[FileGroup], window_files: &[FileGroup]) -> bool { - let picked: HashSet<_> = inputs.iter().flat_map(|fg| fg.file_ids()).collect(); - window_files +/// The picker reorders all files from all sorted runs by `(start asc, end desc)` and +/// enumerates every interval of at most [`MAX_INPUT_FILES`] files as a [`Candidate`]. +/// An interval is eligible when it +/// +/// - holds at least 2 files, +/// - is balanced: no single file dominates it (`largest <= sum of the others`), +/// - makes progress on at least one axis: it reduces the physical file count given +/// the output split threshold `max_output_file_size`, or it resolves at least one +/// overlap between sorted runs. +/// +/// The best interval wins by [`CandidateScore`]; ties keep the earliest interval. +fn pick_count_first( + sorted_runs: Vec>, + max_output_file_size: Option, +) -> Vec { + let files = ordered_files(&sorted_runs); + + let mut best = None; + for left in 0..files.len() { + let mut candidate = Candidate::default(); + let right_bound = left.saturating_add(*MAX_INPUT_FILES).min(files.len()); + let mut participations: Vec = Vec::with_capacity(right_bound - left); + for right in left..right_bound { + candidate.absorb(&files[right], &files[left..right], &mut participations); + if candidate.num_files < 2 + || !candidate.is_balanced() + || !candidate.makes_progress(max_output_file_size) + { + continue; + } + + let score = candidate.score(max_output_file_size); + if best + .as_ref() + .is_none_or(|(best_score, _)| score.is_better_than(best_score)) + { + best = Some((score, &files[left..=right])); + } + } + } + + let Some((_, best)) = best else { + return vec![]; + }; + best.iter().map(|file| file.file.clone()).collect() +} + +/// Flattens sorted runs into files ordered by `(start asc, end desc)`, breaking ties +/// by run and position to keep the order deterministic. +fn ordered_files(sorted_runs: &[SortedRun]) -> Vec> { + let mut files = sorted_runs .iter() - .filter(|fg| !fg.file_ids().iter().any(|id| picked.contains(id))) - .any(|fg| inputs.iter().any(|input| input.overlap_inclusive(fg))) + .enumerate() + .flat_map(|(run_id, run)| { + run.items() + .iter() + .enumerate() + .map(move |(position_in_run, file)| OrderedFile { + file, + run_id, + position_in_run, + }) + }) + .collect::>(); + files.sort_unstable_by(|lhs, rhs| { + let (lhs_start, lhs_end) = lhs.file.range(); + let (rhs_start, rhs_end) = rhs.file.range(); + lhs_start + .cmp(&rhs_start) + .then_with(|| rhs_end.cmp(&lhs_end)) + .then_with(|| lhs.run_id.cmp(&rhs.run_id)) + .then_with(|| lhs.position_in_run.cmp(&rhs.position_in_run)) + }); + files +} + +fn selected_overlaps_unselected(selected: &[FileHandle], window: &Window) -> bool { + // The overall time span of the selection: a file outside it cannot overlap any + // selected file (ranges are inclusive), so it needs no precise overlap check. + let Some((span_start, span_end)) = selected + .iter() + .map(Ranged::range) + .reduce(|(start_a, end_a), (start_b, end_b)| (start_a.min(start_b), end_a.max(end_b))) + else { + return false; + }; + let selected_file_ids = selected + .iter() + .map(FileHandle::file_id) + .collect::>(); + window + .files() + .filter(|file| { + let (start, end) = file.range(); + start <= span_end && span_start <= end + }) + .filter(|file| !selected_file_ids.contains(&file.file_id())) + .any(|unselected| { + selected + .iter() + .any(|selected| selected.overlap_inclusive(unselected)) + }) } #[allow(clippy::too_many_arguments)] @@ -258,7 +439,7 @@ fn log_pick_result( file_num: usize, max_output_file_size: Option, filter_deleted: bool, - inputs: &[FileGroup], + inputs: &[FileHandle], ) { let input_file_str: Vec = inputs .iter() @@ -268,11 +449,11 @@ fn log_pick_result( let end = range.1.to_iso8601_string(); let num_rows = f.num_rows(); format!( - "FileGroup{{id: {:?}, range: ({}, {}), size: {}, num rows: {} }}", - f.file_ids(), + "File{{id: {:?}, range: ({}, {}), size: {}, num rows: {} }}", + f.file_id(), start, end, - ReadableSize(f.size() as u64), + ReadableSize(f.size()), num_rows ) }) @@ -369,9 +550,7 @@ impl Picker for TwcsPicker { struct Window { start: Timestamp, end: Timestamp, - // Mapping from file sequence to file groups. Files with the same sequence is considered - // created from the same compaction task. - files: HashMap, FileGroup>, + files: Vec, time_window: i64, primary_key_range: Option<(bytes::Bytes, bytes::Bytes)>, } @@ -381,11 +560,10 @@ impl Window { fn new_with_file(file: FileHandle) -> Self { let (start, end) = file.time_range(); let primary_key_range = file.primary_key_range(); - let files = HashMap::from([(file.meta_ref().sequence, FileGroup::new_with_file(file))]); Self { start, end, - files, + files: vec![file], time_window: 0, primary_key_range, } @@ -403,19 +581,11 @@ impl Window { self.end = self.end.max(end); self.primary_key_range = merge_primary_key_ranges(self.primary_key_range.take(), file.primary_key_range()); - - match self.files.entry(file.meta_ref().sequence) { - Entry::Occupied(mut o) => { - o.get_mut().add_file(file); - } - Entry::Vacant(v) => { - v.insert(FileGroup::new_with_file(file)); - } - } + self.files.push(file); } - fn files(&self) -> impl Iterator { - self.files.values() + fn files(&self) -> impl Iterator { + self.files.iter() } } @@ -547,6 +717,18 @@ mod tests { use crate::test_util::memtable_util::metadata_for_test; use crate::test_util::scheduler_util::SchedulerEnv; + #[test] + fn test_valid_max_input_files_env_overrides_default() { + assert_eq!(64, parse_max_input_files(Some("64"))); + } + + #[test] + fn test_invalid_max_input_files_env_falls_back_to_default() { + for env_value in [None, Some(""), Some("invalid"), Some("0"), Some("1")] { + assert_eq!(32, parse_max_input_files(env_value)); + } + } + async fn compaction_region_with_expired_sst() -> CompactionRegion { let env = SchedulerEnv::new().await; let metadata = metadata_for_test(); @@ -664,8 +846,7 @@ mod tests { 3, ); let fgs = &windows.get(&0).unwrap().files; - assert_eq!(1, fgs.len()); - assert_eq!(fgs.values().map(|f| f.files().len()).sum::(), 5); + assert_eq!(5, fgs.len()); let files = [FileId::random(); 3]; let windows = assign_to_windows( @@ -679,26 +860,41 @@ mod tests { ); assert_eq!( files[0], - windows.get(&0).unwrap().files().next().unwrap().files()[0] + windows + .get(&0) + .unwrap() + .files() + .next() + .unwrap() .file_id() .file_id() ); assert_eq!( files[1], - windows.get(&3).unwrap().files().next().unwrap().files()[0] + windows + .get(&3) + .unwrap() + .files() + .next() + .unwrap() .file_id() .file_id() ); assert_eq!( files[2], - windows.get(&12).unwrap().files().next().unwrap().files()[0] + windows + .get(&12) + .unwrap() + .files() + .next() + .unwrap() .file_id() .file_id() ); } #[test] - fn test_assign_file_groups_to_windows() { + fn test_assign_files_to_windows() { let files = [ FileId::random(), FileId::random(), @@ -716,25 +912,14 @@ mod tests { 3, ); assert_eq!(windows.len(), 1); - let fgs = &windows.get(&0).unwrap().files; - assert_eq!(2, fgs.len()); + let window_files = &windows.get(&0).unwrap().files; + assert_eq!(4, window_files.len()); assert_eq!( - fgs.get(&NonZeroU64::new(1)) - .unwrap() - .files() + window_files .iter() .map(|f| f.file_id().file_id()) .collect::>(), - [files[0], files[1]].into_iter().collect() - ); - assert_eq!( - fgs.get(&NonZeroU64::new(2)) - .unwrap() - .files() - .iter() - .map(|f| f.file_id().file_id()) - .collect::>(), - [files[2], files[3]].into_iter().collect() + files.into_iter().collect() ); } @@ -751,11 +936,10 @@ mod tests { files[2].set_compacting(true); let mut windows = assign_to_windows(files.iter(), 3); let window0 = windows.remove(&0).unwrap(); - assert_eq!(1, window0.files.len()); + assert_eq!(3, window0.files.len()); let candidates = window0 .files - .into_values() - .flat_map(|fg| fg.into_files()) + .iter() .map(|f| f.file_id().file_id()) .collect::>(); assert_eq!(candidates.len(), 3); @@ -801,12 +985,10 @@ mod tests { assert_eq!(*overlapping, actual_overlapping); let mut file_ranges = actual_window .files - .values() - .flat_map(|f| { - f.files().iter().map(|f| { - let (s, e) = f.time_range(); - (s.value(), e.value()) - }) + .iter() + .map(|f| { + let (s, e) = f.time_range(); + (s.value(), e.value()) }) .collect::>(); file_ranges.sort_unstable_by(|l, r| l.0.cmp(&r.0).then(l.1.cmp(&r.1))); @@ -988,7 +1170,7 @@ mod tests { let active_window = find_latest_window_in_seconds(self.input_files.iter(), self.window_size); let output = TwcsPicker { - trigger_file_num: 4, + trigger_file_num: 2, time_window_seconds: None, max_output_file_size: None, append_mode: false, @@ -1075,7 +1257,7 @@ mod tests { .to_vec(), expected_outputs: vec![ ExpectedOutput { - input_files: vec![2, 4], + input_files: vec![2, 3, 4], output_level: 1, }, ExpectedOutput { @@ -1088,8 +1270,9 @@ mod tests { .await; // Case 3: - // A compaction may split output into several files that have overlapping time ranges and same sequence, - // we should treat these files as one FileGroup. + // A compaction may split output into several files that have overlapping time ranges and + // the same sequence. They are ordinary compaction candidates now: the picker merges the + // overlapping ones in the window that reaches the trigger. let file_ids = (0..6).map(|_| FileId::random()).collect::>(); CompactionPickerTestCase { window_size: 3, @@ -1101,10 +1284,16 @@ mod tests { new_file_handle_with_sequence(file_ids[4], 11, 2990, 0, 3), ] .to_vec(), - expected_outputs: vec![ExpectedOutput { - input_files: vec![0, 1, 4], - output_level: 1, - }], + expected_outputs: vec![ + ExpectedOutput { + input_files: vec![2, 3], + output_level: 1, + }, + ExpectedOutput { + input_files: vec![0, 1, 4], + output_level: 1, + }, + ], } .check() .await; @@ -1160,30 +1349,24 @@ mod tests { let small_file_2 = new_file_handle_with_size_and_sequence(file_ids[2], 0, 999, 0, 3, 800); let large_file_2 = new_file_handle_with_size_and_sequence(file_ids[3], 0, 999, 0, 4, 2000); - // Create file groups (each file is in its own group due to different sequences) - let mut files_to_merge = vec![ - FileGroup::new_with_file(small_file_1), - FileGroup::new_with_file(large_file_1), - FileGroup::new_with_file(small_file_2), - FileGroup::new_with_file(large_file_2), - ]; + let mut files_to_merge = vec![small_file_1, large_file_1, small_file_2, large_file_2]; // Test filtering logic directly let original_count = files_to_merge.len(); // Apply append mode filtering - files_to_merge.retain(|fg| fg.size() <= max_output_file_size as usize); + files_to_merge.retain(|file| file.size() <= max_output_file_size); // Should have filtered out 2 large files, leaving 2 small files assert_eq!(files_to_merge.len(), 2); assert_eq!(original_count, 4); // Verify the remaining files are the small ones - for fg in &files_to_merge { + for file in &files_to_merge { assert!( - fg.size() <= max_output_file_size as usize, + file.size() <= max_output_file_size, "File size {} should be <= {}", - fg.size(), + file.size(), max_output_file_size ); } @@ -1191,7 +1374,7 @@ mod tests { #[tokio::test] async fn test_build_output_multiple_windows_with_zero_runs() { - let file_ids = (0..6).map(|_| FileId::random()).collect::>(); + let file_ids = (0..7).map(|_| FileId::random()).collect::>(); let files = [ // Window 0: Contains 3 files but not forming any runs (not enough files in sequence to reach trigger_file_num) @@ -1202,6 +1385,7 @@ mod tests { new_file_handle_with_sequence(file_ids[3], 3000, 3999, 0, 4), new_file_handle_with_sequence(file_ids[4], 3000, 3999, 0, 5), new_file_handle_with_sequence(file_ids[5], 3000, 3999, 0, 6), + new_file_handle_with_sequence(file_ids[6], 3000, 3999, 0, 7), ]; let windows = assign_to_windows(files.iter(), 3); @@ -1431,11 +1615,11 @@ mod tests { .unwrap(); assert_eq!(1, output.len()); - assert_eq!(output[0].inputs.len(), 2); + assert_eq!(output[0].inputs.len(), num_files); } #[tokio::test] - async fn test_limit_max_input_files() { + async fn test_window_trigger_can_exceed_input_limit() { common_telemetry::init_default_ut_logging(); let num_files = 50; @@ -1460,7 +1644,7 @@ mod tests { let windows = assign_to_windows(files.iter(), 3); let picker = TwcsPicker { - trigger_file_num: 4, + trigger_file_num: num_files, time_window_seconds: Some(3), max_output_file_size: None, append_mode: false, @@ -1475,7 +1659,7 @@ mod tests { .unwrap(); assert_eq!(1, output.len()); - assert_eq!(output[0].inputs.len(), 32); + assert_eq!(output[0].inputs.len(), num_files.min(*MAX_INPUT_FILES)); } #[tokio::test] @@ -1671,5 +1855,394 @@ mod tests { ); } + #[tokio::test] + async fn test_count_first_prefers_more_files_over_smaller_overlap() { + let files = [ + new_file_handle_with_size_and_sequence(FileId::random(), 0, 10, 0, 1, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 20, 30, 0, 2, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 40, 50, 0, 3, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 5, 15, 0, 4, 10), + ]; + let windows = assign_to_windows(files.iter(), 100); + let picker = TwcsPicker { + trigger_file_num: 2, + time_window_seconds: Some(100), + max_output_file_size: None, + append_mode: false, + max_background_tasks: None, + time_range: None, + }; + + let output = picker + .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None) + .await + .unwrap(); + + assert_eq!(1, output.len()); + assert_eq!(4, output[0].inputs.len()); + } + + #[tokio::test] + async fn test_count_first_trigger_counts_physical_ssts() { + let files = [ + new_file_handle_with_size_and_sequence(FileId::random(), 0, 10, 0, 1, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 0, 10, 0, 1, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 20, 30, 0, 2, 10), + ]; + let windows = assign_to_windows(files.iter(), 100); + let picker = TwcsPicker { + trigger_file_num: 3, + time_window_seconds: Some(100), + max_output_file_size: None, + append_mode: false, + max_background_tasks: None, + time_range: None, + }; + + let output = picker + .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None) + .await + .unwrap(); + + assert_eq!(1, output.len()); + assert_eq!(3, output[0].inputs.len()); + } + + #[tokio::test] + async fn test_count_first_does_not_compact_overlap_below_trigger() { + let files = [ + new_file_handle_with_size_and_sequence(FileId::random(), 0, 20, 0, 1, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 10, 30, 0, 2, 10), + ]; + let windows = assign_to_windows(files.iter(), 100); + let picker = TwcsPicker { + trigger_file_num: 3, + time_window_seconds: Some(100), + max_output_file_size: None, + append_mode: false, + max_background_tasks: None, + time_range: None, + }; + + let output = picker + .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None) + .await + .unwrap(); + + assert!(output.is_empty()); + } + + #[tokio::test] + async fn test_filter_deleted_is_false_when_selected_files_overlap_unselected_file() { + let mut files = (0..32) + .map(|idx| { + new_file_handle_with_size_and_sequence( + FileId::random(), + idx * 10, + idx * 10 + 9, + 0, + idx as u64 + 1, + 10, + ) + }) + .collect::>(); + files.push(new_file_handle_with_size_and_sequence( + FileId::random(), + 0, + 320, + 0, + 33, + 10, + )); + let windows = assign_to_windows(files.iter(), 1000); + let picker = TwcsPicker { + trigger_file_num: 2, + time_window_seconds: Some(1000), + max_output_file_size: None, + append_mode: false, + max_background_tasks: None, + time_range: None, + }; + + let output = picker + .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None) + .await + .unwrap(); + + assert_eq!(1, output.len()); + assert_eq!(DEFAULT_MAX_INPUT_FILES, output[0].inputs.len()); + assert!(!output[0].filter_deleted); + } + + fn new_file(start: i64, end: i64, sequence: u64, file_size: u64) -> FileHandle { + new_file_handle_with_size_and_sequence(FileId::random(), start, end, 0, sequence, file_size) + } + + fn picked_ranges(files: &[FileHandle]) -> Vec<(i64, i64)> { + files + .iter() + .map(|file| { + let (start, end) = file.range(); + (start.value(), end.value()) + }) + .collect() + } + + #[test] + fn test_count_first_rejects_dominant_historical_file() { + let picked = pick_count_first( + vec![SortedRun::from(vec![ + new_file(0, 9, 1, 400), + new_file(20, 29, 2, 100), + ])], + None, + ); + + assert!(picked.is_empty()); + } + + #[test] + fn test_count_first_accepts_balanced_historical_file() { + let picked = pick_count_first( + vec![SortedRun::from(vec![ + new_file(0, 9, 1, 400), + new_file(20, 29, 2, 100), + new_file(40, 49, 3, 100), + new_file(60, 69, 4, 100), + new_file(80, 89, 5, 100), + ])], + None, + ); + + assert_eq!( + vec![(0, 9), (20, 29), (40, 49), (60, 69), (80, 89)], + picked_ranges(&picked) + ); + } + + #[test] + fn test_count_first_finds_smaller_balanced_interval_when_larger_one_is_unbalanced() { + let picked = pick_count_first( + vec![SortedRun::from(vec![ + new_file(0, 9, 1, 1000), + new_file(20, 29, 2, 100), + new_file(40, 49, 3, 100), + new_file(60, 69, 4, 100), + new_file(80, 89, 5, 100), + ])], + None, + ); + + assert_eq!( + vec![(20, 29), (40, 49), (60, 69), (80, 89)], + picked_ranges(&picked) + ); + } + + #[test] + fn test_count_first_prefers_overlap_participants_when_file_counts_match() { + let first_run = (0..DEFAULT_MAX_INPUT_FILES) + .map(|idx| { + let start = idx as i64 * 20; + let end = if idx + 1 == DEFAULT_MAX_INPUT_FILES { + 700 + } else { + start + 9 + }; + new_file(start, end, idx as u64 + 1, 10) + }) + .collect::>(); + let overlapping = new_file(690, 710, 100, 10); + + let picked = pick_count_first( + vec![ + SortedRun::from(first_run), + SortedRun::from(vec![overlapping]), + ], + None, + ); + let ranges = picked_ranges(&picked); + + assert_eq!(DEFAULT_MAX_INPUT_FILES, ranges.len()); + assert!(!ranges.contains(&(0, 9))); + assert!(ranges.contains(&(690, 710))); + } + + #[test] + fn test_count_first_prefers_smaller_bytes_when_file_counts_match() { + let files = (0..=DEFAULT_MAX_INPUT_FILES) + .map(|idx| { + let start = idx as i64 * 20; + let size = if idx == 0 { + 100 + } else if idx == DEFAULT_MAX_INPUT_FILES { + 1 + } else { + 10 + }; + new_file(start, start + 9, idx as u64 + 1, size) + }) + .collect::>(); + + let picked = pick_count_first(vec![SortedRun::from(files)], None); + let ranges = picked_ranges(&picked); + + assert_eq!(DEFAULT_MAX_INPUT_FILES, ranges.len()); + assert_eq!(Some(&(20, 29)), ranges.first()); + assert_eq!(Some(&(640, 649)), ranges.last()); + } + + #[test] + fn test_count_first_skips_pure_rewrite_without_progress() { + // Four non-overlapping files whose combined size splits back into at least + // four outputs: compacting them reduces neither the file count nor any + // overlap, so there is nothing worth picking. + let picked = pick_count_first( + vec![SortedRun::from(vec![ + new_file(0, 9, 1, 600), + new_file(20, 29, 2, 600), + new_file(40, 49, 3, 600), + new_file(60, 69, 4, 600), + ])], + Some(512), + ); + + assert!(picked.is_empty()); + } + + #[test] + fn test_count_first_allows_overlap_resolution_without_file_reduction() { + // Two large overlapping files from different runs: compacting them cannot + // reduce the file count (the output splits back into just as many files), + // but it resolves the overlap and merges two sorted runs into one. + let picked = pick_count_first( + vec![ + SortedRun::from(vec![new_file(0, 19, 1, 600)]), + SortedRun::from(vec![new_file(10, 29, 2, 600)]), + ], + Some(512), + ); + + assert_eq!(vec![(0, 19), (10, 29)], picked_ranges(&picked)); + } + + #[test] + fn test_count_first_prefers_guaranteed_reduction_over_pure_rewrite() { + // A window mixing large and small files: the large triple on its own is a + // pure rewrite, and every interval containing large files rewrites more + // bytes for the same predicted reduction, so the small triple wins. + let picked = pick_count_first( + vec![SortedRun::from(vec![ + new_file(0, 9, 1, 600), + new_file(10, 19, 2, 600), + new_file(20, 29, 3, 600), + new_file(30, 39, 4, 10), + new_file(40, 49, 5, 10), + new_file(50, 59, 6, 10), + ])], + Some(512), + ); + + assert_eq!(vec![(30, 39), (40, 49), (50, 59)], picked_ranges(&picked)); + } + + #[test] + fn test_count_first_prefers_earlier_time_on_exact_tie() { + let files = (0..=DEFAULT_MAX_INPUT_FILES) + .map(|idx| { + let start = idx as i64 * 20; + new_file(start, start + 9, idx as u64 + 1, 10) + }) + .collect::>(); + + let picked = pick_count_first(vec![SortedRun::from(files)], None); + let ranges = picked_ranges(&picked); + + assert_eq!(DEFAULT_MAX_INPUT_FILES, ranges.len()); + assert_eq!(Some(&(0, 9)), ranges.first()); + assert_eq!(Some(&(620, 629)), ranges.last()); + } + + #[test] + fn test_count_first_keeps_each_interleaved_run_contiguous() { + let picked = pick_count_first( + vec![ + SortedRun::from(vec![ + new_file(0, 9, 1, 1), + new_file(20, 29, 2, 1), + new_file(40, 49, 3, 1), + ]), + SortedRun::from(vec![new_file(10, 19, 4, 1), new_file(30, 39, 5, 1)]), + ], + None, + ); + + assert_eq!( + vec![(0, 9), (10, 19), (20, 29), (30, 39), (40, 49)], + picked_ranges(&picked) + ); + } + // TODO(hl): TTL tester that checks if get_expired_ssts function works as expected. + + // Reproduces the leftover-small-files deadlock observed in the compaction bench: + // an "EngineFull" flush emits several SSTs sharing one sequence. They must be ordinary + // candidates so a contiguous candidate can span them instead of treating them as an + // indivisible block that fragments the timeline. + #[test] + fn test_count_first_candidate_spans_same_sequence_files() { + let picked = pick_count_first( + vec![SortedRun::from(vec![ + new_file(0, 9, 1, 10), + // One EngineFull flush output: three SSTs with the same sequence. + new_file(10, 19, 2, 10), + new_file(10, 19, 2, 10), + new_file(10, 19, 2, 10), + new_file(20, 29, 3, 10), + ])], + None, + ); + + assert_eq!( + vec![(0, 9), (10, 19), (10, 19), (10, 19), (20, 29)], + picked_ranges(&picked) + ); + } + + // End-to-end repro of the compaction-bench leftover files: a window where + // EngineFull flush groups (multiple SSTs sharing a sequence) sit between + // singleton flushes. The picker must be able to merge the whole contiguous + // range; otherwise these files can never be compacted once the active window + // moves on. + #[tokio::test] + async fn test_count_first_merges_window_with_interleaved_flush_groups() { + let files = [ + // Singleton flush. + new_file_handle_with_size_and_sequence(FileId::random(), 0, 9, 0, 1, 10), + // EngineFull flush: three SSTs sharing sequence 2. + new_file_handle_with_size_and_sequence(FileId::random(), 10, 19, 0, 2, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 10, 19, 0, 2, 10), + new_file_handle_with_size_and_sequence(FileId::random(), 10, 19, 0, 2, 10), + // Singleton flush. + new_file_handle_with_size_and_sequence(FileId::random(), 20, 29, 0, 3, 10), + ]; + let windows = assign_to_windows(files.iter(), 100); + let picker = TwcsPicker { + trigger_file_num: 3, + time_window_seconds: Some(100), + max_output_file_size: None, + append_mode: false, + max_background_tasks: None, + time_range: None, + }; + + let output = picker + .build_output_with_time_range(RegionId::from_u64(1), windows, Some(0), None) + .await + .unwrap(); + + assert_eq!(1, output.len()); + // All five SSTs of the window are merged in one compaction. + assert_eq!(5, output[0].inputs.len()); + } } diff --git a/src/mito2/src/engine/compaction_test.rs b/src/mito2/src/engine/compaction_test.rs index 1fc3a8f785d..04c97bdfbba 100644 --- a/src/mito2/src/engine/compaction_test.rs +++ b/src/mito2/src/engine/compaction_test.rs @@ -768,6 +768,7 @@ async fn test_truncate_waits_for_non_cancellable_compaction_commit() { .await; let create = CreateRequestBuilder::new() .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.trigger_file_num", "4") .build(); let column_schemas = create .column_metadatas @@ -779,7 +780,9 @@ async fn test_truncate_waits_for_non_cancellable_compaction_commit() { .await .unwrap(); put_and_flush(&engine, region_id, &column_schemas, 0..10).await; - put_and_flush(&engine, region_id, &column_schemas, 5..20).await; + put_and_flush(&engine, region_id, &column_schemas, 10..20).await; + put_and_flush(&engine, region_id, &column_schemas, 20..30).await; + put_and_flush(&engine, region_id, &column_schemas, 30..40).await; let commit_guard = gate.arm_commit(); let compact_engine = engine.clone(); @@ -803,7 +806,7 @@ async fn test_truncate_waits_for_non_cancellable_compaction_commit() { RegionRequest::Truncate(RegionTruncateRequest::ByTimeRanges { time_ranges: vec![( Timestamp::new_millisecond(0), - Timestamp::new_millisecond(19_000), + Timestamp::new_millisecond(39_000), )], }), ) @@ -981,12 +984,10 @@ async fn test_compaction_region_with_format(flat_format: bool) { // [20....29] // -[15.........29]- (delete) // [15.....24] - // Output: - // [0..9] - // [10............29] (contains delete) - // [15....24] + // Count-first compaction consumes the first 4 SSTs as soon as the trigger is reached. + // The compacted output and final flush leave 2 SSTs. assert_eq!( - 3, + 2, scanner.num_files(), "unexpected files: {:?}", scanner.file_ids() @@ -1530,6 +1531,7 @@ async fn test_readonly_during_compaction_with_format(flat_format: bool) { let request = CreateRequestBuilder::new() .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.trigger_file_num", "4") .build(); let column_schemas = request @@ -1542,7 +1544,9 @@ async fn test_readonly_during_compaction_with_format(flat_format: bool) { .await .unwrap(); let listener_guard = CompactionListenerGuard::new(listener.clone()); - // Flush 2 SSTs for compaction. + // Flush 4 balanced SSTs for compaction. + put_and_flush(&engine, region_id, &column_schemas, 0..10).await; + put_and_flush(&engine, region_id, &column_schemas, 5..20).await; put_and_flush(&engine, region_id, &column_schemas, 0..10).await; put_and_flush(&engine, region_id, &column_schemas, 5..20).await; @@ -1573,7 +1577,7 @@ async fn test_readonly_during_compaction_with_format(flat_format: bool) { .await .unwrap(); assert_eq!( - 2, + 4, scanner.num_files(), "unexpected files: {:?}", scanner.file_ids() @@ -1614,6 +1618,7 @@ async fn test_local_compaction_cancellation_notifies_before_pending_ddl_dispatch let request = CreateRequestBuilder::new() .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.trigger_file_num", "4") .build(); let column_schemas = request .column_metadatas @@ -1628,6 +1633,8 @@ async fn test_local_compaction_cancellation_notifies_before_pending_ddl_dispatch let merge_guard = gate.arm_merge(); put_and_flush(&engine, region_id, &column_schemas, 0..10).await; put_and_flush(&engine, region_id, &column_schemas, 5..20).await; + put_and_flush(&engine, region_id, &column_schemas, 0..10).await; + put_and_flush(&engine, region_id, &column_schemas, 5..20).await; tokio::time::timeout(Duration::from_secs(5), gate.wait_until_merge_entered()) .await diff --git a/src/mito2/src/engine/merge_mode_test.rs b/src/mito2/src/engine/merge_mode_test.rs index 40a87642ae4..afb984e77bd 100644 --- a/src/mito2/src/engine/merge_mode_test.rs +++ b/src/mito2/src/engine/merge_mode_test.rs @@ -127,6 +127,7 @@ async fn test_merge_mode_compaction_with_format(flat_format: bool) { let request = CreateRequestBuilder::new() .field_num(2) .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.trigger_file_num", "3") .insert_option("merge_mode", "last_non_null") .build(); let table_dir = request.table_dir.clone(); @@ -208,7 +209,7 @@ async fn test_merge_mode_compaction_with_format(flat_format: bool) { .scanner(region_id, ScanRequest::default()) .await .unwrap(); - assert_eq!(2, scanner.num_files()); + assert_eq!(1, scanner.num_files()); assert_eq!(1, scanner.num_memtables()); scanner.set_target_partitions(2); let stream = scanner.scan().await.unwrap(); diff --git a/src/mito2/src/gc/worker_test.rs b/src/mito2/src/gc/worker_test.rs index cd8b875e9e8..e0b03484996 100644 --- a/src/mito2/src/gc/worker_test.rs +++ b/src/mito2/src/gc/worker_test.rs @@ -399,7 +399,8 @@ async fn test_gc_worker_basic_compact() { let region = engine.get_region(region_id).unwrap(); let manifest = region.manifest_ctx.manifest().await; - assert_eq!(manifest.removed_files.removed_files[0].files.len(), 3); + let removed_file_num = manifest.removed_files.removed_files[0].files.len(); + assert!(removed_file_num > 0); let version = manifest.manifest_version; @@ -413,7 +414,10 @@ async fn test_gc_worker_basic_compact() { let gc_worker = create_gc_worker(&engine, regions, &file_ref_manifest, true).await; let report = gc_worker.run().await.unwrap(); - assert_eq!(report.deleted_files.get(®ion_id).unwrap().len(), 3,); + assert_eq!( + report.deleted_files.get(®ion_id).unwrap().len(), + removed_file_num + ); assert!(report.need_retry_regions.is_empty()); } @@ -476,7 +480,7 @@ async fn test_gc_worker_compact_with_ref() { let region = engine.get_region(region_id).unwrap(); let manifest = region.manifest_ctx.manifest().await; - assert_eq!(manifest.removed_files.removed_files[0].files.len(), 3); + assert!(!manifest.removed_files.removed_files[0].files.is_empty()); let version = manifest.manifest_version; diff --git a/src/mito2/src/worker/handle_compaction.rs b/src/mito2/src/worker/handle_compaction.rs index 15dc884a1e5..1f8a66887cd 100644 --- a/src/mito2/src/worker/handle_compaction.rs +++ b/src/mito2/src/worker/handle_compaction.rs @@ -131,6 +131,11 @@ impl RegionWorkerLoop { return; } let execution = request.execution.clone(); + // Whether this execution reduced the SST file count. The scheduler keeps + // draining while compaction makes progress; a rewrite that did not reduce + // files (e.g. its output was split into more files than its input) must end + // the chain to avoid a no-progress loop. + let made_progress = request.edit.files_to_remove.len() > request.edit.files_to_add.len(); region.version_control.apply_edit( Some(request.edit.clone()), @@ -167,6 +172,7 @@ impl RegionWorkerLoop { &execution, ®ion.manifest_ctx, self.schema_metadata_manager.clone(), + made_progress, ) .await; match transition { From de6a903cb196d1ad01aa314f8b138f53f8c27ff7 Mon Sep 17 00:00:00 2001 From: zhaiAohan <1573612521@qq.com> Date: Sun, 30 Aug 2026 05:26:00 +0000 Subject: [PATCH 19/22] feat: Add built-in daemon mode for standalone service (#8960) * feat: Add built-in daemon mode for standalone service Add an optional --daemon/-d flag to `greptimedb standalone start` to run the server as a background daemon detached from the shell session, similar to Redis `daemonize yes`. - Use the daemonize crate (Unix only) to fork/setsid and detach from the controlling terminal before the Tokio runtime is built. - stdin/stdout/stderr are redirected to /dev/null; logs continue to go to data_home/logs/ via the existing tracing appender. - On non-Unix platforms, --daemon is a no-op with a warning (runs in foreground). Closes #7314 Signed-off-by: tian1220A <1573612521@qq.com> * test(cmd): add integration test for standalone daemon mode Verify that standalone start --daemon detaches from the shell and brings up the HTTP listener, guarding against regressions where the daemon blocks in the foreground or crashes after forking. Signed-off-by: tian1220A <1573612521@qq.com> * fix(cmd): gate --daemon flag with #[cfg(unix)] and fix rustfmt Address review feedback to not expose --daemon/-d on non-Unix platforms. The daemon field is now gated with #[cfg(unix)], and a cross-platform is_daemon() accessor returns false on non-Unix so maybe_daemonize() stays unconditional. Also fix the rustfmt formatting issues that caused fmt-check to fail. Signed-off-by: tian1220A <1573612521@qq.com> * chore(cmd): declare daemonize dependency directly in cmd daemonize is only used by the cmd crate. Move it out of the workspace dependencies and declare it directly in the unix-only dependencies of src/cmd/Cargo.toml. Signed-off-by: tian1220A <1573612521@qq.com> --------- Signed-off-by: tian1220A <1573612521@qq.com> --- Cargo.lock | 10 +++ src/cmd/Cargo.toml | 1 + src/cmd/src/bin/greptime.rs | 58 ++++++++++++++-- src/cmd/src/error.rs | 10 ++- src/cmd/src/standalone.rs | 39 +++++++++++ src/cmd/tests/standalone_daemon.rs | 107 +++++++++++++++++++++++++++++ 6 files changed, 217 insertions(+), 8 deletions(-) create mode 100644 src/cmd/tests/standalone_daemon.rs diff --git a/Cargo.lock b/Cargo.lock index d77ec340a8d..8fe6d74dd2f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2167,6 +2167,7 @@ dependencies = [ "common-time", "common-version", "common-wal", + "daemonize", "datafusion", "datafusion-common", "datafusion-physical-plan", @@ -3570,6 +3571,15 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "daemonize" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab8bfdaacb3c887a54d41bdf48d3af8873b3f5566469f8ba21b92057509f116e" +dependencies = [ + "libc", +] + [[package]] name = "darling" version = "0.14.4" diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml index 579e5b1454c..8cb8cff4300 100644 --- a/src/cmd/Cargo.toml +++ b/src/cmd/Cargo.toml @@ -118,6 +118,7 @@ tonic.workspace = true tracing-appender.workspace = true [target.'cfg(unix)'.dependencies] +daemonize = "0.5" pprof = { version = "0.14", features = [ "flamegraph", ] } diff --git a/src/cmd/src/bin/greptime.rs b/src/cmd/src/bin/greptime.rs index f5f4a590487..2fa75ff4171 100644 --- a/src/cmd/src/bin/greptime.rs +++ b/src/cmd/src/bin/greptime.rs @@ -17,6 +17,8 @@ use clap::{Parser, Subcommand}; use cmd::datanode::builder::InstanceBuilder; +#[cfg(unix)] +use cmd::error::DaemonizeSnafu; use cmd::error::{InitTlsProviderSnafu, Result}; use cmd::options::GlobalOptions; use cmd::{App, cli, datanode, flownode, frontend, metasrv, standalone, user}; @@ -73,20 +75,24 @@ static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; #[cfg(debug_assertions)] fn main() -> Result<()> { use snafu::ResultExt; + + let cli = Command::parse(); + maybe_daemonize(&cli)?; + // Set the stack size to 8MB for the thread so it wouldn't overflow on large stack usage in debug mode // see https://github.com/GreptimeTeam/greptimedb/pull/4317 // and https://github.com/rust-lang/rust/issues/34283 std::thread::Builder::new() .name("main_spawn".to_string()) .stack_size(8 * 1024 * 1024) - .spawn(|| { + .spawn(move || { { tokio::runtime::Builder::new_multi_thread() .thread_stack_size(8 * 1024 * 1024) .enable_all() .build() .expect("Failed building the Runtime") - .block_on(main_body()) + .block_on(main_body(cli)) } }) .context(cmd::error::SpawnThreadSnafu)? @@ -95,15 +101,53 @@ fn main() -> Result<()> { } #[cfg(not(debug_assertions))] -#[tokio::main] -async fn main() -> Result<()> { - main_body().await +fn main() -> Result<()> { + let cli = Command::parse(); + maybe_daemonize(&cli)?; + + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .expect("Failed building the Runtime") + .block_on(main_body(cli)) } -async fn main_body() -> Result<()> { +async fn main_body(cli: Command) -> Result<()> { setup_human_panic(); install_default_crypto_provider().map_err(|msg| InitTlsProviderSnafu { msg }.build())?; - start(Command::parse()).await + start(cli).await +} + +/// If `--daemon` was passed for `standalone start`, fork into the background +/// before any Tokio runtime is built. On non-Unix platforms the `--daemon` +/// flag does not exist, so this is a no-op. +fn maybe_daemonize(cli: &Command) -> Result<()> { + let daemon = matches!(&cli.subcmd, SubCommand::Standalone(cmd) if cmd.is_daemon()); + if !daemon { + return Ok(()); + } + + #[cfg(unix)] + { + let cwd = std::env::current_dir().map_err(|e| { + DaemonizeSnafu { + msg: format!("failed to read current working directory: {e}"), + } + .build() + })?; + + daemonize::Daemonize::new() + .working_directory(cwd) + .start() + .map_err(|e| { + DaemonizeSnafu { + msg: format!("failed to daemonize: {e}"), + } + .build() + })?; + } + + Ok(()) } async fn start(cli: Command) -> Result<()> { diff --git a/src/cmd/src/error.rs b/src/cmd/src/error.rs index b80828d13ac..82273c0b67d 100644 --- a/src/cmd/src/error.rs +++ b/src/cmd/src/error.rs @@ -162,6 +162,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Failed to daemonize: {msg}"))] + Daemonize { + msg: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Unsupported selector type: {}", selector_type))] UnsupportedSelectorType { selector_type: String, @@ -361,7 +368,8 @@ impl ErrorExt for Error { Error::SerdeJson { .. } | Error::FileIo { .. } | Error::SpawnThread { .. } - | Error::InitTlsProvider { .. } => StatusCode::Unexpected, + | Error::InitTlsProvider { .. } + | Error::Daemonize { .. } => StatusCode::Unexpected, Error::Other { source, .. } => source.status_code(), diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index 2b5bd5e8cec..d628461ba75 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -152,6 +152,11 @@ impl Command { ) -> Result> { self.subcmd.load_options(global_options) } + + /// Whether the `standalone start` command requested daemonization. + pub fn is_daemon(&self) -> bool { + self.subcmd.is_daemon() + } } #[derive(Parser)] @@ -174,6 +179,12 @@ impl SubCommand { SubCommand::Start(cmd) => cmd.load_options(global_options), } } + + fn is_daemon(&self) -> bool { + match self { + SubCommand::Start(cmd) => cmd.is_daemon(), + } + } } pub struct Instance { @@ -291,9 +302,24 @@ pub struct StartCommand { /// The working home directory of this standalone instance. #[clap(long)] data_home: Option, + /// Run in the background as a daemon. + #[cfg(unix)] + #[clap(short, long)] + daemon: bool, } impl StartCommand { + /// Whether the `standalone start` command requested daemonization. + #[cfg(unix)] + pub(crate) fn is_daemon(&self) -> bool { + self.daemon + } + + #[cfg(not(unix))] + pub(crate) fn is_daemon(&self) -> bool { + false + } + /// Load the GreptimeDB options from various sources (command line, config file or env). pub fn load_options( &self, @@ -1356,6 +1382,19 @@ mod tests { assert_eq!(command.grpc_bind_addr.as_deref(), Some("127.0.0.1:34001")); } + #[cfg(unix)] + #[test] + fn test_parse_daemon_flag() { + let command = StartCommand::try_parse_from(["standalone", "--daemon"]).unwrap(); + assert!(command.is_daemon()); + + let command = StartCommand::try_parse_from(["standalone", "-d"]).unwrap(); + assert!(command.is_daemon()); + + let command = StartCommand::try_parse_from(["standalone"]).unwrap(); + assert!(!command.is_daemon()); + } + #[test] fn test_help_uses_grpc_option_names() { let mut cmd = StartCommand::command(); diff --git a/src/cmd/tests/standalone_daemon.rs b/src/cmd/tests/standalone_daemon.rs new file mode 100644 index 00000000000..2f4d1cd068e --- /dev/null +++ b/src/cmd/tests/standalone_daemon.rs @@ -0,0 +1,107 @@ +// 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::net::TcpStream; +use std::process::{Command, Stdio}; +use std::time::{Duration, Instant}; + +/// Stops the orphaned background daemon when the test finishes (including on +/// panic). The daemon's command line contains the unique `--data-home` path, so +/// we match on it to locate the process. +struct DaemonGuard(String); + +impl Drop for DaemonGuard { + fn drop(&mut self) { + let _ = Command::new("pkill") + .args(["-f", &self.0]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status(); + } +} + +#[test] +#[cfg(unix)] +fn test_standalone_daemon_mode_detaches_from_shell() { + let bin = env!("CARGO_BIN_EXE_greptime"); + let data_home = tempfile::tempdir().unwrap(); + let data_home_path = data_home.path().to_str().unwrap().to_string(); + + let http_addr = format!("127.0.0.1:{}", common_test_util::ports::get_port()); + let grpc_addr = format!("127.0.0.1:{}", common_test_util::ports::get_port()); + let mysql_addr = format!("127.0.0.1:{}", common_test_util::ports::get_port()); + let postgres_addr = format!("127.0.0.1:{}", common_test_util::ports::get_port()); + + let mut child = Command::new(bin) + .args([ + "standalone", + "start", + "--daemon", + "--data-home", + &data_home_path, + "--http-addr", + &http_addr, + "--grpc-bind-addr", + &grpc_addr, + "--mysql-addr", + &mysql_addr, + "--postgres-addr", + &postgres_addr, + ]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .unwrap(); + + let _guard = DaemonGuard(data_home_path); + + // When daemonization works, the parent forks and exits right away. When it + // does not, the process stays in the foreground and never returns. + let detach_timeout = Duration::from_secs(30); + let started = Instant::now(); + loop { + match child.try_wait().unwrap() { + Some(status) => { + assert!( + status.success(), + "daemon start should exit successfully, got {status:?}" + ); + break; + } + None => { + assert!( + started.elapsed() < detach_timeout, + "daemon mode should detach and return quickly, but it blocked for over {detach_timeout:?}" + ); + std::thread::sleep(Duration::from_millis(200)); + } + } + } + + // A successful detach is not enough: a broken implementation could fork, + // have the parent exit, then crash the child. Poll the HTTP address to + // confirm the daemon actually came up and is listening. + let ready_timeout = Duration::from_secs(30); + let started = Instant::now(); + loop { + if TcpStream::connect(&http_addr).is_ok() { + break; + } + assert!( + started.elapsed() < ready_timeout, + "daemon should come up and listen on {http_addr}, but it was unreachable for {ready_timeout:?}" + ); + std::thread::sleep(Duration::from_millis(200)); + } +} From d32cd77505d3cb304d66796b515c45ab42e739c6 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Sun, 30 Aug 2026 13:24:17 +0000 Subject: [PATCH 20/22] fix: postgres describe for more statements (#8974) * fix: postgres describe for more statements Signed-off-by: Ning Sun * fix: cover more show statements Signed-off-by: Ning Sun * fix: address review comments - add missing `clippy::too_many_arguments` allow on `query_from_information_schema_dataframe` (CI clippy failure) - take `&ShowKind` in the information-schema dataframe helper so `kind` is no longer cloned at every call site; only the WHERE arm (which needs an owned expression for `sql_to_expr`) clones internally - document why re-applying TQL explain formats never overwrites an existing value (per-query context state) Signed-off-by: Ning Sun * chore: trim comments to essentials Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- src/common/recordbatch/src/cursor.rs | 1 - src/frontend/src/instance.rs | 101 ++++++++ src/frontend/src/lib.rs | 2 + src/operator/src/statement.rs | 2 +- src/operator/src/statement/admin.rs | 171 ++++++++----- src/query/src/analyze.rs | 10 +- src/query/src/sql.rs | 249 +++++++++++++++---- src/servers/src/postgres/handler.rs | 83 ++++--- src/session/src/lib.rs | 5 +- tests-integration/tests/sql.rs | 343 +++++++++++++++++++++++++++ 10 files changed, 825 insertions(+), 142 deletions(-) diff --git a/src/common/recordbatch/src/cursor.rs b/src/common/recordbatch/src/cursor.rs index 60ae9d54fa2..125a68f4f13 100644 --- a/src/common/recordbatch/src/cursor.rs +++ b/src/common/recordbatch/src/cursor.rs @@ -47,7 +47,6 @@ impl RecordBatchStreamCursor { } } - /// Returns the schema of the underlying record batch stream. pub fn schema(&self) -> SchemaRef { self.schema.clone() } diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 412ff495f65..90650014e6a 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -59,6 +59,7 @@ use common_recordbatch::error::StreamTimeoutSnafu; use common_telemetry::logging::SlowQueryOptions; use common_telemetry::{debug, error, tracing}; use dashmap::DashMap; +use datafusion::dataframe::DataFrame; use datafusion::physical_plan::ExecutionPlan; use datafusion_expr::LogicalPlan; use futures::{Stream, StreamExt, future}; @@ -847,6 +848,22 @@ impl Instance { query_interceptor.pre_execute(stmt.as_ref(), Some(&plan), query_ctx.clone())?; + // TQL EXPLAIN/ANALYZE formats are consumed from the query context at + // execution time (see `optimize_physical_plan`); re-apply the side + // effect of `plan_tql` that was lost when the plan was built during + // Describe. `explain_format` is per-query state, so this never + // overwrites anything. + if let Some(Statement::Tql(tql)) = &stmt { + let format = match tql { + Tql::Explain(explain) => explain.format.as_ref(), + Tql::Analyze(analyze) => analyze.format.as_ref(), + Tql::Eval(_) => None, + }; + if let Some(format) = format { + query_ctx.set_explain_format(format.to_string()); + } + } + let query = stmt .as_ref() .map(|s| s.to_string()) @@ -928,6 +945,59 @@ impl Instance { vec![result] } + /// Builds the [`DataFrame`] for an information-schema-backed `SHOW` + /// statement; `None` for other statements. The future is boxed to keep + /// `do_describe_inner`'s state machine small. + fn show_statement_dataframe<'a>( + &'a self, + stmt: &'a Statement, + query_ctx: &'a QueryContextRef, + ) -> Pin>> + Send + 'a>> { + Box::pin(async move { + let engine = &self.query_engine; + let catalog_manager = self.catalog_manager(); + let ctx = query_ctx.clone(); + let dataframe = match stmt { + Statement::ShowDatabases(show) => { + query::sql::show_databases_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowTables(show) => { + query::sql::show_tables_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowViews(show) => { + query::sql::show_views_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowFlows(show) => { + query::sql::show_flows_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowColumns(show) => { + query::sql::show_columns_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowTableStatus(show) => { + query::sql::show_table_status_dataframe(show, engine, catalog_manager, ctx) + .await + } + Statement::ShowCharset(kind) => { + query::sql::show_charsets_dataframe(kind, engine, catalog_manager, ctx).await + } + Statement::ShowCollation(kind) => { + query::sql::show_collations_dataframe(kind, engine, catalog_manager, ctx).await + } + Statement::ShowIndex(show) => { + query::sql::show_index_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowRegion(show) => { + query::sql::show_region_dataframe(show, engine, catalog_manager, ctx).await + } + Statement::ShowProcesslist(show) => { + query::sql::show_processlist_dataframe(show, engine, catalog_manager, ctx).await + } + _ => return None, + }; + Some(dataframe) + }) + } + async fn do_describe_inner( &self, stmt: Statement, @@ -947,6 +1017,37 @@ impl Instance { let plannable = is_inner_plannable(&stmt) || matches!(&stmt, Statement::Explain(explain) if is_inner_plannable(explain.statement.as_ref())); + if let Statement::Tql(tql) = stmt { + // TQL produces a logical plan; describe it from the plan so the + // extended-protocol RowDescription matches the executed DataRows. + self.check_sql_permission(&Statement::Tql(tql.clone()), &query_ctx) + .await?; + let plan = self.statement_executor.plan_tql(tql, &query_ctx).await?; + return self + .query_engine + .describe(plan, query_ctx) + .await + .map(Some) + .context(error::DescribeStatementSnafu); + } + + // Describe SHOW statements from the same projection the executor builds. + if let Some(dataframe) = self + .show_statement_dataframe(&stmt, &query_ctx) + .await + .transpose() + .context(PlanStatementSnafu)? + { + self.check_sql_permission(&stmt, &query_ctx).await?; + let plan = dataframe.into_unoptimized_plan(); + return self + .query_engine + .describe(plan, query_ctx) + .await + .map(Some) + .context(error::DescribeStatementSnafu); + } + if plannable { self.check_sql_permission(&stmt, &query_ctx).await?; diff --git a/src/frontend/src/lib.rs b/src/frontend/src/lib.rs index c1702360739..1468fe7c0b2 100644 --- a/src/frontend/src/lib.rs +++ b/src/frontend/src/lib.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +#![recursion_limit = "256"] + pub mod error; pub mod events; pub mod frontend; diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 478f06e83f7..21c3ed61464 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -84,7 +84,7 @@ use table::table_reference::TableReference; pub use self::admin::{ AdminEventRecorderHandle, AdminFunctionLayer, AdminFunctionLayerRef, AdminFunctionRecordingLayer, AdminFunctionRequest, AdminFunctionResponse, AdminFunctionService, - AdminFunctionServiceRef, + AdminFunctionServiceRef, admin_output_schema, }; use self::set::{ set_bytea_output, set_datestyle, set_intervalstyle, set_timezone, validate_client_encoding, diff --git a/src/operator/src/statement/admin.rs b/src/operator/src/statement/admin.rs index 2b4a909e3cd..ff8b26915cc 100644 --- a/src/operator/src/statement/admin.rs +++ b/src/operator/src/statement/admin.rs @@ -19,6 +19,7 @@ use std::sync::Arc; use common_function::function::FunctionContext; use common_function::function_registry::{FUNCTION_REGISTRY, get_admin_function}; +use common_function::state::FunctionState; use common_query::Output; use common_recordbatch::{RecordBatch, RecordBatches}; use common_sql::convert::sql_value_to_value; @@ -77,6 +78,103 @@ struct CoreAdminFunctionService { query_engine: query::QueryEngineRef, } +/// Parts of an `ADMIN` call needed both for execution and schema derivation. +struct ResolvedAdminFunction { + admin_udf: datafusion_expr::ScalarUDF, + fn_name: String, + args: Vec, + arg_types: Vec, + ret_type: ArrowDataType, +} + +/// Resolves the function, parses its literal arguments and derives its +/// return type, without executing it. +fn resolve_admin_function( + stmt: &Admin, + query_ctx: &QueryContextRef, + state: Arc, +) -> Result { + let Admin::Func(func) = stmt; + // the function name should be in lower case. + let func_name = func.name.to_string().to_lowercase(); + let factory = get_admin_function(&func_name) + .or_else(|| FUNCTION_REGISTRY.get_function(&func_name)) + .context(error::AdminFunctionNotFoundSnafu { + name: func_name.clone(), + })?; + + let func_ctx = FunctionContext { + query_ctx: query_ctx.clone(), + state, + }; + + let admin_udf = factory.provide(func_ctx); + admin_udf + .as_async() + .context(error::AdminFunctionNotFoundSnafu { name: func_name })?; + + let fn_name = admin_udf.name().to_string(); + let signature = admin_udf.signature(); + + // Parse function arguments + let FunctionArguments::List(args) = &func.args else { + return error::BuildAdminFunctionArgsSnafu { + msg: format!("unsupported function args {} for {}", func.args, fn_name), + } + .fail(); + }; + let arg_values = args + .args + .iter() + .map(|arg| { + let FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(value))) = arg else { + return error::BuildAdminFunctionArgsSnafu { + msg: format!("unsupported function arg {arg} for {}", fn_name), + } + .fail(); + }; + Ok(&value.value) + }) + .collect::>>()?; + + let args = args_to_vector(&signature.type_signature, &arg_values, query_ctx)?; + let arg_types = args + .iter() + .map(|arg| arg.data_type().as_arrow_type()) + .collect::>(); + let ret_type = + admin_udf + .return_type(&arg_types) + .map_err(|e| error::Error::BuildAdminFunctionArgs { + msg: format!( + "Failed to get return type of admin function {}: {}", + fn_name, e + ), + })?; + + Ok(ResolvedAdminFunction { + admin_udf, + fn_name, + args, + arg_types, + ret_type, + }) +} + +/// Output schema of an `ADMIN` statement, mirroring what +/// [`CoreAdminFunctionService::execute`] produces. `None` if the function +/// or arguments are unresolvable; execution will then surface the error. +pub fn admin_output_schema(stmt: &Admin, query_ctx: &QueryContextRef) -> Option { + let resolved = + resolve_admin_function(stmt, query_ctx, Arc::new(FunctionState::default())).ok()?; + Some(Schema::new(vec![ColumnSchema::new( + // Use statement as the result column name + stmt.to_string(), + ConcreteDataType::from_arrow_type(&resolved.ret_type), + false, + )])) +} + impl CoreAdminFunctionService { fn new(query_engine: query::QueryEngineRef) -> Self { Self { query_engine } @@ -88,62 +186,23 @@ impl CoreAdminFunctionService { 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(); - let factory = get_admin_function(&func_name) - .or_else(|| FUNCTION_REGISTRY.get_function(&func_name)) - .context(error::AdminFunctionNotFoundSnafu { - name: func_name.clone(), - })?; - - let func_ctx = FunctionContext { - query_ctx: query_ctx.clone(), - state: self.query_engine.engine_state().function_state(), - }; - - let admin_udf = factory.provide(func_ctx); + let resolved = resolve_admin_function( + &stmt, + &query_ctx, + self.query_engine.engine_state().function_state(), + )?; + let ResolvedAdminFunction { + admin_udf, + fn_name, + args, + arg_types, + ret_type, + } = resolved; let admin_async_fn = admin_udf .as_async() - .context(error::AdminFunctionNotFoundSnafu { name: func_name })?; - - let fn_name = admin_udf.name(); - let signature = admin_udf.signature(); - - // Parse function arguments - let FunctionArguments::List(args) = &func.args else { - return error::BuildAdminFunctionArgsSnafu { - msg: format!("unsupported function args {} for {}", func.args, fn_name), - } - .fail(); - }; - let arg_values = args - .args - .iter() - .map(|arg| { - let FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(value))) = arg else { - return error::BuildAdminFunctionArgsSnafu { - msg: format!("unsupported function arg {arg} for {}", fn_name), - } - .fail(); - }; - Ok(&value.value) - }) - .collect::>>()?; - - let args = args_to_vector(&signature.type_signature, &arg_values, &query_ctx)?; - let arg_types = args - .iter() - .map(|arg| arg.data_type().as_arrow_type()) - .collect::>(); - let ret_type = admin_udf.return_type(&arg_types).map_err(|e| { - error::Error::BuildAdminFunctionArgs { - msg: format!( - "Failed to get return type of admin function {}: {}", - fn_name, e - ), - } - })?; + .context(error::AdminFunctionNotFoundSnafu { + name: fn_name.clone(), + })?; // Convert arguments to DataFusion ColumnarValue format let columnar_args: Vec = args @@ -174,9 +233,7 @@ impl CoreAdminFunctionService { let result_columnar = admin_async_fn .invoke_async_with_args(func_args) .await - .with_context(|_| ExecuteAdminFunctionSnafu { - msg: fn_name.to_string(), - })?; + .with_context(|_| ExecuteAdminFunctionSnafu { msg: fn_name })?; // Convert result back to VectorRef let result_columnar: common_query::prelude::ColumnarValue = diff --git a/src/query/src/analyze.rs b/src/query/src/analyze.rs index b1522fc05d5..2ab3ec5176e 100644 --- a/src/query/src/analyze.rs +++ b/src/query/src/analyze.rs @@ -46,13 +46,9 @@ const STAGE: &str = "stage"; const NODE: &str = "node"; const PLAN: &str = "plan"; -/// Returns the fixed output schema of [`DistAnalyzeExec`]: -/// (`stage`: UInt32, `node`: UInt32, `plan`: Utf8). -/// -/// Exposed so protocol-level `Describe` handlers can advertise the schema -/// clients will actually receive (the physical plan is rewritten in -/// `optimize_physical_plan` and its schema differs from the DataFusion -/// `Analyze` logical plan schema). +/// Fixed output schema of [`DistAnalyzeExec`], for `Describe` handlers: +/// execution rewrites the plan in `optimize_physical_plan`, so this schema +/// differs from the logical `Analyze` plan's. pub fn dist_analyze_output_schema() -> SchemaRef { SchemaRef::new(Schema::new(vec![ Field::new(STAGE, DataType::UInt32, true), diff --git a/src/query/src/sql.rs b/src/query/src/sql.rs index b2d979c6664..93e74ca33ec 100644 --- a/src/query/src/sql.rs +++ b/src/query/src/sql.rs @@ -41,6 +41,7 @@ use common_recordbatch::RecordBatches; use common_recordbatch::adapter::RecordBatchStreamAdapter; use common_time::Timestamp; use common_time::timezone::get_timezone; +use datafusion::dataframe::DataFrame; use datafusion::prelude::SessionContext; use datafusion_expr::{Expr, SortExpr, col, lit}; use datatypes::prelude::*; @@ -109,7 +110,7 @@ const INDEX_KEY_NAME_COLUMN: &str = "Key_name"; const INDEX_SEQ_IN_INDEX_COLUMN: &str = "Seq_in_index"; const INDEX_COLUMN_NAME_COLUMN: &str = "Column_name"; -static DESCRIBE_TABLE_OUTPUT_SCHEMA: Lazy> = Lazy::new(|| { +pub static DESCRIBE_TABLE_OUTPUT_SCHEMA: Lazy> = Lazy::new(|| { Arc::new(Schema::new(vec![ ColumnSchema::new( COLUMN_NAME_COLUMN, @@ -178,6 +179,18 @@ pub async fn show_databases( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { + let dataframe = + show_databases_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW DATABASES` without executing it. +pub async fn show_databases_dataframe( + stmt: &ShowDatabases, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { let projects = if stmt.full { vec![ (schemata::SCHEMA_NAME, SCHEMAS_COLUMN), @@ -191,7 +204,7 @@ pub async fn show_databases( let like_field = Some(schemata::SCHEMA_NAME); let sort = vec![col(schemata::SCHEMA_NAME).sort(true, true)]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -201,7 +214,7 @@ pub async fn show_databases( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -249,6 +262,45 @@ async fn query_from_information_schema_table( sort: Vec, kind: ShowKind, ) -> Result { + let dataframe = query_from_information_schema_dataframe( + query_engine, + catalog_manager, + query_ctx, + table_name, + select, + projects, + filters, + like_field, + sort, + &kind, + ) + .await?; + dataframe_to_output(dataframe).await +} + +async fn dataframe_to_output(dataframe: DataFrame) -> Result { + let stream = dataframe.execute_stream().await?; + Ok(Output::new_with_stream(Box::pin( + RecordBatchStreamAdapter::try_new(stream).context(error::CreateRecordBatchSnafu)?, + ))) +} + +/// Builds the [`DataFrame`] for a `SHOW` statement without executing it, +/// so `Describe` handlers can derive the output schema from the same +/// projection the executor uses. +#[allow(clippy::too_many_arguments)] +async fn query_from_information_schema_dataframe( + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, + table_name: &str, + select: Vec, + projects: Vec<(&str, &str)>, + filters: Vec, + like_field: Option<&str>, + sort: Vec, + kind: &ShowKind, +) -> Result { let table = catalog_manager .table( query_ctx.current_catalog(), @@ -336,7 +388,7 @@ async fn query_from_information_schema_table( .expect("Must be the datafusion planner"); let filter = planner - .sql_to_expr(filter, dataframe.schema(), false, query_ctx) + .sql_to_expr(filter.clone(), dataframe.schema(), false, query_ctx) .await?; // Apply the `where` clause filters @@ -344,11 +396,7 @@ async fn query_from_information_schema_table( } }; - let stream = dataframe.execute_stream().await?; - - Ok(Output::new_with_stream(Box::pin( - RecordBatchStreamAdapter::try_new(stream).context(error::CreateRecordBatchSnafu)?, - ))) + Ok(dataframe) } /// Execute `SHOW COLUMNS` statement. @@ -358,8 +406,19 @@ pub async fn show_columns( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { - let schema_name = if let Some(database) = stmt.database { - database + let dataframe = show_columns_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW COLUMNS` without executing it. +pub async fn show_columns_dataframe( + stmt: &ShowColumns, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { + let schema_name = if let Some(database) = &stmt.database { + database.clone() } else { query_ctx.current_schema() }; @@ -397,7 +456,7 @@ pub async fn show_columns( let like_field = Some(columns::COLUMN_NAME); let sort = vec![col(columns::COLUMN_NAME).sort(true, true)]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -407,7 +466,7 @@ pub async fn show_columns( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -419,8 +478,19 @@ pub async fn show_index( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { - let schema_name = if let Some(database) = stmt.database { - database + let dataframe = show_index_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW INDEX` without executing it. +pub async fn show_index_dataframe( + stmt: &ShowIndex, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { + let schema_name = if let Some(database) = &stmt.database { + database.clone() } else { query_ctx.current_schema() }; @@ -471,7 +541,7 @@ pub async fn show_index( col(statistics::SEQ_IN_INDEX).sort(true, true), ]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -481,7 +551,7 @@ pub async fn show_index( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -493,8 +563,19 @@ pub async fn show_region( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { - let schema_name = if let Some(database) = stmt.database { - database + let dataframe = show_region_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW REGION` without executing it. +pub async fn show_region_dataframe( + stmt: &ShowRegion, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { + let schema_name = if let Some(database) = &stmt.database { + database.clone() } else { query_ctx.current_schema() }; @@ -517,7 +598,7 @@ pub async fn show_region( col(columns::PEER_ID).sort(true, true), ]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -527,7 +608,7 @@ pub async fn show_region( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -539,8 +620,19 @@ pub async fn show_tables( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { - let schema_name = if let Some(database) = stmt.database { - database + let dataframe = show_tables_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for [`ShowTables`] without executing it. +pub async fn show_tables_dataframe( + stmt: &ShowTables, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { + let schema_name = if let Some(database) = &stmt.database { + database.clone() } else { query_ctx.current_schema() }; @@ -564,15 +656,16 @@ pub async fn show_tables( // Transform the WHERE clause for backward compatibility: // Replace "Tables" with "Tables_in_{schema}" to support old queries - let kind = match stmt.kind { - ShowKind::Where(mut filter) => { + let rewritten_kind = match &stmt.kind { + ShowKind::Where(filter) => { + let mut filter = filter.clone(); replace_column_in_expr(&mut filter, "Tables", &tables_column); ShowKind::Where(filter) } - other => other, + kind => kind.clone(), }; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -582,7 +675,7 @@ pub async fn show_tables( filters, like_field, sort, - kind, + &rewritten_kind, ) .await } @@ -594,8 +687,20 @@ pub async fn show_table_status( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { - let schema_name = if let Some(database) = stmt.database { - database + let dataframe = + show_table_status_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for [`ShowTableStatus`] without executing it. +pub async fn show_table_status_dataframe( + stmt: &ShowTableStatus, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { + let schema_name = if let Some(database) = &stmt.database { + database.clone() } else { query_ctx.current_schema() }; @@ -629,7 +734,7 @@ pub async fn show_table_status( let like_field = Some(tables::TABLE_NAME); let sort = vec![col(tables::TABLE_NAME).sort(true, true)]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -639,7 +744,7 @@ pub async fn show_table_status( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -651,6 +756,18 @@ pub async fn show_collations( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { + let dataframe = + show_collations_dataframe(&kind, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW COLLATION` without executing it. +pub async fn show_collations_dataframe( + kind: &ShowKind, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { // Refer to https://dev.mysql.com/doc/refman/8.0/en/show-collation.html let projects = vec![ ("collation_name", "Collation"), @@ -665,7 +782,7 @@ pub async fn show_collations( let like_field = Some("collation_name"); let sort = vec![]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -687,6 +804,18 @@ pub async fn show_charsets( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { + let dataframe = + show_charsets_dataframe(&kind, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW CHARSET` without executing it. +pub async fn show_charsets_dataframe( + kind: &ShowKind, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { // Refer to https://dev.mysql.com/doc/refman/8.0/en/show-character-set.html let projects = vec![ ("character_set_name", "Charset"), @@ -699,7 +828,7 @@ pub async fn show_charsets( let like_field = Some("character_set_name"); let sort = vec![]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -916,8 +1045,19 @@ pub async fn show_views( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { - let schema_name = if let Some(database) = stmt.database { - database + let dataframe = show_views_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for [`ShowViews`] without executing it. +pub async fn show_views_dataframe( + stmt: &ShowViews, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { + let schema_name = if let Some(database) = &stmt.database { + database.clone() } else { query_ctx.current_schema() }; @@ -930,7 +1070,7 @@ pub async fn show_views( let like_field = Some(tables::TABLE_NAME); let sort = vec![col(tables::TABLE_NAME).sort(true, true)]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -940,7 +1080,7 @@ pub async fn show_views( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -952,12 +1092,23 @@ pub async fn show_flows( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { + let dataframe = show_flows_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for [`ShowFlows`] without executing it. +pub async fn show_flows_dataframe( + stmt: &ShowFlows, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { let projects = vec![(flows::FLOW_NAME, FLOWS_COLUMN)]; let filters = vec![col(flows::TABLE_CATALOG).eq(lit(query_ctx.current_catalog()))]; let like_field = Some(flows::FLOW_NAME); let sort = vec![col(flows::FLOW_NAME).sort(true, true)]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx, @@ -967,7 +1118,7 @@ pub async fn show_flows( filters, like_field, sort, - stmt.kind, + &stmt.kind, ) .await } @@ -1340,6 +1491,18 @@ pub async fn show_processlist( catalog_manager: &CatalogManagerRef, query_ctx: QueryContextRef, ) -> Result { + let dataframe = + show_processlist_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?; + dataframe_to_output(dataframe).await +} + +/// Builds the [`DataFrame`] for `SHOW PROCESSLIST` without executing it. +pub async fn show_processlist_dataframe( + stmt: &ShowProcessList, + query_engine: &QueryEngineRef, + catalog_manager: &CatalogManagerRef, + query_ctx: QueryContextRef, +) -> Result { let projects = if stmt.full { vec![ (process_list::ID, "Id"), @@ -1367,17 +1530,17 @@ pub async fn show_processlist( }; let like_field = None; let sort = vec![col("id").sort(true, true)]; - query_from_information_schema_table( + query_from_information_schema_dataframe( query_engine, catalog_manager, query_ctx.clone(), "process_list", vec![], - projects.clone(), + projects, filters, like_field, sort, - ShowKind::All, + &ShowKind::All, ) .await } diff --git a/src/servers/src/postgres/handler.rs b/src/servers/src/postgres/handler.rs index 433eb11cdad..aca4362bd34 100644 --- a/src/servers/src/postgres/handler.rs +++ b/src/servers/src/postgres/handler.rs @@ -28,6 +28,7 @@ use datafusion_pg_catalog::sql::PostgresCompatibilityParser; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{Schema, SchemaRef}; use futures::{Sink, SinkExt, Stream, StreamExt, future, stream}; +use operator::statement::admin_output_schema; use pgwire::api::portal::{Format, Portal}; use pgwire::api::query::{ExtendedQueryHandler, SimpleQueryHandler}; use pgwire::api::results::{ @@ -43,6 +44,7 @@ use pgwire::messages::data::DataRow; use query::dist_analyze_output_schema; use query::planner::DfLogicalPlanner; use query::query_engine::DescribeResult; +use query::sql::DESCRIBE_TABLE_OUTPUT_SCHEMA; use session::Session; use session::context::QueryContextRef; use snafu::ResultExt; @@ -555,11 +557,8 @@ fn describe_fields( session: &Arc, ) -> PgWireResult> { match sql_plan { - // EXPLAIN ANALYZE: at execution time the physical plan is replaced with - // GreptimeDB's `DistAnalyzeExec` (see `optimize_physical_plan`), whose - // output schema (stage/node/plan) differs from the DataFusion `Analyze` - // logical plan schema (plan_type/plan). Describe with the schema the - // client will actually receive so the DataRow field count matches. + // Execution swaps in DistAnalyzeExec (stage/node/plan), whose schema + // differs from the logical `Analyze` plan's (plan_type/plan). SqlPlan::Plan(LogicalPlan::Analyze(_), _) => { let schema: Schema = Schema::try_from(dist_analyze_output_schema()).map_err(convert_err)?; @@ -658,17 +657,6 @@ fn describe_fields( ), ]), - // single column show statements - SqlPlan::Statement( - Statement::ShowTables(_) | Statement::ShowFlows(_) | Statement::ShowViews(_), - _, - ) => Ok(vec![FieldInfo::new( - "name".to_string(), - None, - None, - Type::TEXT, - format.format_for(0), - )]), #[cfg(feature = "enterprise")] SqlPlan::Statement(Statement::ShowTriggers(_), _) => Ok(vec![FieldInfo::new( "name".to_string(), @@ -690,22 +678,59 @@ fn describe_fields( Ok(vec![]) } } - // FETCH cursor: return the cursor's schema so the RowDescription - // matches the DataRow field count sent during Execute. + // Single column named after the variable (see `query::sql::show_variable`). + SqlPlan::Statement(Statement::ShowVariables(show), _) => Ok(vec![FieldInfo::new( + show.variable.to_string().to_uppercase(), + None, + None, + Type::TEXT, + format.format_for(0), + )]), + // Mirrors `query::sql::show_status` (currently always empty). + SqlPlan::Statement(Statement::ShowStatus(_), _) => Ok(vec![ + FieldInfo::new( + "Variable_name".to_string(), + None, + None, + Type::TEXT, + format.format_for(0), + ), + FieldInfo::new( + "Value".to_string(), + None, + None, + Type::TEXT, + format.format_for(1), + ), + ]), + SqlPlan::Statement(Statement::ShowSearchPath(_), _) => Ok(vec![FieldInfo::new( + "search_path".to_string(), + None, + None, + Type::TEXT, + format.format_for(0), + )]), + // Mirrors `query::sql::describe_table`. + SqlPlan::Statement(Statement::DescribeTable(_), _) => { + schema_to_pg(&DESCRIBE_TABLE_OUTPUT_SCHEMA, format, None).map_err(convert_err) + } + // Single column typed with the function's return type (see + // `operator::statement::admin_output_schema`). + SqlPlan::Statement(Statement::Admin(admin), _) => { + let query_ctx = session.new_query_context(); + match admin_output_schema(admin, &query_ctx) { + Some(schema) => schema_to_pg(&schema, format, None).map_err(convert_err), + // Unresolvable; execution will surface the error. + None => Ok(vec![]), + } + } + // Describe from the declared cursor's schema. SqlPlan::Statement(Statement::FetchCursor(fetch), _) => { let cursor_name = fetch.cursor_name.to_string(); match session.get_cursor(&cursor_name) { - Some(cursor) => { - // `cursor.schema()` is the GreptimeDB `SchemaRef` captured - // when the cursor was declared; `schema_to_pg` accepts it - // directly (no DataFusion -> GreptimeDB conversion needed). - schema_to_pg(&cursor.schema(), format, None).map_err(convert_err) - } - None => { - // Cursor not found (e.g. DECLARE hasn't executed yet in - // an extended-protocol batch). Return NoData as a fallback. - Ok(vec![]) - } + Some(cursor) => schema_to_pg(&cursor.schema(), format, None).map_err(convert_err), + // Cursor not declared yet; execution will error. + None => Ok(vec![]), } } _ => { diff --git a/src/session/src/lib.rs b/src/session/src/lib.rs index 44f3ee66891..c1414fc5a3b 100644 --- a/src/session/src/lib.rs +++ b/src/session/src/lib.rs @@ -111,10 +111,7 @@ impl Session { .into() } - /// Returns the cursor with the given name, if it exists. - /// - /// Cursors are stored in the session's mutable inner data and shared across - /// all query contexts created from this session. + /// Cursors are shared across query contexts created from this session. pub fn get_cursor(&self, name: &str) -> Option> { let guard = self.mutable_inner.read().unwrap(); guard.cursors.get(name).cloned() diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index 32a2881dd78..c95adb483ec 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -91,6 +91,8 @@ macro_rules! sql_tests { test_mysql_prepare_stmt_insert_timestamp, test_mysql_prepare_stmt_timezone, test_mysql_federated_prepare_stmt, + test_mysql_prepare_tql_and_show, + test_postgres_extended_query_row_returning_statements, test_declare_fetch_close_cursor, test_alter_update_on, ); @@ -1487,6 +1489,280 @@ fn assert_pg_numeric_range_error(error: tokio_postgres::Error) { assert_eq!("numeric_value_out_of_range", error.message()); } +pub async fn test_postgres_extended_query_row_returning_statements(store_type: StorageType) { + // Regression test for the tokio-postgres >= 0.7.14 DataRow/RowDescription + // mismatch: statements answered with NoData at Describe but emitting + // DataRows at Execute must describe their real output schema. + let (mut guard, fe_pg_server) = setup_pg_server(store_type, "test_pg_extended_row_stmts").await; + let addr = fe_pg_server.bind_addr().unwrap().to_string(); + + let (client, connection) = tokio_postgres::connect(&format!("postgres://{addr}/public"), NoTls) + .await + .unwrap(); + + let (tx, rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + connection.await.unwrap(); + tx.send(()).unwrap(); + }); + + client + .execute( + "CREATE TABLE demo_metrics (ts timestamp time index, val double, host string primary key skipping index)", + &[], + ) + .await + .unwrap(); + client + .execute( + "INSERT INTO demo_metrics (ts, host, val) VALUES (1000, 'host-a', 1.0), (2000, 'host-b', 2.0)", + &[], + ) + .await + .unwrap(); + + // ---- SHOW DATABASES: single `Database` column ---- + let rows = client.query("SHOW DATABASES", &[]).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(1, rows[0].columns().len()); + assert_eq!("Database", rows[0].columns()[0].name()); + assert!(rows.iter().any(|r| r.get::<_, String>(0) == "public")); + + // ---- SHOW FULL DATABASES: `Database` + `Options` columns ---- + let rows = client.query("SHOW FULL DATABASES", &[]).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(2, rows[0].columns().len()); + assert_eq!("Database", rows[0].columns()[0].name()); + assert_eq!("Options", rows[0].columns()[1].name()); + assert!(rows.iter().any(|r| r.get::<_, String>(0) == "public")); + + // ---- SHOW TABLES: single `Tables_in_` column ---- + let rows = client.query("SHOW TABLES", &[]).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(1, rows[0].columns().len()); + assert_eq!("Tables_in_public", rows[0].columns()[0].name()); + assert!(rows.iter().any(|r| r.get::<_, String>(0) == "demo_metrics")); + + // ---- SHOW FULL TABLES: `Tables_in_` + `Table_type` columns ---- + let rows = client.query("SHOW FULL TABLES", &[]).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(2, rows[0].columns().len()); + assert_eq!("Tables_in_public", rows[0].columns()[0].name()); + assert_eq!("Table_type", rows[0].columns()[1].name()); + + // ---- SHOW VIEWS / SHOW FLOWS: empty results, still described with one column ---- + let rows = client.query("SHOW VIEWS", &[]).await.unwrap(); + assert!(rows.is_empty()); + let stmt = client.prepare("SHOW VIEWS").await.unwrap(); + assert_eq!(1, stmt.columns().len()); + assert_eq!("Views", stmt.columns()[0].name()); + let rows = client.query("SHOW FLOWS", &[]).await.unwrap(); + assert!(rows.is_empty()); + let stmt = client.prepare("SHOW FLOWS").await.unwrap(); + assert_eq!(1, stmt.columns().len()); + assert_eq!("Flows", stmt.columns()[0].name()); + + // ---- SHOW TABLE STATUS: fixed eighteen-column schema ---- + let rows = client.query("SHOW TABLE STATUS", &[]).await.unwrap(); + assert!(!rows.is_empty()); + let names: Vec<&str> = rows[0].columns().iter().map(|c| c.name()).collect(); + assert_eq!( + vec![ + "Name", + "Engine", + "Version", + "Row_format", + "Rows", + "Avg_row_length", + "Data_length", + "Max_data_length", + "Index_length", + "Data_free", + "Auto_increment", + "Create_time", + "Update_time", + "Check_time", + "Collation", + "Checksum", + "Create_options", + "Comment", + ], + names + ); + + // ---- SHOW COLUMNS / SHOW FULL COLUMNS ---- + let rows = client + .query("SHOW COLUMNS FROM demo_metrics", &[]) + .await + .unwrap(); + assert_eq!(3, rows.len()); + let names: Vec<&str> = rows[0].columns().iter().map(|c| c.name()).collect(); + assert_eq!( + vec![ + "Field", + "Type", + "Null", + "Key", + "Default", + "Extra", + "Greptime_type" + ], + names + ); + let rows = client + .query("SHOW FULL COLUMNS FROM demo_metrics", &[]) + .await + .unwrap(); + assert_eq!(10, rows[0].columns().len()); + + // ---- SHOW CHARSET / SHOW COLLATION ---- + let rows = client.query("SHOW CHARSET", &[]).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(4, rows[0].columns().len()); + let rows = client.query("SHOW COLLATION", &[]).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(6, rows[0].columns().len()); + + // ---- SHOW INDEX: fixed fifteen-column schema ---- + let rows = client + .query("SHOW INDEX IN demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + let names: Vec<&str> = rows[0].columns().iter().map(|c| c.name()).collect(); + assert_eq!( + vec![ + "Table", + "Non_unique", + "Key_name", + "Seq_in_index", + "Column_name", + "Collation", + "Cardinality", + "Sub_part", + "Packed", + "Null", + "Index_type", + "Comment", + "Index_comment", + "Visible", + "Expression", + ], + names + ); + + // ---- SHOW REGION ---- + let rows = client + .query("SHOW REGION IN demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + assert_eq!(4, rows[0].columns().len()); + + // ---- SHOW SEARCH_PATH: single string column ---- + let rows = client.query("SHOW SEARCH_PATH", &[]).await.unwrap(); + assert_eq!(1, rows.len()); + assert_eq!(1, rows[0].columns().len()); + assert_eq!("search_path", rows[0].columns()[0].name()); + assert_eq!("public", rows[0].get::<_, String>(0)); + + // ---- SHOW VARIABLES: single column named after the variable ---- + let rows = client.query("SHOW VARIABLES timezone", &[]).await.unwrap(); + assert_eq!(1, rows.len()); + assert_eq!(1, rows[0].columns().len()); + assert_eq!("TIMEZONE", rows[0].columns()[0].name()); + let _ = rows[0].get::<_, String>(0); + + // ---- DESCRIBE TABLE: fixed six-column string schema ---- + let rows = client + .query("DESCRIBE TABLE demo_metrics", &[]) + .await + .unwrap(); + assert_eq!(3, rows.len()); + let names: Vec<&str> = rows[0].columns().iter().map(|c| c.name()).collect(); + assert_eq!( + vec!["Column", "Type", "Key", "Null", "Default", "Semantic Type"], + names + ); + // first column of each row is the column name; ensure values decode as TEXT + let columns: Vec = rows.iter().map(|r| r.get::<_, String>(0)).collect(); + assert!(columns.contains(&"ts".to_string())); + assert!(columns.contains(&"val".to_string())); + assert!(columns.contains(&"host".to_string())); + + // ---- ADMIN function: single column named after the statement ---- + let rows = client + .query("ADMIN flush_table('demo_metrics')", &[]) + .await + .unwrap(); + assert_eq!(1, rows.len()); + assert_eq!(1, rows[0].columns().len()); + assert!(rows[0].columns()[0].name().contains("flush_table")); + + // ---- TQL EVAL / EXPLAIN / ANALYZE: described from the planned query ---- + let rows = client + .query("TQL EVAL (0, 3000, '1s') demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + let names: Vec<&str> = rows[0].columns().iter().map(|c| c.name()).collect(); + assert_eq!(vec!["ts", "val", "host"], names); + + let rows = client + .query("TQL EXPLAIN (0, 3000, '1s') demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + assert_eq!(2, rows[0].columns().len()); + + let rows = client + .query("TQL ANALYZE (0, 3000, '1s') demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + assert_eq!(3, rows[0].columns().len()); + + // FORMAT JSON variants keep working through the plan-based execution path + let rows = client + .query("TQL EXPLAIN FORMAT JSON (0, 3000, '1s') demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + let rows = client + .query("TQL ANALYZE FORMAT JSON (0, 3000, '1s') demo_metrics", &[]) + .await + .unwrap(); + assert!(!rows.is_empty()); + + // ---- the same statements also work over the simple query protocol ---- + for sql in [ + "SHOW DATABASES", + "SHOW FULL TABLES", + "SHOW TABLE STATUS", + "SHOW COLUMNS FROM demo_metrics", + "SHOW CHARSET", + "SHOW COLLATION", + "SHOW INDEX IN demo_metrics", + "SHOW REGION IN demo_metrics", + "SHOW SEARCH_PATH", + "DESCRIBE TABLE demo_metrics", + "ADMIN flush_table('demo_metrics')", + "TQL EVAL (0, 3000, '1s') demo_metrics", + ] { + let msgs = client.simple_query(sql).await.unwrap(); + assert!( + msgs.iter().any(|m| matches!(m, SimpleQueryMessage::Row(_))), + "simple query {sql} should return rows" + ); + } + + drop(client); + rx.await.unwrap(); + + let _ = fe_pg_server.shutdown().await; + guard.remove_all().await; +} + pub async fn test_postgres_explain_bind_parameter(store_type: StorageType) { // Regression test for #8029: EXPLAIN / EXPLAIN ANALYZE must accept bind // parameters over the Postgres extended query protocol. @@ -1890,6 +2166,73 @@ pub async fn test_mysql_federated_prepare_stmt(store_type: StorageType) { guard.remove_all().await; } +pub async fn test_mysql_prepare_tql_and_show(store_type: StorageType) { + // `do_describe` now plans TQL and information-schema-backed SHOW + // statements, so MySQL prepared statements derive their column metadata + // from the planned query and execute the plan directly. This exercises + // that path end-to-end. + common_telemetry::init_default_ut_logging(); + + let (mut guard, fe_mysql_server) = + setup_mysql_server(store_type, "test_mysql_prepare_tql_and_show").await; + let addr = fe_mysql_server.bind_addr().unwrap().to_string(); + + let pool = MySqlPoolOptions::new() + .max_connections(2) + .connect(&format!("mysql://{addr}/public")) + .await + .unwrap(); + + sqlx::query( + "CREATE TABLE demo_metrics (ts timestamp time index, val double, host string primary key)", + ) + .execute(&pool) + .await + .unwrap(); + sqlx::query("INSERT INTO demo_metrics (ts, host, val) VALUES (1000, 'host-a', 1.0), (2000, 'host-b', 2.0)") + .execute(&pool) + .await + .unwrap(); + + // sqlx::query uses the binary prepared statement protocol + // (COM_STMT_PREPARE + COM_STMT_EXECUTE). + let rows = sqlx::query("TQL EVAL (0, 3000, '1s') demo_metrics") + .fetch_all(&pool) + .await + .unwrap(); + assert!(!rows.is_empty()); + assert_eq!(3, rows[0].columns().len()); + + let rows = sqlx::query("TQL ANALYZE (0, 3000, '1s') demo_metrics") + .fetch_all(&pool) + .await + .unwrap(); + assert!(!rows.is_empty()); + + // SHOW statements share the same describe path; prepared SHOW FULL + // TABLES must report both columns. + let rows = sqlx::query("SHOW TABLES").fetch_all(&pool).await.unwrap(); + assert!(!rows.is_empty()); + assert_eq!(1, rows[0].columns().len()); + + let rows = sqlx::query("SHOW FULL TABLES") + .fetch_all(&pool) + .await + .unwrap(); + assert!(!rows.is_empty()); + assert_eq!(2, rows[0].columns().len()); + + let rows = sqlx::query("SHOW DATABASES") + .fetch_all(&pool) + .await + .unwrap(); + assert!(!rows.is_empty()); + assert_eq!(1, rows[0].columns().len()); + + let _ = fe_mysql_server.shutdown().await; + guard.remove_all().await; +} + pub async fn test_postgres_array_types(store_type: StorageType) { let (mut guard, fe_pg_server) = setup_pg_server(store_type, "test_postgres_array_types").await; let addr = fe_pg_server.bind_addr().unwrap().to_string(); From 6504af641ef3e2a6269c4fc64a3ff80e843766b0 Mon Sep 17 00:00:00 2001 From: localhost Date: Mon, 31 Aug 2026 06:18:20 +0000 Subject: [PATCH 21/22] fix: increase system disk size to 50 GiB for ECS instances (#8986) --- .github/runner-scale-sets/query-regression/README.md | 6 +++--- .github/scripts/aliyun-ecs-runner-provision.py | 11 +++++++---- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/.github/runner-scale-sets/query-regression/README.md b/.github/runner-scale-sets/query-regression/README.md index ce447dbdd0a..c0e34a4c40d 100644 --- a/.github/runner-scale-sets/query-regression/README.md +++ b/.github/runner-scale-sets/query-regression/README.md @@ -49,7 +49,7 @@ Configuration lives in repository variables/secrets: | vars | `ALIYUN_ECS_INSTANCE_TYPE` | Dedicated (non-burstable, non-shared) instance family. Prefer 32 GiB (e.g. `ecs.g8i.2xlarge`); `ecs.c9i.2xlarge` is 8c16g and nightly thin-LTO of greptime peaks above that. Both base and candidate clusters run on the same machine, so noisy neighbors break thresholds. | | vars | `QUERY_REGRESSION_ECS_IMAGE_ID` | Custom image built by `ecs-image/build-ecs-image.py`. | -The system disk is 40 GiB, which covers the image, a 16 GiB swapfile, the +The system disk is 50 GiB, which covers the image, a 16 GiB swapfile, the checkout, and cold build caches (target dir, cargo registry, sccache). ENOSPC stops the runner itself from writing logs, which GitHub reports as `The operation was canceled` with no telemetry, indistinguishable from a @@ -164,7 +164,7 @@ reintroduce those actions unless the image contract changes. Each run provisions its own ECS instance, so overlapping dispatches proceed in parallel. There is no workflow `concurrency` group. All Cargo state (`CARGO_HOME` including registry/git, `CARGO_TARGET_DIR`, sccache, cache -metadata) lives on the 40 GiB system disk and is discarded with the VM. +metadata) lives on the 50 GiB system disk and is discarded with the VM. `RUSTUP_HOME=/opt/rustup` and `/opt/cargo/bin` are image-owned. The runner sets `RUSTC_WRAPPER=/usr/local/bin/sccache`, `SCCACHE_DIR=/home/runner/.cache/sccache`, `SCCACHE_CACHE_SIZE=10G`, and @@ -173,7 +173,7 @@ at 10G. Base and candidate builds share the target on that disk; Cargo fingerprints invalidate source and dependency changes. The workflow reports `du` / `df` in telemetry. It does not try to reclaim -space across runs: a cold 40 GiB disk that fills up fails the build. +space across runs: a cold 50 GiB disk that fills up fails the build. ## Future optional phases diff --git a/.github/scripts/aliyun-ecs-runner-provision.py b/.github/scripts/aliyun-ecs-runner-provision.py index 74a69681f3d..307ee336ebf 100644 --- a/.github/scripts/aliyun-ecs-runner-provision.py +++ b/.github/scripts/aliyun-ecs-runner-provision.py @@ -80,6 +80,9 @@ POLL_INTERVAL_SECONDS = 5 # to the cold build caches. SWAP_FILE = "/swapfile" SWAP_SIZE_GIB = 16 +# Image (~12G), 16G swap, checkout, base+candidate data homes, and cold +# build caches. Deleted with the instance. +SYSTEM_DISK_GIB = 50 @dataclass(frozen=True) @@ -367,10 +370,10 @@ def run_instance(client, config: ProvisionConfig, user_data: str) -> str: spot_strategy="NoSpot", internet_charge_type="PayByTraffic", internet_max_bandwidth_out=100, - # Holds the image (~12G), the 16G swapfile, the source checkout, the - # base+candidate cluster data homes, and cold build caches. Deleted - # with the instance. - system_disk=ecs_models.RunInstancesRequestSystemDisk(category="cloud_essd", size="40"), + system_disk=ecs_models.RunInstancesRequestSystemDisk( + category="cloud_essd", + size=str(SYSTEM_DISK_GIB), + ), user_data=user_data, tag=[ ecs_models.RunInstancesRequestTag(key=MANAGED_BY_TAG_KEY, value=MANAGED_BY_TAG_VALUE), From 51b94bb73f2e3d6335ba8b570b1b89c14dc0aaff Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 31 Aug 2026 07:02:44 +0000 Subject: [PATCH 22/22] fix(cmd): gate daemon integration test on Unix (#8987) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/cmd/tests/standalone_daemon.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/cmd/tests/standalone_daemon.rs b/src/cmd/tests/standalone_daemon.rs index 2f4d1cd068e..5b7c7427d5f 100644 --- a/src/cmd/tests/standalone_daemon.rs +++ b/src/cmd/tests/standalone_daemon.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +#![cfg(unix)] + use std::net::TcpStream; use std::process::{Command, Stdio}; use std::time::{Duration, Instant};