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

servers/prom_remote_write/
row_builder.rs
