From 1eb4b8ed4f8dc5c000aeb1c4b7be834faa07d6b7 Mon Sep 17 00:00:00 2001 From: zyy17 Date: Thu, 23 Jan 2025 16:50:50 +0800 Subject: [PATCH] refactor: support to flatten json object in `greptime_identity` pipeline (#5358) * refactor: support to flatten json object in greptime_identity pipeline * refactor: add GreptimeIdentityPipelineParams to configure greptime_identity pipeline * refactor: pass greptime identity pipeline params by one header kv * refactor: code review * refactor: make pipeline params more general for all internal pipelines * chore: remove axum deps from pipeline * fix: clippy errors * chore: fix and add test * test: adopt api change for test client --------- Co-authored-by: shuiyisong Co-authored-by: Ning Sun Co-authored-by: Ning Sun --- src/pipeline/src/etl/error.rs | 8 + .../src/etl/transform/transformer/greptime.rs | 209 +++++++++++++++++- src/pipeline/src/lib.rs | 4 +- src/servers/src/elasticsearch.rs | 8 +- src/servers/src/http/event.rs | 17 +- tests-integration/tests/http.rs | 53 ++++- 6 files changed, 281 insertions(+), 18 deletions(-) diff --git a/src/pipeline/src/etl/error.rs b/src/pipeline/src/etl/error.rs index e19aaad8396..999345fb1e2 100644 --- a/src/pipeline/src/etl/error.rs +++ b/src/pipeline/src/etl/error.rs @@ -594,6 +594,14 @@ pub enum Error { TablePartRequiredForDispatcherRule, #[snafu(display("value is required for dispatcher rule"))] ValueRequiredForDispatcherRule, + #[snafu(display( + "Reached max nested levels when flattening JSON object: {max_nested_levels}" + ))] + ReachedMaxNestedLevels { + max_nested_levels: usize, + #[snafu(implicit)] + location: Location, + }, } pub type Result = std::result::Result; diff --git a/src/pipeline/src/etl/transform/transformer/greptime.rs b/src/pipeline/src/etl/transform/transformer/greptime.rs index 0ae15d0d50b..dedb07e842d 100644 --- a/src/pipeline/src/etl/transform/transformer/greptime.rs +++ b/src/pipeline/src/etl/transform/transformer/greptime.rs @@ -25,11 +25,12 @@ use api::v1::{ColumnDataType, ColumnDataTypeExtension, JsonTypeExtension, Semant use coerce::{coerce_columns, coerce_value}; use greptime_proto::v1::{ColumnSchema, Row, Rows, Value as GreptimeValue}; use itertools::Itertools; -use serde_json::{Map, Number}; +use serde_json::{Map, Number, Value as JsonValue}; use crate::etl::error::{ - IdentifyPipelineColumnTypeMismatchSnafu, Result, TransformColumnNameMustBeUniqueSnafu, - TransformEmptySnafu, TransformMultipleTimestampIndexSnafu, TransformTimestampIndexCountSnafu, + IdentifyPipelineColumnTypeMismatchSnafu, ReachedMaxNestedLevelsSnafu, Result, + TransformColumnNameMustBeUniqueSnafu, TransformEmptySnafu, + TransformMultipleTimestampIndexSnafu, TransformTimestampIndexCountSnafu, UnsupportedNumberTypeSnafu, }; use crate::etl::field::{InputFieldInfo, OneInputOneOutputField}; @@ -37,7 +38,11 @@ use crate::etl::transform::index::Index; use crate::etl::transform::{Transform, Transformer, Transforms}; use crate::etl::value::{Timestamp, Value}; +/// The header key that contains the pipeline params. +pub const GREPTIME_PIPELINE_PARAMS_HEADER: &str = "x-greptime-pipeline-params"; + const DEFAULT_GREPTIME_TIMESTAMP_COLUMN: &str = "greptime_timestamp"; +const DEFAULT_MAX_NESTED_LEVELS_FOR_JSON_FLATTENING: usize = 10; /// fields not in the columns will be discarded /// to prevent automatic column creation in GreptimeDB @@ -47,6 +52,37 @@ pub struct GreptimeTransformer { schema: Vec, } +/// Parameters that can be used to configure the greptime pipelines. +#[derive(Debug, Clone, Default)] +pub struct GreptimePipelineParams { + /// The options for configuring the greptime pipelines. + pub options: HashMap, +} + +impl GreptimePipelineParams { + /// Create a `GreptimePipelineParams` from params string which is from the http header with key `x-greptime-pipeline-params` + /// The params is in the format of `key1=value1&key2=value2`,for example: + /// x-greptime-pipeline-params: flatten_json_object=true + pub fn from_params(params: Option<&str>) -> Self { + let options = params + .unwrap_or_default() + .split('&') + .filter_map(|s| s.split_once('=')) + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect::>(); + + Self { options } + } + + /// Whether to flatten the JSON object. + pub fn flatten_json_object(&self) -> bool { + self.options + .get("flatten_json_object") + .map(|v| v == "true") + .unwrap_or(false) + } +} + impl GreptimeTransformer { /// Add a default timestamp column to the transforms fn add_greptime_timestamp_column(transforms: &mut Transforms) { @@ -365,12 +401,18 @@ fn json_value_to_row( fn identity_pipeline_inner<'a>( array: Vec, tag_column_names: Option>, + params: &GreptimePipelineParams, ) -> Result { let mut rows = Vec::with_capacity(array.len()); let mut schema_info = SchemaInfo::default(); for value in array { if let serde_json::Value::Object(map) = value { - let row = json_value_to_row(&mut schema_info, map)?; + let object = if params.flatten_json_object() { + flatten_json_object(map, DEFAULT_MAX_NESTED_LEVELS_FOR_JSON_FLATTENING)? + } else { + map + }; + let row = json_value_to_row(&mut schema_info, object)?; rows.push(row); } } @@ -420,22 +462,78 @@ fn identity_pipeline_inner<'a>( pub fn identity_pipeline( array: Vec, table: Option>, + params: &GreptimePipelineParams, ) -> Result { match table { Some(table) => { let table_info = table.table_info(); let tag_column_names = table_info.meta.row_key_column_names(); - identity_pipeline_inner(array, Some(tag_column_names)) + identity_pipeline_inner(array, Some(tag_column_names), params) } - None => identity_pipeline_inner(array, None::>), + None => identity_pipeline_inner(array, None::>, params), } } +/// Consumes the JSON object and consumes it into a single-level object. +/// +/// The `max_nested_levels` parameter is used to limit the nested levels of the JSON object. +/// The error will be returned if the nested levels is greater than the `max_nested_levels`. +pub fn flatten_json_object( + object: Map, + max_nested_levels: usize, +) -> Result> { + let mut flattened = Map::new(); + + if !object.is_empty() { + // it will use recursion to flatten the object. + do_flatten_json_object(&mut flattened, None, object, 1, max_nested_levels)?; + } + + Ok(flattened) +} + +fn do_flatten_json_object( + dest: &mut Map, + base: Option<&str>, + object: Map, + current_level: usize, + max_nested_levels: usize, +) -> Result<()> { + // For safety, we do not allow the depth to be greater than the max_object_depth. + if current_level > max_nested_levels { + return ReachedMaxNestedLevelsSnafu { max_nested_levels }.fail(); + } + + for (key, value) in object { + let new_key = base.map_or_else(|| key.clone(), |base_key| format!("{base_key}.{key}")); + + match value { + JsonValue::Object(object) => { + do_flatten_json_object( + dest, + Some(&new_key), + object, + current_level + 1, + max_nested_levels, + )?; + } + // For other types, we will directly insert them into as JSON type. + _ => { + dest.insert(new_key, value); + } + } + } + + Ok(()) +} + #[cfg(test)] mod tests { use api::v1::SemanticType; - use crate::etl::transform::transformer::greptime::identity_pipeline_inner; + use crate::etl::transform::transformer::greptime::{ + flatten_json_object, identity_pipeline_inner, GreptimePipelineParams, + }; use crate::identity_pipeline; #[test] @@ -461,7 +559,7 @@ mod tests { "gaga": "gaga" }), ]; - let rows = identity_pipeline(array, None); + let rows = identity_pipeline(array, None, &GreptimePipelineParams::default()); assert!(rows.is_err()); assert_eq!( rows.err().unwrap().to_string(), @@ -489,7 +587,7 @@ mod tests { "gaga": "gaga" }), ]; - let rows = identity_pipeline(array, None); + let rows = identity_pipeline(array, None, &GreptimePipelineParams::default()); assert!(rows.is_err()); assert_eq!( rows.err().unwrap().to_string(), @@ -517,7 +615,7 @@ mod tests { "gaga": "gaga" }), ]; - let rows = identity_pipeline(array, None); + let rows = identity_pipeline(array, None, &GreptimePipelineParams::default()); assert!(rows.is_ok()); let rows = rows.unwrap(); assert_eq!(rows.schema.len(), 8); @@ -547,7 +645,11 @@ mod tests { }), ]; let tag_column_names = ["name".to_string(), "address".to_string()]; - let rows = identity_pipeline_inner(array, Some(tag_column_names.iter())); + let rows = identity_pipeline_inner( + array, + Some(tag_column_names.iter()), + &GreptimePipelineParams::default(), + ); assert!(rows.is_ok()); let rows = rows.unwrap(); assert_eq!(rows.schema.len(), 8); @@ -579,4 +681,89 @@ mod tests { ); } } + + #[test] + fn test_flatten() { + let test_cases = vec![ + // Basic case. + ( + serde_json::json!( + { + "a": { + "b": { + "c": [1, 2, 3] + } + }, + "d": [ + "foo", + "bar" + ], + "e": { + "f": [7, 8, 9], + "g": { + "h": 123, + "i": "hello", + "j": { + "k": true + } + } + } + } + ), + 10, + Some(serde_json::json!( + { + "a.b.c": [1,2,3], + "d": ["foo","bar"], + "e.f": [7,8,9], + "e.g.h": 123, + "e.g.i": "hello", + "e.g.j.k": true + } + )), + ), + // Test the case where the object has more than 3 nested levels. + ( + serde_json::json!( + { + "a": { + "b": { + "c": { + "d": [1, 2, 3] + } + } + }, + "e": [ + "foo", + "bar" + ] + } + ), + 3, + None, + ), + ]; + + for (input, max_depth, expected) in test_cases { + let flattened_object = + flatten_json_object(input.as_object().unwrap().clone(), max_depth); + match flattened_object { + Ok(flattened_object) => { + assert_eq!(&flattened_object, expected.unwrap().as_object().unwrap()) + } + Err(_) => assert_eq!(None, expected), + } + } + } + + #[test] + fn test_greptime_pipeline_params() { + let params = Some("flatten_json_object=true"); + let pipeline_params = GreptimePipelineParams::from_params(params); + assert!(pipeline_params.flatten_json_object()); + + let params = None; + let pipeline_params = GreptimePipelineParams::from_params(params); + assert!(!pipeline_params.flatten_json_object()); + } } diff --git a/src/pipeline/src/lib.rs b/src/pipeline/src/lib.rs index 49ecea41c44..edb6ce1f587 100644 --- a/src/pipeline/src/lib.rs +++ b/src/pipeline/src/lib.rs @@ -19,7 +19,9 @@ mod metrics; pub use etl::error::Result; pub use etl::processor::Processor; -pub use etl::transform::transformer::greptime::SchemaInfo; +pub use etl::transform::transformer::greptime::{ + GreptimePipelineParams, SchemaInfo, GREPTIME_PIPELINE_PARAMS_HEADER, +}; pub use etl::transform::transformer::identity_pipeline; pub use etl::transform::{GreptimeTransformer, Transformer}; pub use etl::value::{Array, Map, Value}; diff --git a/src/servers/src/elasticsearch.rs b/src/servers/src/elasticsearch.rs index 48638de2262..e9f1204e25b 100644 --- a/src/servers/src/elasticsearch.rs +++ b/src/servers/src/elasticsearch.rs @@ -91,9 +91,10 @@ pub async fn handle_bulk_api( Query(params): Query, Extension(query_ctx): Extension, TypedHeader(_content_type): TypedHeader, + headers: HeaderMap, payload: String, ) -> impl IntoResponse { - do_handle_bulk_api(log_state, None, params, query_ctx, payload).await + do_handle_bulk_api(log_state, None, params, query_ctx, headers, payload).await } /// Process `/${index}/_bulk` API requests. Only support to create logs. @@ -105,9 +106,10 @@ pub async fn handle_bulk_api_with_index( Query(params): Query, Extension(query_ctx): Extension, TypedHeader(_content_type): TypedHeader, + headers: HeaderMap, payload: String, ) -> impl IntoResponse { - do_handle_bulk_api(log_state, Some(index), params, query_ctx, payload).await + do_handle_bulk_api(log_state, Some(index), params, query_ctx, headers, payload).await } async fn do_handle_bulk_api( @@ -115,6 +117,7 @@ async fn do_handle_bulk_api( index: Option, params: LogIngesterQueryParams, mut query_ctx: QueryContext, + headers: HeaderMap, payload: String, ) -> impl IntoResponse { let start = Instant::now(); @@ -164,6 +167,7 @@ async fn do_handle_bulk_api( None, requests, Arc::new(query_ctx), + headers, ) .await { diff --git a/src/servers/src/http/event.rs b/src/servers/src/http/event.rs index f77220d8382..00b177b2c09 100644 --- a/src/servers/src/http/event.rs +++ b/src/servers/src/http/event.rs @@ -20,7 +20,7 @@ use api::v1::{RowInsertRequest, RowInsertRequests, Rows}; use async_trait::async_trait; use axum::extract::{FromRequest, Multipart, Path, Query, Request, State}; use axum::http::header::CONTENT_TYPE; -use axum::http::StatusCode; +use axum::http::{HeaderMap, StatusCode}; use axum::response::{IntoResponse, Response}; use axum::{Extension, Json}; use axum_extra::TypedHeader; @@ -32,7 +32,9 @@ use headers::ContentType; use lazy_static::lazy_static; use pipeline::error::PipelineTransformSnafu; use pipeline::util::to_pipeline_version; -use pipeline::{GreptimeTransformer, PipelineVersion}; +use pipeline::{ + GreptimePipelineParams, GreptimeTransformer, PipelineVersion, GREPTIME_PIPELINE_PARAMS_HEADER, +}; use serde::{Deserialize, Serialize}; use serde_json::{json, Deserializer, Map, Value}; use session::context::{Channel, QueryContext, QueryContextRef}; @@ -479,6 +481,7 @@ pub async fn log_ingester( Query(query_params): Query, Extension(mut query_ctx): Extension, TypedHeader(content_type): TypedHeader, + headers: HeaderMap, payload: String, ) -> Result { // validate source and payload @@ -523,6 +526,7 @@ pub async fn log_ingester( values: value, }], query_ctx, + headers, ) .await } @@ -552,19 +556,26 @@ pub(crate) async fn ingest_logs_inner( version: PipelineVersion, log_ingest_requests: Vec, query_ctx: QueryContextRef, + headers: HeaderMap, ) -> Result { let db = query_ctx.get_db_string(); let exec_timer = std::time::Instant::now(); let mut insert_requests = Vec::with_capacity(log_ingest_requests.len()); + let pipeline_params = GreptimePipelineParams::from_params( + headers + .get(GREPTIME_PIPELINE_PARAMS_HEADER) + .and_then(|v| v.to_str().ok()), + ); + for request in log_ingest_requests { let transformed_data: Rows = if pipeline_name == GREPTIME_INTERNAL_IDENTITY_PIPELINE_NAME { let table = state .get_table(&request.table, &query_ctx) .await .context(CatalogSnafu)?; - pipeline::identity_pipeline(request.values, table) + pipeline::identity_pipeline(request.values, table, &pipeline_params) .context(PipelineTransformSnafu) .context(PipelineSnafu)? } else { diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 8e1c4e116bb..76c17ac5cdc 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -90,6 +90,7 @@ macro_rules! http_tests { test_test_pipeline_api, test_plain_text_ingestion, test_identify_pipeline, + test_identify_pipeline_with_flatten, test_otlp_metrics, test_otlp_traces, @@ -1256,7 +1257,8 @@ transform: pub async fn test_identify_pipeline(store_type: StorageType) { common_telemetry::init_default_ut_logging(); - let (app, mut guard) = setup_test_http_app_with_frontend(store_type, "test_pipeline_api").await; + let (app, mut guard) = + setup_test_http_app_with_frontend(store_type, "test_identify_pipeline").await; // handshake let client = TestClient::new(app).await; @@ -1309,6 +1311,55 @@ pub async fn test_identify_pipeline(store_type: StorageType) { guard.remove_all().await; } +pub async fn test_identify_pipeline_with_flatten(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = + setup_test_http_app_with_frontend(store_type, "test_identify_pipeline_with_flatten").await; + + let client = TestClient::new(app).await; + let body = r#"{"__time__":1453809242,"__topic__":"","__source__":"10.170.***.***","ip":"10.200.**.***","time":"26/Jan/2016:19:54:02 +0800","url":"POST/PutData?Category=YunOsAccountOpLog&AccessKeyId=&Date=Fri%2C%2028%20Jun%202013%2006%3A53%3A30%20GMT&Topic=raw&Signature=HTTP/1.1","status":"200","user-agent":"aliyun-sdk-java","custom_map":{"value_a":["a","b","c"],"value_b":"b"}}"#; + + let res = send_req( + &client, + vec![ + ( + HeaderName::from_static("content-type"), + HeaderValue::from_static("application/json"), + ), + ( + HeaderName::from_static("x-greptime-pipeline-params"), + HeaderValue::from_static("flatten_json_object=true"), + ), + ], + "/v1/events/logs?table=logs&pipeline_name=greptime_identity", + body.as_bytes().to_vec(), + false, + ) + .await; + + assert_eq!(StatusCode::OK, res.status()); + + let expected = r#"[["__source__","String","","YES","","FIELD"],["__time__","Int64","","YES","","FIELD"],["__topic__","String","","YES","","FIELD"],["custom_map.value_a","Json","","YES","","FIELD"],["custom_map.value_b","String","","YES","","FIELD"],["ip","String","","YES","","FIELD"],["status","String","","YES","","FIELD"],["time","String","","YES","","FIELD"],["url","String","","YES","","FIELD"],["user-agent","String","","YES","","FIELD"],["greptime_timestamp","TimestampNanosecond","PRI","NO","","TIMESTAMP"]]"#; + validate_data( + "test_identify_pipeline_with_flatten_desc_logs", + &client, + "desc logs", + expected, + ) + .await; + + let expected = "[[[\"a\",\"b\",\"c\"]]]"; + validate_data( + "test_identify_pipeline_with_flatten_select_json", + &client, + "select `custom_map.value_a` from logs", + expected, + ) + .await; + + guard.remove_all().await; +} + pub async fn test_test_pipeline_api(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = setup_test_http_app_with_frontend(store_type, "test_pipeline_api").await;