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 <xixing.sys@gmail.com>
Co-authored-by: Ning Sun <sunng@protonmail.com>
Co-authored-by: Ning Sun <sunning@greptime.com>

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<T> = std::result::Result<T, Error>;

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<ColumnSchema>,
 }
 
+/// 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<String, String>,
+}
+
+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::<HashMap<String, String>>();
+
+        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<serde_json::Value>,
     tag_column_names: Option<impl Iterator<Item = &'a String>>,
+    params: &GreptimePipelineParams,
 ) -> Result<Rows> {
     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<serde_json::Value>,
     table: Option<Arc<table::Table>>,
+    params: &GreptimePipelineParams,
 ) -> Result<Rows> {
     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::<std::iter::Empty<&String>>),
+        None => identity_pipeline_inner(array, None::<std::iter::Empty<&String>>, 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<String, JsonValue>,
+    max_nested_levels: usize,
+) -> Result<Map<String, JsonValue>> {
+    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<String, JsonValue>,
+    base: Option<&str>,
+    object: Map<String, JsonValue>,
+    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());
+    }
 }

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};

src/servers/src/elasticsearch.rs

@@ -91,9 +91,10 @@ pub async fn handle_bulk_api(
     Query(params): Query<LogIngesterQueryParams>,
     Extension(query_ctx): Extension<QueryContext>,
     TypedHeader(_content_type): TypedHeader<ContentType>,
+    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<LogIngesterQueryParams>,
     Extension(query_ctx): Extension<QueryContext>,
     TypedHeader(_content_type): TypedHeader<ContentType>,
+    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<String>,
     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
     {

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<LogIngesterQueryParams>,
     Extension(mut query_ctx): Extension<QueryContext>,
     TypedHeader(content_type): TypedHeader<ContentType>,
+    headers: HeaderMap,
     payload: String,
 ) -> Result<HttpResponse> {
     // 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<LogIngestRequest>,
     query_ctx: QueryContextRef,
+    headers: HeaderMap,
 ) -> Result<HttpResponse> {
     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 {

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=<yourAccessKeyId>&Date=Fri%2C%2028%20Jun%202013%2006%3A53%3A30%20GMT&Topic=raw&Signature=<yourSignature>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;