Skip to main content

servers/
http.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::convert::Infallible;
17use std::fmt::Display;
18use std::net::SocketAddr;
19use std::sync::{Arc, Mutex as StdMutex};
20use std::time::Duration;
21
22use async_trait::async_trait;
23use auth::UserProviderRef;
24use axum::extract::{DefaultBodyLimit, Request, State};
25use axum::http::StatusCode as HttpStatusCode;
26use axum::middleware::Next;
27use axum::response::{IntoResponse, Response};
28use axum::routing::Route;
29use axum::serve::ListenerExt;
30use axum::{Extension, Router, middleware, routing};
31use common_base::readable_size::ReadableSize;
32use common_recordbatch::RecordBatch;
33use common_telemetry::{error, info};
34use common_time::Timestamp;
35use common_time::timestamp::TimeUnit;
36use datatypes::data_type::DataType;
37use datatypes::schema::SchemaRef;
38use event::{LogState, LogValidatorRef};
39use futures::FutureExt;
40use http::{HeaderValue, Method};
41use serde::{Deserialize, Serialize};
42use serde_json::Value;
43use session::context::QueryContext;
44use snafu::{ResultExt, ensure};
45use tokio::sync::Mutex;
46use tokio::sync::oneshot::{self, Sender};
47use tonic::codegen::Service;
48use tower::{Layer, ServiceBuilder};
49use tower_http::compression::CompressionLayer;
50use tower_http::cors::{AllowOrigin, Any, CorsLayer};
51use tower_http::decompression::RequestDecompressionLayer;
52use tower_http::trace::TraceLayer;
53
54use self::authorize::AuthState;
55use self::result::table_result::TableResponse;
56use crate::batcher::BatchingProtocol;
57use crate::batcher::logical_table::LogicalTablePendingRowsBatcher;
58use crate::elasticsearch;
59use crate::error::{
60    AddressBindSnafu, AlreadyStartedSnafu, Error, InternalIoSnafu, InvalidHeaderValueSnafu, Result,
61};
62use crate::http::influxdb::{influxdb_health, influxdb_ping, influxdb_write_v1, influxdb_write_v2};
63use crate::http::otlp::OtlpState;
64use crate::http::prom_store::PromStoreState;
65use crate::http::prometheus::{
66    build_info_query, format_query, instant_query, label_values_query, labels_query,
67    metadata_query, parse_query, range_query, series_query,
68};
69use crate::http::result::arrow_result::ArrowResponse;
70use crate::http::result::csv_result::CsvResponse;
71use crate::http::result::error_result::ErrorResponse;
72use crate::http::result::greptime_result_v1::GreptimedbV1Response;
73use crate::http::result::influxdb_result_v1::InfluxdbV1Response;
74use crate::http::result::json_result::JsonResponse;
75use crate::http::result::null_result::NullResponse;
76use crate::interceptor::LogIngestInterceptorRef;
77use crate::metrics::http_metrics_layer;
78use crate::metrics_handler::MetricsHandler;
79use crate::prometheus_handler::PrometheusHandlerRef;
80use crate::query_handler::sql::ServerSqlQueryHandlerRef;
81use crate::query_handler::{
82    DashboardHandlerRef, InfluxdbLineProtocolHandlerRef, JaegerQueryHandlerRef, LogQueryHandlerRef,
83    OpenTelemetryProtocolHandlerRef, OpentsdbProtocolHandlerRef, PipelineHandlerRef,
84    PromStoreProtocolHandlerRef,
85};
86use crate::request_memory_limiter::ServerMemoryLimiter;
87use crate::server::Server;
88
89pub mod authorize;
90#[cfg(feature = "dashboard")]
91mod dashboard;
92pub mod dyn_log;
93pub mod dyn_trace;
94pub mod event;
95pub mod extractor;
96pub mod handler;
97pub mod header;
98pub mod influxdb;
99pub mod jaeger;
100pub mod logs;
101pub mod loki;
102pub mod mem_prof;
103mod memory_limit;
104pub mod opentsdb;
105pub mod otlp;
106pub mod pprof;
107pub mod prom_store;
108pub mod prometheus;
109pub mod result;
110pub mod splunk;
111mod timeout;
112pub mod utils;
113mod workload_scheduler;
114
115pub use memory_limit::decoded_body_accounting_middleware;
116use result::HttpOutputWriter;
117pub(crate) use timeout::DynamicTimeoutLayer;
118
119mod client_ip;
120use crate::prom_remote_write::validation::PromValidationMode;
121mod hints;
122mod read_preference;
123mod skip_wal;
124#[cfg(any(test, feature = "testing"))]
125pub mod test_helpers;
126
127pub const HTTP_API_VERSION: &str = "v1";
128pub const HTTP_API_PREFIX: &str = "/v1/";
129pub const HTTP_API_PREFIX_WITHOUT_TRAILING_SLASH: &str = "/v1";
130
131/// Provides an extra router to merge into an HTTP server.
132pub trait ExtraHttpRouterProvider: Send + Sync {
133    /// Returns the extra HTTP router.
134    fn router(&self) -> Router;
135}
136
137pub type ExtraHttpRouterProviderRef = Arc<dyn ExtraHttpRouterProvider>;
138
139/// Collection of extra HTTP router providers.
140#[derive(Clone, Default)]
141pub struct ExtraHttpRouterProviders {
142    providers: Vec<ExtraHttpRouterProviderRef>,
143}
144
145impl ExtraHttpRouterProviders {
146    /// Creates an empty provider collection.
147    pub fn new() -> Self {
148        Self::default()
149    }
150
151    /// Adds an extra HTTP router provider.
152    pub fn add(&mut self, provider: ExtraHttpRouterProviderRef) {
153        self.providers.push(provider);
154    }
155
156    /// Returns all registered providers.
157    pub fn iter(&self) -> impl Iterator<Item = &dyn ExtraHttpRouterProvider> {
158        self.providers.iter().map(|x| x.as_ref())
159    }
160}
161
162/// Default http body limit (64M).
163const DEFAULT_BODY_LIMIT: ReadableSize = ReadableSize::mb(64);
164/// Default address port for the public HTTP API server.
165const DEFAULT_HTTP_API_ADDR_PORT: u16 = 4006;
166
167/// Authorization header
168pub const AUTHORIZATION_HEADER: &str = "x-greptime-auth";
169
170// TODO(fys): This is a temporary workaround, it will be improved later
171// these APIs will bypass authentication.
172pub static PUBLIC_API_PREFIX: [&str; 4] = [
173    "/v1/influxdb/ping",
174    "/v1/influxdb/health",
175    "/v1/health",
176    "/v1/splunk/services/collector/health",
177];
178
179/// Listener mode for an [`HttpServer`].
180///
181/// Both modes share the *same* complete route registry; the [`HttpServerKind::Api`]
182/// listener additionally installs an outer guard that only exposes the `/v1` APIs
183/// and the dashboard, returning `404` for every other path.
184#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
185pub(crate) enum HttpServerKind {
186    /// Serves the complete historical HTTP surface (both `v1` APIs and internal
187    /// interfaces such as health, status, metrics, config and debug).
188    #[default]
189    Full,
190    /// Serves only the `v1` interfaces plus the dashboard.
191    Api,
192}
193
194#[derive(Default)]
195pub struct HttpServer {
196    router: StdMutex<Router>,
197    shutdown_tx: Mutex<Option<Sender<()>>>,
198    user_provider: Option<UserProviderRef>,
199    memory_limiter: ServerMemoryLimiter,
200
201    // server configs
202    options: HttpOptions,
203    batching_protocols: Vec<BatchingProtocol>,
204    logical_batching_protocols: Vec<BatchingProtocol>,
205    bind_addr: Option<SocketAddr>,
206    /// What this server instance exposes. See [`HttpServerKind`].
207    kind: HttpServerKind,
208}
209
210/// Returns `true` when `path` equals `root` or is nested directly under it
211/// (`root/...`). Uses an exact slash boundary so `/v1` does **not** match `/v10`
212/// or `/v1-internal`.
213pub(crate) fn is_namespace(path: &str, root: &str) -> bool {
214    path == root
215        || path
216            .strip_prefix(root)
217            .is_some_and(|suffix| suffix.starts_with('/'))
218}
219
220/// Paths visible on the dedicated API listener: the `/v1` APIs and the dashboard.
221pub fn is_api_listener_path(path: &str) -> bool {
222    is_namespace(path, HTTP_API_PREFIX_WITHOUT_TRAILING_SLASH) || is_namespace(path, "/dashboard")
223}
224
225#[derive(Clone)]
226struct LogicalBatchingProtocols(Vec<BatchingProtocol>);
227
228/// Sets a local-only write selector after authentication creates the context.
229pub async fn set_http_write_batching(
230    State(protocol): State<BatchingProtocol>,
231    mut req: Request,
232    next: Next,
233) -> Response {
234    let enabled = req
235        .extensions()
236        .get::<Arc<Vec<BatchingProtocol>>>()
237        .is_some_and(|protocols| protocols.contains(&protocol));
238    let logical_enabled = req
239        .extensions()
240        .get::<LogicalBatchingProtocols>()
241        .is_some_and(|protocols| protocols.0.contains(&protocol));
242    if let Some(ctx) = req.extensions_mut().get_mut::<QueryContext>() {
243        ctx.set_batching_enabled(enabled);
244        ctx.set_logical_batching_enabled(logical_enabled);
245    }
246    next.run(req).await
247}
248
249/// Outer guard for the API listener. Rejects any path outside the API surface
250/// with `404 Not Found` before it can reach a handler (and before auth side
251/// effects). Reachability is kept separate from authentication.
252async fn enforce_api_surface(req: Request, next: Next) -> Response {
253    if !is_api_listener_path(req.uri().path()) {
254        return HttpStatusCode::NOT_FOUND.into_response();
255    }
256    next.run(req).await
257}
258
259impl HttpServer {
260    /// A short human-readable label for this server instance, used in logs.
261    fn kind(&self) -> &'static str {
262        match self.kind {
263            HttpServerKind::Api => "HTTP API",
264            HttpServerKind::Full => "HTTP",
265        }
266    }
267}
268
269#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
270#[serde(default)]
271pub struct HttpOptions {
272    pub addr: String,
273
274    /// Request timeout; zero disables it. Frontend raises a nonzero timeout to at least
275    /// the largest active synchronous Prom or shared table batch flush interval plus one second.
276    #[serde(with = "humantime_serde")]
277    pub timeout: Duration,
278
279    #[serde(skip)]
280    pub disable_dashboard: bool,
281
282    pub body_limit: ReadableSize,
283
284    pub cors_allowed_origins: Vec<String>,
285
286    pub enable_cors: bool,
287
288    /// Whether to start the dedicated public HTTP **API** server, which serves
289    /// only the `v1` interfaces plus the dashboard. It shares every other
290    /// `[http]` option with the main server and only differs by its bound
291    /// address (see `api_server_addr`).
292    pub enable_api_server: bool,
293    /// The address to bind the dedicated HTTP API server, in the same form as
294    /// `addr` (e.g. `127.0.0.1:4006`).
295    pub api_server_addr: String,
296}
297
298impl Default for HttpOptions {
299    fn default() -> Self {
300        Self {
301            addr: "127.0.0.1:4000".to_string(),
302            timeout: Duration::from_secs(0),
303            disable_dashboard: false,
304            body_limit: DEFAULT_BODY_LIMIT,
305            cors_allowed_origins: Vec::new(),
306            enable_cors: true,
307            enable_api_server: false,
308            api_server_addr: format!("127.0.0.1:{}", DEFAULT_HTTP_API_ADDR_PORT),
309        }
310    }
311}
312
313#[derive(Debug, Serialize, Deserialize, Eq, PartialEq)]
314pub struct ColumnSchema {
315    name: String,
316    data_type: String,
317}
318
319impl ColumnSchema {
320    pub fn new(name: String, data_type: String) -> ColumnSchema {
321        ColumnSchema { name, data_type }
322    }
323}
324
325#[derive(Debug, Serialize, Deserialize, Eq, PartialEq)]
326pub struct OutputSchema {
327    column_schemas: Vec<ColumnSchema>,
328}
329
330impl OutputSchema {
331    pub fn new(columns: Vec<ColumnSchema>) -> OutputSchema {
332        OutputSchema {
333            column_schemas: columns,
334        }
335    }
336}
337
338impl From<SchemaRef> for OutputSchema {
339    fn from(schema: SchemaRef) -> OutputSchema {
340        OutputSchema {
341            column_schemas: schema
342                .column_schemas()
343                .iter()
344                .map(|cs| ColumnSchema {
345                    name: cs.name.clone(),
346                    data_type: cs.data_type.name(),
347                })
348                .collect(),
349        }
350    }
351}
352
353#[derive(Debug, Serialize, Deserialize, Eq, PartialEq)]
354pub struct HttpRecordsOutput {
355    schema: OutputSchema,
356    rows: Vec<Vec<Value>>,
357    // total_rows is equal to rows.len() in most cases,
358    // the Dashboard query result may be truncated, so we need to return the total_rows.
359    #[serde(default)]
360    total_rows: usize,
361
362    // plan level execution metrics
363    #[serde(skip_serializing_if = "HashMap::is_empty")]
364    #[serde(default)]
365    metrics: HashMap<String, Value>,
366}
367
368impl HttpRecordsOutput {
369    pub fn num_rows(&self) -> usize {
370        self.rows.len()
371    }
372
373    pub fn num_cols(&self) -> usize {
374        self.schema.column_schemas.len()
375    }
376
377    pub fn schema(&self) -> &OutputSchema {
378        &self.schema
379    }
380
381    pub fn rows(&self) -> &Vec<Vec<Value>> {
382        &self.rows
383    }
384}
385
386impl HttpRecordsOutput {
387    pub fn try_new(
388        schema: SchemaRef,
389        recordbatches: Vec<RecordBatch>,
390    ) -> std::result::Result<HttpRecordsOutput, Error> {
391        if recordbatches.is_empty() {
392            Ok(HttpRecordsOutput {
393                schema: OutputSchema::from(schema),
394                rows: vec![],
395                total_rows: 0,
396                metrics: Default::default(),
397            })
398        } else {
399            let num_rows = recordbatches.iter().map(|r| r.num_rows()).sum::<usize>();
400            let mut rows = Vec::with_capacity(num_rows);
401
402            for recordbatch in recordbatches {
403                let mut writer = HttpOutputWriter::new(schema.num_columns(), None);
404                writer.write(recordbatch, &mut rows)?;
405            }
406
407            Ok(HttpRecordsOutput {
408                schema: OutputSchema::from(schema),
409                total_rows: rows.len(),
410                rows,
411                metrics: Default::default(),
412            })
413        }
414    }
415}
416
417#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
418#[serde(rename_all = "lowercase")]
419pub enum GreptimeQueryOutput {
420    AffectedRows(usize),
421    Records(HttpRecordsOutput),
422}
423
424/// It allows the results of SQL queries to be presented in different formats.
425#[derive(Default, Debug, Clone, Copy, PartialEq, Eq)]
426pub enum ResponseFormat {
427    Arrow,
428    // (with_names, with_types)
429    Csv(bool, bool),
430    Table,
431    #[default]
432    GreptimedbV1,
433    InfluxdbV1,
434    Json,
435    Null,
436}
437
438impl ResponseFormat {
439    pub fn parse(s: &str) -> Option<Self> {
440        match s {
441            "arrow" => Some(ResponseFormat::Arrow),
442            "csv" => Some(ResponseFormat::Csv(false, false)),
443            "csvwithnames" => Some(ResponseFormat::Csv(true, false)),
444            "csvwithnamesandtypes" => Some(ResponseFormat::Csv(true, true)),
445            "table" => Some(ResponseFormat::Table),
446            "greptimedb_v1" => Some(ResponseFormat::GreptimedbV1),
447            "influxdb_v1" => Some(ResponseFormat::InfluxdbV1),
448            "json" => Some(ResponseFormat::Json),
449            "null" => Some(ResponseFormat::Null),
450            _ => None,
451        }
452    }
453
454    pub fn as_str(&self) -> &'static str {
455        match self {
456            ResponseFormat::Arrow => "arrow",
457            ResponseFormat::Csv(_, _) => "csv",
458            ResponseFormat::Table => "table",
459            ResponseFormat::GreptimedbV1 => "greptimedb_v1",
460            ResponseFormat::InfluxdbV1 => "influxdb_v1",
461            ResponseFormat::Json => "json",
462            ResponseFormat::Null => "null",
463        }
464    }
465}
466
467#[derive(Debug, Clone, Copy, PartialEq, Eq)]
468pub enum Epoch {
469    Nanosecond,
470    Microsecond,
471    Millisecond,
472    Second,
473}
474
475impl Epoch {
476    pub fn parse(s: &str) -> Option<Epoch> {
477        // Both u and µ indicate microseconds.
478        // epoch = [ns,u,µ,ms,s],
479        // For details, see the Influxdb documents.
480        // https://docs.influxdata.com/influxdb/v1/tools/api/#query-string-parameters-1
481        match s {
482            "ns" => Some(Epoch::Nanosecond),
483            "u" | "µ" => Some(Epoch::Microsecond),
484            "ms" => Some(Epoch::Millisecond),
485            "s" => Some(Epoch::Second),
486            _ => None, // just returns None for other cases
487        }
488    }
489
490    pub fn convert_timestamp(&self, ts: Timestamp) -> Option<Timestamp> {
491        match self {
492            Epoch::Nanosecond => ts.convert_to(TimeUnit::Nanosecond),
493            Epoch::Microsecond => ts.convert_to(TimeUnit::Microsecond),
494            Epoch::Millisecond => ts.convert_to(TimeUnit::Millisecond),
495            Epoch::Second => ts.convert_to(TimeUnit::Second),
496        }
497    }
498}
499
500impl Display for Epoch {
501    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
502        match self {
503            Epoch::Nanosecond => write!(f, "Epoch::Nanosecond"),
504            Epoch::Microsecond => write!(f, "Epoch::Microsecond"),
505            Epoch::Millisecond => write!(f, "Epoch::Millisecond"),
506            Epoch::Second => write!(f, "Epoch::Second"),
507        }
508    }
509}
510
511#[derive(Serialize, Deserialize, Debug)]
512pub enum HttpResponse {
513    Arrow(ArrowResponse),
514    Csv(CsvResponse),
515    Table(TableResponse),
516    Error(ErrorResponse),
517    GreptimedbV1(GreptimedbV1Response),
518    InfluxdbV1(InfluxdbV1Response),
519    Json(JsonResponse),
520    Null(NullResponse),
521}
522
523impl HttpResponse {
524    pub fn with_execution_time(self, execution_time: u64) -> Self {
525        match self {
526            HttpResponse::Arrow(resp) => resp.with_execution_time(execution_time).into(),
527            HttpResponse::Csv(resp) => resp.with_execution_time(execution_time).into(),
528            HttpResponse::Table(resp) => resp.with_execution_time(execution_time).into(),
529            HttpResponse::GreptimedbV1(resp) => resp.with_execution_time(execution_time).into(),
530            HttpResponse::InfluxdbV1(resp) => resp.with_execution_time(execution_time).into(),
531            HttpResponse::Json(resp) => resp.with_execution_time(execution_time).into(),
532            HttpResponse::Null(resp) => resp.with_execution_time(execution_time).into(),
533            HttpResponse::Error(resp) => resp.with_execution_time(execution_time).into(),
534        }
535    }
536
537    pub fn with_limit(self, limit: usize) -> Self {
538        match self {
539            HttpResponse::Csv(resp) => resp.with_limit(limit).into(),
540            HttpResponse::Table(resp) => resp.with_limit(limit).into(),
541            HttpResponse::GreptimedbV1(resp) => resp.with_limit(limit).into(),
542            HttpResponse::Json(resp) => resp.with_limit(limit).into(),
543            _ => self,
544        }
545    }
546}
547
548pub fn process_with_limit(
549    mut outputs: Vec<GreptimeQueryOutput>,
550    limit: usize,
551) -> Vec<GreptimeQueryOutput> {
552    outputs
553        .drain(..)
554        .map(|data| match data {
555            GreptimeQueryOutput::Records(mut records) => {
556                if records.rows.len() > limit {
557                    records.rows.truncate(limit);
558                    records.total_rows = limit;
559                }
560                GreptimeQueryOutput::Records(records)
561            }
562            _ => data,
563        })
564        .collect()
565}
566
567impl IntoResponse for HttpResponse {
568    fn into_response(self) -> Response {
569        match self {
570            HttpResponse::Arrow(resp) => resp.into_response(),
571            HttpResponse::Csv(resp) => resp.into_response(),
572            HttpResponse::Table(resp) => resp.into_response(),
573            HttpResponse::GreptimedbV1(resp) => resp.into_response(),
574            HttpResponse::InfluxdbV1(resp) => resp.into_response(),
575            HttpResponse::Json(resp) => resp.into_response(),
576            HttpResponse::Null(resp) => resp.into_response(),
577            HttpResponse::Error(resp) => resp.into_response(),
578        }
579    }
580}
581
582impl From<ArrowResponse> for HttpResponse {
583    fn from(value: ArrowResponse) -> Self {
584        HttpResponse::Arrow(value)
585    }
586}
587
588impl From<CsvResponse> for HttpResponse {
589    fn from(value: CsvResponse) -> Self {
590        HttpResponse::Csv(value)
591    }
592}
593
594impl From<TableResponse> for HttpResponse {
595    fn from(value: TableResponse) -> Self {
596        HttpResponse::Table(value)
597    }
598}
599
600impl From<ErrorResponse> for HttpResponse {
601    fn from(value: ErrorResponse) -> Self {
602        HttpResponse::Error(value)
603    }
604}
605
606impl From<GreptimedbV1Response> for HttpResponse {
607    fn from(value: GreptimedbV1Response) -> Self {
608        HttpResponse::GreptimedbV1(value)
609    }
610}
611
612impl From<InfluxdbV1Response> for HttpResponse {
613    fn from(value: InfluxdbV1Response) -> Self {
614        HttpResponse::InfluxdbV1(value)
615    }
616}
617
618impl From<JsonResponse> for HttpResponse {
619    fn from(value: JsonResponse) -> Self {
620        HttpResponse::Json(value)
621    }
622}
623
624impl From<NullResponse> for HttpResponse {
625    fn from(value: NullResponse) -> Self {
626        HttpResponse::Null(value)
627    }
628}
629
630#[derive(Clone)]
631pub struct ApiState {
632    pub sql_handler: ServerSqlQueryHandlerRef,
633}
634
635#[derive(Clone)]
636pub struct GreptimeOptionsConfigState {
637    pub greptime_config_options: String,
638}
639
640#[derive(Clone)]
641pub struct DashboardState {
642    pub handler: DashboardHandlerRef,
643}
644
645pub struct HttpServerBuilder {
646    options: HttpOptions,
647    batching_protocols: Vec<BatchingProtocol>,
648    logical_batching_protocols: Vec<BatchingProtocol>,
649    user_provider: Option<UserProviderRef>,
650    router: Router,
651    memory_limiter: ServerMemoryLimiter,
652}
653
654impl HttpServerBuilder {
655    pub fn new(options: HttpOptions) -> Self {
656        Self {
657            options,
658            batching_protocols: Vec::new(),
659            logical_batching_protocols: Vec::new(),
660            user_provider: None,
661            router: Router::new(),
662            memory_limiter: ServerMemoryLimiter::default(),
663        }
664    }
665
666    /// Selects HTTP protocols allowed to use logical-table batching.
667    pub fn with_logical_batching_protocols(mut self, protocols: Vec<BatchingProtocol>) -> Self {
668        self.logical_batching_protocols = protocols;
669        self
670    }
671
672    /// Selects HTTP protocols allowed to use ordinary-table batching.
673    pub fn with_batching_protocols(mut self, protocols: Vec<BatchingProtocol>) -> Self {
674        self.batching_protocols = protocols;
675        self
676    }
677
678    /// Set a global memory limiter for all server protocols.
679    pub fn with_memory_limiter(mut self, limiter: ServerMemoryLimiter) -> Self {
680        self.memory_limiter = limiter;
681        self
682    }
683
684    pub fn with_sql_handler(self, sql_handler: ServerSqlQueryHandlerRef) -> Self {
685        let batch_body_limit = if self.options.body_limit.0 == 0 {
686            4 * 1024 * 1024
687        } else {
688            (self.options.body_limit.0 as usize).min(4 * 1024 * 1024)
689        };
690        let sql_router = HttpServer::route_sql(ApiState { sql_handler }, batch_body_limit);
691
692        Self {
693            router: self
694                .router
695                .nest(&format!("/{HTTP_API_VERSION}"), sql_router),
696            ..self
697        }
698    }
699
700    pub fn with_logs_handler(self, logs_handler: LogQueryHandlerRef) -> Self {
701        let logs_router = HttpServer::route_logs(logs_handler);
702
703        Self {
704            router: self
705                .router
706                .nest(&format!("/{HTTP_API_VERSION}"), logs_router),
707            ..self
708        }
709    }
710
711    pub fn with_opentsdb_handler(self, handler: OpentsdbProtocolHandlerRef) -> Self {
712        Self {
713            router: self.router.nest(
714                &format!("/{HTTP_API_VERSION}/opentsdb"),
715                HttpServer::route_opentsdb(handler),
716            ),
717            ..self
718        }
719    }
720
721    pub fn with_influxdb_handler(self, handler: InfluxdbLineProtocolHandlerRef) -> Self {
722        Self {
723            router: self.router.nest(
724                &format!("/{HTTP_API_VERSION}/influxdb"),
725                HttpServer::route_influxdb(handler, self.memory_limiter.clone()),
726            ),
727            ..self
728        }
729    }
730
731    pub fn with_prom_handler(
732        self,
733        handler: PromStoreProtocolHandlerRef,
734        pipeline_handler: Option<PipelineHandlerRef>,
735        prom_store_with_metric_engine: bool,
736        prom_validation_mode: PromValidationMode,
737        pending_rows_batcher: Option<Arc<LogicalTablePendingRowsBatcher>>,
738    ) -> Self {
739        let state = PromStoreState {
740            prom_store_handler: handler,
741            pipeline_handler,
742            prom_store_with_metric_engine,
743            prom_validation_mode,
744            pending_rows_batcher,
745            memory_limiter: self.memory_limiter.clone(),
746        };
747
748        Self {
749            router: self.router.nest(
750                &format!("/{HTTP_API_VERSION}/prometheus"),
751                HttpServer::route_prom(state),
752            ),
753            ..self
754        }
755    }
756
757    pub fn with_prometheus_handler(self, handler: PrometheusHandlerRef) -> Self {
758        Self {
759            router: self.router.nest(
760                &format!("/{HTTP_API_VERSION}/prometheus/api/v1"),
761                HttpServer::route_prometheus(handler),
762            ),
763            ..self
764        }
765    }
766
767    pub fn with_otlp_handler(
768        self,
769        handler: OpenTelemetryProtocolHandlerRef,
770        with_metric_engine: bool,
771    ) -> Self {
772        Self {
773            router: self.router.nest(
774                &format!("/{HTTP_API_VERSION}/otlp"),
775                HttpServer::route_otlp(handler, with_metric_engine, self.memory_limiter.clone()),
776            ),
777            ..self
778        }
779    }
780
781    pub fn with_user_provider(self, user_provider: UserProviderRef) -> Self {
782        Self {
783            user_provider: Some(user_provider),
784            ..self
785        }
786    }
787
788    pub fn with_metrics_handler(self, handler: MetricsHandler) -> Self {
789        Self {
790            router: self.router.merge(HttpServer::route_metrics(handler)),
791            ..self
792        }
793    }
794
795    pub fn with_log_ingest_handler(
796        self,
797        handler: PipelineHandlerRef,
798        validator: Option<LogValidatorRef>,
799        ingest_interceptor: Option<LogIngestInterceptorRef<Error>>,
800    ) -> Self {
801        let log_state = LogState {
802            log_handler: handler,
803            log_validator: validator,
804            ingest_interceptor,
805        };
806
807        let router = self.router.nest(
808            &format!("/{HTTP_API_VERSION}"),
809            HttpServer::route_pipelines(log_state.clone(), self.memory_limiter.clone()),
810        );
811        // deprecated since v0.11.0. Use `/logs` and `/pipelines` instead.
812        let router = router.nest(
813            &format!("/{HTTP_API_VERSION}/events"),
814            #[allow(deprecated)]
815            HttpServer::route_log_deprecated(log_state.clone(), self.memory_limiter.clone()),
816        );
817
818        let router = router.nest(
819            &format!("/{HTTP_API_VERSION}/loki"),
820            HttpServer::route_loki(log_state.clone(), self.memory_limiter.clone()),
821        );
822
823        let router = router.nest(
824            &format!("/{HTTP_API_VERSION}/elasticsearch"),
825            HttpServer::route_elasticsearch(log_state.clone(), self.memory_limiter.clone()),
826        );
827
828        let router = router.nest(
829            &format!("/{HTTP_API_VERSION}/elasticsearch/"),
830            Router::new()
831                .route("/", routing::get(elasticsearch::handle_get_version))
832                .with_state(log_state.clone()),
833        );
834
835        let router = router.nest(
836            &format!("/{HTTP_API_VERSION}/splunk"),
837            HttpServer::route_splunk(log_state, self.memory_limiter.clone()),
838        );
839
840        Self { router, ..self }
841    }
842
843    pub fn with_greptime_config_options(self, opts: String) -> Self {
844        let config_router = HttpServer::route_config(GreptimeOptionsConfigState {
845            greptime_config_options: opts,
846        });
847
848        Self {
849            router: self.router.merge(config_router),
850            ..self
851        }
852    }
853
854    pub fn with_jaeger_handler(self, handler: JaegerQueryHandlerRef) -> Self {
855        Self {
856            router: self.router.nest(
857                &format!("/{HTTP_API_VERSION}/jaeger"),
858                HttpServer::route_jaeger(handler),
859            ),
860            ..self
861        }
862    }
863
864    pub fn with_dashboard_handler(self, handler: DashboardHandlerRef) -> Self {
865        Self {
866            router: self.router.nest(
867                &format!("/{HTTP_API_VERSION}/dashboards"),
868                HttpServer::route_dashboard(handler, self.memory_limiter.clone()),
869            ),
870            ..self
871        }
872    }
873
874    pub fn with_extra_router(self, router: Router) -> Self {
875        Self {
876            router: self.router.merge(router),
877            ..self
878        }
879    }
880
881    pub fn add_layer<L>(self, layer: L) -> Self
882    where
883        L: Layer<Route> + Clone + Send + Sync + 'static,
884        L::Service: Service<Request> + Clone + Send + Sync + 'static,
885        <L::Service as Service<Request>>::Response: IntoResponse + 'static,
886        <L::Service as Service<Request>>::Error: Into<Infallible> + 'static,
887        <L::Service as Service<Request>>::Future: Send + 'static,
888    {
889        Self {
890            router: self.router.layer(layer),
891            ..self
892        }
893    }
894
895    pub fn build(self) -> HttpServer {
896        HttpServer {
897            options: self.options,
898            batching_protocols: self.batching_protocols.clone(),
899            logical_batching_protocols: self.logical_batching_protocols.clone(),
900            user_provider: self.user_provider,
901            shutdown_tx: Mutex::new(None),
902            router: StdMutex::new(self.router),
903            bind_addr: None,
904            memory_limiter: self.memory_limiter,
905            kind: HttpServerKind::Full,
906        }
907    }
908
909    /// Builds the full HTTP server (serves the `v1` interfaces **and** the
910    /// internal interfaces such as health, status, metrics, config and debug).
911    /// When [`HttpOptions::enable_api_server`] is `true`, it also builds a
912    /// dedicated HTTP **API** server that serves only the `v1` interfaces plus
913    /// the dashboard.
914    ///
915    /// Both servers are built from the **same** complete route registry (so every
916    /// built-in, plugin and downstream route registered via `with_extra_router()`
917    /// is present on both). The API server additionally installs an outer guard
918    /// that exposes only the `/v1` and `/dashboard` namespaces, returning `404`
919    /// for everything else. This means downstream projects (e.g. enterprise) do
920    /// not need to classify their routes as public/internal: non-`/v1`
921    /// operational routes are hidden from the API listener automatically.
922    ///
923    /// The first returned server is bound to `self.options.addr` (e.g. port
924    /// `4000`) and keeps the historical single-server behavior. The optional
925    /// second server is bound to `api_server_addr` (e.g. `127.0.0.1:4006`). The
926    /// API server inherits **every** option from `[http]` except the bound address.
927    pub fn build_servers(self) -> (HttpServer, Option<HttpServer>) {
928        let api_enabled = self.options.enable_api_server;
929        let api_addr = self.options.api_server_addr.clone();
930
931        let internal = HttpServer {
932            options: self.options,
933            batching_protocols: self.batching_protocols.clone(),
934            logical_batching_protocols: self.logical_batching_protocols.clone(),
935            user_provider: self.user_provider.clone(),
936            shutdown_tx: Mutex::new(None),
937            router: StdMutex::new(self.router.clone()),
938            bind_addr: None,
939            memory_limiter: self.memory_limiter.clone(),
940            kind: HttpServerKind::Full,
941        };
942
943        let api = if api_enabled {
944            // The API server shares every option with the full server except the
945            // bound address.
946            let api_options = HttpOptions {
947                addr: api_addr,
948                ..internal.options.clone()
949            };
950            Some(HttpServer {
951                options: api_options,
952                batching_protocols: self.batching_protocols,
953                logical_batching_protocols: self.logical_batching_protocols,
954                user_provider: self.user_provider.clone(),
955                shutdown_tx: Mutex::new(None),
956                router: StdMutex::new(self.router),
957                bind_addr: None,
958                memory_limiter: self.memory_limiter,
959                kind: HttpServerKind::Api,
960            })
961        } else {
962            None
963        };
964
965        (internal, api)
966    }
967}
968
969impl HttpServer {
970    /// Builds the complete route registry shared by every HTTP server: all the
971    /// `v1` protocol routes, the root/health/status/metrics/config interfaces,
972    /// the dashboard and (via `build()`) the debug routes.
973    ///
974    /// Both [`HttpServerKind::Full`] and [`HttpServerKind::Api`] listeners are
975    /// built from this same router. The API listener hides the non-`v1` paths
976    /// later, in [`HttpServer::build`], via an outer guard. Keeping one router
977    /// means downstream routes added through `with_extra_router()` are present
978    /// on both listeners automatically, and the API guard decides exposure.
979    pub fn make_app(&self) -> Router {
980        let mut router = self.router.lock().unwrap().clone();
981
982        router = router
983            .route(
984                &format!("/{HTTP_API_VERSION}/health"),
985                routing::get(handler::health).post(handler::health),
986            )
987            .route("/", routing::get(handler::index))
988            .route(
989                "/health",
990                routing::get(handler::health).post(handler::health),
991            )
992            .route(
993                "/ready",
994                routing::get(handler::health).post(handler::health),
995            )
996            .route("/status", routing::get(handler::status));
997
998        // The dashboard runs on top of the `v1` APIs, so it is exposed on every
999        // HTTP server (the API listener's namespace guard allows `/dashboard`).
1000        #[cfg(feature = "dashboard")]
1001        {
1002            if !self.options.disable_dashboard {
1003                info!("Enable dashboard service at '/dashboard'");
1004                // redirect /dashboard to /dashboard/
1005                router = router.route(
1006                    "/dashboard",
1007                    routing::get(|uri: axum::http::uri::Uri| async move {
1008                        let path = uri.path();
1009                        let query = uri.query().map(|q| format!("?{}", q)).unwrap_or_default();
1010
1011                        let new_uri = format!("{}/{}", path, query);
1012                        axum::response::Redirect::permanent(&new_uri)
1013                    }),
1014                );
1015
1016                // "/dashboard" and "/dashboard/" are two different paths in Axum.
1017                // We cannot nest "/dashboard/", because we already mapping "/dashboard/{*x}" while nesting "/dashboard".
1018                // So we explicitly route "/dashboard/" here.
1019                router = router
1020                    .route(
1021                        "/dashboard/",
1022                        routing::get(dashboard::static_handler).post(dashboard::static_handler),
1023                    )
1024                    .route(
1025                        "/dashboard/{*x}",
1026                        routing::get(dashboard::static_handler).post(dashboard::static_handler),
1027                    );
1028            }
1029        }
1030
1031        // Add a layer to collect HTTP metrics for axum.
1032        router = router.route_layer(middleware::from_fn(http_metrics_layer));
1033
1034        router
1035    }
1036
1037    /// Attaches middlewares and debug routes to the router.
1038    /// Callers should call this method after [HttpServer::make_app()].
1039    pub fn build(&self, router: Router) -> Result<Router> {
1040        let timeout_layer = if self.options.timeout != Duration::default() {
1041            Some(
1042                ServiceBuilder::new().layer(
1043                    DynamicTimeoutLayer::new(self.options.timeout)
1044                        .with_status_code_fn(Self::request_timeout_status_code),
1045                ),
1046            )
1047        } else {
1048            info!("HTTP server timeout is disabled");
1049            None
1050        };
1051        let body_limit_layer = if self.options.body_limit != ReadableSize(0) {
1052            Some(
1053                ServiceBuilder::new()
1054                    .layer(DefaultBodyLimit::max(self.options.body_limit.0 as usize)),
1055            )
1056        } else {
1057            info!("HTTP server body limit is disabled");
1058            None
1059        };
1060        let cors_layer = if self.options.enable_cors {
1061            Some(
1062                CorsLayer::new()
1063                    .allow_methods([
1064                        Method::GET,
1065                        Method::POST,
1066                        Method::PUT,
1067                        Method::DELETE,
1068                        Method::HEAD,
1069                    ])
1070                    .allow_origin(cors_allow_origin(&self.options.cors_allowed_origins)?)
1071                    .allow_headers(Any),
1072            )
1073        } else {
1074            info!("HTTP server cross-origin is disabled");
1075            None
1076        };
1077
1078        let router = router
1079            // middlewares
1080            .layer(
1081                ServiceBuilder::new()
1082                    // disable on failure tracing. because printing out isn't very helpful,
1083                    // and we have impl IntoResponse for Error. It will print out more detailed error messages
1084                    .layer(TraceLayer::new_for_http().on_failure(()))
1085                    .option_layer(cors_layer)
1086                    .option_layer(timeout_layer)
1087                    .option_layer(body_limit_layer)
1088                    // memory limit layer - must be before body is consumed
1089                    .layer(middleware::from_fn_with_state(
1090                        self.memory_limiter.clone(),
1091                        memory_limit::memory_limit_middleware,
1092                    ))
1093                    // auth layer
1094                    .layer(middleware::from_fn_with_state(
1095                        AuthState::new(self.user_provider.clone()),
1096                        authorize::check_http_auth,
1097                    ))
1098                    .layer(Extension(Arc::new(self.batching_protocols.clone())))
1099                    .layer(Extension(LogicalBatchingProtocols(
1100                        self.logical_batching_protocols.clone(),
1101                    )))
1102                    .layer(middleware::from_fn(hints::extract_hints))
1103                    .layer(middleware::from_fn(client_ip::log_error_with_client_ip))
1104                    .layer(middleware::from_fn(
1105                        read_preference::extract_read_preference,
1106                    ))
1107                    .layer(middleware::from_fn(skip_wal::extract_skip_wal)),
1108            );
1109
1110        // Debug handlers are part of the complete router; the API listener hides
1111        // them through the outer namespace guard applied below.
1112        let router = router.nest(
1113            "/debug",
1114            Router::new()
1115                // handler for changing log level dynamically
1116                .route("/log_level", routing::post(dyn_log::dyn_log_handler))
1117                .route(
1118                    "/workload_scheduler",
1119                    routing::get(workload_scheduler::get_status_handler),
1120                )
1121                .route(
1122                    "/workload_scheduler/enabled",
1123                    routing::post(workload_scheduler::set_enabled_handler),
1124                )
1125                .route(
1126                    "/workload_scheduler/weights",
1127                    routing::post(workload_scheduler::set_weights_handler),
1128                )
1129                .route("/enable_trace", routing::post(dyn_trace::dyn_trace_handler))
1130                .nest(
1131                    "/prof",
1132                    Router::new()
1133                        .route("/cpu", routing::post(pprof::pprof_handler))
1134                        .route("/mem", routing::post(mem_prof::mem_prof_handler))
1135                        .route("/mem/symbol", routing::post(mem_prof::symbolicate_handler))
1136                        .route(
1137                            "/mem/activate",
1138                            routing::post(mem_prof::activate_heap_prof_handler),
1139                        )
1140                        .route(
1141                            "/mem/deactivate",
1142                            routing::post(mem_prof::deactivate_heap_prof_handler),
1143                        )
1144                        .route(
1145                            "/mem/status",
1146                            routing::get(mem_prof::heap_prof_status_handler),
1147                        ) // jemalloc gdump flag status and toggle
1148                        .route(
1149                            "/mem/gdump",
1150                            routing::get(mem_prof::gdump_status_handler)
1151                                .post(mem_prof::gdump_toggle_handler),
1152                        ),
1153                ),
1154        );
1155
1156        // Seal the API listener last: this is the outermost layer, applied after
1157        // all middleware, debug routes and any caller/plugin routes, so that only
1158        // `/v1` and `/dashboard` paths can reach a handler. Excluded paths return
1159        // 404 before auth or any other side effect.
1160        if self.kind == HttpServerKind::Api {
1161            Ok(router.layer(middleware::from_fn(enforce_api_surface)))
1162        } else {
1163            Ok(router)
1164        }
1165    }
1166
1167    fn request_timeout_status_code(request: &Request) -> HttpStatusCode {
1168        if request.uri().path() == "/v1/prometheus/write" {
1169            HttpStatusCode::GATEWAY_TIMEOUT
1170        } else {
1171            HttpStatusCode::REQUEST_TIMEOUT
1172        }
1173    }
1174
1175    fn route_metrics<S>(metrics_handler: MetricsHandler) -> Router<S> {
1176        Router::new()
1177            .route("/metrics", routing::get(handler::metrics))
1178            .with_state(metrics_handler)
1179    }
1180
1181    fn route_loki<S>(log_state: LogState, memory_limiter: ServerMemoryLimiter) -> Router<S> {
1182        Router::new()
1183            .route("/api/v1/push", routing::post(loki::loki_ingest))
1184            .layer(middleware::from_fn_with_state(
1185                memory_limiter,
1186                memory_limit::decoded_body_accounting_middleware,
1187            ))
1188            .layer(
1189                ServiceBuilder::new()
1190                    .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)),
1191            )
1192            .layer(middleware::from_fn_with_state(
1193                BatchingProtocol::Loki,
1194                set_http_write_batching,
1195            ))
1196            .with_state(log_state)
1197    }
1198
1199    fn route_splunk<S>(log_state: LogState, memory_limiter: ServerMemoryLimiter) -> Router<S> {
1200        Router::new()
1201            .route(
1202                "/services/collector/health",
1203                routing::get(splunk::handle_health),
1204            )
1205            .route(
1206                "/services/collector/health/1.0",
1207                routing::get(splunk::handle_health),
1208            )
1209            // The event endpoint plus its base and versioned aliases all serve
1210            // the same handler (Splunk JSON event protocol).
1211            .route(
1212                "/services/collector/event",
1213                routing::post(splunk::handle_event),
1214            )
1215            .route("/services/collector", routing::post(splunk::handle_event))
1216            .route(
1217                "/services/collector/event/1.0",
1218                routing::post(splunk::handle_event),
1219            )
1220            // The raw endpoint (plain-text body, one event per line) plus its
1221            // versioned alias.
1222            .route("/services/collector/raw", routing::post(splunk::handle_raw))
1223            .route(
1224                "/services/collector/raw/1.0",
1225                routing::post(splunk::handle_raw),
1226            )
1227            .layer(middleware::from_fn_with_state(
1228                memory_limiter,
1229                memory_limit::decoded_body_accounting_middleware,
1230            ))
1231            .layer(
1232                ServiceBuilder::new()
1233                    .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)),
1234            )
1235            .layer(middleware::from_fn_with_state(
1236                BatchingProtocol::Splunk,
1237                set_http_write_batching,
1238            ))
1239            .with_state(log_state)
1240    }
1241
1242    fn route_elasticsearch<S>(
1243        log_state: LogState,
1244        memory_limiter: ServerMemoryLimiter,
1245    ) -> Router<S> {
1246        Router::new()
1247            // Return fake responsefor HEAD '/' request.
1248            .route(
1249                "/",
1250                routing::head((HttpStatusCode::OK, elasticsearch::elasticsearch_headers())),
1251            )
1252            // Return fake response for Elasticsearch version request.
1253            .route("/", routing::get(elasticsearch::handle_get_version))
1254            // Return fake response for Elasticsearch license request.
1255            .route("/_license", routing::get(elasticsearch::handle_get_license))
1256            .route("/_bulk", routing::post(elasticsearch::handle_bulk_api))
1257            .route(
1258                "/{index}/_bulk",
1259                routing::post(elasticsearch::handle_bulk_api_with_index),
1260            )
1261            // Return fake response for Elasticsearch ilm request.
1262            .route(
1263                "/_ilm/policy/{*path}",
1264                routing::any((
1265                    HttpStatusCode::OK,
1266                    elasticsearch::elasticsearch_headers(),
1267                    axum::Json(serde_json::json!({})),
1268                )),
1269            )
1270            // Return fake response for Elasticsearch index template request.
1271            .route(
1272                "/_index_template/{*path}",
1273                routing::any((
1274                    HttpStatusCode::OK,
1275                    elasticsearch::elasticsearch_headers(),
1276                    axum::Json(serde_json::json!({})),
1277                )),
1278            )
1279            // Return fake response for Elasticsearch ingest pipeline request.
1280            // See: https://www.elastic.co/guide/en/elasticsearch/reference/8.8/put-pipeline-api.html.
1281            .route(
1282                "/_ingest/{*path}",
1283                routing::any((
1284                    HttpStatusCode::OK,
1285                    elasticsearch::elasticsearch_headers(),
1286                    axum::Json(serde_json::json!({})),
1287                )),
1288            )
1289            // Return fake response for Elasticsearch nodes discovery request.
1290            // See: https://www.elastic.co/guide/en/elasticsearch/reference/8.8/cluster.html.
1291            .route(
1292                "/_nodes/{*path}",
1293                routing::any((
1294                    HttpStatusCode::OK,
1295                    elasticsearch::elasticsearch_headers(),
1296                    axum::Json(serde_json::json!({})),
1297                )),
1298            )
1299            // Return fake response for Logstash APIs requests.
1300            // See: https://www.elastic.co/guide/en/elasticsearch/reference/8.8/logstash-apis.html
1301            .route(
1302                "/logstash/{*path}",
1303                routing::any((
1304                    HttpStatusCode::OK,
1305                    elasticsearch::elasticsearch_headers(),
1306                    axum::Json(serde_json::json!({})),
1307                )),
1308            )
1309            .route(
1310                "/_logstash/{*path}",
1311                routing::any((
1312                    HttpStatusCode::OK,
1313                    elasticsearch::elasticsearch_headers(),
1314                    axum::Json(serde_json::json!({})),
1315                )),
1316            )
1317            .layer(middleware::from_fn_with_state(
1318                memory_limiter,
1319                memory_limit::decoded_body_accounting_middleware,
1320            ))
1321            .layer(ServiceBuilder::new().layer(RequestDecompressionLayer::new()))
1322            .layer(middleware::from_fn_with_state(
1323                BatchingProtocol::Elasticsearch,
1324                set_http_write_batching,
1325            ))
1326            .with_state(log_state)
1327    }
1328
1329    #[deprecated(since = "0.11.0", note = "Use `route_pipelines()` instead.")]
1330    fn route_log_deprecated<S>(
1331        log_state: LogState,
1332        memory_limiter: ServerMemoryLimiter,
1333    ) -> Router<S> {
1334        Router::new()
1335            .route("/logs", routing::post(event::log_ingester))
1336            .route(
1337                "/pipelines/{pipeline_name}",
1338                routing::get(event::query_pipeline),
1339            )
1340            .route(
1341                "/pipelines/{pipeline_name}",
1342                routing::post(event::add_pipeline),
1343            )
1344            .route(
1345                "/pipelines/{pipeline_name}",
1346                routing::delete(event::delete_pipeline),
1347            )
1348            .route("/pipelines/dryrun", routing::post(event::pipeline_dryrun))
1349            .layer(middleware::from_fn_with_state(
1350                memory_limiter,
1351                memory_limit::decoded_body_accounting_middleware,
1352            ))
1353            .layer(
1354                ServiceBuilder::new()
1355                    .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)),
1356            )
1357            .layer(middleware::from_fn_with_state(
1358                BatchingProtocol::Logs,
1359                set_http_write_batching,
1360            ))
1361            .with_state(log_state)
1362    }
1363
1364    fn route_pipelines<S>(log_state: LogState, memory_limiter: ServerMemoryLimiter) -> Router<S> {
1365        Router::new()
1366            .route("/ingest", routing::post(event::log_ingester))
1367            .route(
1368                "/pipelines/{pipeline_name}",
1369                routing::get(event::query_pipeline),
1370            )
1371            .route(
1372                "/pipelines/{pipeline_name}/ddl",
1373                routing::get(event::query_pipeline_ddl),
1374            )
1375            .route(
1376                "/pipelines/{pipeline_name}",
1377                routing::post(event::add_pipeline),
1378            )
1379            .route(
1380                "/pipelines/{pipeline_name}",
1381                routing::delete(event::delete_pipeline),
1382            )
1383            .route("/pipelines/_dryrun", routing::post(event::pipeline_dryrun))
1384            .layer(middleware::from_fn_with_state(
1385                memory_limiter,
1386                memory_limit::decoded_body_accounting_middleware,
1387            ))
1388            .layer(
1389                ServiceBuilder::new()
1390                    .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)),
1391            )
1392            .layer(middleware::from_fn_with_state(
1393                BatchingProtocol::Logs,
1394                set_http_write_batching,
1395            ))
1396            .with_state(log_state)
1397    }
1398
1399    fn route_sql<S>(api_state: ApiState, batch_body_limit: usize) -> Router<S> {
1400        Router::new()
1401            .route(
1402                "/capabilities",
1403                routing::get(|State(state): State<ApiState>| async move {
1404                    let mut capabilities =
1405                        serde_json::json!({"metric_packed_import": 1, "metric_packed_export": 1});
1406                    if state.sql_handler.supports_metric_batch_ddl() {
1407                        capabilities["metric_batch_ddl"] = serde_json::json!(1);
1408                    }
1409                    axum::Json(capabilities)
1410                }),
1411            )
1412            .route(
1413                "/ddl/logical-tables",
1414                routing::post(handler::create_logical_tables)
1415                    .layer(DefaultBodyLimit::max(batch_body_limit)),
1416            )
1417            .route(
1418                "/sql",
1419                routing::get(handler::sql).post(handler::sql).layer(
1420                    middleware::from_fn_with_state(
1421                        BatchingProtocol::HttpSql,
1422                        set_http_write_batching,
1423                    ),
1424                ),
1425            )
1426            .route(
1427                "/sql/parse",
1428                routing::get(handler::sql_parse).post(handler::sql_parse),
1429            )
1430            .route(
1431                "/sql/format",
1432                routing::get(handler::sql_format).post(handler::sql_format),
1433            )
1434            .route(
1435                "/promql",
1436                routing::get(handler::promql).post(handler::promql),
1437            )
1438            .route(
1439                "/sql/analyze/stream",
1440                routing::post(handler::sql_analyze_stream),
1441            )
1442            .with_state(api_state)
1443    }
1444
1445    fn route_logs<S>(log_handler: LogQueryHandlerRef) -> Router<S> {
1446        Router::new()
1447            .route("/logs", routing::get(logs::logs).post(logs::logs))
1448            .with_state(log_handler)
1449    }
1450
1451    /// Route Prometheus [HTTP API].
1452    ///
1453    /// [HTTP API]: https://prometheus.io/docs/prometheus/latest/querying/api/
1454    pub fn route_prometheus<S>(prometheus_handler: PrometheusHandlerRef) -> Router<S> {
1455        Router::new()
1456            .route(
1457                "/format_query",
1458                routing::post(format_query).get(format_query),
1459            )
1460            .route("/status/buildinfo", routing::get(build_info_query))
1461            .route("/query", routing::post(instant_query).get(instant_query))
1462            .route("/query_range", routing::post(range_query).get(range_query))
1463            .route("/labels", routing::post(labels_query).get(labels_query))
1464            .route("/metadata", routing::get(metadata_query))
1465            .route("/series", routing::post(series_query).get(series_query))
1466            .route("/parse_query", routing::post(parse_query).get(parse_query))
1467            .route(
1468                "/label/{label_name}/values",
1469                routing::get(label_values_query),
1470            )
1471            .layer(ServiceBuilder::new().layer(CompressionLayer::new()))
1472            .with_state(prometheus_handler)
1473    }
1474
1475    /// Route Prometheus remote [read] and [write] API. In other places the related modules are
1476    /// called `prom_store`.
1477    ///
1478    /// [read]: https://prometheus.io/docs/prometheus/latest/querying/remote_read_api/
1479    /// [write]: https://prometheus.io/docs/concepts/remote_write_spec/
1480    fn route_prom<S>(state: PromStoreState) -> Router<S> {
1481        let write = routing::post(prom_store::remote_write);
1482        let write = if state.prom_store_with_metric_engine {
1483            write
1484        } else {
1485            write.layer(middleware::from_fn_with_state(
1486                BatchingProtocol::Prom,
1487                set_http_write_batching,
1488            ))
1489        };
1490        Router::new()
1491            .route("/read", routing::post(prom_store::remote_read))
1492            .route("/write", write)
1493            .with_state(state)
1494    }
1495
1496    fn route_influxdb<S>(
1497        influxdb_handler: InfluxdbLineProtocolHandlerRef,
1498        memory_limiter: ServerMemoryLimiter,
1499    ) -> Router<S> {
1500        Router::new()
1501            .route("/write", routing::post(influxdb_write_v1))
1502            .route("/api/v2/write", routing::post(influxdb_write_v2))
1503            .layer(middleware::from_fn_with_state(
1504                memory_limiter,
1505                memory_limit::decoded_body_accounting_middleware,
1506            ))
1507            .layer(
1508                ServiceBuilder::new()
1509                    .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)),
1510            )
1511            .route("/ping", routing::get(influxdb_ping))
1512            .route("/health", routing::get(influxdb_health))
1513            .layer(middleware::from_fn_with_state(
1514                BatchingProtocol::Influxdb,
1515                set_http_write_batching,
1516            ))
1517            .with_state(influxdb_handler)
1518    }
1519
1520    fn route_opentsdb<S>(opentsdb_handler: OpentsdbProtocolHandlerRef) -> Router<S> {
1521        Router::new()
1522            .route("/api/put", routing::post(opentsdb::put))
1523            .layer(middleware::from_fn_with_state(
1524                BatchingProtocol::Opentsdb,
1525                set_http_write_batching,
1526            ))
1527            .with_state(opentsdb_handler)
1528    }
1529
1530    fn route_otlp<S>(
1531        otlp_handler: OpenTelemetryProtocolHandlerRef,
1532        with_metric_engine: bool,
1533        memory_limiter: ServerMemoryLimiter,
1534    ) -> Router<S> {
1535        Router::new()
1536            .route("/v1/metrics", routing::post(otlp::metrics))
1537            .route("/v1/traces", routing::post(otlp::traces))
1538            .route("/v1/logs", routing::post(otlp::logs))
1539            .layer(middleware::from_fn_with_state(
1540                memory_limiter,
1541                memory_limit::decoded_body_accounting_middleware,
1542            ))
1543            .layer(
1544                ServiceBuilder::new()
1545                    .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)),
1546            )
1547            .layer(middleware::from_fn_with_state(
1548                BatchingProtocol::Otlp,
1549                set_http_write_batching,
1550            ))
1551            .with_state(OtlpState {
1552                with_metric_engine,
1553                handler: otlp_handler,
1554            })
1555    }
1556
1557    fn route_config<S>(state: GreptimeOptionsConfigState) -> Router<S> {
1558        Router::new()
1559            .route("/config", routing::get(handler::config))
1560            .with_state(state)
1561    }
1562
1563    fn route_jaeger<S>(handler: JaegerQueryHandlerRef) -> Router<S> {
1564        Router::new()
1565            .route("/api/services", routing::get(jaeger::handle_get_services))
1566            .route(
1567                "/api/services/{service_name}/operations",
1568                routing::get(jaeger::handle_get_operations_by_service),
1569            )
1570            .route(
1571                "/api/operations",
1572                routing::get(jaeger::handle_get_operations),
1573            )
1574            .route("/api/traces", routing::get(jaeger::handle_find_traces))
1575            .route(
1576                "/api/traces/{trace_id}",
1577                routing::get(jaeger::handle_get_trace),
1578            )
1579            .with_state(handler)
1580    }
1581
1582    #[cfg(feature = "dashboard")]
1583    fn route_dashboard<S>(
1584        handler: DashboardHandlerRef,
1585        memory_limiter: ServerMemoryLimiter,
1586    ) -> Router<S> {
1587        use crate::http::dashboard::{add_dashboard, delete_dashboard, list_dashboards};
1588
1589        Router::new()
1590            .route("/", routing::get(list_dashboards))
1591            .route("/{dashboard_name}", routing::post(add_dashboard))
1592            .route("/{dashboard_name}", routing::delete(delete_dashboard))
1593            .layer(middleware::from_fn_with_state(
1594                memory_limiter,
1595                memory_limit::decoded_body_accounting_middleware,
1596            ))
1597            .layer(
1598                ServiceBuilder::new()
1599                    .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)),
1600            )
1601            .with_state(DashboardState { handler })
1602    }
1603
1604    #[cfg(not(feature = "dashboard"))]
1605    fn route_dashboard<S>(
1606        handler: DashboardHandlerRef,
1607        _memory_limiter: ServerMemoryLimiter,
1608    ) -> Router<S> {
1609        Router::new().with_state(DashboardState { handler })
1610    }
1611}
1612
1613pub const HTTP_SERVER: &str = "HTTP_SERVER";
1614pub const HTTP_API_SERVER: &str = "HTTP_API_SERVER";
1615
1616/// An empty list allows any origin.
1617pub(crate) fn cors_allow_origin(origins: &[String]) -> Result<AllowOrigin> {
1618    if origins.is_empty() {
1619        return Ok(AllowOrigin::any());
1620    }
1621    let origins = origins
1622        .iter()
1623        .map(|origin| HeaderValue::from_str(origin).context(InvalidHeaderValueSnafu))
1624        .collect::<Result<Vec<_>>>()?;
1625    Ok(AllowOrigin::list(origins))
1626}
1627
1628#[async_trait]
1629impl Server for HttpServer {
1630    async fn shutdown(&self) -> Result<()> {
1631        let mut shutdown_tx = self.shutdown_tx.lock().await;
1632        if let Some(tx) = shutdown_tx.take()
1633            && tx.send(()).is_err()
1634        {
1635            info!("Receiver dropped, the HTTP server has already exited");
1636        }
1637        info!("Shutdown {}", self.kind());
1638
1639        Ok(())
1640    }
1641
1642    async fn start(&mut self, listening: SocketAddr) -> Result<()> {
1643        let (tx, rx) = oneshot::channel();
1644        let serve = {
1645            let mut shutdown_tx = self.shutdown_tx.lock().await;
1646            ensure!(
1647                shutdown_tx.is_none(),
1648                AlreadyStartedSnafu {
1649                    server: self.kind()
1650                }
1651            );
1652
1653            let app = self.build(self.make_app())?;
1654            let listener = tokio::net::TcpListener::bind(listening)
1655                .await
1656                .context(AddressBindSnafu { addr: listening })?
1657                .tap_io(|tcp_stream| {
1658                    if let Err(e) = tcp_stream.set_nodelay(true) {
1659                        error!(e; "Failed to set TCP_NODELAY on incoming connection");
1660                    }
1661                });
1662            let serve = axum::serve(
1663                listener,
1664                app.into_make_service_with_connect_info::<SocketAddr>(),
1665            );
1666
1667            // FIXME(yingwen): Support keepalive.
1668            // See:
1669            // - https://github.com/tokio-rs/axum/discussions/2939
1670            // - https://stackoverflow.com/questions/73069718/how-do-i-keep-alive-tokiotcpstream-in-rust
1671            // let server = axum::Server::try_bind(&listening)
1672            //     .with_context(|_| AddressBindSnafu { addr: listening })?
1673            //     .tcp_nodelay(true)
1674            //     // Enable TCP keepalive to close the dangling established connections.
1675            //     // It's configured to let the keepalive probes first send after the connection sits
1676            //     // idle for 59 minutes, and then send every 10 seconds for 6 times.
1677            //     // So the connection will be closed after roughly 1 hour.
1678            //     .tcp_keepalive(Some(Duration::from_secs(59 * 60)))
1679            //     .tcp_keepalive_interval(Some(Duration::from_secs(10)))
1680            //     .tcp_keepalive_retries(Some(6))
1681            //     .serve(app.into_make_service());
1682
1683            *shutdown_tx = Some(tx);
1684
1685            serve
1686        };
1687        let listening = serve.local_addr().context(InternalIoSnafu)?;
1688        info!("{} server is bound to {}", self.kind(), listening);
1689
1690        common_runtime::spawn_global(async move {
1691            if let Err(e) = serve
1692                .with_graceful_shutdown(rx.map(drop))
1693                .await
1694                .context(InternalIoSnafu)
1695            {
1696                error!(e; "Failed to shutdown http server");
1697            }
1698        });
1699
1700        self.bind_addr = Some(listening);
1701        Ok(())
1702    }
1703
1704    fn name(&self) -> &str {
1705        match self.kind {
1706            HttpServerKind::Api => HTTP_API_SERVER,
1707            HttpServerKind::Full => HTTP_SERVER,
1708        }
1709    }
1710
1711    fn bind_addr(&self) -> Option<SocketAddr> {
1712        self.bind_addr
1713    }
1714
1715    fn as_any(&self) -> &dyn std::any::Any {
1716        self
1717    }
1718}
1719
1720#[cfg(test)]
1721mod test {
1722    use std::future::pending;
1723    use std::io::Cursor;
1724    use std::sync::Arc;
1725
1726    use arrow_ipc::reader::StreamReader;
1727    use arrow_schema::DataType;
1728    use axum::http::StatusCode;
1729    use axum::routing::{get, post};
1730    use common_query::{Output, OutputData};
1731    use common_recordbatch::RecordBatches;
1732    use datafusion_expr::LogicalPlan;
1733    use datatypes::prelude::*;
1734    use datatypes::schema::{ColumnSchema, Schema};
1735    use datatypes::vectors::{StringVector, UInt32Vector};
1736    use header::constants::GREPTIME_DB_HEADER_TIMEOUT;
1737    use query::parser::PromQuery;
1738    use query::query_engine::DescribeResult;
1739    use session::context::QueryContextRef;
1740    use sql::statements::statement::Statement;
1741    use tokio::sync::mpsc;
1742    use tokio::time::Instant;
1743
1744    use crate::http::test_helpers::TestClient;
1745    use crate::http::*;
1746    use crate::prom_remote_write::validation::validate_label_name;
1747    use crate::query_handler::sql::SqlQueryHandler;
1748
1749    struct DummyInstance {
1750        _tx: mpsc::Sender<(String, Vec<u8>)>,
1751    }
1752
1753    #[async_trait]
1754    impl SqlQueryHandler for DummyInstance {
1755        async fn do_query(&self, _: &str, _: QueryContextRef) -> Vec<Result<Output>> {
1756            unimplemented!()
1757        }
1758
1759        async fn do_analyze_stream_query(&self, _: &str, _: QueryContextRef) -> Result<Output> {
1760            let stream = common_recordbatch::RecordBatches::empty().as_stream();
1761            Ok(Output::new(OutputData::Stream(stream), Default::default()))
1762        }
1763
1764        async fn do_promql_query(&self, _: &PromQuery, _: QueryContextRef) -> Vec<Result<Output>> {
1765            unimplemented!()
1766        }
1767
1768        async fn do_exec_plan(
1769            &self,
1770            _plan: LogicalPlan,
1771            _stmt: Option<Statement>,
1772            _query_ctx: QueryContextRef,
1773        ) -> Result<Output> {
1774            unimplemented!()
1775        }
1776
1777        async fn do_describe(
1778            &self,
1779            _stmt: sql::statements::statement::Statement,
1780            _query_ctx: QueryContextRef,
1781        ) -> Result<Option<DescribeResult>> {
1782            unimplemented!()
1783        }
1784
1785        async fn is_valid_schema(&self, _catalog: &str, _schema: &str) -> Result<bool> {
1786            Ok(true)
1787        }
1788    }
1789
1790    async fn forever() {
1791        pending().await
1792    }
1793
1794    fn make_test_app(tx: mpsc::Sender<(String, Vec<u8>)>) -> Router {
1795        make_test_app_custom(tx, HttpOptions::default())
1796    }
1797
1798    fn make_test_app_custom(tx: mpsc::Sender<(String, Vec<u8>)>, options: HttpOptions) -> Router {
1799        let instance = Arc::new(DummyInstance { _tx: tx });
1800        let server = HttpServerBuilder::new(options)
1801            .with_sql_handler(instance.clone())
1802            .build();
1803        let app = server
1804            .make_app()
1805            .route("/test/timeout", get(forever))
1806            .route("/v1/prometheus/write", post(forever));
1807        server.build(app).unwrap()
1808    }
1809
1810    fn make_split_builder() -> HttpServerBuilder {
1811        let (tx, _rx) = mpsc::channel(100);
1812        let instance = Arc::new(DummyInstance { _tx: tx });
1813        // Enable the dedicated API server so the split behavior can be exercised.
1814        let options = HttpOptions {
1815            enable_api_server: true,
1816            ..HttpOptions::default()
1817        };
1818        HttpServerBuilder::new(options)
1819            .with_sql_handler(instance)
1820            .with_metrics_handler(MetricsHandler)
1821            .with_greptime_config_options("dummy = \"value\"".to_string())
1822    }
1823
1824    #[tokio::test]
1825    pub async fn test_http_api_options_defaults() {
1826        let opts = HttpOptions::default();
1827        // The API server is disabled by default.
1828        assert!(!opts.enable_api_server);
1829        assert_eq!(opts.api_server_addr, "127.0.0.1:4006");
1830    }
1831
1832    #[tokio::test]
1833    pub async fn test_build_serves_all_routes() {
1834        // The single (internal/full) server built via `build()` keeps serving both
1835        // the public v1 interfaces and the internal ones. This is the historical
1836        // single-server behavior and must not change.
1837        let server = make_split_builder().build();
1838        assert_eq!(server.name(), HTTP_SERVER);
1839
1840        let app = server.build(server.make_app()).unwrap();
1841        let client = TestClient::new(app).await;
1842
1843        for path in ["/v1/health", "/health", "/status", "/metrics", "/config"] {
1844            let res = client.get(path).send().await;
1845            assert_eq!(
1846                res.status(),
1847                StatusCode::OK,
1848                "internal/full server should serve {path}"
1849            );
1850        }
1851    }
1852
1853    #[tokio::test]
1854    pub async fn test_build_servers_separates_api_and_internal_routes() {
1855        // `build_servers` produces the internal/full server (everything) plus a
1856        // dedicated API server that only serves the v1 interfaces (and dashboard).
1857        let (internal, api) = make_split_builder().build_servers();
1858        let api = api.expect("API server is explicitly enabled in make_split_builder");
1859        assert_eq!(internal.name(), HTTP_SERVER);
1860        assert_eq!(api.name(), HTTP_API_SERVER);
1861
1862        let internal_app = internal.build(internal.make_app()).unwrap();
1863        let api_app = api.build(api.make_app()).unwrap();
1864
1865        let internal_client = TestClient::new(internal_app).await;
1866        let api_client = TestClient::new(api_app).await;
1867
1868        // Both servers serve the v1 interfaces (e.g. the v1 health check).
1869        assert_eq!(
1870            internal_client.get("/v1/health").send().await.status(),
1871            StatusCode::OK
1872        );
1873        assert_eq!(
1874            api_client.get("/v1/health").send().await.status(),
1875            StatusCode::OK
1876        );
1877
1878        // The internal-only interfaces are served by the internal/full server...
1879        for path in ["/health", "/status", "/metrics", "/config"] {
1880            assert_eq!(
1881                internal_client.get(path).send().await.status(),
1882                StatusCode::OK,
1883                "internal/full server should serve {path}"
1884            );
1885            // ...and NOT by the API-only server.
1886            assert_eq!(
1887                api_client.get(path).send().await.status(),
1888                StatusCode::NOT_FOUND,
1889                "API server should NOT serve {path}"
1890            );
1891        }
1892    }
1893
1894    #[test]
1895    fn test_build_servers_api_inherits_http_options() {
1896        // The API server must inherit every `[http]` option except the bound addr.
1897        let http_opts = HttpOptions {
1898            timeout: Duration::from_secs(42),
1899            body_limit: ReadableSize::mb(128),
1900            cors_allowed_origins: vec!["https://example.com".to_string()],
1901            enable_api_server: true,
1902            ..HttpOptions::default()
1903        };
1904        let (internal, api) = HttpServerBuilder::new(http_opts.clone()).build_servers();
1905        let api = api.expect("API server is explicitly enabled");
1906
1907        // Only the bound address differs.
1908        assert_eq!(internal.options.addr, http_opts.addr);
1909        assert_eq!(api.options.addr, "127.0.0.1:4006");
1910        // Everything else is inherited from `[http]`.
1911        assert_eq!(api.options.timeout, http_opts.timeout);
1912        assert_eq!(api.options.body_limit, http_opts.body_limit);
1913        assert_eq!(
1914            api.options.cors_allowed_origins,
1915            http_opts.cors_allowed_origins
1916        );
1917    }
1918
1919    #[test]
1920    fn test_is_api_listener_path() {
1921        // `/v1` namespace (exact slash boundary).
1922        assert!(is_api_listener_path("/v1"));
1923        assert!(is_api_listener_path("/v1/"));
1924        assert!(is_api_listener_path("/v1/sql"));
1925        assert!(!is_api_listener_path("/v10/sql"));
1926        assert!(!is_api_listener_path("/v1-internal"));
1927        // `/dashboard` namespace.
1928        assert!(is_api_listener_path("/dashboard"));
1929        assert!(is_api_listener_path("/dashboard/app.js"));
1930        assert!(!is_api_listener_path("/dashboard-admin"));
1931        // Operational paths are never API-visible.
1932        assert!(!is_api_listener_path("/metrics"));
1933        assert!(!is_api_listener_path("/status/plugin"));
1934        assert!(!is_api_listener_path("/health"));
1935    }
1936
1937    #[tokio::test]
1938    async fn packed_capability_requires_auth_on_full_and_api_listeners() {
1939        let (tx, _rx) = mpsc::channel(1);
1940        let provider =
1941            auth::static_user_provider_from_option("static_user_provider:cmd:user=password")
1942                .unwrap();
1943        let (full, api) = HttpServerBuilder::new(HttpOptions {
1944            enable_api_server: true,
1945            ..Default::default()
1946        })
1947        .with_sql_handler(Arc::new(DummyInstance { _tx: tx }))
1948        .with_user_provider(Arc::new(provider))
1949        .build_servers();
1950        for server in [full, api.unwrap()] {
1951            let client = TestClient::new(server.build(server.make_app()).unwrap()).await;
1952            assert_eq!(
1953                client.get("/v1/capabilities").send().await.status(),
1954                StatusCode::UNAUTHORIZED
1955            );
1956            assert_eq!(
1957                client
1958                    .post("/v1/ddl/logical-tables")
1959                    .form(&[("sql", "SELECT 1")])
1960                    .send()
1961                    .await
1962                    .status(),
1963                StatusCode::UNAUTHORIZED
1964            );
1965            let response = client
1966                .get("/v1/capabilities")
1967                .header("Authorization", "Basic dXNlcjpwYXNzd29yZA==")
1968                .send()
1969                .await;
1970            assert_eq!(response.status(), StatusCode::OK);
1971            assert_eq!(
1972                response.json::<serde_json::Value>().await,
1973                serde_json::json!({"metric_packed_import": 1, "metric_packed_export": 1})
1974            );
1975        }
1976    }
1977
1978    #[tokio::test]
1979    async fn logical_ddl_form_limits_and_unsupported_handler() {
1980        use axum::body::{Body, to_bytes};
1981        use tower::ServiceExt;
1982
1983        async fn post_form(app: Router, body: String) -> (StatusCode, serde_json::Value) {
1984            let response = app
1985                .oneshot(
1986                    Request::builder()
1987                        .method(Method::POST)
1988                        .uri("/v1/ddl/logical-tables")
1989                        .header("Content-Type", "application/x-www-form-urlencoded")
1990                        .body(Body::from(body))
1991                        .unwrap(),
1992                )
1993                .await
1994                .unwrap();
1995            let status = response.status();
1996            let body = to_bytes(response.into_body(), 4096).await.unwrap();
1997            (status, serde_json::from_slice(&body).unwrap())
1998        }
1999
2000        for limit in [0, 64] {
2001            let (tx, _rx) = mpsc::channel(1);
2002            let server = HttpServerBuilder::new(HttpOptions {
2003                body_limit: ReadableSize(limit),
2004                ..Default::default()
2005            })
2006            .with_sql_handler(Arc::new(DummyInstance { _tx: tx }))
2007            .build();
2008            let app = server.build(server.make_app()).unwrap();
2009            let client = TestClient::new(app.clone()).await;
2010            assert_eq!(
2011                client.get("/v1/ddl/logical-tables").send().await.status(),
2012                StatusCode::METHOD_NOT_ALLOWED
2013            );
2014            let small = client
2015                .post("/v1/ddl/logical-tables")
2016                .form(&[("sql", "CREATE TABLE t")])
2017                .send()
2018                .await;
2019            assert_eq!(
2020                small.json::<serde_json::Value>().await["code"],
2021                common_error::status_code::StatusCode::Unsupported as u32
2022            );
2023            let (_, body) = post_form(
2024                app.clone(),
2025                if limit == 0 {
2026                    format!("sql={}", "%20".repeat(1024 * 1024))
2027                } else {
2028                    "sql=".to_string() + &" ".repeat(65)
2029                },
2030            )
2031            .await;
2032            let code = body["code"].clone();
2033            assert_eq!(
2034                code,
2035                if limit == 0 {
2036                    common_error::status_code::StatusCode::Unsupported as u32
2037                } else {
2038                    common_error::status_code::StatusCode::InvalidArguments as u32
2039                }
2040            );
2041            for body in [
2042                format!("sql={}", "x".repeat(1024 * 1024 + 1)),
2043                format!("sql={}", "%C3%A9".repeat(512 * 1024 + 1)),
2044                format!("sql=CREATE+TABLE+t&padding={}", "x".repeat(4 * 1024 * 1024)),
2045            ] {
2046                let (status, body) = post_form(app.clone(), body).await;
2047                assert_eq!(status, StatusCode::BAD_REQUEST);
2048                assert_eq!(
2049                    body["code"],
2050                    common_error::status_code::StatusCode::InvalidArguments as u32
2051                );
2052            }
2053        }
2054    }
2055
2056    #[tokio::test]
2057    pub async fn test_extra_router_respects_api_surface() {
2058        // Downstream routes added through `with_extra_router()` (e.g. enterprise
2059        // plugins or the `RouterConfigurator`) are present on the full listener
2060        // and on the API listener's router, but the API listener's outer guard
2061        // hides every non-`/v1` path automatically. This is what keeps the
2062        // feature correct for greptimedb-enterprise with no downstream changes.
2063        let (tx, _rx) = mpsc::channel(100);
2064        let instance = Arc::new(DummyInstance { _tx: tx });
2065        let options = HttpOptions {
2066            enable_api_server: true,
2067            ..HttpOptions::default()
2068        };
2069        let builder = HttpServerBuilder::new(options)
2070            .with_sql_handler(instance)
2071            .with_extra_router(
2072                Router::new()
2073                    .route("/status/plugin", routing::get(|| async { "ok" }))
2074                    .route("/v1/plugin", routing::get(|| async { "ok" })),
2075            );
2076        let (full, api) = builder.build_servers();
2077        let api = api.expect("API server is explicitly enabled");
2078
2079        let full_client = TestClient::new(full.build(full.make_app()).unwrap()).await;
2080        let api_client = TestClient::new(api.build(api.make_app()).unwrap()).await;
2081
2082        // A non-`/v1` route is reachable on the full listener...
2083        assert_eq!(
2084            full_client.get("/status/plugin").send().await.status(),
2085            StatusCode::OK
2086        );
2087        // ...and hidden (404) on the API listener, without the caller classifying it.
2088        assert_eq!(
2089            api_client.get("/status/plugin").send().await.status(),
2090            StatusCode::NOT_FOUND
2091        );
2092
2093        // A `/v1` route is reachable on both.
2094        assert_eq!(
2095            full_client.get("/v1/plugin").send().await.status(),
2096            StatusCode::OK
2097        );
2098        assert_eq!(
2099            api_client.get("/v1/plugin").send().await.status(),
2100            StatusCode::OK
2101        );
2102
2103        // Slash-boundary correctness on the API listener.
2104        assert_eq!(
2105            api_client.get("/v10/plugin").send().await.status(),
2106            StatusCode::NOT_FOUND
2107        );
2108    }
2109
2110    #[test]
2111    fn test_build_servers_api_disabled_by_default() {
2112        // The API server is disabled by default, so `build_servers` returns no
2113        // second server unless `enable_api_server` is set.
2114        let (_internal, api) = HttpServerBuilder::new(HttpOptions::default()).build_servers();
2115        assert!(api.is_none());
2116    }
2117
2118    #[tokio::test]
2119    pub async fn test_cors() {
2120        // cors is on by default
2121        let (tx, _rx) = mpsc::channel(100);
2122        let app = make_test_app(tx);
2123        let client = TestClient::new(app).await;
2124
2125        let res = client.get("/health").send().await;
2126
2127        assert_eq!(res.status(), StatusCode::OK);
2128        assert_eq!(
2129            res.headers()
2130                .get(http::header::ACCESS_CONTROL_ALLOW_ORIGIN)
2131                .expect("expect cors header origin"),
2132            "*"
2133        );
2134
2135        let res = client.get("/v1/health").send().await;
2136
2137        assert_eq!(res.status(), StatusCode::OK);
2138        assert_eq!(
2139            res.headers()
2140                .get(http::header::ACCESS_CONTROL_ALLOW_ORIGIN)
2141                .expect("expect cors header origin"),
2142            "*"
2143        );
2144
2145        let res = client
2146            .options("/health")
2147            .header("Access-Control-Request-Headers", "x-greptime-auth")
2148            .header("Access-Control-Request-Method", "DELETE")
2149            .header("Origin", "https://example.com")
2150            .send()
2151            .await;
2152        assert_eq!(res.status(), StatusCode::OK);
2153        assert_eq!(
2154            res.headers()
2155                .get(http::header::ACCESS_CONTROL_ALLOW_ORIGIN)
2156                .expect("expect cors header origin"),
2157            "*"
2158        );
2159        assert_eq!(
2160            res.headers()
2161                .get(http::header::ACCESS_CONTROL_ALLOW_HEADERS)
2162                .expect("expect cors header headers"),
2163            "*"
2164        );
2165        assert_eq!(
2166            res.headers()
2167                .get(http::header::ACCESS_CONTROL_ALLOW_METHODS)
2168                .expect("expect cors header methods"),
2169            "GET,POST,PUT,DELETE,HEAD"
2170        );
2171    }
2172
2173    #[tokio::test]
2174    pub async fn test_cors_custom_origins() {
2175        // cors is on by default
2176        let (tx, _rx) = mpsc::channel(100);
2177        let origin = "https://example.com";
2178
2179        let options = HttpOptions {
2180            cors_allowed_origins: vec![origin.to_string()],
2181            ..Default::default()
2182        };
2183
2184        let app = make_test_app_custom(tx, options);
2185        let client = TestClient::new(app).await;
2186
2187        let res = client.get("/health").header("Origin", origin).send().await;
2188
2189        assert_eq!(res.status(), StatusCode::OK);
2190        assert_eq!(
2191            res.headers()
2192                .get(http::header::ACCESS_CONTROL_ALLOW_ORIGIN)
2193                .expect("expect cors header origin"),
2194            origin
2195        );
2196
2197        let res = client
2198            .get("/health")
2199            .header("Origin", "https://notallowed.com")
2200            .send()
2201            .await;
2202
2203        assert_eq!(res.status(), StatusCode::OK);
2204        assert!(
2205            !res.headers()
2206                .contains_key(http::header::ACCESS_CONTROL_ALLOW_ORIGIN)
2207        );
2208    }
2209
2210    #[tokio::test]
2211    pub async fn test_cors_disabled() {
2212        // cors is on by default
2213        let (tx, _rx) = mpsc::channel(100);
2214
2215        let options = HttpOptions {
2216            enable_cors: false,
2217            ..Default::default()
2218        };
2219
2220        let app = make_test_app_custom(tx, options);
2221        let client = TestClient::new(app).await;
2222
2223        let res = client.get("/health").send().await;
2224
2225        assert_eq!(res.status(), StatusCode::OK);
2226        assert!(
2227            !res.headers()
2228                .contains_key(http::header::ACCESS_CONTROL_ALLOW_ORIGIN)
2229        );
2230    }
2231
2232    #[test]
2233    fn test_http_options_default() {
2234        let default = HttpOptions::default();
2235        assert_eq!("127.0.0.1:4000".to_string(), default.addr);
2236        assert_eq!(Duration::from_secs(0), default.timeout)
2237    }
2238
2239    #[tokio::test]
2240    async fn test_http_options_legacy_analyze_stream_config_is_ignored() {
2241        let options: HttpOptions = serde_json::from_value(serde_json::json!({
2242            "addr": "127.0.0.1:4000",
2243            "timeout": "0s",
2244            "body_limit": "64MiB",
2245            "cors_allowed_origins": [],
2246            "enable_cors": true,
2247            "experimental_enable_explain_analyze_stream": false,
2248            "enable_api_server": false,
2249            "api_server_addr": "127.0.0.1:4006"
2250        }))
2251        .unwrap();
2252        let serialized = serde_json::to_string(&options).unwrap();
2253        assert!(!serialized.contains("experimental_enable_explain_analyze_stream"));
2254
2255        let (tx, _rx) = mpsc::channel(100);
2256        let app = make_test_app_custom(tx, options);
2257        let client = TestClient::new(app).await;
2258        let response = client
2259            .post("/v1/sql/analyze/stream")
2260            .form(&handler::SqlQuery {
2261                sql: Some("EXPLAIN ANALYZE VERBOSE SELECT 1".to_string()),
2262                ..Default::default()
2263            })
2264            .send()
2265            .await;
2266        assert_ne!(response.status(), StatusCode::NOT_FOUND);
2267    }
2268
2269    #[tokio::test]
2270    async fn test_http_server_request_timeout() {
2271        common_telemetry::init_default_ut_logging();
2272
2273        let (tx, _rx) = mpsc::channel(100);
2274        let options = HttpOptions {
2275            timeout: Duration::from_millis(10),
2276            ..Default::default()
2277        };
2278        let app = make_test_app_custom(tx, options);
2279        let client = TestClient::new(app).await;
2280        let res = client.get("/test/timeout").send().await;
2281        assert_eq!(res.status(), StatusCode::REQUEST_TIMEOUT);
2282
2283        let res = client.post("/v1/prometheus/write").send().await;
2284        assert_eq!(res.status(), StatusCode::GATEWAY_TIMEOUT);
2285
2286        let now = Instant::now();
2287        let res = client
2288            .get("/test/timeout")
2289            .header(GREPTIME_DB_HEADER_TIMEOUT, "20ms")
2290            .send()
2291            .await;
2292        assert_eq!(res.status(), StatusCode::REQUEST_TIMEOUT);
2293        let elapsed = now.elapsed();
2294        assert!(elapsed > Duration::from_millis(15));
2295
2296        tokio::time::timeout(
2297            Duration::from_millis(15),
2298            client
2299                .get("/test/timeout")
2300                .header(GREPTIME_DB_HEADER_TIMEOUT, "0s")
2301                .send(),
2302        )
2303        .await
2304        .unwrap_err();
2305
2306        tokio::time::timeout(
2307            Duration::from_millis(15),
2308            client
2309                .get("/test/timeout")
2310                .header(
2311                    GREPTIME_DB_HEADER_TIMEOUT,
2312                    humantime::format_duration(Duration::default()).to_string(),
2313                )
2314                .send(),
2315        )
2316        .await
2317        .unwrap_err();
2318    }
2319
2320    #[tokio::test]
2321    async fn test_schema_for_empty_response() {
2322        let column_schemas = vec![
2323            ColumnSchema::new("numbers", ConcreteDataType::uint32_datatype(), false),
2324            ColumnSchema::new("strings", ConcreteDataType::string_datatype(), true),
2325        ];
2326        let schema = Arc::new(Schema::new(column_schemas));
2327
2328        let recordbatches = RecordBatches::try_new(schema.clone(), vec![]).unwrap();
2329        let outputs = vec![Ok(Output::new_with_record_batches(recordbatches))];
2330
2331        let json_resp = GreptimedbV1Response::from_output(outputs).await;
2332        if let HttpResponse::GreptimedbV1(json_resp) = json_resp {
2333            let json_output = &json_resp.output[0];
2334            if let GreptimeQueryOutput::Records(r) = json_output {
2335                assert_eq!(r.num_rows(), 0);
2336                assert_eq!(r.num_cols(), 2);
2337                assert_eq!(r.schema.column_schemas[0].name, "numbers");
2338                assert_eq!(r.schema.column_schemas[0].data_type, "UInt32");
2339            } else {
2340                panic!("invalid output type");
2341            }
2342        } else {
2343            panic!("invalid format")
2344        }
2345    }
2346
2347    #[tokio::test]
2348    async fn test_recordbatches_conversion() {
2349        let column_schemas = vec![
2350            ColumnSchema::new("numbers", ConcreteDataType::uint32_datatype(), false),
2351            ColumnSchema::new("strings", ConcreteDataType::string_datatype(), true),
2352        ];
2353        let schema = Arc::new(Schema::new(column_schemas));
2354        let columns: Vec<VectorRef> = vec![
2355            Arc::new(UInt32Vector::from_slice(vec![1, 2, 3, 4])),
2356            Arc::new(StringVector::from(vec![
2357                None,
2358                Some("hello"),
2359                Some("greptime"),
2360                None,
2361            ])),
2362        ];
2363        let recordbatch = RecordBatch::new(schema.clone(), columns).unwrap();
2364
2365        for format in [
2366            ResponseFormat::GreptimedbV1,
2367            ResponseFormat::InfluxdbV1,
2368            ResponseFormat::Csv(true, true),
2369            ResponseFormat::Table,
2370            ResponseFormat::Arrow,
2371            ResponseFormat::Json,
2372            ResponseFormat::Null,
2373        ] {
2374            let recordbatches =
2375                RecordBatches::try_new(schema.clone(), vec![recordbatch.clone()]).unwrap();
2376            let outputs = vec![Ok(Output::new_with_record_batches(recordbatches))];
2377            let json_resp = match format {
2378                ResponseFormat::Arrow => ArrowResponse::from_output(outputs, None).await,
2379                ResponseFormat::Csv(with_names, with_types) => {
2380                    CsvResponse::from_output(outputs, with_names, with_types).await
2381                }
2382                ResponseFormat::Table => TableResponse::from_output(outputs).await,
2383                ResponseFormat::GreptimedbV1 => GreptimedbV1Response::from_output(outputs).await,
2384                ResponseFormat::InfluxdbV1 => InfluxdbV1Response::from_output(outputs, None).await,
2385                ResponseFormat::Json => JsonResponse::from_output(outputs).await,
2386                ResponseFormat::Null => NullResponse::from_output(outputs).await,
2387            };
2388
2389            match json_resp {
2390                HttpResponse::GreptimedbV1(resp) => {
2391                    let json_output = &resp.output[0];
2392                    if let GreptimeQueryOutput::Records(r) = json_output {
2393                        assert_eq!(r.num_rows(), 4);
2394                        assert_eq!(r.num_cols(), 2);
2395                        assert_eq!(r.schema.column_schemas[0].name, "numbers");
2396                        assert_eq!(r.schema.column_schemas[0].data_type, "UInt32");
2397                        assert_eq!(r.rows[0][0], serde_json::Value::from(1));
2398                        assert_eq!(r.rows[0][1], serde_json::Value::Null);
2399                    } else {
2400                        panic!("invalid output type");
2401                    }
2402                }
2403                HttpResponse::InfluxdbV1(resp) => {
2404                    let json_output = &resp.results()[0];
2405                    assert_eq!(json_output.num_rows(), 4);
2406                    assert_eq!(json_output.num_cols(), 2);
2407                    assert_eq!(json_output.series[0].columns.clone()[0], "numbers");
2408                    assert_eq!(
2409                        json_output.series[0].values[0][0],
2410                        serde_json::Value::from(1)
2411                    );
2412                    assert_eq!(json_output.series[0].values[0][1], serde_json::Value::Null);
2413                }
2414                HttpResponse::Csv(resp) => {
2415                    let output = &resp.output()[0];
2416                    if let GreptimeQueryOutput::Records(r) = output {
2417                        assert_eq!(r.num_rows(), 4);
2418                        assert_eq!(r.num_cols(), 2);
2419                        assert_eq!(r.schema.column_schemas[0].name, "numbers");
2420                        assert_eq!(r.schema.column_schemas[0].data_type, "UInt32");
2421                        assert_eq!(r.rows[0][0], serde_json::Value::from(1));
2422                        assert_eq!(r.rows[0][1], serde_json::Value::Null);
2423                    } else {
2424                        panic!("invalid output type");
2425                    }
2426                }
2427
2428                HttpResponse::Table(resp) => {
2429                    let output = &resp.output()[0];
2430                    if let GreptimeQueryOutput::Records(r) = output {
2431                        assert_eq!(r.num_rows(), 4);
2432                        assert_eq!(r.num_cols(), 2);
2433                        assert_eq!(r.schema.column_schemas[0].name, "numbers");
2434                        assert_eq!(r.schema.column_schemas[0].data_type, "UInt32");
2435                        assert_eq!(r.rows[0][0], serde_json::Value::from(1));
2436                        assert_eq!(r.rows[0][1], serde_json::Value::Null);
2437                    } else {
2438                        panic!("invalid output type");
2439                    }
2440                }
2441
2442                HttpResponse::Arrow(resp) => {
2443                    let output = resp.data;
2444                    let mut reader = StreamReader::try_new(Cursor::new(output), None)
2445                        .expect("Arrow reader error");
2446                    let schema = reader.schema();
2447                    assert_eq!(schema.fields[0].name(), "numbers");
2448                    assert_eq!(schema.fields[0].data_type(), &DataType::UInt32);
2449                    assert_eq!(schema.fields[1].name(), "strings");
2450                    assert_eq!(schema.fields[1].data_type(), &DataType::Utf8);
2451
2452                    let rb = reader.next().unwrap().expect("read record batch failed");
2453                    assert_eq!(rb.num_columns(), 2);
2454                    assert_eq!(rb.num_rows(), 4);
2455                }
2456
2457                HttpResponse::Json(resp) => {
2458                    let output = &resp.output()[0];
2459                    if let GreptimeQueryOutput::Records(r) = output {
2460                        assert_eq!(r.num_rows(), 4);
2461                        assert_eq!(r.num_cols(), 2);
2462                        assert_eq!(r.schema.column_schemas[0].name, "numbers");
2463                        assert_eq!(r.schema.column_schemas[0].data_type, "UInt32");
2464                        assert_eq!(r.rows[0][0], serde_json::Value::from(1));
2465                        assert_eq!(r.rows[0][1], serde_json::Value::Null);
2466                    } else {
2467                        panic!("invalid output type");
2468                    }
2469                }
2470
2471                HttpResponse::Null(resp) => {
2472                    assert_eq!(resp.rows(), 4);
2473                }
2474
2475                HttpResponse::Error(err) => unreachable!("{err:?}"),
2476            }
2477        }
2478    }
2479
2480    #[test]
2481    fn test_response_format_misc() {
2482        assert_eq!(ResponseFormat::default(), ResponseFormat::GreptimedbV1);
2483        assert_eq!(ResponseFormat::parse("arrow"), Some(ResponseFormat::Arrow));
2484        assert_eq!(
2485            ResponseFormat::parse("csv"),
2486            Some(ResponseFormat::Csv(false, false))
2487        );
2488        assert_eq!(
2489            ResponseFormat::parse("csvwithnames"),
2490            Some(ResponseFormat::Csv(true, false))
2491        );
2492        assert_eq!(
2493            ResponseFormat::parse("csvwithnamesandtypes"),
2494            Some(ResponseFormat::Csv(true, true))
2495        );
2496        assert_eq!(ResponseFormat::parse("table"), Some(ResponseFormat::Table));
2497        assert_eq!(
2498            ResponseFormat::parse("greptimedb_v1"),
2499            Some(ResponseFormat::GreptimedbV1)
2500        );
2501        assert_eq!(
2502            ResponseFormat::parse("influxdb_v1"),
2503            Some(ResponseFormat::InfluxdbV1)
2504        );
2505        assert_eq!(ResponseFormat::parse("json"), Some(ResponseFormat::Json));
2506        assert_eq!(ResponseFormat::parse("null"), Some(ResponseFormat::Null));
2507
2508        // invalid formats
2509        assert_eq!(ResponseFormat::parse("invalid"), None);
2510        assert_eq!(ResponseFormat::parse(""), None);
2511        assert_eq!(ResponseFormat::parse("CSV"), None); // Case sensitive
2512
2513        // as str
2514        assert_eq!(ResponseFormat::Arrow.as_str(), "arrow");
2515        assert_eq!(ResponseFormat::Csv(false, false).as_str(), "csv");
2516        assert_eq!(ResponseFormat::Csv(true, true).as_str(), "csv");
2517        assert_eq!(ResponseFormat::Table.as_str(), "table");
2518        assert_eq!(ResponseFormat::GreptimedbV1.as_str(), "greptimedb_v1");
2519        assert_eq!(ResponseFormat::InfluxdbV1.as_str(), "influxdb_v1");
2520        assert_eq!(ResponseFormat::Json.as_str(), "json");
2521        assert_eq!(ResponseFormat::Null.as_str(), "null");
2522        assert_eq!(ResponseFormat::default().as_str(), "greptimedb_v1");
2523    }
2524
2525    #[test]
2526    fn test_decode_label_name_strict() {
2527        let strict = PromValidationMode::Strict;
2528
2529        // Valid Prometheus label names
2530        assert!(strict.decode_label_name(b"__name__").is_ok());
2531        assert!(strict.decode_label_name(b"job").is_ok());
2532        assert!(strict.decode_label_name(b"instance").is_ok());
2533        assert!(strict.decode_label_name(b"_private").is_ok());
2534        assert!(strict.decode_label_name(b"label_with_underscores").is_ok());
2535        assert!(strict.decode_label_name(b"abc123").is_ok());
2536        assert!(strict.decode_label_name(b"A").is_ok());
2537        assert!(strict.decode_label_name(b"_").is_ok());
2538
2539        // Invalid: starts with digit
2540        assert!(strict.decode_label_name(b"0abc").is_err());
2541        assert!(strict.decode_label_name(b"123").is_err());
2542
2543        // Invalid: contains special characters
2544        assert!(strict.decode_label_name(b"label-name").is_err());
2545        assert!(strict.decode_label_name(b"label.name").is_err());
2546        assert!(strict.decode_label_name(b"label name").is_err());
2547        assert!(strict.decode_label_name(b"label/name").is_err());
2548
2549        // Invalid: empty
2550        assert!(strict.decode_label_name(b"").is_err());
2551
2552        // Invalid: non-ASCII UTF-8
2553        assert!(strict.decode_label_name("ラベル".as_bytes()).is_err());
2554
2555        // Invalid UTF-8 bytes should fail
2556        assert!(strict.decode_label_name(&[0xff, 0xfe]).is_err());
2557    }
2558
2559    #[test]
2560    fn test_decode_label_name_lossy() {
2561        let lossy = PromValidationMode::Lossy;
2562
2563        // Label name validation is always enforced.
2564        assert!(lossy.decode_label_name(b"__name__").is_ok());
2565        assert!(lossy.decode_label_name(b"label-name").is_err());
2566        assert!(lossy.decode_label_name(b"0abc").is_err());
2567
2568        // Invalid UTF-8 bytes fail the label-name byte check.
2569        assert!(lossy.decode_label_name(&[0xff, 0xfe]).is_err());
2570    }
2571
2572    #[test]
2573    fn test_decode_label_name_unchecked() {
2574        let unchecked = PromValidationMode::Unchecked;
2575
2576        // Label name validation is always enforced.
2577        assert!(unchecked.decode_label_name(b"__name__").is_ok());
2578        assert!(unchecked.decode_label_name(b"label-name").is_err());
2579        assert!(unchecked.decode_label_name(b"0abc").is_err());
2580    }
2581
2582    #[test]
2583    fn test_is_valid_prom_label_name_bytes() {
2584        assert!(validate_label_name(b"__name__"));
2585        assert!(validate_label_name(b"job"));
2586        assert!(validate_label_name(b"_"));
2587        assert!(validate_label_name(b"A"));
2588        assert!(validate_label_name(b"abc123"));
2589        assert!(validate_label_name(b"_leading_underscore"));
2590
2591        assert!(!validate_label_name(b""));
2592        assert!(!validate_label_name(b"0starts_with_digit"));
2593        assert!(!validate_label_name(b"has-dash"));
2594        assert!(!validate_label_name(b"has.dot"));
2595        assert!(!validate_label_name(b"has space"));
2596        assert!(!validate_label_name(&[0xff, 0xfe]));
2597    }
2598}
2599
2600#[cfg(test)]
2601mod batching_tests {
2602    use std::sync::{Arc, Mutex};
2603
2604    use axum::http::StatusCode;
2605    use common_query::Output;
2606    use session::context::QueryContextRef;
2607
2608    use crate::batcher::BatchingProtocol;
2609    use crate::error::Result as ServerResult;
2610    use crate::http::test_helpers::TestClient;
2611    use crate::http::{HttpOptions, HttpServerBuilder};
2612    use crate::influxdb::InfluxdbRequest;
2613    use crate::opentsdb::codec::DataPoint;
2614    use crate::query_handler::{InfluxdbLineProtocolHandler, OpentsdbProtocolHandler};
2615
2616    #[tokio::test]
2617    async fn test_batching_selectors_are_independent() {
2618        use axum::routing::post;
2619        use axum::{Extension, Json, Router, middleware};
2620        use session::context::{QueryContext, QueryContextBuilder};
2621
2622        use crate::http::{LogicalBatchingProtocols, set_http_write_batching};
2623        for table in [false, true] {
2624            for logical in [false, true] {
2625                let app = Router::new()
2626                    .route(
2627                        "/",
2628                        post(|Extension(ctx): Extension<QueryContext>| async move {
2629                            Json([ctx.batching_enabled(), ctx.logical_batching_enabled()])
2630                        }),
2631                    )
2632                    .route_layer(middleware::from_fn_with_state(
2633                        BatchingProtocol::Otlp,
2634                        set_http_write_batching,
2635                    ))
2636                    .layer(Extension(QueryContextBuilder::default().build()))
2637                    .layer(Extension(Arc::new(if table {
2638                        vec![BatchingProtocol::Otlp]
2639                    } else {
2640                        vec![]
2641                    })))
2642                    .layer(Extension(LogicalBatchingProtocols(if logical {
2643                        vec![BatchingProtocol::Otlp]
2644                    } else {
2645                        vec![]
2646                    })));
2647                let client = TestClient::new(app).await;
2648                let actual: [bool; 2] = client.post("/").send().await.json().await;
2649                assert_eq!(actual, [table, logical]);
2650            }
2651        }
2652    }
2653
2654    #[derive(Default)]
2655    struct RecordingWriteHandler {
2656        selections: Mutex<Vec<bool>>,
2657    }
2658
2659    #[async_trait::async_trait]
2660    impl OpentsdbProtocolHandler for RecordingWriteHandler {
2661        async fn preflight(&self, _: &[DataPoint], _: QueryContextRef) -> ServerResult<()> {
2662            Ok(())
2663        }
2664
2665        async fn exec(&self, points: Vec<DataPoint>, ctx: QueryContextRef) -> ServerResult<usize> {
2666            self.selections.lock().unwrap().push(ctx.batching_enabled());
2667            Ok(points.len())
2668        }
2669    }
2670
2671    #[async_trait::async_trait]
2672    impl InfluxdbLineProtocolHandler for RecordingWriteHandler {
2673        async fn exec(&self, _: InfluxdbRequest, ctx: QueryContextRef) -> ServerResult<Output> {
2674            self.selections.lock().unwrap().push(ctx.batching_enabled());
2675            Ok(Output::new_with_affected_rows(1))
2676        }
2677    }
2678
2679    #[tokio::test]
2680    async fn test_real_influx_and_opentsdb_routes_keep_selection() {
2681        for protocols in [
2682            vec![],
2683            vec![BatchingProtocol::Influxdb],
2684            vec![BatchingProtocol::Opentsdb],
2685            vec![BatchingProtocol::Influxdb, BatchingProtocol::Opentsdb],
2686        ] {
2687            let handler = Arc::new(RecordingWriteHandler::default());
2688            let influx_enabled = protocols.contains(&BatchingProtocol::Influxdb);
2689            let opentsdb_enabled = protocols.contains(&BatchingProtocol::Opentsdb);
2690            let server = HttpServerBuilder::new(HttpOptions::default())
2691                .with_batching_protocols(protocols)
2692                .with_influxdb_handler(handler.clone())
2693                .with_opentsdb_handler(handler.clone())
2694                .build();
2695            let client = TestClient::new(server.build(server.make_app()).unwrap()).await;
2696            for path in [
2697                "/v1/influxdb/write",
2698                "/v1/influxdb/api/v2/write?bucket=public",
2699            ] {
2700                assert_eq!(
2701                    client
2702                        .post(path)
2703                        .body("cpu value=1 42")
2704                        .send()
2705                        .await
2706                        .status(),
2707                    StatusCode::NO_CONTENT
2708                );
2709            }
2710            let point =
2711                serde_json::json!({"metric":"cpu", "timestamp":42, "value":1, "tags":{"host":"a"}});
2712            for query in ["", "?summary", "?details"] {
2713                let response = client
2714                    .post(&format!("/v1/opentsdb/api/put{query}"))
2715                    .json(&point)
2716                    .send()
2717                    .await;
2718                assert_eq!(
2719                    response.status(),
2720                    if query.is_empty() {
2721                        StatusCode::NO_CONTENT
2722                    } else {
2723                        StatusCode::OK
2724                    }
2725                );
2726            }
2727            assert_eq!(
2728                *handler.selections.lock().unwrap(),
2729                vec![
2730                    influx_enabled,
2731                    influx_enabled,
2732                    opentsdb_enabled,
2733                    opentsdb_enabled,
2734                    opentsdb_enabled
2735                ]
2736            );
2737        }
2738    }
2739}