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            .route(
1580                "/api/v3/services",
1581                routing::get(jaeger::v3::handle_get_services),
1582            )
1583            .route(
1584                "/api/v3/operations",
1585                routing::get(jaeger::v3::handle_get_operations),
1586            )
1587            .route(
1588                "/api/v3/traces",
1589                routing::get(jaeger::v3::handle_find_traces),
1590            )
1591            .route(
1592                "/api/v3/trace-summaries",
1593                routing::get(jaeger::v3::handle_find_trace_summaries),
1594            )
1595            .route(
1596                "/api/v3/traces/{trace_id}",
1597                routing::get(jaeger::v3::handle_get_trace),
1598            )
1599            .with_state(handler)
1600    }
1601
1602    #[cfg(feature = "dashboard")]
1603    fn route_dashboard<S>(
1604        handler: DashboardHandlerRef,
1605        memory_limiter: ServerMemoryLimiter,
1606    ) -> Router<S> {
1607        use crate::http::dashboard::{add_dashboard, delete_dashboard, list_dashboards};
1608
1609        Router::new()
1610            .route("/", routing::get(list_dashboards))
1611            .route("/{dashboard_name}", routing::post(add_dashboard))
1612            .route("/{dashboard_name}", routing::delete(delete_dashboard))
1613            .layer(middleware::from_fn_with_state(
1614                memory_limiter,
1615                memory_limit::decoded_body_accounting_middleware,
1616            ))
1617            .layer(
1618                ServiceBuilder::new()
1619                    .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)),
1620            )
1621            .with_state(DashboardState { handler })
1622    }
1623
1624    #[cfg(not(feature = "dashboard"))]
1625    fn route_dashboard<S>(
1626        handler: DashboardHandlerRef,
1627        _memory_limiter: ServerMemoryLimiter,
1628    ) -> Router<S> {
1629        Router::new().with_state(DashboardState { handler })
1630    }
1631}
1632
1633pub const HTTP_SERVER: &str = "HTTP_SERVER";
1634pub const HTTP_API_SERVER: &str = "HTTP_API_SERVER";
1635
1636/// An empty list allows any origin.
1637pub(crate) fn cors_allow_origin(origins: &[String]) -> Result<AllowOrigin> {
1638    if origins.is_empty() {
1639        return Ok(AllowOrigin::any());
1640    }
1641    let origins = origins
1642        .iter()
1643        .map(|origin| HeaderValue::from_str(origin).context(InvalidHeaderValueSnafu))
1644        .collect::<Result<Vec<_>>>()?;
1645    Ok(AllowOrigin::list(origins))
1646}
1647
1648#[async_trait]
1649impl Server for HttpServer {
1650    async fn shutdown(&self) -> Result<()> {
1651        let mut shutdown_tx = self.shutdown_tx.lock().await;
1652        if let Some(tx) = shutdown_tx.take()
1653            && tx.send(()).is_err()
1654        {
1655            info!("Receiver dropped, the HTTP server has already exited");
1656        }
1657        info!("Shutdown {}", self.kind());
1658
1659        Ok(())
1660    }
1661
1662    async fn start(&mut self, listening: SocketAddr) -> Result<()> {
1663        let (tx, rx) = oneshot::channel();
1664        let serve = {
1665            let mut shutdown_tx = self.shutdown_tx.lock().await;
1666            ensure!(
1667                shutdown_tx.is_none(),
1668                AlreadyStartedSnafu {
1669                    server: self.kind()
1670                }
1671            );
1672
1673            let app = self.build(self.make_app())?;
1674            let listener = tokio::net::TcpListener::bind(listening)
1675                .await
1676                .context(AddressBindSnafu { addr: listening })?
1677                .tap_io(|tcp_stream| {
1678                    if let Err(e) = tcp_stream.set_nodelay(true) {
1679                        error!(e; "Failed to set TCP_NODELAY on incoming connection");
1680                    }
1681                });
1682            let serve = axum::serve(
1683                listener,
1684                app.into_make_service_with_connect_info::<SocketAddr>(),
1685            );
1686
1687            // FIXME(yingwen): Support keepalive.
1688            // See:
1689            // - https://github.com/tokio-rs/axum/discussions/2939
1690            // - https://stackoverflow.com/questions/73069718/how-do-i-keep-alive-tokiotcpstream-in-rust
1691            // let server = axum::Server::try_bind(&listening)
1692            //     .with_context(|_| AddressBindSnafu { addr: listening })?
1693            //     .tcp_nodelay(true)
1694            //     // Enable TCP keepalive to close the dangling established connections.
1695            //     // It's configured to let the keepalive probes first send after the connection sits
1696            //     // idle for 59 minutes, and then send every 10 seconds for 6 times.
1697            //     // So the connection will be closed after roughly 1 hour.
1698            //     .tcp_keepalive(Some(Duration::from_secs(59 * 60)))
1699            //     .tcp_keepalive_interval(Some(Duration::from_secs(10)))
1700            //     .tcp_keepalive_retries(Some(6))
1701            //     .serve(app.into_make_service());
1702
1703            *shutdown_tx = Some(tx);
1704
1705            serve
1706        };
1707        let listening = serve.local_addr().context(InternalIoSnafu)?;
1708        info!("{} server is bound to {}", self.kind(), listening);
1709
1710        common_runtime::spawn_global(async move {
1711            if let Err(e) = serve
1712                .with_graceful_shutdown(rx.map(drop))
1713                .await
1714                .context(InternalIoSnafu)
1715            {
1716                error!(e; "Failed to shutdown http server");
1717            }
1718        });
1719
1720        self.bind_addr = Some(listening);
1721        Ok(())
1722    }
1723
1724    fn name(&self) -> &str {
1725        match self.kind {
1726            HttpServerKind::Api => HTTP_API_SERVER,
1727            HttpServerKind::Full => HTTP_SERVER,
1728        }
1729    }
1730
1731    fn bind_addr(&self) -> Option<SocketAddr> {
1732        self.bind_addr
1733    }
1734
1735    fn as_any(&self) -> &dyn std::any::Any {
1736        self
1737    }
1738}
1739
1740#[cfg(test)]
1741mod test {
1742    use std::future::pending;
1743    use std::io::Cursor;
1744    use std::sync::Arc;
1745
1746    use arrow_ipc::reader::StreamReader;
1747    use arrow_schema::DataType;
1748    use axum::http::StatusCode;
1749    use axum::routing::{get, post};
1750    use common_query::{Output, OutputData};
1751    use common_recordbatch::RecordBatches;
1752    use datafusion_expr::LogicalPlan;
1753    use datatypes::prelude::*;
1754    use datatypes::schema::{ColumnSchema, Schema};
1755    use datatypes::vectors::{StringVector, UInt32Vector};
1756    use header::constants::GREPTIME_DB_HEADER_TIMEOUT;
1757    use query::parser::PromQuery;
1758    use query::query_engine::DescribeResult;
1759    use session::context::QueryContextRef;
1760    use sql::statements::statement::Statement;
1761    use tokio::sync::mpsc;
1762    use tokio::time::Instant;
1763
1764    use crate::http::test_helpers::TestClient;
1765    use crate::http::*;
1766    use crate::prom_remote_write::validation::validate_label_name;
1767    use crate::query_handler::sql::SqlQueryHandler;
1768
1769    struct DummyInstance {
1770        _tx: mpsc::Sender<(String, Vec<u8>)>,
1771    }
1772
1773    #[async_trait]
1774    impl SqlQueryHandler for DummyInstance {
1775        async fn do_query(&self, _: &str, _: QueryContextRef) -> Vec<Result<Output>> {
1776            unimplemented!()
1777        }
1778
1779        async fn do_analyze_stream_query(&self, _: &str, _: QueryContextRef) -> Result<Output> {
1780            let stream = common_recordbatch::RecordBatches::empty().as_stream();
1781            Ok(Output::new(OutputData::Stream(stream), Default::default()))
1782        }
1783
1784        async fn do_promql_query(&self, _: &PromQuery, _: QueryContextRef) -> Vec<Result<Output>> {
1785            unimplemented!()
1786        }
1787
1788        async fn do_exec_plan(
1789            &self,
1790            _plan: LogicalPlan,
1791            _stmt: Option<Statement>,
1792            _query_ctx: QueryContextRef,
1793        ) -> Result<Output> {
1794            unimplemented!()
1795        }
1796
1797        async fn do_describe(
1798            &self,
1799            _stmt: sql::statements::statement::Statement,
1800            _query_ctx: QueryContextRef,
1801        ) -> Result<Option<DescribeResult>> {
1802            unimplemented!()
1803        }
1804
1805        async fn is_valid_schema(&self, _catalog: &str, _schema: &str) -> Result<bool> {
1806            Ok(true)
1807        }
1808    }
1809
1810    async fn forever() {
1811        pending().await
1812    }
1813
1814    fn make_test_app(tx: mpsc::Sender<(String, Vec<u8>)>) -> Router {
1815        make_test_app_custom(tx, HttpOptions::default())
1816    }
1817
1818    fn make_test_app_custom(tx: mpsc::Sender<(String, Vec<u8>)>, options: HttpOptions) -> Router {
1819        let instance = Arc::new(DummyInstance { _tx: tx });
1820        let server = HttpServerBuilder::new(options)
1821            .with_sql_handler(instance.clone())
1822            .build();
1823        let app = server
1824            .make_app()
1825            .route("/test/timeout", get(forever))
1826            .route("/v1/prometheus/write", post(forever));
1827        server.build(app).unwrap()
1828    }
1829
1830    fn make_split_builder() -> HttpServerBuilder {
1831        let (tx, _rx) = mpsc::channel(100);
1832        let instance = Arc::new(DummyInstance { _tx: tx });
1833        // Enable the dedicated API server so the split behavior can be exercised.
1834        let options = HttpOptions {
1835            enable_api_server: true,
1836            ..HttpOptions::default()
1837        };
1838        HttpServerBuilder::new(options)
1839            .with_sql_handler(instance)
1840            .with_metrics_handler(MetricsHandler)
1841            .with_greptime_config_options("dummy = \"value\"".to_string())
1842    }
1843
1844    #[tokio::test]
1845    pub async fn test_http_api_options_defaults() {
1846        let opts = HttpOptions::default();
1847        // The API server is disabled by default.
1848        assert!(!opts.enable_api_server);
1849        assert_eq!(opts.api_server_addr, "127.0.0.1:4006");
1850    }
1851
1852    #[tokio::test]
1853    pub async fn test_build_serves_all_routes() {
1854        // The single (internal/full) server built via `build()` keeps serving both
1855        // the public v1 interfaces and the internal ones. This is the historical
1856        // single-server behavior and must not change.
1857        let server = make_split_builder().build();
1858        assert_eq!(server.name(), HTTP_SERVER);
1859
1860        let app = server.build(server.make_app()).unwrap();
1861        let client = TestClient::new(app).await;
1862
1863        for path in ["/v1/health", "/health", "/status", "/metrics", "/config"] {
1864            let res = client.get(path).send().await;
1865            assert_eq!(
1866                res.status(),
1867                StatusCode::OK,
1868                "internal/full server should serve {path}"
1869            );
1870        }
1871    }
1872
1873    #[tokio::test]
1874    pub async fn test_build_servers_separates_api_and_internal_routes() {
1875        // `build_servers` produces the internal/full server (everything) plus a
1876        // dedicated API server that only serves the v1 interfaces (and dashboard).
1877        let (internal, api) = make_split_builder().build_servers();
1878        let api = api.expect("API server is explicitly enabled in make_split_builder");
1879        assert_eq!(internal.name(), HTTP_SERVER);
1880        assert_eq!(api.name(), HTTP_API_SERVER);
1881
1882        let internal_app = internal.build(internal.make_app()).unwrap();
1883        let api_app = api.build(api.make_app()).unwrap();
1884
1885        let internal_client = TestClient::new(internal_app).await;
1886        let api_client = TestClient::new(api_app).await;
1887
1888        // Both servers serve the v1 interfaces (e.g. the v1 health check).
1889        assert_eq!(
1890            internal_client.get("/v1/health").send().await.status(),
1891            StatusCode::OK
1892        );
1893        assert_eq!(
1894            api_client.get("/v1/health").send().await.status(),
1895            StatusCode::OK
1896        );
1897
1898        // The internal-only interfaces are served by the internal/full server...
1899        for path in ["/health", "/status", "/metrics", "/config"] {
1900            assert_eq!(
1901                internal_client.get(path).send().await.status(),
1902                StatusCode::OK,
1903                "internal/full server should serve {path}"
1904            );
1905            // ...and NOT by the API-only server.
1906            assert_eq!(
1907                api_client.get(path).send().await.status(),
1908                StatusCode::NOT_FOUND,
1909                "API server should NOT serve {path}"
1910            );
1911        }
1912    }
1913
1914    #[test]
1915    fn test_build_servers_api_inherits_http_options() {
1916        // The API server must inherit every `[http]` option except the bound addr.
1917        let http_opts = HttpOptions {
1918            timeout: Duration::from_secs(42),
1919            body_limit: ReadableSize::mb(128),
1920            cors_allowed_origins: vec!["https://example.com".to_string()],
1921            enable_api_server: true,
1922            ..HttpOptions::default()
1923        };
1924        let (internal, api) = HttpServerBuilder::new(http_opts.clone()).build_servers();
1925        let api = api.expect("API server is explicitly enabled");
1926
1927        // Only the bound address differs.
1928        assert_eq!(internal.options.addr, http_opts.addr);
1929        assert_eq!(api.options.addr, "127.0.0.1:4006");
1930        // Everything else is inherited from `[http]`.
1931        assert_eq!(api.options.timeout, http_opts.timeout);
1932        assert_eq!(api.options.body_limit, http_opts.body_limit);
1933        assert_eq!(
1934            api.options.cors_allowed_origins,
1935            http_opts.cors_allowed_origins
1936        );
1937    }
1938
1939    #[test]
1940    fn test_is_api_listener_path() {
1941        // `/v1` namespace (exact slash boundary).
1942        assert!(is_api_listener_path("/v1"));
1943        assert!(is_api_listener_path("/v1/"));
1944        assert!(is_api_listener_path("/v1/sql"));
1945        assert!(!is_api_listener_path("/v10/sql"));
1946        assert!(!is_api_listener_path("/v1-internal"));
1947        // `/dashboard` namespace.
1948        assert!(is_api_listener_path("/dashboard"));
1949        assert!(is_api_listener_path("/dashboard/app.js"));
1950        assert!(!is_api_listener_path("/dashboard-admin"));
1951        // Operational paths are never API-visible.
1952        assert!(!is_api_listener_path("/metrics"));
1953        assert!(!is_api_listener_path("/status/plugin"));
1954        assert!(!is_api_listener_path("/health"));
1955    }
1956
1957    #[tokio::test]
1958    async fn packed_capability_requires_auth_on_full_and_api_listeners() {
1959        let (tx, _rx) = mpsc::channel(1);
1960        let provider =
1961            auth::static_user_provider_from_option("static_user_provider:cmd:user=password")
1962                .unwrap();
1963        let (full, api) = HttpServerBuilder::new(HttpOptions {
1964            enable_api_server: true,
1965            ..Default::default()
1966        })
1967        .with_sql_handler(Arc::new(DummyInstance { _tx: tx }))
1968        .with_user_provider(Arc::new(provider))
1969        .build_servers();
1970        for server in [full, api.unwrap()] {
1971            let client = TestClient::new(server.build(server.make_app()).unwrap()).await;
1972            assert_eq!(
1973                client.get("/v1/capabilities").send().await.status(),
1974                StatusCode::UNAUTHORIZED
1975            );
1976            assert_eq!(
1977                client
1978                    .post("/v1/ddl/logical-tables")
1979                    .form(&[("sql", "SELECT 1")])
1980                    .send()
1981                    .await
1982                    .status(),
1983                StatusCode::UNAUTHORIZED
1984            );
1985            let response = client
1986                .get("/v1/capabilities")
1987                .header("Authorization", "Basic dXNlcjpwYXNzd29yZA==")
1988                .send()
1989                .await;
1990            assert_eq!(response.status(), StatusCode::OK);
1991            assert_eq!(
1992                response.json::<serde_json::Value>().await,
1993                serde_json::json!({"metric_packed_import": 1, "metric_packed_export": 1})
1994            );
1995        }
1996    }
1997
1998    #[tokio::test]
1999    async fn logical_ddl_form_limits_and_unsupported_handler() {
2000        use axum::body::{Body, to_bytes};
2001        use tower::ServiceExt;
2002
2003        async fn post_form(app: Router, body: String) -> (StatusCode, serde_json::Value) {
2004            let response = app
2005                .oneshot(
2006                    Request::builder()
2007                        .method(Method::POST)
2008                        .uri("/v1/ddl/logical-tables")
2009                        .header("Content-Type", "application/x-www-form-urlencoded")
2010                        .body(Body::from(body))
2011                        .unwrap(),
2012                )
2013                .await
2014                .unwrap();
2015            let status = response.status();
2016            let body = to_bytes(response.into_body(), 4096).await.unwrap();
2017            (status, serde_json::from_slice(&body).unwrap())
2018        }
2019
2020        for limit in [0, 64] {
2021            let (tx, _rx) = mpsc::channel(1);
2022            let server = HttpServerBuilder::new(HttpOptions {
2023                body_limit: ReadableSize(limit),
2024                ..Default::default()
2025            })
2026            .with_sql_handler(Arc::new(DummyInstance { _tx: tx }))
2027            .build();
2028            let app = server.build(server.make_app()).unwrap();
2029            let client = TestClient::new(app.clone()).await;
2030            assert_eq!(
2031                client.get("/v1/ddl/logical-tables").send().await.status(),
2032                StatusCode::METHOD_NOT_ALLOWED
2033            );
2034            let small = client
2035                .post("/v1/ddl/logical-tables")
2036                .form(&[("sql", "CREATE TABLE t")])
2037                .send()
2038                .await;
2039            assert_eq!(
2040                small.json::<serde_json::Value>().await["code"],
2041                common_error::status_code::StatusCode::Unsupported as u32
2042            );
2043            let (_, body) = post_form(
2044                app.clone(),
2045                if limit == 0 {
2046                    format!("sql={}", "%20".repeat(1024 * 1024))
2047                } else {
2048                    "sql=".to_string() + &" ".repeat(65)
2049                },
2050            )
2051            .await;
2052            let code = body["code"].clone();
2053            assert_eq!(
2054                code,
2055                if limit == 0 {
2056                    common_error::status_code::StatusCode::Unsupported as u32
2057                } else {
2058                    common_error::status_code::StatusCode::InvalidArguments as u32
2059                }
2060            );
2061            for body in [
2062                format!("sql={}", "x".repeat(1024 * 1024 + 1)),
2063                format!("sql={}", "%C3%A9".repeat(512 * 1024 + 1)),
2064                format!("sql=CREATE+TABLE+t&padding={}", "x".repeat(4 * 1024 * 1024)),
2065            ] {
2066                let (status, body) = post_form(app.clone(), body).await;
2067                assert_eq!(status, StatusCode::BAD_REQUEST);
2068                assert_eq!(
2069                    body["code"],
2070                    common_error::status_code::StatusCode::InvalidArguments as u32
2071                );
2072            }
2073        }
2074    }
2075
2076    #[tokio::test]
2077    pub async fn test_extra_router_respects_api_surface() {
2078        // Downstream routes added through `with_extra_router()` (e.g. enterprise
2079        // plugins or the `RouterConfigurator`) are present on the full listener
2080        // and on the API listener's router, but the API listener's outer guard
2081        // hides every non-`/v1` path automatically. This is what keeps the
2082        // feature correct for greptimedb-enterprise with no downstream changes.
2083        let (tx, _rx) = mpsc::channel(100);
2084        let instance = Arc::new(DummyInstance { _tx: tx });
2085        let options = HttpOptions {
2086            enable_api_server: true,
2087            ..HttpOptions::default()
2088        };
2089        let builder = HttpServerBuilder::new(options)
2090            .with_sql_handler(instance)
2091            .with_extra_router(
2092                Router::new()
2093                    .route("/status/plugin", routing::get(|| async { "ok" }))
2094                    .route("/v1/plugin", routing::get(|| async { "ok" })),
2095            );
2096        let (full, api) = builder.build_servers();
2097        let api = api.expect("API server is explicitly enabled");
2098
2099        let full_client = TestClient::new(full.build(full.make_app()).unwrap()).await;
2100        let api_client = TestClient::new(api.build(api.make_app()).unwrap()).await;
2101
2102        // A non-`/v1` route is reachable on the full listener...
2103        assert_eq!(
2104            full_client.get("/status/plugin").send().await.status(),
2105            StatusCode::OK
2106        );
2107        // ...and hidden (404) on the API listener, without the caller classifying it.
2108        assert_eq!(
2109            api_client.get("/status/plugin").send().await.status(),
2110            StatusCode::NOT_FOUND
2111        );
2112
2113        // A `/v1` route is reachable on both.
2114        assert_eq!(
2115            full_client.get("/v1/plugin").send().await.status(),
2116            StatusCode::OK
2117        );
2118        assert_eq!(
2119            api_client.get("/v1/plugin").send().await.status(),
2120            StatusCode::OK
2121        );
2122
2123        // Slash-boundary correctness on the API listener.
2124        assert_eq!(
2125            api_client.get("/v10/plugin").send().await.status(),
2126            StatusCode::NOT_FOUND
2127        );
2128    }
2129
2130    #[test]
2131    fn test_build_servers_api_disabled_by_default() {
2132        // The API server is disabled by default, so `build_servers` returns no
2133        // second server unless `enable_api_server` is set.
2134        let (_internal, api) = HttpServerBuilder::new(HttpOptions::default()).build_servers();
2135        assert!(api.is_none());
2136    }
2137
2138    #[tokio::test]
2139    pub async fn test_cors() {
2140        // cors is on by default
2141        let (tx, _rx) = mpsc::channel(100);
2142        let app = make_test_app(tx);
2143        let client = TestClient::new(app).await;
2144
2145        let res = client.get("/health").send().await;
2146
2147        assert_eq!(res.status(), StatusCode::OK);
2148        assert_eq!(
2149            res.headers()
2150                .get(http::header::ACCESS_CONTROL_ALLOW_ORIGIN)
2151                .expect("expect cors header origin"),
2152            "*"
2153        );
2154
2155        let res = client.get("/v1/health").send().await;
2156
2157        assert_eq!(res.status(), StatusCode::OK);
2158        assert_eq!(
2159            res.headers()
2160                .get(http::header::ACCESS_CONTROL_ALLOW_ORIGIN)
2161                .expect("expect cors header origin"),
2162            "*"
2163        );
2164
2165        let res = client
2166            .options("/health")
2167            .header("Access-Control-Request-Headers", "x-greptime-auth")
2168            .header("Access-Control-Request-Method", "DELETE")
2169            .header("Origin", "https://example.com")
2170            .send()
2171            .await;
2172        assert_eq!(res.status(), StatusCode::OK);
2173        assert_eq!(
2174            res.headers()
2175                .get(http::header::ACCESS_CONTROL_ALLOW_ORIGIN)
2176                .expect("expect cors header origin"),
2177            "*"
2178        );
2179        assert_eq!(
2180            res.headers()
2181                .get(http::header::ACCESS_CONTROL_ALLOW_HEADERS)
2182                .expect("expect cors header headers"),
2183            "*"
2184        );
2185        assert_eq!(
2186            res.headers()
2187                .get(http::header::ACCESS_CONTROL_ALLOW_METHODS)
2188                .expect("expect cors header methods"),
2189            "GET,POST,PUT,DELETE,HEAD"
2190        );
2191    }
2192
2193    #[tokio::test]
2194    pub async fn test_cors_custom_origins() {
2195        // cors is on by default
2196        let (tx, _rx) = mpsc::channel(100);
2197        let origin = "https://example.com";
2198
2199        let options = HttpOptions {
2200            cors_allowed_origins: vec![origin.to_string()],
2201            ..Default::default()
2202        };
2203
2204        let app = make_test_app_custom(tx, options);
2205        let client = TestClient::new(app).await;
2206
2207        let res = client.get("/health").header("Origin", origin).send().await;
2208
2209        assert_eq!(res.status(), StatusCode::OK);
2210        assert_eq!(
2211            res.headers()
2212                .get(http::header::ACCESS_CONTROL_ALLOW_ORIGIN)
2213                .expect("expect cors header origin"),
2214            origin
2215        );
2216
2217        let res = client
2218            .get("/health")
2219            .header("Origin", "https://notallowed.com")
2220            .send()
2221            .await;
2222
2223        assert_eq!(res.status(), StatusCode::OK);
2224        assert!(
2225            !res.headers()
2226                .contains_key(http::header::ACCESS_CONTROL_ALLOW_ORIGIN)
2227        );
2228    }
2229
2230    #[tokio::test]
2231    pub async fn test_cors_disabled() {
2232        // cors is on by default
2233        let (tx, _rx) = mpsc::channel(100);
2234
2235        let options = HttpOptions {
2236            enable_cors: false,
2237            ..Default::default()
2238        };
2239
2240        let app = make_test_app_custom(tx, options);
2241        let client = TestClient::new(app).await;
2242
2243        let res = client.get("/health").send().await;
2244
2245        assert_eq!(res.status(), StatusCode::OK);
2246        assert!(
2247            !res.headers()
2248                .contains_key(http::header::ACCESS_CONTROL_ALLOW_ORIGIN)
2249        );
2250    }
2251
2252    #[test]
2253    fn test_http_options_default() {
2254        let default = HttpOptions::default();
2255        assert_eq!("127.0.0.1:4000".to_string(), default.addr);
2256        assert_eq!(Duration::from_secs(0), default.timeout)
2257    }
2258
2259    #[tokio::test]
2260    async fn test_http_options_legacy_analyze_stream_config_is_ignored() {
2261        let options: HttpOptions = serde_json::from_value(serde_json::json!({
2262            "addr": "127.0.0.1:4000",
2263            "timeout": "0s",
2264            "body_limit": "64MiB",
2265            "cors_allowed_origins": [],
2266            "enable_cors": true,
2267            "experimental_enable_explain_analyze_stream": false,
2268            "enable_api_server": false,
2269            "api_server_addr": "127.0.0.1:4006"
2270        }))
2271        .unwrap();
2272        let serialized = serde_json::to_string(&options).unwrap();
2273        assert!(!serialized.contains("experimental_enable_explain_analyze_stream"));
2274
2275        let (tx, _rx) = mpsc::channel(100);
2276        let app = make_test_app_custom(tx, options);
2277        let client = TestClient::new(app).await;
2278        let response = client
2279            .post("/v1/sql/analyze/stream")
2280            .form(&handler::SqlQuery {
2281                sql: Some("EXPLAIN ANALYZE VERBOSE SELECT 1".to_string()),
2282                ..Default::default()
2283            })
2284            .send()
2285            .await;
2286        assert_ne!(response.status(), StatusCode::NOT_FOUND);
2287    }
2288
2289    #[tokio::test]
2290    async fn test_http_server_request_timeout() {
2291        common_telemetry::init_default_ut_logging();
2292
2293        let (tx, _rx) = mpsc::channel(100);
2294        let options = HttpOptions {
2295            timeout: Duration::from_millis(10),
2296            ..Default::default()
2297        };
2298        let app = make_test_app_custom(tx, options);
2299        let client = TestClient::new(app).await;
2300        let res = client.get("/test/timeout").send().await;
2301        assert_eq!(res.status(), StatusCode::REQUEST_TIMEOUT);
2302
2303        let res = client.post("/v1/prometheus/write").send().await;
2304        assert_eq!(res.status(), StatusCode::GATEWAY_TIMEOUT);
2305
2306        let now = Instant::now();
2307        let res = client
2308            .get("/test/timeout")
2309            .header(GREPTIME_DB_HEADER_TIMEOUT, "20ms")
2310            .send()
2311            .await;
2312        assert_eq!(res.status(), StatusCode::REQUEST_TIMEOUT);
2313        let elapsed = now.elapsed();
2314        assert!(elapsed > Duration::from_millis(15));
2315
2316        tokio::time::timeout(
2317            Duration::from_millis(15),
2318            client
2319                .get("/test/timeout")
2320                .header(GREPTIME_DB_HEADER_TIMEOUT, "0s")
2321                .send(),
2322        )
2323        .await
2324        .unwrap_err();
2325
2326        tokio::time::timeout(
2327            Duration::from_millis(15),
2328            client
2329                .get("/test/timeout")
2330                .header(
2331                    GREPTIME_DB_HEADER_TIMEOUT,
2332                    humantime::format_duration(Duration::default()).to_string(),
2333                )
2334                .send(),
2335        )
2336        .await
2337        .unwrap_err();
2338    }
2339
2340    #[tokio::test]
2341    async fn test_schema_for_empty_response() {
2342        let column_schemas = vec![
2343            ColumnSchema::new("numbers", ConcreteDataType::uint32_datatype(), false),
2344            ColumnSchema::new("strings", ConcreteDataType::string_datatype(), true),
2345        ];
2346        let schema = Arc::new(Schema::new(column_schemas));
2347
2348        let recordbatches = RecordBatches::try_new(schema.clone(), vec![]).unwrap();
2349        let outputs = vec![Ok(Output::new_with_record_batches(recordbatches))];
2350
2351        let json_resp = GreptimedbV1Response::from_output(outputs).await;
2352        if let HttpResponse::GreptimedbV1(json_resp) = json_resp {
2353            let json_output = &json_resp.output[0];
2354            if let GreptimeQueryOutput::Records(r) = json_output {
2355                assert_eq!(r.num_rows(), 0);
2356                assert_eq!(r.num_cols(), 2);
2357                assert_eq!(r.schema.column_schemas[0].name, "numbers");
2358                assert_eq!(r.schema.column_schemas[0].data_type, "UInt32");
2359            } else {
2360                panic!("invalid output type");
2361            }
2362        } else {
2363            panic!("invalid format")
2364        }
2365    }
2366
2367    #[tokio::test]
2368    async fn test_recordbatches_conversion() {
2369        let column_schemas = vec![
2370            ColumnSchema::new("numbers", ConcreteDataType::uint32_datatype(), false),
2371            ColumnSchema::new("strings", ConcreteDataType::string_datatype(), true),
2372        ];
2373        let schema = Arc::new(Schema::new(column_schemas));
2374        let columns: Vec<VectorRef> = vec![
2375            Arc::new(UInt32Vector::from_slice(vec![1, 2, 3, 4])),
2376            Arc::new(StringVector::from(vec![
2377                None,
2378                Some("hello"),
2379                Some("greptime"),
2380                None,
2381            ])),
2382        ];
2383        let recordbatch = RecordBatch::new(schema.clone(), columns).unwrap();
2384
2385        for format in [
2386            ResponseFormat::GreptimedbV1,
2387            ResponseFormat::InfluxdbV1,
2388            ResponseFormat::Csv(true, true),
2389            ResponseFormat::Table,
2390            ResponseFormat::Arrow,
2391            ResponseFormat::Json,
2392            ResponseFormat::Null,
2393        ] {
2394            let recordbatches =
2395                RecordBatches::try_new(schema.clone(), vec![recordbatch.clone()]).unwrap();
2396            let outputs = vec![Ok(Output::new_with_record_batches(recordbatches))];
2397            let json_resp = match format {
2398                ResponseFormat::Arrow => ArrowResponse::from_output(outputs, None).await,
2399                ResponseFormat::Csv(with_names, with_types) => {
2400                    CsvResponse::from_output(outputs, with_names, with_types).await
2401                }
2402                ResponseFormat::Table => TableResponse::from_output(outputs).await,
2403                ResponseFormat::GreptimedbV1 => GreptimedbV1Response::from_output(outputs).await,
2404                ResponseFormat::InfluxdbV1 => InfluxdbV1Response::from_output(outputs, None).await,
2405                ResponseFormat::Json => JsonResponse::from_output(outputs).await,
2406                ResponseFormat::Null => NullResponse::from_output(outputs).await,
2407            };
2408
2409            match json_resp {
2410                HttpResponse::GreptimedbV1(resp) => {
2411                    let json_output = &resp.output[0];
2412                    if let GreptimeQueryOutput::Records(r) = json_output {
2413                        assert_eq!(r.num_rows(), 4);
2414                        assert_eq!(r.num_cols(), 2);
2415                        assert_eq!(r.schema.column_schemas[0].name, "numbers");
2416                        assert_eq!(r.schema.column_schemas[0].data_type, "UInt32");
2417                        assert_eq!(r.rows[0][0], serde_json::Value::from(1));
2418                        assert_eq!(r.rows[0][1], serde_json::Value::Null);
2419                    } else {
2420                        panic!("invalid output type");
2421                    }
2422                }
2423                HttpResponse::InfluxdbV1(resp) => {
2424                    let json_output = &resp.results()[0];
2425                    assert_eq!(json_output.num_rows(), 4);
2426                    assert_eq!(json_output.num_cols(), 2);
2427                    assert_eq!(json_output.series[0].columns.clone()[0], "numbers");
2428                    assert_eq!(
2429                        json_output.series[0].values[0][0],
2430                        serde_json::Value::from(1)
2431                    );
2432                    assert_eq!(json_output.series[0].values[0][1], serde_json::Value::Null);
2433                }
2434                HttpResponse::Csv(resp) => {
2435                    let output = &resp.output()[0];
2436                    if let GreptimeQueryOutput::Records(r) = output {
2437                        assert_eq!(r.num_rows(), 4);
2438                        assert_eq!(r.num_cols(), 2);
2439                        assert_eq!(r.schema.column_schemas[0].name, "numbers");
2440                        assert_eq!(r.schema.column_schemas[0].data_type, "UInt32");
2441                        assert_eq!(r.rows[0][0], serde_json::Value::from(1));
2442                        assert_eq!(r.rows[0][1], serde_json::Value::Null);
2443                    } else {
2444                        panic!("invalid output type");
2445                    }
2446                }
2447
2448                HttpResponse::Table(resp) => {
2449                    let output = &resp.output()[0];
2450                    if let GreptimeQueryOutput::Records(r) = output {
2451                        assert_eq!(r.num_rows(), 4);
2452                        assert_eq!(r.num_cols(), 2);
2453                        assert_eq!(r.schema.column_schemas[0].name, "numbers");
2454                        assert_eq!(r.schema.column_schemas[0].data_type, "UInt32");
2455                        assert_eq!(r.rows[0][0], serde_json::Value::from(1));
2456                        assert_eq!(r.rows[0][1], serde_json::Value::Null);
2457                    } else {
2458                        panic!("invalid output type");
2459                    }
2460                }
2461
2462                HttpResponse::Arrow(resp) => {
2463                    let output = resp.data;
2464                    let mut reader = StreamReader::try_new(Cursor::new(output), None)
2465                        .expect("Arrow reader error");
2466                    let schema = reader.schema();
2467                    assert_eq!(schema.fields[0].name(), "numbers");
2468                    assert_eq!(schema.fields[0].data_type(), &DataType::UInt32);
2469                    assert_eq!(schema.fields[1].name(), "strings");
2470                    assert_eq!(schema.fields[1].data_type(), &DataType::Utf8);
2471
2472                    let rb = reader.next().unwrap().expect("read record batch failed");
2473                    assert_eq!(rb.num_columns(), 2);
2474                    assert_eq!(rb.num_rows(), 4);
2475                }
2476
2477                HttpResponse::Json(resp) => {
2478                    let output = &resp.output()[0];
2479                    if let GreptimeQueryOutput::Records(r) = output {
2480                        assert_eq!(r.num_rows(), 4);
2481                        assert_eq!(r.num_cols(), 2);
2482                        assert_eq!(r.schema.column_schemas[0].name, "numbers");
2483                        assert_eq!(r.schema.column_schemas[0].data_type, "UInt32");
2484                        assert_eq!(r.rows[0][0], serde_json::Value::from(1));
2485                        assert_eq!(r.rows[0][1], serde_json::Value::Null);
2486                    } else {
2487                        panic!("invalid output type");
2488                    }
2489                }
2490
2491                HttpResponse::Null(resp) => {
2492                    assert_eq!(resp.rows(), 4);
2493                }
2494
2495                HttpResponse::Error(err) => unreachable!("{err:?}"),
2496            }
2497        }
2498    }
2499
2500    #[test]
2501    fn test_response_format_misc() {
2502        assert_eq!(ResponseFormat::default(), ResponseFormat::GreptimedbV1);
2503        assert_eq!(ResponseFormat::parse("arrow"), Some(ResponseFormat::Arrow));
2504        assert_eq!(
2505            ResponseFormat::parse("csv"),
2506            Some(ResponseFormat::Csv(false, false))
2507        );
2508        assert_eq!(
2509            ResponseFormat::parse("csvwithnames"),
2510            Some(ResponseFormat::Csv(true, false))
2511        );
2512        assert_eq!(
2513            ResponseFormat::parse("csvwithnamesandtypes"),
2514            Some(ResponseFormat::Csv(true, true))
2515        );
2516        assert_eq!(ResponseFormat::parse("table"), Some(ResponseFormat::Table));
2517        assert_eq!(
2518            ResponseFormat::parse("greptimedb_v1"),
2519            Some(ResponseFormat::GreptimedbV1)
2520        );
2521        assert_eq!(
2522            ResponseFormat::parse("influxdb_v1"),
2523            Some(ResponseFormat::InfluxdbV1)
2524        );
2525        assert_eq!(ResponseFormat::parse("json"), Some(ResponseFormat::Json));
2526        assert_eq!(ResponseFormat::parse("null"), Some(ResponseFormat::Null));
2527
2528        // invalid formats
2529        assert_eq!(ResponseFormat::parse("invalid"), None);
2530        assert_eq!(ResponseFormat::parse(""), None);
2531        assert_eq!(ResponseFormat::parse("CSV"), None); // Case sensitive
2532
2533        // as str
2534        assert_eq!(ResponseFormat::Arrow.as_str(), "arrow");
2535        assert_eq!(ResponseFormat::Csv(false, false).as_str(), "csv");
2536        assert_eq!(ResponseFormat::Csv(true, true).as_str(), "csv");
2537        assert_eq!(ResponseFormat::Table.as_str(), "table");
2538        assert_eq!(ResponseFormat::GreptimedbV1.as_str(), "greptimedb_v1");
2539        assert_eq!(ResponseFormat::InfluxdbV1.as_str(), "influxdb_v1");
2540        assert_eq!(ResponseFormat::Json.as_str(), "json");
2541        assert_eq!(ResponseFormat::Null.as_str(), "null");
2542        assert_eq!(ResponseFormat::default().as_str(), "greptimedb_v1");
2543    }
2544
2545    #[test]
2546    fn test_decode_label_name_strict() {
2547        let strict = PromValidationMode::Strict;
2548
2549        // Valid Prometheus label names
2550        assert!(strict.decode_label_name(b"__name__").is_ok());
2551        assert!(strict.decode_label_name(b"job").is_ok());
2552        assert!(strict.decode_label_name(b"instance").is_ok());
2553        assert!(strict.decode_label_name(b"_private").is_ok());
2554        assert!(strict.decode_label_name(b"label_with_underscores").is_ok());
2555        assert!(strict.decode_label_name(b"abc123").is_ok());
2556        assert!(strict.decode_label_name(b"A").is_ok());
2557        assert!(strict.decode_label_name(b"_").is_ok());
2558
2559        // Invalid: starts with digit
2560        assert!(strict.decode_label_name(b"0abc").is_err());
2561        assert!(strict.decode_label_name(b"123").is_err());
2562
2563        // Invalid: contains special characters
2564        assert!(strict.decode_label_name(b"label-name").is_err());
2565        assert!(strict.decode_label_name(b"label.name").is_err());
2566        assert!(strict.decode_label_name(b"label name").is_err());
2567        assert!(strict.decode_label_name(b"label/name").is_err());
2568
2569        // Invalid: empty
2570        assert!(strict.decode_label_name(b"").is_err());
2571
2572        // Invalid: non-ASCII UTF-8
2573        assert!(strict.decode_label_name("ラベル".as_bytes()).is_err());
2574
2575        // Invalid UTF-8 bytes should fail
2576        assert!(strict.decode_label_name(&[0xff, 0xfe]).is_err());
2577    }
2578
2579    #[test]
2580    fn test_decode_label_name_lossy() {
2581        let lossy = PromValidationMode::Lossy;
2582
2583        // Label name validation is always enforced.
2584        assert!(lossy.decode_label_name(b"__name__").is_ok());
2585        assert!(lossy.decode_label_name(b"label-name").is_err());
2586        assert!(lossy.decode_label_name(b"0abc").is_err());
2587
2588        // Invalid UTF-8 bytes fail the label-name byte check.
2589        assert!(lossy.decode_label_name(&[0xff, 0xfe]).is_err());
2590    }
2591
2592    #[test]
2593    fn test_decode_label_name_unchecked() {
2594        let unchecked = PromValidationMode::Unchecked;
2595
2596        // Label name validation is always enforced.
2597        assert!(unchecked.decode_label_name(b"__name__").is_ok());
2598        assert!(unchecked.decode_label_name(b"label-name").is_err());
2599        assert!(unchecked.decode_label_name(b"0abc").is_err());
2600    }
2601
2602    #[test]
2603    fn test_is_valid_prom_label_name_bytes() {
2604        assert!(validate_label_name(b"__name__"));
2605        assert!(validate_label_name(b"job"));
2606        assert!(validate_label_name(b"_"));
2607        assert!(validate_label_name(b"A"));
2608        assert!(validate_label_name(b"abc123"));
2609        assert!(validate_label_name(b"_leading_underscore"));
2610
2611        assert!(!validate_label_name(b""));
2612        assert!(!validate_label_name(b"0starts_with_digit"));
2613        assert!(!validate_label_name(b"has-dash"));
2614        assert!(!validate_label_name(b"has.dot"));
2615        assert!(!validate_label_name(b"has space"));
2616        assert!(!validate_label_name(&[0xff, 0xfe]));
2617    }
2618}
2619
2620#[cfg(test)]
2621mod batching_tests {
2622    use std::sync::{Arc, Mutex};
2623
2624    use axum::http::StatusCode;
2625    use common_query::Output;
2626    use session::context::QueryContextRef;
2627
2628    use crate::batcher::BatchingProtocol;
2629    use crate::error::Result as ServerResult;
2630    use crate::http::test_helpers::TestClient;
2631    use crate::http::{HttpOptions, HttpServerBuilder};
2632    use crate::influxdb::InfluxdbRequest;
2633    use crate::opentsdb::codec::DataPoint;
2634    use crate::query_handler::{InfluxdbLineProtocolHandler, OpentsdbProtocolHandler};
2635
2636    #[tokio::test]
2637    async fn test_batching_selectors_are_independent() {
2638        use axum::routing::post;
2639        use axum::{Extension, Json, Router, middleware};
2640        use session::context::{QueryContext, QueryContextBuilder};
2641
2642        use crate::http::{LogicalBatchingProtocols, set_http_write_batching};
2643        for table in [false, true] {
2644            for logical in [false, true] {
2645                let app = Router::new()
2646                    .route(
2647                        "/",
2648                        post(|Extension(ctx): Extension<QueryContext>| async move {
2649                            Json([ctx.batching_enabled(), ctx.logical_batching_enabled()])
2650                        }),
2651                    )
2652                    .route_layer(middleware::from_fn_with_state(
2653                        BatchingProtocol::Otlp,
2654                        set_http_write_batching,
2655                    ))
2656                    .layer(Extension(QueryContextBuilder::default().build()))
2657                    .layer(Extension(Arc::new(if table {
2658                        vec![BatchingProtocol::Otlp]
2659                    } else {
2660                        vec![]
2661                    })))
2662                    .layer(Extension(LogicalBatchingProtocols(if logical {
2663                        vec![BatchingProtocol::Otlp]
2664                    } else {
2665                        vec![]
2666                    })));
2667                let client = TestClient::new(app).await;
2668                let actual: [bool; 2] = client.post("/").send().await.json().await;
2669                assert_eq!(actual, [table, logical]);
2670            }
2671        }
2672    }
2673
2674    #[derive(Default)]
2675    struct RecordingWriteHandler {
2676        selections: Mutex<Vec<bool>>,
2677    }
2678
2679    #[async_trait::async_trait]
2680    impl OpentsdbProtocolHandler for RecordingWriteHandler {
2681        async fn preflight(&self, _: &[DataPoint], _: QueryContextRef) -> ServerResult<()> {
2682            Ok(())
2683        }
2684
2685        async fn exec(&self, points: Vec<DataPoint>, ctx: QueryContextRef) -> ServerResult<usize> {
2686            self.selections.lock().unwrap().push(ctx.batching_enabled());
2687            Ok(points.len())
2688        }
2689    }
2690
2691    #[async_trait::async_trait]
2692    impl InfluxdbLineProtocolHandler for RecordingWriteHandler {
2693        async fn exec(&self, _: InfluxdbRequest, ctx: QueryContextRef) -> ServerResult<Output> {
2694            self.selections.lock().unwrap().push(ctx.batching_enabled());
2695            Ok(Output::new_with_affected_rows(1))
2696        }
2697    }
2698
2699    #[tokio::test]
2700    async fn test_real_influx_and_opentsdb_routes_keep_selection() {
2701        for protocols in [
2702            vec![],
2703            vec![BatchingProtocol::Influxdb],
2704            vec![BatchingProtocol::Opentsdb],
2705            vec![BatchingProtocol::Influxdb, BatchingProtocol::Opentsdb],
2706        ] {
2707            let handler = Arc::new(RecordingWriteHandler::default());
2708            let influx_enabled = protocols.contains(&BatchingProtocol::Influxdb);
2709            let opentsdb_enabled = protocols.contains(&BatchingProtocol::Opentsdb);
2710            let server = HttpServerBuilder::new(HttpOptions::default())
2711                .with_batching_protocols(protocols)
2712                .with_influxdb_handler(handler.clone())
2713                .with_opentsdb_handler(handler.clone())
2714                .build();
2715            let client = TestClient::new(server.build(server.make_app()).unwrap()).await;
2716            for path in [
2717                "/v1/influxdb/write",
2718                "/v1/influxdb/api/v2/write?bucket=public",
2719            ] {
2720                assert_eq!(
2721                    client
2722                        .post(path)
2723                        .body("cpu value=1 42")
2724                        .send()
2725                        .await
2726                        .status(),
2727                    StatusCode::NO_CONTENT
2728                );
2729            }
2730            let point =
2731                serde_json::json!({"metric":"cpu", "timestamp":42, "value":1, "tags":{"host":"a"}});
2732            for query in ["", "?summary", "?details"] {
2733                let response = client
2734                    .post(&format!("/v1/opentsdb/api/put{query}"))
2735                    .json(&point)
2736                    .send()
2737                    .await;
2738                assert_eq!(
2739                    response.status(),
2740                    if query.is_empty() {
2741                        StatusCode::NO_CONTENT
2742                    } else {
2743                        StatusCode::OK
2744                    }
2745                );
2746            }
2747            assert_eq!(
2748                *handler.selections.lock().unwrap(),
2749                vec![
2750                    influx_enabled,
2751                    influx_enabled,
2752                    opentsdb_enabled,
2753                    opentsdb_enabled,
2754                    opentsdb_enabled
2755                ]
2756            );
2757        }
2758    }
2759}