servers/http/
prom_store.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::sync::Arc;
16
17use api::prom_store::remote::ReadRequest;
18use axum::Extension;
19use axum::body::Bytes;
20use axum::extract::{Query, State};
21use axum::http::{HeaderValue, StatusCode, header};
22use axum::response::IntoResponse;
23use axum_extra::TypedHeader;
24use common_catalog::consts::DEFAULT_SCHEMA_NAME;
25use common_query::prelude::GREPTIME_PHYSICAL_TABLE;
26use common_telemetry::tracing;
27use hyper::HeaderMap;
28use pipeline::PipelineDefinition;
29use pipeline::util::to_pipeline_version;
30use prost::Message;
31use serde::{Deserialize, Serialize};
32use session::context::{Channel, QueryContext};
33use snafu::prelude::*;
34use table::requests::{
35    METADATA_QUALITY_INFERRED, SEMANTIC_METRIC_METADATA_QUALITY, SEMANTIC_SIGNAL_TYPE,
36    SEMANTIC_SOURCE, SIGNAL_TYPE_METRIC, SOURCE_PROMETHEUS,
37};
38
39use crate::error::{self, InternalSnafu, PipelineSnafu, Result};
40use crate::http::extractor::PipelineInfo;
41use crate::http::header::{GREPTIME_DB_HEADER_METRICS, write_cost_header_map};
42use crate::pending_rows_batcher::PendingRowsBatcher;
43use crate::prom_remote_write::decode::PromSeriesProcessor;
44use crate::prom_remote_write::decode_remote_write_request;
45use crate::prom_remote_write::validation::PromValidationMode;
46use crate::prom_store::{extract_schema_from_read_request, snappy_decompress};
47use crate::query_handler::{PipelineHandlerRef, PromStoreProtocolHandlerRef, PromStoreResponse};
48
49pub const PHYSICAL_TABLE_PARAM: &str = "physical_table";
50pub const DEFAULT_ENCODING: &str = "snappy";
51pub const VM_ENCODING: &str = "zstd";
52pub const VM_PROTO_VERSION: &str = "1";
53
54#[derive(Clone)]
55pub struct PromStoreState {
56    pub prom_store_handler: PromStoreProtocolHandlerRef,
57    pub pipeline_handler: Option<PipelineHandlerRef>,
58    pub prom_store_with_metric_engine: bool,
59    pub prom_validation_mode: PromValidationMode,
60    pub pending_rows_batcher: Option<Arc<PendingRowsBatcher>>,
61}
62
63#[derive(Debug, Serialize, Deserialize)]
64pub struct RemoteWriteQuery {
65    pub db: Option<String>,
66    /// Specify which physical table to use for storing metrics.
67    /// This only works on remote write requests.
68    pub physical_table: Option<String>,
69    /// For VictoriaMetrics modified remote write protocol
70    pub get_vm_proto_version: Option<String>,
71}
72
73impl Default for RemoteWriteQuery {
74    fn default() -> RemoteWriteQuery {
75        Self {
76            db: Some(DEFAULT_SCHEMA_NAME.to_string()),
77            physical_table: Some(GREPTIME_PHYSICAL_TABLE.to_string()),
78            get_vm_proto_version: None,
79        }
80    }
81}
82
83#[axum_macros::debug_handler]
84#[tracing::instrument(
85    skip_all,
86    fields(protocol = "prometheus", request_type = "remote_write")
87)]
88pub async fn remote_write(
89    State(state): State<PromStoreState>,
90    Query(params): Query<RemoteWriteQuery>,
91    Extension(mut query_ctx): Extension<QueryContext>,
92    pipeline_info: PipelineInfo,
93    content_encoding: TypedHeader<headers::ContentEncoding>,
94    body: Bytes,
95) -> Result<impl IntoResponse> {
96    let PromStoreState {
97        prom_store_handler,
98        pipeline_handler,
99        prom_store_with_metric_engine,
100        prom_validation_mode,
101        pending_rows_batcher,
102    } = state;
103
104    if let Some(_vm_handshake) = params.get_vm_proto_version {
105        return Ok(VM_PROTO_VERSION.into_response());
106    }
107
108    let db = params.db.clone().unwrap_or_default();
109    query_ctx.set_channel(Channel::Prometheus);
110    let physical_table = params
111        .physical_table
112        .clone()
113        .unwrap_or_else(|| GREPTIME_PHYSICAL_TABLE.to_string());
114    query_ctx.set_extension(PHYSICAL_TABLE_PARAM, physical_table.clone());
115    // Stamp the Prometheus metric identity here, before `as_req_iter` splits into the
116    // batched and direct write paths, so both inherit it (the batched path bypasses
117    // `PromStoreProtocolHandler::write`). Prom RW v1 metadata is weak, so the type is
118    // inferred from naming.
119    query_ctx.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC);
120    query_ctx.set_extension(SEMANTIC_SOURCE, SOURCE_PROMETHEUS);
121    query_ctx.set_extension(SEMANTIC_METRIC_METADATA_QUALITY, METADATA_QUALITY_INFERRED);
122    let query_ctx = Arc::new(query_ctx);
123    let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_WRITE_ELAPSED
124        .with_label_values(&[db.as_str()])
125        .start_timer();
126
127    let is_zstd = content_encoding.contains(VM_ENCODING);
128
129    let mut processor = PromSeriesProcessor::default_processor();
130
131    if let Some(pipeline_name) = pipeline_info.pipeline_name {
132        let pipeline_def = PipelineDefinition::from_name(
133            &pipeline_name,
134            to_pipeline_version(pipeline_info.pipeline_version.as_deref())
135                .context(PipelineSnafu)?,
136            None,
137        )
138        .context(PipelineSnafu)?;
139        let pipeline_handler = pipeline_handler.context(InternalSnafu {
140            err_msg: "pipeline handler is not set".to_string(),
141        })?;
142
143        processor.set_pipeline(pipeline_handler, query_ctx.clone(), pipeline_def);
144    }
145
146    let mut req = decode_remote_write_request(is_zstd, body, prom_validation_mode, &mut processor)?;
147
148    let req = if processor.use_pipeline {
149        processor.exec_pipeline().await?
150    } else {
151        req.as_insert_requests()
152    };
153
154    if prom_store_with_metric_engine && let Some(batcher) = pending_rows_batcher {
155        for (temp_ctx, reqs) in req.as_req_iter(query_ctx) {
156            prom_store_handler
157                .pre_write(&reqs, temp_ctx.clone())
158                .await?;
159            let rows = batcher.submit(reqs, temp_ctx).await?;
160            crate::metrics::PROM_STORE_REMOTE_WRITE_SAMPLES
161                .with_label_values(&[db.as_str()])
162                .inc_by(rows);
163        }
164        return Ok((StatusCode::NO_CONTENT, write_cost_header_map(0)).into_response());
165    }
166
167    let mut cost = 0;
168    for (temp_ctx, reqs) in req.as_req_iter(query_ctx) {
169        let cnt: u64 = reqs
170            .inserts
171            .iter()
172            .filter_map(|s| s.rows.as_ref().map(|r| r.rows.len() as u64))
173            .sum();
174        let output = prom_store_handler
175            .write(reqs, temp_ctx, prom_store_with_metric_engine)
176            .await?;
177        crate::metrics::PROM_STORE_REMOTE_WRITE_SAMPLES
178            .with_label_values(&[db.as_str()])
179            .inc_by(cnt);
180        cost += output.meta.cost;
181    }
182
183    Ok((StatusCode::NO_CONTENT, write_cost_header_map(cost)).into_response())
184}
185
186impl IntoResponse for PromStoreResponse {
187    fn into_response(self) -> axum::response::Response {
188        let mut header_map = HeaderMap::new();
189        header_map.insert(&header::CONTENT_TYPE, self.content_type);
190        header_map.insert(&header::CONTENT_ENCODING, self.content_encoding);
191
192        let metrics = if self.resp_metrics.is_empty() {
193            None
194        } else {
195            serde_json::to_string(&self.resp_metrics).ok()
196        };
197        if let Some(m) = metrics.and_then(|m| HeaderValue::from_str(&m).ok()) {
198            header_map.insert(&GREPTIME_DB_HEADER_METRICS, m);
199        }
200
201        (header_map, self.body).into_response()
202    }
203}
204
205#[axum_macros::debug_handler]
206#[tracing::instrument(
207    skip_all,
208    fields(protocol = "prometheus", request_type = "remote_read")
209)]
210pub async fn remote_read(
211    State(state): State<PromStoreState>,
212    Query(params): Query<RemoteWriteQuery>,
213    Extension(mut query_ctx): Extension<QueryContext>,
214    body: Bytes,
215) -> Result<PromStoreResponse> {
216    let db = params.db.clone().unwrap_or_default();
217    query_ctx.set_channel(Channel::Prometheus);
218
219    let request = decode_remote_read_request(body).await?;
220
221    // Extract schema from special labels and set it in query context
222    if let Some(schema) = extract_schema_from_read_request(&request) {
223        query_ctx.set_current_schema(&schema);
224    }
225
226    let query_ctx = Arc::new(query_ctx);
227    let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_READ_ELAPSED
228        .with_label_values(&[db.as_str()])
229        .start_timer();
230
231    state.prom_store_handler.read(request, query_ctx).await
232}
233
234async fn decode_remote_read_request(body: Bytes) -> Result<ReadRequest> {
235    let buf = snappy_decompress(&body[..])?;
236
237    ReadRequest::decode(&buf[..]).context(error::DecodePromRemoteRequestSnafu)
238}
servers/http/prom_store.rs

servers/http/
prom_store.rs
