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

servers/http/
otlp.rs
