servers/http/

handler.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::HashMap;
use std::future;
use std::panic::AssertUnwindSafe;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};

use axum::extract::rejection::FormRejection;
use axum::extract::{Json, Query, State};
use axum::response::sse::{Event, KeepAlive, Sse};
use axum::response::{IntoResponse, Response};
use axum::{Extension, Form};
use common_catalog::parse_catalog_and_schema_from_db_string;
use common_error::ext::ErrorExt;
use common_error::status_code::StatusCode;
use common_plugins::GREPTIME_EXEC_WRITE_COST;
use common_query::{Output, OutputData};
use common_recordbatch::error::Result as RecordBatchResult;
use common_recordbatch::{RecordBatch, RecordBatchStreamWrapper, RecordBatches, util};
use common_telemetry::tracing;
use datafusion::physical_plan::ExecutionPlan;
use futures::{FutureExt, StreamExt};
use query::parser::{DEFAULT_LOOKBACK_STRING, PromQuery};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use session::context::{Channel, QueryContext, QueryContextRef};
use snafu::ResultExt;
use sql::dialect::GreptimeDbDialect;
use sql::parser::{ParseOptions, ParserContext};
use sql::statements::statement::Statement;
use tokio::sync::{Notify, watch};

use crate::error::{CollectRecordbatchSnafu, FailedToParseQuerySnafu, InvalidQuerySnafu, Result};
use crate::http::header::collect_plan_metrics;
use crate::http::result::arrow_result::ArrowResponse;
use crate::http::result::csv_result::CsvResponse;
use crate::http::result::error_result::ErrorResponse;
use crate::http::result::greptime_result_v1::GreptimedbV1Response;
use crate::http::result::influxdb_result_v1::InfluxdbV1Response;
use crate::http::result::json_result::JsonResponse;
use crate::http::result::null_result::NullResponse;
use crate::http::result::table_result::TableResponse;
use crate::http::{
    ApiState, Epoch, GreptimeOptionsConfigState, GreptimeQueryOutput, HttpRecordsOutput,
    HttpResponse, ResponseFormat,
};
use crate::metrics_handler::MetricsHandler;
use crate::query_handler::sql::ServerSqlQueryHandlerRef;

#[derive(Debug, Default, Serialize, Deserialize)]
pub struct SqlQuery {
    pub db: Option<String>,
    pub sql: Option<String>,
    // (Optional) result format: [`greptimedb_v1`, `influxdb_v1`, `csv`,
    // `arrow`],
    // the default value is `greptimedb_v1`
    pub format: Option<String>,
    // Returns epoch timestamps with the specified precision.
    // Both u and µ indicate microseconds.
    // epoch = [ns,u,µ,ms,s],
    //
    // TODO(jeremy): currently, only InfluxDB result format is supported,
    // and all columns of the `Timestamp` type will be converted to their
    // specified time precision. Maybe greptimedb format can support this
    // param too.
    pub epoch: Option<String>,
    pub limit: Option<usize>,
    // For arrow output
    pub compression: Option<String>,
    pub snapshot_interval_ms: Option<u64>,
}

const DEFAULT_ANALYZE_SNAPSHOT_INTERVAL_MS: u64 = 5000;
const MIN_ANALYZE_SNAPSHOT_INTERVAL_MS: u64 = 1000;
const MAX_ANALYZE_SNAPSHOT_INTERVAL_MS: u64 = 60000;

#[derive(Serialize)]
struct AnalyzeStreamPayload {
    seq: u64,
    state: &'static str,
    partial: bool,
    elapsed_ms: u64,
    #[serde(skip_serializing_if = "Option::is_none")]
    metrics: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    output: Option<GreptimeQueryOutput>,
    #[serde(skip_serializing_if = "Option::is_none")]
    reason: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    code: Option<u32>,
}

#[derive(Clone, Debug)]
#[doc(hidden)]
pub struct AnalyzeStreamMessage {
    pub event_name: &'static str,
    pub payload: String,
}

struct AnalyzeStreamWorkerGuard {
    cancel: watch::Sender<bool>,
    handle: common_runtime::JoinHandle<()>,
}

impl Drop for AnalyzeStreamWorkerGuard {
    fn drop(&mut self) {
        let _ = self.cancel.send(true);
        self.handle.abort();
    }
}

struct AnalyzeStreamBodyState {
    latest_metrics: watch::Receiver<Option<String>>,
    terminal: watch::Receiver<Option<AnalyzeStreamMessage>>,
    notify: Arc<Notify>,
    // Keeping the guard in the body state makes dropping the response cancel and
    // abort the worker instead of leaving an owned query stream detached.
    _worker: AnalyzeStreamWorkerGuard,
    done: bool,
}

/// Handler to execute sql
#[axum_macros::debug_handler]
#[tracing::instrument(skip_all, fields(protocol = "http", request_type = "sql"))]
pub async fn sql(
    State(state): State<ApiState>,
    Query(query_params): Query<SqlQuery>,
    Extension(mut query_ctx): Extension<QueryContext>,
    Form(form_params): Form<SqlQuery>,
) -> HttpResponse {
    let start = Instant::now();
    let sql_handler = &state.sql_handler;
    if let Some(db) = &query_params.db.or(form_params.db) {
        let (catalog, schema) = parse_catalog_and_schema_from_db_string(db);
        query_ctx.set_current_catalog(&catalog);
        query_ctx.set_current_schema(&schema);
    }
    let db = query_ctx.get_db_string();

    query_ctx.set_channel(Channel::HttpSql);
    let query_ctx = Arc::new(query_ctx);

    let _timer = crate::metrics::METRIC_HTTP_SQL_ELAPSED
        .with_label_values(&[db.as_str()])
        .start_timer();

    let sql = query_params.sql.or(form_params.sql);
    let format = query_params
        .format
        .or(form_params.format)
        .map(|s| s.to_lowercase())
        .map(|s| ResponseFormat::parse(s.as_str()).unwrap_or(ResponseFormat::GreptimedbV1))
        .unwrap_or(ResponseFormat::GreptimedbV1);
    let epoch = query_params
        .epoch
        .or(form_params.epoch)
        .map(|s| s.to_lowercase())
        .map(|s| Epoch::parse(s.as_str()).unwrap_or(Epoch::Millisecond));

    let result = if let Some(sql) = &sql {
        if let Some((status, msg)) = validate_schema(sql_handler.clone(), query_ctx.clone()).await {
            Err((status, msg))
        } else {
            Ok(sql_handler.do_query(sql, query_ctx.clone()).await)
        }
    } else {
        Err((
            StatusCode::InvalidArguments,
            "sql parameter is required.".to_string(),
        ))
    };

    let outputs = match result {
        Err((status, msg)) => {
            return HttpResponse::Error(
                ErrorResponse::from_error_message(status, msg)
                    .with_execution_time(start.elapsed().as_millis() as u64),
            );
        }
        Ok(outputs) => outputs,
    };

    let outputs = match query_params.limit {
        Some(limit)
            if matches!(
                format,
                ResponseFormat::Csv(..)
                    | ResponseFormat::Table
                    | ResponseFormat::GreptimedbV1
                    | ResponseFormat::Json
            ) =>
        {
            outputs
                .into_iter()
                .map(|output| output.and_then(|output| limit_output_rows(output, limit)))
                .collect()
        }
        _ => outputs,
    };

    let resp = match format {
        ResponseFormat::Arrow => {
            ArrowResponse::from_output(outputs, query_params.compression).await
        }
        ResponseFormat::Csv(with_names, with_types) => {
            CsvResponse::from_output(outputs, with_names, with_types).await
        }
        ResponseFormat::Table => TableResponse::from_output(outputs).await,
        ResponseFormat::GreptimedbV1 => GreptimedbV1Response::from_output(outputs).await,
        ResponseFormat::InfluxdbV1 => InfluxdbV1Response::from_output(outputs, epoch).await,
        ResponseFormat::Json => JsonResponse::from_output(outputs).await,
        ResponseFormat::Null => NullResponse::from_output(outputs).await,
    };

    resp.with_execution_time(start.elapsed().as_millis() as u64)
}

/// Limits each statement's response before collecting batches and converting rows.
fn limit_output_rows(output: Output, limit: usize) -> Result<Output> {
    let mut remaining = limit;
    let data = match output.data {
        OutputData::AffectedRows(rows) => OutputData::AffectedRows(rows),
        OutputData::RecordBatches(batches) => {
            let schema = batches.schema();
            let batches = batches
                .into_iter()
                .filter_map(|batch| take_response_rows(batch, &mut remaining).transpose())
                .collect::<RecordBatchResult<Vec<_>>>()
                .context(CollectRecordbatchSnafu)?;
            OutputData::RecordBatches(
                RecordBatches::try_new(schema, batches).context(CollectRecordbatchSnafu)?,
            )
        }
        OutputData::Stream(stream) => {
            let schema = stream.schema();
            // This is a response limit. Drain the input even after reaching it so
            // execution errors and terminal query metrics are still observed.
            let stream = stream.filter_map(move |batch| {
                future::ready(
                    batch
                        .and_then(|batch| take_response_rows(batch, &mut remaining))
                        .transpose(),
                )
            });
            OutputData::Stream(Box::pin(RecordBatchStreamWrapper::new(schema, stream)))
        }
    };
    Ok(Output::new(data, output.meta))
}

fn take_response_rows(
    batch: RecordBatch,
    remaining: &mut usize,
) -> RecordBatchResult<Option<RecordBatch>> {
    let num_rows = batch.num_rows().min(*remaining);
    *remaining -= num_rows;
    if num_rows == 0 {
        Ok(None)
    } else if num_rows == batch.num_rows() {
        Ok(Some(batch))
    } else {
        batch.slice(0, num_rows).map(Some)
    }
}

/// Handler to stream partial `EXPLAIN ANALYZE VERBOSE` metrics as SSE.
///
/// This endpoint is POST-only SSE, so browser `EventSource` does
/// not apply. Each `metrics` event carries a complete snapshot (not a delta);
/// large snapshots are throttled but never truncated. `final`, `canceled`, and
/// `error` are terminal events. If the client disconnects it won't receive a
/// `canceled` event, but the production frontend stream is dropped and
/// best-effort cancels the underlying query.
#[axum_macros::debug_handler]
#[tracing::instrument(
    skip_all,
    fields(protocol = "http", request_type = "sql_analyze_stream")
)]
pub async fn sql_analyze_stream(
    State(state): State<ApiState>,
    Query(query_params): Query<SqlQuery>,
    Extension(mut query_ctx): Extension<QueryContext>,
    form_params: std::result::Result<Form<SqlQuery>, FormRejection>,
) -> Response {
    let start = Instant::now();
    let form_params = match form_params {
        Ok(Form(params)) => params,
        Err(err) => {
            if err.status() != axum::http::StatusCode::UNSUPPORTED_MEDIA_TYPE {
                return ErrorResponse::from_error_message(
                    StatusCode::InvalidArguments,
                    err.body_text(),
                )
                .with_execution_time(start.elapsed().as_millis() as u64)
                .into_response();
            }
            SqlQuery::default()
        }
    };
    let sql_handler = &state.sql_handler;
    if let Some(db) = &query_params.db.or(form_params.db) {
        let (catalog, schema) = parse_catalog_and_schema_from_db_string(db);
        query_ctx.set_current_catalog(&catalog);
        query_ctx.set_current_schema(&schema);
    }
    query_ctx.set_channel(Channel::HttpSql);
    query_ctx.enable_live_analyze_metrics();
    let query_ctx = Arc::new(query_ctx);

    let Some(sql) = query_params.sql.or(form_params.sql) else {
        return ErrorResponse::from_error_message(
            StatusCode::InvalidArguments,
            "sql parameter is required.".to_string(),
        )
        .with_execution_time(start.elapsed().as_millis() as u64)
        .into_response();
    };
    if let Some((status, msg)) = validate_schema(sql_handler.clone(), query_ctx.clone()).await {
        return ErrorResponse::from_error_message(status, msg)
            .with_execution_time(start.elapsed().as_millis() as u64)
            .into_response();
    }

    let interval_ms = query_params
        .snapshot_interval_ms
        .or(form_params.snapshot_interval_ms)
        .unwrap_or(DEFAULT_ANALYZE_SNAPSHOT_INTERVAL_MS)
        .clamp(
            MIN_ANALYZE_SNAPSHOT_INTERVAL_MS,
            MAX_ANALYZE_SNAPSHOT_INTERVAL_MS,
        );

    let output = match state
        .sql_handler
        .do_analyze_stream_query(&sql, query_ctx.clone())
        .await
    {
        Ok(output) => output,
        Err(err) => {
            return ErrorResponse::from_error(err)
                .with_execution_time(start.elapsed().as_millis() as u64)
                .into_response();
        }
    };

    let plan = output.meta.plan.clone();
    let OutputData::Stream(stream) = output.data else {
        return ErrorResponse::from_error_message(
            StatusCode::InvalidArguments,
            "analyze stream query must return a stream".to_string(),
        )
        .with_execution_time(start.elapsed().as_millis() as u64)
        .into_response();
    };
    let schema = stream.schema();

    let (metrics_tx, metrics_rx) = watch::channel::<Option<String>>(None);
    let (terminal_tx, terminal_rx) = watch::channel::<Option<AnalyzeStreamMessage>>(None);
    let notify = Arc::new(Notify::new());
    let (cancel_tx, mut cancel_rx) = watch::channel(false);
    let worker_notify = notify.clone();
    let panic_notify = notify.clone();
    let sequence = Arc::new(AtomicU64::new(0));
    let worker_sequence = sequence.clone();
    let worker_terminal_tx = terminal_tx.clone();
    let worker = common_runtime::spawn_global(async move {
        let worker_result = AssertUnwindSafe(async move {
            let mut stream = stream;
            let mut batches = Vec::new();
            let mut current_interval_ms = interval_ms;
            let tick = tokio::time::sleep(Duration::from_millis(current_interval_ms));
            tokio::pin!(tick);

            loop {
                tokio::select! {
                    _ = cancel_rx.changed() => return,
                    item = stream.next() => {
                        match item {
                            Some(Ok(next_batch)) => batches.push(next_batch),
                            Some(Err(err)) => {
                                let status = err.status_code();
                                let event_name = if status == StatusCode::Cancelled { "canceled" } else { "error" };
                                let (payload, _) = make_analyze_payload(AnalyzePayloadArgs {
                                    seq: worker_sequence.load(Ordering::Relaxed),
                                    state: event_name,
                                    partial: false,
                                    start,
                                    plan: plan.as_ref(),
                                    output: None,
                                    reason: Some(err.output_msg()),
                                    code: Some(status as u32),
                                });
                                send_analyze_terminal(&terminal_tx, &worker_notify, event_name, payload);
                                return;
                            }
                            None => {
                                let output = HttpRecordsOutput::try_new(schema.clone(), batches)
                                    .map(GreptimeQueryOutput::Records);
                                let (event_name, payload) = make_final_analyze_event(
                                    output.map_err(|err| (err.output_msg(), err.status_code() as u32)),
                                    worker_sequence.load(Ordering::Relaxed),
                                    start,
                                    plan.as_ref(),
                                );
                                send_analyze_terminal(&terminal_tx, &worker_notify, event_name, payload);
                                return;
                            }
                        }
                    }
                    _ = &mut tick, if plan.is_some() => {
                        let (payload, payload_bytes) = make_analyze_payload(AnalyzePayloadArgs {
                            seq: worker_sequence.load(Ordering::Relaxed),
                            state: "metrics",
                            partial: true,
                            start,
                            plan: plan.as_ref(),
                            output: None,
                            reason: None,
                            code: None,
                        });
                        current_interval_ms = adaptive_interval_ms(payload_bytes, interval_ms);
                        worker_sequence.fetch_add(1, Ordering::Relaxed);
                        if metrics_tx.send(Some(payload)).is_err() {
                            return;
                        }
                        worker_notify.notify_one();
                        tick.as_mut().reset(tokio::time::Instant::now() + Duration::from_millis(current_interval_ms));
                    }
                }
            }
        })
        .catch_unwind()
        .await;

        if worker_result.is_err() {
            tracing::debug!("analyze stream worker panicked");
            let (payload, _) = make_analyze_payload(AnalyzePayloadArgs {
                seq: sequence.load(Ordering::Relaxed),
                state: "error",
                partial: false,
                start,
                plan: None,
                output: None,
                reason: Some("analyze stream worker panicked".to_string()),
                code: Some(StatusCode::Internal as u32),
            });
            send_analyze_terminal(&worker_terminal_tx, &panic_notify, "error", payload);
        }
    });

    let sse_stream = analyze_stream_body(metrics_rx, terminal_rx, notify, cancel_tx, worker);

    Sse::new(sse_stream)
        .keep_alive(KeepAlive::new().interval(Duration::from_secs(15)))
        .into_response()
}

#[doc(hidden)]
pub fn analyze_stream_body(
    metrics: watch::Receiver<Option<String>>,
    terminal: watch::Receiver<Option<AnalyzeStreamMessage>>,
    notify: Arc<Notify>,
    cancel: watch::Sender<bool>,
    worker: common_runtime::JoinHandle<()>,
) -> impl futures::Stream<Item = std::result::Result<Event, std::convert::Infallible>> {
    futures::stream::unfold(
        AnalyzeStreamBodyState {
            latest_metrics: metrics,
            terminal,
            notify,
            _worker: AnalyzeStreamWorkerGuard {
                cancel,
                handle: worker,
            },
            done: false,
        },
        |mut state| async move {
            if state.done {
                return None;
            }
            loop {
                let notify = Arc::clone(&state.notify);
                let notified = notify.notified();
                tokio::pin!(notified);
                // Register before checking the slots to avoid a lost wakeup.
                notified.as_mut().enable();
                let latest = {
                    let latest = state.latest_metrics.borrow_and_update();
                    latest.has_changed().then(|| latest.clone())
                };
                if let Some(Some(payload)) = latest {
                    return Some((Ok(Event::default().event("metrics").data(payload)), state));
                }
                // Inspect the terminal slot directly because a closed watch channel
                // can otherwise hide a value published while a worker was unwinding.
                if state.terminal.borrow().is_some() {
                    let terminal = { state.terminal.borrow_and_update().clone() };
                    if let Some(AnalyzeStreamMessage {
                        event_name,
                        payload,
                    }) = terminal
                    {
                        state.done = true;
                        return Some((Ok(Event::default().event(event_name).data(payload)), state));
                    }
                }
                notified.await;
            }
        },
    )
}

#[doc(hidden)]
pub fn send_analyze_terminal(
    terminal_tx: &watch::Sender<Option<AnalyzeStreamMessage>>,
    notify: &Notify,
    event_name: &'static str,
    payload: String,
) {
    if terminal_tx.send_if_modified(|terminal| {
        if terminal.is_none() {
            *terminal = Some(AnalyzeStreamMessage {
                event_name,
                payload,
            });
            true
        } else {
            false
        }
    }) {
        notify.notify_one();
    }
}

fn adaptive_interval_ms(payload_bytes: usize, requested_ms: u64) -> u64 {
    if payload_bytes >= 10 * 1024 * 1024 {
        requested_ms.max(30_000)
    } else if payload_bytes >= 1024 * 1024 {
        requested_ms.max(10_000)
    } else {
        requested_ms
    }
}

fn make_final_analyze_event(
    output: std::result::Result<GreptimeQueryOutput, (String, u32)>,
    seq: u64,
    start: Instant,
    plan: Option<&Arc<dyn ExecutionPlan>>,
) -> (&'static str, String) {
    match output {
        Ok(output) => (
            "final",
            make_analyze_payload(AnalyzePayloadArgs {
                seq,
                state: "final",
                partial: false,
                start,
                plan,
                output: Some(output),
                reason: None,
                code: None,
            })
            .0,
        ),
        Err((reason, code)) => (
            "error",
            make_analyze_payload(AnalyzePayloadArgs {
                seq,
                state: "error",
                partial: false,
                start,
                plan,
                output: None,
                reason: Some(reason),
                code: Some(code),
            })
            .0,
        ),
    }
}

struct AnalyzePayloadArgs<'a> {
    seq: u64,
    state: &'static str,
    partial: bool,
    start: Instant,
    plan: Option<&'a Arc<dyn ExecutionPlan>>,
    output: Option<GreptimeQueryOutput>,
    reason: Option<String>,
    code: Option<u32>,
}

fn make_analyze_payload(args: AnalyzePayloadArgs<'_>) -> (String, usize) {
    let AnalyzePayloadArgs {
        seq,
        state,
        partial,
        start,
        plan,
        output,
        reason,
        code,
    } = args;
    // Periodic snapshots are compact; terminal snapshots retain verbose plan details.
    let metrics =
        plan.and_then(|plan| query::analyze_plan_metrics_to_json_value(plan, !partial).ok());
    let payload = AnalyzeStreamPayload {
        seq,
        state,
        partial,
        elapsed_ms: start.elapsed().as_millis() as u64,
        metrics,
        output,
        reason,
        code,
    };
    let payload_string = serde_json::to_string(&payload).unwrap_or_else(|e| {
        serde_json::json!({
            "seq": seq,
            "state": "error",
            "partial": false,
            "reason": format!("Failed to serialize SSE payload: {e}"),
        })
        .to_string()
    });
    let payload_bytes = payload_string.len();
    (payload_string, payload_bytes)
}

/// Handler to parse sql
#[axum_macros::debug_handler]
#[tracing::instrument(skip_all, fields(protocol = "http", request_type = "sql"))]
pub async fn sql_parse(
    Query(query_params): Query<SqlQuery>,
    Form(form_params): Form<SqlQuery>,
) -> Result<Json<Vec<Statement>>> {
    let Some(sql) = query_params.sql.or(form_params.sql) else {
        return InvalidQuerySnafu {
            reason: "sql parameter is required.",
        }
        .fail();
    };

    let stmts =
        ParserContext::create_with_dialect(&sql, &GreptimeDbDialect {}, ParseOptions::default())
            .context(FailedToParseQuerySnafu)?;

    Ok(stmts.into())
}

#[derive(Debug, Serialize, Deserialize)]
pub struct SqlFormatResponse {
    pub formatted: String,
}

/// Handler to format sql string
#[axum_macros::debug_handler]
#[tracing::instrument(skip_all, fields(protocol = "http", request_type = "sql_format"))]
pub async fn sql_format(
    Query(query_params): Query<SqlQuery>,
    Form(form_params): Form<SqlQuery>,
) -> axum::response::Response {
    let Some(sql) = query_params.sql.or(form_params.sql) else {
        let resp = ErrorResponse::from_error_message(
            StatusCode::InvalidArguments,
            "sql parameter is required.".to_string(),
        );
        return HttpResponse::Error(resp).into_response();
    };

    // Parse using GreptimeDB dialect then reconstruct statements via Display
    let stmts = match ParserContext::create_with_dialect(
        &sql,
        &GreptimeDbDialect {},
        ParseOptions::default(),
    ) {
        Ok(v) => v,
        Err(e) => return HttpResponse::Error(ErrorResponse::from_error(e)).into_response(),
    };

    let mut parts: Vec<String> = Vec::with_capacity(stmts.len());
    for stmt in stmts {
        let mut s = format!("{stmt}");
        if !s.trim_end().ends_with(';') {
            s.push(';');
        }
        parts.push(s);
    }

    let formatted = parts.join("\n");
    Json(SqlFormatResponse { formatted }).into_response()
}

/// Create a response from query result
pub async fn from_output(
    outputs: Vec<crate::error::Result<Output>>,
) -> std::result::Result<(Vec<GreptimeQueryOutput>, HashMap<String, Value>), ErrorResponse> {
    // TODO(sunng87): this api response structure cannot represent error well.
    //  It hides successful execution results from error response
    let mut results = Vec::with_capacity(outputs.len());
    let mut merge_map = HashMap::new();

    for out in outputs {
        match out {
            Ok(o) => match o.data {
                OutputData::AffectedRows(rows) => {
                    results.push(GreptimeQueryOutput::AffectedRows(rows));
                    if o.meta.cost > 0 {
                        merge_map.insert(GREPTIME_EXEC_WRITE_COST.to_string(), o.meta.cost as u64);
                    }
                }
                OutputData::Stream(stream) => {
                    let schema = stream.schema().clone();
                    // TODO(sunng87): streaming response
                    let mut http_record_output = match util::collect(stream).await {
                        Ok(rows) => match HttpRecordsOutput::try_new(schema, rows) {
                            Ok(rows) => rows,
                            Err(err) => {
                                return Err(ErrorResponse::from_error(err));
                            }
                        },
                        Err(err) => {
                            return Err(ErrorResponse::from_error(err));
                        }
                    };
                    if let Some(physical_plan) = o.meta.plan {
                        let mut result_map = HashMap::new();

                        let mut tmp = vec![&mut merge_map, &mut result_map];
                        collect_plan_metrics(&physical_plan, &mut tmp);
                        let re = result_map
                            .into_iter()
                            .map(|(k, v)| (k, Value::from(v)))
                            .collect::<HashMap<String, Value>>();
                        http_record_output.metrics.extend(re);
                    }
                    results.push(GreptimeQueryOutput::Records(http_record_output))
                }
                OutputData::RecordBatches(rbs) => {
                    match HttpRecordsOutput::try_new(rbs.schema(), rbs.take()) {
                        Ok(rows) => {
                            results.push(GreptimeQueryOutput::Records(rows));
                        }
                        Err(err) => {
                            return Err(ErrorResponse::from_error(err));
                        }
                    }
                }
            },

            Err(err) => {
                return Err(ErrorResponse::from_error(err));
            }
        }
    }

    let merge_map = merge_map
        .into_iter()
        .map(|(k, v)| (k, Value::from(v)))
        .collect();

    Ok((results, merge_map))
}

#[derive(Debug, Default, Serialize, Deserialize)]
pub struct PromqlQuery {
    pub query: String,
    pub start: String,
    pub end: String,
    pub step: String,
    pub lookback: Option<String>,
    pub db: Option<String>,
    // (Optional) result format: [`greptimedb_v1`, `influxdb_v1`, `csv`,
    // `arrow`],
    // the default value is `greptimedb_v1`
    pub format: Option<String>,
    // For arrow output
    pub compression: Option<String>,
    // Returns epoch timestamps with the specified precision.
    // Both u and µ indicate microseconds.
    // epoch = [ns,u,µ,ms,s],
    //
    // For influx output only
    //
    // TODO(jeremy): currently, only InfluxDB result format is supported,
    // and all columns of the `Timestamp` type will be converted to their
    // specified time precision. Maybe greptimedb format can support this
    // param too.
    pub epoch: Option<String>,
}

impl From<PromqlQuery> for PromQuery {
    fn from(query: PromqlQuery) -> Self {
        PromQuery {
            query: query.query,
            start: query.start,
            end: query.end,
            step: query.step,
            lookback: query
                .lookback
                .unwrap_or_else(|| DEFAULT_LOOKBACK_STRING.to_string()),
            // TODO(dennis): support alias from http params or parse from query.query
            alias: None,
        }
    }
}

/// Handler to execute promql
#[axum_macros::debug_handler]
#[tracing::instrument(skip_all, fields(protocol = "http", request_type = "promql"))]
pub async fn promql(
    State(state): State<ApiState>,
    Query(params): Query<PromqlQuery>,
    Extension(mut query_ctx): Extension<QueryContext>,
) -> Response {
    let sql_handler = &state.sql_handler;
    let exec_start = Instant::now();
    let db = query_ctx.get_db_string();

    query_ctx.set_channel(Channel::Promql);
    let query_ctx = Arc::new(query_ctx);

    let _timer = crate::metrics::METRIC_HTTP_PROMQL_ELAPSED
        .with_label_values(&[db.as_str()])
        .start_timer();

    let resp = if let Some((status, msg)) =
        validate_schema(sql_handler.clone(), query_ctx.clone()).await
    {
        let resp = ErrorResponse::from_error_message(status, msg);
        HttpResponse::Error(resp)
    } else {
        let format = params
            .format
            .as_ref()
            .map(|s| s.to_lowercase())
            .map(|s| ResponseFormat::parse(s.as_str()).unwrap_or(ResponseFormat::GreptimedbV1))
            .unwrap_or(ResponseFormat::GreptimedbV1);
        let epoch = params
            .epoch
            .as_ref()
            .map(|s| s.to_lowercase())
            .map(|s| Epoch::parse(s.as_str()).unwrap_or(Epoch::Millisecond));
        let compression = params.compression.clone();

        let prom_query = params.into();
        let outputs = sql_handler.do_promql_query(&prom_query, query_ctx).await;

        match format {
            ResponseFormat::Arrow => ArrowResponse::from_output(outputs, compression).await,
            ResponseFormat::Csv(with_names, with_types) => {
                CsvResponse::from_output(outputs, with_names, with_types).await
            }
            ResponseFormat::Table => TableResponse::from_output(outputs).await,
            ResponseFormat::GreptimedbV1 => GreptimedbV1Response::from_output(outputs).await,
            ResponseFormat::InfluxdbV1 => InfluxdbV1Response::from_output(outputs, epoch).await,
            ResponseFormat::Json => JsonResponse::from_output(outputs).await,
            ResponseFormat::Null => NullResponse::from_output(outputs).await,
        }
    };

    resp.with_execution_time(exec_start.elapsed().as_millis() as u64)
        .into_response()
}

/// Handler to export metrics
#[axum_macros::debug_handler]
pub async fn metrics(
    State(state): State<MetricsHandler>,
    Query(_params): Query<HashMap<String, String>>,
) -> String {
    // A default ProcessCollector is registered automatically in prometheus.
    // We do not need to explicitly collect process-related data.
    // But ProcessCollector only support on linux.

    #[cfg(not(windows))]
    if let Some(c) = crate::metrics::jemalloc::JEMALLOC_COLLECTOR.as_ref()
        && let Err(e) = c.update()
    {
        common_telemetry::error!(e; "Failed to update jemalloc metrics");
    }
    state.render()
}

#[derive(Debug, Serialize, Deserialize)]
pub struct HealthQuery {}

#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct HealthResponse {}

/// Handler to export healthy check
///
/// Currently simply return status "200 OK" (default) with an empty json payload "{}"
#[axum_macros::debug_handler]
pub async fn health(Query(_params): Query<HealthQuery>) -> Json<HealthResponse> {
    Json(HealthResponse {})
}

#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct StatusResponse<'a> {
    pub commit: &'a str,
    pub branch: &'a str,
    pub rustc_version: &'a str,
    pub hostname: String,
    pub version: &'a str,
}

/// Handler to expose information info about runtime, build, etc.
#[axum_macros::debug_handler]
pub async fn status() -> Json<StatusResponse<'static>> {
    let hostname = hostname::get()
        .map(|s| s.to_string_lossy().to_string())
        .unwrap_or_else(|_| "unknown".to_string());
    let build_info = common_version::build_info();
    Json(StatusResponse {
        commit: build_info.commit,
        branch: build_info.branch,
        rustc_version: build_info.rustc,
        hostname,
        version: build_info.version,
    })
}

/// Handler to expose configuration information info about runtime, build, etc.
#[axum_macros::debug_handler]
pub async fn config(State(state): State<GreptimeOptionsConfigState>) -> Response {
    (axum::http::StatusCode::OK, state.greptime_config_options).into_response()
}

async fn validate_schema(
    sql_handler: ServerSqlQueryHandlerRef,
    query_ctx: QueryContextRef,
) -> Option<(StatusCode, String)> {
    match sql_handler
        .is_valid_schema(query_ctx.current_catalog(), &query_ctx.current_schema())
        .await
    {
        Ok(true) => None,
        Ok(false) => Some((
            StatusCode::DatabaseNotFound,
            format!("Database not found: {}", query_ctx.get_db_string()),
        )),
        Err(e) => Some((
            StatusCode::Internal,
            format!(
                "Error checking database: {}, {}",
                query_ctx.get_db_string(),
                e.output_msg(),
            ),
        )),
    }
}

pub async fn index() -> axum::response::Html<String> {
    let name = common_version::product_name();
    let version = common_version::version();
    axum::response::Html(format!(
        r#"<!DOCTYPE html>
<html>
<head><title>{name}</title></head>
<body>
<h1>{name}</h1>
<p>Version: {version}</p>
<ul>
<li><a href="/dashboard">Dashboard UI</a></li>
<li><a href="/v1/health">Health</a> (JSON)</li>
<li><a href="/status">Status</a> (JSON)</li>
<li><a href="/metrics">Metrics</a> (For Prometheus Scrape)</li>
<li><a href="/config">Config</a> (TXT)</li>
</ul>
</body>
</html>"#,
    ))
}

#[cfg(test)]
mod tests {
    use std::sync::Mutex;
    use std::sync::atomic::AtomicUsize;

    use async_trait::async_trait;
    use common_query::OutputMeta;
    use datafusion_expr::LogicalPlan;
    use datatypes::prelude::ConcreteDataType;
    use datatypes::schema::{ColumnSchema, Schema};
    use datatypes::vectors::{BinaryVector, UInt32Vector, VectorRef};
    use futures::stream;
    use query::query_engine::DescribeResult;

    use super::*;
    use crate::query_handler::sql::SqlQueryHandler;

    struct TestSqlQueryHandler {
        outputs: Mutex<Vec<Result<Output>>>,
    }

    #[async_trait]
    impl SqlQueryHandler for TestSqlQueryHandler {
        async fn do_query(&self, _: &str, _: QueryContextRef) -> Vec<Result<Output>> {
            std::mem::take(&mut *self.outputs.lock().unwrap())
        }

        async fn do_analyze_stream_query(&self, _: &str, _: QueryContextRef) -> Result<Output> {
            unimplemented!()
        }

        async fn do_exec_plan(
            &self,
            _: LogicalPlan,
            _: Option<Statement>,
            _: QueryContextRef,
        ) -> Result<Output> {
            unimplemented!()
        }

        async fn do_promql_query(&self, _: &PromQuery, _: QueryContextRef) -> Vec<Result<Output>> {
            unimplemented!()
        }

        async fn do_describe(
            &self,
            _: Statement,
            _: QueryContextRef,
        ) -> Result<Option<DescribeResult>> {
            unimplemented!()
        }

        async fn is_valid_schema(&self, _: &str, _: &str) -> Result<bool> {
            Ok(true)
        }
    }

    async fn sql_response(
        outputs: Vec<Result<Output>>,
        format: &str,
        limit: Option<usize>,
    ) -> HttpResponse {
        sql(
            State(ApiState {
                sql_handler: Arc::new(TestSqlQueryHandler {
                    outputs: Mutex::new(outputs),
                }),
            }),
            Query(SqlQuery {
                sql: Some("select number from numbers".to_string()),
                format: Some(format.to_string()),
                limit,
                ..Default::default()
            }),
            Extension(QueryContext::with_db_name(None)),
            Form(SqlQuery::default()),
        )
        .await
    }

    fn number_batches() -> RecordBatches {
        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
            "number",
            ConcreteDataType::uint32_datatype(),
            false,
        )]));
        let batches = [vec![], vec![0, 1], vec![], vec![2, 3, 4], vec![]]
            .into_iter()
            .map(|values| {
                let columns: Vec<VectorRef> = vec![Arc::new(UInt32Vector::from_slice(values))];
                RecordBatch::new(schema.clone(), columns).unwrap()
            })
            .collect();
        RecordBatches::try_new(schema, batches).unwrap()
    }

    fn number_output(streaming: bool) -> Output {
        let batches = number_batches();
        if streaming {
            Output::new_with_stream(batches.as_stream())
        } else {
            Output::new_with_record_batches(batches)
        }
    }

    #[tokio::test]
    async fn test_sql_response_limit_across_batches_and_formats() {
        for format in [
            "greptimedb_v1",
            "json",
            "csv",
            "csvWithNames",
            "csvWithNamesAndTypes",
            "table",
        ] {
            for streaming in [false, true] {
                for limit in [
                    None,
                    Some(0),
                    Some(1),
                    Some(2),
                    Some(3),
                    Some(5),
                    Some(10),
                    Some(usize::MAX),
                ] {
                    let response =
                        sql_response(vec![Ok(number_output(streaming))], format, limit).await;
                    let output = match &response {
                        HttpResponse::GreptimedbV1(response) => response.output(),
                        HttpResponse::Json(response) => response.output(),
                        HttpResponse::Csv(response) => response.output(),
                        HttpResponse::Table(response) => response.output(),
                        _ => panic!("unexpected response: {response:?}"),
                    };
                    let GreptimeQueryOutput::Records(records) = &output[0] else {
                        panic!("expected records");
                    };
                    let expected_rows: Vec<_> = (0..limit.unwrap_or(5).min(5))
                        .map(|number| vec![Value::from(number)])
                        .collect();
                    assert_eq!(records.rows(), &expected_rows);
                    assert_eq!(records.total_rows, expected_rows.len());
                    assert_eq!(records.schema.column_schemas[0].name, "number");
                }
            }
        }
    }

    #[tokio::test]
    async fn test_sql_response_limit_is_per_statement_and_preserves_write_cost() {
        let affected = Output::new(
            OutputData::AffectedRows(7),
            OutputMeta {
                cost: 42,
                ..Default::default()
            },
        );
        let response = sql_response(
            vec![
                Ok(number_output(true)),
                Ok(affected),
                Ok(number_output(false)),
            ],
            "greptimedb_v1",
            Some(1),
        )
        .await;
        let HttpResponse::GreptimedbV1(response) = response else {
            panic!("expected greptimedb response");
        };
        assert_eq!(response.output().len(), 3);
        for index in [0, 2] {
            let GreptimeQueryOutput::Records(records) = &response.output()[index] else {
                panic!("expected records");
            };
            assert_eq!(records.rows(), &vec![vec![Value::from(0)]]);
        }
        assert!(matches!(
            response.output()[1],
            GreptimeQueryOutput::AffectedRows(7)
        ));
        assert_eq!(
            response.resp_metrics[GREPTIME_EXEC_WRITE_COST],
            Value::from(42)
        );
    }

    #[tokio::test]
    async fn test_sql_response_limit_does_not_affect_other_formats() {
        for format in ["arrow", "influxdb_v1", "null"] {
            let unlimited = sql_response(vec![Ok(number_output(true))], format, None).await;
            let limited = sql_response(vec![Ok(number_output(true))], format, Some(0)).await;
            assert_eq!(
                serde_json::to_value(unlimited.with_execution_time(0)).unwrap(),
                serde_json::to_value(limited.with_execution_time(0)).unwrap(),
            );
        }
    }

    #[tokio::test]
    async fn test_sql_response_limit_drains_input_and_propagates_late_errors() {
        for limit in [0, 1, 3] {
            for late_error in [false, true] {
                let batches = number_batches();
                let schema = batches.schema();
                let mut input: Vec<_> = batches.into_iter().map(Ok).collect();
                if late_error {
                    input.push(
                        common_recordbatch::error::CreateRecordBatchesSnafu {
                            reason: "late stream error",
                        }
                        .fail(),
                    );
                }
                let expected_batches = input.len();
                let observed = Arc::new(AtomicUsize::new(0));
                let counter = observed.clone();
                let stream = stream::iter(input).inspect(move |_| {
                    counter.fetch_add(1, Ordering::Relaxed);
                });
                let output = Output::new_with_stream(Box::pin(RecordBatchStreamWrapper::new(
                    schema, stream,
                )));
                let response = sql_response(vec![Ok(output)], "greptimedb_v1", Some(limit)).await;
                assert_eq!(observed.load(Ordering::Relaxed), expected_batches);
                assert_eq!(matches!(response, HttpResponse::Error(_)), late_error);
            }
        }
    }

    #[tokio::test]
    async fn test_sql_response_limit_skips_conversion_of_discarded_rows() {
        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
            "payload",
            ConcreteDataType::json_datatype(),
            false,
        )]));
        let columns: Vec<VectorRef> = vec![Arc::new(BinaryVector::from(vec![
            datatypes::types::parse_string_to_jsonb(r#"{"ok":1}"#).unwrap(),
            b"invalid jsonb".to_vec(),
        ]))];
        let batch = RecordBatch::new(schema.clone(), columns).unwrap();
        for streaming in [false, true] {
            for limit in [0, 1, 2] {
                let batches = RecordBatches::try_new(schema.clone(), vec![batch.clone()]).unwrap();
                let output = if streaming {
                    Output::new_with_stream(batches.as_stream())
                } else {
                    Output::new_with_record_batches(batches)
                };
                let response = sql_response(vec![Ok(output)], "greptimedb_v1", Some(limit)).await;
                // Only retained rows should be decoded into HTTP response values.
                assert_eq!(matches!(response, HttpResponse::Error(_)), limit == 2);
            }
        }
    }

    #[test]
    fn test_final_analyze_event_uses_error_event_for_conversion_error() {
        let (event_name, payload) = make_final_analyze_event(
            Err((
                "conversion failed".to_string(),
                StatusCode::InvalidArguments as u32,
            )),
            7,
            Instant::now(),
            None,
        );

        assert_eq!(event_name, "error");
        let value: Value = serde_json::from_str(&payload).unwrap();
        assert_eq!(value["state"], "error");
        assert_eq!(value["reason"], "conversion failed");
    }
}