Skip to main content

servers/http/
handler.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16use std::panic::AssertUnwindSafe;
17use std::sync::Arc;
18use std::sync::atomic::{AtomicU64, Ordering};
19use std::time::{Duration, Instant};
20
21use axum::extract::rejection::FormRejection;
22use axum::extract::{Json, Query, State};
23use axum::response::sse::{Event, KeepAlive, Sse};
24use axum::response::{IntoResponse, Response};
25use axum::{Extension, Form};
26use common_catalog::parse_catalog_and_schema_from_db_string;
27use common_error::ext::ErrorExt;
28use common_error::status_code::StatusCode;
29use common_plugins::GREPTIME_EXEC_WRITE_COST;
30use common_query::{Output, OutputData};
31use common_recordbatch::util;
32use common_telemetry::tracing;
33use datafusion::physical_plan::ExecutionPlan;
34use futures::{FutureExt, StreamExt};
35use query::parser::{DEFAULT_LOOKBACK_STRING, PromQuery};
36use serde::{Deserialize, Serialize};
37use serde_json::Value;
38use session::context::{Channel, QueryContext, QueryContextRef};
39use snafu::ResultExt;
40use sql::dialect::GreptimeDbDialect;
41use sql::parser::{ParseOptions, ParserContext};
42use sql::statements::statement::Statement;
43use tokio::sync::{Notify, watch};
44
45use crate::error::{FailedToParseQuerySnafu, InvalidQuerySnafu, Result};
46use crate::http::header::collect_plan_metrics;
47use crate::http::result::arrow_result::ArrowResponse;
48use crate::http::result::csv_result::CsvResponse;
49use crate::http::result::error_result::ErrorResponse;
50use crate::http::result::greptime_result_v1::GreptimedbV1Response;
51use crate::http::result::influxdb_result_v1::InfluxdbV1Response;
52use crate::http::result::json_result::JsonResponse;
53use crate::http::result::null_result::NullResponse;
54use crate::http::result::table_result::TableResponse;
55use crate::http::{
56    ApiState, Epoch, GreptimeOptionsConfigState, GreptimeQueryOutput, HttpRecordsOutput,
57    HttpResponse, ResponseFormat,
58};
59use crate::metrics_handler::MetricsHandler;
60use crate::query_handler::sql::ServerSqlQueryHandlerRef;
61
62#[derive(Debug, Default, Serialize, Deserialize)]
63pub struct SqlQuery {
64    pub db: Option<String>,
65    pub sql: Option<String>,
66    // (Optional) result format: [`greptimedb_v1`, `influxdb_v1`, `csv`,
67    // `arrow`],
68    // the default value is `greptimedb_v1`
69    pub format: Option<String>,
70    // Returns epoch timestamps with the specified precision.
71    // Both u and µ indicate microseconds.
72    // epoch = [ns,u,µ,ms,s],
73    //
74    // TODO(jeremy): currently, only InfluxDB result format is supported,
75    // and all columns of the `Timestamp` type will be converted to their
76    // specified time precision. Maybe greptimedb format can support this
77    // param too.
78    pub epoch: Option<String>,
79    pub limit: Option<usize>,
80    // For arrow output
81    pub compression: Option<String>,
82    pub snapshot_interval_ms: Option<u64>,
83}
84
85const DEFAULT_ANALYZE_SNAPSHOT_INTERVAL_MS: u64 = 5000;
86const MIN_ANALYZE_SNAPSHOT_INTERVAL_MS: u64 = 1000;
87const MAX_ANALYZE_SNAPSHOT_INTERVAL_MS: u64 = 60000;
88
89#[derive(Serialize)]
90struct AnalyzeStreamPayload {
91    seq: u64,
92    state: &'static str,
93    partial: bool,
94    elapsed_ms: u64,
95    #[serde(skip_serializing_if = "Option::is_none")]
96    metrics: Option<Value>,
97    #[serde(skip_serializing_if = "Option::is_none")]
98    output: Option<GreptimeQueryOutput>,
99    #[serde(skip_serializing_if = "Option::is_none")]
100    reason: Option<String>,
101    #[serde(skip_serializing_if = "Option::is_none")]
102    code: Option<u32>,
103}
104
105#[derive(Clone, Debug)]
106#[doc(hidden)]
107pub struct AnalyzeStreamMessage {
108    pub event_name: &'static str,
109    pub payload: String,
110}
111
112struct AnalyzeStreamWorkerGuard {
113    cancel: watch::Sender<bool>,
114    handle: common_runtime::JoinHandle<()>,
115}
116
117impl Drop for AnalyzeStreamWorkerGuard {
118    fn drop(&mut self) {
119        let _ = self.cancel.send(true);
120        self.handle.abort();
121    }
122}
123
124struct AnalyzeStreamBodyState {
125    latest_metrics: watch::Receiver<Option<String>>,
126    terminal: watch::Receiver<Option<AnalyzeStreamMessage>>,
127    notify: Arc<Notify>,
128    // Keeping the guard in the body state makes dropping the response cancel and
129    // abort the worker instead of leaving an owned query stream detached.
130    _worker: AnalyzeStreamWorkerGuard,
131    done: bool,
132}
133
134/// Handler to execute sql
135#[axum_macros::debug_handler]
136#[tracing::instrument(skip_all, fields(protocol = "http", request_type = "sql"))]
137pub async fn sql(
138    State(state): State<ApiState>,
139    Query(query_params): Query<SqlQuery>,
140    Extension(mut query_ctx): Extension<QueryContext>,
141    Form(form_params): Form<SqlQuery>,
142) -> HttpResponse {
143    let start = Instant::now();
144    let sql_handler = &state.sql_handler;
145    if let Some(db) = &query_params.db.or(form_params.db) {
146        let (catalog, schema) = parse_catalog_and_schema_from_db_string(db);
147        query_ctx.set_current_catalog(&catalog);
148        query_ctx.set_current_schema(&schema);
149    }
150    let db = query_ctx.get_db_string();
151
152    query_ctx.set_channel(Channel::HttpSql);
153    let query_ctx = Arc::new(query_ctx);
154
155    let _timer = crate::metrics::METRIC_HTTP_SQL_ELAPSED
156        .with_label_values(&[db.as_str()])
157        .start_timer();
158
159    let sql = query_params.sql.or(form_params.sql);
160    let format = query_params
161        .format
162        .or(form_params.format)
163        .map(|s| s.to_lowercase())
164        .map(|s| ResponseFormat::parse(s.as_str()).unwrap_or(ResponseFormat::GreptimedbV1))
165        .unwrap_or(ResponseFormat::GreptimedbV1);
166    let epoch = query_params
167        .epoch
168        .or(form_params.epoch)
169        .map(|s| s.to_lowercase())
170        .map(|s| Epoch::parse(s.as_str()).unwrap_or(Epoch::Millisecond));
171
172    let result = if let Some(sql) = &sql {
173        if let Some((status, msg)) = validate_schema(sql_handler.clone(), query_ctx.clone()).await {
174            Err((status, msg))
175        } else {
176            Ok(sql_handler.do_query(sql, query_ctx.clone()).await)
177        }
178    } else {
179        Err((
180            StatusCode::InvalidArguments,
181            "sql parameter is required.".to_string(),
182        ))
183    };
184
185    let outputs = match result {
186        Err((status, msg)) => {
187            return HttpResponse::Error(
188                ErrorResponse::from_error_message(status, msg)
189                    .with_execution_time(start.elapsed().as_millis() as u64),
190            );
191        }
192        Ok(outputs) => outputs,
193    };
194
195    let mut resp = match format {
196        ResponseFormat::Arrow => {
197            ArrowResponse::from_output(outputs, query_params.compression).await
198        }
199        ResponseFormat::Csv(with_names, with_types) => {
200            CsvResponse::from_output(outputs, with_names, with_types).await
201        }
202        ResponseFormat::Table => TableResponse::from_output(outputs).await,
203        ResponseFormat::GreptimedbV1 => GreptimedbV1Response::from_output(outputs).await,
204        ResponseFormat::InfluxdbV1 => InfluxdbV1Response::from_output(outputs, epoch).await,
205        ResponseFormat::Json => JsonResponse::from_output(outputs).await,
206        ResponseFormat::Null => NullResponse::from_output(outputs).await,
207    };
208
209    if let Some(limit) = query_params.limit {
210        resp = resp.with_limit(limit);
211    }
212    resp.with_execution_time(start.elapsed().as_millis() as u64)
213}
214
215/// Handler to stream partial `EXPLAIN ANALYZE VERBOSE` metrics as SSE.
216///
217/// This endpoint is POST-only SSE, so browser `EventSource` does
218/// not apply. Each `metrics` event carries a complete snapshot (not a delta);
219/// large snapshots are throttled but never truncated. `final`, `canceled`, and
220/// `error` are terminal events. If the client disconnects it won't receive a
221/// `canceled` event, but the production frontend stream is dropped and
222/// best-effort cancels the underlying query.
223#[axum_macros::debug_handler]
224#[tracing::instrument(
225    skip_all,
226    fields(protocol = "http", request_type = "sql_analyze_stream")
227)]
228pub async fn sql_analyze_stream(
229    State(state): State<ApiState>,
230    Query(query_params): Query<SqlQuery>,
231    Extension(mut query_ctx): Extension<QueryContext>,
232    form_params: std::result::Result<Form<SqlQuery>, FormRejection>,
233) -> Response {
234    let start = Instant::now();
235    let form_params = match form_params {
236        Ok(Form(params)) => params,
237        Err(err) => {
238            if err.status() != axum::http::StatusCode::UNSUPPORTED_MEDIA_TYPE {
239                return ErrorResponse::from_error_message(
240                    StatusCode::InvalidArguments,
241                    err.body_text(),
242                )
243                .with_execution_time(start.elapsed().as_millis() as u64)
244                .into_response();
245            }
246            SqlQuery::default()
247        }
248    };
249    let sql_handler = &state.sql_handler;
250    if let Some(db) = &query_params.db.or(form_params.db) {
251        let (catalog, schema) = parse_catalog_and_schema_from_db_string(db);
252        query_ctx.set_current_catalog(&catalog);
253        query_ctx.set_current_schema(&schema);
254    }
255    query_ctx.set_channel(Channel::HttpSql);
256    query_ctx.enable_live_analyze_metrics();
257    let query_ctx = Arc::new(query_ctx);
258
259    let Some(sql) = query_params.sql.or(form_params.sql) else {
260        return ErrorResponse::from_error_message(
261            StatusCode::InvalidArguments,
262            "sql parameter is required.".to_string(),
263        )
264        .with_execution_time(start.elapsed().as_millis() as u64)
265        .into_response();
266    };
267    if let Some((status, msg)) = validate_schema(sql_handler.clone(), query_ctx.clone()).await {
268        return ErrorResponse::from_error_message(status, msg)
269            .with_execution_time(start.elapsed().as_millis() as u64)
270            .into_response();
271    }
272
273    let interval_ms = query_params
274        .snapshot_interval_ms
275        .or(form_params.snapshot_interval_ms)
276        .unwrap_or(DEFAULT_ANALYZE_SNAPSHOT_INTERVAL_MS)
277        .clamp(
278            MIN_ANALYZE_SNAPSHOT_INTERVAL_MS,
279            MAX_ANALYZE_SNAPSHOT_INTERVAL_MS,
280        );
281
282    let output = match state
283        .sql_handler
284        .do_analyze_stream_query(&sql, query_ctx.clone())
285        .await
286    {
287        Ok(output) => output,
288        Err(err) => {
289            return ErrorResponse::from_error(err)
290                .with_execution_time(start.elapsed().as_millis() as u64)
291                .into_response();
292        }
293    };
294
295    let plan = output.meta.plan.clone();
296    let OutputData::Stream(stream) = output.data else {
297        return ErrorResponse::from_error_message(
298            StatusCode::InvalidArguments,
299            "analyze stream query must return a stream".to_string(),
300        )
301        .with_execution_time(start.elapsed().as_millis() as u64)
302        .into_response();
303    };
304    let schema = stream.schema();
305
306    let (metrics_tx, metrics_rx) = watch::channel::<Option<String>>(None);
307    let (terminal_tx, terminal_rx) = watch::channel::<Option<AnalyzeStreamMessage>>(None);
308    let notify = Arc::new(Notify::new());
309    let (cancel_tx, mut cancel_rx) = watch::channel(false);
310    let worker_notify = notify.clone();
311    let panic_notify = notify.clone();
312    let sequence = Arc::new(AtomicU64::new(0));
313    let worker_sequence = sequence.clone();
314    let worker_terminal_tx = terminal_tx.clone();
315    let worker = common_runtime::spawn_global(async move {
316        let worker_result = AssertUnwindSafe(async move {
317            let mut stream = stream;
318            let mut batches = Vec::new();
319            let mut current_interval_ms = interval_ms;
320            let tick = tokio::time::sleep(Duration::from_millis(current_interval_ms));
321            tokio::pin!(tick);
322
323            loop {
324                tokio::select! {
325                    _ = cancel_rx.changed() => return,
326                    item = stream.next() => {
327                        match item {
328                            Some(Ok(next_batch)) => batches.push(next_batch),
329                            Some(Err(err)) => {
330                                let status = err.status_code();
331                                let event_name = if status == StatusCode::Cancelled { "canceled" } else { "error" };
332                                let (payload, _) = make_analyze_payload(AnalyzePayloadArgs {
333                                    seq: worker_sequence.load(Ordering::Relaxed),
334                                    state: event_name,
335                                    partial: false,
336                                    start,
337                                    plan: plan.as_ref(),
338                                    output: None,
339                                    reason: Some(err.output_msg()),
340                                    code: Some(status as u32),
341                                });
342                                send_analyze_terminal(&terminal_tx, &worker_notify, event_name, payload);
343                                return;
344                            }
345                            None => {
346                                let output = HttpRecordsOutput::try_new(schema.clone(), batches)
347                                    .map(GreptimeQueryOutput::Records);
348                                let (event_name, payload) = make_final_analyze_event(
349                                    output.map_err(|err| (err.output_msg(), err.status_code() as u32)),
350                                    worker_sequence.load(Ordering::Relaxed),
351                                    start,
352                                    plan.as_ref(),
353                                );
354                                send_analyze_terminal(&terminal_tx, &worker_notify, event_name, payload);
355                                return;
356                            }
357                        }
358                    }
359                    _ = &mut tick, if plan.is_some() => {
360                        let (payload, payload_bytes) = make_analyze_payload(AnalyzePayloadArgs {
361                            seq: worker_sequence.load(Ordering::Relaxed),
362                            state: "metrics",
363                            partial: true,
364                            start,
365                            plan: plan.as_ref(),
366                            output: None,
367                            reason: None,
368                            code: None,
369                        });
370                        current_interval_ms = adaptive_interval_ms(payload_bytes, interval_ms);
371                        worker_sequence.fetch_add(1, Ordering::Relaxed);
372                        if metrics_tx.send(Some(payload)).is_err() {
373                            return;
374                        }
375                        worker_notify.notify_one();
376                        tick.as_mut().reset(tokio::time::Instant::now() + Duration::from_millis(current_interval_ms));
377                    }
378                }
379            }
380        })
381        .catch_unwind()
382        .await;
383
384        if worker_result.is_err() {
385            tracing::debug!("analyze stream worker panicked");
386            let (payload, _) = make_analyze_payload(AnalyzePayloadArgs {
387                seq: sequence.load(Ordering::Relaxed),
388                state: "error",
389                partial: false,
390                start,
391                plan: None,
392                output: None,
393                reason: Some("analyze stream worker panicked".to_string()),
394                code: Some(StatusCode::Internal as u32),
395            });
396            send_analyze_terminal(&worker_terminal_tx, &panic_notify, "error", payload);
397        }
398    });
399
400    let sse_stream = analyze_stream_body(metrics_rx, terminal_rx, notify, cancel_tx, worker);
401
402    Sse::new(sse_stream)
403        .keep_alive(KeepAlive::new().interval(Duration::from_secs(15)))
404        .into_response()
405}
406
407#[doc(hidden)]
408pub fn analyze_stream_body(
409    metrics: watch::Receiver<Option<String>>,
410    terminal: watch::Receiver<Option<AnalyzeStreamMessage>>,
411    notify: Arc<Notify>,
412    cancel: watch::Sender<bool>,
413    worker: common_runtime::JoinHandle<()>,
414) -> impl futures::Stream<Item = std::result::Result<Event, std::convert::Infallible>> {
415    futures::stream::unfold(
416        AnalyzeStreamBodyState {
417            latest_metrics: metrics,
418            terminal,
419            notify,
420            _worker: AnalyzeStreamWorkerGuard {
421                cancel,
422                handle: worker,
423            },
424            done: false,
425        },
426        |mut state| async move {
427            if state.done {
428                return None;
429            }
430            loop {
431                let notify = Arc::clone(&state.notify);
432                let notified = notify.notified();
433                tokio::pin!(notified);
434                // Register before checking the slots to avoid a lost wakeup.
435                notified.as_mut().enable();
436                let latest = {
437                    let latest = state.latest_metrics.borrow_and_update();
438                    latest.has_changed().then(|| latest.clone())
439                };
440                if let Some(Some(payload)) = latest {
441                    return Some((Ok(Event::default().event("metrics").data(payload)), state));
442                }
443                // Inspect the terminal slot directly because a closed watch channel
444                // can otherwise hide a value published while a worker was unwinding.
445                if state.terminal.borrow().is_some() {
446                    let terminal = { state.terminal.borrow_and_update().clone() };
447                    if let Some(AnalyzeStreamMessage {
448                        event_name,
449                        payload,
450                    }) = terminal
451                    {
452                        state.done = true;
453                        return Some((Ok(Event::default().event(event_name).data(payload)), state));
454                    }
455                }
456                notified.await;
457            }
458        },
459    )
460}
461
462#[doc(hidden)]
463pub fn send_analyze_terminal(
464    terminal_tx: &watch::Sender<Option<AnalyzeStreamMessage>>,
465    notify: &Notify,
466    event_name: &'static str,
467    payload: String,
468) {
469    if terminal_tx.send_if_modified(|terminal| {
470        if terminal.is_none() {
471            *terminal = Some(AnalyzeStreamMessage {
472                event_name,
473                payload,
474            });
475            true
476        } else {
477            false
478        }
479    }) {
480        notify.notify_one();
481    }
482}
483
484fn adaptive_interval_ms(payload_bytes: usize, requested_ms: u64) -> u64 {
485    if payload_bytes >= 10 * 1024 * 1024 {
486        requested_ms.max(30_000)
487    } else if payload_bytes >= 1024 * 1024 {
488        requested_ms.max(10_000)
489    } else {
490        requested_ms
491    }
492}
493
494fn make_final_analyze_event(
495    output: std::result::Result<GreptimeQueryOutput, (String, u32)>,
496    seq: u64,
497    start: Instant,
498    plan: Option<&Arc<dyn ExecutionPlan>>,
499) -> (&'static str, String) {
500    match output {
501        Ok(output) => (
502            "final",
503            make_analyze_payload(AnalyzePayloadArgs {
504                seq,
505                state: "final",
506                partial: false,
507                start,
508                plan,
509                output: Some(output),
510                reason: None,
511                code: None,
512            })
513            .0,
514        ),
515        Err((reason, code)) => (
516            "error",
517            make_analyze_payload(AnalyzePayloadArgs {
518                seq,
519                state: "error",
520                partial: false,
521                start,
522                plan,
523                output: None,
524                reason: Some(reason),
525                code: Some(code),
526            })
527            .0,
528        ),
529    }
530}
531
532struct AnalyzePayloadArgs<'a> {
533    seq: u64,
534    state: &'static str,
535    partial: bool,
536    start: Instant,
537    plan: Option<&'a Arc<dyn ExecutionPlan>>,
538    output: Option<GreptimeQueryOutput>,
539    reason: Option<String>,
540    code: Option<u32>,
541}
542
543fn make_analyze_payload(args: AnalyzePayloadArgs<'_>) -> (String, usize) {
544    let AnalyzePayloadArgs {
545        seq,
546        state,
547        partial,
548        start,
549        plan,
550        output,
551        reason,
552        code,
553    } = args;
554    // Periodic snapshots are compact; terminal snapshots retain verbose plan details.
555    let metrics =
556        plan.and_then(|plan| query::analyze_plan_metrics_to_json_value(plan, !partial).ok());
557    let payload = AnalyzeStreamPayload {
558        seq,
559        state,
560        partial,
561        elapsed_ms: start.elapsed().as_millis() as u64,
562        metrics,
563        output,
564        reason,
565        code,
566    };
567    let payload_string = serde_json::to_string(&payload).unwrap_or_else(|e| {
568        serde_json::json!({
569            "seq": seq,
570            "state": "error",
571            "partial": false,
572            "reason": format!("Failed to serialize SSE payload: {e}"),
573        })
574        .to_string()
575    });
576    let payload_bytes = payload_string.len();
577    (payload_string, payload_bytes)
578}
579
580/// Handler to parse sql
581#[axum_macros::debug_handler]
582#[tracing::instrument(skip_all, fields(protocol = "http", request_type = "sql"))]
583pub async fn sql_parse(
584    Query(query_params): Query<SqlQuery>,
585    Form(form_params): Form<SqlQuery>,
586) -> Result<Json<Vec<Statement>>> {
587    let Some(sql) = query_params.sql.or(form_params.sql) else {
588        return InvalidQuerySnafu {
589            reason: "sql parameter is required.",
590        }
591        .fail();
592    };
593
594    let stmts =
595        ParserContext::create_with_dialect(&sql, &GreptimeDbDialect {}, ParseOptions::default())
596            .context(FailedToParseQuerySnafu)?;
597
598    Ok(stmts.into())
599}
600
601#[derive(Debug, Serialize, Deserialize)]
602pub struct SqlFormatResponse {
603    pub formatted: String,
604}
605
606/// Handler to format sql string
607#[axum_macros::debug_handler]
608#[tracing::instrument(skip_all, fields(protocol = "http", request_type = "sql_format"))]
609pub async fn sql_format(
610    Query(query_params): Query<SqlQuery>,
611    Form(form_params): Form<SqlQuery>,
612) -> axum::response::Response {
613    let Some(sql) = query_params.sql.or(form_params.sql) else {
614        let resp = ErrorResponse::from_error_message(
615            StatusCode::InvalidArguments,
616            "sql parameter is required.".to_string(),
617        );
618        return HttpResponse::Error(resp).into_response();
619    };
620
621    // Parse using GreptimeDB dialect then reconstruct statements via Display
622    let stmts = match ParserContext::create_with_dialect(
623        &sql,
624        &GreptimeDbDialect {},
625        ParseOptions::default(),
626    ) {
627        Ok(v) => v,
628        Err(e) => return HttpResponse::Error(ErrorResponse::from_error(e)).into_response(),
629    };
630
631    let mut parts: Vec<String> = Vec::with_capacity(stmts.len());
632    for stmt in stmts {
633        let mut s = format!("{stmt}");
634        if !s.trim_end().ends_with(';') {
635            s.push(';');
636        }
637        parts.push(s);
638    }
639
640    let formatted = parts.join("\n");
641    Json(SqlFormatResponse { formatted }).into_response()
642}
643
644/// Create a response from query result
645pub async fn from_output(
646    outputs: Vec<crate::error::Result<Output>>,
647) -> std::result::Result<(Vec<GreptimeQueryOutput>, HashMap<String, Value>), ErrorResponse> {
648    // TODO(sunng87): this api response structure cannot represent error well.
649    //  It hides successful execution results from error response
650    let mut results = Vec::with_capacity(outputs.len());
651    let mut merge_map = HashMap::new();
652
653    for out in outputs {
654        match out {
655            Ok(o) => match o.data {
656                OutputData::AffectedRows(rows) => {
657                    results.push(GreptimeQueryOutput::AffectedRows(rows));
658                    if o.meta.cost > 0 {
659                        merge_map.insert(GREPTIME_EXEC_WRITE_COST.to_string(), o.meta.cost as u64);
660                    }
661                }
662                OutputData::Stream(stream) => {
663                    let schema = stream.schema().clone();
664                    // TODO(sunng87): streaming response
665                    let mut http_record_output = match util::collect(stream).await {
666                        Ok(rows) => match HttpRecordsOutput::try_new(schema, rows) {
667                            Ok(rows) => rows,
668                            Err(err) => {
669                                return Err(ErrorResponse::from_error(err));
670                            }
671                        },
672                        Err(err) => {
673                            return Err(ErrorResponse::from_error(err));
674                        }
675                    };
676                    if let Some(physical_plan) = o.meta.plan {
677                        let mut result_map = HashMap::new();
678
679                        let mut tmp = vec![&mut merge_map, &mut result_map];
680                        collect_plan_metrics(&physical_plan, &mut tmp);
681                        let re = result_map
682                            .into_iter()
683                            .map(|(k, v)| (k, Value::from(v)))
684                            .collect::<HashMap<String, Value>>();
685                        http_record_output.metrics.extend(re);
686                    }
687                    results.push(GreptimeQueryOutput::Records(http_record_output))
688                }
689                OutputData::RecordBatches(rbs) => {
690                    match HttpRecordsOutput::try_new(rbs.schema(), rbs.take()) {
691                        Ok(rows) => {
692                            results.push(GreptimeQueryOutput::Records(rows));
693                        }
694                        Err(err) => {
695                            return Err(ErrorResponse::from_error(err));
696                        }
697                    }
698                }
699            },
700
701            Err(err) => {
702                return Err(ErrorResponse::from_error(err));
703            }
704        }
705    }
706
707    let merge_map = merge_map
708        .into_iter()
709        .map(|(k, v)| (k, Value::from(v)))
710        .collect();
711
712    Ok((results, merge_map))
713}
714
715#[derive(Debug, Default, Serialize, Deserialize)]
716pub struct PromqlQuery {
717    pub query: String,
718    pub start: String,
719    pub end: String,
720    pub step: String,
721    pub lookback: Option<String>,
722    pub db: Option<String>,
723    // (Optional) result format: [`greptimedb_v1`, `influxdb_v1`, `csv`,
724    // `arrow`],
725    // the default value is `greptimedb_v1`
726    pub format: Option<String>,
727    // For arrow output
728    pub compression: Option<String>,
729    // Returns epoch timestamps with the specified precision.
730    // Both u and µ indicate microseconds.
731    // epoch = [ns,u,µ,ms,s],
732    //
733    // For influx output only
734    //
735    // TODO(jeremy): currently, only InfluxDB result format is supported,
736    // and all columns of the `Timestamp` type will be converted to their
737    // specified time precision. Maybe greptimedb format can support this
738    // param too.
739    pub epoch: Option<String>,
740}
741
742impl From<PromqlQuery> for PromQuery {
743    fn from(query: PromqlQuery) -> Self {
744        PromQuery {
745            query: query.query,
746            start: query.start,
747            end: query.end,
748            step: query.step,
749            lookback: query
750                .lookback
751                .unwrap_or_else(|| DEFAULT_LOOKBACK_STRING.to_string()),
752            // TODO(dennis): support alias from http params or parse from query.query
753            alias: None,
754        }
755    }
756}
757
758/// Handler to execute promql
759#[axum_macros::debug_handler]
760#[tracing::instrument(skip_all, fields(protocol = "http", request_type = "promql"))]
761pub async fn promql(
762    State(state): State<ApiState>,
763    Query(params): Query<PromqlQuery>,
764    Extension(mut query_ctx): Extension<QueryContext>,
765) -> Response {
766    let sql_handler = &state.sql_handler;
767    let exec_start = Instant::now();
768    let db = query_ctx.get_db_string();
769
770    query_ctx.set_channel(Channel::Promql);
771    let query_ctx = Arc::new(query_ctx);
772
773    let _timer = crate::metrics::METRIC_HTTP_PROMQL_ELAPSED
774        .with_label_values(&[db.as_str()])
775        .start_timer();
776
777    let resp = if let Some((status, msg)) =
778        validate_schema(sql_handler.clone(), query_ctx.clone()).await
779    {
780        let resp = ErrorResponse::from_error_message(status, msg);
781        HttpResponse::Error(resp)
782    } else {
783        let format = params
784            .format
785            .as_ref()
786            .map(|s| s.to_lowercase())
787            .map(|s| ResponseFormat::parse(s.as_str()).unwrap_or(ResponseFormat::GreptimedbV1))
788            .unwrap_or(ResponseFormat::GreptimedbV1);
789        let epoch = params
790            .epoch
791            .as_ref()
792            .map(|s| s.to_lowercase())
793            .map(|s| Epoch::parse(s.as_str()).unwrap_or(Epoch::Millisecond));
794        let compression = params.compression.clone();
795
796        let prom_query = params.into();
797        let outputs = sql_handler.do_promql_query(&prom_query, query_ctx).await;
798
799        match format {
800            ResponseFormat::Arrow => ArrowResponse::from_output(outputs, compression).await,
801            ResponseFormat::Csv(with_names, with_types) => {
802                CsvResponse::from_output(outputs, with_names, with_types).await
803            }
804            ResponseFormat::Table => TableResponse::from_output(outputs).await,
805            ResponseFormat::GreptimedbV1 => GreptimedbV1Response::from_output(outputs).await,
806            ResponseFormat::InfluxdbV1 => InfluxdbV1Response::from_output(outputs, epoch).await,
807            ResponseFormat::Json => JsonResponse::from_output(outputs).await,
808            ResponseFormat::Null => NullResponse::from_output(outputs).await,
809        }
810    };
811
812    resp.with_execution_time(exec_start.elapsed().as_millis() as u64)
813        .into_response()
814}
815
816/// Handler to export metrics
817#[axum_macros::debug_handler]
818pub async fn metrics(
819    State(state): State<MetricsHandler>,
820    Query(_params): Query<HashMap<String, String>>,
821) -> String {
822    // A default ProcessCollector is registered automatically in prometheus.
823    // We do not need to explicitly collect process-related data.
824    // But ProcessCollector only support on linux.
825
826    #[cfg(not(windows))]
827    if let Some(c) = crate::metrics::jemalloc::JEMALLOC_COLLECTOR.as_ref()
828        && let Err(e) = c.update()
829    {
830        common_telemetry::error!(e; "Failed to update jemalloc metrics");
831    }
832    state.render()
833}
834
835#[derive(Debug, Serialize, Deserialize)]
836pub struct HealthQuery {}
837
838#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
839pub struct HealthResponse {}
840
841/// Handler to export healthy check
842///
843/// Currently simply return status "200 OK" (default) with an empty json payload "{}"
844#[axum_macros::debug_handler]
845pub async fn health(Query(_params): Query<HealthQuery>) -> Json<HealthResponse> {
846    Json(HealthResponse {})
847}
848
849#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
850pub struct StatusResponse<'a> {
851    pub commit: &'a str,
852    pub branch: &'a str,
853    pub rustc_version: &'a str,
854    pub hostname: String,
855    pub version: &'a str,
856}
857
858/// Handler to expose information info about runtime, build, etc.
859#[axum_macros::debug_handler]
860pub async fn status() -> Json<StatusResponse<'static>> {
861    let hostname = hostname::get()
862        .map(|s| s.to_string_lossy().to_string())
863        .unwrap_or_else(|_| "unknown".to_string());
864    let build_info = common_version::build_info();
865    Json(StatusResponse {
866        commit: build_info.commit,
867        branch: build_info.branch,
868        rustc_version: build_info.rustc,
869        hostname,
870        version: build_info.version,
871    })
872}
873
874/// Handler to expose configuration information info about runtime, build, etc.
875#[axum_macros::debug_handler]
876pub async fn config(State(state): State<GreptimeOptionsConfigState>) -> Response {
877    (axum::http::StatusCode::OK, state.greptime_config_options).into_response()
878}
879
880async fn validate_schema(
881    sql_handler: ServerSqlQueryHandlerRef,
882    query_ctx: QueryContextRef,
883) -> Option<(StatusCode, String)> {
884    match sql_handler
885        .is_valid_schema(query_ctx.current_catalog(), &query_ctx.current_schema())
886        .await
887    {
888        Ok(true) => None,
889        Ok(false) => Some((
890            StatusCode::DatabaseNotFound,
891            format!("Database not found: {}", query_ctx.get_db_string()),
892        )),
893        Err(e) => Some((
894            StatusCode::Internal,
895            format!(
896                "Error checking database: {}, {}",
897                query_ctx.get_db_string(),
898                e.output_msg(),
899            ),
900        )),
901    }
902}
903
904pub async fn index() -> axum::response::Html<String> {
905    let name = common_version::product_name();
906    let version = common_version::version();
907    axum::response::Html(format!(
908        r#"<!DOCTYPE html>
909<html>
910<head><title>{name}</title></head>
911<body>
912<h1>{name}</h1>
913<p>Version: {version}</p>
914<ul>
915<li><a href="/dashboard">Dashboard UI</a></li>
916<li><a href="/v1/health">Health</a> (JSON)</li>
917<li><a href="/status">Status</a> (JSON)</li>
918<li><a href="/metrics">Metrics</a> (For Prometheus Scrape)</li>
919<li><a href="/config">Config</a> (TXT)</li>
920</ul>
921</body>
922</html>"#,
923    ))
924}
925
926#[cfg(test)]
927mod tests {
928    use super::*;
929
930    #[test]
931    fn test_final_analyze_event_uses_error_event_for_conversion_error() {
932        let (event_name, payload) = make_final_analyze_event(
933            Err((
934                "conversion failed".to_string(),
935                StatusCode::InvalidArguments as u32,
936            )),
937            7,
938            Instant::now(),
939            None,
940        );
941
942        assert_eq!(event_name, "error");
943        let value: Value = serde_json::from_str(&payload).unwrap();
944        assert_eq!(value["state"], "error");
945        assert_eq!(value["reason"], "conversion failed");
946    }
947}