diff --git a/src/servers/src/error.rs b/src/servers/src/error.rs index 8a3c554058..825e2296c0 100644 --- a/src/servers/src/error.rs +++ b/src/servers/src/error.rs @@ -292,6 +292,14 @@ pub enum Error { error: snap::Error, }, + #[snafu(display("Failed to decompress snappy Loki request"))] + DecompressSnappyLokiRequest { + #[snafu(implicit)] + location: Location, + #[snafu(source)] + error: snap::Error, + }, + #[snafu(display("Failed to decompress zstd prometheus remote request"))] DecompressZstdPromRemoteRequest { #[snafu(implicit)] @@ -756,6 +764,7 @@ impl ErrorExt for Error { | UnsupportedJsonContentType { .. } | CompressPromRemoteRequest { .. } | DecompressSnappyPromRemoteRequest { .. } + | DecompressSnappyLokiRequest { .. } | DecompressZstdPromRemoteRequest { .. } | InvalidPromRemoteRequest { .. } | InvalidFlightTicket { .. } diff --git a/src/servers/src/http/event.rs b/src/servers/src/http/event.rs index dc468a9b75..f08c76bcdb 100644 --- a/src/servers/src/http/event.rs +++ b/src/servers/src/http/event.rs @@ -38,6 +38,7 @@ use mime_guess::mime; use operator::expr_helper::{create_table_expr_by_column_schemas, expr_to_create}; use pipeline::util::to_pipeline_version; use pipeline::{ContextReq, GreptimePipelineParams, PipelineContext, PipelineDefinition}; +use prometheus::{HistogramVec, IntCounterVec}; use serde::{Deserialize, Serialize}; use serde_json::{Deserializer, Map, Value as JsonValue, json}; use session::context::{Channel, QueryContext, QueryContextRef}; @@ -920,9 +921,9 @@ pub(crate) async fn ingest_logs_inner( query_ctx: QueryContextRef, pipeline_params: GreptimePipelineParams, ) -> Result { - let db = query_ctx.get_db_string(); - let exec_timer = std::time::Instant::now(); - + // Keep the timer boundary before pipeline execution to preserve existing + // ingestion elapsed metrics. + let exec_timer = Instant::now(); let mut req = ContextReq::default(); let pipeline_ctx = PipelineContext::new(&pipeline, &pipeline_params, query_ctx.channel()); @@ -933,10 +934,31 @@ pub(crate) async fn ingest_logs_inner( req.merge(requests); } - let mut outputs = Vec::new(); + execute_log_context_req( + handler, + req, + query_ctx, + exec_timer, + &METRIC_HTTP_LOGS_INGESTION_COUNTER, + &METRIC_HTTP_LOGS_INGESTION_ELAPSED, + ) + .await +} + +pub(crate) async fn execute_log_context_req( + handler: PipelineHandlerRef, + ctx_req: ContextReq, + query_ctx: QueryContextRef, + exec_timer: Instant, + counter: &IntCounterVec, + elapsed: &HistogramVec, +) -> Result { + let db = query_ctx.get_db_string(); + + let mut outputs = Vec::with_capacity(ctx_req.map_len()); let mut total_rows: u64 = 0; let mut fail = false; - for (temp_ctx, act_req) in req.as_req_iter(query_ctx) { + for (temp_ctx, act_req) in ctx_req.as_req_iter(query_ctx) { let output = handler.insert(act_req, temp_ctx).await; if let Ok(Output { @@ -951,16 +973,15 @@ pub(crate) async fn ingest_logs_inner( outputs.push(output); } + // Record one aggregate metric sample for the whole ingestion request. if total_rows > 0 { - METRIC_HTTP_LOGS_INGESTION_COUNTER - .with_label_values(&[db.as_str()]) - .inc_by(total_rows); - METRIC_HTTP_LOGS_INGESTION_ELAPSED + counter.with_label_values(&[db.as_str()]).inc_by(total_rows); + elapsed .with_label_values(&[db.as_str(), METRIC_SUCCESS_VALUE]) .observe(exec_timer.elapsed().as_secs_f64()); } if fail { - METRIC_HTTP_LOGS_INGESTION_ELAPSED + elapsed .with_label_values(&[db.as_str(), METRIC_FAILURE_VALUE]) .observe(exec_timer.elapsed().as_secs_f64()); } diff --git a/src/servers/src/http/loki.md b/src/servers/src/http/loki.md new file mode 100644 index 0000000000..3860a4bc44 --- /dev/null +++ b/src/servers/src/http/loki.md @@ -0,0 +1,108 @@ +# Loki Push Ingestion + +This note describes the HTTP Loki push path implemented in `loki.rs`. + +## Entry Point + +The route is registered by `HttpServer::route_loki`: + +```text +POST /v1/loki/api/v1/push +``` + +The handler is `loki_ingest`. It receives the authenticated `QueryContext`, marks +the channel and semantic source as Loki log ingestion, resolves the target table, +and converts the request body into a `ContextReq`. + +The target table is selected by `x-greptime-log-table-name`. If the header is +absent, Loki writes to `loki_logs`. + +## Supported Payloads + +Payload decoding is selected by `Content-Type`. + +| Content-Type | Parser | Notes | +| --- | --- | --- | +| `application/json` | `LokiJsonParser` | Expects a top-level `streams` array. | +| `application/x-protobuf` | `LokiPbParser` | Expects a Snappy-compressed Loki `PushRequest`. | + +The HTTP route also has request decompression middleware. That is separate from +the Snappy decoding required by Loki protobuf push bodies. + +Both parsers produce `LokiMiddleItem` values: + +```text +JSON stream values -> LokiMiddleItem +protobuf entries -> LokiMiddleItem> +``` + +Malformed inner streams or entries are warned and skipped. Invalid top-level +payloads, unsupported content types, protobuf decode errors, and protobuf Snappy +decode errors return request errors. + +## Direct Ingestion + +When no pipeline header is present, Loki uses the direct-write path: + +```text +bytes -> extract_item -> RowInsertRequest -> ContextReq +``` + +The initial schema is: + +| Column | Type | Semantic type | +| --- | --- | --- | +| `greptime_timestamp()` | timestamp nanosecond | timestamp | +| `line` | string | field | +| `structured_metadata` | JSON binary | field | + +Loki labels are added as string tag columns. Since later entries may introduce +new labels, previously built rows are padded with nulls after the final schema is +known. + +Structured metadata is stored as one JSON binary value in the +`structured_metadata` field. + +## Pipeline Ingestion + +When `x-greptime-log-pipeline-name` or `x-greptime-pipeline-name` is present, +Loki uses the pipeline path: + +```text +bytes -> extract_item -> PipelineIngestRequest -> run_pipeline -> ContextReq +``` + +Pipeline version and parameters come from the shared pipeline headers: + +| Header | Meaning | +| --- | --- | +| `x-greptime-log-pipeline-version` / `x-greptime-pipeline-version` | Pipeline version. | +| `x-greptime-pipeline-params` | Pipeline parameters. | + +Each Loki entry is converted into a VRL object before pipeline execution: + +| Source | Pipeline field | +| --- | --- | +| Entry timestamp | `greptime_timestamp()` | +| Entry line | `loki_line` | +| Structured metadata key `k` | `loki_metadata_k` | +| Label key `k` | `loki_label_k` | + +## Execution And Metrics + +Both direct and pipeline paths produce a `ContextReq`. The shared +`execute_log_context_req` helper performs inserts through `PipelineHandler::insert`, +collects all outputs, builds a `GreptimedbV1Response`, and records aggregate +Loki ingestion metrics. + +The elapsed timer starts before payload conversion, so it includes parsing, +optional pipeline execution, and insertion. + +## Important Contracts + +- The public endpoint and response format are unchanged from the Loki push API. +- `LogState.log_validator` is not used by this Loki-specific handler. +- Protobuf push bodies must be Snappy-compressed even when there is no HTTP + `Content-Encoding`. +- Invalid Snappy data must return a controlled Loki decompression error rather + than panic. diff --git a/src/servers/src/http/loki.rs b/src/servers/src/http/loki.rs index 81063e371d..4e12eeb931 100644 --- a/src/servers/src/http/loki.rs +++ b/src/servers/src/http/loki.rs @@ -27,7 +27,6 @@ use axum_extra::TypedHeader; use bytes::Bytes; use chrono::DateTime; use common_query::prelude::greptime_timestamp; -use common_query::{Output, OutputData}; use common_telemetry::{error, warn}; use headers::ContentType; use jsonb::Value; @@ -35,28 +34,29 @@ use lazy_static::lazy_static; use loki_proto::logproto::LabelPairAdapter; use loki_proto::prost_types::Timestamp as LokiTimestamp; use pipeline::util::to_pipeline_version; -use pipeline::{ContextReq, PipelineContext, PipelineDefinition, SchemaInfo}; +use pipeline::{ + ContextReq, GreptimePipelineParams, PipelineContext, PipelineDefinition, SchemaInfo, +}; use prost::Message; use quoted_string::test_utils::TestSpec; -use session::context::{Channel, QueryContext}; +use session::context::{Channel, QueryContext, QueryContextRef}; use snafu::{OptionExt, ResultExt, ensure}; +use snap::raw::Decoder; use table::requests::{SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_LOG, SOURCE_LOKI}; use vrl::value::{KeyString, Value as VrlValue}; use crate::error::{ - DecodeLokiRequestSnafu, InvalidLokiLabelsSnafu, InvalidLokiPayloadSnafu, ParseJsonSnafu, - PipelineSnafu, Result, UnsupportedContentTypeSnafu, + DecodeLokiRequestSnafu, DecompressSnappyLokiRequestSnafu, InvalidLokiLabelsSnafu, + InvalidLokiPayloadSnafu, ParseJsonSnafu, PipelineSnafu, Result, UnsupportedContentTypeSnafu, }; use crate::http::HttpResponse; -use crate::http::event::{JSON_CONTENT_TYPE, LogState, PB_CONTENT_TYPE, PipelineIngestRequest}; -use crate::http::extractor::{LogTableName, PipelineInfo}; -use crate::http::result::greptime_result_v1::GreptimedbV1Response; -use crate::metrics::{ - METRIC_FAILURE_VALUE, METRIC_LOKI_LOGS_INGESTION_COUNTER, METRIC_LOKI_LOGS_INGESTION_ELAPSED, - METRIC_SUCCESS_VALUE, +use crate::http::event::{ + JSON_CONTENT_TYPE, LogState, PB_CONTENT_TYPE, PipelineIngestRequest, execute_log_context_req, }; +use crate::http::extractor::{LogTableName, PipelineInfo}; +use crate::metrics::{METRIC_LOKI_LOGS_INGESTION_COUNTER, METRIC_LOKI_LOGS_INGESTION_ELAPSED}; use crate::pipeline::run_pipeline; -use crate::prom_store; +use crate::query_handler::PipelineHandlerRef; const LOKI_TABLE_NAME: &str = "loki_logs"; const LOKI_LINE_COLUMN: &str = "line"; @@ -115,90 +115,30 @@ pub async fn loki_ingest( ctx.set_extension(SEMANTIC_SOURCE, SOURCE_LOKI); let ctx = Arc::new(ctx); let table_name = table_name.unwrap_or_else(|| LOKI_TABLE_NAME.to_string()); - let db = ctx.get_db_string(); - let db_str = db.as_str(); + let handler = log_state.log_handler; + // Preserve the old elapsed metric boundary: it includes parsing, optional + // pipeline execution, and insertion. let exec_timer = Instant::now(); - let handler = log_state.log_handler; + let ctx_req = build_loki_context_req( + &handler, + content_type, + table_name, + pipeline_info, + bytes, + &ctx, + ) + .await?; - let ctx_req = if let Some(pipeline_name) = pipeline_info.pipeline_name { - // go pipeline - let version = to_pipeline_version(pipeline_info.pipeline_version.as_deref()) - .context(PipelineSnafu)?; - let def = - PipelineDefinition::from_name(&pipeline_name, version, None).context(PipelineSnafu)?; - let pipeline_ctx = - PipelineContext::new(&def, &pipeline_info.pipeline_params, Channel::Loki); - - let v = extract_item::(content_type, bytes)? - .map(|i| i.map) - .collect::>(); - - let req = PipelineIngestRequest { - table: table_name, - values: v, - }; - - run_pipeline(&handler, &pipeline_ctx, req, &ctx, true).await? - } else { - // init schemas - let mut schema_info = SchemaInfo::from_schema_list(LOKI_INIT_SCHEMAS.clone()); - let mut rows = Vec::with_capacity(256); - for loki_row in extract_item::(content_type, bytes)? { - let mut row = init_row( - schema_info.schema.len(), - loki_row.ts, - loki_row.line, - loki_row.structured_metadata, - ); - process_labels(&mut schema_info, &mut row, loki_row.labels); - rows.push(row); - } - - let schemas = schema_info.column_schemas()?; - // fill Null for missing values - for row in rows.iter_mut() { - row.resize(schemas.len(), GreptimeValue::default()); - } - let rows = Rows { - rows: rows.into_iter().map(|values| Row { values }).collect(), - schema: schemas, - }; - let ins_req = RowInsertRequest { - table_name, - rows: Some(rows), - }; - - ContextReq::default_opt_with_reqs(vec![ins_req]) - }; - - let mut outputs = Vec::with_capacity(ctx_req.map_len()); - for (temp_ctx, req) in ctx_req.as_req_iter(ctx) { - let output = handler.insert(req, temp_ctx).await; - - if let Ok(Output { - data: OutputData::AffectedRows(rows), - meta: _, - }) = &output - { - METRIC_LOKI_LOGS_INGESTION_COUNTER - .with_label_values(&[db_str]) - .inc_by(*rows as u64); - METRIC_LOKI_LOGS_INGESTION_ELAPSED - .with_label_values(&[db_str, METRIC_SUCCESS_VALUE]) - .observe(exec_timer.elapsed().as_secs_f64()); - } else { - METRIC_LOKI_LOGS_INGESTION_ELAPSED - .with_label_values(&[db_str, METRIC_FAILURE_VALUE]) - .observe(exec_timer.elapsed().as_secs_f64()); - } - outputs.push(output); - } - - let response = GreptimedbV1Response::from_output(outputs) - .await - .with_execution_time(exec_timer.elapsed().as_millis() as u64); - Ok(response) + execute_log_context_req( + handler, + ctx_req, + ctx, + exec_timer, + &METRIC_LOKI_LOGS_INGESTION_COUNTER, + &METRIC_LOKI_LOGS_INGESTION_ELAPSED, + ) + .await } /// This is the holder of the loki lines parsed from json or protobuf. @@ -226,38 +166,115 @@ pub struct LokiPipeline { pub map: VrlValue, } -/// This is the flow of the Loki ingestion. -/// +--------+ -/// | bytes | -/// +--------+ -/// | -/// +----------------------+----------------------+ -/// | | | -/// | JSON content type | PB content type | -/// +----------------------+----------------------+ -/// | | | -/// | JsonStreamItem | PbStreamItem | -/// | stream: serde_json | stream: adapter | -/// +----------------------+----------------------+ -/// | | | -/// | MiddleItem | MiddleItem | -/// +----------------------+----------------------+ -/// \ / -/// \ / -/// \ / -/// +----------------------+ -/// | MiddleItem | -/// +----------------------+ -/// | -/// +----------------+----------------+ -/// | | -/// +------------------+ +---------------------+ -/// | LokiRawItem | | LokiPipelineItem | -/// +------------------+ +---------------------+ -/// | | -/// +------------------+ +---------------------+ -/// | Loki ingest | | run_pipeline | -/// +------------------+ +---------------------+ +struct LokiPipelineContextReq { + content_type: ContentType, + table_name: String, + pipeline_name: String, + pipeline_version: Option, + pipeline_params: GreptimePipelineParams, + bytes: Bytes, +} + +async fn build_loki_context_req( + handler: &PipelineHandlerRef, + content_type: ContentType, + table_name: String, + pipeline_info: PipelineInfo, + bytes: Bytes, + ctx: &QueryContextRef, +) -> Result { + // A pipeline header switches Loki into the generic pipeline path; without + // it, Loki writes directly to the target log table. + match pipeline_info.pipeline_name { + Some(pipeline_name) => { + let pipeline_req = LokiPipelineContextReq { + content_type, + table_name, + pipeline_name, + pipeline_version: pipeline_info.pipeline_version, + pipeline_params: pipeline_info.pipeline_params, + bytes, + }; + build_loki_pipeline_context_req(handler, pipeline_req, ctx).await + } + None => { + let req = build_loki_raw_insert_request(content_type, table_name, bytes)?; + Ok(ContextReq::default_opt_with_reqs(vec![req])) + } + } +} + +async fn build_loki_pipeline_context_req( + handler: &PipelineHandlerRef, + pipeline_req: LokiPipelineContextReq, + ctx: &QueryContextRef, +) -> Result { + let LokiPipelineContextReq { + content_type, + table_name, + pipeline_name, + pipeline_version, + pipeline_params, + bytes, + } = pipeline_req; + + let version = to_pipeline_version(pipeline_version.as_deref()).context(PipelineSnafu)?; + let def = + PipelineDefinition::from_name(&pipeline_name, version, None).context(PipelineSnafu)?; + let pipeline_ctx = PipelineContext::new(&def, &pipeline_params, Channel::Loki); + + let values = extract_item::(content_type, bytes)? + .map(|item| item.map) + .collect::>(); + + let req = PipelineIngestRequest { + table: table_name, + values, + }; + + run_pipeline(handler, &pipeline_ctx, req, ctx, true).await +} + +fn build_loki_raw_insert_request( + content_type: ContentType, + table_name: String, + bytes: Bytes, +) -> Result { + let mut schema_info = SchemaInfo::from_schema_list(LOKI_INIT_SCHEMAS.clone()); + let mut rows = Vec::with_capacity(256); + for loki_row in extract_item::(content_type, bytes)? { + let mut row = init_row( + schema_info.schema.len(), + loki_row.ts, + loki_row.line, + loki_row.structured_metadata, + ); + process_labels(&mut schema_info, &mut row, loki_row.labels); + rows.push(row); + } + + let schemas = schema_info.column_schemas()?; + // Labels can introduce new tag columns after earlier rows were built. + for row in rows.iter_mut() { + row.resize(schemas.len(), GreptimeValue::default()); + } + let rows = Rows { + rows: rows.into_iter().map(|values| Row { values }).collect(), + schema: schemas, + }; + + Ok(RowInsertRequest { + table_name, + rows: Some(rows), + }) +} + +/// Extract Loki entries from the supported wire format into the caller's +/// destination type. +/// +/// JSON push bodies become `LokiMiddleItem`, protobuf push bodies +/// become `LokiMiddleItem>`, and the generic `Into` +/// conversion selects either direct-write `LokiRawItem` or pipeline `LokiPipeline`. fn extract_item(content_type: ContentType, bytes: Bytes) -> Result>> where LokiMiddleItem: Into, @@ -408,6 +425,93 @@ impl Iterator for JsonStreamItem { } } +type LokiPipelineMap = BTreeMap; + +fn vrl_metadata_to_jsonb(structured_metadata: Option) -> Vec { + // JSON push structured metadata arrives as a VRL object. + let structured_metadata = structured_metadata + .and_then(|metadata| match metadata { + VrlValue::Object(metadata) => Some(metadata), + _ => None, + }) + .map(|metadata| { + metadata + .into_iter() + .filter_map(|(key, value)| match value { + VrlValue::Bytes(bytes) => Some(( + key.into(), + Value::String(String::from_utf8_lossy(&bytes).to_string().into()), + )), + _ => None, + }) + .collect::>() + }) + .unwrap_or_default(); + + Value::Object(structured_metadata).to_vec() +} + +fn label_pair_metadata_to_jsonb(structured_metadata: Option>) -> Vec { + // Protobuf push structured metadata arrives as Loki label pairs. + let structured_metadata = structured_metadata + .unwrap_or_default() + .into_iter() + .map(|metadata| (metadata.name, Value::String(metadata.value.into()))) + .collect::>(); + + Value::Object(structured_metadata).to_vec() +} + +fn new_loki_pipeline_map(ts: i64, line: String) -> LokiPipelineMap { + let mut map = BTreeMap::new(); + map.insert( + KeyString::from(greptime_timestamp()), + VrlValue::Timestamp(DateTime::from_timestamp_nanos(ts)), + ); + map.insert( + KeyString::from(LOKI_LINE_COLUMN_NAME), + VrlValue::Bytes(line.into()), + ); + map +} + +fn append_vrl_pipeline_metadata(map: &mut LokiPipelineMap, structured_metadata: Option) { + if let Some(VrlValue::Object(metadata)) = structured_metadata { + for (key, value) in metadata { + map.insert( + KeyString::from(format!("{}{}", LOKI_PIPELINE_METADATA_PREFIX, key)), + value, + ); + } + } +} + +fn append_label_pair_pipeline_metadata( + map: &mut LokiPipelineMap, + structured_metadata: Option>, +) { + for metadata in structured_metadata.unwrap_or_default() { + map.insert( + KeyString::from(format!( + "{}{}", + LOKI_PIPELINE_METADATA_PREFIX, metadata.name + )), + VrlValue::Bytes(metadata.value.into()), + ); + } +} + +fn append_pipeline_labels(map: &mut LokiPipelineMap, labels: Option>) { + if let Some(labels) = labels { + for (key, value) in labels { + map.insert( + KeyString::from(format!("{}{}", LOKI_PIPELINE_LABEL_PREFIX, key)), + VrlValue::Bytes(value.into()), + ); + } + } +} + impl From> for LokiRawItem { fn from(val: LokiMiddleItem) -> Self { let LokiMiddleItem { @@ -417,29 +521,10 @@ impl From> for LokiRawItem { labels, } = val; - let structured_metadata = structured_metadata - .and_then(|m| match m { - VrlValue::Object(m) => Some(m), - _ => None, - }) - .map(|m| { - m.into_iter() - .filter_map(|(k, v)| match v { - VrlValue::Bytes(bytes) => Some(( - k.into(), - Value::String(String::from_utf8_lossy(&bytes).to_string().into()), - )), - _ => None, - }) - .collect::>() - }) - .unwrap_or_default(); - let structured_metadata = Value::Object(structured_metadata).to_vec(); - LokiRawItem { ts, line, - structured_metadata, + structured_metadata: vrl_metadata_to_jsonb(structured_metadata), labels, } } @@ -454,32 +539,9 @@ impl From> for LokiPipeline { labels, } = value; - let mut map = BTreeMap::new(); - map.insert( - KeyString::from(greptime_timestamp()), - VrlValue::Timestamp(DateTime::from_timestamp_nanos(ts)), - ); - map.insert( - KeyString::from(LOKI_LINE_COLUMN_NAME), - VrlValue::Bytes(line.into()), - ); - - if let Some(VrlValue::Object(m)) = structured_metadata { - for (k, v) in m { - map.insert( - KeyString::from(format!("{}{}", LOKI_PIPELINE_METADATA_PREFIX, k)), - v, - ); - } - } - if let Some(v) = labels { - v.into_iter().for_each(|(k, v)| { - map.insert( - KeyString::from(format!("{}{}", LOKI_PIPELINE_LABEL_PREFIX, k)), - VrlValue::Bytes(v.into()), - ); - }); - } + let mut map = new_loki_pipeline_map(ts, line); + append_vrl_pipeline_metadata(&mut map, structured_metadata); + append_pipeline_labels(&mut map, labels); LokiPipeline { map: VrlValue::Object(map), @@ -493,7 +555,7 @@ pub struct LokiPbParser { impl LokiPbParser { pub fn from_bytes(bytes: Bytes) -> Result { - let decompressed = prom_store::snappy_decompress(&bytes).unwrap(); + let decompressed = snappy_decompress_loki_request(&bytes)?; let req = loki_proto::logproto::PushRequest::decode(&decompressed[..]) .context(DecodeLokiRequestSnafu)?; @@ -503,6 +565,15 @@ impl LokiPbParser { } } +fn snappy_decompress_loki_request(buf: &[u8]) -> Result> { + // Loki's protobuf push body is Snappy-compressed independent of HTTP + // content-encoding, so keep this decode step explicit. + let mut decoder = Decoder::new(); + decoder + .decompress_vec(buf) + .context(DecompressSnappyLokiRequestSnafu) +} + impl Iterator for LokiPbParser { type Item = PbStreamItem; @@ -562,17 +633,10 @@ impl From>> for LokiRawItem { labels, } = val; - let structured_metadata = structured_metadata - .unwrap_or_default() - .into_iter() - .map(|d| (d.name, Value::String(d.value.into()))) - .collect::>(); - let structured_metadata = Value::Object(structured_metadata).to_vec(); - LokiRawItem { ts, line, - structured_metadata, + structured_metadata: label_pair_metadata_to_jsonb(structured_metadata), labels, } } @@ -587,34 +651,9 @@ impl From>> for LokiPipeline { labels, } = value; - let mut map = BTreeMap::new(); - map.insert( - KeyString::from(greptime_timestamp()), - VrlValue::Timestamp(DateTime::from_timestamp_nanos(ts)), - ); - map.insert( - KeyString::from(LOKI_LINE_COLUMN_NAME), - VrlValue::Bytes(line.into()), - ); - - structured_metadata - .unwrap_or_default() - .into_iter() - .for_each(|d| { - map.insert( - KeyString::from(format!("{}{}", LOKI_PIPELINE_METADATA_PREFIX, d.name)), - VrlValue::Bytes(d.value.into()), - ); - }); - - if let Some(v) = labels { - v.into_iter().for_each(|(k, v)| { - map.insert( - KeyString::from(format!("{}{}", LOKI_PIPELINE_LABEL_PREFIX, k)), - VrlValue::Bytes(v.into()), - ); - }); - } + let mut map = new_loki_pipeline_map(ts, line); + append_label_pair_pipeline_metadata(&mut map, structured_metadata); + append_pipeline_labels(&mut map, labels); LokiPipeline { map: VrlValue::Object(map), @@ -772,10 +811,51 @@ fn process_labels( mod tests { use std::collections::BTreeMap; + use bytes::Bytes; + use loki_proto::logproto::{EntryAdapter, PushRequest, StreamAdapter}; use loki_proto::prost_types::Timestamp; + use prost::Message; - use crate::error::Error::InvalidLokiLabels; - use crate::http::loki::{parse_loki_labels, prost_ts_to_nano}; + use super::*; + use crate::error::Error::{DecompressSnappyLokiRequest, InvalidLokiLabels}; + use crate::prom_store::snappy_compress; + + const JSON_PAYLOAD: &[u8] = br#"{ + "streams": [ + { + "stream": { + "job": "api", + "namespace": "prod" + }, + "values": [ + ["1731748568804293888", "line one", {"trace_id": "abc"}] + ] + }, + { + "stream": { + "job": "worker", + "pod": "worker-0" + }, + "values": [ + ["1731748568804293889", "line two"] + ] + } + ] + }"#; + + fn row_string_value(row: &Row, index: usize) -> Option<&str> { + match row.values[index].value_data.as_ref() { + Some(ValueData::StringValue(value)) => Some(value.as_str()), + _ => None, + } + } + + fn pipeline_bytes_value(map: &BTreeMap, key: &str) -> Option { + match map.get(&KeyString::from(key))? { + VrlValue::Bytes(value) => Some(String::from_utf8_lossy(value.as_ref()).to_string()), + _ => None, + } + } #[test] fn test_ts_to_nano() { @@ -789,6 +869,134 @@ mod tests { assert_eq!(prost_ts_to_nano(&ts), 1731748568804293888); } + #[test] + fn test_json_direct_ingest_builds_schema_and_pads_rows() { + let request = build_loki_raw_insert_request( + JSON_CONTENT_TYPE.clone(), + "custom_loki".to_string(), + Bytes::from_static(JSON_PAYLOAD), + ) + .unwrap(); + + assert_eq!(request.table_name, "custom_loki"); + let rows = request.rows.unwrap(); + let column_names = rows + .schema + .iter() + .map(|schema| schema.column_name.as_str()) + .collect::>(); + assert_eq!( + column_names, + vec![ + greptime_timestamp(), + LOKI_LINE_COLUMN, + LOKI_STRUCTURED_METADATA_COLUMN, + "job", + "namespace", + "pod", + ] + ); + assert_eq!(rows.schema[3].semantic_type, SemanticType::Tag as i32); + assert_eq!(rows.schema[4].semantic_type, SemanticType::Tag as i32); + assert_eq!(rows.schema[5].semantic_type, SemanticType::Tag as i32); + assert_eq!(rows.rows.len(), 2); + + let first = &rows.rows[0]; + assert_eq!(first.values.len(), rows.schema.len()); + assert_eq!(row_string_value(first, 1), Some("line one")); + assert_eq!(row_string_value(first, 3), Some("api")); + assert_eq!(row_string_value(first, 4), Some("prod")); + assert!(first.values[5].value_data.is_none()); + + let second = &rows.rows[1]; + assert_eq!(second.values.len(), rows.schema.len()); + assert_eq!(row_string_value(second, 1), Some("line two")); + assert_eq!(row_string_value(second, 3), Some("worker")); + assert!(second.values[4].value_data.is_none()); + assert_eq!(row_string_value(second, 5), Some("worker-0")); + } + + #[test] + fn test_json_pipeline_conversion_names_loki_fields() { + let items = extract_item::( + JSON_CONTENT_TYPE.clone(), + Bytes::from_static(JSON_PAYLOAD), + ) + .unwrap() + .collect::>(); + + assert_eq!(items.len(), 2); + let VrlValue::Object(map) = &items[0].map else { + panic!("expected pipeline object"); + }; + assert!(matches!( + map.get(&KeyString::from(greptime_timestamp())), + Some(VrlValue::Timestamp(_)) + )); + assert_eq!( + pipeline_bytes_value(map, LOKI_LINE_COLUMN_NAME), + Some("line one".to_string()) + ); + assert_eq!( + pipeline_bytes_value(map, "loki_label_job"), + Some("api".to_string()) + ); + assert_eq!( + pipeline_bytes_value(map, "loki_label_namespace"), + Some("prod".to_string()) + ); + assert_eq!( + pipeline_bytes_value(map, "loki_metadata_trace_id"), + Some("abc".to_string()) + ); + } + + #[test] + fn test_protobuf_parser_decodes_snappy_push_request() { + let request = PushRequest { + streams: vec![StreamAdapter { + labels: r#"{job="api"}"#.to_string(), + entries: vec![EntryAdapter { + timestamp: Some(Timestamp { + seconds: 1731748568, + nanos: 804293888, + }), + line: "line one".to_string(), + structured_metadata: vec![LabelPairAdapter { + name: "trace_id".to_string(), + value: "abc".to_string(), + }], + ..Default::default() + }], + ..Default::default() + }], + }; + let bytes = snappy_compress(&request.encode_to_vec()).unwrap(); + + let items = extract_item::(PB_CONTENT_TYPE.clone(), Bytes::from(bytes)) + .unwrap() + .collect::>(); + + assert_eq!(items.len(), 1); + assert_eq!(items[0].ts, 1731748568804293888); + assert_eq!(items[0].line, "line one"); + assert_eq!( + items[0].labels.as_ref().unwrap().get("job"), + Some(&"api".to_string()) + ); + assert!(!items[0].structured_metadata.is_empty()); + } + + #[test] + fn test_protobuf_parser_rejects_invalid_snappy_payload() { + let err = match LokiPbParser::from_bytes(Bytes::from_static(b"not-snappy")) { + Ok(_) => panic!("expected invalid snappy payload to fail"), + Err(err) => err, + }; + + assert!(matches!(err, DecompressSnappyLokiRequest { .. })); + } + #[test] fn test_parse_loki_labels() { let mut expected = BTreeMap::new();