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::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};
48
49fn is_json_content_type(content_type: Option<&ContentType>) -> bool {
50    match content_type {
51        None => false,
52        Some(ct) => {
53            let mime: mime::Mime = ct.clone().into();
54            mime.subtype() == mime::JSON
55        }
56    }
57}
58
59fn content_type_to_string(content_type: Option<&TypedHeader<ContentType>>) -> String {
60    content_type
61        .map(|h| h.0.to_string())
62        .unwrap_or_else(|| "not specified".to_string())
63}
64
65#[derive(Clone)]
66pub struct OtlpState {
67    pub with_metric_engine: bool,
68    pub handler: OpenTelemetryProtocolHandlerRef,
69}
70
71#[axum_macros::debug_handler]
72#[tracing::instrument(skip_all, fields(protocol = "otlp", request_type = "metrics"))]
73pub async fn metrics(
74    State(state): State<OtlpState>,
75    Extension(mut query_ctx): Extension<QueryContext>,
76    http_opts: OtlpMetricOptions,
77    content_type: Option<TypedHeader<ContentType>>,
78    bytes: Bytes,
79) -> Result<OtlpResponse<ExportMetricsServiceResponse>> {
80    if is_json_content_type(content_type.as_ref().map(|h| &h.0)) {
81        return error::UnsupportedJsonContentTypeSnafu {}.fail();
82    }
83
84    let db = query_ctx.get_db_string();
85    query_ctx.set_channel(Channel::Otlp);
86
87    let _timer = crate::metrics::METRIC_HTTP_OPENTELEMETRY_METRICS_ELAPSED
88        .with_label_values(&[db.as_str()])
89        .start_timer();
90    let request = ExportMetricsServiceRequest::decode(bytes).with_context(|_| {
91        error::DecodeOtlpRequestSnafu {
92            content_type: content_type_to_string(content_type.as_ref()),
93        }
94    })?;
95
96    let OtlpState {
97        with_metric_engine,
98        handler,
99    } = state;
100
101    query_ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx {
102        promote_all_resource_attrs: http_opts.promote_all_resource_attrs,
103        resource_attrs: http_opts.resource_attrs,
104        promote_scope_attrs: http_opts.promote_scope_attrs,
105        with_metric_engine,
106        // set is_legacy later
107        is_legacy: false,
108        metric_type: MetricType::Init,
109    }));
110    let query_ctx = Arc::new(query_ctx);
111
112    handler
113        .metrics(request, query_ctx)
114        .await
115        .map(|o| OtlpResponse {
116            resp_body: ExportMetricsServiceResponse {
117                partial_success: None,
118            },
119            write_cost: o.meta.cost,
120        })
121}
122
123#[axum_macros::debug_handler]
124#[tracing::instrument(skip_all, fields(protocol = "otlp", request_type = "traces"))]
125pub async fn traces(
126    State(state): State<OtlpState>,
127    TraceTableName(table_name): TraceTableName,
128    pipeline_info: PipelineInfo,
129    Extension(mut query_ctx): Extension<QueryContext>,
130    content_type: Option<TypedHeader<ContentType>>,
131    bytes: Bytes,
132) -> Result<OtlpResponse<ExportTraceServiceResponse>> {
133    if is_json_content_type(content_type.as_ref().map(|h| &h.0)) {
134        return error::UnsupportedJsonContentTypeSnafu {}.fail();
135    }
136
137    let db = query_ctx.get_db_string();
138    let table_name = table_name.unwrap_or_else(|| TRACE_TABLE_NAME.to_string());
139
140    query_ctx.set_channel(Channel::Otlp);
141    query_ctx.set_extension(TRACE_TABLE_NAME_SESSION_KEY, &table_name);
142
143    let query_ctx = Arc::new(query_ctx);
144    let _timer = crate::metrics::METRIC_HTTP_OPENTELEMETRY_TRACES_ELAPSED
145        .with_label_values(&[db.as_str()])
146        .start_timer();
147    let request = ExportTraceServiceRequest::decode(bytes).with_context(|_| {
148        error::DecodeOtlpRequestSnafu {
149            content_type: content_type_to_string(content_type.as_ref()),
150        }
151    })?;
152
153    let pipeline = PipelineWay::from_name_and_default(
154        pipeline_info.pipeline_name.as_deref(),
155        pipeline_info.pipeline_version.as_deref(),
156        None,
157    )
158    .context(PipelineSnafu)?;
159
160    let pipeline_params = pipeline_info.pipeline_params;
161
162    let OtlpState { handler, .. } = state;
163
164    // here we use nightly feature `trait_upcasting` to convert handler to
165    // pipeline_handler
166    let pipeline_handler: Arc<dyn PipelineHandler + Send + Sync> = handler.clone();
167
168    handler
169        .traces(
170            pipeline_handler,
171            request,
172            pipeline,
173            pipeline_params,
174            table_name,
175            query_ctx,
176        )
177        .await
178        .map(|outcome| OtlpResponse {
179            resp_body: ExportTraceServiceResponse {
180                partial_success: outcome.error_message.map(|error_message| {
181                    ExportTracePartialSuccess {
182                        rejected_spans: outcome.rejected_spans as i64,
183                        error_message,
184                    }
185                }),
186            },
187            write_cost: outcome.write_cost,
188        })
189}
190
191#[axum_macros::debug_handler]
192#[tracing::instrument(skip_all, fields(protocol = "otlp", request_type = "logs"))]
193pub async fn logs(
194    State(state): State<OtlpState>,
195    Extension(mut query_ctx): Extension<QueryContext>,
196    pipeline_info: PipelineInfo,
197    LogTableName(tablename): LogTableName,
198    SelectInfoWrapper(select_info): SelectInfoWrapper,
199    content_type: Option<TypedHeader<ContentType>>,
200    bytes: Bytes,
201) -> Result<OtlpResponse<ExportLogsServiceResponse>> {
202    if is_json_content_type(content_type.as_ref().map(|h| &h.0)) {
203        return error::UnsupportedJsonContentTypeSnafu {}.fail();
204    }
205
206    let tablename = tablename.unwrap_or_else(|| "opentelemetry_logs".to_string());
207    let db = query_ctx.get_db_string();
208    query_ctx.set_channel(Channel::Otlp);
209    let query_ctx = Arc::new(query_ctx);
210    let _timer = METRIC_HTTP_OPENTELEMETRY_LOGS_ELAPSED
211        .with_label_values(&[db.as_str()])
212        .start_timer();
213    let request = ExportLogsServiceRequest::decode(bytes).with_context(|_| {
214        error::DecodeOtlpRequestSnafu {
215            content_type: content_type_to_string(content_type.as_ref()),
216        }
217    })?;
218
219    let pipeline = PipelineWay::from_name_and_default(
220        pipeline_info.pipeline_name.as_deref(),
221        pipeline_info.pipeline_version.as_deref(),
222        Some(PipelineWay::OtlpLogDirect(Box::new(select_info))),
223    )
224    .context(PipelineSnafu)?;
225    let pipeline_params = pipeline_info.pipeline_params;
226
227    let OtlpState { handler, .. } = state;
228
229    // here we use nightly feature `trait_upcasting` to convert handler to
230    // pipeline_handler
231    let pipeline_handler: Arc<dyn PipelineHandler + Send + Sync> = handler.clone();
232    handler
233        .logs(
234            pipeline_handler,
235            request,
236            pipeline,
237            pipeline_params,
238            tablename,
239            query_ctx,
240        )
241        .await
242        .map(|o| OtlpResponse {
243            resp_body: ExportLogsServiceResponse {
244                partial_success: None,
245            },
246            write_cost: o.iter().map(|o| o.meta.cost).sum(),
247        })
248}
249
250pub struct OtlpResponse<T: Message> {
251    resp_body: T,
252    write_cost: usize,
253}
254
255impl<T: Message> IntoResponse for OtlpResponse<T> {
256    fn into_response(self) -> axum::response::Response {
257        let mut header_map = write_cost_header_map(self.write_cost);
258        header_map.insert(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.clone());
259
260        (header_map, self.resp_body.encode_to_vec()).into_response()
261    }
262}
servers/http/otlp.rs

servers/http/
otlp.rs
