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