frontend/

instance.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.

pub mod builder;
mod dashboard;
mod grpc;
mod influxdb;
mod jaeger;
mod log_handler;
mod logs;
mod opentsdb;
mod otlp;
pub mod prom_store;
mod promql;
mod region_query;
pub mod standalone;

use std::pin::Pin;
use std::sync::atomic::AtomicBool;
use std::sync::{Arc, atomic};
use std::time::{Duration, SystemTime};

use async_stream::stream;
use async_trait::async_trait;
use auth::{PermissionChecker, PermissionCheckerRef, PermissionReq};
use catalog::CatalogManagerRef;
use catalog::process_manager::{
    ProcessManagerRef, QueryStatement as CatalogQueryStatement, SlowQueryTimer,
};
use client::OutputData;
use common_base::Plugins;
use common_base::cancellation::CancellableFuture;
use common_error::ext::{BoxedError, ErrorExt};
use common_event_recorder::EventRecorderRef;
use common_meta::cache_invalidator::CacheInvalidatorRef;
use common_meta::key::TableMetadataManagerRef;
use common_meta::key::table_name::TableNameKey;
use common_meta::node_manager::NodeManagerRef;
use common_meta::procedure_executor::ProcedureExecutorRef;
use common_query::Output;
use common_recordbatch::RecordBatchStreamWrapper;
use common_recordbatch::error::StreamTimeoutSnafu;
use common_telemetry::logging::SlowQueryOptions;
use common_telemetry::{debug, error, tracing};
use dashmap::DashMap;
use datafusion_expr::LogicalPlan;
use futures::{Stream, StreamExt};
use lazy_static::lazy_static;
use operator::delete::DeleterRef;
use operator::insert::InserterRef;
use operator::statement::{StatementExecutor, StatementExecutorRef};
use partition::manager::PartitionRuleManagerRef;
use pipeline::pipeline_operator::PipelineOperator;
use prometheus::HistogramTimer;
use promql_parser::label::Matcher;
use query::QueryEngineRef;
use query::metrics::OnDone;
use query::parser::{PromQuery, QueryLanguageParser, QueryStatement};
use query::query_engine::DescribeResult;
use query::query_engine::options::{QueryOptions, validate_catalog_and_schema};
use servers::error::{
    self as server_error, AuthSnafu, CommonMetaSnafu, ExecuteQuerySnafu,
    OtlpMetricModeIncompatibleSnafu, ParsePromQLSnafu, UnexpectedResultSnafu,
};
use servers::interceptor::{
    PromQueryInterceptor, PromQueryInterceptorRef, SqlQueryInterceptor, SqlQueryInterceptorRef,
};
use servers::otlp::metrics::legacy_normalize_otlp_name;
use servers::prometheus_handler::PrometheusHandler;
use servers::query_handler::sql::SqlQueryHandler;
use session::context::{Channel, QueryContextRef};
use session::table_name::table_idents_to_full_name;
use snafu::prelude::*;
use sql::ast::ObjectNamePartExt;
use sql::dialect::Dialect;
use sql::parser::{ParseOptions, ParserContext};
use sql::statements::comment::CommentObject;
use sql::statements::copy::{CopyDatabase, CopyTable};
use sql::statements::statement::Statement;
use sql::statements::tql::Tql;
use sqlparser::ast::{AnalyzeFormat, ObjectName};
pub use standalone::StandaloneDatanodeManager;
use table::requests::{OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM};
use tracing::Span;

use crate::error::{
    self, Error, ExecLogicalPlanSnafu, ExecutePromqlSnafu, ExternalSnafu, InvalidSqlSnafu,
    ParseSqlSnafu, PermissionSnafu, PlanStatementSnafu, Result, SqlExecInterceptedSnafu,
    StatementTimeoutSnafu, TableOperationSnafu,
};
use crate::service_config::InfluxdbMergeMode;
use crate::stream_wrapper::CancellableStreamWrapper;

lazy_static! {
    static ref OTLP_LEGACY_DEFAULT_VALUE: String = "legacy".to_string();
}

/// The frontend instance contains necessary components, and implements many
/// traits, like [`servers::query_handler::grpc::GrpcQueryHandler`],
/// [`servers::query_handler::sql::SqlQueryHandler`], etc.
#[derive(Clone)]
pub struct Instance {
    frontend_peer_addr: String,
    catalog_manager: CatalogManagerRef,
    pipeline_operator: Arc<PipelineOperator>,
    statement_executor: Arc<StatementExecutor>,
    query_engine: QueryEngineRef,
    plugins: Plugins,
    inserter: InserterRef,
    deleter: DeleterRef,
    table_metadata_manager: TableMetadataManagerRef,
    event_recorder: Option<EventRecorderRef>,
    process_manager: ProcessManagerRef,
    slow_query_options: SlowQueryOptions,
    influxdb_default_merge_mode: InfluxdbMergeMode,
    trace_ingest_chunk_size: usize,
    suspend: Arc<AtomicBool>,

    // cache for otlp metrics
    // first layer key: db-string
    // key: direct input metric name
    // value: if runs in legacy mode
    otlp_metrics_table_legacy_cache: DashMap<String, DashMap<String, bool>>,
}

impl Instance {
    pub fn frontend_peer_addr(&self) -> &str {
        &self.frontend_peer_addr
    }

    pub fn catalog_manager(&self) -> &CatalogManagerRef {
        &self.catalog_manager
    }

    pub fn query_engine(&self) -> &QueryEngineRef {
        &self.query_engine
    }

    pub fn plugins(&self) -> &Plugins {
        &self.plugins
    }

    pub fn statement_executor(&self) -> &StatementExecutorRef {
        &self.statement_executor
    }

    pub fn table_metadata_manager(&self) -> &TableMetadataManagerRef {
        &self.table_metadata_manager
    }

    pub fn inserter(&self) -> &InserterRef {
        &self.inserter
    }

    pub fn process_manager(&self) -> &ProcessManagerRef {
        &self.process_manager
    }

    pub fn node_manager(&self) -> &NodeManagerRef {
        self.inserter.node_manager()
    }

    pub fn partition_manager(&self) -> &PartitionRuleManagerRef {
        self.inserter.partition_manager()
    }

    pub fn cache_invalidator(&self) -> &CacheInvalidatorRef {
        self.statement_executor.cache_invalidator()
    }

    pub fn procedure_executor(&self) -> &ProcedureExecutorRef {
        self.statement_executor.procedure_executor()
    }

    pub fn suspend_state(&self) -> Arc<AtomicBool> {
        self.suspend.clone()
    }

    pub(crate) fn is_suspended(&self) -> bool {
        self.suspend.load(atomic::Ordering::Relaxed)
    }
}

fn parse_stmt(sql: &str, dialect: &(dyn Dialect + Send + Sync)) -> Result<Vec<Statement>> {
    ParserContext::create_with_dialect(sql, dialect, ParseOptions::default()).context(ParseSqlSnafu)
}

fn validate_analyze_stream_statement(stmt: &mut Statement) -> Result<()> {
    let Statement::Explain(explain) = stmt else {
        return InvalidSqlSnafu {
            err_msg: "only EXPLAIN ANALYZE VERBOSE statement is supported",
        }
        .fail();
    };
    ensure!(
        explain.analyze && explain.verbose,
        InvalidSqlSnafu {
            err_msg: "statement must be EXPLAIN ANALYZE VERBOSE"
        }
    );
    match explain.format {
        None | Some(AnalyzeFormat::JSON) => {
            // Keep explicit FORMAT JSON accepted, but pass JSON through
            // QueryContext.explain_format instead of the statement to avoid the
            // planner's current `EXPLAIN VERBOSE with FORMAT` limitation.
            explain.format = None;
            Ok(())
        }
        Some(_) => InvalidSqlSnafu {
            err_msg: "only FORMAT JSON is supported for analyze stream",
        }
        .fail(),
    }
}

impl Instance {
    fn statement_slow_query_timer(
        &self,
        stmt: &Statement,
        schema_name: String,
    ) -> Option<SlowQueryTimer> {
        if !stmt.is_readonly() || !self.slow_query_options.enable {
            return None;
        }

        self.event_recorder.clone().map(|event_recorder| {
            SlowQueryTimer::new(
                CatalogQueryStatement::Sql(stmt.clone()),
                schema_name,
                self.slow_query_options.threshold,
                self.slow_query_options.sample_ratio,
                self.slow_query_options.record_type,
                event_recorder,
            )
        })
    }

    async fn query_statement(&self, stmt: Statement, query_ctx: QueryContextRef) -> Result<Output> {
        check_permission(self.plugins.clone(), &stmt, &query_ctx)?;

        let query_interceptor = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
        let query_interceptor = query_interceptor.as_ref();

        if should_track_statement_process(&stmt) {
            let catalog_name = query_ctx.current_catalog().to_string();
            let schema_name = query_ctx.current_schema();
            let slow_query_timer = self.statement_slow_query_timer(&stmt, schema_name.clone());

            let ticket = self.process_manager.register_query(
                catalog_name,
                vec![schema_name],
                stmt.to_string(),
                query_ctx.conn_info().to_string(),
                Some(query_ctx.process_id()),
                slow_query_timer,
            );

            let query_fut = self.exec_statement_with_timeout(stmt, query_ctx, query_interceptor);

            CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
                .await
                .map_err(|_| error::CancelledSnafu.build())?
                .map(|output| {
                    let Output { meta, data } = output;

                    let data = match data {
                        OutputData::Stream(stream) => OutputData::Stream(Box::pin(
                            CancellableStreamWrapper::new(stream, ticket),
                        )),
                        other => other,
                    };
                    Output { data, meta }
                })
        } else {
            self.exec_statement_with_timeout(stmt, query_ctx, query_interceptor)
                .await
        }
    }

    async fn exec_statement_with_timeout(
        &self,
        stmt: Statement,
        query_ctx: QueryContextRef,
        query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
    ) -> Result<Output> {
        let timeout = derive_timeout(&stmt, &query_ctx);
        match timeout {
            Some(timeout) => {
                let start = tokio::time::Instant::now();
                let output = tokio::time::timeout(
                    timeout,
                    self.exec_statement(stmt, query_ctx, query_interceptor),
                )
                .await
                .map_err(|_| StatementTimeoutSnafu.build())??;
                // compute remaining timeout
                let remaining_timeout = timeout.checked_sub(start.elapsed()).unwrap_or_default();
                attach_timeout(output, remaining_timeout)
            }
            None => {
                self.exec_statement(stmt, query_ctx, query_interceptor)
                    .await
            }
        }
    }

    async fn exec_statement(
        &self,
        stmt: Statement,
        query_ctx: QueryContextRef,
        query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
    ) -> Result<Output> {
        match stmt {
            Statement::Query(_) | Statement::Explain(_) | Statement::Delete(_) => {
                // TODO: remove this when format is supported in datafusion
                if let Statement::Explain(explain) = &stmt
                    && let Some(format) = explain.format()
                {
                    query_ctx.set_explain_format(format.to_string());
                }

                self.plan_and_exec_sql(stmt, &query_ctx, query_interceptor)
                    .await
            }
            Statement::Tql(tql) => {
                self.plan_and_exec_tql(&query_ctx, query_interceptor, tql)
                    .await
            }
            _ => {
                query_interceptor.pre_execute(Some(&stmt), None, query_ctx.clone())?;
                self.statement_executor
                    .execute_sql(stmt, query_ctx)
                    .await
                    .context(TableOperationSnafu)
            }
        }
    }

    async fn plan_and_exec_sql(
        &self,
        stmt: Statement,
        query_ctx: &QueryContextRef,
        query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
    ) -> Result<Output> {
        let stmt = QueryStatement::Sql(stmt);
        let plan = self
            .statement_executor
            .plan(&stmt, query_ctx.clone())
            .await?;
        let QueryStatement::Sql(stmt) = stmt else {
            unreachable!()
        };
        query_interceptor.pre_execute(Some(&stmt), Some(&plan), query_ctx.clone())?;

        self.statement_executor
            .exec_plan(plan, query_ctx.clone())
            .await
            .context(TableOperationSnafu)
    }

    async fn plan_and_exec_tql(
        &self,
        query_ctx: &QueryContextRef,
        query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
        tql: Tql,
    ) -> Result<Output> {
        let plan = self
            .statement_executor
            .plan_tql(tql.clone(), query_ctx)
            .await?;
        query_interceptor.pre_execute(
            Some(&Statement::Tql(tql)),
            Some(&plan),
            query_ctx.clone(),
        )?;
        self.statement_executor
            .exec_plan(plan, query_ctx.clone())
            .await
            .context(TableOperationSnafu)
    }

    async fn check_otlp_legacy(
        &self,
        names: &[&String],
        ctx: QueryContextRef,
    ) -> server_error::Result<bool> {
        let db_string = ctx.get_db_string();
        // fast cache check
        let cache = self
            .otlp_metrics_table_legacy_cache
            .entry(db_string.clone())
            .or_default();
        if let Some(flag) = fast_legacy_check(&cache, names)? {
            return Ok(flag);
        }
        // release cache reference to avoid lock contention
        drop(cache);

        let catalog = ctx.current_catalog();
        let schema = ctx.current_schema();

        // query legacy table names
        let normalized_names = names
            .iter()
            .map(|n| legacy_normalize_otlp_name(n))
            .collect::<Vec<_>>();
        let table_names = normalized_names
            .iter()
            .map(|n| TableNameKey::new(catalog, &schema, n))
            .collect::<Vec<_>>();
        let table_values = self
            .table_metadata_manager()
            .table_name_manager()
            .batch_get(table_names)
            .await
            .context(CommonMetaSnafu)?;
        let table_ids = table_values
            .into_iter()
            .filter_map(|v| v.map(|vi| vi.table_id()))
            .collect::<Vec<_>>();

        // means no existing table is found, use new mode
        if table_ids.is_empty() {
            let cache = self
                .otlp_metrics_table_legacy_cache
                .entry(db_string)
                .or_default();
            names.iter().for_each(|name| {
                cache.insert((*name).clone(), false);
            });
            return Ok(false);
        }

        // has existing table, check table options
        let table_infos = self
            .table_metadata_manager()
            .table_info_manager()
            .batch_get(&table_ids)
            .await
            .context(CommonMetaSnafu)?;
        let options = table_infos
            .values()
            .map(|info| {
                info.table_info
                    .meta
                    .options
                    .extra_options
                    .get(OTLP_METRIC_COMPAT_KEY)
                    .unwrap_or(&OTLP_LEGACY_DEFAULT_VALUE)
            })
            .collect::<Vec<_>>();
        let cache = self
            .otlp_metrics_table_legacy_cache
            .entry(db_string)
            .or_default();
        if !options.is_empty() {
            // check value consistency
            let has_prom = options.iter().any(|opt| *opt == OTLP_METRIC_COMPAT_PROM);
            let has_legacy = options
                .iter()
                .any(|opt| *opt == OTLP_LEGACY_DEFAULT_VALUE.as_str());
            ensure!(!(has_prom && has_legacy), OtlpMetricModeIncompatibleSnafu);
            let flag = has_legacy;
            names.iter().for_each(|name| {
                cache.insert((*name).clone(), flag);
            });
            Ok(flag)
        } else {
            // no table info, use new mode
            names.iter().for_each(|name| {
                cache.insert((*name).clone(), false);
            });
            Ok(false)
        }
    }
}

fn fast_legacy_check(
    cache: &DashMap<String, bool>,
    names: &[&String],
) -> server_error::Result<Option<bool>> {
    let hit_cache = names
        .iter()
        .filter_map(|name| cache.get(*name))
        .collect::<Vec<_>>();
    if !hit_cache.is_empty() {
        let hit_legacy = hit_cache.iter().any(|en| *en.value());
        let hit_prom = hit_cache.iter().any(|en| !*en.value());

        // hit but have true and false, means both legacy and new mode are used
        // we cannot handle this case, so return error
        // add doc links in err msg later
        ensure!(!(hit_legacy && hit_prom), OtlpMetricModeIncompatibleSnafu);

        let flag = hit_legacy;
        // drop hit_cache to release references before inserting to avoid deadlock
        drop(hit_cache);

        // set cache for all names
        names.iter().for_each(|name| {
            if !cache.contains_key(*name) {
                cache.insert((*name).clone(), flag);
            }
        });
        Ok(Some(flag))
    } else {
        Ok(None)
    }
}

/// If the relevant variables are set, the timeout is enforced for all PostgreSQL statements.
/// For MySQL, it applies only to read-only statements.
fn derive_timeout(stmt: &Statement, query_ctx: &QueryContextRef) -> Option<Duration> {
    let query_timeout = query_ctx.query_timeout()?;
    if query_timeout.is_zero() {
        return None;
    }
    match query_ctx.channel() {
        Channel::Mysql if stmt.is_readonly() => Some(query_timeout),
        Channel::Postgres => Some(query_timeout),
        _ => None,
    }
}

/// Derives timeout for plan execution.
fn derive_timeout_for_plan(plan: &LogicalPlan, query_ctx: &QueryContextRef) -> Option<Duration> {
    let query_timeout = query_ctx.query_timeout()?;
    if query_timeout.is_zero() {
        return None;
    }
    match query_ctx.channel() {
        Channel::Mysql if is_readonly_plan(plan) => Some(query_timeout),
        Channel::Postgres => Some(query_timeout),
        _ => None,
    }
}

fn attach_timeout(output: Output, mut timeout: Duration) -> Result<Output> {
    if timeout.is_zero() {
        return StatementTimeoutSnafu.fail();
    }

    let output = match output.data {
        OutputData::AffectedRows(_) | OutputData::RecordBatches(_) => output,
        OutputData::Stream(mut stream) => {
            let schema = stream.schema();
            let s = Box::pin(stream! {
                let mut start = tokio::time::Instant::now();
                while let Some(item) = tokio::time::timeout(timeout, stream.next()).await.map_err(|_| StreamTimeoutSnafu.build())? {
                    yield item;

                    let now = tokio::time::Instant::now();
                    timeout = timeout.checked_sub(now - start).unwrap_or(Duration::ZERO);
                    start = now;
                    // tokio::time::timeout may not return an error immediately when timeout is 0.
                    if timeout.is_zero() {
                        StreamTimeoutSnafu.fail()?;
                    }
                }
            }) as Pin<Box<dyn Stream<Item = _> + Send>>;
            let stream = RecordBatchStreamWrapper {
                schema,
                stream: s,
                output_ordering: None,
                metrics: Default::default(),
                span: Span::current(),
            };
            Output::new(OutputData::Stream(Box::pin(stream)), output.meta)
        }
    };

    Ok(output)
}

impl Instance {
    #[tracing::instrument(skip_all, name = "SqlQueryHandler::do_analyze_stream_query")]
    async fn do_analyze_stream_query_inner(
        &self,
        query: &str,
        query_ctx: QueryContextRef,
    ) -> Result<Output> {
        ensure!(!self.is_suspended(), error::SuspendedSnafu);

        let query_interceptor_opt = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
        let query_interceptor = query_interceptor_opt.as_ref();
        let query = query_interceptor.pre_parsing(query, query_ctx.clone())?;
        let mut stmts = parse_stmt(query.as_ref(), query_ctx.sql_dialect())
            .and_then(|stmts| query_interceptor.post_parsing(stmts, query_ctx.clone()))?;

        ensure!(
            stmts.len() == 1,
            InvalidSqlSnafu {
                err_msg: "only single EXPLAIN ANALYZE VERBOSE statement is supported"
            }
        );
        let mut stmt = stmts.remove(0);
        validate_analyze_stream_statement(&mut stmt)?;
        query_ctx.set_explain_format(AnalyzeFormat::JSON.to_string());

        let checker_ref = self.plugins.get::<PermissionCheckerRef>();
        checker_ref
            .as_ref()
            .check_permission_with_context(
                query_ctx.current_user(),
                PermissionReq::SqlStatement(&stmt),
                Some(&query_ctx.current_schema()),
            )
            .context(PermissionSnafu)?;
        check_permission(self.plugins.clone(), &stmt, &query_ctx)?;
        let catalog_name = query_ctx.current_catalog().to_string();
        let schema_name = query_ctx.current_schema();
        let slow_query_timer = self.statement_slow_query_timer(&stmt, schema_name.clone());
        let ticket = self.process_manager.register_query(
            catalog_name,
            vec![schema_name],
            stmt.to_string(),
            query_ctx.conn_info().to_string(),
            Some(query_ctx.process_id()),
            slow_query_timer,
        );
        let query_fut =
            self.exec_statement_with_timeout(stmt, query_ctx.clone(), query_interceptor);
        let output = CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
            .await
            .map_err(|_| error::CancelledSnafu.build())??;
        let Output { meta, data } = output;
        let data = match data {
            OutputData::Stream(stream) => OutputData::Stream(Box::pin(
                CancellableStreamWrapper::new_cancel_on_drop(stream, ticket),
            )),
            other => other,
        };
        query_interceptor.post_execute(Output { data, meta }, query_ctx)
    }

    #[tracing::instrument(skip_all, name = "SqlQueryHandler::do_query")]
    async fn do_query_inner(&self, query: &str, query_ctx: QueryContextRef) -> Vec<Result<Output>> {
        if self.is_suspended() {
            return vec![error::SuspendedSnafu {}.fail()];
        }

        let query_interceptor_opt = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
        let query_interceptor = query_interceptor_opt.as_ref();
        let query = match query_interceptor.pre_parsing(query, query_ctx.clone()) {
            Ok(q) => q,
            Err(e) => return vec![Err(e)],
        };

        let checker_ref = self.plugins.get::<PermissionCheckerRef>();
        let checker = checker_ref.as_ref();

        match parse_stmt(query.as_ref(), query_ctx.sql_dialect())
            .and_then(|stmts| query_interceptor.post_parsing(stmts, query_ctx.clone()))
        {
            Ok(stmts) => {
                if stmts.is_empty() {
                    return vec![
                        InvalidSqlSnafu {
                            err_msg: "empty statements",
                        }
                        .fail(),
                    ];
                }

                let mut results = Vec::with_capacity(stmts.len());
                for stmt in stmts {
                    if let Err(e) = checker
                        .check_permission_with_context(
                            query_ctx.current_user(),
                            PermissionReq::SqlStatement(&stmt),
                            Some(&query_ctx.current_schema()),
                        )
                        .context(PermissionSnafu)
                    {
                        results.push(Err(e));
                        break;
                    }

                    match self.query_statement(stmt.clone(), query_ctx.clone()).await {
                        Ok(output) => {
                            let output_result =
                                query_interceptor.post_execute(output, query_ctx.clone());
                            results.push(output_result);
                        }
                        Err(e) => {
                            if e.status_code().should_log_error() {
                                error!(e; "Failed to execute query: {stmt}");
                            } else {
                                debug!("Failed to execute query: {stmt}, {e}");
                            }
                            results.push(Err(e));
                            break;
                        }
                    }
                }
                results
            }
            Err(e) => {
                vec![Err(e)]
            }
        }
    }

    async fn exec_plan(&self, plan: LogicalPlan, query_ctx: QueryContextRef) -> Result<Output> {
        self.query_engine
            .execute(plan, query_ctx)
            .await
            .context(ExecLogicalPlanSnafu)
    }

    async fn exec_plan_with_timeout(
        &self,
        plan: LogicalPlan,
        query_ctx: QueryContextRef,
    ) -> Result<Output> {
        let timeout = derive_timeout_for_plan(&plan, &query_ctx);
        match timeout {
            Some(timeout) => {
                let start = tokio::time::Instant::now();
                let output = tokio::time::timeout(timeout, self.exec_plan(plan, query_ctx))
                    .await
                    .map_err(|_| StatementTimeoutSnafu.build())??;
                let remaining_timeout = timeout.checked_sub(start.elapsed()).unwrap_or_default();
                attach_timeout(output, remaining_timeout)
            }
            None => self.exec_plan(plan, query_ctx).await,
        }
    }

    async fn do_exec_plan_inner(
        &self,
        plan: LogicalPlan,
        stmt: Option<Statement>,
        query_ctx: QueryContextRef,
    ) -> Result<Output> {
        ensure!(!self.is_suspended(), error::SuspendedSnafu);

        let query_interceptor_opt = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
        let query_interceptor = query_interceptor_opt.as_ref();

        query_interceptor.pre_execute(stmt.as_ref(), Some(&plan), query_ctx.clone())?;

        let query = stmt
            .as_ref()
            .map(|s| s.to_string())
            .unwrap_or_else(|| plan.display_indent().to_string());

        let plan_is_readonly = is_readonly_plan(&plan);
        let result = if should_track_plan_process(stmt.as_ref(), &plan) {
            let catalog_name = query_ctx.current_catalog().to_string();
            let schema_name = query_ctx.current_schema();
            let slow_query_timer = if plan_is_readonly {
                self.slow_query_options
                    .enable
                    .then(|| self.event_recorder.clone())
                    .flatten()
                    .map(|event_recorder| {
                        SlowQueryTimer::new(
                            CatalogQueryStatement::Plan(query.clone()),
                            schema_name.clone(),
                            self.slow_query_options.threshold,
                            self.slow_query_options.sample_ratio,
                            self.slow_query_options.record_type,
                            event_recorder,
                        )
                    })
            } else {
                None
            };

            let ticket = self.process_manager.register_query(
                catalog_name,
                vec![schema_name],
                query,
                query_ctx.conn_info().to_string(),
                Some(query_ctx.process_id()),
                slow_query_timer,
            );

            let query_fut = self.exec_plan_with_timeout(plan, query_ctx.clone());

            CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
                .await
                .map_err(|_| error::CancelledSnafu.build())?
                .map(|output| {
                    let Output { meta, data } = output;

                    let data = match data {
                        OutputData::Stream(stream) => OutputData::Stream(Box::pin(
                            CancellableStreamWrapper::new(stream, ticket),
                        )),
                        other => other,
                    };
                    Output { data, meta }
                })
        } else {
            self.exec_plan_with_timeout(plan, query_ctx.clone()).await
        };

        result.and_then(|output| query_interceptor.post_execute(output, query_ctx))
    }

    #[tracing::instrument(skip_all, name = "SqlQueryHandler::do_promql_query")]
    async fn do_promql_query_inner(
        &self,
        query: &PromQuery,
        query_ctx: QueryContextRef,
    ) -> Vec<Result<Output>> {
        if self.is_suspended() {
            return vec![error::SuspendedSnafu {}.fail()];
        }

        // check will be done in prometheus handler's do_query
        let result = PrometheusHandler::do_query(self, query, query_ctx)
            .await
            .with_context(|_| ExecutePromqlSnafu {
                query: format!("{query:?}"),
            });
        vec![result]
    }

    async fn do_describe_inner(
        &self,
        stmt: Statement,
        query_ctx: QueryContextRef,
    ) -> Result<Option<DescribeResult>> {
        ensure!(!self.is_suspended(), error::SuspendedSnafu);

        // EXPLAIN / EXPLAIN ANALYZE wrap an inner statement; describe them when the
        // wrapped statement is something we already plan (so that bind parameters
        // in the inner query get their types inferred). See #8029.
        let is_inner_plannable = |s: &Statement| {
            matches!(
                s,
                Statement::Insert(_) | Statement::Query(_) | Statement::Delete(_)
            )
        };
        let plannable = is_inner_plannable(&stmt)
            || matches!(&stmt, Statement::Explain(explain) if is_inner_plannable(explain.statement.as_ref()));

        if plannable {
            self.plugins
                .get::<PermissionCheckerRef>()
                .as_ref()
                .check_permission_with_context(
                    query_ctx.current_user(),
                    PermissionReq::SqlStatement(&stmt),
                    Some(&query_ctx.current_schema()),
                )
                .context(PermissionSnafu)?;

            let plan = self
                .query_engine
                .planner()
                .plan(&QueryStatement::Sql(stmt), query_ctx.clone())
                .await
                .context(PlanStatementSnafu)?;
            self.query_engine
                .describe(plan, query_ctx)
                .await
                .map(Some)
                .context(error::DescribeStatementSnafu)
        } else {
            Ok(None)
        }
    }

    async fn is_valid_schema_inner(&self, catalog: &str, schema: &str) -> Result<bool> {
        self.catalog_manager
            .schema_exists(catalog, schema, None)
            .await
            .context(error::CatalogSnafu)
    }
}

#[async_trait]
impl SqlQueryHandler for Instance {
    async fn do_query(
        &self,
        query: &str,
        query_ctx: QueryContextRef,
    ) -> Vec<server_error::Result<Output>> {
        self.do_query_inner(query, query_ctx)
            .await
            .into_iter()
            .map(|result| result.map_err(BoxedError::new).context(ExecuteQuerySnafu))
            .collect()
    }

    async fn do_analyze_stream_query(
        &self,
        query: &str,
        query_ctx: QueryContextRef,
    ) -> server_error::Result<Output> {
        self.do_analyze_stream_query_inner(query, query_ctx)
            .await
            .map_err(BoxedError::new)
            .context(ExecuteQuerySnafu)
    }

    async fn do_exec_plan(
        &self,
        plan: LogicalPlan,
        stmt: Option<Statement>,
        query_ctx: QueryContextRef,
    ) -> server_error::Result<Output> {
        self.do_exec_plan_inner(plan, stmt, query_ctx)
            .await
            .map_err(BoxedError::new)
            .context(server_error::ExecutePlanSnafu)
    }

    async fn do_promql_query(
        &self,
        query: &PromQuery,
        query_ctx: QueryContextRef,
    ) -> Vec<server_error::Result<Output>> {
        self.do_promql_query_inner(query, query_ctx)
            .await
            .into_iter()
            .map(|result| result.map_err(BoxedError::new).context(ExecuteQuerySnafu))
            .collect()
    }

    async fn do_describe(
        &self,
        stmt: Statement,
        query_ctx: QueryContextRef,
    ) -> server_error::Result<Option<DescribeResult>> {
        self.do_describe_inner(stmt, query_ctx)
            .await
            .map_err(BoxedError::new)
            .context(server_error::DescribeStatementSnafu)
    }

    async fn is_valid_schema(&self, catalog: &str, schema: &str) -> server_error::Result<bool> {
        self.is_valid_schema_inner(catalog, schema)
            .await
            .map_err(BoxedError::new)
            .context(server_error::CheckDatabaseValiditySnafu)
    }
}

/// Attaches a timer to the output and observes it once the output is exhausted.
pub fn attach_timer(output: Output, timer: HistogramTimer) -> Output {
    match output.data {
        OutputData::AffectedRows(_) | OutputData::RecordBatches(_) => output,
        OutputData::Stream(stream) => {
            let stream = OnDone::new(stream, move || {
                timer.observe_duration();
            });
            Output::new(OutputData::Stream(Box::pin(stream)), output.meta)
        }
    }
}

#[async_trait]
impl PrometheusHandler for Instance {
    #[tracing::instrument(skip_all)]
    async fn do_query(
        &self,
        query: &PromQuery,
        query_ctx: QueryContextRef,
    ) -> server_error::Result<Output> {
        let interceptor = self
            .plugins
            .get::<PromQueryInterceptorRef<server_error::Error>>();

        self.plugins
            .get::<PermissionCheckerRef>()
            .as_ref()
            .check_permission(query_ctx.current_user(), PermissionReq::PromQuery)
            .context(AuthSnafu)?;

        let stmt = QueryLanguageParser::parse_promql(query, &query_ctx).with_context(|_| {
            ParsePromQLSnafu {
                query: query.clone(),
            }
        })?;

        let plan = self
            .statement_executor
            .plan(&stmt, query_ctx.clone())
            .await
            .map_err(BoxedError::new)
            .context(ExecuteQuerySnafu)?;

        let QueryStatement::Promql(eval_stmt, _) = &stmt else {
            unreachable!("query is parsed from promql");
        };

        interceptor.pre_execute(query, &eval_stmt.expr, Some(&plan), query_ctx.clone())?;

        // Take the EvalStmt from the original QueryStatement and use it to create the CatalogQueryStatement.
        let query_statement = if let QueryStatement::Promql(eval_stmt, alias) = stmt {
            CatalogQueryStatement::Promql(eval_stmt, alias)
        } else {
            // It should not happen since the query is already parsed successfully.
            return UnexpectedResultSnafu {
                reason: "The query should always be promql.".to_string(),
            }
            .fail();
        };
        let raw_query = query_statement.to_string();

        let slow_query_timer = self
            .slow_query_options
            .enable
            .then(|| self.event_recorder.clone())
            .flatten()
            .map(|event_recorder| {
                SlowQueryTimer::new(
                    query_statement,
                    query_ctx.current_schema(),
                    self.slow_query_options.threshold,
                    self.slow_query_options.sample_ratio,
                    self.slow_query_options.record_type,
                    event_recorder,
                )
            });

        let ticket = self.process_manager.register_query(
            query_ctx.current_catalog().to_string(),
            vec![query_ctx.current_schema()],
            raw_query,
            query_ctx.conn_info().to_string(),
            Some(query_ctx.process_id()),
            slow_query_timer,
        );

        let query_fut = self.statement_executor.exec_plan(plan, query_ctx.clone());

        let output = CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
            .await
            .map_err(|_| servers::error::CancelledSnafu.build())?
            .map(|output| {
                let Output { meta, data } = output;
                let data = match data {
                    OutputData::Stream(stream) => {
                        OutputData::Stream(Box::pin(CancellableStreamWrapper::new(stream, ticket)))
                    }
                    other => other,
                };
                Output { data, meta }
            })
            .map_err(BoxedError::new)
            .context(ExecuteQuerySnafu)?;

        Ok(interceptor.post_execute(output, query_ctx)?)
    }

    async fn query_metric_names(
        &self,
        matchers: Vec<Matcher>,
        ctx: &QueryContextRef,
    ) -> server_error::Result<Vec<String>> {
        self.handle_query_metric_names(matchers, ctx)
            .await
            .map_err(BoxedError::new)
            .context(ExecuteQuerySnafu)
    }

    async fn query_label_values(
        &self,
        metric: String,
        label_name: String,
        matchers: Vec<Matcher>,
        start: SystemTime,
        end: SystemTime,
        ctx: &QueryContextRef,
    ) -> server_error::Result<Vec<String>> {
        self.handle_query_label_values(metric, label_name, matchers, start, end, ctx)
            .await
            .map_err(BoxedError::new)
            .context(ExecuteQuerySnafu)
    }

    fn catalog_manager(&self) -> CatalogManagerRef {
        self.catalog_manager.clone()
    }
}

/// Validate `stmt.database` permission if it's presented.
macro_rules! validate_db_permission {
    ($stmt: expr, $query_ctx: expr) => {
        if let Some(database) = &$stmt.database {
            validate_catalog_and_schema($query_ctx.current_catalog(), database, $query_ctx)
                .map_err(BoxedError::new)
                .context(SqlExecInterceptedSnafu)?;
        }
    };
}

pub fn check_permission(
    plugins: Plugins,
    stmt: &Statement,
    query_ctx: &QueryContextRef,
) -> Result<()> {
    let need_validate = plugins
        .get::<QueryOptions>()
        .map(|opts| opts.disallow_cross_catalog_query)
        .unwrap_or_default();

    if !need_validate {
        return Ok(());
    }

    match stmt {
        // Will be checked in execution.
        // TODO(dennis): add a hook for admin commands.
        Statement::Admin(_) => {}
        // These are executed by query engine, and will be checked there.
        Statement::Query(_)
        | Statement::Explain(_)
        | Statement::Tql(_)
        | Statement::Delete(_)
        | Statement::DeclareCursor(_)
        | Statement::Copy(sql::statements::copy::Copy::CopyQueryTo(_)) => {}
        // database ops won't be checked
        Statement::CreateDatabase(_)
        | Statement::ShowDatabases(_)
        | Statement::DropDatabase(_)
        | Statement::AlterDatabase(_)
        | Statement::DropFlow(_)
        | Statement::Use(_) => {}
        #[cfg(feature = "enterprise")]
        Statement::DropTrigger(_) => {}
        Statement::ShowCreateDatabase(stmt) => {
            validate_database(&stmt.database_name, query_ctx)?;
        }
        Statement::ShowCreateTable(stmt) => {
            validate_param(&stmt.table_name, query_ctx)?;
        }
        Statement::ShowCreateFlow(stmt) => {
            validate_flow(&stmt.flow_name, query_ctx)?;
        }
        #[cfg(feature = "enterprise")]
        Statement::ShowCreateTrigger(stmt) => {
            validate_param(&stmt.trigger_name, query_ctx)?;
        }
        Statement::ShowCreateView(stmt) => {
            validate_param(&stmt.view_name, query_ctx)?;
        }
        Statement::CreateExternalTable(stmt) => {
            validate_param(&stmt.name, query_ctx)?;
        }
        Statement::CreateFlow(stmt) => {
            // TODO: should also validate source table name here?
            validate_param(&stmt.sink_table_name, query_ctx)?;
        }
        #[cfg(feature = "enterprise")]
        Statement::CreateTrigger(stmt) => {
            validate_param(&stmt.trigger_name, query_ctx)?;
        }
        Statement::CreateView(stmt) => {
            validate_param(&stmt.name, query_ctx)?;
        }
        Statement::AlterTable(stmt) => {
            validate_param(stmt.table_name(), query_ctx)?;
        }
        #[cfg(feature = "enterprise")]
        Statement::AlterTrigger(_) => {}
        // set/show variable now only alter/show variable in session
        Statement::SetVariables(_) | Statement::ShowVariables(_) => {}
        // show charset and show collation won't be checked
        Statement::ShowCharset(_) | Statement::ShowCollation(_) => {}

        Statement::Comment(comment) => match &comment.object {
            CommentObject::Table(table) => validate_param(table, query_ctx)?,
            CommentObject::Column { table, .. } => validate_param(table, query_ctx)?,
            CommentObject::Flow(flow) => validate_flow(flow, query_ctx)?,
        },

        Statement::Insert(insert) => {
            let name = insert.table_name().context(ParseSqlSnafu)?;
            validate_param(name, query_ctx)?;
        }
        Statement::CreateTable(stmt) => {
            validate_param(&stmt.name, query_ctx)?;
        }
        Statement::CreateTableLike(stmt) => {
            validate_param(&stmt.table_name, query_ctx)?;
            validate_param(&stmt.source_name, query_ctx)?;
        }
        Statement::DropTable(drop_stmt) => {
            for table_name in drop_stmt.table_names() {
                validate_param(table_name, query_ctx)?;
            }
        }
        Statement::DropView(stmt) => {
            validate_param(&stmt.view_name, query_ctx)?;
        }
        Statement::ShowTables(stmt) => {
            validate_db_permission!(stmt, query_ctx);
        }
        Statement::ShowTableStatus(stmt) => {
            validate_db_permission!(stmt, query_ctx);
        }
        Statement::ShowColumns(stmt) => {
            validate_db_permission!(stmt, query_ctx);
        }
        Statement::ShowIndex(stmt) => {
            validate_db_permission!(stmt, query_ctx);
        }
        Statement::ShowRegion(stmt) => {
            validate_db_permission!(stmt, query_ctx);
        }
        Statement::ShowViews(stmt) => {
            validate_db_permission!(stmt, query_ctx);
        }
        Statement::ShowFlows(stmt) => {
            validate_db_permission!(stmt, query_ctx);
        }
        #[cfg(feature = "enterprise")]
        Statement::ShowTriggers(_stmt) => {
            // The trigger is organized based on the catalog dimension, so there
            // is no need to check the permission of the database(schema).
        }
        Statement::ShowStatus(_stmt) => {}
        Statement::ShowSearchPath(_stmt) => {}
        Statement::DescribeTable(stmt) => {
            validate_param(stmt.name(), query_ctx)?;
        }
        Statement::Copy(sql::statements::copy::Copy::CopyTable(stmt)) => match stmt {
            CopyTable::To(copy_table_to) => validate_param(&copy_table_to.table_name, query_ctx)?,
            CopyTable::From(copy_table_from) => {
                validate_param(&copy_table_from.table_name, query_ctx)?
            }
        },
        Statement::Copy(sql::statements::copy::Copy::CopyDatabase(copy_database)) => {
            match copy_database {
                CopyDatabase::To(stmt) => validate_database(&stmt.database_name, query_ctx)?,
                CopyDatabase::From(stmt) => validate_database(&stmt.database_name, query_ctx)?,
            }
        }
        Statement::TruncateTable(stmt) => {
            validate_param(stmt.table_name(), query_ctx)?;
        }
        // cursor operations are always allowed once it's created
        Statement::FetchCursor(_) | Statement::CloseCursor(_) => {}
        // User can only kill process in their own catalog.
        Statement::Kill(_) => {}
        // SHOW PROCESSLIST
        Statement::ShowProcesslist(_) => {}
    }
    Ok(())
}

fn validate_param(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
    let (catalog, schema, _) = table_idents_to_full_name(name, query_ctx)
        .map_err(BoxedError::new)
        .context(ExternalSnafu)?;

    validate_catalog_and_schema(&catalog, &schema, query_ctx)
        .map_err(BoxedError::new)
        .context(SqlExecInterceptedSnafu)
}

fn validate_flow(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
    let catalog = match &name.0[..] {
        [_flow] => query_ctx.current_catalog().to_string(),
        [catalog, _flow] => catalog.to_string_unquoted(),
        _ => {
            return InvalidSqlSnafu {
                err_msg: format!(
                    "expect flow name to be <catalog>.<flow_name> or <flow_name>, actual: {name}",
                ),
            }
            .fail();
        }
    };

    let schema = query_ctx.current_schema();

    validate_catalog_and_schema(&catalog, &schema, query_ctx)
        .map_err(BoxedError::new)
        .context(SqlExecInterceptedSnafu)
}

fn validate_database(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
    let (catalog, schema) = match &name.0[..] {
        [schema] => (
            query_ctx.current_catalog().to_string(),
            schema.to_string_unquoted(),
        ),
        [catalog, schema] => (catalog.to_string_unquoted(), schema.to_string_unquoted()),
        _ => InvalidSqlSnafu {
            err_msg: format!(
                "expect database name to be <catalog>.<schema> or <schema>, actual: {name}",
            ),
        }
        .fail()?,
    };

    validate_catalog_and_schema(&catalog, &schema, query_ctx)
        .map_err(BoxedError::new)
        .context(SqlExecInterceptedSnafu)
}

fn is_readonly_plan(plan: &LogicalPlan) -> bool {
    !matches!(plan, LogicalPlan::Dml(_) | LogicalPlan::Ddl(_))
}

fn should_track_statement_process(stmt: &Statement) -> bool {
    stmt.is_readonly()
        || matches!(stmt, Statement::Insert(insert) if insert.has_non_values_query_source())
}

fn should_track_plan_process(stmt: Option<&Statement>, plan: &LogicalPlan) -> bool {
    is_readonly_plan(plan)
        || matches!(stmt, Some(Statement::Insert(insert)) if insert.has_non_values_query_source())
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::future::Future;
    use std::pin::Pin;
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::sync::{Arc, Barrier};
    use std::task::{Context, Poll};
    use std::thread;
    use std::time::{Duration, Instant};

    use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse};
    use catalog::process_manager::ProcessManager;
    use common_base::Plugins;
    use common_error::ext::{BoxedError, PlainError};
    use common_error::status_code::StatusCode;
    use common_meta::cache::LayeredCacheRegistryBuilder;
    use common_meta::kv_backend::memory::MemoryKvBackend;
    use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor};
    use common_meta::rpc::ddl::{SubmitDdlTaskRequest, SubmitDdlTaskResponse};
    use common_meta::rpc::procedure::{
        MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse,
    };
    use common_query::Output;
    use common_recordbatch::{
        OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream,
    };
    use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
    use datafusion_expr::dml::InsertOp;
    use datafusion_expr::{LogicalPlanBuilder, LogicalTableSource};
    use datatypes::prelude::ConcreteDataType;
    use datatypes::schema::{ColumnSchema, Schema as GtSchema, SchemaRef as GtSchemaRef};
    use query::query_engine::options::QueryOptions;
    use session::context::{Channel, ConnInfo, QueryContext, QueryContextBuilder};
    use snafu::{Location, Snafu};
    use sql::dialect::GreptimeDbDialect;
    use store_api::data_source::DataSource;
    use store_api::storage::ScanRequest;
    use strfmt::Format;
    use table::metadata::{FilterPushDownType, TableInfo, TableInfoBuilder, TableMetaBuilder};
    use table::test_util::EmptyTable;
    use table::{Table, TableRef};
    use tokio::sync::{mpsc, oneshot};

    use super::*;
    use crate::frontend::FrontendOptions;
    use crate::instance::builder::FrontendBuilder;

    fn parse_test_sql(sql: &str) -> Vec<Statement> {
        parse_stmt(sql, &GreptimeDbDialect {}).unwrap()
    }

    #[test]
    fn test_validate_analyze_stream_statement_strictness() {
        for sql in [
            "select 1",
            "explain analyze select 1",
            "explain analyze verbose format text select 1",
            "explain analyze verbose format graphviz select 1",
        ] {
            let mut stmts = parse_test_sql(sql);
            assert!(
                validate_analyze_stream_statement(&mut stmts[0]).is_err(),
                "{sql}"
            );
        }

        for sql in [
            "explain analyze verbose select 1",
            "explain analyze verbose format json select 1",
        ] {
            let mut stmts = parse_test_sql(sql);
            assert!(
                validate_analyze_stream_statement(&mut stmts[0]).is_ok(),
                "{sql}"
            );
            let Statement::Explain(explain) = &stmts[0] else {
                unreachable!();
            };
            assert!(explain.format.is_none());
        }

        assert_eq!(
            parse_test_sql("explain analyze verbose select 1; select 2").len(),
            2
        );
    }

    #[derive(Debug, Snafu)]
    enum TestError {
        #[snafu(display("Failed to build test cache registry"))]
        BuildCacheRegistry {
            source: cache::error::Error,
            #[snafu(implicit)]
            location: Location,
        },

        #[snafu(display("Failed to build test table meta for table: {table_name}"))]
        BuildTableMeta {
            table_name: String,
            source: table::metadata::TableMetaBuilderError,
            #[snafu(implicit)]
            location: Location,
        },

        #[snafu(display("Failed to build test table info for table: {table_name}"))]
        BuildTableInfo {
            table_name: String,
            source: table::metadata::TableInfoBuilderError,
            #[snafu(implicit)]
            location: Location,
        },

        #[snafu(display("Failed to register test table: {table_name}"))]
        RegisterTable {
            table_name: String,
            source: catalog::error::Error,
            #[snafu(implicit)]
            location: Location,
        },

        #[snafu(display("Failed to build test frontend instance"))]
        BuildFrontend {
            source: crate::error::Error,
            #[snafu(implicit)]
            location: Location,
        },

        #[snafu(display("Expected exactly one output for SQL `{sql}`, got {actual}"))]
        UnexpectedOutputCount {
            sql: String,
            actual: usize,
            #[snafu(implicit)]
            location: Location,
        },

        #[snafu(display("Failed to execute SQL `{sql}`"))]
        ExecuteSql {
            sql: String,
            source: crate::error::Error,
            #[snafu(implicit)]
            location: Location,
        },

        #[snafu(display("Timed out waiting for insert-select start notification"))]
        InsertStartTimeout {
            source: tokio::time::error::Elapsed,
            #[snafu(implicit)]
            location: Location,
        },

        #[snafu(display("Insert-select start notification channel closed"))]
        InsertStartChannelClosed {
            #[snafu(implicit)]
            location: Location,
        },

        #[snafu(display("Failed to release blocking insert-select interceptor"))]
        ReleaseBlockedInsert {
            #[snafu(implicit)]
            location: Location,
        },

        #[snafu(display("Timed out waiting for insert-select source to be polled"))]
        SourcePollTimeout {
            source: tokio::time::error::Elapsed,
            #[snafu(implicit)]
            location: Location,
        },

        #[snafu(display("Insert-select source poll notification channel closed"))]
        SourcePollChannelClosed {
            source: oneshot::error::RecvError,
            #[snafu(implicit)]
            location: Location,
        },

        #[snafu(display("Timed out waiting for insert task to finish"))]
        InsertTaskTimeout {
            source: tokio::time::error::Elapsed,
            #[snafu(implicit)]
            location: Location,
        },

        #[snafu(display("Insert task panicked"))]
        InsertTaskPanic {
            source: tokio::task::JoinError,
            #[snafu(implicit)]
            location: Location,
        },

        #[snafu(display("Expected insert-select to be cancelled"))]
        InsertSelectNotCancelled {
            #[snafu(implicit)]
            location: Location,
        },
    }

    type TestResult<T> = std::result::Result<T, TestError>;

    fn parse_one_sql(sql: &str) -> Statement {
        parse_stmt(sql, &GreptimeDbDialect {}).unwrap().remove(0)
    }

    fn test_query_ctx(process_id: u32) -> QueryContextRef {
        Arc::new(
            QueryContextBuilder::default()
                .channel(Channel::Mysql)
                .conn_info(ConnInfo::new(None, Channel::Mysql))
                .process_id(process_id)
                .build(),
        )
    }

    struct BlockingInsertSelectInterceptor {
        started_tx: mpsc::UnboundedSender<()>,
        finish_rx: std::sync::Mutex<Option<oneshot::Receiver<()>>>,
    }

    impl BlockingInsertSelectInterceptor {
        fn new(started_tx: mpsc::UnboundedSender<()>, finish_rx: oneshot::Receiver<()>) -> Self {
            Self {
                started_tx,
                finish_rx: std::sync::Mutex::new(Some(finish_rx)),
            }
        }
    }

    impl SqlQueryInterceptor for BlockingInsertSelectInterceptor {
        type Error = Error;

        fn pre_execute(
            &self,
            statement: Option<&Statement>,
            _plan: Option<&LogicalPlan>,
            _query_ctx: QueryContextRef,
        ) -> Result<()> {
            let Some(Statement::Insert(insert)) = statement else {
                return Ok(());
            };
            if !insert.has_non_values_query_source() {
                return Ok(());
            }

            let finish_rx = self.finish_rx.lock().unwrap().take().unwrap();
            let _ = self.started_tx.send(());
            tokio::task::block_in_place(|| {
                tokio::runtime::Handle::current()
                    .block_on(finish_rx)
                    .unwrap();
            });
            Ok(())
        }
    }

    struct PendingRecordBatchStream {
        schema: GtSchemaRef,
        polled_tx: Option<oneshot::Sender<()>>,
        _finish_tx: oneshot::Sender<()>,
        finish_rx: Pin<Box<oneshot::Receiver<()>>>,
    }

    impl RecordBatchStream for PendingRecordBatchStream {
        fn schema(&self) -> GtSchemaRef {
            self.schema.clone()
        }

        fn output_ordering(&self) -> Option<&[OrderOption]> {
            None
        }

        fn metrics(&self) -> Option<common_recordbatch::adapter::RecordBatchMetrics> {
            None
        }
    }

    impl Stream for PendingRecordBatchStream {
        type Item = common_recordbatch::error::Result<RecordBatch>;

        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
            if let Some(polled_tx) = self.polled_tx.take() {
                let _ = polled_tx.send(());
            }

            match self.finish_rx.as_mut().poll(cx) {
                Poll::Ready(_) => Poll::Ready(None),
                Poll::Pending => Poll::Pending,
            }
        }
    }

    impl Unpin for PendingRecordBatchStream {}

    struct PendingDataSource {
        schema: GtSchemaRef,
        polled_tx: std::sync::Mutex<Option<oneshot::Sender<()>>>,
    }

    impl DataSource for PendingDataSource {
        fn get_stream(
            &self,
            _request: ScanRequest,
        ) -> std::result::Result<SendableRecordBatchStream, BoxedError> {
            let (finish_tx, finish_rx) = oneshot::channel();
            let mut polled_tx = self.polled_tx.lock().map_err(|_| {
                BoxedError::new(PlainError::new(
                    "pending data source lock poisoned".to_string(),
                    StatusCode::Unexpected,
                ))
            })?;
            Ok(Box::pin(PendingRecordBatchStream {
                schema: self.schema.clone(),
                polled_tx: polled_tx.take(),
                _finish_tx: finish_tx,
                finish_rx: Box::pin(finish_rx),
            }))
        }
    }

    struct NoopProcedureExecutor;

    #[async_trait::async_trait]
    impl ProcedureExecutor for NoopProcedureExecutor {
        async fn submit_ddl_task(
            &self,
            _ctx: &ExecutorContext,
            _request: SubmitDdlTaskRequest,
        ) -> common_meta::error::Result<SubmitDdlTaskResponse> {
            common_meta::error::UnsupportedSnafu {
                operation: "submit_ddl_task",
            }
            .fail()
        }

        async fn migrate_region(
            &self,
            _ctx: &ExecutorContext,
            _request: MigrateRegionRequest,
        ) -> common_meta::error::Result<MigrateRegionResponse> {
            common_meta::error::UnsupportedSnafu {
                operation: "migrate_region",
            }
            .fail()
        }

        async fn reconcile(
            &self,
            _ctx: &ExecutorContext,
            _request: ReconcileRequest,
        ) -> common_meta::error::Result<ReconcileResponse> {
            common_meta::error::UnsupportedSnafu {
                operation: "reconcile",
            }
            .fail()
        }

        async fn query_procedure_state(
            &self,
            _ctx: &ExecutorContext,
            _pid: &str,
        ) -> common_meta::error::Result<ProcedureStateResponse> {
            common_meta::error::UnsupportedSnafu {
                operation: "query_procedure_state",
            }
            .fail()
        }

        async fn list_procedures(
            &self,
            _ctx: &ExecutorContext,
        ) -> common_meta::error::Result<ProcedureDetailResponse> {
            common_meta::error::UnsupportedSnafu {
                operation: "list_procedures",
            }
            .fail()
        }
    }

    fn test_cache_registry(
        kv_backend: common_meta::kv_backend::KvBackendRef,
    ) -> TestResult<common_meta::cache::LayeredCacheRegistryRef> {
        Ok(Arc::new(
            cache::with_default_composite_cache_registry(
                LayeredCacheRegistryBuilder::default()
                    .add_cache_registry(cache::build_fundamental_cache_registry(kv_backend)),
            )
            .context(BuildCacheRegistrySnafu)?
            .build(),
        ))
    }

    fn test_table_info(table_id: u32, table_name: &str) -> TestResult<TableInfo> {
        let schema = Arc::new(GtSchema::new(vec![
            ColumnSchema::new("id", ConcreteDataType::int32_datatype(), false),
            ColumnSchema::new(
                "ts",
                ConcreteDataType::timestamp_millisecond_datatype(),
                false,
            )
            .with_time_index(true),
        ]));
        let table_meta = TableMetaBuilder::empty()
            .schema(schema)
            .primary_key_indices(vec![0])
            .value_indices(vec![1])
            .next_column_id(1024)
            .build()
            .with_context(|_| BuildTableMetaSnafu {
                table_name: table_name.to_string(),
            })?;

        TableInfoBuilder::new(table_name, table_meta)
            .table_id(table_id)
            .build()
            .with_context(|_| BuildTableInfoSnafu {
                table_name: table_name.to_string(),
            })
    }

    fn test_table(table_id: u32, table_name: &str) -> TestResult<table::TableRef> {
        let table_info = test_table_info(table_id, table_name)?;
        Ok(EmptyTable::from_table_info(&table_info))
    }

    fn pending_table(
        table_id: u32,
        table_name: &str,
        polled_tx: oneshot::Sender<()>,
    ) -> TestResult<table::TableRef> {
        let table_info = test_table_info(table_id, table_name)?;
        let data_source = Arc::new(PendingDataSource {
            schema: table_info.meta.schema.clone(),
            polled_tx: std::sync::Mutex::new(Some(polled_tx)),
        });

        Ok(Arc::new(Table::new(
            Arc::new(table_info),
            FilterPushDownType::Unsupported,
            data_source,
        )))
    }

    async fn test_instance_with_tables(
        source_table: TableRef,
        target_table: TableRef,
    ) -> TestResult<Instance> {
        test_instance_with_plugins(source_table, target_table, Plugins::new()).await
    }

    async fn test_instance_with_insert_select_interceptor(
        interceptor: SqlQueryInterceptorRef<Error>,
    ) -> TestResult<Instance> {
        let plugins = Plugins::new();
        plugins.insert::<SqlQueryInterceptorRef<Error>>(interceptor);

        test_instance_with_plugins(
            test_table(1024, "source")?,
            test_table(1025, "target")?,
            plugins,
        )
        .await
    }

    async fn test_instance_with_plugins(
        source_table: TableRef,
        target_table: TableRef,
        plugins: Plugins,
    ) -> TestResult<Instance> {
        let kv_backend = Arc::new(MemoryKvBackend::new());
        let process_manager = Arc::new(ProcessManager::new("test-frontend".to_string(), None));
        let catalog_manager = catalog::memory::MemoryCatalogManager::new_with_table(source_table);
        let target_table_name = "target";
        catalog_manager
            .register_table_sync(catalog::RegisterTableRequest {
                catalog: "greptime".to_string(),
                schema: "public".to_string(),
                table_name: target_table_name.to_string(),
                table_id: 1025,
                table: target_table,
            })
            .with_context(|_| RegisterTableSnafu {
                table_name: target_table_name.to_string(),
            })?;
        catalog_manager.register_process_list_table(process_manager.clone());

        let cache_registry = test_cache_registry(kv_backend.clone())?;

        FrontendBuilder::new(
            FrontendOptions::default(),
            kv_backend,
            cache_registry,
            catalog_manager,
            Arc::new(client::client_manager::NodeClients::default()),
            Arc::new(NoopProcedureExecutor),
            process_manager,
        )
        .with_plugin(plugins)
        .try_build()
        .await
        .context(BuildFrontendSnafu)
    }

    async fn execute_one_sql(
        instance: &Instance,
        sql: &str,
        query_ctx: QueryContextRef,
    ) -> TestResult<Output> {
        let mut results = instance.do_query_inner(sql, query_ctx).await;
        ensure!(
            results.len() == 1,
            UnexpectedOutputCountSnafu {
                sql: sql.to_string(),
                actual: results.len(),
            }
        );
        results.remove(0).with_context(|_| ExecuteSqlSnafu {
            sql: sql.to_string(),
        })
    }

    #[test]
    fn test_fast_legacy_check_deadlock_prevention() {
        // Create a DashMap to simulate the cache
        let cache = DashMap::new();

        // Pre-populate cache with some entries
        cache.insert("metric1".to_string(), true); // legacy mode
        cache.insert("metric2".to_string(), false); // prom mode
        cache.insert("metric3".to_string(), true); // legacy mode

        // Test case 1: Normal operation with cache hits
        let metric1 = "metric1".to_string();
        let metric4 = "metric4".to_string();
        let names1 = vec![&metric1, &metric4];
        let result = fast_legacy_check(&cache, &names1);
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), Some(true)); // should return legacy mode

        // Verify that metric4 was added to cache
        assert!(cache.contains_key("metric4"));
        assert!(*cache.get("metric4").unwrap().value());

        // Test case 2: No cache hits
        let metric5 = "metric5".to_string();
        let metric6 = "metric6".to_string();
        let names2 = vec![&metric5, &metric6];
        let result = fast_legacy_check(&cache, &names2);
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), None); // should return None as no cache hits

        // Test case 3: Incompatible modes should return error
        let cache_incompatible = DashMap::new();
        cache_incompatible.insert("metric1".to_string(), true); // legacy
        cache_incompatible.insert("metric2".to_string(), false); // prom
        let metric1_test = "metric1".to_string();
        let metric2_test = "metric2".to_string();
        let names3 = vec![&metric1_test, &metric2_test];
        let result = fast_legacy_check(&cache_incompatible, &names3);
        assert!(result.is_err()); // should error due to incompatible modes

        // Test case 4: Intensive concurrent access to test deadlock prevention
        // This test specifically targets the scenario where multiple threads
        // access the same cache entries simultaneously
        let cache_concurrent = Arc::new(DashMap::new());
        cache_concurrent.insert("shared_metric".to_string(), true);

        let num_threads = 8;
        let operations_per_thread = 100;
        let barrier = Arc::new(Barrier::new(num_threads));
        let success_flag = Arc::new(AtomicBool::new(true));

        let handles: Vec<_> = (0..num_threads)
            .map(|thread_id| {
                let cache_clone = Arc::clone(&cache_concurrent);
                let barrier_clone = Arc::clone(&barrier);
                let success_flag_clone = Arc::clone(&success_flag);

                thread::spawn(move || {
                    // Wait for all threads to be ready
                    barrier_clone.wait();

                    let start_time = Instant::now();
                    for i in 0..operations_per_thread {
                        // Each operation references existing cache entry and adds new ones
                        let shared_metric = "shared_metric".to_string();
                        let new_metric = format!("thread_{}_metric_{}", thread_id, i);
                        let names = vec![&shared_metric, &new_metric];

                        match fast_legacy_check(&cache_clone, &names) {
                            Ok(_) => {}
                            Err(_) => {
                                success_flag_clone.store(false, Ordering::Relaxed);
                                return;
                            }
                        }

                        // If the test takes too long, it likely means deadlock
                        if start_time.elapsed() > Duration::from_secs(10) {
                            success_flag_clone.store(false, Ordering::Relaxed);
                            return;
                        }
                    }
                })
            })
            .collect();

        // Join all threads with timeout
        let start_time = Instant::now();
        for (i, handle) in handles.into_iter().enumerate() {
            let join_result = handle.join();

            // Check if we're taking too long (potential deadlock)
            if start_time.elapsed() > Duration::from_secs(30) {
                panic!("Test timed out - possible deadlock detected!");
            }

            if join_result.is_err() {
                panic!("Thread {} panicked during execution", i);
            }
        }

        // Verify all operations completed successfully
        assert!(
            success_flag.load(Ordering::Relaxed),
            "Some operations failed"
        );

        // Verify that many new entries were added (proving operations completed)
        let final_count = cache_concurrent.len();
        assert!(
            final_count > 1 + num_threads * operations_per_thread / 2,
            "Expected more cache entries, got {}",
            final_count
        );
    }

    #[test]
    fn test_should_track_statement_process() {
        assert!(should_track_statement_process(&parse_one_sql(
            "SELECT * FROM demo"
        )));
        assert!(should_track_statement_process(&parse_one_sql(
            "INSERT INTO demo SELECT * FROM source"
        )));
        assert!(!should_track_statement_process(&parse_one_sql(
            "INSERT INTO demo VALUES (1)"
        )));
        assert!(!should_track_statement_process(&parse_one_sql(
            "INSERT INTO demo VALUES (now())"
        )));
    }

    #[test]
    fn test_should_track_plan_process() {
        let select_stmt = parse_one_sql("SELECT * FROM demo");
        let insert_select_stmt = parse_one_sql("INSERT INTO demo SELECT * FROM source");
        let insert_values_stmt = parse_one_sql("INSERT INTO demo VALUES (now())");

        let empty_plan = LogicalPlanBuilder::empty(false).build().unwrap();
        assert!(should_track_plan_process(Some(&select_stmt), &empty_plan));
        assert!(should_track_plan_process(
            Some(&insert_select_stmt),
            &insert_dml_plan()
        ));
        assert!(!should_track_plan_process(
            Some(&insert_values_stmt),
            &insert_dml_plan()
        ));
        assert!(!should_track_plan_process(None, &insert_dml_plan()));
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_insert_select_is_visible_in_show_processlist() -> TestResult<()> {
        let insert_sql = "INSERT INTO target SELECT * FROM source";
        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
        let (finish_tx, finish_rx) = oneshot::channel();
        let interceptor = Arc::new(BlockingInsertSelectInterceptor::new(started_tx, finish_rx));
        let instance = Arc::new(test_instance_with_insert_select_interceptor(interceptor).await?);

        let insert_task = tokio::spawn({
            let instance = instance.clone();
            async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
        });

        tokio::time::timeout(Duration::from_secs(5), started_rx.recv())
            .await
            .context(InsertStartTimeoutSnafu)?
            .context(InsertStartChannelClosedSnafu)?;

        let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
        let process_list = output.data.pretty_print().await;
        assert!(
            process_list.contains(insert_sql),
            "process list did not contain running insert:\n{process_list}"
        );

        finish_tx
            .send(())
            .map_err(|_| ReleaseBlockedInsertSnafu.build())?;
        insert_task.await.context(InsertTaskPanicSnafu)??;

        Ok(())
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_kill_query_cancels_insert_select() -> TestResult<()> {
        assert_kill_cancels_insert_select("KILL QUERY 4242").await
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_kill_process_id_cancels_insert_select() -> TestResult<()> {
        assert_kill_cancels_insert_select("KILL 'test-frontend/4242'").await
    }

    async fn assert_kill_cancels_insert_select(kill_sql: &str) -> TestResult<()> {
        let insert_sql = "INSERT INTO target SELECT * FROM source";
        let (source_polled_tx, source_polled_rx) = oneshot::channel();
        let instance = Arc::new(
            test_instance_with_tables(
                pending_table(1024, "source", source_polled_tx)?,
                test_table(1025, "target")?,
            )
            .await?,
        );

        let insert_task = tokio::spawn({
            let instance = instance.clone();
            async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
        });

        tokio::time::timeout(Duration::from_secs(5), source_polled_rx)
            .await
            .context(SourcePollTimeoutSnafu)?
            .context(SourcePollChannelClosedSnafu)?;

        let output = execute_one_sql(&instance, kill_sql, test_query_ctx(43)).await?;
        assert!(matches!(output.data, OutputData::AffectedRows(1)));

        let insert_result = tokio::time::timeout(Duration::from_secs(5), insert_task)
            .await
            .context(InsertTaskTimeoutSnafu)?
            .context(InsertTaskPanicSnafu)?;
        let err = match insert_result {
            Ok(_) => return InsertSelectNotCancelledSnafu.fail(),
            Err(TestError::ExecuteSql { source, .. }) => source,
            Err(err) => return Err(err),
        };
        assert_eq!(StatusCode::Cancelled, err.status_code());

        let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
        let process_list = output.data.pretty_print().await;
        assert!(
            !process_list.contains(insert_sql),
            "process list still contains killed insert:\n{process_list}"
        );

        Ok(())
    }

    fn insert_dml_plan() -> LogicalPlan {
        let schema = SchemaRef::new(Schema::new(vec![Field::new(
            "value",
            DataType::Int64,
            true,
        )]));
        let target = Arc::new(LogicalTableSource::new(schema));
        let input = LogicalPlanBuilder::empty(false).build().unwrap();

        LogicalPlanBuilder::insert_into(input, "demo", target, InsertOp::Append)
            .unwrap()
            .build()
            .unwrap()
    }

    #[test]
    fn test_exec_validation() {
        let query_ctx = QueryContext::arc();
        let plugins: Plugins = Plugins::new();
        plugins.insert(QueryOptions {
            disallow_cross_catalog_query: true,
        });

        let sql = r#"
        SELECT * FROM demo;
        EXPLAIN SELECT * FROM demo;
        CREATE DATABASE test_database;
        SHOW DATABASES;
        "#;
        let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
        assert_eq!(stmts.len(), 4);
        for stmt in stmts {
            let re = check_permission(plugins.clone(), &stmt, &query_ctx);
            re.unwrap();
        }

        let sql = r#"
        SHOW CREATE TABLE demo;
        ALTER TABLE demo ADD COLUMN new_col INT;
        "#;
        let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
        assert_eq!(stmts.len(), 2);
        for stmt in stmts {
            let re = check_permission(plugins.clone(), &stmt, &query_ctx);
            re.unwrap();
        }

        fn replace_test(template_sql: &str, plugins: Plugins, query_ctx: &QueryContextRef) {
            // test right
            let right = vec![("", ""), ("", "public."), ("greptime.", "public.")];
            for (catalog, schema) in right {
                let sql = do_fmt(template_sql, catalog, schema);
                do_test(&sql, plugins.clone(), query_ctx, true);
            }

            let wrong = vec![
                ("wrongcatalog.", "public."),
                ("wrongcatalog.", "wrongschema."),
            ];
            for (catalog, schema) in wrong {
                let sql = do_fmt(template_sql, catalog, schema);
                do_test(&sql, plugins.clone(), query_ctx, false);
            }
        }

        fn do_fmt(template: &str, catalog: &str, schema: &str) -> String {
            let vars = HashMap::from([
                ("catalog".to_string(), catalog),
                ("schema".to_string(), schema),
            ]);
            template.format(&vars).unwrap()
        }

        fn do_test(sql: &str, plugins: Plugins, query_ctx: &QueryContextRef, is_ok: bool) {
            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
            let re = check_permission(plugins, stmt, query_ctx);
            if is_ok {
                re.unwrap();
            } else {
                assert!(re.is_err());
            }
        }

        // test insert
        let sql = "INSERT INTO {catalog}{schema}monitor(host) VALUES ('host1');";
        replace_test(sql, plugins.clone(), &query_ctx);

        // test create table
        let sql = r#"CREATE TABLE {catalog}{schema}demo(
                            host STRING,
                            ts TIMESTAMP,
                            TIME INDEX (ts),
                            PRIMARY KEY(host)
                        ) engine=mito;"#;
        replace_test(sql, plugins.clone(), &query_ctx);

        // test drop table
        let sql = "DROP TABLE {catalog}{schema}demo;";
        replace_test(sql, plugins.clone(), &query_ctx);

        // test show tables
        let sql = "SHOW TABLES FROM public";
        let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
        check_permission(plugins.clone(), &stmt[0], &query_ctx).unwrap();

        let sql = "SHOW TABLES FROM private";
        let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
        let re = check_permission(plugins.clone(), &stmt[0], &query_ctx);
        assert!(re.is_ok());

        // test describe table
        let sql = "DESC TABLE {catalog}{schema}demo;";
        replace_test(sql, plugins.clone(), &query_ctx);

        let comment_flow_cases = [
            ("COMMENT ON FLOW my_flow IS 'comment';", true),
            ("COMMENT ON FLOW greptime.my_flow IS 'comment';", true),
            ("COMMENT ON FLOW wrongcatalog.my_flow IS 'comment';", false),
        ];
        for (sql, is_ok) in comment_flow_cases {
            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
            let result = check_permission(plugins.clone(), stmt, &query_ctx);
            assert_eq!(result.is_ok(), is_ok);
        }

        let show_flow_cases = [
            ("SHOW CREATE FLOW my_flow;", true),
            ("SHOW CREATE FLOW greptime.my_flow;", true),
            ("SHOW CREATE FLOW wrongcatalog.my_flow;", false),
        ];
        for (sql, is_ok) in show_flow_cases {
            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
            let result = check_permission(plugins.clone(), stmt, &query_ctx);
            assert_eq!(result.is_ok(), is_ok);
        }
    }
}