Skip to main content

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::error::{self, AuthSnafu, Result as ServerResult};
46use servers::http::header::{CONTENT_ENCODING_SNAPPY, CONTENT_TYPE_PROTOBUF, collect_plan_metrics};
47use servers::http::prom_store::PHYSICAL_TABLE_PARAM;
48use servers::interceptor::{PromStoreProtocolInterceptor, PromStoreProtocolInterceptorRef};
49use servers::pending_rows_batcher::PendingRowsSchemaAlterer;
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            prepared.push(self.prepare_prom_store_write(request, ctx).await?);
531        }
532
533        let mut outputs = Vec::with_capacity(prepared.len());
534        for (request, ctx) in prepared {
535            let output = self.write_prepared(request, ctx, with_metric_engine).await;
536            let failed = output.is_err();
537            outputs.push(output);
538            if failed {
539                break;
540            }
541        }
542        Ok(outputs)
543    }
544
545    #[instrument(skip_all, fields(table_name))]
546    async fn read(
547        &self,
548        request: ReadRequest,
549        ctx: QueryContextRef,
550    ) -> ServerResult<PromStoreResponse> {
551        self.plugins
552            .get::<PermissionCheckerRef>()
553            .as_ref()
554            .check_permission(ctx.current_user(), PermissionReq::Action(PROM_STORE_READ))
555            .context(AuthSnafu)?;
556
557        let interceptor_ref = self
558            .plugins
559            .get::<PromStoreProtocolInterceptorRef<servers::error::Error>>();
560        interceptor_ref.pre_read(&request, ctx.clone())?;
561        let ctx = Arc::new(ctx.fork());
562
563        let table_names = request
564            .queries
565            .iter()
566            .map(prom_store::table_name)
567            .collect::<ServerResult<Vec<_>>>()?;
568        let query_targets = request
569            .queries
570            .iter()
571            .zip(&table_names)
572            .map(|(query, table_name)| resolve_remote_query_target(&ctx, query, table_name))
573            .collect::<Vec<_>>();
574        let disallow_cross_catalog_query = self
575            .plugins
576            .get::<QueryOptions>()
577            .map(|opts| opts.disallow_cross_catalog_query)
578            .unwrap_or_default();
579        if disallow_cross_catalog_query {
580            for target in &query_targets {
581                validate_catalog_and_schema(&target.catalog, &target.schema, &ctx)
582                    .map_err(BoxedError::new)
583                    .context(error::ExecuteQuerySnafu)?;
584            }
585        }
586        let targets = self
587            .resolve_query_permission_targets(
588                PermissionTableTargets::resolved(query_targets.clone()),
589                &ctx,
590            )
591            .await?;
592        self.check_table_permission(&ctx, PermissionReq::Action(PROM_STORE_READ), targets)
593            .context(AuthSnafu)?;
594
595        let response_type = negotiate_response_type(&request.accepted_response_types)?;
596
597        // TODO(dennis): use read_hints to speedup query if possible
598        let results = self
599            .handle_remote_queries(ctx, &request.queries, &query_targets)
600            .await?;
601
602        match response_type {
603            ResponseType::Samples => {
604                let mut query_results = Vec::with_capacity(results.len());
605                let mut map = HashMap::new();
606                for result in results {
607                    let RemoteQueryOutput {
608                        table_name,
609                        timestamp_column_name,
610                        value_column_name,
611                        output,
612                    } = result;
613                    let plan = output.meta.plan.clone();
614                    query_results.push(
615                        to_query_result(
616                            &table_name,
617                            &timestamp_column_name,
618                            &value_column_name,
619                            output,
620                        )
621                        .await?,
622                    );
623                    if let Some(ref plan) = plan {
624                        collect_plan_metrics(plan, &mut [&mut map]);
625                    }
626                }
627
628                let response = ReadResponse {
629                    results: query_results,
630                };
631
632                let resp_metrics = map
633                    .into_iter()
634                    .map(|(k, v)| (k, v.into()))
635                    .collect::<HashMap<_, _>>();
636
637                // TODO(dennis): may consume too much memory, adds flow control
638                Ok(PromStoreResponse {
639                    content_type: CONTENT_TYPE_PROTOBUF.clone(),
640                    content_encoding: CONTENT_ENCODING_SNAPPY.clone(),
641                    resp_metrics,
642                    body: prom_store::snappy_compress(&response.encode_to_vec())?,
643                })
644            }
645            ResponseType::StreamedXorChunks => error::NotSupportedSnafu {
646                feat: "streamed remote read",
647            }
648            .fail(),
649        }
650    }
651}
652
653impl Instance {
654    async fn create_metric_physical_table_if_missing(
655        &self,
656        catalog: &str,
657        schema: &str,
658        physical_table: &str,
659        ctx: QueryContextRef,
660    ) -> ServerResult<()> {
661        let table = self
662            .catalog_manager()
663            .table(catalog, schema, physical_table, Some(ctx.as_ref()))
664            .await
665            .map_err(BoxedError::new)
666            .context(error::ExecuteGrpcQuerySnafu)?;
667        if table.is_some() {
668            return Ok(());
669        }
670
671        let table_ref = TableReference::full(catalog, schema, physical_table);
672        let default_schema = vec![
673            api::v1::ColumnSchema {
674                column_name: common_query::prelude::greptime_timestamp().to_string(),
675                datatype: api::v1::ColumnDataType::TimestampMillisecond as i32,
676                semantic_type: api::v1::SemanticType::Timestamp as i32,
677                datatype_extension: None,
678                options: None,
679            },
680            api::v1::ColumnSchema {
681                column_name: common_query::prelude::greptime_value().to_string(),
682                datatype: api::v1::ColumnDataType::Float64 as i32,
683                semantic_type: api::v1::SemanticType::Field as i32,
684                datatype_extension: None,
685                options: None,
686            },
687        ];
688        let mut create_table_expr = build_create_table_expr(
689            &table_ref,
690            &default_schema,
691            common_catalog::consts::default_engine(),
692        )
693        .map_err(BoxedError::new)
694        .context(error::ExecuteGrpcQuerySnafu)?;
695        create_table_expr.engine = METRIC_ENGINE_NAME.to_string();
696        fill_metric_physical_table_options(&mut create_table_expr.table_options);
697
698        self.statement_executor
699            .create_table_inner(&mut create_table_expr, None, ctx, TriggerReason::AutoCreate)
700            .await
701            .map_err(BoxedError::new)
702            .context(error::ExecuteGrpcQuerySnafu)?;
703
704        Ok(())
705    }
706}
707
708/// This handler is mainly used for `frontend` or `standalone` to directly import
709/// the metrics collected by itself, thereby avoiding importing metrics through the network,
710/// thus reducing compression and network transmission overhead,
711/// so only implement `PromStoreProtocolHandler::write` method.
712pub struct ExportMetricHandler {
713    inserter: InserterRef,
714    statement_executor: Arc<StatementExecutor>,
715}
716
717impl ExportMetricHandler {
718    pub fn new_handler(
719        inserter: InserterRef,
720        statement_executor: Arc<StatementExecutor>,
721    ) -> PromStoreProtocolHandlerRef {
722        Arc::new(Self {
723            inserter,
724            statement_executor,
725        })
726    }
727}
728
729#[async_trait]
730impl PromStoreProtocolHandler for ExportMetricHandler {
731    async fn pre_write(
732        &self,
733        _request: &RowInsertRequests,
734        _ctx: QueryContextRef,
735    ) -> ServerResult<()> {
736        Ok(())
737    }
738
739    async fn write_prepared(
740        &self,
741        request: RowInsertRequests,
742        ctx: QueryContextRef,
743        _: bool,
744    ) -> ServerResult<Output> {
745        self.inserter
746            .handle_metric_row_inserts(
747                request,
748                ctx,
749                &self.statement_executor,
750                GREPTIME_PHYSICAL_TABLE.to_string(),
751            )
752            .await
753            .map_err(BoxedError::new)
754            .context(error::ExecuteGrpcQuerySnafu)
755    }
756
757    async fn write(
758        &self,
759        request: RowInsertRequests,
760        ctx: QueryContextRef,
761        with_metric_engine: bool,
762    ) -> ServerResult<Output> {
763        self.write_prepared(request, ctx, with_metric_engine).await
764    }
765
766    async fn write_all(
767        &self,
768        requests: Vec<(QueryContextRef, RowInsertRequests)>,
769        with_metric_engine: bool,
770    ) -> ServerResult<Vec<ServerResult<Output>>> {
771        let mut outputs = Vec::with_capacity(requests.len());
772        for (ctx, request) in requests {
773            let output = self.write_prepared(request, ctx, with_metric_engine).await;
774            let failed = output.is_err();
775            outputs.push(output);
776            if failed {
777                break;
778            }
779        }
780        Ok(outputs)
781    }
782
783    async fn read(
784        &self,
785        _request: ReadRequest,
786        _ctx: QueryContextRef,
787    ) -> ServerResult<PromStoreResponse> {
788        unreachable!();
789    }
790}
791
792#[cfg(test)]
793mod tests {
794    use std::sync::Arc;
795
796    use api::prom_store::remote::LabelMatcher;
797    use session::context::QueryContext;
798
799    use super::*;
800
801    #[test]
802    fn test_auto_create_table_type_for_prom_remote_write_metric_engine() {
803        let mut query_ctx = QueryContext::with(
804            common_catalog::consts::DEFAULT_CATALOG_NAME,
805            common_catalog::consts::DEFAULT_SCHEMA_NAME,
806        );
807        query_ctx.set_extension(PHYSICAL_TABLE_PARAM, "metric_physical".to_string());
808        let ctx = Arc::new(query_ctx);
809
810        let create_type = auto_create_table_type_for_prom_remote_write(&ctx, true);
811        match create_type {
812            AutoCreateTableType::Logical(physical) => assert_eq!(physical, "metric_physical"),
813            _ => panic!("expected logical table create type"),
814        }
815    }
816
817    #[test]
818    fn test_auto_create_table_type_for_prom_remote_write_without_metric_engine() {
819        let ctx = Arc::new(QueryContext::with(
820            common_catalog::consts::DEFAULT_CATALOG_NAME,
821            common_catalog::consts::DEFAULT_SCHEMA_NAME,
822        ));
823
824        let create_type = auto_create_table_type_for_prom_remote_write(&ctx, false);
825        match create_type {
826            AutoCreateTableType::Physical => {}
827            _ => panic!("expected physical table create type"),
828        }
829    }
830
831    #[test]
832    fn test_required_physical_table_for_create_type() {
833        let logical = AutoCreateTableType::Logical("phy_table".to_string());
834        assert_eq!(
835            Some("phy_table"),
836            required_physical_table_for_create_type(&logical)
837        );
838
839        let physical = AutoCreateTableType::Physical;
840        assert_eq!(None, required_physical_table_for_create_type(&physical));
841    }
842
843    #[test]
844    fn test_metric_physical_table_options_forces_flat_sst_format() {
845        let mut table_options = HashMap::new();
846
847        fill_metric_physical_table_options(&mut table_options);
848
849        assert_eq!(
850            Some("flat"),
851            table_options.get(SST_FORMAT_KEY).map(String::as_str)
852        );
853        assert_eq!(
854            Some("true"),
855            table_options
856                .get(PHYSICAL_TABLE_METADATA_KEY)
857                .map(String::as_str)
858        );
859    }
860
861    fn query_with_database(database: &str) -> Query {
862        Query {
863            matchers: vec![LabelMatcher {
864                name: servers::prom_store::DATABASE_LABEL.to_string(),
865                value: database.to_string(),
866                r#type: api::prom_store::remote::label_matcher::Type::Eq as i32,
867            }],
868            ..Default::default()
869        }
870    }
871
872    #[test]
873    fn test_resolve_remote_query_target() {
874        let ctx = Arc::new(QueryContext::with("request_catalog", "request_schema"));
875
876        assert_eq!(
877            PermissionTableTarget::new("request_catalog", "request_schema", "fallback_metric"),
878            resolve_remote_query_target(&ctx, &Query::default(), "fallback_metric")
879        );
880        assert_eq!(
881            PermissionTableTarget::new("request_catalog", "selected_schema", "schema_metric"),
882            resolve_remote_query_target(
883                &ctx,
884                &query_with_database("selected_schema"),
885                "schema_metric"
886            )
887        );
888        assert_eq!(
889            PermissionTableTarget::new("selected_catalog", "selected_schema", "catalog_metric"),
890            resolve_remote_query_target(
891                &ctx,
892                &query_with_database("selected_catalog-selected_schema"),
893                "catalog_metric"
894            )
895        );
896    }
897
898    #[test]
899    fn test_resolve_remote_query_target_preserves_context() {
900        let ctx = Arc::new(QueryContext::with("request_catalog", "request_schema"));
901
902        resolve_remote_query_target(
903            &ctx,
904            &query_with_database("selected_catalog-selected_schema"),
905            "metric",
906        );
907
908        assert_eq!("request_catalog", ctx.current_catalog());
909        assert_eq!("request_schema", ctx.current_schema());
910    }
911}