From d90cca4b752b34d7ccdb54fe8445ce3bf832e5e0 Mon Sep 17 00:00:00 2001 From: Gaurav Shokeen Date: Fri, 7 Aug 2026 19:07:00 +0530 Subject: [PATCH] fix(prometheus): custom column remote reads (#8659) * fix(prometheus): custom column remote reads Resolve timestamp and value column names from the table schema and carry them through query planning and result conversion. Add a remote-read regression test covering custom_ts and custom_value. Signed-off-by: grezzko * fix: resolve remote-read value columns safely Prefer the sole field for custom schemas and greptime_value for multi-field tables. Reject ambiguous schemas and add regression tests. Signed-off-by: grezzko --------- Signed-off-by: grezzko Signed-off-by: Lei, HUANG Co-authored-by: Lei, HUANG --- src/frontend/src/error.rs | 11 ++ src/frontend/src/instance.rs | 213 +++++++++++++++++++++++- src/frontend/src/instance/prom_store.rs | 107 ++++++++++-- src/servers/src/prom_store.rs | 90 +++++++--- 4 files changed, 378 insertions(+), 43 deletions(-) diff --git a/src/frontend/src/error.rs b/src/frontend/src/error.rs index 144fb96ab9..9a70b71661 100644 --- a/src/frontend/src/error.rs +++ b/src/frontend/src/error.rs @@ -164,6 +164,16 @@ pub enum Error { location: Location, }, + #[snafu(display( + "Ambiguous value column in table '{table_name}', candidates: {field_columns:?}" + ))] + AmbiguousValueColumn { + table_name: String, + field_columns: Vec, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to collect recordbatch"))] CollectRecordbatch { #[snafu(implicit)] @@ -377,6 +387,7 @@ impl ErrorExt for Error { | Error::IllegalPrimaryKeysDef { .. } | Error::SchemaExists { .. } | Error::ColumnNotFound { .. } + | Error::AmbiguousValueColumn { .. } | Error::UnsupportedFormat { .. } | Error::IllegalAuthConfig { .. } | Error::ColumnNoneDefaultValue { .. } diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 14acf87787..23090ef932 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -1675,7 +1675,9 @@ mod tests { use std::time::Duration; use api::prom_store::remote::label_matcher::Type as PromMatcherType; - use api::prom_store::remote::{LabelMatcher, Query as RemoteQuery, ReadRequest}; + use api::prom_store::remote::{ + Label, LabelMatcher, Query as RemoteQuery, ReadRequest, ReadResponse, Sample, + }; use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse}; use auth::{ DASHBOARD_DELETE, DASHBOARD_QUERY, DASHBOARD_SAVE, JAEGER_QUERY, PIPELINE_DELETE, @@ -1684,7 +1686,7 @@ mod tests { use catalog::process_manager::{ProcessManager, QueryStatement, SlowQueryTimer}; use common_base::Plugins; use common_catalog::consts::DEFAULT_PRIVATE_SCHEMA_NAME; - use common_error::ext::{BoxedError, PlainError}; + use common_error::ext::{BoxedError, ErrorExt, PlainError}; use common_error::status_code::StatusCode; use common_event_recorder::{Event, EventRecorder, EventTypeFilter, EventTypeFilterRef}; use common_frontend::slow_query_event::SlowQueryEvent; @@ -1695,6 +1697,7 @@ mod tests { use common_meta::rpc::procedure::{ MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse, }; + use common_query::prelude::greptime_value; use common_query::{Output, OutputMeta}; use common_recordbatch::{ OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream, @@ -1706,8 +1709,12 @@ mod tests { use datafusion_expr::{LogicalPlanBuilder, LogicalTableSource}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema as GtSchema, SchemaRef as GtSchemaRef}; - use datatypes::vectors::{StringVector, TimestampNanosecondVector, VectorRef}; + use datatypes::vectors::{ + Float64Vector, StringVector, TimestampMillisecondVector, TimestampNanosecondVector, + VectorRef, + }; use log_query::LogQuery; + use prost::Message; use query::query_engine::options::QueryOptions; use servers::query_handler::{ DashboardHandler, JaegerQueryHandler, LogQueryHandler, PipelineHandler, PipelineHandlerRef, @@ -2726,6 +2733,206 @@ mod tests { assert_eq!(CheckedAction { action, targets }, checker.take_check()); } + #[tokio::test] + async fn test_prom_remote_read_with_custom_timestamp_and_value_columns() -> TestResult<()> { + let schema = Arc::new(GtSchema::new(vec![ + ColumnSchema::new( + "custom_ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("custom_value", ConcreteDataType::float64_datatype(), false), + ])); + let recordbatch = RecordBatch::new( + schema, + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as VectorRef, + Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as VectorRef, + ], + ) + .unwrap(); + let instance = test_instance_with_tables( + MemTable::table("custom_metric", recordbatch), + test_table(1025, "target")?, + ) + .await?; + + let response = PromStoreProtocolHandler::read( + &instance, + ReadRequest { + queries: vec![RemoteQuery { + start_timestamp_ms: 1500, + end_timestamp_ms: 2500, + matchers: vec![LabelMatcher { + r#type: PromMatcherType::Eq as i32, + name: servers::prom_store::METRIC_NAME_LABEL.to_string(), + value: "custom_metric".to_string(), + }], + ..Default::default() + }], + ..Default::default() + }, + test_query_ctx(1), + ) + .await + .unwrap(); + let body = servers::prom_store::snappy_decompress(&response.body).unwrap(); + let response = ReadResponse::decode(body.as_slice()).unwrap(); + + assert_eq!(1, response.results.len()); + assert_eq!(1, response.results[0].timeseries.len()); + let timeseries = &response.results[0].timeseries[0]; + assert_eq!( + vec![Label { + name: servers::prom_store::METRIC_NAME_LABEL.to_string(), + value: "custom_metric".to_string(), + }], + timeseries.labels + ); + assert_eq!( + vec![Sample { + value: 2.0, + timestamp: 2000, + }], + timeseries.samples + ); + + Ok(()) + } + + #[tokio::test] + async fn test_prom_remote_read_prefers_default_value_column() -> TestResult<()> { + let schema = Arc::new(GtSchema::new(vec![ + ColumnSchema::new( + "custom_ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("extra_field", ConcreteDataType::float64_datatype(), false), + ColumnSchema::new( + greptime_value(), + ConcreteDataType::float64_datatype(), + false, + ), + ])); + let recordbatch = RecordBatch::new( + schema, + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as VectorRef, + Arc::new(Float64Vector::from_vec(vec![99.0, 99.0, 99.0])) as VectorRef, + Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as VectorRef, + ], + ) + .unwrap(); + let instance = test_instance_with_tables( + MemTable::table("multi_field_metric", recordbatch), + test_table(1025, "target")?, + ) + .await?; + + let response = PromStoreProtocolHandler::read( + &instance, + ReadRequest { + queries: vec![RemoteQuery { + start_timestamp_ms: 1500, + end_timestamp_ms: 2500, + matchers: vec![LabelMatcher { + r#type: PromMatcherType::Eq as i32, + name: servers::prom_store::METRIC_NAME_LABEL.to_string(), + value: "multi_field_metric".to_string(), + }], + ..Default::default() + }], + ..Default::default() + }, + test_query_ctx(1), + ) + .await + .unwrap(); + let body = servers::prom_store::snappy_decompress(&response.body).unwrap(); + let response = ReadResponse::decode(body.as_slice()).unwrap(); + + assert_eq!(1, response.results.len()); + assert_eq!(1, response.results[0].timeseries.len()); + let timeseries = &response.results[0].timeseries[0]; + assert_eq!( + vec![ + Label { + name: servers::prom_store::METRIC_NAME_LABEL.to_string(), + value: "multi_field_metric".to_string(), + }, + Label { + name: "extra_field".to_string(), + value: "99".to_string(), + }, + ], + timeseries.labels + ); + assert_eq!( + vec![Sample { + value: 2.0, + timestamp: 2000, + }], + timeseries.samples + ); + + Ok(()) + } + + #[tokio::test] + async fn test_prom_remote_read_rejects_ambiguous_value_columns() -> TestResult<()> { + let schema = Arc::new(GtSchema::new(vec![ + ColumnSchema::new( + "custom_ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("field_a", ConcreteDataType::float64_datatype(), false), + ColumnSchema::new("field_b", ConcreteDataType::float64_datatype(), false), + ])); + let recordbatch = RecordBatch::new( + schema, + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![1000])) as VectorRef, + Arc::new(Float64Vector::from_vec(vec![1.0])) as VectorRef, + Arc::new(Float64Vector::from_vec(vec![2.0])) as VectorRef, + ], + ) + .unwrap(); + let instance = test_instance_with_tables( + MemTable::table("ambiguous_metric", recordbatch), + test_table(1025, "target")?, + ) + .await?; + + let err = PromStoreProtocolHandler::read( + &instance, + ReadRequest { + queries: vec![RemoteQuery { + matchers: vec![LabelMatcher { + r#type: PromMatcherType::Eq as i32, + name: servers::prom_store::METRIC_NAME_LABEL.to_string(), + value: "ambiguous_metric".to_string(), + }], + ..Default::default() + }], + ..Default::default() + }, + test_query_ctx(1), + ) + .await + .err() + .expect("ambiguous value columns should fail remote read"); + + assert_eq!(StatusCode::InvalidArguments, err.status_code()); + assert!(format!("{err:?}").contains("Ambiguous value column")); + + Ok(()) + } + #[tokio::test] async fn test_event_recorder_is_exposed() -> TestResult<()> { let instance = diff --git a/src/frontend/src/instance/prom_store.rs b/src/frontend/src/instance/prom_store.rs index 7c86f3316a..1bfb5f5643 100644 --- a/src/frontend/src/instance/prom_store.rs +++ b/src/frontend/src/instance/prom_store.rs @@ -32,7 +32,7 @@ use common_catalog::{format_full_table_name, parse_optional_catalog_and_schema_f use common_error::ext::BoxedError; use common_meta::rpc::ddl::TriggerReason; use common_query::Output; -use common_query::prelude::GREPTIME_PHYSICAL_TABLE; +use common_query::prelude::{GREPTIME_PHYSICAL_TABLE, greptime_value}; use common_recordbatch::RecordBatches; use common_telemetry::{debug, tracing}; use operator::insert::{ @@ -54,17 +54,25 @@ use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt}; use store_api::metric_engine_consts::{METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY}; use store_api::mito_engine_options::SST_FORMAT_KEY; +use table::TableRef; use table::table_reference::TableReference; use tracing::instrument; use crate::error::{ - CatalogSnafu, ExecLogicalPlanSnafu, PromStoreRemoteQueryPlanSnafu, ReadTableSnafu, Result, - TableNotFoundSnafu, + AmbiguousValueColumnSnafu, CatalogSnafu, ColumnNotFoundSnafu, ExecLogicalPlanSnafu, + PromStoreRemoteQueryPlanSnafu, ReadTableSnafu, Result, TableNotFoundSnafu, }; use crate::instance::Instance; const SAMPLES_RESPONSE_TYPE: i32 = ResponseType::Samples as i32; +struct RemoteQueryOutput { + table_name: String, + timestamp_column_name: String, + value_column_name: String, + output: Output, +} + fn auto_create_table_type_for_prom_remote_write( ctx: &QueryContextRef, with_metric_engine: bool, @@ -141,7 +149,12 @@ fn resolve_remote_query_target( } #[instrument(skip_all, fields(table_name))] -async fn to_query_result(table_name: &str, output: Output) -> ServerResult { +async fn to_query_result( + table_name: &str, + timestamp_column_name: &str, + value_column_name: &str, + output: Output, +) -> ServerResult { let OutputData::Stream(stream) = output.data else { unreachable!() }; @@ -149,10 +162,41 @@ async fn to_query_result(table_name: &str, output: Output) -> ServerResult Result { + let columns = table + .field_columns() + .map(|column| column.name) + .collect::>(); + + match columns.as_slice() { + [] => ColumnNotFoundSnafu { + msg: format!("value field in table '{table_name}'"), + } + .fail(), + + [only] => Ok(only.clone()), + + columns if columns.iter().any(|name| name == greptime_value()) => { + Ok(greptime_value().to_string()) + } + + columns => AmbiguousValueColumnSnafu { + table_name: table_name.to_string(), + field_columns: columns.to_vec(), + } + .fail(), + } +} + impl Instance { #[tracing::instrument(skip_all)] async fn handle_remote_query( @@ -162,7 +206,7 @@ impl Instance { schema_name: &str, table_name: &str, query: &Query, - ) -> Result { + ) -> Result { let table = self .catalog_manager .table(catalog_name, schema_name, table_name, Some(ctx)) @@ -172,6 +216,17 @@ impl Instance { table_name: format_full_table_name(catalog_name, schema_name, table_name), })?; + let timestamp_column_name = table + .schema() + .timestamp_column() + .with_context(|| ColumnNotFoundSnafu { + msg: format!("time index in table '{table_name}'"), + })? + .name + .clone(); + + let value_column_name = resolve_column_names(table_name, &table)?; + let dataframe = self .query_engine .read_table(table) @@ -179,8 +234,8 @@ impl Instance { table_name: format_full_table_name(catalog_name, schema_name, table_name), })?; - let logical_plan = - prom_store::query_to_plan(dataframe, query).context(PromStoreRemoteQueryPlanSnafu)?; + let logical_plan = prom_store::query_to_plan(dataframe, query, ×tamp_column_name) + .context(PromStoreRemoteQueryPlanSnafu)?; debug!( "Prometheus remote read, table: {}, logical plan: {}", @@ -188,10 +243,18 @@ impl Instance { logical_plan.display_indent(), ); - self.query_engine + let output = self + .query_engine .execute(logical_plan, ctx.clone()) .await - .context(ExecLogicalPlanSnafu) + .context(ExecLogicalPlanSnafu)?; + + Ok(RemoteQueryOutput { + table_name: table_name.to_string(), + timestamp_column_name, + value_column_name, + output, + }) } #[tracing::instrument(skip_all)] @@ -200,17 +263,17 @@ impl Instance { ctx: QueryContextRef, queries: &[Query], query_targets: &[PermissionTableTarget], - ) -> ServerResult> { + ) -> ServerResult> { let mut results = Vec::with_capacity(queries.len()); for (query, target) in queries.iter().zip(query_targets) { - let output = self + let result = self .handle_remote_query(&ctx, &target.catalog, &target.schema, &target.table, query) .await .map_err(BoxedError::new) .context(error::ExecuteQuerySnafu)?; - results.push((target.table.clone(), output)); + results.push(result); } Ok(results) } @@ -527,9 +590,23 @@ impl PromStoreProtocolHandler for Instance { ResponseType::Samples => { let mut query_results = Vec::with_capacity(results.len()); let mut map = HashMap::new(); - for (table_name, output) in results { + for result in results { + let RemoteQueryOutput { + table_name, + timestamp_column_name, + value_column_name, + output, + } = result; let plan = output.meta.plan.clone(); - query_results.push(to_query_result(&table_name, output).await?); + query_results.push( + to_query_result( + &table_name, + ×tamp_column_name, + &value_column_name, + output, + ) + .await?, + ); if let Some(ref plan) = plan { collect_plan_metrics(plan, &mut [&mut map]); } diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index 772689f0e0..33707b2935 100644 --- a/src/servers/src/prom_store.rs +++ b/src/servers/src/prom_store.rs @@ -137,7 +137,11 @@ pub fn extract_schema_from_query(query: &Query) -> Option { /// Create a DataFrame from a remote Query #[tracing::instrument(skip_all)] -pub fn query_to_plan(dataframe: DataFrame, q: &Query) -> Result { +pub fn query_to_plan( + dataframe: DataFrame, + q: &Query, + timestamp_column_name: &str, +) -> Result { let start_timestamp_ms = q.start_timestamp_ms; let end_timestamp_ms = q.end_timestamp_ms; @@ -145,8 +149,9 @@ pub fn query_to_plan(dataframe: DataFrame, q: &Query) -> Result { let mut conditions = Vec::with_capacity(label_matches.len() + 1); - conditions.push(col(greptime_timestamp()).gt_eq(lit_timestamp_millisecond(start_timestamp_ms))); - conditions.push(col(greptime_timestamp()).lt_eq(lit_timestamp_millisecond(end_timestamp_ms))); + conditions + .push(col(timestamp_column_name).gt_eq(lit_timestamp_millisecond(start_timestamp_ms))); + conditions.push(col(timestamp_column_name).lt_eq(lit_timestamp_millisecond(end_timestamp_ms))); for m in label_matches { let name = &m.name; @@ -261,14 +266,18 @@ struct LabelColumn<'a> { values: LabelValues<'a>, } -fn label_columns(recordbatch: &RecordBatch) -> Result>> { +fn label_columns<'a>( + recordbatch: &'a RecordBatch, + timestamp_column_name: &str, + value_column_name: &str, +) -> Result>> { recordbatch .schema .column_schemas() .iter() .enumerate() .filter(|(_, column_schema)| { - column_schema.name != greptime_timestamp() && column_schema.name != greptime_value() + column_schema.name != timestamp_column_name && column_schema.name != value_column_name }) .map(|(index, column_schema)| { let array = recordbatch.column(index); @@ -354,24 +363,31 @@ fn new_timeseries(table: &str, columns: &[LabelColumn<'_>], row: usize) -> TimeS pub fn recordbatches_to_timeseries( table_name: &str, + timestamp_column_name: &str, + value_column_name: &str, recordbatches: RecordBatches, ) -> Result> { Ok(recordbatches .take() .into_iter() - .map(|x| recordbatch_to_timeseries(table_name, x)) + .map(|x| recordbatch_to_timeseries(table_name, timestamp_column_name, value_column_name, x)) .collect::>>()? .into_iter() .flatten() .collect()) } -fn recordbatch_to_timeseries(table: &str, recordbatch: RecordBatch) -> Result> { - let ts_column = recordbatch.column_by_name(greptime_timestamp()).context( - error::InvalidPromRemoteReadQueryResultSnafu { - msg: "missing greptime_timestamp column in query result", - }, - )?; +fn recordbatch_to_timeseries( + table: &str, + timestamp_column_name: &str, + value_column_name: &str, + recordbatch: RecordBatch, +) -> Result> { + let ts_column = recordbatch + .column_by_name(timestamp_column_name) + .with_context(|| error::InvalidPromRemoteReadQueryResultSnafu { + msg: format!("missing timestamp column '{timestamp_column_name}' in query result"), + })?; let ts_column = ts_column .as_primitive_opt::() .with_context(|| error::InvalidPromRemoteReadQueryResultSnafu { @@ -383,11 +399,11 @@ fn recordbatch_to_timeseries(table: &str, recordbatch: RecordBatch) -> Result() .with_context(|| error::InvalidPromRemoteReadQueryResultSnafu { @@ -397,7 +413,7 @@ fn recordbatch_to_timeseries(table: &str, recordbatch: RecordBatch) -> Result = Vec::new(); let mut timeseries_by_hash: HashMap> = HashMap::new(); let mut previous_timeseries: Option = None; @@ -834,7 +850,7 @@ mod tests { let table_provider = Arc::new(DfTableProviderAdapter::new(table)); let dataframe = ctx.read_table(table_provider.clone()).unwrap(); - let plan = query_to_plan(dataframe, &q).unwrap(); + let plan = query_to_plan(dataframe, &q, greptime_timestamp()).unwrap(); let display_string = format!("{}", plan.display_indent()); let ts_col = greptime_timestamp(); @@ -868,7 +884,7 @@ mod tests { }; let dataframe = ctx.read_table(table_provider).unwrap(); - let plan = query_to_plan(dataframe, &q).unwrap(); + let plan = query_to_plan(dataframe, &q, greptime_timestamp()).unwrap(); let display_string = format!("{}", plan.display_indent()); let ts_col = greptime_timestamp(); @@ -1056,7 +1072,13 @@ mod tests { ) .unwrap(); - let timeseries = recordbatches_to_timeseries("metric1", recordbatches).unwrap(); + let timeseries = recordbatches_to_timeseries( + "metric1", + greptime_timestamp(), + greptime_value(), + recordbatches, + ) + .unwrap(); assert_eq!(2, timeseries.len()); assert_eq!( @@ -1131,7 +1153,7 @@ mod tests { ) .unwrap(); let recordbatch = RecordBatch::from_df_record_batch(schema.clone(), batch); - let columns = label_columns(&recordbatch).unwrap(); + let columns = label_columns(&recordbatch, greptime_timestamp(), greptime_value()).unwrap(); assert!(matches!( columns[0].values, LabelValues::DictionaryUtf8 { .. } @@ -1139,7 +1161,13 @@ mod tests { drop(columns); let recordbatches = RecordBatches::try_new(schema, vec![recordbatch]).unwrap(); - let timeseries = recordbatches_to_timeseries("metric1", recordbatches).unwrap(); + let timeseries = recordbatches_to_timeseries( + "metric1", + greptime_timestamp(), + greptime_value(), + recordbatches, + ) + .unwrap(); assert_eq!(3, timeseries.len()); assert_eq!( @@ -1207,7 +1235,13 @@ mod tests { ) .unwrap(); - let timeseries = recordbatch_to_timeseries("metric1", recordbatch).unwrap(); + let timeseries = recordbatch_to_timeseries( + "metric1", + greptime_timestamp(), + greptime_value(), + recordbatch, + ) + .unwrap(); // The result stays sorted by labels as it was with the previous BTreeMap. assert_eq!("host1", timeseries[0].labels[1].value); @@ -1258,7 +1292,13 @@ mod tests { ) .unwrap(); - let timeseries = recordbatch_to_timeseries("metric1", recordbatch).unwrap(); + let timeseries = recordbatch_to_timeseries( + "metric1", + greptime_timestamp(), + greptime_value(), + recordbatch, + ) + .unwrap(); assert_eq!(2, timeseries.len()); assert_eq!(