feat: support decode gzip if influxdb write specify it

src/servers/Cargo.toml

@@ -48,6 +48,7 @@ datafusion-expr.workspace = true
 datatypes.workspace = true
 derive_builder.workspace = true
 digest = "0.10"
+flate2 = "1.0"
 futures = "0.3"
 hashbrown = "0.14"
 headers = "0.3"

src/servers/src/error.rs

@@ -298,6 +298,13 @@ pub enum Error {
         location: Location,
     },
 
+    #[snafu(display("Error decoding gzip stream"))]
+    InvalidGzip {
+        #[snafu(source)]
+        error: std::io::Error,
+        location: Location,
+    },
+
     #[snafu(display("Error accessing catalog"))]
     CatalogError { source: catalog::error::Error },
 
@@ -529,7 +536,7 @@ impl ErrorExt for Error {
             DatabaseNotFound { .. } => StatusCode::DatabaseNotFound,
             #[cfg(feature = "mem-prof")]
             DumpProfileData { source, .. } => source.status_code(),
-            InvalidFlushArgument { .. } => StatusCode::InvalidArguments,
+            InvalidFlushArgument { .. } | InvalidGzip { .. } => StatusCode::InvalidArguments,
 
             ReplacePreparedStmtParams { source, .. }
             | GetPreparedStmtParams { source, .. }

src/servers/src/http/influxdb.rs

@@ -15,15 +15,20 @@
 use std::collections::HashMap;
 
 use axum::extract::{Query, State};
-use axum::http::StatusCode;
+use axum::http::{HeaderMap, StatusCode};
 use axum::response::IntoResponse;
 use axum::Extension;
+use bytes::Bytes;
 use common_catalog::consts::DEFAULT_SCHEMA_NAME;
 use common_grpc::writer::Precision;
 use common_telemetry::tracing;
 use session::context::QueryContextRef;
+use snafu::ResultExt;
 
-use crate::error::{Result, TimePrecisionSnafu};
+use crate::error::{
+    InvalidGzipSnafu, InvalidQuerySnafu, InvalidUtf8ValueSnafu, InvisibleASCIISnafu, Result,
+    TimePrecisionSnafu,
+};
 use crate::influxdb::InfluxdbRequest;
 use crate::query_handler::InfluxdbLineProtocolHandlerRef;
 
@@ -39,13 +44,29 @@ pub async fn influxdb_health() -> Result<impl IntoResponse> {
     Ok(StatusCode::OK)
 }
 
+fn is_gzip(headers: &HeaderMap) -> Result<bool> {
+    match headers
+        .get("Content-Encoding")
+        .map(|val| val.to_str().context(InvisibleASCIISnafu))
+        .transpose()?
+    {
+        None | Some("identity") => Ok(false),
+        Some("gzip") => Ok(true),
+        Some(enc) => InvalidQuerySnafu {
+            reason: format!("unacceptable content-encoding: {enc}"),
+        }
+        .fail(),
+    }
+}
+
 #[axum_macros::debug_handler]
 #[tracing::instrument(skip_all, fields(protocol = "influxdb", request_type = "write_v1"))]
 pub async fn influxdb_write_v1(
     State(handler): State<InfluxdbLineProtocolHandlerRef>,
     Query(mut params): Query<HashMap<String, String>>,
     Extension(query_ctx): Extension<QueryContextRef>,
-    lines: String,
+    headers: HeaderMap,
+    lines: Bytes,
 ) -> Result<impl IntoResponse> {
     let db = params
         .remove("db")
@@ -56,7 +77,9 @@ pub async fn influxdb_write_v1(
         .map(|val| parse_time_precision(val))
         .transpose()?;
 
-    influxdb_write(&db, precision, lines, handler, query_ctx).await
+    let gzip = is_gzip(&headers)?;
+
+    influxdb_write(&db, precision, gzip, lines, handler, query_ctx).await
 }
 
 #[axum_macros::debug_handler]
@@ -65,7 +88,8 @@ pub async fn influxdb_write_v2(
     State(handler): State<InfluxdbLineProtocolHandlerRef>,
     Query(mut params): Query<HashMap<String, String>>,
     Extension(query_ctx): Extension<QueryContextRef>,
-    lines: String,
+    headers: HeaderMap,
+    lines: Bytes,
 ) -> Result<impl IntoResponse> {
     let db = match (params.remove("db"), params.remove("bucket")) {
         (_, Some(bucket)) => bucket.clone(),
@@ -78,20 +102,36 @@ pub async fn influxdb_write_v2(
         .map(|val| parse_time_precision(val))
         .transpose()?;
 
-    influxdb_write(&db, precision, lines, handler, query_ctx).await
+    let gzip = is_gzip(&headers)?;
+
+    influxdb_write(&db, precision, gzip, lines, handler, query_ctx).await
 }
 
 pub async fn influxdb_write(
     db: &str,
     precision: Option<Precision>,
-    lines: String,
+    gzip: bool,
+    lines: Bytes,
     handler: InfluxdbLineProtocolHandlerRef,
     ctx: QueryContextRef,
 ) -> Result<impl IntoResponse> {
     let _timer = crate::metrics::METRIC_HTTP_INFLUXDB_WRITE_ELAPSED
         .with_label_values(&[db])
         .start_timer();
 
+    let lines = if gzip {
+        // Unzip the gzip-encoded content
+        use std::io::Read;
+        let mut decoder = flate2::read::GzDecoder::new(&lines[..]);
+        let mut decoded_data = Vec::new();
+        decoder
+            .read_to_end(&mut decoded_data)
+            .context(InvalidGzipSnafu)?;
+        String::from_utf8(decoded_data).context(InvalidUtf8ValueSnafu)?
+    } else {
+        String::from_utf8(lines.into()).context(InvalidUtf8ValueSnafu)?
+    };
+
     let request = InfluxdbRequest { precision, lines };
     handler.exec(request, ctx).await?;