frontend/instance/
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::collections::HashMap;
16use std::sync::Arc;
17
18use api::prom_store::remote::read_request::ResponseType;
19use api::prom_store::remote::{Query, QueryResult, ReadRequest, ReadResponse};
20use api::v1::alter_table_expr::Kind;
21use api::v1::{
22    AddColumn, AddColumns, AlterTableExpr, ColumnDataType, ColumnDef, CreateTableExpr,
23    RowInsertRequests, SemanticType,
24};
25use async_trait::async_trait;
26use auth::{
27    PROM_STORE_READ, PROM_STORE_WRITE, PermissionChecker, PermissionCheckerRef, PermissionReq,
28    PermissionTableTarget, PermissionTableTargets,
29};
30use client::OutputData;
31use common_catalog::{format_full_table_name, parse_optional_catalog_and_schema_from_db_string};
32use common_error::ext::BoxedError;
33use common_meta::rpc::ddl::TriggerReason;
34use common_query::Output;
35use common_query::prelude::{GREPTIME_PHYSICAL_TABLE, greptime_value};
36use common_recordbatch::RecordBatches;
37use common_telemetry::{debug, tracing};
38use operator::insert::{
39    AutoCreateTableType, InserterRef, PerTableSemanticIndex, apply_per_table_semantic_options,
40    build_create_table_expr, fill_table_options_for_create, parse_per_table_semantic_index,
41};
42use operator::statement::StatementExecutor;
43use prost::Message;
44use query::query_engine::options::{QueryOptions, validate_catalog_and_schema};
45use servers::batcher::logical_table::PendingRowsSchemaAlterer;
46use servers::error::{self, AuthSnafu, Result as ServerResult};
47use servers::http::header::{CONTENT_ENCODING_SNAPPY, CONTENT_TYPE_PROTOBUF, collect_plan_metrics};
48use servers::http::prom_store::PHYSICAL_TABLE_PARAM;
49use servers::interceptor::{PromStoreProtocolInterceptor, PromStoreProtocolInterceptorRef};
50use servers::prom_store;
51use servers::query_handler::{
52    PromStoreProtocolHandler, PromStoreProtocolHandlerRef, PromStoreResponse,
53};
54use session::context::QueryContextRef;
55use snafu::{OptionExt, ResultExt};
56use store_api::metric_engine_consts::{METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY};
57use store_api::mito_engine_options::SST_FORMAT_KEY;
58use table::TableRef;
59use table::table_reference::TableReference;
60use tracing::instrument;
61
62use crate::error::{
63    AmbiguousValueColumnSnafu, CatalogSnafu, ColumnNotFoundSnafu, ExecLogicalPlanSnafu,
64    PromStoreRemoteQueryPlanSnafu, ReadTableSnafu, Result, TableNotFoundSnafu,
65};
66use crate::instance::Instance;
67
68const SAMPLES_RESPONSE_TYPE: i32 = ResponseType::Samples as i32;
69
70struct RemoteQueryOutput {
71    table_name: String,
72    timestamp_column_name: String,
73    value_column_name: String,
74    output: Output,
75}
76
77fn auto_create_table_type_for_prom_remote_write(
78    ctx: &QueryContextRef,
79    with_metric_engine: bool,
80) -> AutoCreateTableType {
81    if with_metric_engine {
82        let physical_table = ctx
83            .extension(PHYSICAL_TABLE_PARAM)
84            .unwrap_or(GREPTIME_PHYSICAL_TABLE)
85            .to_string();
86        AutoCreateTableType::Logical(physical_table)
87    } else {
88        AutoCreateTableType::Physical
89    }
90}
91
92fn required_physical_table_for_create_type(create_type: &AutoCreateTableType) -> Option<&str> {
93    match create_type {
94        AutoCreateTableType::Logical(physical_table) => Some(physical_table.as_str()),
95        _ => None,
96    }
97}
98
99fn fill_metric_physical_table_options(table_options: &mut HashMap<String, String>) {
100    // We always enforce flat format in this ingestion path.
101    table_options.insert(SST_FORMAT_KEY.to_string(), "flat".to_string());
102    table_options.insert(PHYSICAL_TABLE_METADATA_KEY.to_string(), "true".to_string());
103}
104
105#[inline]
106fn is_supported(response_type: i32) -> bool {
107    // Only supports samples response right now
108    response_type == SAMPLES_RESPONSE_TYPE
109}
110
111/// Negotiating the content type of the remote read response.
112///
113/// Response types are taken from the list in the FIFO order. If no response type in `accepted_response_types` is
114/// implemented by server, error is returned.
115/// For request that do not contain `accepted_response_types` field the SAMPLES response type will be used.
116fn negotiate_response_type(accepted_response_types: &[i32]) -> ServerResult<ResponseType> {
117    if accepted_response_types.is_empty() {
118        return Ok(ResponseType::Samples);
119    }
120
121    let response_type = accepted_response_types
122        .iter()
123        .find(|t| is_supported(**t))
124        .with_context(|| error::NotSupportedSnafu {
125            feat: format!(
126                "server does not support any of the requested response types: {accepted_response_types:?}",
127            ),
128        })?;
129
130    // It's safe to unwrap here, we known that it should be SAMPLES_RESPONSE_TYPE
131    Ok(ResponseType::try_from(*response_type).unwrap())
132}
133
134fn resolve_remote_query_target(
135    ctx: &QueryContextRef,
136    query: &Query,
137    table_name: &str,
138) -> PermissionTableTarget {
139    match prom_store::extract_schema_from_query(query) {
140        Some(database) => {
141            let (catalog, schema) = parse_optional_catalog_and_schema_from_db_string(&database);
142            PermissionTableTarget::new(
143                catalog.unwrap_or_else(|| ctx.current_catalog().to_string()),
144                schema,
145                table_name,
146            )
147        }
148        None => PermissionTableTarget::new(ctx.current_catalog(), ctx.current_schema(), table_name),
149    }
150}
151
152#[instrument(skip_all, fields(table_name))]
153async fn to_query_result(
154    table_name: &str,
155    timestamp_column_name: &str,
156    value_column_name: &str,
157    output: Output,
158) -> ServerResult<QueryResult> {
159    let OutputData::Stream(stream) = output.data else {
160        unreachable!()
161    };
162    let recordbatches = RecordBatches::try_collect(stream)
163        .await
164        .context(error::CollectRecordbatchSnafu)?;
165    Ok(QueryResult {
166        timeseries: prom_store::recordbatches_to_timeseries(
167            table_name,
168            timestamp_column_name,
169            value_column_name,
170            recordbatches,
171        )?,
172    })
173}
174
175fn resolve_column_names(table_name: &str, table: &TableRef) -> Result<String> {
176    let columns = table
177        .field_columns()
178        .map(|column| column.name)
179        .collect::<Vec<_>>();
180
181    match columns.as_slice() {
182        [] => ColumnNotFoundSnafu {
183            msg: format!("value field in table '{table_name}'"),
184        }
185        .fail(),
186
187        [only] => Ok(only.clone()),
188
189        columns if columns.iter().any(|name| name == greptime_value()) => {
190            Ok(greptime_value().to_string())
191        }
192
193        columns => AmbiguousValueColumnSnafu {
194            table_name: table_name.to_string(),
195            field_columns: columns.to_vec(),
196        }
197        .fail(),
198    }
199}
200
201impl Instance {
202    #[tracing::instrument(skip_all)]
203    async fn handle_remote_query(
204        &self,
205        ctx: &QueryContextRef,
206        catalog_name: &str,
207        schema_name: &str,
208        table_name: &str,
209        query: &Query,
210    ) -> Result<RemoteQueryOutput> {
211        let table = self
212            .catalog_manager
213            .table(catalog_name, schema_name, table_name, Some(ctx))
214            .await
215            .context(CatalogSnafu)?
216            .with_context(|| TableNotFoundSnafu {
217                table_name: format_full_table_name(catalog_name, schema_name, table_name),
218            })?;
219
220        let timestamp_column_name = table
221            .schema()
222            .timestamp_column()
223            .with_context(|| ColumnNotFoundSnafu {
224                msg: format!("time index in table '{table_name}'"),
225            })?
226            .name
227            .clone();
228
229        let value_column_name = resolve_column_names(table_name, &table)?;
230
231        let dataframe = self
232            .query_engine
233            .read_table(table)
234            .with_context(|_| ReadTableSnafu {
235                table_name: format_full_table_name(catalog_name, schema_name, table_name),
236            })?;
237
238        let logical_plan = prom_store::query_to_plan(dataframe, query, &timestamp_column_name)
239            .context(PromStoreRemoteQueryPlanSnafu)?;
240
241        debug!(
242            "Prometheus remote read, table: {}, logical plan: {}",
243            table_name,
244            logical_plan.display_indent(),
245        );
246
247        let output = self
248            .query_engine
249            .execute(logical_plan, ctx.clone())
250            .await
251            .context(ExecLogicalPlanSnafu)?;
252
253        Ok(RemoteQueryOutput {
254            table_name: table_name.to_string(),
255            timestamp_column_name,
256            value_column_name,
257            output,
258        })
259    }
260
261    #[tracing::instrument(skip_all)]
262    async fn handle_remote_queries(
263        &self,
264        ctx: QueryContextRef,
265        queries: &[Query],
266        query_targets: &[PermissionTableTarget],
267    ) -> ServerResult<Vec<RemoteQueryOutput>> {
268        let mut results = Vec::with_capacity(queries.len());
269
270        for (query, target) in queries.iter().zip(query_targets) {
271            let result = self
272                .handle_remote_query(&ctx, &target.catalog, &target.schema, &target.table, query)
273                .await
274                .map_err(BoxedError::new)
275                .context(error::ExecuteQuerySnafu)?;
276
277            results.push(result);
278        }
279        Ok(results)
280    }
281}
282
283#[async_trait]
284impl PendingRowsSchemaAlterer for Instance {
285    async fn create_tables_if_missing_batch(
286        &self,
287        catalog: &str,
288        schema: &str,
289        tables: &[(&str, &[api::v1::ColumnSchema])],
290        with_metric_engine: bool,
291        ctx: QueryContextRef,
292    ) -> ServerResult<()> {
293        if tables.is_empty() {
294            return Ok(());
295        }
296
297        let create_type = auto_create_table_type_for_prom_remote_write(&ctx, with_metric_engine);
298        if let Some(physical_table) = required_physical_table_for_create_type(&create_type) {
299            self.create_metric_physical_table_if_missing(
300                catalog,
301                schema,
302                physical_table,
303                ctx.clone(),
304            )
305            .await?;
306        }
307
308        let engine = if matches!(create_type, AutoCreateTableType::Logical(_)) {
309            METRIC_ENGINE_NAME
310        } else {
311            common_catalog::consts::default_engine()
312        };
313
314        // Check which tables actually still need to be created (may have been
315        // concurrently created by another request).
316        let mut create_exprs: Vec<CreateTableExpr> = Vec::with_capacity(tables.len());
317        let mut per_table_semantics: Option<Option<PerTableSemanticIndex>> = None;
318        for &(table_name, request_schema) in tables {
319            let existing = self
320                .catalog_manager()
321                .table(catalog, schema, table_name, Some(ctx.as_ref()))
322                .await
323                .map_err(BoxedError::new)
324                .context(error::ExecuteGrpcQuerySnafu)?;
325            if existing.is_some() {
326                continue;
327            }
328
329            let table_ref = TableReference::full(catalog, schema, table_name);
330            let mut create_table_expr = build_create_table_expr(&table_ref, request_schema, engine)
331                .map_err(BoxedError::new)
332                .context(error::ExecuteGrpcQuerySnafu)?;
333
334            let mut table_options = std::collections::HashMap::with_capacity(4);
335            fill_table_options_for_create(&mut table_options, &create_type, &ctx);
336            // The batched create path bypasses the operator's auto-create, so
337            // fold the per-table semantic index in here too.
338            let semantic_index = per_table_semantics
339                .get_or_insert_with(|| parse_per_table_semantic_index(&ctx))
340                .as_ref();
341            apply_per_table_semantic_options(
342                &mut table_options,
343                semantic_index,
344                schema,
345                table_name,
346            );
347            create_table_expr.table_options.extend(table_options);
348            create_exprs.push(create_table_expr);
349        }
350
351        if create_exprs.is_empty() {
352            return Ok(());
353        }
354
355        match create_type {
356            AutoCreateTableType::Logical(_) => {
357                // Use the batch API for logical tables.
358                self.statement_executor
359                    .create_logical_tables(&create_exprs, ctx, TriggerReason::AutoCreate)
360                    .await
361                    .map_err(BoxedError::new)
362                    .context(error::ExecuteGrpcQuerySnafu)?;
363            }
364            AutoCreateTableType::Physical => {
365                // Physical tables don't have a batch DDL path; create one at a time.
366                for mut expr in create_exprs {
367                    expr.table_options
368                        .insert(SST_FORMAT_KEY.to_string(), "flat".to_string());
369                    self.statement_executor
370                        .create_table_inner(&mut expr, None, ctx.clone(), TriggerReason::AutoCreate)
371                        .await
372                        .map_err(BoxedError::new)
373                        .context(error::ExecuteGrpcQuerySnafu)?;
374                }
375            }
376            create_type => {
377                return error::InvalidPromRemoteRequestSnafu {
378                    msg: format!(
379                        "prom remote write only supports logical or physical auto-create: {}",
380                        create_type.as_str()
381                    ),
382                }
383                .fail();
384            }
385        }
386
387        Ok(())
388    }
389
390    async fn add_missing_prom_tag_columns_batch(
391        &self,
392        catalog: &str,
393        schema: &str,
394        tables: &[(&str, &[String])],
395        ctx: QueryContextRef,
396    ) -> ServerResult<()> {
397        if tables.is_empty() {
398            return Ok(());
399        }
400
401        let alter_exprs: Vec<AlterTableExpr> = tables
402            .iter()
403            .filter(|(_, columns)| !columns.is_empty())
404            .map(|&(table_name, columns)| {
405                let add_columns = AddColumns {
406                    add_columns: columns
407                        .iter()
408                        .map(|column_name| AddColumn {
409                            column_def: Some(ColumnDef {
410                                name: column_name.clone(),
411                                data_type: ColumnDataType::String as i32,
412                                is_nullable: true,
413                                semantic_type: SemanticType::Tag as i32,
414                                comment: String::new(),
415                                ..Default::default()
416                            }),
417                            location: None,
418                            add_if_not_exists: true,
419                        })
420                        .collect(),
421                };
422
423                AlterTableExpr {
424                    catalog_name: catalog.to_string(),
425                    schema_name: schema.to_string(),
426                    table_name: table_name.to_string(),
427                    kind: Some(Kind::AddColumns(add_columns)),
428                }
429            })
430            .collect();
431
432        if alter_exprs.is_empty() {
433            return Ok(());
434        }
435
436        self.statement_executor
437            .alter_logical_tables(alter_exprs, ctx, TriggerReason::AutoAlter)
438            .await
439            .map_err(BoxedError::new)
440            .context(error::ExecuteGrpcQuerySnafu)?;
441
442        Ok(())
443    }
444}
445
446impl Instance {
447    async fn prepare_prom_store_write(
448        &self,
449        request: RowInsertRequests,
450        ctx: QueryContextRef,
451    ) -> ServerResult<(RowInsertRequests, QueryContextRef)> {
452        PromStoreProtocolHandler::pre_write(self, &request, ctx.clone()).await?;
453        Ok((request, Arc::new(ctx.fork())))
454    }
455
456    async fn execute_prom_store_write(
457        &self,
458        request: RowInsertRequests,
459        ctx: QueryContextRef,
460        with_metric_engine: bool,
461    ) -> ServerResult<Output> {
462        let output = if with_metric_engine {
463            let physical_table = ctx
464                .extension(PHYSICAL_TABLE_PARAM)
465                .unwrap_or(GREPTIME_PHYSICAL_TABLE)
466                .to_string();
467            self.handle_metric_row_inserts(request, ctx.clone(), physical_table.clone())
468                .await
469                .map_err(BoxedError::new)
470                .context(error::ExecuteGrpcQuerySnafu)?
471        } else {
472            self.handle_row_inserts(request, ctx.clone(), true, true)
473                .await
474                .map_err(BoxedError::new)
475                .context(error::ExecuteGrpcQuerySnafu)?
476        };
477
478        Ok(output)
479    }
480}
481
482#[async_trait]
483impl PromStoreProtocolHandler for Instance {
484    async fn pre_write(
485        &self,
486        request: &RowInsertRequests,
487        ctx: QueryContextRef,
488    ) -> ServerResult<()> {
489        self.plugins
490            .get::<PermissionCheckerRef>()
491            .as_ref()
492            .check_permission(ctx.current_user(), PermissionReq::Action(PROM_STORE_WRITE))
493            .context(AuthSnafu)?;
494        let interceptor_ref = self
495            .plugins
496            .get::<PromStoreProtocolInterceptorRef<servers::error::Error>>();
497        interceptor_ref.pre_write(request, ctx.clone())?;
498        self.check_row_insert_permission(request, &ctx, PermissionReq::Action(PROM_STORE_WRITE))
499            .context(AuthSnafu)?;
500        Ok(())
501    }
502
503    async fn write_prepared(
504        &self,
505        request: RowInsertRequests,
506        ctx: QueryContextRef,
507        with_metric_engine: bool,
508    ) -> ServerResult<Output> {
509        self.execute_prom_store_write(request, ctx, with_metric_engine)
510            .await
511    }
512
513    async fn write(
514        &self,
515        request: RowInsertRequests,
516        ctx: QueryContextRef,
517        with_metric_engine: bool,
518    ) -> ServerResult<Output> {
519        let (request, ctx) = self.prepare_prom_store_write(request, ctx).await?;
520        self.write_prepared(request, ctx, with_metric_engine).await
521    }
522
523    async fn write_all(
524        &self,
525        requests: Vec<(QueryContextRef, RowInsertRequests)>,
526        with_metric_engine: bool,
527    ) -> ServerResult<Vec<ServerResult<Output>>> {
528        let mut prepared = Vec::with_capacity(requests.len());
529        for (ctx, request) in requests {
530            let (request, ctx) = self.prepare_prom_store_write(request, ctx).await?;
531            prepared.push((ctx, request));
532        }
533        operator::insert::admit_row_insert_batches(&mut prepared)
534            .await
535            .map_err(BoxedError::new)
536            .context(error::ExecuteGrpcQuerySnafu)?;
537
538        let mut outputs = Vec::with_capacity(prepared.len());
539        for (ctx, request) in prepared {
540            let output = self.write_prepared(request, ctx, with_metric_engine).await;
541            let failed = output.is_err();
542            outputs.push(output);
543            if failed {
544                break;
545            }
546        }
547        Ok(outputs)
548    }
549
550    #[instrument(skip_all, fields(table_name))]
551    async fn read(
552        &self,
553        request: ReadRequest,
554        ctx: QueryContextRef,
555    ) -> ServerResult<PromStoreResponse> {
556        self.plugins
557            .get::<PermissionCheckerRef>()
558            .as_ref()
559            .check_permission(ctx.current_user(), PermissionReq::Action(PROM_STORE_READ))
560            .context(AuthSnafu)?;
561
562        let interceptor_ref = self
563            .plugins
564            .get::<PromStoreProtocolInterceptorRef<servers::error::Error>>();
565        interceptor_ref.pre_read(&request, ctx.clone())?;
566        let ctx = Arc::new(ctx.fork());
567
568        let table_names = request
569            .queries
570            .iter()
571            .map(prom_store::table_name)
572            .collect::<ServerResult<Vec<_>>>()?;
573        let query_targets = request
574            .queries
575            .iter()
576            .zip(&table_names)
577            .map(|(query, table_name)| resolve_remote_query_target(&ctx, query, table_name))
578            .collect::<Vec<_>>();
579        let disallow_cross_catalog_query = self
580            .plugins
581            .get::<QueryOptions>()
582            .map(|opts| opts.disallow_cross_catalog_query)
583            .unwrap_or_default();
584        if disallow_cross_catalog_query {
585            for target in &query_targets {
586                validate_catalog_and_schema(&target.catalog, &target.schema, &ctx)
587                    .map_err(BoxedError::new)
588                    .context(error::ExecuteQuerySnafu)?;
589            }
590        }
591        let targets = self
592            .resolve_query_permission_targets(
593                PermissionTableTargets::resolved(query_targets.clone()),
594                &ctx,
595            )
596            .await?;
597        self.check_table_permission(&ctx, PermissionReq::Action(PROM_STORE_READ), targets)
598            .context(AuthSnafu)?;
599
600        let response_type = negotiate_response_type(&request.accepted_response_types)?;
601
602        // TODO(dennis): use read_hints to speedup query if possible
603        let results = self
604            .handle_remote_queries(ctx, &request.queries, &query_targets)
605            .await?;
606
607        match response_type {
608            ResponseType::Samples => {
609                let mut query_results = Vec::with_capacity(results.len());
610                let mut map = HashMap::new();
611                for result in results {
612                    let RemoteQueryOutput {
613                        table_name,
614                        timestamp_column_name,
615                        value_column_name,
616                        output,
617                    } = result;
618                    let plan = output.meta.plan.clone();
619                    query_results.push(
620                        to_query_result(
621                            &table_name,
622                            &timestamp_column_name,
623                            &value_column_name,
624                            output,
625                        )
626                        .await?,
627                    );
628                    if let Some(ref plan) = plan {
629                        collect_plan_metrics(plan, &mut [&mut map]);
630                    }
631                }
632
633                let response = ReadResponse {
634                    results: query_results,
635                };
636
637                let resp_metrics = map
638                    .into_iter()
639                    .map(|(k, v)| (k, v.into()))
640                    .collect::<HashMap<_, _>>();
641
642                // TODO(dennis): may consume too much memory, adds flow control
643                Ok(PromStoreResponse {
644                    content_type: CONTENT_TYPE_PROTOBUF.clone(),
645                    content_encoding: CONTENT_ENCODING_SNAPPY.clone(),
646                    resp_metrics,
647                    body: prom_store::snappy_compress(&response.encode_to_vec())?,
648                })
649            }
650            ResponseType::StreamedXorChunks => error::NotSupportedSnafu {
651                feat: "streamed remote read",
652            }
653            .fail(),
654        }
655    }
656}
657
658impl Instance {
659    async fn create_metric_physical_table_if_missing(
660        &self,
661        catalog: &str,
662        schema: &str,
663        physical_table: &str,
664        ctx: QueryContextRef,
665    ) -> ServerResult<()> {
666        let table = self
667            .catalog_manager()
668            .table(catalog, schema, physical_table, Some(ctx.as_ref()))
669            .await
670            .map_err(BoxedError::new)
671            .context(error::ExecuteGrpcQuerySnafu)?;
672        if table.is_some() {
673            return Ok(());
674        }
675
676        let table_ref = TableReference::full(catalog, schema, physical_table);
677        let default_schema = vec![
678            api::v1::ColumnSchema {
679                column_name: common_query::prelude::greptime_timestamp().to_string(),
680                datatype: api::v1::ColumnDataType::TimestampMillisecond as i32,
681                semantic_type: api::v1::SemanticType::Timestamp as i32,
682                datatype_extension: None,
683                options: None,
684            },
685            api::v1::ColumnSchema {
686                column_name: common_query::prelude::greptime_value().to_string(),
687                datatype: api::v1::ColumnDataType::Float64 as i32,
688                semantic_type: api::v1::SemanticType::Field as i32,
689                datatype_extension: None,
690                options: None,
691            },
692        ];
693        let mut create_table_expr = build_create_table_expr(
694            &table_ref,
695            &default_schema,
696            common_catalog::consts::default_engine(),
697        )
698        .map_err(BoxedError::new)
699        .context(error::ExecuteGrpcQuerySnafu)?;
700        create_table_expr.engine = METRIC_ENGINE_NAME.to_string();
701        fill_metric_physical_table_options(&mut create_table_expr.table_options);
702
703        self.statement_executor
704            .create_table_inner(&mut create_table_expr, None, ctx, TriggerReason::AutoCreate)
705            .await
706            .map_err(BoxedError::new)
707            .context(error::ExecuteGrpcQuerySnafu)?;
708
709        Ok(())
710    }
711}
712
713/// This handler is mainly used for `frontend` or `standalone` to directly import
714/// the metrics collected by itself, thereby avoiding importing metrics through the network,
715/// thus reducing compression and network transmission overhead,
716/// so only implement `PromStoreProtocolHandler::write` method.
717pub struct ExportMetricHandler {
718    inserter: InserterRef,
719    statement_executor: Arc<StatementExecutor>,
720}
721
722impl ExportMetricHandler {
723    pub fn new_handler(
724        inserter: InserterRef,
725        statement_executor: Arc<StatementExecutor>,
726    ) -> PromStoreProtocolHandlerRef {
727        Arc::new(Self {
728            inserter,
729            statement_executor,
730        })
731    }
732}
733
734#[async_trait]
735impl PromStoreProtocolHandler for ExportMetricHandler {
736    async fn pre_write(
737        &self,
738        _request: &RowInsertRequests,
739        _ctx: QueryContextRef,
740    ) -> ServerResult<()> {
741        Ok(())
742    }
743
744    async fn write_prepared(
745        &self,
746        request: RowInsertRequests,
747        ctx: QueryContextRef,
748        _: bool,
749    ) -> ServerResult<Output> {
750        self.inserter
751            .handle_metric_row_inserts(
752                request,
753                ctx,
754                &self.statement_executor,
755                GREPTIME_PHYSICAL_TABLE.to_string(),
756            )
757            .await
758            .map_err(BoxedError::new)
759            .context(error::ExecuteGrpcQuerySnafu)
760    }
761
762    async fn write(
763        &self,
764        request: RowInsertRequests,
765        ctx: QueryContextRef,
766        with_metric_engine: bool,
767    ) -> ServerResult<Output> {
768        self.write_prepared(request, ctx, with_metric_engine).await
769    }
770
771    async fn write_all(
772        &self,
773        mut requests: Vec<(QueryContextRef, RowInsertRequests)>,
774        with_metric_engine: bool,
775    ) -> ServerResult<Vec<ServerResult<Output>>> {
776        operator::insert::admit_row_insert_batches(&mut requests)
777            .await
778            .map_err(BoxedError::new)
779            .context(error::ExecuteGrpcQuerySnafu)?;
780        let mut outputs = Vec::with_capacity(requests.len());
781        for (ctx, request) in requests {
782            let output = self.write_prepared(request, ctx, with_metric_engine).await;
783            let failed = output.is_err();
784            outputs.push(output);
785            if failed {
786                break;
787            }
788        }
789        Ok(outputs)
790    }
791
792    async fn read(
793        &self,
794        _request: ReadRequest,
795        _ctx: QueryContextRef,
796    ) -> ServerResult<PromStoreResponse> {
797        unreachable!();
798    }
799}
800
801#[cfg(test)]
802mod tests {
803    use std::sync::Arc;
804
805    use api::prom_store::remote::LabelMatcher;
806    use session::context::QueryContext;
807
808    use super::*;
809
810    #[test]
811    fn test_auto_create_table_type_for_prom_remote_write_metric_engine() {
812        let mut query_ctx = QueryContext::with(
813            common_catalog::consts::DEFAULT_CATALOG_NAME,
814            common_catalog::consts::DEFAULT_SCHEMA_NAME,
815        );
816        query_ctx.set_extension(PHYSICAL_TABLE_PARAM, "metric_physical".to_string());
817        let ctx = Arc::new(query_ctx);
818
819        let create_type = auto_create_table_type_for_prom_remote_write(&ctx, true);
820        match create_type {
821            AutoCreateTableType::Logical(physical) => assert_eq!(physical, "metric_physical"),
822            _ => panic!("expected logical table create type"),
823        }
824    }
825
826    #[test]
827    fn test_auto_create_table_type_for_prom_remote_write_without_metric_engine() {
828        let ctx = Arc::new(QueryContext::with(
829            common_catalog::consts::DEFAULT_CATALOG_NAME,
830            common_catalog::consts::DEFAULT_SCHEMA_NAME,
831        ));
832
833        let create_type = auto_create_table_type_for_prom_remote_write(&ctx, false);
834        match create_type {
835            AutoCreateTableType::Physical => {}
836            _ => panic!("expected physical table create type"),
837        }
838    }
839
840    #[test]
841    fn test_required_physical_table_for_create_type() {
842        let logical = AutoCreateTableType::Logical("phy_table".to_string());
843        assert_eq!(
844            Some("phy_table"),
845            required_physical_table_for_create_type(&logical)
846        );
847
848        let physical = AutoCreateTableType::Physical;
849        assert_eq!(None, required_physical_table_for_create_type(&physical));
850    }
851
852    #[test]
853    fn test_metric_physical_table_options_forces_flat_sst_format() {
854        let mut table_options = HashMap::new();
855
856        fill_metric_physical_table_options(&mut table_options);
857
858        assert_eq!(
859            Some("flat"),
860            table_options.get(SST_FORMAT_KEY).map(String::as_str)
861        );
862        assert_eq!(
863            Some("true"),
864            table_options
865                .get(PHYSICAL_TABLE_METADATA_KEY)
866                .map(String::as_str)
867        );
868    }
869
870    fn query_with_database(database: &str) -> Query {
871        Query {
872            matchers: vec![LabelMatcher {
873                name: servers::prom_store::DATABASE_LABEL.to_string(),
874                value: database.to_string(),
875                r#type: api::prom_store::remote::label_matcher::Type::Eq as i32,
876            }],
877            ..Default::default()
878        }
879    }
880
881    #[test]
882    fn test_resolve_remote_query_target() {
883        let ctx = Arc::new(QueryContext::with("request_catalog", "request_schema"));
884
885        assert_eq!(
886            PermissionTableTarget::new("request_catalog", "request_schema", "fallback_metric"),
887            resolve_remote_query_target(&ctx, &Query::default(), "fallback_metric")
888        );
889        assert_eq!(
890            PermissionTableTarget::new("request_catalog", "selected_schema", "schema_metric"),
891            resolve_remote_query_target(
892                &ctx,
893                &query_with_database("selected_schema"),
894                "schema_metric"
895            )
896        );
897        assert_eq!(
898            PermissionTableTarget::new("selected_catalog", "selected_schema", "catalog_metric"),
899            resolve_remote_query_target(
900                &ctx,
901                &query_with_database("selected_catalog-selected_schema"),
902                "catalog_metric"
903            )
904        );
905    }
906
907    #[test]
908    fn test_resolve_remote_query_target_preserves_context() {
909        let ctx = Arc::new(QueryContext::with("request_catalog", "request_schema"));
910
911        resolve_remote_query_target(
912            &ctx,
913            &query_with_database("selected_catalog-selected_schema"),
914            "metric",
915        );
916
917        assert_eq!("request_catalog", ctx.current_catalog());
918        assert_eq!("request_schema", ctx.current_schema());
919    }
920}
frontend/instance/prom_store.rs

frontend/instance/
prom_store.rs
