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

fix(prometheus): custom column remote reads (#8659)

Browse Source 
* 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 <me@gauravshokeen.com>

* 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 <me@gauravshokeen.com>

---------

Signed-off-by: grezzko <me@gauravshokeen.com>
Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>
Co-authored-by: Lei, HUANG <mrsatangel@gmail.com>
This commit is contained in:
Gaurav Shokeen
2026-08-07 19:07:00 +05:30
committed by GitHub
parent 1f1c9270a8
commit d90cca4b75
4 changed files with 378 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/frontend/src/error.rs
+11
View File
@@ -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<String>,
#[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 { .. }
src/frontend/src/instance.rs
+210 -3
View File
@@ -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 =
src/frontend/src/instance/prom_store.rs
+92 -15
View File
@@ -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<QueryResult> {
async fn to_query_result(
table_name: &str,
timestamp_column_name: &str,
value_column_name: &str,
output: Output,
) -> ServerResult<QueryResult> {
let OutputData::Stream(stream) = output.data else {
unreachable!()
};
@@ -149,10 +162,41 @@ async fn to_query_result(table_name: &str, output: Output) -> ServerResult<Query
.await
.context(error::CollectRecordbatchSnafu)?;
Ok(QueryResult {
timeseries: prom_store::recordbatches_to_timeseries(table_name, recordbatches)?,
timeseries: prom_store::recordbatches_to_timeseries(
table_name,
timestamp_column_name,
value_column_name,
recordbatches,
)?,
})
}
fn resolve_column_names(table_name: &str, table: &TableRef) -> Result<String> {
let columns = table
.field_columns()
.map(|column| column.name)
.collect::<Vec<_>>();
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<Output> {
) -> Result<RemoteQueryOutput> {
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, &timestamp_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<Vec<(String, Output)>> {
) -> ServerResult<Vec<RemoteQueryOutput>> {
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,
&timestamp_column_name,
&value_column_name,
output,
)
.await?,
);
if let Some(ref plan) = plan {
collect_plan_metrics(plan, &mut [&mut map]);
}
src/servers/src/prom_store.rs
+65 -25
View File
@@ -137,7 +137,11 @@ pub fn extract_schema_from_query(query: &Query) -> Option<String> {
/// Create a DataFrame from a remote Query
#[tracing::instrument(skip_all)]
pub fn query_to_plan(dataframe: DataFrame, q: &Query) -> Result<LogicalPlan> {
pub fn query_to_plan(
dataframe: DataFrame,
q: &Query,
timestamp_column_name: &str,
) -> Result<LogicalPlan> {
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<LogicalPlan> {
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<Vec<LabelColumn<'_>>> {
fn label_columns<'a>(
recordbatch: &'a RecordBatch,
timestamp_column_name: &str,
value_column_name: &str,
) -> Result<Vec<LabelColumn<'a>>> {
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<Vec<TimeSeries>> {
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::<Result<Vec<_>>>()?
.into_iter()
.flatten()
.collect())
}
fn recordbatch_to_timeseries(table: &str, recordbatch: RecordBatch) -> Result<Vec<TimeSeries>> {
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<Vec<TimeSeries>> {
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::<TimestampMillisecondType>()
.with_context(|| error::InvalidPromRemoteReadQueryResultSnafu {
@@ -383,11 +399,11 @@ fn recordbatch_to_timeseries(table: &str, recordbatch: RecordBatch) -> Result<Ve
// TODO: Add native-histogram encoding when Prometheus Remote Read support is prioritized.
// The current path intentionally returns scalar samples only.
let field_column = recordbatch.column_by_name(greptime_value()).context(
error::InvalidPromRemoteReadQueryResultSnafu {
msg: "missing greptime_value column in query result",
},
)?;
let field_column = recordbatch
.column_by_name(value_column_name)
.with_context(|| error::InvalidPromRemoteReadQueryResultSnafu {
msg: format!("missing value column '{value_column_name}' in query result"),
})?;
let field_column = field_column
.as_primitive_opt::<Float64Type>()
.with_context(|| error::InvalidPromRemoteReadQueryResultSnafu {
@@ -397,7 +413,7 @@ fn recordbatch_to_timeseries(table: &str, recordbatch: RecordBatch) -> Result<Ve
),
})?;
let columns = label_columns(&recordbatch)?;
let columns = label_columns(&recordbatch, timestamp_column_name, value_column_name)?;
let mut timeseries: Vec<TimeSeries> = Vec::new();
let mut timeseries_by_hash: HashMap<u64, Vec<usize>> = HashMap::new();
let mut previous_timeseries: Option<usize> = 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!(