servers/postgres/copy_in/
otlp.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//! The `otlp` format for COPY FROM STDIN: one binary OTLP payload for
16//! any of the three signals — `ExportMetricsServiceRequest`,
17//! `ExportTraceServiceRequest` or `ExportLogsServiceRequest`.
18//!
19//! `COPY FROM STDIN WITH (FORMAT otlp [, signal 'traces'] [, <options>])`
20//! is a GreptimeDB-specific, **binary** format (wire format code 1).
21//! Like the influxdb format it is table-less and does not decode here:
22//! the whole payload is buffered until `CopyDone` — exactly like an
23//! OTLP/HTTP request body — then decoded and completed through the
24//! frontend's `OpenTelemetryProtocolHandler`, the same path the
25//! `/v1/otlp/v1/{metrics,traces,logs}` endpoints use, under the
26//! connection's regular postgres query context.
27//!
28//! The `WITH` options carry what the HTTP endpoints take as request
29//! headers; values already provided by the postgres session (database,
30//! user, permissions) are not repeated. Option semantics, defaults and
31//! parsing helpers are shared with the HTTP extractors (`truthy`,
32//! `SelectInfo`, `GreptimePipelineParams`, `PipelineWay`), so the two
33//! ingestion routes cannot drift.
34//!
35//! One COPY statement is one bounded ingestion request (see
36//! [`crate::postgres::copy_in::COPY_IN_MAX_PAYLOAD_BYTES`]).
37
38use ahash::HashSet;
39
40use pgwire::error::PgWireResult;
41use pipeline::truthy;
42use session::protocol_ctx::OtlpMetricTranslationStrategy;
43
44use crate::postgres::copy_in::bad_copy_data;
45use crate::postgres::copy_in::codec::CopyFormat;
46use crate::postgres::copy_in::options::{CopyOptionSet, CopyOptionValue};
47
48/// The format's validated options: the signal and its header-equivalent
49/// configuration, carried to the flush that completes the copy.
50#[derive(Debug, Clone)]
51pub(super) struct OtlpOptions {
52    pub(super) signal: OtlpSignal,
53}
54
55#[derive(Debug, Clone)]
56pub(super) enum OtlpSignal {
57    /// Metric options mirroring the `x-greptime-otlp-metric-*` headers.
58    Metrics {
59        /// Persist all resource attributes; `resource_attrs` then acts
60        /// as an exclude list, otherwise as an include list.
61        promote_all_resource_attrs: bool,
62        resource_attrs: HashSet<String>,
63        promote_scope_attrs: bool,
64        translation_strategy: OtlpMetricTranslationStrategy,
65    },
66    /// Trace options mirroring `x-greptime-trace-table-name` and the
67    /// pipeline headers. The raw pipeline-params string is parsed by the
68    /// shared `GreptimePipelineParams::from_params` when the copy
69    /// completes, like the HTTP request does.
70    Traces {
71        table_name: Option<String>,
72        pipeline_name: Option<String>,
73        pipeline_version: Option<String>,
74        pipeline_params: Option<String>,
75    },
76    /// Log options mirroring `x-greptime-log-table-name`,
77    /// `x-greptime-log-extract-keys` and the pipeline headers.
78    Logs {
79        table_name: Option<String>,
80        pipeline_name: Option<String>,
81        pipeline_version: Option<String>,
82        pipeline_params: Option<String>,
83        extract_keys: Option<String>,
84    },
85}
86
87pub(super) fn build_otlp_codec(mut options: CopyOptionSet) -> PgWireResult<CopyFormat> {
88    let signal = options.take_string("signal")?;
89    let metric_options = (
90        take_truthy(&mut options, "promote_all_resource_attrs")?,
91        options.take_string("resource_attrs")?,
92        take_truthy(&mut options, "promote_scope_attrs")?,
93        options.take_string("translation_strategy")?,
94    );
95    let trace_log_options = (
96        options.take_string("table_name")?,
97        options.take_string("pipeline_name")?,
98        options.take_string("pipeline_version")?,
99        options.take_string("pipeline_params")?,
100    );
101    let extract_keys = options.take_string("extract_keys")?;
102    options.reject_leftovers("otlp")?;
103
104    let unexpected = |option: &str| {
105        bad_copy_data(format!(
106            "COPY option \"{option}\" does not apply to the selected signal"
107        ))
108    };
109    let signal = match signal.as_deref() {
110        None | Some("metrics") => {
111            if trace_log_options != (None, None, None, None) || extract_keys.is_some() {
112                return Err(unexpected("table_name/pipeline/extract_keys"));
113            }
114            let (promote_all, resource_attrs, promote_scope, strategy) = metric_options;
115            let promote_all = promote_all.unwrap_or(false);
116            let promote_scope = promote_scope.unwrap_or(false);
117            let translation_strategy = match strategy.as_deref() {
118                None => OtlpMetricTranslationStrategy::default(),
119                // Same values as the HTTP header (`VALUES`), parsed by
120                // the same `FromStr`.
121                Some(value) => value.parse().map_err(|_| {
122                    bad_copy_data(format!(
123                        "invalid translation_strategy \"{value}\", expected one of: {}",
124                        OtlpMetricTranslationStrategy::VALUES.join(", ")
125                    ))
126                })?,
127            };
128            let resource_attrs = resource_attrs
129                // Same ';'-separated list and trimming as the header
130                // extractor.
131                .map(|s| s.split(';').map(|s| s.trim().to_string()).collect())
132                .unwrap_or_default();
133            OtlpSignal::Metrics {
134                promote_all_resource_attrs: promote_all,
135                resource_attrs,
136                promote_scope_attrs: promote_scope,
137                translation_strategy,
138            }
139        }
140        Some("traces") | Some("logs") => {
141            if metric_options != (None, None, None, None) {
142                return Err(unexpected("promote_*/resource_attrs/translation_strategy"));
143            }
144            let (table_name, pipeline_name, pipeline_version, pipeline_params) = trace_log_options;
145            if signal.as_deref() == Some("traces") {
146                if extract_keys.is_some() {
147                    // `extract_keys` only applies to logs.
148                    return Err(unexpected("extract_keys"));
149                }
150                OtlpSignal::Traces {
151                    table_name,
152                    pipeline_name,
153                    pipeline_version,
154                    pipeline_params,
155                }
156            } else {
157                OtlpSignal::Logs {
158                    table_name,
159                    pipeline_name,
160                    pipeline_version,
161                    pipeline_params,
162                    extract_keys,
163                }
164            }
165        }
166        Some(unknown) => {
167            return Err(bad_copy_data(format!(
168                "invalid signal \"{unknown}\", expected one of: metrics, traces, logs"
169            )));
170        }
171    };
172    Ok(CopyFormat::Otlp(OtlpOptions { signal }))
173}
174
175/// Takes a boolean option with the HTTP header's exact truthy
176/// spellings (`true`, `1`, `yes`, `on`, `t`), plus the COPY forms
177/// (bare flag or `TRUE`/`FALSE`), so both routes accept the same values.
178fn take_truthy(options: &mut CopyOptionSet, key: &str) -> PgWireResult<Option<bool>> {
179    match options.take_value(key) {
180        None => Ok(None),
181        Some(CopyOptionValue::Flag) => Ok(Some(true)),
182        Some(CopyOptionValue::Boolean(b)) => Ok(Some(b)),
183        Some(CopyOptionValue::String(v)) | Some(CopyOptionValue::Ident(v)) => Ok(Some(truthy(v))),
184        Some(other) => Err(bad_copy_data(format!(
185            "COPY option {key} requires a boolean value, got {}",
186            other.kind()
187        ))),
188    }
189}
190
191#[cfg(test)]
192mod tests {
193    use super::*;
194    use crate::postgres::copy_in::options::CopyOptionValue;
195
196    fn opts(entries: &[(&str, CopyOptionValue)]) -> CopyOptionSet {
197        let mut set = CopyOptionSet::default();
198        for (k, v) in entries {
199            set.insert((*k).into(), v.clone()).unwrap();
200        }
201        set
202    }
203
204    #[test]
205    fn test_signal_defaults_and_options() {
206        // Default signal is metrics.
207        assert!(matches!(
208            build_otlp_codec(CopyOptionSet::default()).unwrap(),
209            CopyFormat::Otlp(OtlpOptions {
210                signal: OtlpSignal::Metrics { .. }
211            })
212        ));
213
214        // Metric options parse like the headers.
215        let format = build_otlp_codec(opts(&[
216            ("promote_all_resource_attrs", CopyOptionValue::Flag),
217            (
218                "resource_attrs",
219                CopyOptionValue::String("service.name; host.id ".into()),
220            ),
221            ("promote_scope_attrs", CopyOptionValue::Flag),
222            (
223                "translation_strategy",
224                CopyOptionValue::String("NoTranslation".into()),
225            ),
226        ]))
227        .unwrap();
228        let CopyFormat::Otlp(OtlpOptions {
229            signal:
230                OtlpSignal::Metrics {
231                    promote_all_resource_attrs,
232                    resource_attrs,
233                    promote_scope_attrs,
234                    translation_strategy,
235                },
236        }) = format
237        else {
238            panic!("expected metrics");
239        };
240        assert!(promote_all_resource_attrs && promote_scope_attrs);
241        assert_eq!(
242            resource_attrs,
243            ["service.name", "host.id"]
244                .into_iter()
245                .map(String::from)
246                .collect()
247        );
248        assert!(matches!(
249            translation_strategy,
250            OtlpMetricTranslationStrategy::NoTranslation
251        ));
252
253        // Traces and logs accept their header equivalents.
254        assert!(matches!(
255            build_otlp_codec(opts(&[(
256                "signal",
257                CopyOptionValue::String("traces".into())
258            )]))
259            .unwrap(),
260            CopyFormat::Otlp(OtlpOptions {
261                signal: OtlpSignal::Traces { .. }
262            })
263        ));
264        assert!(matches!(
265            build_otlp_codec(opts(&[
266                ("signal", CopyOptionValue::String("logs".into())),
267                ("table_name", CopyOptionValue::String("my_logs".into())),
268                ("extract_keys", CopyOptionValue::String("a,b".into())),
269            ]))
270            .unwrap(),
271            CopyFormat::Otlp(OtlpOptions {
272                signal: OtlpSignal::Logs { .. }
273            })
274        ));
275    }
276
277    #[test]
278    fn test_rejects_unknown_and_mismatched_options() {
279        // Unknown signal.
280        let set = opts(&[("signal", CopyOptionValue::String("spans".into()))]);
281        assert!(build_otlp_codec(set).is_err());
282
283        // Unknown option.
284        let set = opts(&[("precision", CopyOptionValue::String("ns".into()))]);
285        assert!(build_otlp_codec(set).is_err());
286
287        // Signal-specific options do not cross.
288        let set = opts(&[
289            ("signal", CopyOptionValue::String("traces".into())),
290            ("promote_scope_attrs", CopyOptionValue::Flag),
291        ]);
292        assert!(build_otlp_codec(set).is_err());
293        let set = opts(&[
294            ("signal", CopyOptionValue::String("metrics".into())),
295            ("table_name", CopyOptionValue::String("t".into())),
296        ]);
297        assert!(build_otlp_codec(set).is_err());
298        let set = opts(&[
299            ("signal", CopyOptionValue::String("traces".into())),
300            (
301                "extract_keys",
302                CopyOptionValue::String("attributes.foo".into()),
303            ),
304        ]);
305        assert!(build_otlp_codec(set).is_err());
306    }
307}
servers/postgres/copy_in/otlp.rs

servers/postgres/copy_in/
otlp.rs
