frontend/instance/

otlp.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

mod trace_ingest;
pub mod trace_semconv;
pub mod trace_types;

use std::collections::HashSet;
use std::sync::Arc;

use api::v1::RowInsertRequests;
use async_trait::async_trait;
use auth::{
    OTLP_WRITE, PermissionChecker, PermissionCheckerRef, PermissionReq, PermissionTableTarget,
    PermissionTableTargets,
};
use client::Output;
use common_catalog::consts::{trace_operations_table_name, trace_services_table_name};
use common_error::ext::BoxedError;
use common_query::prelude::GREPTIME_PHYSICAL_TABLE;
use common_query::{OutputData, OutputMeta};
use common_telemetry::{tracing, warn};
use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest;
use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest;
use operator::insert::{Inserter, admit_row_insert_batches, admit_write};
use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest;
use pipeline::{GreptimePipelineParams, PipelineWay};
use servers::error::{self, AuthSnafu, Result as ServerResult};
use servers::http::prom_store::PHYSICAL_TABLE_PARAM;
use servers::interceptor::{OpenTelemetryProtocolInterceptor, OpenTelemetryProtocolInterceptorRef};
use servers::otlp;
use servers::otlp::trace::span::TraceSpanGroup;
use servers::query_handler::{
    MetricsIngestOutcome, OpenTelemetryProtocolHandler, PipelineHandlerRef, TraceIngestOutcome,
};
use session::context::QueryContextRef;
use session::protocol_ctx::OtlpMetricCtx;
use snafu::ResultExt;
use table::requests::{
    OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM, SEMANTIC_PER_TABLE_INDEX_KEY,
    SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_LOG, SIGNAL_TYPE_METRIC,
    SOURCE_OPENTELEMETRY,
};

pub(crate) use self::trace_ingest::TraceAuxCache;
use self::trace_ingest::trace_conventions;
use crate::instance::Instance;
use crate::metrics::{OTLP_LOGS_ROWS, OTLP_METRICS_ROWS, OTLP_RESOURCE_INFO_WRITE_ERRORS};

fn trace_permission_targets(
    table_name: &str,
    groups: &[TraceSpanGroup],
    ctx: &QueryContextRef,
) -> PermissionTableTargets {
    let catalog = ctx.current_catalog();
    let schema = ctx.current_schema();
    if catalog.is_empty() || schema.is_empty() || table_name.is_empty() {
        return PermissionTableTargets::Unresolved;
    }

    let has_spans = groups.iter().any(|group| !group.spans.is_empty());
    if !has_spans {
        return PermissionTableTargets::resolved(Vec::new());
    }

    let mut targets = vec![PermissionTableTarget::new(catalog, &schema, table_name)];
    if groups
        .iter()
        .flat_map(|group| &group.spans)
        .any(|span| span.service_name.is_some())
    {
        targets.extend([
            PermissionTableTarget::new(catalog, &schema, trace_services_table_name(table_name)),
            PermissionTableTarget::new(catalog, &schema, trace_operations_table_name(table_name)),
        ]);
    }

    PermissionTableTargets::resolved(targets)
}

#[async_trait]
impl OpenTelemetryProtocolHandler for Instance {
    #[tracing::instrument(skip_all)]
    async fn metrics(
        &self,
        request: ExportMetricsServiceRequest,
        ctx: QueryContextRef,
    ) -> ServerResult<MetricsIngestOutcome> {
        self.plugins
            .get::<PermissionCheckerRef>()
            .as_ref()
            .check_permission(ctx.current_user(), PermissionReq::Action(OTLP_WRITE))
            .context(AuthSnafu)?;

        let interceptor_ref = self
            .plugins
            .get::<OpenTelemetryProtocolInterceptorRef<servers::error::Error>>();
        interceptor_ref.pre_execute(ctx.clone())?;
        let ctx = Arc::new(ctx.fork());

        let mut metric_ctx = ctx
            .protocol_ctx()
            .get_otlp_metric_ctx()
            .cloned()
            .unwrap_or_default();
        metric_ctx.resource_info = self.otlp_resource_info;

        let mut prepared = otlp::metrics::prepare_metrics(request, &mut metric_ctx)?;
        // Check all destinations before preflight can expose schemas or discard points.
        let targets = {
            let catalog = ctx.current_catalog();
            let schema = ctx.current_schema();
            let mut table_names = HashSet::new();
            prepared
                .requests
                .iter()
                .filter(|request| table_names.insert(request.table_name.as_str()))
                .map(|request| PermissionTableTarget::new(catalog, &schema, &request.table_name))
                .collect()
        };
        self.check_table_permission(
            &ctx,
            PermissionReq::Action(OTLP_WRITE),
            PermissionTableTargets::resolved(targets),
        )
        .context(AuthSnafu)?;

        let preflight = self
            .inserter
            .validate_and_normalize_otlp_metrics(
                &mut prepared.requests,
                &ctx,
                metric_ctx.with_metric_engine,
                ctx.extension(PHYSICAL_TABLE_PARAM)
                    .unwrap_or(GREPTIME_PHYSICAL_TABLE),
            )
            .await
            .map_err(BoxedError::new)
            .context(error::ExecuteGrpcQuerySnafu)?;
        let rejections = prepared.point_rejections(preflight.rejections)?;
        let otlp::metrics::MetricsConversion {
            requests,
            rows,
            semantic_index,
            resource_info,
            mut outcome,
        } = prepared.finish(rejections, &preflight.physical_tables)?;
        if outcome.rejected_data_points > 0 {
            warn!(
                "Rejected {} OTLP metrics data points: {}",
                outcome.rejected_data_points,
                outcome.error_message.as_deref().unwrap_or_default()
            );
        }
        if outcome.accepted_data_points == 0 {
            return Ok(outcome);
        }

        let ctx = admit_write(rows as u64, &ctx)
            .await
            .map_err(BoxedError::new)
            .context(error::ExecuteGrpcQuerySnafu)?;
        OTLP_METRICS_ROWS.inc_by(rows as u64);

        let ctx = {
            let mut c = (*ctx).clone();
            c.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC);
            c.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY);
            // Per-table metric specifics + resource/scope lineage ride this
            // internal channel; the auto-create path folds them per schema and
            // table name.
            if let Some(index) = semantic_index.encode(&c.current_schema()) {
                c.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index);
            }
            c.set_extension(OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM.to_string());
            Arc::new(c)
        };

        let output = self
            .handle_otlp_metric_row_inserts(requests, ctx.clone(), &metric_ctx)
            .await?;
        outcome.write_cost = output.meta.cost;

        // Derived enrichment follows the accepted metric submission, which may
        // still be queued in asynchronous mode. Failures remain warning-only
        // to avoid retrying metric data the server already accepted.
        if let Some(resource_info) = resource_info {
            let written = match self.check_row_insert_permission(
                &resource_info,
                &ctx,
                PermissionReq::Action(OTLP_WRITE),
            ) {
                Ok(_) => self
                    .handle_row_inserts(resource_info, ctx, false, false)
                    .await
                    .map_err(BoxedError::new)
                    .map_err(|e| e.to_string()),
                Err(e) => Err(e.to_string()),
            };
            match written {
                Ok(descriptor_output) => outcome.write_cost += descriptor_output.meta.cost,
                Err(e) => {
                    OTLP_RESOURCE_INFO_WRITE_ERRORS.inc();
                    warn!("Failed to write the OTLP resource descriptor table: {e}");
                    outcome.error_message.get_or_insert(format!(
                        "metric data was accepted, but writing the resource \
                         descriptor table `{}` failed: {e}",
                        otlp::metrics::OTEL_RESOURCE_INFO_TABLE_NAME
                    ));
                }
            }
        }

        Ok(outcome)
    }

    #[tracing::instrument(skip_all)]
    async fn traces(
        &self,
        _pipeline_handler: PipelineHandlerRef,
        request: ExportTraceServiceRequest,
        pipeline: PipelineWay,
        _pipeline_params: GreptimePipelineParams,
        table_name: String,
        ctx: QueryContextRef,
    ) -> ServerResult<TraceIngestOutcome> {
        self.plugins
            .get::<PermissionCheckerRef>()
            .as_ref()
            .check_permission(ctx.current_user(), PermissionReq::Action(OTLP_WRITE))
            .context(AuthSnafu)?;

        let interceptor_ref = self
            .plugins
            .get::<OpenTelemetryProtocolInterceptorRef<servers::error::Error>>();
        interceptor_ref.pre_execute(ctx.clone())?;
        let ctx = Arc::new(ctx.fork());

        // `schema_url` is consumed by `parse`, so derive conventions first.
        let conventions = trace_conventions(&request);
        let spans = otlp::trace::span::parse(request);
        let targets = trace_permission_targets(&table_name, &spans, &ctx);
        self.check_table_permission(&ctx, PermissionReq::Action(OTLP_WRITE), targets)
            .context(AuthSnafu)?;
        // Count the external spans once, before chunking, retries, or derived tables.
        let rows = spans.iter().map(|group| group.spans.len() as u64).sum();
        let ctx = admit_write(rows, &ctx)
            .await
            .map_err(BoxedError::new)
            .context(error::ExecuteGrpcQuerySnafu)?;
        self.ingest_trace_spans(&pipeline, table_name, spans, &conventions, ctx)
            .await
    }

    #[tracing::instrument(skip_all)]
    async fn logs(
        &self,
        pipeline_handler: PipelineHandlerRef,
        request: ExportLogsServiceRequest,
        pipeline: PipelineWay,
        pipeline_params: GreptimePipelineParams,
        table_name: String,
        ctx: QueryContextRef,
    ) -> ServerResult<Vec<Output>> {
        self.plugins
            .get::<PermissionCheckerRef>()
            .as_ref()
            .check_permission(ctx.current_user(), PermissionReq::Action(OTLP_WRITE))
            .context(AuthSnafu)?;

        let interceptor_ref = self
            .plugins
            .get::<OpenTelemetryProtocolInterceptorRef<servers::error::Error>>();
        interceptor_ref.pre_execute(ctx.clone())?;
        let ctx = Arc::new(ctx.fork());

        // `as_req_iter` clones this ctx into each `temp_ctx`, so identity set here
        // reaches the context that drives table auto-create.
        let ctx = {
            let mut c = (*ctx).clone();
            c.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_LOG);
            c.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY);
            Arc::new(c)
        };

        let opt_req = otlp::logs::to_grpc_insert_requests(
            request,
            pipeline,
            pipeline_params,
            table_name,
            &ctx,
            pipeline_handler,
        )
        .await?;

        let mut batches = opt_req.as_req_iter(ctx).collect::<Vec<_>>();
        for (temp_ctx, requests) in &batches {
            self.check_row_insert_permission(requests, temp_ctx, PermissionReq::Action(OTLP_WRITE))
                .context(AuthSnafu)?;
        }
        admit_row_insert_batches(&mut batches)
            .await
            .map_err(BoxedError::new)
            .context(error::ExecuteGrpcQuerySnafu)?;

        let mut outputs = Vec::with_capacity(batches.len());
        for (temp_ctx, requests) in batches {
            let cnt = requests
                .inserts
                .iter()
                .filter_map(|r| r.rows.as_ref().map(|r| r.rows.len()))
                .sum::<usize>();

            let o = self
                .handle_log_inserts(requests, temp_ctx)
                .await
                .inspect(|_| OTLP_LOGS_ROWS.inc_by(cnt as u64))
                .map_err(BoxedError::new)
                .context(error::ExecuteGrpcQuerySnafu)?;
            outputs.push(o);
        }

        Ok(outputs)
    }
}

impl Instance {
    /// Inserts converted OTLP metrics through the eligible batching or ordinary write path.
    /// The caller must check table permissions and admit the complete request first.
    pub async fn handle_otlp_metric_row_inserts(
        &self,
        requests: RowInsertRequests,
        ctx: QueryContextRef,
        metric_ctx: &OtlpMetricCtx,
    ) -> ServerResult<Output> {
        let physical_table = ctx
            .extension(PHYSICAL_TABLE_PARAM)
            .unwrap_or(GREPTIME_PHYSICAL_TABLE)
            .to_string();
        // The bulk path converts each request's time index unit to the
        // destination table's unit during batch alignment, so no pre-gate
        // alignment is needed here.
        let batcher = self
            .logical_batcher()
            .filter(|_| ctx.logical_batching_enabled() && metric_ctx.with_metric_engine);
        let batcher = if batcher.is_some()
            && self
                .inserter
                .can_batch_metric_rows(&requests, &ctx, &physical_table)
                .await
                .map_err(BoxedError::new)
                .context(error::ExecuteGrpcQuerySnafu)?
        {
            batcher
        } else {
            None
        };

        // OTLP tables have one sample field in both the ordinary and physical paths.
        if let Some(batcher) = batcher {
            let (rows, cost) = batcher
                .submit_with(requests, ctx.clone(), |mut requests| {
                    let ctx = ctx.clone();
                    async move {
                        Inserter::meter_row_inserts(&mut requests, &ctx)
                            .await
                            .map_err(BoxedError::new)
                            .context(error::ExecuteGrpcQuerySnafu)
                    }
                })
                .await?;
            Ok(Output::new(
                OutputData::AffectedRows(rows as usize),
                OutputMeta::new_with_cost(cost as _),
            ))
        } else if !metric_ctx.with_metric_engine {
            self.handle_row_inserts(requests, ctx.clone(), false, true)
                .await
                .map_err(BoxedError::new)
                .context(error::ExecuteGrpcQuerySnafu)
        } else {
            self.handle_metric_row_inserts(requests, ctx.clone(), physical_table)
                .await
                .map_err(BoxedError::new)
                .context(error::ExecuteGrpcQuerySnafu)
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use auth::{PermissionTableTarget, PermissionTableTargets};
    use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest;
    use opentelemetry_proto::tonic::common::v1::{AnyValue, KeyValue, any_value};
    use opentelemetry_proto::tonic::resource::v1::Resource;
    use opentelemetry_proto::tonic::trace::v1::{ResourceSpans, ScopeSpans, Span};
    use session::context::QueryContext;

    use super::trace_permission_targets;

    #[test]
    fn test_trace_permission_targets() {
        let request = ExportTraceServiceRequest {
            resource_spans: vec![ResourceSpans {
                resource: Some(Resource {
                    attributes: vec![KeyValue {
                        key: "service.name".to_string(),
                        value: Some(AnyValue {
                            value: Some(any_value::Value::StringValue("frontend".to_string())),
                        }),
                        ..Default::default()
                    }],
                    ..Default::default()
                }),
                scope_spans: vec![ScopeSpans {
                    spans: vec![Span::default()],
                    ..Default::default()
                }],
                ..Default::default()
            }],
        };
        let groups = servers::otlp::trace::span::parse(request);
        let ctx = Arc::new(QueryContext::with("greptime", "public"));

        assert_eq!(
            PermissionTableTargets::Resolved(vec![
                PermissionTableTarget::new("greptime", "public", "traces"),
                PermissionTableTarget::new("greptime", "public", "traces_services"),
                PermissionTableTarget::new("greptime", "public", "traces_operations"),
            ]),
            trace_permission_targets("traces", &groups, &ctx)
        );
        assert_eq!(
            PermissionTableTargets::Unresolved,
            trace_permission_targets("", &groups, &ctx)
        );
    }
}