frontend/instance/

prom_store.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashMap;
use std::sync::Arc;

use api::prom_store::remote::read_request::ResponseType;
use api::prom_store::remote::{Query, QueryResult, ReadRequest, ReadResponse};
use api::v1::alter_table_expr::Kind;
use api::v1::{
    AddColumn, AddColumns, AlterTableExpr, ColumnDataType, ColumnDef, CreateTableExpr,
    RowInsertRequests, SemanticType,
};
use async_trait::async_trait;
use auth::{
    PROM_STORE_READ, PROM_STORE_WRITE, PermissionChecker, PermissionCheckerRef, PermissionReq,
    PermissionTableTarget, PermissionTableTargets,
};
use client::OutputData;
use common_catalog::{format_full_table_name, parse_optional_catalog_and_schema_from_db_string};
use common_error::ext::BoxedError;
use common_meta::rpc::ddl::TriggerReason;
use common_query::Output;
use common_query::prelude::{GREPTIME_PHYSICAL_TABLE, greptime_value};
use common_recordbatch::RecordBatches;
use common_telemetry::{debug, tracing};
use datatypes::data_type::ConcreteDataType;
use datatypes::schema::Schema;
use operator::insert::{
    AutoCreateTableType, InserterRef, PerTableSemanticIndex, apply_per_table_semantic_options,
    build_create_table_expr, fill_table_options_for_create, parse_per_table_semantic_index,
};
use operator::statement::StatementExecutor;
use prost::Message;
use query::query_engine::options::{QueryOptions, validate_catalog_and_schema};
use servers::batcher::logical_table::PendingRowsSchemaAlterer;
use servers::error::{self, AuthSnafu, Result as ServerResult};
use servers::http::header::{CONTENT_ENCODING_SNAPPY, CONTENT_TYPE_PROTOBUF, collect_plan_metrics};
use servers::http::prom_store::PHYSICAL_TABLE_PARAM;
use servers::interceptor::{PromStoreProtocolInterceptor, PromStoreProtocolInterceptorRef};
use servers::prom_store;
use servers::query_handler::{
    PromStoreProtocolHandler, PromStoreProtocolHandlerRef, PromStoreResponse,
};
use session::context::QueryContextRef;
use snafu::{OptionExt, ResultExt, ensure};
use store_api::metric_engine_consts::{METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY};
use store_api::mito_engine_options::SST_FORMAT_KEY;
use table::TableRef;
use table::table_reference::TableReference;
use tracing::instrument;

use crate::error::{
    AmbiguousValueColumnSnafu, CatalogSnafu, ColumnNotFoundSnafu, ExecLogicalPlanSnafu,
    InvalidPromRemoteReadSchemaSnafu, PromStoreRemoteQueryPlanSnafu, ReadTableSnafu, Result,
};
use crate::instance::Instance;

const SAMPLES_RESPONSE_TYPE: i32 = ResponseType::Samples as i32;

struct RemoteQueryOutput {
    table_name: String,
    timestamp_column_name: String,
    value_column_name: String,
    output: Output,
}

fn auto_create_table_type_for_prom_remote_write(
    ctx: &QueryContextRef,
    with_metric_engine: bool,
) -> AutoCreateTableType {
    if with_metric_engine {
        let physical_table = ctx
            .extension(PHYSICAL_TABLE_PARAM)
            .unwrap_or(GREPTIME_PHYSICAL_TABLE)
            .to_string();
        AutoCreateTableType::Logical(physical_table)
    } else {
        AutoCreateTableType::Physical
    }
}

fn required_physical_table_for_create_type(create_type: &AutoCreateTableType) -> Option<&str> {
    match create_type {
        AutoCreateTableType::Logical(physical_table) => Some(physical_table.as_str()),
        _ => None,
    }
}

fn fill_metric_physical_table_options(table_options: &mut HashMap<String, String>) {
    // We always enforce flat format in this ingestion path.
    table_options.insert(SST_FORMAT_KEY.to_string(), "flat".to_string());
    table_options.insert(PHYSICAL_TABLE_METADATA_KEY.to_string(), "true".to_string());
}

#[inline]
fn is_supported(response_type: i32) -> bool {
    // Only supports samples response right now
    response_type == SAMPLES_RESPONSE_TYPE
}

/// Negotiating the content type of the remote read response.
///
/// Response types are taken from the list in the FIFO order. If no response type in `accepted_response_types` is
/// implemented by server, error is returned.
/// For request that do not contain `accepted_response_types` field the SAMPLES response type will be used.
fn negotiate_response_type(accepted_response_types: &[i32]) -> ServerResult<ResponseType> {
    if accepted_response_types.is_empty() {
        return Ok(ResponseType::Samples);
    }

    let response_type = accepted_response_types
        .iter()
        .find(|t| is_supported(**t))
        .with_context(|| error::NotSupportedSnafu {
            feat: format!(
                "server does not support any of the requested response types: {accepted_response_types:?}",
            ),
        })?;

    // It's safe to unwrap here, we known that it should be SAMPLES_RESPONSE_TYPE
    Ok(ResponseType::try_from(*response_type).unwrap())
}

fn resolve_remote_query_target(
    ctx: &QueryContextRef,
    query: &Query,
    table_name: &str,
) -> PermissionTableTarget {
    match prom_store::extract_schema_from_query(query) {
        Some(database) => {
            let (catalog, schema) = parse_optional_catalog_and_schema_from_db_string(&database);
            PermissionTableTarget::new(
                catalog.unwrap_or_else(|| ctx.current_catalog().to_string()),
                schema,
                table_name,
            )
        }
        None => PermissionTableTarget::new(ctx.current_catalog(), ctx.current_schema(), table_name),
    }
}

#[instrument(skip_all, fields(table_name))]
async fn to_query_result(
    table_name: &str,
    timestamp_column_name: &str,
    value_column_name: &str,
    output: Output,
) -> ServerResult<QueryResult> {
    let OutputData::Stream(stream) = output.data else {
        unreachable!()
    };
    let recordbatches = RecordBatches::try_collect(stream)
        .await
        .context(error::CollectRecordbatchSnafu)?;
    Ok(QueryResult {
        timeseries: prom_store::recordbatches_to_timeseries(
            table_name,
            timestamp_column_name,
            value_column_name,
            recordbatches,
        )?,
    })
}

fn resolve_value_column_name(table_name: &str, table: &TableRef) -> Result<String> {
    let columns = table
        .field_columns()
        .map(|column| column.name)
        .collect::<Vec<_>>();

    match columns.as_slice() {
        [] => ColumnNotFoundSnafu {
            msg: format!("value field in table '{table_name}'"),
        }
        .fail(),

        [only] => Ok(only.clone()),

        columns if columns.iter().any(|name| name == greptime_value()) => {
            Ok(greptime_value().to_string())
        }

        columns => AmbiguousValueColumnSnafu {
            table_name: table_name.to_string(),
            field_columns: columns.to_vec(),
        }
        .fail(),
    }
}

fn ensure_string_label_columns(
    schema: &Schema,
    timestamp_column_name: &str,
    value_column_name: &str,
) -> Result<()> {
    for column in schema.column_schemas() {
        if column.name == timestamp_column_name || column.name == value_column_name {
            continue;
        }
        // All remaining columns are emitted as labels, including non-key columns.
        ensure!(
            matches!(column.data_type, ConcreteDataType::String(_)),
            InvalidPromRemoteReadSchemaSnafu {
                detail: format!(
                    "Prometheus remote read only supports string labels: column '{}' has type {:?}",
                    column.name, column.data_type
                ),
            }
        );
    }
    Ok(())
}

impl Instance {
    #[tracing::instrument(skip_all)]
    async fn handle_remote_query(
        &self,
        ctx: &QueryContextRef,
        catalog_name: &str,
        schema_name: &str,
        table_name: &str,
        query: &Query,
    ) -> Result<Option<RemoteQueryOutput>> {
        let Some(table) = self
            .catalog_manager
            .table(catalog_name, schema_name, table_name, Some(ctx))
            .await
            .context(CatalogSnafu)?
        else {
            debug!(
                catalog = catalog_name,
                schema = schema_name,
                table = table_name,
                "Table not found during remote read, returning empty result"
            );
            return Ok(None);
        };

        let timestamp_column_name = table
            .schema()
            .timestamp_column()
            .with_context(|| ColumnNotFoundSnafu {
                msg: format!("time index in table '{table_name}'"),
            })?
            .name
            .clone();

        let value_column_name = resolve_value_column_name(table_name, &table)?;

        ensure_string_label_columns(&table.schema(), &timestamp_column_name, &value_column_name)?;

        let dataframe = self
            .query_engine
            .read_table(table)
            .with_context(|_| ReadTableSnafu {
                table_name: format_full_table_name(catalog_name, schema_name, table_name),
            })?;

        let logical_plan = prom_store::query_to_plan(dataframe, query, &timestamp_column_name)
            .context(PromStoreRemoteQueryPlanSnafu)?;

        debug!(
            "Prometheus remote read, table: {}, logical plan: {}",
            table_name,
            logical_plan.display_indent(),
        );

        let output = self
            .query_engine
            .execute(logical_plan, ctx.clone())
            .await
            .context(ExecLogicalPlanSnafu)?;

        Ok(Some(RemoteQueryOutput {
            table_name: table_name.to_string(),
            timestamp_column_name,
            value_column_name,
            output,
        }))
    }

    #[tracing::instrument(skip_all)]
    async fn handle_remote_queries(
        &self,
        ctx: QueryContextRef,
        queries: &[Query],
        query_targets: &[PermissionTableTarget],
    ) -> ServerResult<Vec<Option<RemoteQueryOutput>>> {
        let mut results = Vec::with_capacity(queries.len());

        for (query, target) in queries.iter().zip(query_targets) {
            let result = self
                .handle_remote_query(&ctx, &target.catalog, &target.schema, &target.table, query)
                .await
                .map_err(BoxedError::new)
                .context(error::ExecuteQuerySnafu)?;

            results.push(result);
        }
        Ok(results)
    }
}

#[async_trait]
impl PendingRowsSchemaAlterer for Instance {
    async fn create_tables_if_missing_batch(
        &self,
        catalog: &str,
        schema: &str,
        tables: &[(&str, &[api::v1::ColumnSchema])],
        with_metric_engine: bool,
        ctx: QueryContextRef,
    ) -> ServerResult<()> {
        if tables.is_empty() {
            return Ok(());
        }

        let create_type = auto_create_table_type_for_prom_remote_write(&ctx, with_metric_engine);
        if let Some(physical_table) = required_physical_table_for_create_type(&create_type) {
            self.create_metric_physical_table_if_missing(
                catalog,
                schema,
                physical_table,
                ctx.clone(),
            )
            .await?;
        }

        let engine = if matches!(create_type, AutoCreateTableType::Logical(_)) {
            METRIC_ENGINE_NAME
        } else {
            common_catalog::consts::default_engine()
        };

        // Check which tables actually still need to be created (may have been
        // concurrently created by another request).
        let mut create_exprs: Vec<CreateTableExpr> = Vec::with_capacity(tables.len());
        let mut per_table_semantics: Option<Option<PerTableSemanticIndex>> = None;
        for &(table_name, request_schema) in tables {
            let existing = self
                .catalog_manager()
                .table(catalog, schema, table_name, Some(ctx.as_ref()))
                .await
                .map_err(BoxedError::new)
                .context(error::ExecuteGrpcQuerySnafu)?;
            if existing.is_some() {
                continue;
            }

            let table_ref = TableReference::full(catalog, schema, table_name);
            let mut create_table_expr = build_create_table_expr(&table_ref, request_schema, engine)
                .map_err(BoxedError::new)
                .context(error::ExecuteGrpcQuerySnafu)?;

            let mut table_options = std::collections::HashMap::with_capacity(4);
            fill_table_options_for_create(&mut table_options, &create_type, &ctx);
            // The batched create path bypasses the operator's auto-create, so
            // fold the per-table semantic index in here too.
            let semantic_index = per_table_semantics
                .get_or_insert_with(|| parse_per_table_semantic_index(&ctx))
                .as_ref();
            apply_per_table_semantic_options(
                &mut table_options,
                semantic_index,
                schema,
                table_name,
            );
            create_table_expr.table_options.extend(table_options);
            create_exprs.push(create_table_expr);
        }

        if create_exprs.is_empty() {
            return Ok(());
        }

        match create_type {
            AutoCreateTableType::Logical(_) => {
                // Use the batch API for logical tables.
                self.statement_executor
                    .create_logical_tables(&create_exprs, ctx, TriggerReason::AutoCreate)
                    .await
                    .map_err(BoxedError::new)
                    .context(error::ExecuteGrpcQuerySnafu)?;
            }
            AutoCreateTableType::Physical => {
                // Physical tables don't have a batch DDL path; create one at a time.
                for mut expr in create_exprs {
                    expr.table_options
                        .insert(SST_FORMAT_KEY.to_string(), "flat".to_string());
                    self.statement_executor
                        .create_table_inner(&mut expr, None, ctx.clone(), TriggerReason::AutoCreate)
                        .await
                        .map_err(BoxedError::new)
                        .context(error::ExecuteGrpcQuerySnafu)?;
                }
            }
            create_type => {
                return error::InvalidPromRemoteRequestSnafu {
                    msg: format!(
                        "prom remote write only supports logical or physical auto-create: {}",
                        create_type.as_str()
                    ),
                }
                .fail();
            }
        }

        Ok(())
    }

    async fn add_missing_prom_tag_columns_batch(
        &self,
        catalog: &str,
        schema: &str,
        tables: &[(&str, &[String])],
        ctx: QueryContextRef,
    ) -> ServerResult<()> {
        if tables.is_empty() {
            return Ok(());
        }

        let alter_exprs: Vec<AlterTableExpr> = tables
            .iter()
            .filter(|(_, columns)| !columns.is_empty())
            .map(|&(table_name, columns)| {
                let add_columns = AddColumns {
                    add_columns: columns
                        .iter()
                        .map(|column_name| AddColumn {
                            column_def: Some(ColumnDef {
                                name: column_name.clone(),
                                data_type: ColumnDataType::String as i32,
                                is_nullable: true,
                                semantic_type: SemanticType::Tag as i32,
                                comment: String::new(),
                                ..Default::default()
                            }),
                            location: None,
                            add_if_not_exists: true,
                        })
                        .collect(),
                };

                AlterTableExpr {
                    catalog_name: catalog.to_string(),
                    schema_name: schema.to_string(),
                    table_name: table_name.to_string(),
                    kind: Some(Kind::AddColumns(add_columns)),
                }
            })
            .collect();

        if alter_exprs.is_empty() {
            return Ok(());
        }

        self.statement_executor
            .alter_logical_tables(alter_exprs, ctx, TriggerReason::AutoAlter)
            .await
            .map_err(BoxedError::new)
            .context(error::ExecuteGrpcQuerySnafu)?;

        Ok(())
    }
}

impl Instance {
    async fn prepare_prom_store_write(
        &self,
        request: RowInsertRequests,
        ctx: QueryContextRef,
    ) -> ServerResult<(RowInsertRequests, QueryContextRef)> {
        PromStoreProtocolHandler::pre_write(self, &request, ctx.clone()).await?;
        Ok((request, Arc::new(ctx.fork())))
    }

    async fn execute_prom_store_write(
        &self,
        request: RowInsertRequests,
        ctx: QueryContextRef,
        with_metric_engine: bool,
    ) -> ServerResult<Output> {
        let output = if with_metric_engine {
            let physical_table = ctx
                .extension(PHYSICAL_TABLE_PARAM)
                .unwrap_or(GREPTIME_PHYSICAL_TABLE)
                .to_string();
            self.handle_metric_row_inserts(request, ctx.clone(), physical_table.clone())
                .await
                .map_err(BoxedError::new)
                .context(error::ExecuteGrpcQuerySnafu)?
        } else {
            self.handle_row_inserts(request, ctx.clone(), true, true)
                .await
                .map_err(BoxedError::new)
                .context(error::ExecuteGrpcQuerySnafu)?
        };

        Ok(output)
    }
}

#[async_trait]
impl PromStoreProtocolHandler for Instance {
    async fn prepare_metric_write(
        &self,
        request: RowInsertRequests,
        ctx: QueryContextRef,
    ) -> ServerResult<operator::insert::PreparedMetricInserts> {
        let physical_table = ctx
            .extension(PHYSICAL_TABLE_PARAM)
            .unwrap_or(GREPTIME_PHYSICAL_TABLE)
            .to_string();
        self.inserter
            .prepare_metric_row_inserts(request, ctx, &self.statement_executor, physical_table)
            .await
            .map_err(BoxedError::new)
            .context(error::ExecuteGrpcQuerySnafu)
    }

    async fn write_prepared_metric(
        &self,
        request: operator::insert::PreparedMetricInserts,
        ctx: QueryContextRef,
    ) -> ServerResult<Output> {
        self.inserter
            .handle_prepared_metric_inserts(request, ctx)
            .await
            .map_err(BoxedError::new)
            .context(error::ExecuteGrpcQuerySnafu)
    }

    async fn check_exemplar_sources(
        &self,
        sources: &std::collections::BTreeSet<String>,
        ctx: QueryContextRef,
    ) -> ServerResult<()> {
        let targets = sources
            .iter()
            .map(|name| {
                PermissionTableTarget::new(ctx.current_catalog(), ctx.current_schema(), name)
            })
            .collect();
        self.check_table_permission(
            &ctx,
            PermissionReq::Action(PROM_STORE_WRITE),
            PermissionTableTargets::resolved(targets),
        )
        .context(AuthSnafu)?;
        Ok(())
    }

    async fn write_exemplars_prepared(
        &self,
        request: RowInsertRequests,
        ctx: QueryContextRef,
    ) -> ServerResult<Output> {
        self.write_exemplars_with_sources(request, ctx, &[]).await
    }

    async fn write_exemplars_with_sources(
        &self,
        request: RowInsertRequests,
        ctx: QueryContextRef,
        sources: &[table::metadata::TableInfoRef],
    ) -> ServerResult<Output> {
        self.inserter
            .handle_exemplar_inserts(request, ctx, &self.statement_executor, sources)
            .await
            .map_err(BoxedError::new)
            .context(error::ExecuteGrpcQuerySnafu)
    }

    async fn pre_write(
        &self,
        request: &RowInsertRequests,
        ctx: QueryContextRef,
    ) -> ServerResult<()> {
        self.plugins
            .get::<PermissionCheckerRef>()
            .as_ref()
            .check_permission(ctx.current_user(), PermissionReq::Action(PROM_STORE_WRITE))
            .context(AuthSnafu)?;
        let interceptor_ref = self
            .plugins
            .get::<PromStoreProtocolInterceptorRef<servers::error::Error>>();
        interceptor_ref.pre_write(request, ctx.clone())?;
        self.check_row_insert_permission(request, &ctx, PermissionReq::Action(PROM_STORE_WRITE))
            .context(AuthSnafu)?;
        Ok(())
    }

    async fn write_prepared(
        &self,
        request: RowInsertRequests,
        ctx: QueryContextRef,
        with_metric_engine: bool,
    ) -> ServerResult<Output> {
        self.execute_prom_store_write(request, ctx, with_metric_engine)
            .await
    }

    async fn write(
        &self,
        request: RowInsertRequests,
        ctx: QueryContextRef,
        with_metric_engine: bool,
    ) -> ServerResult<Output> {
        let (request, ctx) = self.prepare_prom_store_write(request, ctx).await?;
        self.write_prepared(request, ctx, with_metric_engine).await
    }

    async fn write_all(
        &self,
        requests: Vec<(QueryContextRef, RowInsertRequests)>,
        with_metric_engine: bool,
    ) -> ServerResult<Vec<ServerResult<Output>>> {
        let mut prepared = Vec::with_capacity(requests.len());
        for (ctx, request) in requests {
            let (request, ctx) = self.prepare_prom_store_write(request, ctx).await?;
            prepared.push((ctx, request));
        }
        operator::insert::admit_row_insert_batches(&mut prepared)
            .await
            .map_err(BoxedError::new)
            .context(error::ExecuteGrpcQuerySnafu)?;

        let mut outputs = Vec::with_capacity(prepared.len());
        for (ctx, request) in prepared {
            let output = self.write_prepared(request, ctx, with_metric_engine).await;
            let failed = output.is_err();
            outputs.push(output);
            if failed {
                break;
            }
        }
        Ok(outputs)
    }

    #[instrument(skip_all, fields(table_name))]
    async fn read(
        &self,
        request: ReadRequest,
        ctx: QueryContextRef,
    ) -> ServerResult<PromStoreResponse> {
        self.plugins
            .get::<PermissionCheckerRef>()
            .as_ref()
            .check_permission(ctx.current_user(), PermissionReq::Action(PROM_STORE_READ))
            .context(AuthSnafu)?;

        let interceptor_ref = self
            .plugins
            .get::<PromStoreProtocolInterceptorRef<servers::error::Error>>();
        interceptor_ref.pre_read(&request, ctx.clone())?;
        let ctx = Arc::new(ctx.fork());

        let table_names = request
            .queries
            .iter()
            .map(prom_store::table_name)
            .collect::<ServerResult<Vec<_>>>()?;
        let query_targets = request
            .queries
            .iter()
            .zip(&table_names)
            .map(|(query, table_name)| resolve_remote_query_target(&ctx, query, table_name))
            .collect::<Vec<_>>();
        let disallow_cross_catalog_query = self
            .plugins
            .get::<QueryOptions>()
            .map(|opts| opts.disallow_cross_catalog_query)
            .unwrap_or_default();
        if disallow_cross_catalog_query {
            for target in &query_targets {
                validate_catalog_and_schema(&target.catalog, &target.schema, &ctx)
                    .map_err(BoxedError::new)
                    .context(error::ExecuteQuerySnafu)?;
            }
        }
        let targets = self
            .resolve_query_permission_targets(
                PermissionTableTargets::resolved(query_targets.clone()),
                &ctx,
            )
            .await?;
        self.check_table_permission(&ctx, PermissionReq::Action(PROM_STORE_READ), targets)
            .context(AuthSnafu)?;

        let response_type = negotiate_response_type(&request.accepted_response_types)?;

        for query in &request.queries {
            prom_store::validate_matchers(&query.matchers)?;
        }

        // query_to_plan intersects valid time hints with each query's outer range.
        // Other hints do not change raw samples or labels.
        let results = self
            .handle_remote_queries(ctx, &request.queries, &query_targets)
            .await?;

        match response_type {
            ResponseType::Samples => {
                let mut query_results = Vec::with_capacity(results.len());
                let mut map = HashMap::new();
                for result in results {
                    let Some(result) = result else {
                        query_results.push(QueryResult { timeseries: vec![] });
                        continue;
                    };
                    let RemoteQueryOutput {
                        table_name,
                        timestamp_column_name,
                        value_column_name,
                        output,
                    } = result;
                    let plan = output.meta.plan.clone();
                    query_results.push(
                        to_query_result(
                            &table_name,
                            &timestamp_column_name,
                            &value_column_name,
                            output,
                        )
                        .await?,
                    );
                    if let Some(ref plan) = plan {
                        collect_plan_metrics(plan, &mut [&mut map]);
                    }
                }

                let response = ReadResponse {
                    results: query_results,
                };

                let resp_metrics = map
                    .into_iter()
                    .map(|(k, v)| (k, v.into()))
                    .collect::<HashMap<_, _>>();

                // TODO(dennis): may consume too much memory, adds flow control
                Ok(PromStoreResponse {
                    content_type: CONTENT_TYPE_PROTOBUF.clone(),
                    content_encoding: CONTENT_ENCODING_SNAPPY.clone(),
                    resp_metrics,
                    body: prom_store::snappy_compress(&response.encode_to_vec())?,
                })
            }
            ResponseType::StreamedXorChunks => error::NotSupportedSnafu {
                feat: "streamed remote read",
            }
            .fail(),
        }
    }
}

impl Instance {
    async fn create_metric_physical_table_if_missing(
        &self,
        catalog: &str,
        schema: &str,
        physical_table: &str,
        ctx: QueryContextRef,
    ) -> ServerResult<()> {
        let table = self
            .catalog_manager()
            .table(catalog, schema, physical_table, Some(ctx.as_ref()))
            .await
            .map_err(BoxedError::new)
            .context(error::ExecuteGrpcQuerySnafu)?;
        if table.is_some() {
            return Ok(());
        }

        let table_ref = TableReference::full(catalog, schema, physical_table);
        let default_schema = vec![
            api::v1::ColumnSchema {
                column_name: common_query::prelude::greptime_timestamp().to_string(),
                datatype: api::v1::ColumnDataType::TimestampMillisecond as i32,
                semantic_type: api::v1::SemanticType::Timestamp as i32,
                datatype_extension: None,
                options: None,
            },
            api::v1::ColumnSchema {
                column_name: common_query::prelude::greptime_value().to_string(),
                datatype: api::v1::ColumnDataType::Float64 as i32,
                semantic_type: api::v1::SemanticType::Field as i32,
                datatype_extension: None,
                options: None,
            },
        ];
        let mut create_table_expr = build_create_table_expr(
            &table_ref,
            &default_schema,
            common_catalog::consts::default_engine(),
        )
        .map_err(BoxedError::new)
        .context(error::ExecuteGrpcQuerySnafu)?;
        create_table_expr.engine = METRIC_ENGINE_NAME.to_string();
        fill_metric_physical_table_options(&mut create_table_expr.table_options);

        self.statement_executor
            .create_table_inner(&mut create_table_expr, None, ctx, TriggerReason::AutoCreate)
            .await
            .map_err(BoxedError::new)
            .context(error::ExecuteGrpcQuerySnafu)?;

        Ok(())
    }
}

/// This handler is mainly used for `frontend` or `standalone` to directly import
/// the metrics collected by itself, thereby avoiding importing metrics through the network,
/// thus reducing compression and network transmission overhead,
/// so only implement `PromStoreProtocolHandler::write` method.
pub struct ExportMetricHandler {
    inserter: InserterRef,
    statement_executor: Arc<StatementExecutor>,
}

impl ExportMetricHandler {
    pub fn new_handler(
        inserter: InserterRef,
        statement_executor: Arc<StatementExecutor>,
    ) -> PromStoreProtocolHandlerRef {
        Arc::new(Self {
            inserter,
            statement_executor,
        })
    }
}

#[async_trait]
impl PromStoreProtocolHandler for ExportMetricHandler {
    async fn pre_write(
        &self,
        _request: &RowInsertRequests,
        _ctx: QueryContextRef,
    ) -> ServerResult<()> {
        Ok(())
    }

    async fn write_prepared(
        &self,
        request: RowInsertRequests,
        ctx: QueryContextRef,
        _: bool,
    ) -> ServerResult<Output> {
        self.inserter
            .handle_metric_row_inserts(
                request,
                ctx,
                &self.statement_executor,
                GREPTIME_PHYSICAL_TABLE.to_string(),
            )
            .await
            .map_err(BoxedError::new)
            .context(error::ExecuteGrpcQuerySnafu)
    }

    async fn write(
        &self,
        request: RowInsertRequests,
        ctx: QueryContextRef,
        with_metric_engine: bool,
    ) -> ServerResult<Output> {
        self.write_prepared(request, ctx, with_metric_engine).await
    }

    async fn write_all(
        &self,
        mut requests: Vec<(QueryContextRef, RowInsertRequests)>,
        with_metric_engine: bool,
    ) -> ServerResult<Vec<ServerResult<Output>>> {
        operator::insert::admit_row_insert_batches(&mut requests)
            .await
            .map_err(BoxedError::new)
            .context(error::ExecuteGrpcQuerySnafu)?;
        let mut outputs = Vec::with_capacity(requests.len());
        for (ctx, request) in requests {
            let output = self.write_prepared(request, ctx, with_metric_engine).await;
            let failed = output.is_err();
            outputs.push(output);
            if failed {
                break;
            }
        }
        Ok(outputs)
    }

    async fn read(
        &self,
        _request: ReadRequest,
        _ctx: QueryContextRef,
    ) -> ServerResult<PromStoreResponse> {
        unreachable!();
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use api::prom_store::remote::LabelMatcher;
    use common_error::ext::ErrorExt;
    use common_error::status_code::StatusCode;
    use session::context::QueryContext;

    use super::*;

    #[test]
    fn test_native_histogram_remote_read_response_negotiation() {
        for accepted in [
            vec![],
            vec![ResponseType::Samples as i32],
            vec![
                ResponseType::StreamedXorChunks as i32,
                ResponseType::Samples as i32,
            ],
        ] {
            assert_eq!(
                negotiate_response_type(&accepted).unwrap(),
                ResponseType::Samples
            );
        }
        assert!(negotiate_response_type(&[ResponseType::StreamedXorChunks as i32]).is_err());
    }

    #[test]
    fn test_remote_read_label_types() {
        use datatypes::schema::ColumnSchema;

        for (label_type, supported) in [
            (ConcreteDataType::string_datatype(), true),
            (ConcreteDataType::large_string_datatype(), true),
            (ConcreteDataType::utf8_view_datatype(), true),
            (ConcreteDataType::binary_datatype(), false),
            (ConcreteDataType::binary_view_datatype(), false),
            (ConcreteDataType::json_datatype(), false),
            (ConcreteDataType::vector_datatype(3), false),
            (ConcreteDataType::int64_datatype(), false),
            (ConcreteDataType::boolean_datatype(), false),
            (
                common_query::native_histogram::native_histogram_value_type().clone(),
                false,
            ),
        ] {
            let schema = Schema::new(vec![
                ColumnSchema::new(
                    "ts",
                    ConcreteDataType::timestamp_millisecond_datatype(),
                    false,
                ),
                ColumnSchema::new("value", label_type.clone(), true),
                ColumnSchema::new("label", ConcreteDataType::string_datatype(), true),
            ]);
            // Value columns are validated separately, regardless of their type.
            ensure_string_label_columns(&schema, "ts", "value").unwrap();

            let result = ensure_string_label_columns(&schema, "ts", "label");
            assert_eq!(supported, result.is_ok(), "{label_type:?}");
            if let Err(err) = result {
                assert!(matches!(
                    err,
                    crate::error::Error::InvalidPromRemoteReadSchema { .. }
                ));
                assert_eq!(err.status_code(), StatusCode::InvalidArguments);
                assert!(
                    err.to_string()
                        .starts_with("Invalid Prometheus remote read schema: ")
                );
                assert!(
                    err.to_string()
                        .contains("only supports string labels: column 'value'")
                );
            }
        }
    }

    #[test]
    fn test_auto_create_table_type_for_prom_remote_write_metric_engine() {
        let mut query_ctx = QueryContext::with(
            common_catalog::consts::DEFAULT_CATALOG_NAME,
            common_catalog::consts::DEFAULT_SCHEMA_NAME,
        );
        query_ctx.set_extension(PHYSICAL_TABLE_PARAM, "metric_physical".to_string());
        let ctx = Arc::new(query_ctx);

        let create_type = auto_create_table_type_for_prom_remote_write(&ctx, true);
        match create_type {
            AutoCreateTableType::Logical(physical) => assert_eq!(physical, "metric_physical"),
            _ => panic!("expected logical table create type"),
        }
    }

    #[test]
    fn test_auto_create_table_type_for_prom_remote_write_without_metric_engine() {
        let ctx = Arc::new(QueryContext::with(
            common_catalog::consts::DEFAULT_CATALOG_NAME,
            common_catalog::consts::DEFAULT_SCHEMA_NAME,
        ));

        let create_type = auto_create_table_type_for_prom_remote_write(&ctx, false);
        match create_type {
            AutoCreateTableType::Physical => {}
            _ => panic!("expected physical table create type"),
        }
    }

    #[test]
    fn test_required_physical_table_for_create_type() {
        let logical = AutoCreateTableType::Logical("phy_table".to_string());
        assert_eq!(
            Some("phy_table"),
            required_physical_table_for_create_type(&logical)
        );

        let physical = AutoCreateTableType::Physical;
        assert_eq!(None, required_physical_table_for_create_type(&physical));
    }

    #[test]
    fn test_metric_physical_table_options_forces_flat_sst_format() {
        let mut table_options = HashMap::new();

        fill_metric_physical_table_options(&mut table_options);

        assert_eq!(
            Some("flat"),
            table_options.get(SST_FORMAT_KEY).map(String::as_str)
        );
        assert_eq!(
            Some("true"),
            table_options
                .get(PHYSICAL_TABLE_METADATA_KEY)
                .map(String::as_str)
        );
    }

    fn query_with_database(database: &str) -> Query {
        Query {
            matchers: vec![LabelMatcher {
                name: servers::prom_store::DATABASE_LABEL.to_string(),
                value: database.to_string(),
                r#type: api::prom_store::remote::label_matcher::Type::Eq as i32,
            }],
            ..Default::default()
        }
    }

    #[test]
    fn test_resolve_remote_query_target() {
        let ctx = Arc::new(QueryContext::with("request_catalog", "request_schema"));

        assert_eq!(
            PermissionTableTarget::new("request_catalog", "request_schema", "fallback_metric"),
            resolve_remote_query_target(&ctx, &Query::default(), "fallback_metric")
        );
        assert_eq!(
            PermissionTableTarget::new("request_catalog", "selected_schema", "schema_metric"),
            resolve_remote_query_target(
                &ctx,
                &query_with_database("selected_schema"),
                "schema_metric"
            )
        );
        assert_eq!(
            PermissionTableTarget::new("selected_catalog", "selected_schema", "catalog_metric"),
            resolve_remote_query_target(
                &ctx,
                &query_with_database("selected_catalog-selected_schema"),
                "catalog_metric"
            )
        );
    }

    #[test]
    fn test_resolve_remote_query_target_preserves_context() {
        let ctx = Arc::new(QueryContext::with("request_catalog", "request_schema"));

        resolve_remote_query_target(
            &ctx,
            &query_with_database("selected_catalog-selected_schema"),
            "metric",
        );

        assert_eq!("request_catalog", ctx.current_catalog());
        assert_eq!("request_schema", ctx.current_schema());
    }
}