servers/prom_remote_write/
decode.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
15//! Decoding of Prometheus remote write protobuf payloads.
16
17use std::collections::btree_map::Entry;
18use std::collections::{BTreeMap, HashSet};
19use std::sync::Arc;
20
21use api::prom_store::remote::{Exemplar, Sample};
22use bytes::Buf;
23use common_query::prelude::{greptime_timestamp, greptime_value};
24use pipeline::{ContextReq, GreptimePipelineParams, PipelineContext, PipelineDefinition};
25use prost::DecodeError;
26use prost::encoding::message::merge;
27use prost::encoding::{WireType, decode_key, decode_varint};
28use session::context::QueryContextRef;
29use snafu::OptionExt;
30use vrl::prelude::NotNan;
31use vrl::value::{KeyString, Value as VrlValue};
32
33use crate::error::InternalSnafu;
34use crate::http::event::PipelineIngestRequest;
35use crate::pipeline::run_pipeline;
36use crate::prom_remote_write::row_builder::{PromCtx, TablesBuilder};
37use crate::prom_remote_write::types::{PromLabel, RawBytes};
38use crate::prom_remote_write::validation::{PromValidationMode, validate_label_name};
39#[allow(deprecated)]
40use crate::prom_store::{
41    DATABASE_LABEL_ALT_BYTES, DATABASE_LABEL_BYTES, METRIC_NAME_LABEL_BYTES,
42    PHYSICAL_TABLE_LABEL_ALT_BYTES, PHYSICAL_TABLE_LABEL_BYTES, SCHEMA_LABEL_BYTES,
43};
44use crate::query_handler::PipelineHandlerRef;
45use crate::repeated_field::{Clear, RepeatedField};
46
47#[derive(Default, Debug)]
48pub(crate) struct PromTimeSeries {
49    pub(crate) table_name: String,
50    pub(crate) schema: Option<String>,
51    pub(crate) physical_table: Option<String>,
52
53    pub(crate) labels: RepeatedField<PromLabel>,
54    pub(crate) samples: RepeatedField<Sample>,
55    exemplars: Vec<Exemplar>,
56    label_names: HashSet<RawBytes>,
57    duplicate_labels: bool,
58}
59
60impl Clear for PromTimeSeries {
61    fn clear(&mut self) {
62        self.table_name.clear();
63        self.schema.clear();
64        self.physical_table.clear();
65        // Labels borrow from the request buffer, which may be replaced after this reset.
66        for label in self.labels.iter_mut() {
67            label.clear();
68        }
69        self.labels.clear();
70        self.samples.clear();
71        self.exemplars.clear();
72        self.label_names.clear();
73        self.duplicate_labels = false;
74    }
75}
76
77impl PromTimeSeries {
78    pub fn merge_field(
79        &mut self,
80        tag: u32,
81        wire_type: WireType,
82        buf: &mut &[u8],
83        prom_validation_mode: PromValidationMode,
84    ) -> Result<(), DecodeError> {
85        const STRUCT_NAME: &str = "PromTimeSeries";
86        match tag {
87            1u32 => {
88                let label = self.labels.push_default();
89
90                let len = decode_varint(buf).map_err(|mut error| {
91                    error.push(STRUCT_NAME, "labels");
92                    error
93                })?;
94                let remaining = buf.remaining();
95                if len > remaining as u64 {
96                    return Err(DecodeError::new("buffer underflow"));
97                }
98
99                let limit = remaining - len as usize;
100                while buf.remaining() > limit {
101                    let (tag, wire_type) = decode_key(buf)?;
102                    label.merge_field(tag, wire_type, buf)?;
103                }
104                if buf.remaining() != limit {
105                    return Err(DecodeError::new("delimited length exceeded"));
106                }
107
108                self.duplicate_labels |= !self.label_names.insert(label.name);
109                #[allow(deprecated)]
110                let is_special_label = match label.name {
111                    METRIC_NAME_LABEL_BYTES => {
112                        self.table_name = prom_validation_mode.decode_string(label.value)?;
113                        true
114                    }
115                    SCHEMA_LABEL_BYTES => {
116                        self.schema = Some(prom_validation_mode.decode_string(label.value)?);
117                        true
118                    }
119                    DATABASE_LABEL_BYTES | DATABASE_LABEL_ALT_BYTES => {
120                        if self.schema.is_none() {
121                            self.schema = Some(prom_validation_mode.decode_string(label.value)?);
122                        }
123                        true
124                    }
125                    PHYSICAL_TABLE_LABEL_BYTES | PHYSICAL_TABLE_LABEL_ALT_BYTES => {
126                        self.physical_table =
127                            Some(prom_validation_mode.decode_string(label.value)?);
128                        true
129                    }
130                    _ => false,
131                };
132                if is_special_label {
133                    label.clear();
134                    self.labels.truncate(self.labels.len() - 1);
135                }
136
137                Ok(())
138            }
139            2u32 => {
140                let sample = self.samples.push_default();
141                merge(WireType::LengthDelimited, sample, buf, Default::default()).map_err(
142                    |mut error| {
143                        error.push(STRUCT_NAME, "samples");
144                        error
145                    },
146                )?;
147                Ok(())
148            }
149            3u32 => {
150                let mut exemplar = Exemplar::default();
151                merge(wire_type, &mut exemplar, buf, Default::default())?;
152                self.exemplars.push(exemplar);
153                Ok(())
154            }
155            4u32 => Err(DecodeError::new(
156                "remote write v1 native histogram ingestion is unsupported; use remote write v2",
157            )),
158            _ => prost::encoding::skip_field(wire_type, tag, buf, Default::default()),
159        }
160    }
161
162    fn add_to_table_data<'a>(
163        &mut self,
164        table_builders: &mut TablesBuilder<'a>,
165        prom_validation_mode: PromValidationMode,
166    ) -> Result<(), DecodeError> {
167        let label_num = self.labels.len();
168        let row_num = self.samples.len();
169
170        if !self.exemplars.is_empty() {
171            if self.duplicate_labels || self.table_name.is_empty() {
172                return Err(DecodeError::new(
173                    "exemplar series requires a metric name and unique labels",
174                ));
175            }
176            let inherited = self
177                .labels
178                .iter()
179                .map(|label| {
180                    Ok((
181                        prom_validation_mode
182                            .decode_label_name(label.name)?
183                            .to_string(),
184                        prom_validation_mode.decode_string(label.value)?,
185                    ))
186                })
187                .collect::<Result<Vec<_>, DecodeError>>()?;
188            let inherited = Arc::new(
189                crate::exemplar::labels(inherited, false, validate_label_name)
190                    .map_err(|error| DecodeError::new(error.to_string()))?,
191            );
192            let metric: Arc<str> = self.table_name.as_str().into();
193            for exemplar in self.exemplars.drain(..) {
194                let observation: crate::error::Result<_> = (|| {
195                    let labels = crate::exemplar::labels(
196                        exemplar
197                            .labels
198                            .into_iter()
199                            .map(|label| (label.name, label.value)),
200                        true,
201                        validate_label_name,
202                    )?;
203                    Ok(crate::exemplar::Exemplar {
204                        metric: metric.clone(),
205                        series_labels: inherited.clone(),
206                        timestamp: crate::exemplar::timestamp_from_millis(exemplar.timestamp)?,
207                        value: exemplar.value,
208                        labels,
209                    })
210                })();
211                let observation = match observation {
212                    Ok(observation) => observation,
213                    Err(_) => {
214                        crate::metrics::EXEMPLARS_DROPPED_TOTAL
215                            .with_label_values(&["prometheus_v1"])
216                            .inc();
217                        continue;
218                    }
219                };
220                table_builders
221                    .exemplars
222                    .entry(self.schema.clone())
223                    .or_default()
224                    .push(observation);
225            }
226        }
227        if row_num == 0 {
228            return Ok(());
229        }
230        let prom_ctx = PromCtx {
231            schema: self.schema.take(),
232            physical_table: self.physical_table.take(),
233        };
234
235        let table_data = table_builders.get_or_create_table_builder(
236            prom_ctx,
237            std::mem::take(&mut self.table_name),
238            label_num,
239            row_num,
240        );
241        table_data.add_labels_and_samples(
242            self.labels.as_slice(),
243            self.samples.as_slice(),
244            prom_validation_mode,
245        )?;
246
247        Ok(())
248    }
249}
250
251#[derive(Default, Debug)]
252pub struct PromWriteRequest<'a> {
253    pub(crate) table_data: TablesBuilder<'a>,
254    series: PromTimeSeries,
255}
256
257impl<'a> Clear for PromWriteRequest<'a> {
258    fn clear(&mut self) {
259        self.series.clear();
260        self.table_data.clear();
261    }
262}
263
264impl<'a> PromWriteRequest<'a> {
265    pub fn as_row_insert_requests(&mut self) -> ContextReq {
266        self.table_data.as_insert_requests()
267    }
268
269    pub fn decode(
270        &mut self,
271        buf: Vec<u8>,
272        prom_validation_mode: PromValidationMode,
273        processor: &mut PromSeriesProcessor,
274    ) -> Result<(), DecodeError> {
275        const STRUCT_NAME: &str = "PromWriteRequest";
276        self.clear();
277        self.table_data.set_raw_data(buf);
278        let mut offset = 0;
279        while offset < self.table_data.raw_data.len() {
280            let mut should_add_to_table_data = false;
281            let mut decoded_timeseries = false;
282            {
283                let raw_data = &self.table_data.raw_data;
284                let buf = &mut &raw_data[offset..];
285                let (tag, wire_type) = decode_key(buf)?;
286                if wire_type != WireType::LengthDelimited {
287                    return Err(DecodeError::new(format!(
288                        "invalid wire type: {:?}",
289                        wire_type
290                    )));
291                }
292                match tag {
293                    1u32 => {
294                        let len = decode_varint(buf).map_err(|mut e| {
295                            e.push(STRUCT_NAME, "timeseries");
296                            e
297                        })?;
298                        let remaining = buf.remaining();
299                        if len > remaining as u64 {
300                            return Err(DecodeError::new("buffer underflow"));
301                        }
302
303                        let limit = remaining - len as usize;
304                        while buf.remaining() > limit {
305                            let (tag, wire_type) = decode_key(buf)?;
306                            self.series
307                                .merge_field(tag, wire_type, buf, prom_validation_mode)?;
308                        }
309                        if buf.remaining() != limit {
310                            return Err(DecodeError::new("delimited length exceeded"));
311                        }
312
313                        if processor.use_pipeline && !self.series.exemplars.is_empty() {
314                            return Err(DecodeError::new(
315                                "Remote Write pipelines with exemplars are unsupported",
316                            ));
317                        }
318                        if processor.use_pipeline {
319                            processor.consume_series_to_pipeline_map(
320                                &mut self.series,
321                                prom_validation_mode,
322                            )?;
323                        } else {
324                            should_add_to_table_data = true;
325                        }
326
327                        decoded_timeseries = true;
328                    }
329                    3u32 => {
330                        prost::encoding::skip_field(wire_type, tag, buf, Default::default())?;
331                    }
332                    _ => prost::encoding::skip_field(wire_type, tag, buf, Default::default())?,
333                }
334                offset = raw_data.len() - buf.remaining();
335            }
336
337            if should_add_to_table_data {
338                self.series
339                    .add_to_table_data(&mut self.table_data, prom_validation_mode)?;
340            }
341
342            if decoded_timeseries {
343                self.series.clear();
344            }
345        }
346
347        Ok(())
348    }
349}
350
351/// Hook injected into the PromWriteRequest decoding process.
352pub struct PromSeriesProcessor {
353    pub(crate) use_pipeline: bool,
354    pub(crate) table_values: BTreeMap<String, Vec<VrlValue>>,
355
356    pub(crate) pipeline_handler: Option<PipelineHandlerRef>,
357    pub(crate) query_ctx: Option<QueryContextRef>,
358    pub(crate) pipeline_def: Option<PipelineDefinition>,
359}
360
361impl PromSeriesProcessor {
362    pub fn default_processor() -> Self {
363        Self {
364            use_pipeline: false,
365            table_values: BTreeMap::new(),
366            pipeline_handler: None,
367            query_ctx: None,
368            pipeline_def: None,
369        }
370    }
371
372    pub fn set_pipeline(
373        &mut self,
374        handler: PipelineHandlerRef,
375        query_ctx: QueryContextRef,
376        pipeline_def: PipelineDefinition,
377    ) {
378        self.use_pipeline = true;
379        self.pipeline_handler = Some(handler);
380        self.query_ctx = Some(query_ctx);
381        self.pipeline_def = Some(pipeline_def);
382    }
383
384    pub(crate) fn consume_series_to_pipeline_map(
385        &mut self,
386        series: &mut PromTimeSeries,
387        prom_validation_mode: PromValidationMode,
388    ) -> Result<(), DecodeError> {
389        let mut vec_pipeline_map = Vec::new();
390        let mut pipeline_map = BTreeMap::new();
391        for l in series.labels.iter() {
392            let name = prom_validation_mode.decode_label_name(l.name)?;
393            let value = prom_validation_mode.decode_string(l.value)?;
394            pipeline_map.insert(KeyString::from(name), VrlValue::Bytes(value.into()));
395        }
396
397        let one_sample = series.samples.len() == 1;
398
399        for s in series.samples.iter() {
400            let Ok(value) = NotNan::new(s.value) else {
401                common_telemetry::warn!("Invalid float value: {}", s.value);
402                continue;
403            };
404
405            let timestamp = s.timestamp;
406            pipeline_map.insert(
407                KeyString::from(greptime_timestamp()),
408                VrlValue::Integer(timestamp),
409            );
410            pipeline_map.insert(KeyString::from(greptime_value()), VrlValue::Float(value));
411            if one_sample {
412                vec_pipeline_map.push(VrlValue::Object(pipeline_map));
413                break;
414            } else {
415                vec_pipeline_map.push(VrlValue::Object(pipeline_map.clone()));
416            }
417        }
418
419        let table_name = std::mem::take(&mut series.table_name);
420        match self.table_values.entry(table_name) {
421            Entry::Occupied(mut occupied_entry) => {
422                occupied_entry.get_mut().append(&mut vec_pipeline_map);
423            }
424            Entry::Vacant(vacant_entry) => {
425                vacant_entry.insert(vec_pipeline_map);
426            }
427        }
428
429        Ok(())
430    }
431
432    pub(crate) async fn exec_pipeline(&mut self) -> crate::error::Result<ContextReq> {
433        let handler = self.pipeline_handler.as_ref().context(InternalSnafu {
434            err_msg: "pipeline handler is not set",
435        })?;
436        let pipeline_def = self.pipeline_def.as_ref().context(InternalSnafu {
437            err_msg: "pipeline definition is not set",
438        })?;
439        let pipeline_param = GreptimePipelineParams::default();
440        let query_ctx = self.query_ctx.as_ref().context(InternalSnafu {
441            err_msg: "query context is not set",
442        })?;
443
444        let pipeline_ctx = PipelineContext::new(pipeline_def, &pipeline_param, query_ctx.channel());
445
446        let mut req = ContextReq::default();
447        let table_values = std::mem::take(&mut self.table_values);
448        for (table_name, pipeline_maps) in table_values.into_iter() {
449            let pipeline_req = PipelineIngestRequest {
450                table: table_name,
451                values: pipeline_maps,
452            };
453            let row_req =
454                run_pipeline(handler, &pipeline_ctx, pipeline_req, query_ctx, true).await?;
455            req.merge(row_req);
456        }
457
458        Ok(req)
459    }
460}
461
462#[cfg(test)]
463mod tests {
464    use std::collections::HashMap;
465
466    use api::prom_store::remote::{Histogram, Label, Sample, TimeSeries, WriteRequest};
467    use api::v1::{Row, RowInsertRequests, Rows};
468    use bytes::Bytes;
469    use prost::Message;
470
471    use super::*;
472    use crate::prom_store::to_grpc_row_insert_requests;
473    use crate::repeated_field::Clear;
474
475    fn sort_rows(rows: Rows) -> Rows {
476        let permutation =
477            permutation::sort_by_key(&rows.schema, |schema| schema.column_name.clone());
478        let schema = permutation.apply_slice(&rows.schema);
479        let mut inner_rows = vec![];
480        for row in rows.rows {
481            let values = permutation.apply_slice(&row.values);
482            inner_rows.push(Row { values });
483        }
484        Rows {
485            schema,
486            rows: inner_rows,
487        }
488    }
489
490    fn check_deserialized(
491        prom_write_request: &mut PromWriteRequest,
492        data: &[u8],
493        expected_samples: usize,
494        expected_rows: &RowInsertRequests,
495    ) {
496        let mut p = PromSeriesProcessor::default_processor();
497        prom_write_request.clear();
498        prom_write_request
499            .decode(data.to_owned(), PromValidationMode::Strict, &mut p)
500            .unwrap();
501
502        let req = prom_write_request.as_row_insert_requests();
503
504        let samples = req
505            .ref_all_req()
506            .filter_map(|r| r.rows.as_ref().map(|r| r.rows.len()))
507            .sum::<usize>();
508        let prom_rows = RowInsertRequests {
509            inserts: req.all_req().collect::<Vec<_>>(),
510        };
511
512        assert_eq!(expected_samples, samples);
513        assert_eq!(expected_rows.inserts.len(), prom_rows.inserts.len());
514
515        let expected_rows_map = expected_rows
516            .inserts
517            .iter()
518            .map(|insert| (insert.table_name.clone(), insert.rows.clone().unwrap()))
519            .collect::<HashMap<_, _>>();
520
521        for r in &prom_rows.inserts {
522            let expected_rows = expected_rows_map.get(&r.table_name).unwrap().clone();
523            assert_eq!(sort_rows(expected_rows), sort_rows(r.rows.clone().unwrap()));
524        }
525    }
526
527    #[test]
528    fn test_decode_write_request() {
529        let mut d = std::path::PathBuf::from(
530            std::env::var("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not set"),
531        );
532        d.push("benches");
533        d.push("write_request.pb.data");
534        let data = std::fs::read(d).unwrap();
535
536        let (expected_rows, expected_samples) =
537            to_grpc_row_insert_requests(&WriteRequest::decode(&data[..]).unwrap()).unwrap();
538
539        let mut prom_write_request = PromWriteRequest::default();
540        for _ in 0..3 {
541            check_deserialized(
542                &mut prom_write_request,
543                &data,
544                expected_samples,
545                &expected_rows,
546            );
547        }
548    }
549
550    #[test]
551    fn test_exemplar_decoder_shares_series_data_before_row_expansion() {
552        let metric = "m".repeat(64 * 1024);
553        let label = "x".repeat(64 * 1024);
554        let count = 128;
555        let wire = WriteRequest {
556            timeseries: vec![TimeSeries {
557                labels: vec![
558                    Label {
559                        name: "__name__".into(),
560                        value: metric.clone(),
561                    },
562                    Label {
563                        name: "host".into(),
564                        value: label.clone(),
565                    },
566                ],
567                exemplars: vec![Exemplar::default(); count],
568                ..Default::default()
569            }],
570            ..Default::default()
571        }
572        .encode_to_vec();
573        assert!(wire.len() < 256 * 1024);
574        let mut decoder = PromWriteRequest::default();
575        decoder
576            .decode(
577                wire,
578                PromValidationMode::Strict,
579                &mut PromSeriesProcessor::default_processor(),
580            )
581            .unwrap();
582        let requests = std::mem::take(&mut decoder.table_data.exemplars);
583        drop(decoder);
584
585        let observations = &requests[&None];
586        assert_eq!(observations.len(), count);
587        let first = &observations[0];
588        for observation in observations {
589            assert_eq!(&*observation.metric, metric.as_str());
590            assert_eq!(observation.series_labels["host"], label);
591            // Decoding retains one copy of series strings regardless of exemplar count.
592            assert_eq!(observation.metric.as_ptr(), first.metric.as_ptr());
593            assert_eq!(
594                observation.series_labels["host"].as_ptr(),
595                first.series_labels["host"].as_ptr()
596            );
597        }
598        assert!(
599            crate::exemplar::estimated_row_bytes(requests.values().map(Vec::as_slice)).unwrap()
600                >= count * (metric.len() + label.len())
601        );
602    }
603
604    #[test]
605    fn test_exemplar_decoder_preserves_data_across_requests() {
606        let valid = Exemplar {
607            timestamp: -1,
608            value: 0.5,
609            labels: vec![Label {
610                name: "traceID".into(),
611                value: "abc".into(),
612            }],
613        };
614        let series = TimeSeries {
615            labels: vec![Label {
616                name: "__name__".into(),
617                value: "metric".into(),
618            }],
619            exemplars: vec![valid.clone()],
620            ..Default::default()
621        };
622        let encode = |series| {
623            WriteRequest {
624                timeseries: vec![series],
625                ..Default::default()
626            }
627            .encode_to_vec()
628        };
629        let sample = Sample {
630            timestamp: 2,
631            value: 1.0,
632        };
633        let mut decoder = PromWriteRequest::default();
634        let mut processor = PromSeriesProcessor::default_processor();
635        for invalid in [
636            Exemplar {
637                labels: vec![valid.labels[0].clone(); 2],
638                ..valid.clone()
639            },
640            Exemplar {
641                labels: vec![Label {
642                    name: "traceID".into(),
643                    value: "x".repeat(128),
644                }],
645                ..valid.clone()
646            },
647            Exemplar {
648                timestamp: i64::MAX,
649                ..valid.clone()
650            },
651        ] {
652            for (exemplars, samples, expected) in [
653                (vec![valid.clone()], vec![], 1),
654                (vec![invalid.clone()], vec![], 0),
655                (
656                    vec![invalid.clone(), valid.clone(), invalid, valid.clone()],
657                    vec![sample],
658                    2,
659                ),
660            ] {
661                let sample_count = samples.len();
662                decoder
663                    .decode(
664                        encode(TimeSeries {
665                            exemplars,
666                            samples,
667                            ..series.clone()
668                        }),
669                        PromValidationMode::Strict,
670                        &mut processor,
671                    )
672                    .unwrap();
673                let observations = decoder
674                    .table_data
675                    .exemplars
676                    .get(&None)
677                    .map(Vec::as_slice)
678                    .unwrap_or_default();
679                assert_eq!(observations.len(), expected);
680                for observation in observations {
681                    assert_eq!(observation.timestamp, -1_000_000);
682                    assert_eq!(observation.value, 0.5);
683                    assert_eq!(observation.labels["traceID"], "abc");
684                }
685                assert_eq!(
686                    decoder
687                        .as_row_insert_requests()
688                        .all_req()
689                        .map(|r| r.rows.unwrap().rows.len())
690                        .sum::<usize>(),
691                    sample_count
692                );
693            }
694            for name in ["__name__", "job"] {
695                let mut invalid = series.clone();
696                invalid.labels.extend([
697                    Label {
698                        name: name.into(),
699                        value: "a".into(),
700                    },
701                    Label {
702                        name: name.into(),
703                        value: "b".into(),
704                    },
705                ]);
706                assert!(
707                    decoder
708                        .decode(encode(invalid), PromValidationMode::Strict, &mut processor)
709                        .is_err()
710                );
711            }
712            // A pooled decoder must not leak rejected exemplars into its next request.
713            decoder
714                .decode(
715                    encode(TimeSeries {
716                        samples: vec![sample],
717                        exemplars: vec![],
718                        ..series.clone()
719                    }),
720                    PromValidationMode::Strict,
721                    &mut processor,
722                )
723                .unwrap();
724            assert!(decoder.table_data.exemplars.is_empty());
725            assert_eq!(decoder.as_row_insert_requests().all_req().count(), 1);
726        }
727        processor.use_pipeline = true;
728        assert!(
729            decoder
730                .decode(encode(series), PromValidationMode::Strict, &mut processor)
731                .unwrap_err()
732                .to_string()
733                .contains("pipelines with exemplars")
734        );
735    }
736
737    #[test]
738    fn test_decode_rejects_remote_write_v1_native_histograms() {
739        for samples in [
740            Vec::new(),
741            vec![Sample {
742                value: 1.0,
743                timestamp: 1000,
744            }],
745        ] {
746            let request = WriteRequest {
747                timeseries: vec![TimeSeries {
748                    samples,
749                    histograms: vec![Histogram::default()],
750                    ..Default::default()
751                }],
752                ..Default::default()
753            };
754            let mut processor = PromSeriesProcessor::default_processor();
755            let mut write_request = PromWriteRequest::default();
756            let error = write_request
757                .decode(
758                    request.encode_to_vec(),
759                    PromValidationMode::Strict,
760                    &mut processor,
761                )
762                .unwrap_err();
763
764            assert!(error.to_string().contains(
765                "remote write v1 native histogram ingestion is unsupported; use remote write v2"
766            ));
767            assert_eq!(
768                write_request.as_row_insert_requests().ref_all_req().count(),
769                0
770            );
771        }
772    }
773
774    #[test]
775    fn test_decode_clears_state_after_error() {
776        let label = |name: &str, value: &str| Label {
777            name: name.to_string(),
778            value: value.to_string(),
779        };
780        let sample = Sample {
781            value: 1.0,
782            timestamp: 1000,
783        };
784        let failed_request = WriteRequest {
785            timeseries: vec![
786                TimeSeries {
787                    labels: vec![
788                        label("__name__", "stale_metric"),
789                        label("stale_label", "stale_value"),
790                    ],
791                    samples: vec![sample],
792                    ..Default::default()
793                },
794                TimeSeries {
795                    labels: vec![
796                        label("__name__", "rejected_metric"),
797                        label("rejected_label", "rejected_value"),
798                        label("__schema__", "rejected_schema"),
799                        label("x_greptime_physical_table", "rejected_physical_table"),
800                    ],
801                    samples: vec![sample],
802                    histograms: vec![Histogram::default()],
803                    ..Default::default()
804                },
805            ],
806            ..Default::default()
807        };
808        let successful_request = WriteRequest {
809            timeseries: vec![TimeSeries {
810                labels: vec![
811                    label("__name__", "fresh_metric"),
812                    label("fresh_label", "fresh_value"),
813                ],
814                samples: vec![sample],
815                ..Default::default()
816            }],
817            ..Default::default()
818        };
819        let mut processor = PromSeriesProcessor::default_processor();
820        let mut write_request = PromWriteRequest::default();
821
822        write_request
823            .decode(
824                failed_request.encode_to_vec(),
825                PromValidationMode::Strict,
826                &mut processor,
827            )
828            .unwrap_err();
829        write_request
830            .decode(
831                successful_request.encode_to_vec(),
832                PromValidationMode::Strict,
833                &mut processor,
834            )
835            .unwrap();
836
837        assert_eq!(write_request.table_data.tables.len(), 1);
838        let (prom_ctx, tables) = write_request.table_data.tables.iter().next().unwrap();
839        assert_eq!(prom_ctx.schema, None);
840        assert_eq!(prom_ctx.physical_table, None);
841        assert_eq!(tables.len(), 1);
842        assert!(tables.contains_key("fresh_metric"));
843
844        let requests = write_request
845            .as_row_insert_requests()
846            .all_req()
847            .collect::<Vec<_>>();
848        assert_eq!(requests.len(), 1);
849        assert_eq!(requests[0].table_name, "fresh_metric");
850        let rows = requests[0].rows.as_ref().unwrap();
851        assert_eq!(rows.rows.len(), 1);
852        assert_eq!(
853            rows.schema
854                .iter()
855                .map(|column| column.column_name.as_str())
856                .collect::<Vec<_>>(),
857            vec![greptime_timestamp(), greptime_value(), "fresh_label"]
858        );
859    }
860
861    #[test]
862    fn test_decode_string_strict_mode_valid_utf8() {
863        let valid_utf8 = Bytes::from("hello world");
864        let result = PromValidationMode::Strict.decode_string(&valid_utf8);
865        assert!(result.is_ok());
866        assert_eq!(result.unwrap(), "hello world");
867    }
868
869    #[test]
870    fn test_decode_string_all_modes_ascii() {
871        let ascii = Bytes::from("simple_ascii_123");
872        let strict_result = PromValidationMode::Strict.decode_string(&ascii).unwrap();
873        let lossy_result = PromValidationMode::Lossy.decode_string(&ascii).unwrap();
874        let unchecked_result = PromValidationMode::Unchecked.decode_string(&ascii).unwrap();
875        assert_eq!(strict_result, "simple_ascii_123");
876        assert_eq!(lossy_result, "simple_ascii_123");
877        assert_eq!(unchecked_result, "simple_ascii_123");
878    }
879}
servers/prom_remote_write/decode.rs

servers/prom_remote_write/
decode.rs
