servers/http/
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
15use std::sync::Arc;
16
17use axum::Extension;
18use axum::extract::State;
19use axum::http::{StatusCode, header};
20use axum::response::IntoResponse;
21use axum_extra::TypedHeader;
22use bytes::Bytes;
23use common_catalog::consts::{TRACE_TABLE_NAME, TRACE_TABLE_NAME_SESSION_KEY};
24use common_telemetry::tracing;
25use headers::ContentType;
26use mime_guess::mime;
27use opentelemetry_proto::tonic::collector::logs::v1::{
28    ExportLogsServiceRequest, ExportLogsServiceResponse,
29};
30use opentelemetry_proto::tonic::collector::metrics::v1::{
31    ExportMetricsPartialSuccess, ExportMetricsServiceResponse,
32};
33use opentelemetry_proto::tonic::collector::trace::v1::{
34    ExportTracePartialSuccess, ExportTraceServiceRequest, ExportTraceServiceResponse,
35};
36use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest;
37use pipeline::PipelineWay;
38use prost::Message;
39use session::context::{Channel, QueryContext};
40use session::protocol_ctx::{MetricType, OtlpMetricCtx, ProtocolCtx};
41use snafu::prelude::*;
42
43use crate::error::{self, PipelineSnafu, Result};
44use crate::http::extractor::{
45    LogTableName, OtlpMetricOptions, PipelineInfo, SelectInfoWrapper, TraceTableName,
46};
47use crate::http::header::{CONTENT_TYPE_PROTOBUF, write_cost_header_map};
48use crate::metrics::METRIC_HTTP_OPENTELEMETRY_LOGS_ELAPSED;
49use crate::query_handler::{
50    MetricsIngestOutcome, OpenTelemetryProtocolHandlerRef, PipelineHandler, TraceIngestOutcome,
51};
52
53#[derive(Clone, prost::Message)]
54pub struct GoogleRpcStatus {
55    #[prost(int32, tag = "1")]
56    pub code: i32,
57    #[prost(string, tag = "2")]
58    pub message: String,
59}
60
61fn is_json_content_type(content_type: Option<&ContentType>) -> bool {
62    match content_type {
63        None => false,
64        Some(ct) => {
65            let mime: mime::Mime = ct.clone().into();
66            mime.subtype() == mime::JSON
67        }
68    }
69}
70
71fn content_type_to_string(content_type: Option<&TypedHeader<ContentType>>) -> String {
72    content_type
73        .map(|h| h.0.to_string())
74        .unwrap_or_else(|| "not specified".to_string())
75}
76
77#[derive(Clone)]
78pub struct OtlpState {
79    pub with_metric_engine: bool,
80    pub handler: OpenTelemetryProtocolHandlerRef,
81}
82
83#[axum_macros::debug_handler]
84#[tracing::instrument(skip_all, fields(protocol = "otlp", request_type = "metrics"))]
85pub async fn metrics(
86    State(state): State<OtlpState>,
87    Extension(mut query_ctx): Extension<QueryContext>,
88    http_opts: OtlpMetricOptions,
89    content_type: Option<TypedHeader<ContentType>>,
90    bytes: Bytes,
91) -> Result<OtlpMetricsResponse> {
92    if is_json_content_type(content_type.as_ref().map(|h| &h.0)) {
93        return error::UnsupportedJsonContentTypeSnafu {}.fail();
94    }
95
96    let db = query_ctx.get_db_string();
97    query_ctx.set_channel(Channel::Otlp);
98
99    let _timer = crate::metrics::METRIC_HTTP_OPENTELEMETRY_METRICS_ELAPSED
100        .with_label_values(&[db.as_str()])
101        .start_timer();
102    let request = ExportMetricsServiceRequest::decode(bytes).with_context(|_| {
103        error::DecodeOtlpRequestSnafu {
104            content_type: content_type_to_string(content_type.as_ref()),
105        }
106    })?;
107
108    let OtlpState {
109        with_metric_engine,
110        handler,
111    } = state;
112
113    query_ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx {
114        promote_all_resource_attrs: http_opts.promote_all_resource_attrs,
115        resource_attrs: http_opts.resource_attrs,
116        promote_scope_attrs: http_opts.promote_scope_attrs,
117        with_metric_engine,
118        // set by the frontend from its config
119        resource_info: false,
120        metric_type: MetricType::Init,
121        metric_translation_strategy: http_opts.metric_translation_strategy,
122    }));
123    let query_ctx = Arc::new(query_ctx);
124
125    match handler.metrics(request, query_ctx).await {
126        Ok(outcome) => {
127            if outcome.accepted_data_points == 0 && outcome.rejected_data_points > 0 {
128                Ok(OtlpMetricsResponse::Failure(outcome))
129            } else if outcome.rejected_data_points > 0 || outcome.error_message.is_some() {
130                Ok(OtlpMetricsResponse::PartialSuccess(outcome))
131            } else {
132                Ok(OtlpMetricsResponse::FullSuccess(outcome))
133            }
134        }
135        Err(error::Error::InvalidOtlpMetricInput { reason }) => {
136            Ok(OtlpMetricsResponse::Failure(MetricsIngestOutcome {
137                error_message: Some(reason),
138                ..Default::default()
139            }))
140        }
141        Err(error) => Err(error),
142    }
143}
144
145#[axum_macros::debug_handler]
146#[tracing::instrument(skip_all, fields(protocol = "otlp", request_type = "traces"))]
147pub async fn traces(
148    State(state): State<OtlpState>,
149    TraceTableName(table_name): TraceTableName,
150    pipeline_info: PipelineInfo,
151    Extension(mut query_ctx): Extension<QueryContext>,
152    content_type: Option<TypedHeader<ContentType>>,
153    bytes: Bytes,
154) -> Result<OtlpTraceResponse> {
155    if is_json_content_type(content_type.as_ref().map(|h| &h.0)) {
156        return error::UnsupportedJsonContentTypeSnafu {}.fail();
157    }
158
159    let db = query_ctx.get_db_string();
160    let table_name = table_name.unwrap_or_else(|| TRACE_TABLE_NAME.to_string());
161
162    query_ctx.set_channel(Channel::Otlp);
163    query_ctx.set_extension(TRACE_TABLE_NAME_SESSION_KEY, &table_name);
164
165    let query_ctx = Arc::new(query_ctx);
166    let _timer = crate::metrics::METRIC_HTTP_OPENTELEMETRY_TRACES_ELAPSED
167        .with_label_values(&[db.as_str()])
168        .start_timer();
169    let request = ExportTraceServiceRequest::decode(bytes).with_context(|_| {
170        error::DecodeOtlpRequestSnafu {
171            content_type: content_type_to_string(content_type.as_ref()),
172        }
173    })?;
174
175    let pipeline = PipelineWay::from_name_and_default(
176        pipeline_info.pipeline_name.as_deref(),
177        pipeline_info.pipeline_version.as_deref(),
178        None,
179    )
180    .context(PipelineSnafu)?;
181
182    let pipeline_params = pipeline_info.pipeline_params;
183
184    let OtlpState { handler, .. } = state;
185
186    // here we use nightly feature `trait_upcasting` to convert handler to
187    // pipeline_handler
188    let pipeline_handler: Arc<dyn PipelineHandler + Send + Sync> = handler.clone();
189
190    handler
191        .traces(
192            pipeline_handler,
193            request,
194            pipeline,
195            pipeline_params,
196            table_name,
197            query_ctx,
198        )
199        .await
200        .map(|outcome| {
201            if outcome.accepted_spans == 0 && outcome.rejected_spans > 0 {
202                OtlpTraceResponse::Failure(outcome)
203            } else if outcome.rejected_spans > 0 || outcome.error_message.is_some() {
204                OtlpTraceResponse::PartialSuccess(outcome)
205            } else {
206                OtlpTraceResponse::FullSuccess(outcome)
207            }
208        })
209}
210
211#[axum_macros::debug_handler]
212#[tracing::instrument(skip_all, fields(protocol = "otlp", request_type = "logs"))]
213pub async fn logs(
214    State(state): State<OtlpState>,
215    Extension(mut query_ctx): Extension<QueryContext>,
216    pipeline_info: PipelineInfo,
217    LogTableName(tablename): LogTableName,
218    SelectInfoWrapper(select_info): SelectInfoWrapper,
219    content_type: Option<TypedHeader<ContentType>>,
220    bytes: Bytes,
221) -> Result<OtlpResponse<ExportLogsServiceResponse>> {
222    if is_json_content_type(content_type.as_ref().map(|h| &h.0)) {
223        return error::UnsupportedJsonContentTypeSnafu {}.fail();
224    }
225
226    let tablename = tablename.unwrap_or_else(|| "opentelemetry_logs".to_string());
227    let db = query_ctx.get_db_string();
228    query_ctx.set_channel(Channel::Otlp);
229    let query_ctx = Arc::new(query_ctx);
230    let _timer = METRIC_HTTP_OPENTELEMETRY_LOGS_ELAPSED
231        .with_label_values(&[db.as_str()])
232        .start_timer();
233    let request = ExportLogsServiceRequest::decode(bytes).with_context(|_| {
234        error::DecodeOtlpRequestSnafu {
235            content_type: content_type_to_string(content_type.as_ref()),
236        }
237    })?;
238
239    let pipeline = PipelineWay::from_name_and_default(
240        pipeline_info.pipeline_name.as_deref(),
241        pipeline_info.pipeline_version.as_deref(),
242        Some(PipelineWay::OtlpLogDirect(Box::new(select_info))),
243    )
244    .context(PipelineSnafu)?;
245    let pipeline_params = pipeline_info.pipeline_params;
246
247    let OtlpState { handler, .. } = state;
248
249    // here we use nightly feature `trait_upcasting` to convert handler to
250    // pipeline_handler
251    let pipeline_handler: Arc<dyn PipelineHandler + Send + Sync> = handler.clone();
252    handler
253        .logs(
254            pipeline_handler,
255            request,
256            pipeline,
257            pipeline_params,
258            tablename,
259            query_ctx,
260        )
261        .await
262        .map(|o| OtlpResponse {
263            resp_body: ExportLogsServiceResponse {
264                partial_success: None,
265            },
266            write_cost: o.iter().map(|o| o.meta.cost).sum(),
267        })
268}
269
270pub struct OtlpResponse<T: Message> {
271    resp_body: T,
272    write_cost: usize,
273}
274
275pub enum OtlpMetricsResponse {
276    FullSuccess(MetricsIngestOutcome),
277    PartialSuccess(MetricsIngestOutcome),
278    Failure(MetricsIngestOutcome),
279}
280
281impl IntoResponse for OtlpMetricsResponse {
282    fn into_response(self) -> axum::response::Response {
283        match self {
284            OtlpMetricsResponse::FullSuccess(outcome) => {
285                let mut header_map = write_cost_header_map(outcome.write_cost);
286                header_map.insert(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.clone());
287                let body = ExportMetricsServiceResponse {
288                    partial_success: None,
289                };
290                (header_map, body.encode_to_vec()).into_response()
291            }
292            OtlpMetricsResponse::PartialSuccess(outcome) => {
293                let mut header_map = write_cost_header_map(outcome.write_cost);
294                header_map.insert(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.clone());
295                let body = ExportMetricsServiceResponse {
296                    partial_success: Some(ExportMetricsPartialSuccess {
297                        rejected_data_points: outcome.rejected_data_points,
298                        error_message: outcome.error_message.unwrap_or_default(),
299                    }),
300                };
301                (header_map, body.encode_to_vec()).into_response()
302            }
303            OtlpMetricsResponse::Failure(outcome) => {
304                let status = GoogleRpcStatus {
305                    code: tonic::Code::InvalidArgument as i32,
306                    message: outcome.error_message.unwrap_or_default(),
307                };
308                (
309                    StatusCode::BAD_REQUEST,
310                    [(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.as_ref())],
311                    status.encode_to_vec(),
312                )
313                    .into_response()
314            }
315        }
316    }
317}
318
319impl<T: Message> IntoResponse for OtlpResponse<T> {
320    fn into_response(self) -> axum::response::Response {
321        let mut header_map = write_cost_header_map(self.write_cost);
322        header_map.insert(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.clone());
323
324        (header_map, self.resp_body.encode_to_vec()).into_response()
325    }
326}
327
328pub enum OtlpTraceResponse {
329    FullSuccess(TraceIngestOutcome),
330    PartialSuccess(TraceIngestOutcome),
331    Failure(TraceIngestOutcome),
332}
333
334impl IntoResponse for OtlpTraceResponse {
335    fn into_response(self) -> axum::response::Response {
336        match self {
337            OtlpTraceResponse::FullSuccess(outcome) => {
338                let mut header_map = write_cost_header_map(outcome.write_cost);
339                header_map.insert(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.clone());
340                let body = ExportTraceServiceResponse {
341                    partial_success: None,
342                };
343                (header_map, body.encode_to_vec()).into_response()
344            }
345            OtlpTraceResponse::PartialSuccess(outcome) => {
346                let mut header_map = write_cost_header_map(outcome.write_cost);
347                header_map.insert(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.clone());
348                let body = ExportTraceServiceResponse {
349                    partial_success: outcome.error_message.map(|error_message| {
350                        ExportTracePartialSuccess {
351                            rejected_spans: outcome.rejected_spans as i64,
352                            error_message,
353                        }
354                    }),
355                };
356                (header_map, body.encode_to_vec()).into_response()
357            }
358            OtlpTraceResponse::Failure(outcome) => {
359                let status = GoogleRpcStatus {
360                    code: 0,
361                    message: outcome.error_message.unwrap_or_default(),
362                };
363                (
364                    StatusCode::BAD_REQUEST,
365                    [(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.as_ref())],
366                    status.encode_to_vec(),
367                )
368                    .into_response()
369            }
370        }
371    }
372}
373
374#[cfg(test)]
375mod tests;
servers/http/otlp.rs

servers/http/
otlp.rs
