servers/

pipeline.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.

use std::collections::BTreeMap;
use std::sync::Arc;

use ahash::{HashMap, HashMapExt};
use api::v1::helper::time_index_column_schema;
use api::v1::{ColumnDataType, Row, RowInsertRequest, Rows, Value};
use common_time::timestamp::TimeUnit;
use pipeline::{
    ContextOpt, ContextReq, DispatchedTo, GREPTIME_INTERNAL_IDENTITY_PIPELINE_NAME, Pipeline,
    PipelineContext, PipelineDefinition, PipelineProcessOutput, SchemaInfo, TransformedOutput,
    TransformerMode, identity_pipeline, unwrap_or_continue_if_err,
};
use session::context::{Channel, QueryContextRef};
use snafu::ResultExt;
use vrl::value::Value as VrlValue;

use crate::error::{PipelineSnafu, Result};
use crate::http::event::PipelineIngestRequest;
use crate::metrics::{
    METRIC_FAILURE_VALUE, METRIC_HTTP_LOGS_TRANSFORM_ELAPSED, METRIC_SUCCESS_VALUE,
};
use crate::query_handler::PipelineHandlerRef;

macro_rules! push_to_map {
    ($map:expr, $key:expr, $value:expr, $capacity:expr) => {
        $map.entry($key)
            .or_insert_with(|| Vec::with_capacity($capacity))
            .push($value);
    };
}

/// Never call this on `GreptimeIdentityPipeline` because it's a real pipeline
pub async fn get_pipeline(
    pipeline_def: &PipelineDefinition,
    handler: &PipelineHandlerRef,
    query_ctx: &QueryContextRef,
) -> Result<Arc<Pipeline>> {
    match pipeline_def {
        PipelineDefinition::Resolved(pipeline) => Ok(pipeline.clone()),
        PipelineDefinition::ByNameAndValue((name, version)) => {
            handler
                .get_pipeline(name, *version, query_ctx.clone())
                .await
        }
        _ => {
            unreachable!("Never call get_pipeline on identity.")
        }
    }
}

pub(crate) async fn run_pipeline(
    handler: &PipelineHandlerRef,
    pipeline_ctx: &PipelineContext<'_>,
    pipeline_req: PipelineIngestRequest,
    query_ctx: &QueryContextRef,
    is_top_level: bool,
) -> Result<ContextReq> {
    if pipeline_ctx.pipeline_definition.is_identity() {
        run_identity_pipeline(handler, pipeline_ctx, pipeline_req, query_ctx).await
    } else {
        run_custom_pipeline(handler, pipeline_ctx, pipeline_req, query_ctx, is_top_level).await
    }
}

async fn run_identity_pipeline(
    handler: &PipelineHandlerRef,
    pipeline_ctx: &PipelineContext<'_>,
    pipeline_req: PipelineIngestRequest,
    query_ctx: &QueryContextRef,
) -> Result<ContextReq> {
    let PipelineIngestRequest {
        table: table_name,
        values: data_array,
    } = pipeline_req;
    let table = if pipeline_ctx.channel == Channel::Prometheus {
        None
    } else {
        handler.get_table(&table_name, query_ctx).await?
    };
    identity_pipeline(data_array, table, pipeline_ctx)
        .map(|opt_map| ContextReq::from_opt_map(opt_map, table_name))
        .context(PipelineSnafu)
}

async fn run_custom_pipeline(
    handler: &PipelineHandlerRef,
    pipeline_ctx: &PipelineContext<'_>,
    pipeline_req: PipelineIngestRequest,
    query_ctx: &QueryContextRef,
    is_top_level: bool,
) -> Result<ContextReq> {
    let skip_error = pipeline_ctx.pipeline_param.skip_error();
    let db = query_ctx.get_db_string();
    let pipeline = get_pipeline(pipeline_ctx.pipeline_definition, handler, query_ctx).await?;

    let transform_timer = std::time::Instant::now();

    let PipelineIngestRequest {
        table: table_name,
        values: pipeline_maps,
    } = pipeline_req;
    let arr_len = pipeline_maps.len();
    let mut transformed_map: HashMap<ContextOpt, HashMap<String, Vec<Row>>> = HashMap::new();
    let mut dispatched: BTreeMap<DispatchedTo, Vec<VrlValue>> = BTreeMap::new();

    let mut schema_info = match pipeline.transformer() {
        TransformerMode::GreptimeTransformer(greptime_transformer) => {
            SchemaInfo::from_schema_list(greptime_transformer.schemas().clone())
        }
        TransformerMode::AutoTransform(ts_name, timeunit) => {
            let timeunit = match timeunit {
                TimeUnit::Second => ColumnDataType::TimestampSecond,
                TimeUnit::Millisecond => ColumnDataType::TimestampMillisecond,
                TimeUnit::Microsecond => ColumnDataType::TimestampMicrosecond,
                TimeUnit::Nanosecond => ColumnDataType::TimestampNanosecond,
            };

            let mut schema_info = SchemaInfo::default();
            schema_info
                .schema
                .push(time_index_column_schema(ts_name, timeunit).into());

            schema_info
        }
    };

    let table = handler.get_table(&table_name, query_ctx).await?;
    schema_info.set_table(table);
    let needs_json_settings = matches!(
        pipeline.transformer(),
        TransformerMode::GreptimeTransformer(transformer) if transformer.has_json_transform()
    );

    for pipeline_map in pipeline_maps {
        let result = pipeline
            .process_mut(pipeline_map)
            .inspect_err(|_| {
                METRIC_HTTP_LOGS_TRANSFORM_ELAPSED
                    .with_label_values(&[db.as_str(), METRIC_FAILURE_VALUE])
                    .observe(transform_timer.elapsed().as_secs_f64());
            })
            .context(PipelineSnafu);

        match unwrap_or_continue_if_err!(result, skip_error) {
            PipelineProcessOutput::Processed(value) => {
                if needs_json_settings {
                    // JSON2 coercion must use settings from the final routed table.
                    let values = match &value {
                        VrlValue::Array(values) => values.as_slice(),
                        value => std::slice::from_ref(value),
                    };
                    for value in values.iter().filter(|value| value.is_object()) {
                        let table_suffix = pipeline.resolve_table_suffix(value).unwrap_or_default();
                        if !schema_info.has_table_for_suffix(&table_suffix) {
                            let destination =
                                table_suffix_to_table_name(&table_name, &table_suffix);
                            let table = handler.get_table(&destination, query_ctx).await?;
                            schema_info.set_table_for_suffix(table_suffix, table);
                        }
                    }
                }

                let result = pipeline
                    .transform_mut(value, pipeline_ctx, &mut schema_info)
                    .inspect_err(|_| {
                        METRIC_HTTP_LOGS_TRANSFORM_ELAPSED
                            .with_label_values(&[db.as_str(), METRIC_FAILURE_VALUE])
                            .observe(transform_timer.elapsed().as_secs_f64());
                    })
                    .context(PipelineSnafu);
                let TransformedOutput { rows_by_context } =
                    unwrap_or_continue_if_err!(result, skip_error);

                // Process each ContextOpt group separately
                for (opt, rows_with_suffix) in rows_by_context {
                    let rows_by_suffix = transformed_map.entry(opt).or_default();
                    // Group rows by table suffix within each context.
                    for (row, table_suffix) in rows_with_suffix {
                        rows_by_suffix
                            .entry(table_suffix.unwrap_or_default())
                            .or_insert_with(|| Vec::with_capacity(arr_len))
                            .push(row);
                    }
                }
            }
            PipelineProcessOutput::DispatchedTo(dispatched_to, val) => {
                push_to_map!(dispatched, dispatched_to, val, arr_len);
            }
            PipelineProcessOutput::Filtered => {
                continue;
            }
        }
    }

    let mut results = ContextReq::default();

    // Process transformed outputs. Rows are grouped by context first, then table suffix.
    let column_count = schema_info.schema.len();
    let column_schemas = schema_info.column_schemas()?;
    for (opt, rows_by_suffix) in transformed_map {
        let row_requests = rows_by_suffix.into_iter().map(|(table_suffix, mut rows)| {
            let table_name = table_suffix_to_table_name(&table_name, &table_suffix);

            // Pad rows to match final schema size (schema may have evolved during processing)
            for row in &mut rows {
                row.values.resize(column_count, Value { value_data: None });
            }

            RowInsertRequest {
                rows: Some(Rows {
                    rows,
                    schema: column_schemas.clone(),
                }),
                table_name,
            }
        });

        results.add_rows(opt, row_requests);
    }

    // if current pipeline contains dispatcher and has several rules, we may
    // already accumulated several dispatched rules and rows.
    for (dispatched_to, coll) in dispatched {
        // we generate the new table name according to `table_part` and
        // current custom table name.
        let table_name = dispatched_to.dispatched_to_table_name(&table_name);
        let next_pipeline_name = dispatched_to
            .pipeline
            .as_deref()
            .unwrap_or(GREPTIME_INTERNAL_IDENTITY_PIPELINE_NAME);

        // run pipeline recursively.
        let next_pipeline_def =
            PipelineDefinition::from_name(next_pipeline_name, None, None).context(PipelineSnafu)?;
        let next_pipeline_ctx = PipelineContext::new(
            &next_pipeline_def,
            pipeline_ctx.pipeline_param,
            pipeline_ctx.channel,
        );
        let requests = Box::pin(run_pipeline(
            handler,
            &next_pipeline_ctx,
            PipelineIngestRequest {
                table: table_name,
                values: coll,
            },
            query_ctx,
            false,
        ))
        .await?;

        results.merge(requests);
    }

    if is_top_level {
        METRIC_HTTP_LOGS_TRANSFORM_ELAPSED
            .with_label_values(&[db.as_str(), METRIC_SUCCESS_VALUE])
            .observe(transform_timer.elapsed().as_secs_f64());
    }

    Ok(results)
}

#[inline]
fn table_suffix_to_table_name(table_name: &str, table_suffix: &str) -> String {
    if table_suffix.is_empty() {
        table_name.to_string()
    } else {
        format!("{}{}", table_name, table_suffix)
    }
}