Skip to main content

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