servers/prom_remote_write/
row_builder.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16
17use api::prom_store::remote::Sample;
18use api::v1::helper::{field_column_schema, time_index_column_schema};
19use api::v1::value::ValueData;
20use api::v1::{ColumnDataType, ColumnSchema, Row, RowInsertRequest, Rows, SemanticType, Value};
21use common_query::prelude::{greptime_timestamp, greptime_value};
22use pipeline::{ContextOpt, ContextReq};
23use prost::DecodeError;
24
25use crate::prom_remote_write::PromValidationMode;
26use crate::prom_remote_write::types::PromLabel;
27use crate::prom_remote_write::validation::validate_label_name;
28use crate::repeated_field::Clear;
29
30#[derive(Debug, Default, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
31pub struct PromCtx {
32    pub schema: Option<String>,
33    pub physical_table: Option<String>,
34}
35
36#[derive(Default, Debug)]
37pub struct TablesBuilder<'a> {
38    pub tables: HashMap<PromCtx, HashMap<String, TableBuilder<'a>>>,
39    pub(crate) raw_data: Vec<u8>,
40}
41
42impl<'a> Clear for TablesBuilder<'a> {
43    fn clear(&mut self) {
44        self.tables.clear();
45        self.raw_data.clear();
46    }
47}
48
49impl<'a> TablesBuilder<'a> {
50    pub(crate) fn get_or_create_table_builder(
51        &mut self,
52        prom_ctx: PromCtx,
53        table_name: String,
54        label_num: usize,
55        row_num: usize,
56    ) -> &mut TableBuilder<'a> {
57        self.tables
58            .entry(prom_ctx)
59            .or_default()
60            .entry(table_name)
61            .or_insert_with(|| TableBuilder::with_capacity(label_num + 2, row_num))
62    }
63
64    pub(crate) fn as_insert_requests(&mut self) -> ContextReq {
65        self.tables
66            .drain()
67            .map(|(prom, mut tables)| {
68                let mut opt = ContextOpt::default();
69                if let Some(physical_table) = prom.physical_table {
70                    opt.set_physical_table(physical_table);
71                }
72                if let Some(schema) = prom.schema {
73                    opt.set_schema(schema);
74                }
75
76                let mut ctx_req = ContextReq::default();
77                let reqs = tables
78                    .drain()
79                    .map(|(table_name, mut table)| table.as_row_insert_request(table_name));
80                ctx_req.add_rows(opt, reqs);
81
82                ctx_req
83            })
84            .fold(ContextReq::default(), |mut req, reqs| {
85                req.merge(reqs);
86                req
87            })
88    }
89
90    pub(crate) fn set_raw_data(&mut self, buf: Vec<u8>) {
91        self.raw_data = buf;
92    }
93}
94
95#[derive(Debug)]
96pub struct TableBuilder<'a> {
97    schema: Vec<ColumnSchema>,
98    rows: Vec<Row>,
99    col_indexes: HashMap<&'a [u8], usize>,
100}
101
102impl<'a> Default for TableBuilder<'a> {
103    fn default() -> Self {
104        Self::with_capacity(2, 0)
105    }
106}
107
108impl<'a> TableBuilder<'a> {
109    pub(crate) fn with_capacity(cols: usize, rows: usize) -> Self {
110        let mut col_indexes = HashMap::with_capacity_and_hasher(cols, Default::default());
111        col_indexes.insert(greptime_timestamp().as_bytes(), 0);
112        col_indexes.insert(greptime_value().as_bytes(), 1);
113
114        let mut schema = Vec::with_capacity(cols);
115        schema.push(time_index_column_schema(
116            greptime_timestamp(),
117            ColumnDataType::TimestampMillisecond,
118        ));
119        schema.push(field_column_schema(
120            greptime_value(),
121            ColumnDataType::Float64,
122        ));
123
124        Self {
125            schema,
126            rows: Vec::with_capacity(rows),
127            col_indexes,
128        }
129    }
130
131    pub(crate) fn add_labels_and_samples(
132        &mut self,
133        labels: &[PromLabel],
134        samples: &[Sample],
135        prom_validation_mode: PromValidationMode,
136    ) -> Result<(), DecodeError> {
137        let mut row = vec![Value { value_data: None }; self.col_indexes.len()];
138
139        for PromLabel { name, value } in labels {
140            if !validate_label_name(name) {
141                return Err(DecodeError::new(format!(
142                    "Invalid label name: `{}`",
143                    String::from_utf8_lossy(name)
144                )));
145            }
146            let raw_tag_name = *name;
147            let tag_value = Some(ValueData::StringValue(
148                prom_validation_mode.decode_string(value)?,
149            ));
150            let tag_num = self.col_indexes.len();
151
152            if let Some(e) = self.col_indexes.get_mut(raw_tag_name) {
153                row[*e].value_data = tag_value;
154                continue;
155            }
156
157            let tag_name = unsafe { std::str::from_utf8_unchecked(raw_tag_name) };
158            self.schema.push(ColumnSchema {
159                column_name: tag_name.to_owned(),
160                datatype: ColumnDataType::String as i32,
161                semantic_type: SemanticType::Tag as i32,
162                ..Default::default()
163            });
164            self.col_indexes.insert(raw_tag_name, tag_num);
165
166            row.push(Value {
167                value_data: tag_value,
168            });
169        }
170
171        let Some((last_sample, preceding_samples)) = samples.split_last() else {
172            return Ok(());
173        };
174
175        for sample in preceding_samples {
176            row[0].value_data = Some(ValueData::TimestampMillisecondValue(sample.timestamp));
177            row[1].value_data = Some(ValueData::F64Value(sample.value));
178            self.rows.push(Row {
179                values: row.clone(),
180            });
181        }
182
183        row[0].value_data = Some(ValueData::TimestampMillisecondValue(last_sample.timestamp));
184        row[1].value_data = Some(ValueData::F64Value(last_sample.value));
185        self.rows.push(Row { values: row });
186
187        Ok(())
188    }
189
190    pub fn as_row_insert_request(&mut self, table_name: String) -> RowInsertRequest {
191        let mut rows = std::mem::take(&mut self.rows);
192        let schema = std::mem::take(&mut self.schema);
193        let col_num = schema.len();
194        for row in &mut rows {
195            if row.values.len() < col_num {
196                row.values.resize(col_num, Value { value_data: None });
197            }
198        }
199
200        RowInsertRequest {
201            table_name,
202            rows: Some(Rows { schema, rows }),
203        }
204    }
205
206    pub fn tags(&self) -> impl Iterator<Item = &String> {
207        self.schema
208            .iter()
209            .filter(|v| v.semantic_type == SemanticType::Tag as i32)
210            .map(|c| &c.column_name)
211    }
212}
213
214#[cfg(test)]
215mod tests {
216    use api::prom_store::remote::Sample;
217    use prost::DecodeError;
218
219    use super::*;
220
221    fn assert_sample(row: &Row, timestamp: i64, value: f64) {
222        assert_eq!(
223            Some(&ValueData::TimestampMillisecondValue(timestamp)),
224            row.values[0].value_data.as_ref()
225        );
226        assert_eq!(
227            Some(&ValueData::F64Value(value)),
228            row.values[1].value_data.as_ref()
229        );
230    }
231
232    #[test]
233    fn test_add_labels_and_samples() {
234        let mut builder = TableBuilder::default();
235        builder
236            .add_labels_and_samples(
237                &[PromLabel {
238                    name: b"tag0",
239                    value: b"v0",
240                }],
241                &[],
242                PromValidationMode::Strict,
243            )
244            .unwrap();
245        assert!(builder.rows.is_empty());
246
247        builder
248            .add_labels_and_samples(
249                &[PromLabel {
250                    name: b"tag0",
251                    value: b"v0",
252                }],
253                &[Sample {
254                    value: 1.0,
255                    timestamp: 1,
256                }],
257                PromValidationMode::Strict,
258            )
259            .unwrap();
260
261        builder
262            .add_labels_and_samples(
263                &[
264                    PromLabel {
265                        name: b"tag0",
266                        value: b"v1",
267                    },
268                    PromLabel {
269                        name: b"tag1",
270                        value: b"v2",
271                    },
272                ],
273                &[
274                    Sample {
275                        value: 2.0,
276                        timestamp: 2,
277                    },
278                    Sample {
279                        value: 3.0,
280                        timestamp: 3,
281                    },
282                ],
283                PromValidationMode::Strict,
284            )
285            .unwrap();
286
287        let request = builder.as_row_insert_request("test".to_string());
288        let rows = request.rows.unwrap().rows;
289        assert_eq!(3, rows.len());
290        assert!(rows.iter().all(|row| row.values.len() == 4));
291
292        assert_sample(&rows[0], 1, 1.0);
293        assert_eq!(
294            Some(&ValueData::StringValue("v0".to_string())),
295            rows[0].values[2].value_data.as_ref()
296        );
297        assert!(rows[0].values[3].value_data.is_none());
298
299        assert_sample(&rows[1], 2, 2.0);
300        assert_sample(&rows[2], 3, 3.0);
301        for row in &rows[1..] {
302            assert_eq!(
303                Some(&ValueData::StringValue("v1".to_string())),
304                row.values[2].value_data.as_ref()
305            );
306            assert_eq!(
307                Some(&ValueData::StringValue("v2".to_string())),
308                row.values[3].value_data.as_ref()
309            );
310        }
311    }
312
313    #[test]
314    fn test_table_builder() {
315        let mut builder = TableBuilder::default();
316        let _ = builder.add_labels_and_samples(
317            &[
318                PromLabel {
319                    name: b"tag0",
320                    value: b"v0",
321                },
322                PromLabel {
323                    name: b"tag1",
324                    value: b"v1",
325                },
326            ],
327            &[Sample {
328                value: 0.0,
329                timestamp: 0,
330            }],
331            PromValidationMode::Strict,
332        );
333
334        let _ = builder.add_labels_and_samples(
335            &[
336                PromLabel {
337                    name: b"tag0",
338                    value: b"v0",
339                },
340                PromLabel {
341                    name: b"tag2",
342                    value: b"v2",
343                },
344            ],
345            &[Sample {
346                value: 0.1,
347                timestamp: 1,
348            }],
349            PromValidationMode::Strict,
350        );
351
352        let request = builder.as_row_insert_request("test".to_string());
353        let rows = request.rows.unwrap().rows;
354        assert_eq!(2, rows.len());
355
356        let invalid_utf8_bytes = &[0xFF, 0xFF, 0xFF];
357        let res = builder.add_labels_and_samples(
358            &[PromLabel {
359                name: b"tag0",
360                value: invalid_utf8_bytes,
361            }],
362            &[Sample {
363                value: 0.1,
364                timestamp: 1,
365            }],
366            PromValidationMode::Strict,
367        );
368        assert_eq!(res, Err(DecodeError::new("invalid utf-8")));
369    }
370}
servers/prom_remote_write/row_builder.rs

servers/prom_remote_write/
row_builder.rs
