1pub mod query;
18
19use std::collections::{BTreeMap, BTreeSet};
20use std::sync::Arc;
21
22use api::v1::value::ValueData;
23use api::v1::{ColumnDataType, Row, RowInsertRequests, Value};
24use common_catalog::consts::METRICS_EXEMPLARS_TABLE;
25use common_grpc::precision::Precision;
26use common_query::prelude::{
27 EXEMPLAR_LABELS_COLUMN, EXEMPLAR_METRIC_COLUMN, greptime_timestamp, greptime_value,
28};
29use datatypes::json;
30use metric_engine::row_modifier::TsidGenerator;
31use snafu::{OptionExt, ensure};
32use store_api::metric_engine_consts::{
33 DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME,
34};
35use store_api::storage::consts::is_internal_column;
36
37use crate::error::{InvalidPromRemoteRequestSnafu, Result};
38use crate::row_writer::{self, MultiTableData};
39
40#[derive(Debug, Clone, PartialEq)]
43pub struct Exemplar {
44 pub metric: Arc<str>,
45 pub series_labels: Arc<BTreeMap<String, String>>,
46 pub timestamp: i64,
47 pub value: f64,
48 pub labels: BTreeMap<String, String>,
49}
50
51pub type ExemplarRequests = BTreeMap<Option<String>, Vec<Exemplar>>;
53
54pub fn estimated_row_bytes<'a>(
56 requests: impl IntoIterator<Item = &'a [Exemplar]>,
57) -> Option<usize> {
58 requests.into_iter().try_fold(0usize, |total, exemplars| {
59 let mut tags = BTreeSet::new();
60 let mut payload = 0usize;
61 for exemplar in exemplars {
63 tags.extend(exemplar.series_labels.keys());
64 payload = payload.checked_add(exemplar.metric.len().checked_mul(4)?)?;
65 for (name, value) in exemplar.series_labels.iter().chain(&exemplar.labels) {
66 payload = payload.checked_add(
68 name.len()
69 .checked_add(value.len())?
70 .checked_mul(4)?
71 .checked_add(256)?,
72 )?;
73 }
74 }
75 let columns = tags.len().checked_add(6)?;
77 let values = exemplars
78 .len()
79 .checked_mul(columns)?
80 .checked_mul(2 * std::mem::size_of::<Value>())?;
81 let rows = exemplars
82 .len()
83 .checked_mul(2 * std::mem::size_of::<Row>() + std::mem::size_of::<Exemplar>())?;
84 total
85 .checked_add(values)?
86 .checked_add(rows)?
87 .checked_add(columns.checked_mul(512)?)?
88 .checked_add(payload)
89 })
90}
91
92pub fn labels<I>(
94 labels: I,
95 exemplar_owned: bool,
96 validate_name: impl Fn(&[u8]) -> bool,
97) -> Result<BTreeMap<String, String>>
98where
99 I: IntoIterator<Item = (String, String)>,
100{
101 let mut result = BTreeMap::new();
102 let mut length = 0usize;
103 for (name, value) in labels {
104 ensure!(
105 validate_name(name.as_bytes()),
106 InvalidPromRemoteRequestSnafu {
107 msg: format!("invalid exemplar label name: {name}"),
108 }
109 );
110 if exemplar_owned {
111 ensure!(
112 name != json::JSON2_REMAINDER_FIELD_NAME,
113 InvalidPromRemoteRequestSnafu {
114 msg: "exemplar label name is reserved by JSON2",
115 }
116 );
117 } else {
118 ensure!(
119 ![
120 greptime_timestamp(),
121 greptime_value(),
122 EXEMPLAR_METRIC_COLUMN,
123 EXEMPLAR_LABELS_COLUMN,
124 DATA_SCHEMA_TABLE_ID_COLUMN_NAME,
125 DATA_SCHEMA_TSID_COLUMN_NAME,
126 ]
127 .contains(&name.as_str())
128 && !is_internal_column(&name),
129 InvalidPromRemoteRequestSnafu {
130 msg: format!("inherited exemplar label collides with reserved column: {name}"),
131 }
132 );
133 }
134 length = length
135 .checked_add(name.chars().count())
136 .and_then(|n| n.checked_add(value.chars().count()))
137 .context(InvalidPromRemoteRequestSnafu {
138 msg: "exemplar label length overflow",
139 })?;
140 ensure!(
141 !exemplar_owned || length <= 128,
142 InvalidPromRemoteRequestSnafu {
143 msg: "exemplar label names and values exceed 128 characters",
144 }
145 );
146 ensure!(
147 result.insert(name.clone(), value).is_none(),
148 InvalidPromRemoteRequestSnafu {
149 msg: format!("duplicate exemplar label: {name}"),
150 }
151 );
152 }
153 Ok(result)
154}
155
156pub fn timestamp_from_millis(timestamp: i64) -> Result<i64> {
158 timestamp
159 .checked_mul(1_000_000)
160 .context(InvalidPromRemoteRequestSnafu {
161 msg: "exemplar timestamp is outside the signed nanosecond range",
162 })
163}
164
165pub fn into_rows(exemplars: Vec<Exemplar>) -> Result<(RowInsertRequests, BTreeSet<String>)> {
167 let mut tables = MultiTableData::default();
168 let mut sources = BTreeSet::new();
169 for exemplar in exemplars {
170 ensure!(
171 !exemplar.metric.is_empty() && exemplar.metric.as_ref() != METRICS_EXEMPLARS_TABLE,
172 InvalidPromRemoteRequestSnafu {
173 msg: "invalid or reserved exemplar source metric name"
174 }
175 );
176 sources.insert(exemplar.metric.to_string());
177 let table = tables.get_or_default_table_data(
178 METRICS_EXEMPLARS_TABLE,
179 exemplar.series_labels.len() + 6,
180 1,
181 );
182 let mut row = table.alloc_one_row();
183 let tsid = TsidGenerator::from_sorted_tags(
184 exemplar
185 .series_labels
186 .iter()
187 .map(|(name, value)| (name.as_str(), value.as_str())),
188 );
189 row_writer::write_by_schema(
190 table,
191 std::iter::once((
192 api::v1::helper::tag_column_schema(
193 DATA_SCHEMA_TSID_COLUMN_NAME,
194 ColumnDataType::Uint64,
195 ),
196 Some(ValueData::U64Value(tsid)),
197 )),
198 &mut row,
199 )?;
200 row_writer::write_tag(table, EXEMPLAR_METRIC_COLUMN, &exemplar.metric, &mut row)?;
201 row_writer::write_tags(
202 table,
203 exemplar
204 .series_labels
205 .iter()
206 .map(|(name, value)| (name.as_str(), value.clone())),
207 &mut row,
208 )?;
209 row_writer::write_fields(
210 table,
211 std::iter::once((
212 greptime_value().to_string(),
213 ColumnDataType::Float64,
214 Some(ValueData::F64Value(exemplar.value)),
215 )),
216 &mut row,
217 )?;
218 row_writer::write_ts_to_nanos(
219 table,
220 greptime_timestamp(),
221 Some(exemplar.timestamp),
222 Precision::Nanosecond,
223 &mut row,
224 )?;
225 row_writer::write_by_schema(
226 table,
227 std::iter::once((
228 row_writer::build_json2_column_schema(EXEMPLAR_LABELS_COLUMN)?,
229 Some(row_writer::encode_json2(exemplar.labels)?),
230 )),
231 &mut row,
232 )?;
233 table.add_row(row);
234 }
235 Ok((tables.into_row_insert_requests().0, sources))
236}
237
238#[cfg(test)]
239mod tests {
240 use super::*;
241 use crate::prom_remote_write::validation::validate_label_name;
242
243 #[test]
244 fn test_exemplar_labels_and_timestamp_validation() {
245 assert!(
246 labels(
247 [(String::from("traceID"), String::from("abc"))],
248 true,
249 validate_label_name
250 )
251 .is_ok()
252 );
253 assert!(
254 labels(
255 [
256 (String::from("job"), String::new()),
257 (String::from("job"), String::new())
258 ],
259 false,
260 validate_label_name,
261 )
262 .is_err()
263 );
264 assert!(
265 labels(
266 [(EXEMPLAR_METRIC_COLUMN.into(), "x".into())],
267 false,
268 validate_label_name
269 )
270 .is_err()
271 );
272 assert!(labels([("id".into(), "x".repeat(127))], true, validate_label_name).is_err());
273 for name in ["", "request.method", "请求"] {
274 assert!(labels([(name.into(), "x".into())], true, validate_label_name).is_err());
275 }
276 for name in [
277 DATA_SCHEMA_TABLE_ID_COLUMN_NAME,
278 DATA_SCHEMA_TSID_COLUMN_NAME,
279 ] {
280 assert!(labels([(name.into(), "x".into())], false, validate_label_name).is_err());
281 }
282 assert_eq!(timestamp_from_millis(-1).unwrap(), -1_000_000);
283 assert!(timestamp_from_millis(i64::MAX).is_err());
284 }
285
286 #[test]
287 fn test_exemplar_tsid_uses_only_source_labels() {
288 let exemplar = Exemplar {
289 metric: "cpu".into(),
290 series_labels: BTreeMap::from([("job".into(), "api".into())]).into(),
291 timestamp: 1,
292 value: 2.0,
293 labels: BTreeMap::from([("trace_id".into(), "a".into())]),
294 };
295 let mut other = exemplar.clone();
296 other.metric = "memory".into();
297 other.labels.insert("trace_id".into(), "b".into());
298 let (requests, _) = into_rows(vec![exemplar, other]).unwrap();
299 let rows = requests.inserts[0].rows.as_ref().unwrap();
300 let index = rows
301 .schema
302 .iter()
303 .position(|c| c.column_name == DATA_SCHEMA_TSID_COLUMN_NAME)
304 .unwrap();
305 let expected = ValueData::U64Value(TsidGenerator::from_sorted_tags(
306 [("job", "api")].into_iter(),
307 ));
308 assert!(
309 rows.rows
310 .iter()
311 .all(|row| row.values[index].value_data.as_ref() == Some(&expected))
312 );
313 }
314}