Skip to main content

frontend/
instance.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15pub mod builder;
16mod dashboard;
17mod grpc;
18mod influxdb;
19mod jaeger;
20mod log_handler;
21mod logs;
22mod opentsdb;
23mod otlp;
24pub mod prom_store;
25mod promql;
26mod region_query;
27pub mod standalone;
28
29use std::pin::Pin;
30use std::sync::atomic::AtomicBool;
31use std::sync::{Arc, atomic};
32use std::time::{Duration, SystemTime};
33
34use async_stream::stream;
35use async_trait::async_trait;
36use auth::{PermissionChecker, PermissionCheckerRef, PermissionReq};
37use catalog::CatalogManagerRef;
38use catalog::process_manager::{
39    ProcessManagerRef, QueryStatement as CatalogQueryStatement, SlowQueryTimer,
40};
41use client::OutputData;
42use common_base::Plugins;
43use common_base::cancellation::CancellableFuture;
44use common_error::ext::{BoxedError, ErrorExt};
45use common_event_recorder::EventRecorderRef;
46use common_meta::cache_invalidator::CacheInvalidatorRef;
47use common_meta::key::TableMetadataManagerRef;
48use common_meta::key::table_name::TableNameKey;
49use common_meta::node_manager::NodeManagerRef;
50use common_meta::procedure_executor::ProcedureExecutorRef;
51use common_query::Output;
52use common_recordbatch::RecordBatchStreamWrapper;
53use common_recordbatch::error::StreamTimeoutSnafu;
54use common_telemetry::logging::SlowQueryOptions;
55use common_telemetry::{debug, error, tracing};
56use dashmap::DashMap;
57use datafusion_expr::LogicalPlan;
58use futures::{Stream, StreamExt};
59use lazy_static::lazy_static;
60use operator::delete::DeleterRef;
61use operator::insert::InserterRef;
62use operator::statement::{StatementExecutor, StatementExecutorRef};
63use partition::manager::PartitionRuleManagerRef;
64use pipeline::pipeline_operator::PipelineOperator;
65use prometheus::HistogramTimer;
66use promql_parser::label::Matcher;
67use query::QueryEngineRef;
68use query::metrics::OnDone;
69use query::parser::{PromQuery, QueryLanguageParser, QueryStatement};
70use query::query_engine::DescribeResult;
71use query::query_engine::options::{QueryOptions, validate_catalog_and_schema};
72use servers::error::{
73    self as server_error, AuthSnafu, CommonMetaSnafu, ExecuteQuerySnafu,
74    OtlpMetricModeIncompatibleSnafu, ParsePromQLSnafu, UnexpectedResultSnafu,
75};
76use servers::interceptor::{
77    PromQueryInterceptor, PromQueryInterceptorRef, SqlQueryInterceptor, SqlQueryInterceptorRef,
78};
79use servers::otlp::metrics::legacy_normalize_otlp_name;
80use servers::prometheus_handler::PrometheusHandler;
81use servers::query_handler::sql::SqlQueryHandler;
82use session::context::{Channel, QueryContextRef};
83use session::table_name::table_idents_to_full_name;
84use snafu::prelude::*;
85use sql::ast::ObjectNamePartExt;
86use sql::dialect::Dialect;
87use sql::parser::{ParseOptions, ParserContext};
88use sql::statements::comment::CommentObject;
89use sql::statements::copy::{CopyDatabase, CopyTable};
90use sql::statements::statement::Statement;
91use sql::statements::tql::Tql;
92use sqlparser::ast::{AnalyzeFormat, ObjectName};
93pub use standalone::StandaloneDatanodeManager;
94use table::requests::{OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM};
95use tracing::Span;
96
97use crate::error::{
98    self, CollectRecordbatchSnafu, Error, ExecLogicalPlanSnafu, ExecutePromqlSnafu, ExternalSnafu,
99    InvalidSqlSnafu, ParseSqlSnafu, PermissionSnafu, PlanStatementSnafu, Result,
100    SqlExecInterceptedSnafu, StatementTimeoutSnafu, TableOperationSnafu,
101};
102use crate::service_config::InfluxdbMergeMode;
103use crate::stream_wrapper::CancellableStreamWrapper;
104
105lazy_static! {
106    static ref OTLP_LEGACY_DEFAULT_VALUE: String = "legacy".to_string();
107}
108
109/// The frontend instance contains necessary components, and implements many
110/// traits, like [`servers::query_handler::grpc::GrpcQueryHandler`],
111/// [`servers::query_handler::sql::SqlQueryHandler`], etc.
112#[derive(Clone)]
113pub struct Instance {
114    frontend_peer_addr: String,
115    catalog_manager: CatalogManagerRef,
116    pipeline_operator: Arc<PipelineOperator>,
117    statement_executor: Arc<StatementExecutor>,
118    query_engine: QueryEngineRef,
119    plugins: Plugins,
120    inserter: InserterRef,
121    deleter: DeleterRef,
122    table_metadata_manager: TableMetadataManagerRef,
123    event_recorder: Option<EventRecorderRef>,
124    process_manager: ProcessManagerRef,
125    slow_query_options: SlowQueryOptions,
126    influxdb_default_merge_mode: InfluxdbMergeMode,
127    trace_ingest_chunk_size: usize,
128    suspend: Arc<AtomicBool>,
129
130    // cache for otlp metrics
131    // first layer key: db-string
132    // key: direct input metric name
133    // value: if runs in legacy mode
134    otlp_metrics_table_legacy_cache: DashMap<String, DashMap<String, bool>>,
135}
136
137impl Instance {
138    pub fn frontend_peer_addr(&self) -> &str {
139        &self.frontend_peer_addr
140    }
141
142    pub fn catalog_manager(&self) -> &CatalogManagerRef {
143        &self.catalog_manager
144    }
145
146    pub fn query_engine(&self) -> &QueryEngineRef {
147        &self.query_engine
148    }
149
150    pub fn plugins(&self) -> &Plugins {
151        &self.plugins
152    }
153
154    pub fn statement_executor(&self) -> &StatementExecutorRef {
155        &self.statement_executor
156    }
157
158    pub fn table_metadata_manager(&self) -> &TableMetadataManagerRef {
159        &self.table_metadata_manager
160    }
161
162    pub fn inserter(&self) -> &InserterRef {
163        &self.inserter
164    }
165
166    pub fn process_manager(&self) -> &ProcessManagerRef {
167        &self.process_manager
168    }
169
170    pub fn node_manager(&self) -> &NodeManagerRef {
171        self.inserter.node_manager()
172    }
173
174    pub fn partition_manager(&self) -> &PartitionRuleManagerRef {
175        self.inserter.partition_manager()
176    }
177
178    pub fn cache_invalidator(&self) -> &CacheInvalidatorRef {
179        self.statement_executor.cache_invalidator()
180    }
181
182    pub fn procedure_executor(&self) -> &ProcedureExecutorRef {
183        self.statement_executor.procedure_executor()
184    }
185
186    pub fn suspend_state(&self) -> Arc<AtomicBool> {
187        self.suspend.clone()
188    }
189
190    pub(crate) fn is_suspended(&self) -> bool {
191        self.suspend.load(atomic::Ordering::Relaxed)
192    }
193}
194
195fn parse_stmt(sql: &str, dialect: &(dyn Dialect + Send + Sync)) -> Result<Vec<Statement>> {
196    ParserContext::create_with_dialect(sql, dialect, ParseOptions::default()).context(ParseSqlSnafu)
197}
198
199fn validate_analyze_stream_statement(stmt: &mut Statement) -> Result<()> {
200    let Statement::Explain(explain) = stmt else {
201        return InvalidSqlSnafu {
202            err_msg: "only EXPLAIN ANALYZE VERBOSE statement is supported",
203        }
204        .fail();
205    };
206    ensure!(
207        explain.analyze && explain.verbose,
208        InvalidSqlSnafu {
209            err_msg: "statement must be EXPLAIN ANALYZE VERBOSE"
210        }
211    );
212    match explain.format {
213        None | Some(AnalyzeFormat::JSON) => {
214            // Keep explicit FORMAT JSON accepted, but pass JSON through
215            // QueryContext.explain_format instead of the statement to avoid the
216            // planner's current `EXPLAIN VERBOSE with FORMAT` limitation.
217            explain.format = None;
218            Ok(())
219        }
220        Some(_) => InvalidSqlSnafu {
221            err_msg: "only FORMAT JSON is supported for analyze stream",
222        }
223        .fail(),
224    }
225}
226
227impl Instance {
228    fn statement_slow_query_timer(
229        &self,
230        stmt: &Statement,
231        schema_name: String,
232    ) -> Option<SlowQueryTimer> {
233        if !stmt.is_readonly() || !self.slow_query_options.enable {
234            return None;
235        }
236
237        self.event_recorder.clone().map(|event_recorder| {
238            SlowQueryTimer::new(
239                CatalogQueryStatement::Sql(stmt.clone()),
240                schema_name,
241                self.slow_query_options.threshold,
242                self.slow_query_options.sample_ratio,
243                self.slow_query_options.record_type,
244                event_recorder,
245            )
246        })
247    }
248
249    async fn query_statement(&self, stmt: Statement, query_ctx: QueryContextRef) -> Result<Output> {
250        check_permission(self.plugins.clone(), &stmt, &query_ctx)?;
251
252        let query_interceptor = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
253        let query_interceptor = query_interceptor.as_ref();
254
255        if should_track_statement_process(&stmt) {
256            let catalog_name = query_ctx.current_catalog().to_string();
257            let schema_name = query_ctx.current_schema();
258            let slow_query_timer = self.statement_slow_query_timer(&stmt, schema_name.clone());
259
260            let ticket = self.process_manager.register_query(
261                catalog_name,
262                vec![schema_name],
263                stmt.to_string(),
264                query_ctx.conn_info().to_string(),
265                Some(query_ctx.process_id()),
266                slow_query_timer,
267            );
268
269            let query_fut = self.exec_statement_with_timeout(stmt, query_ctx, query_interceptor);
270
271            CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
272                .await
273                .map_err(|_| error::CancelledSnafu.build())?
274                .map(|output| {
275                    let Output { meta, data } = output;
276
277                    let data = match data {
278                        OutputData::Stream(stream) => OutputData::Stream(Box::pin(
279                            CancellableStreamWrapper::new(stream, ticket),
280                        )),
281                        other => other,
282                    };
283                    Output { data, meta }
284                })
285        } else {
286            self.exec_statement_with_timeout(stmt, query_ctx, query_interceptor)
287                .await
288        }
289    }
290
291    async fn exec_statement_with_timeout(
292        &self,
293        stmt: Statement,
294        query_ctx: QueryContextRef,
295        query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
296    ) -> Result<Output> {
297        let timeout = derive_timeout(&stmt, &query_ctx);
298        match timeout {
299            Some(timeout) => {
300                let start = tokio::time::Instant::now();
301                let output = tokio::time::timeout(
302                    timeout,
303                    self.exec_statement(stmt, query_ctx, query_interceptor),
304                )
305                .await
306                .map_err(|_| StatementTimeoutSnafu.build())??;
307                let output = map_query_output(output)?;
308                // compute remaining timeout
309                let remaining_timeout = timeout.checked_sub(start.elapsed()).unwrap_or_default();
310                attach_timeout(output, remaining_timeout)
311            }
312            None => self
313                .exec_statement(stmt, query_ctx, query_interceptor)
314                .await
315                .and_then(map_query_output),
316        }
317    }
318
319    async fn exec_statement(
320        &self,
321        stmt: Statement,
322        query_ctx: QueryContextRef,
323        query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
324    ) -> Result<Output> {
325        match stmt {
326            Statement::Query(_) | Statement::Explain(_) | Statement::Delete(_) => {
327                // TODO: remove this when format is supported in datafusion
328                if let Statement::Explain(explain) = &stmt
329                    && let Some(format) = explain.format()
330                {
331                    query_ctx.set_explain_format(format.to_string());
332                }
333
334                self.plan_and_exec_sql(stmt, &query_ctx, query_interceptor)
335                    .await
336            }
337            Statement::Tql(tql) => {
338                self.plan_and_exec_tql(&query_ctx, query_interceptor, tql)
339                    .await
340            }
341            _ => {
342                query_interceptor.pre_execute(Some(&stmt), None, query_ctx.clone())?;
343                self.statement_executor
344                    .execute_sql(stmt, query_ctx)
345                    .await
346                    .context(TableOperationSnafu)
347            }
348        }
349    }
350
351    async fn plan_and_exec_sql(
352        &self,
353        stmt: Statement,
354        query_ctx: &QueryContextRef,
355        query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
356    ) -> Result<Output> {
357        let stmt = QueryStatement::Sql(stmt);
358        let plan = self
359            .statement_executor
360            .plan(&stmt, query_ctx.clone())
361            .await?;
362        let QueryStatement::Sql(stmt) = stmt else {
363            unreachable!()
364        };
365        query_interceptor.pre_execute(Some(&stmt), Some(&plan), query_ctx.clone())?;
366
367        self.statement_executor
368            .exec_plan(plan, query_ctx.clone())
369            .await
370            .context(TableOperationSnafu)
371    }
372
373    async fn plan_and_exec_tql(
374        &self,
375        query_ctx: &QueryContextRef,
376        query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
377        tql: Tql,
378    ) -> Result<Output> {
379        let plan = self
380            .statement_executor
381            .plan_tql(tql.clone(), query_ctx)
382            .await?;
383        query_interceptor.pre_execute(
384            Some(&Statement::Tql(tql)),
385            Some(&plan),
386            query_ctx.clone(),
387        )?;
388        self.statement_executor
389            .exec_plan(plan, query_ctx.clone())
390            .await
391            .context(TableOperationSnafu)
392    }
393
394    async fn check_otlp_legacy(
395        &self,
396        names: &[&String],
397        ctx: QueryContextRef,
398    ) -> server_error::Result<bool> {
399        let db_string = ctx.get_db_string();
400        // fast cache check
401        let cache = self
402            .otlp_metrics_table_legacy_cache
403            .entry(db_string.clone())
404            .or_default();
405        if let Some(flag) = fast_legacy_check(&cache, names)? {
406            return Ok(flag);
407        }
408        // release cache reference to avoid lock contention
409        drop(cache);
410
411        let catalog = ctx.current_catalog();
412        let schema = ctx.current_schema();
413
414        // query legacy table names
415        let normalized_names = names
416            .iter()
417            .map(|n| legacy_normalize_otlp_name(n))
418            .collect::<Vec<_>>();
419        let table_names = normalized_names
420            .iter()
421            .map(|n| TableNameKey::new(catalog, &schema, n))
422            .collect::<Vec<_>>();
423        let table_values = self
424            .table_metadata_manager()
425            .table_name_manager()
426            .batch_get(table_names)
427            .await
428            .context(CommonMetaSnafu)?;
429        let table_ids = table_values
430            .into_iter()
431            .filter_map(|v| v.map(|vi| vi.table_id()))
432            .collect::<Vec<_>>();
433
434        // means no existing table is found, use new mode
435        if table_ids.is_empty() {
436            let cache = self
437                .otlp_metrics_table_legacy_cache
438                .entry(db_string)
439                .or_default();
440            names.iter().for_each(|name| {
441                cache.insert((*name).clone(), false);
442            });
443            return Ok(false);
444        }
445
446        // has existing table, check table options
447        let table_infos = self
448            .table_metadata_manager()
449            .table_info_manager()
450            .batch_get(&table_ids)
451            .await
452            .context(CommonMetaSnafu)?;
453        let options = table_infos
454            .values()
455            .map(|info| {
456                info.table_info
457                    .meta
458                    .options
459                    .extra_options
460                    .get(OTLP_METRIC_COMPAT_KEY)
461                    .unwrap_or(&OTLP_LEGACY_DEFAULT_VALUE)
462            })
463            .collect::<Vec<_>>();
464        let cache = self
465            .otlp_metrics_table_legacy_cache
466            .entry(db_string)
467            .or_default();
468        if !options.is_empty() {
469            // check value consistency
470            let has_prom = options.iter().any(|opt| *opt == OTLP_METRIC_COMPAT_PROM);
471            let has_legacy = options
472                .iter()
473                .any(|opt| *opt == OTLP_LEGACY_DEFAULT_VALUE.as_str());
474            ensure!(!(has_prom && has_legacy), OtlpMetricModeIncompatibleSnafu);
475            let flag = has_legacy;
476            names.iter().for_each(|name| {
477                cache.insert((*name).clone(), flag);
478            });
479            Ok(flag)
480        } else {
481            // no table info, use new mode
482            names.iter().for_each(|name| {
483                cache.insert((*name).clone(), false);
484            });
485            Ok(false)
486        }
487    }
488}
489
490fn fast_legacy_check(
491    cache: &DashMap<String, bool>,
492    names: &[&String],
493) -> server_error::Result<Option<bool>> {
494    let hit_cache = names
495        .iter()
496        .filter_map(|name| cache.get(*name))
497        .collect::<Vec<_>>();
498    if !hit_cache.is_empty() {
499        let hit_legacy = hit_cache.iter().any(|en| *en.value());
500        let hit_prom = hit_cache.iter().any(|en| !*en.value());
501
502        // hit but have true and false, means both legacy and new mode are used
503        // we cannot handle this case, so return error
504        // add doc links in err msg later
505        ensure!(!(hit_legacy && hit_prom), OtlpMetricModeIncompatibleSnafu);
506
507        let flag = hit_legacy;
508        // drop hit_cache to release references before inserting to avoid deadlock
509        drop(hit_cache);
510
511        // set cache for all names
512        names.iter().for_each(|name| {
513            if !cache.contains_key(*name) {
514                cache.insert((*name).clone(), flag);
515            }
516        });
517        Ok(Some(flag))
518    } else {
519        Ok(None)
520    }
521}
522
523/// If the relevant variables are set, the timeout is enforced for all PostgreSQL statements.
524/// For MySQL, it applies only to read-only statements.
525fn derive_timeout(stmt: &Statement, query_ctx: &QueryContextRef) -> Option<Duration> {
526    let query_timeout = query_ctx.query_timeout()?;
527    if query_timeout.is_zero() {
528        return None;
529    }
530    match query_ctx.channel() {
531        Channel::Mysql if stmt.is_readonly() => Some(query_timeout),
532        Channel::Postgres => Some(query_timeout),
533        _ => None,
534    }
535}
536
537/// Derives timeout for plan execution.
538fn derive_timeout_for_plan(plan: &LogicalPlan, query_ctx: &QueryContextRef) -> Option<Duration> {
539    let query_timeout = query_ctx.query_timeout()?;
540    if query_timeout.is_zero() {
541        return None;
542    }
543    match query_ctx.channel() {
544        Channel::Mysql if is_readonly_plan(plan) => Some(query_timeout),
545        Channel::Postgres => Some(query_timeout),
546        _ => None,
547    }
548}
549
550fn attach_timeout(output: Output, mut timeout: Duration) -> Result<Output> {
551    if timeout.is_zero() {
552        return StatementTimeoutSnafu.fail();
553    }
554
555    let output = match output.data {
556        OutputData::AffectedRows(_) | OutputData::RecordBatches(_) => output,
557        OutputData::Stream(mut stream) => {
558            let schema = stream.schema();
559            let s = Box::pin(stream! {
560                let mut start = tokio::time::Instant::now();
561                while let Some(item) = tokio::time::timeout(timeout, stream.next()).await.map_err(|_| StreamTimeoutSnafu.build())? {
562                    yield item;
563
564                    let now = tokio::time::Instant::now();
565                    timeout = timeout.checked_sub(now - start).unwrap_or(Duration::ZERO);
566                    start = now;
567                    // tokio::time::timeout may not return an error immediately when timeout is 0.
568                    if timeout.is_zero() {
569                        StreamTimeoutSnafu.fail()?;
570                    }
571                }
572            }) as Pin<Box<dyn Stream<Item = _> + Send>>;
573            let stream = RecordBatchStreamWrapper {
574                schema,
575                stream: s,
576                output_ordering: None,
577                metrics: Default::default(),
578                span: Span::current(),
579            };
580            Output::new(OutputData::Stream(Box::pin(stream)), output.meta)
581        }
582    };
583
584    Ok(output)
585}
586
587impl Instance {
588    #[tracing::instrument(skip_all, name = "SqlQueryHandler::do_analyze_stream_query")]
589    async fn do_analyze_stream_query_inner(
590        &self,
591        query: &str,
592        query_ctx: QueryContextRef,
593    ) -> Result<Output> {
594        ensure!(!self.is_suspended(), error::SuspendedSnafu);
595
596        let query_interceptor_opt = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
597        let query_interceptor = query_interceptor_opt.as_ref();
598        let query = query_interceptor.pre_parsing(query, query_ctx.clone())?;
599        let mut stmts = parse_stmt(query.as_ref(), query_ctx.sql_dialect())
600            .and_then(|stmts| query_interceptor.post_parsing(stmts, query_ctx.clone()))?;
601
602        ensure!(
603            stmts.len() == 1,
604            InvalidSqlSnafu {
605                err_msg: "only single EXPLAIN ANALYZE VERBOSE statement is supported"
606            }
607        );
608        let mut stmt = stmts.remove(0);
609        validate_analyze_stream_statement(&mut stmt)?;
610        query_ctx.set_explain_format(AnalyzeFormat::JSON.to_string());
611
612        let checker_ref = self.plugins.get::<PermissionCheckerRef>();
613        checker_ref
614            .as_ref()
615            .check_permission_with_context(
616                query_ctx.current_user(),
617                PermissionReq::SqlStatement(&stmt),
618                Some(&query_ctx.current_schema()),
619            )
620            .context(PermissionSnafu)?;
621        check_permission(self.plugins.clone(), &stmt, &query_ctx)?;
622        let catalog_name = query_ctx.current_catalog().to_string();
623        let schema_name = query_ctx.current_schema();
624        let slow_query_timer = self.statement_slow_query_timer(&stmt, schema_name.clone());
625        let ticket = self.process_manager.register_query(
626            catalog_name,
627            vec![schema_name],
628            stmt.to_string(),
629            query_ctx.conn_info().to_string(),
630            Some(query_ctx.process_id()),
631            slow_query_timer,
632        );
633        let query_fut =
634            self.exec_statement_with_timeout(stmt, query_ctx.clone(), query_interceptor);
635        let output = CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
636            .await
637            .map_err(|_| error::CancelledSnafu.build())??;
638        let Output { meta, data } = output;
639        let data = match data {
640            OutputData::Stream(stream) => OutputData::Stream(Box::pin(
641                CancellableStreamWrapper::new_cancel_on_drop(stream, ticket),
642            )),
643            other => other,
644        };
645        query_interceptor.post_execute(Output { data, meta }, query_ctx)
646    }
647
648    #[tracing::instrument(skip_all, name = "SqlQueryHandler::do_query")]
649    async fn do_query_inner(&self, query: &str, query_ctx: QueryContextRef) -> Vec<Result<Output>> {
650        if self.is_suspended() {
651            return vec![error::SuspendedSnafu {}.fail()];
652        }
653
654        let query_interceptor_opt = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
655        let query_interceptor = query_interceptor_opt.as_ref();
656        let query = match query_interceptor.pre_parsing(query, query_ctx.clone()) {
657            Ok(q) => q,
658            Err(e) => return vec![Err(e)],
659        };
660
661        let checker_ref = self.plugins.get::<PermissionCheckerRef>();
662        let checker = checker_ref.as_ref();
663
664        match parse_stmt(query.as_ref(), query_ctx.sql_dialect())
665            .and_then(|stmts| query_interceptor.post_parsing(stmts, query_ctx.clone()))
666        {
667            Ok(stmts) => {
668                if stmts.is_empty() {
669                    return vec![
670                        InvalidSqlSnafu {
671                            err_msg: "empty statements",
672                        }
673                        .fail(),
674                    ];
675                }
676
677                let mut results = Vec::with_capacity(stmts.len());
678                for stmt in stmts {
679                    if let Err(e) = checker
680                        .check_permission_with_context(
681                            query_ctx.current_user(),
682                            PermissionReq::SqlStatement(&stmt),
683                            Some(&query_ctx.current_schema()),
684                        )
685                        .context(PermissionSnafu)
686                    {
687                        results.push(Err(e));
688                        break;
689                    }
690
691                    match self.query_statement(stmt.clone(), query_ctx.clone()).await {
692                        Ok(output) => {
693                            let output_result =
694                                query_interceptor.post_execute(output, query_ctx.clone());
695                            results.push(output_result);
696                        }
697                        Err(e) => {
698                            if e.status_code().should_log_error() {
699                                error!(e; "Failed to execute query: {stmt}");
700                            } else {
701                                debug!("Failed to execute query: {stmt}, {e}");
702                            }
703                            results.push(Err(e));
704                            break;
705                        }
706                    }
707                }
708                results
709            }
710            Err(e) => {
711                vec![Err(e)]
712            }
713        }
714    }
715
716    async fn exec_plan(&self, plan: LogicalPlan, query_ctx: QueryContextRef) -> Result<Output> {
717        self.query_engine
718            .execute(plan, query_ctx)
719            .await
720            .context(ExecLogicalPlanSnafu)
721    }
722
723    async fn exec_plan_with_timeout(
724        &self,
725        plan: LogicalPlan,
726        query_ctx: QueryContextRef,
727    ) -> Result<Output> {
728        let timeout = derive_timeout_for_plan(&plan, &query_ctx);
729        match timeout {
730            Some(timeout) => {
731                let start = tokio::time::Instant::now();
732                let output = tokio::time::timeout(timeout, self.exec_plan(plan, query_ctx))
733                    .await
734                    .map_err(|_| StatementTimeoutSnafu.build())??;
735                let output = map_query_output(output)?;
736                let remaining_timeout = timeout.checked_sub(start.elapsed()).unwrap_or_default();
737                attach_timeout(output, remaining_timeout)
738            }
739            None => self
740                .exec_plan(plan, query_ctx)
741                .await
742                .and_then(map_query_output),
743        }
744    }
745
746    async fn do_exec_plan_inner(
747        &self,
748        plan: LogicalPlan,
749        stmt: Option<Statement>,
750        query_ctx: QueryContextRef,
751    ) -> Result<Output> {
752        ensure!(!self.is_suspended(), error::SuspendedSnafu);
753
754        let query_interceptor_opt = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
755        let query_interceptor = query_interceptor_opt.as_ref();
756
757        query_interceptor.pre_execute(stmt.as_ref(), Some(&plan), query_ctx.clone())?;
758
759        let query = stmt
760            .as_ref()
761            .map(|s| s.to_string())
762            .unwrap_or_else(|| plan.display_indent().to_string());
763
764        let plan_is_readonly = is_readonly_plan(&plan);
765        let result = if should_track_plan_process(stmt.as_ref(), &plan) {
766            let catalog_name = query_ctx.current_catalog().to_string();
767            let schema_name = query_ctx.current_schema();
768            let slow_query_timer = if plan_is_readonly {
769                self.slow_query_options
770                    .enable
771                    .then(|| self.event_recorder.clone())
772                    .flatten()
773                    .map(|event_recorder| {
774                        SlowQueryTimer::new(
775                            CatalogQueryStatement::Plan(query.clone()),
776                            schema_name.clone(),
777                            self.slow_query_options.threshold,
778                            self.slow_query_options.sample_ratio,
779                            self.slow_query_options.record_type,
780                            event_recorder,
781                        )
782                    })
783            } else {
784                None
785            };
786
787            let ticket = self.process_manager.register_query(
788                catalog_name,
789                vec![schema_name],
790                query,
791                query_ctx.conn_info().to_string(),
792                Some(query_ctx.process_id()),
793                slow_query_timer,
794            );
795
796            let query_fut = self.exec_plan_with_timeout(plan, query_ctx.clone());
797
798            CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
799                .await
800                .map_err(|_| error::CancelledSnafu.build())?
801                .map(|output| {
802                    let Output { meta, data } = output;
803
804                    let data = match data {
805                        OutputData::Stream(stream) => OutputData::Stream(Box::pin(
806                            CancellableStreamWrapper::new(stream, ticket),
807                        )),
808                        other => other,
809                    };
810                    Output { data, meta }
811                })
812        } else {
813            self.exec_plan_with_timeout(plan, query_ctx.clone()).await
814        };
815
816        result.and_then(|output| query_interceptor.post_execute(output, query_ctx))
817    }
818
819    #[tracing::instrument(skip_all, name = "SqlQueryHandler::do_promql_query")]
820    async fn do_promql_query_inner(
821        &self,
822        query: &PromQuery,
823        query_ctx: QueryContextRef,
824    ) -> Vec<Result<Output>> {
825        if self.is_suspended() {
826            return vec![error::SuspendedSnafu {}.fail()];
827        }
828
829        // check will be done in prometheus handler's do_query
830        let result = PrometheusHandler::do_query(self, query, query_ctx)
831            .await
832            .with_context(|_| ExecutePromqlSnafu {
833                query: format!("{query:?}"),
834            });
835        vec![result]
836    }
837
838    async fn do_describe_inner(
839        &self,
840        stmt: Statement,
841        query_ctx: QueryContextRef,
842    ) -> Result<Option<DescribeResult>> {
843        ensure!(!self.is_suspended(), error::SuspendedSnafu);
844
845        // EXPLAIN / EXPLAIN ANALYZE wrap an inner statement; describe them when the
846        // wrapped statement is something we already plan (so that bind parameters
847        // in the inner query get their types inferred). See #8029.
848        let is_inner_plannable = |s: &Statement| {
849            matches!(
850                s,
851                Statement::Insert(_) | Statement::Query(_) | Statement::Delete(_)
852            )
853        };
854        let plannable = is_inner_plannable(&stmt)
855            || matches!(&stmt, Statement::Explain(explain) if is_inner_plannable(explain.statement.as_ref()));
856
857        if plannable {
858            self.plugins
859                .get::<PermissionCheckerRef>()
860                .as_ref()
861                .check_permission_with_context(
862                    query_ctx.current_user(),
863                    PermissionReq::SqlStatement(&stmt),
864                    Some(&query_ctx.current_schema()),
865                )
866                .context(PermissionSnafu)?;
867
868            let plan = self
869                .query_engine
870                .planner()
871                .plan(&QueryStatement::Sql(stmt), query_ctx.clone())
872                .await
873                .context(PlanStatementSnafu)?;
874            self.query_engine
875                .describe(plan, query_ctx)
876                .await
877                .map(Some)
878                .context(error::DescribeStatementSnafu)
879        } else {
880            Ok(None)
881        }
882    }
883
884    async fn is_valid_schema_inner(&self, catalog: &str, schema: &str) -> Result<bool> {
885        self.catalog_manager
886            .schema_exists(catalog, schema, None)
887            .await
888            .context(error::CatalogSnafu)
889    }
890}
891
892#[async_trait]
893impl SqlQueryHandler for Instance {
894    async fn do_query(
895        &self,
896        query: &str,
897        query_ctx: QueryContextRef,
898    ) -> Vec<server_error::Result<Output>> {
899        self.do_query_inner(query, query_ctx)
900            .await
901            .into_iter()
902            .map(|result| result.map_err(BoxedError::new).context(ExecuteQuerySnafu))
903            .collect()
904    }
905
906    async fn do_analyze_stream_query(
907        &self,
908        query: &str,
909        query_ctx: QueryContextRef,
910    ) -> server_error::Result<Output> {
911        self.do_analyze_stream_query_inner(query, query_ctx)
912            .await
913            .map_err(BoxedError::new)
914            .context(ExecuteQuerySnafu)
915    }
916
917    async fn do_exec_plan(
918        &self,
919        plan: LogicalPlan,
920        stmt: Option<Statement>,
921        query_ctx: QueryContextRef,
922    ) -> server_error::Result<Output> {
923        self.do_exec_plan_inner(plan, stmt, query_ctx)
924            .await
925            .map_err(BoxedError::new)
926            .context(server_error::ExecutePlanSnafu)
927    }
928
929    async fn do_promql_query(
930        &self,
931        query: &PromQuery,
932        query_ctx: QueryContextRef,
933    ) -> Vec<server_error::Result<Output>> {
934        self.do_promql_query_inner(query, query_ctx)
935            .await
936            .into_iter()
937            .map(|result| result.map_err(BoxedError::new).context(ExecuteQuerySnafu))
938            .collect()
939    }
940
941    async fn do_describe(
942        &self,
943        stmt: Statement,
944        query_ctx: QueryContextRef,
945    ) -> server_error::Result<Option<DescribeResult>> {
946        self.do_describe_inner(stmt, query_ctx)
947            .await
948            .map_err(BoxedError::new)
949            .context(server_error::DescribeStatementSnafu)
950    }
951
952    async fn is_valid_schema(&self, catalog: &str, schema: &str) -> server_error::Result<bool> {
953        self.is_valid_schema_inner(catalog, schema)
954            .await
955            .map_err(BoxedError::new)
956            .context(server_error::CheckDatabaseValiditySnafu)
957    }
958}
959
960/// Expands scan-time dictionaries only when a query result leaves the frontend.
961pub(crate) fn map_query_output(output: Output) -> Result<Output> {
962    output
963        .map_dictionary_to_values()
964        .context(CollectRecordbatchSnafu)
965}
966
967/// Attaches a timer to the output and observes it once the output is exhausted.
968pub fn attach_timer(output: Output, timer: HistogramTimer) -> Output {
969    match output.data {
970        OutputData::AffectedRows(_) | OutputData::RecordBatches(_) => output,
971        OutputData::Stream(stream) => {
972            let stream = OnDone::new(stream, move || {
973                timer.observe_duration();
974            });
975            Output::new(OutputData::Stream(Box::pin(stream)), output.meta)
976        }
977    }
978}
979
980#[async_trait]
981impl PrometheusHandler for Instance {
982    #[tracing::instrument(skip_all)]
983    async fn do_query(
984        &self,
985        query: &PromQuery,
986        query_ctx: QueryContextRef,
987    ) -> server_error::Result<Output> {
988        let interceptor = self
989            .plugins
990            .get::<PromQueryInterceptorRef<server_error::Error>>();
991
992        self.plugins
993            .get::<PermissionCheckerRef>()
994            .as_ref()
995            .check_permission(query_ctx.current_user(), PermissionReq::PromQuery)
996            .context(AuthSnafu)?;
997
998        let stmt = QueryLanguageParser::parse_promql(query, &query_ctx).with_context(|_| {
999            ParsePromQLSnafu {
1000                query: query.clone(),
1001            }
1002        })?;
1003
1004        let plan = self
1005            .statement_executor
1006            .plan(&stmt, query_ctx.clone())
1007            .await
1008            .map_err(BoxedError::new)
1009            .context(ExecuteQuerySnafu)?;
1010
1011        let QueryStatement::Promql(eval_stmt, _) = &stmt else {
1012            unreachable!("query is parsed from promql");
1013        };
1014
1015        interceptor.pre_execute(query, &eval_stmt.expr, Some(&plan), query_ctx.clone())?;
1016
1017        // Take the EvalStmt from the original QueryStatement and use it to create the CatalogQueryStatement.
1018        let query_statement = if let QueryStatement::Promql(eval_stmt, alias) = stmt {
1019            CatalogQueryStatement::Promql(eval_stmt, alias)
1020        } else {
1021            // It should not happen since the query is already parsed successfully.
1022            return UnexpectedResultSnafu {
1023                reason: "The query should always be promql.".to_string(),
1024            }
1025            .fail();
1026        };
1027        let raw_query = query_statement.to_string();
1028
1029        let slow_query_timer = self
1030            .slow_query_options
1031            .enable
1032            .then(|| self.event_recorder.clone())
1033            .flatten()
1034            .map(|event_recorder| {
1035                SlowQueryTimer::new(
1036                    query_statement,
1037                    query_ctx.current_schema(),
1038                    self.slow_query_options.threshold,
1039                    self.slow_query_options.sample_ratio,
1040                    self.slow_query_options.record_type,
1041                    event_recorder,
1042                )
1043            });
1044
1045        let ticket = self.process_manager.register_query(
1046            query_ctx.current_catalog().to_string(),
1047            vec![query_ctx.current_schema()],
1048            raw_query,
1049            query_ctx.conn_info().to_string(),
1050            Some(query_ctx.process_id()),
1051            slow_query_timer,
1052        );
1053
1054        let query_fut = self.statement_executor.exec_plan(plan, query_ctx.clone());
1055
1056        let output = CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
1057            .await
1058            .map_err(|_| servers::error::CancelledSnafu.build())?
1059            .map_err(BoxedError::new)
1060            .context(ExecuteQuerySnafu)?;
1061        let output = map_query_output(output)
1062            .map_err(BoxedError::new)
1063            .context(ExecuteQuerySnafu)?;
1064        let Output { meta, data } = output;
1065        let data = match data {
1066            OutputData::Stream(stream) => {
1067                OutputData::Stream(Box::pin(CancellableStreamWrapper::new(stream, ticket)))
1068            }
1069            other => other,
1070        };
1071        let output = Output { data, meta };
1072        Ok(interceptor.post_execute(output, query_ctx)?)
1073    }
1074
1075    async fn query_metric_names(
1076        &self,
1077        matchers: Vec<Matcher>,
1078        ctx: &QueryContextRef,
1079    ) -> server_error::Result<Vec<String>> {
1080        self.handle_query_metric_names(matchers, ctx)
1081            .await
1082            .map_err(BoxedError::new)
1083            .context(ExecuteQuerySnafu)
1084    }
1085
1086    async fn query_label_values(
1087        &self,
1088        metric: String,
1089        label_name: String,
1090        matchers: Vec<Matcher>,
1091        start: SystemTime,
1092        end: SystemTime,
1093        ctx: &QueryContextRef,
1094    ) -> server_error::Result<Vec<String>> {
1095        self.handle_query_label_values(metric, label_name, matchers, start, end, ctx)
1096            .await
1097            .map_err(BoxedError::new)
1098            .context(ExecuteQuerySnafu)
1099    }
1100
1101    fn catalog_manager(&self) -> CatalogManagerRef {
1102        self.catalog_manager.clone()
1103    }
1104}
1105
1106/// Validate `stmt.database` permission if it's presented.
1107macro_rules! validate_db_permission {
1108    ($stmt: expr, $query_ctx: expr) => {
1109        if let Some(database) = &$stmt.database {
1110            validate_catalog_and_schema($query_ctx.current_catalog(), database, $query_ctx)
1111                .map_err(BoxedError::new)
1112                .context(SqlExecInterceptedSnafu)?;
1113        }
1114    };
1115}
1116
1117pub fn check_permission(
1118    plugins: Plugins,
1119    stmt: &Statement,
1120    query_ctx: &QueryContextRef,
1121) -> Result<()> {
1122    let need_validate = plugins
1123        .get::<QueryOptions>()
1124        .map(|opts| opts.disallow_cross_catalog_query)
1125        .unwrap_or_default();
1126
1127    if !need_validate {
1128        return Ok(());
1129    }
1130
1131    match stmt {
1132        // Will be checked in execution.
1133        // TODO(dennis): add a hook for admin commands.
1134        Statement::Admin(_) => {}
1135        // These are executed by query engine, and will be checked there.
1136        Statement::Query(_)
1137        | Statement::Explain(_)
1138        | Statement::Tql(_)
1139        | Statement::Delete(_)
1140        | Statement::DeclareCursor(_)
1141        | Statement::Copy(sql::statements::copy::Copy::CopyQueryTo(_)) => {}
1142        // database ops won't be checked
1143        Statement::CreateDatabase(_)
1144        | Statement::ShowDatabases(_)
1145        | Statement::DropDatabase(_)
1146        | Statement::AlterDatabase(_)
1147        | Statement::DropFlow(_)
1148        | Statement::Use(_) => {}
1149        #[cfg(feature = "enterprise")]
1150        Statement::DropTrigger(_) => {}
1151        Statement::ShowCreateDatabase(stmt) => {
1152            validate_database(&stmt.database_name, query_ctx)?;
1153        }
1154        Statement::ShowCreateTable(stmt) => {
1155            validate_param(&stmt.table_name, query_ctx)?;
1156        }
1157        Statement::ShowCreateFlow(stmt) => {
1158            validate_flow(&stmt.flow_name, query_ctx)?;
1159        }
1160        #[cfg(feature = "enterprise")]
1161        Statement::ShowCreateTrigger(stmt) => {
1162            validate_param(&stmt.trigger_name, query_ctx)?;
1163        }
1164        Statement::ShowCreateView(stmt) => {
1165            validate_param(&stmt.view_name, query_ctx)?;
1166        }
1167        Statement::CreateExternalTable(stmt) => {
1168            validate_param(&stmt.name, query_ctx)?;
1169        }
1170        Statement::CreateFlow(stmt) => {
1171            // TODO: should also validate source table name here?
1172            validate_param(&stmt.sink_table_name, query_ctx)?;
1173        }
1174        #[cfg(feature = "enterprise")]
1175        Statement::CreateTrigger(stmt) => {
1176            validate_param(&stmt.trigger_name, query_ctx)?;
1177        }
1178        Statement::CreateView(stmt) => {
1179            validate_param(&stmt.name, query_ctx)?;
1180        }
1181        Statement::AlterTable(stmt) => {
1182            validate_param(stmt.table_name(), query_ctx)?;
1183        }
1184        #[cfg(feature = "enterprise")]
1185        Statement::AlterTrigger(_) => {}
1186        // set/show variable now only alter/show variable in session
1187        Statement::SetVariables(_) | Statement::ShowVariables(_) => {}
1188        // show charset and show collation won't be checked
1189        Statement::ShowCharset(_) | Statement::ShowCollation(_) => {}
1190
1191        Statement::Comment(comment) => match &comment.object {
1192            CommentObject::Table(table) => validate_param(table, query_ctx)?,
1193            CommentObject::Column { table, .. } => validate_param(table, query_ctx)?,
1194            CommentObject::Flow(flow) => validate_flow(flow, query_ctx)?,
1195        },
1196
1197        Statement::Insert(insert) => {
1198            let name = insert.table_name().context(ParseSqlSnafu)?;
1199            validate_param(name, query_ctx)?;
1200        }
1201        Statement::CreateTable(stmt) => {
1202            validate_param(&stmt.name, query_ctx)?;
1203        }
1204        Statement::CreateTableLike(stmt) => {
1205            validate_param(&stmt.table_name, query_ctx)?;
1206            validate_param(&stmt.source_name, query_ctx)?;
1207        }
1208        Statement::DropTable(drop_stmt) => {
1209            for table_name in drop_stmt.table_names() {
1210                validate_param(table_name, query_ctx)?;
1211            }
1212        }
1213        Statement::UndropTable(stmt) => {
1214            validate_param(stmt.table_name(), query_ctx)?;
1215        }
1216        Statement::DropView(stmt) => {
1217            validate_param(&stmt.view_name, query_ctx)?;
1218        }
1219        Statement::ShowTables(stmt) => {
1220            validate_db_permission!(stmt, query_ctx);
1221        }
1222        Statement::ShowTableStatus(stmt) => {
1223            validate_db_permission!(stmt, query_ctx);
1224        }
1225        Statement::ShowColumns(stmt) => {
1226            validate_db_permission!(stmt, query_ctx);
1227        }
1228        Statement::ShowIndex(stmt) => {
1229            validate_db_permission!(stmt, query_ctx);
1230        }
1231        Statement::ShowRegion(stmt) => {
1232            validate_db_permission!(stmt, query_ctx);
1233        }
1234        Statement::ShowViews(stmt) => {
1235            validate_db_permission!(stmt, query_ctx);
1236        }
1237        Statement::ShowFlows(stmt) => {
1238            validate_db_permission!(stmt, query_ctx);
1239        }
1240        #[cfg(feature = "enterprise")]
1241        Statement::ShowTriggers(_stmt) => {
1242            // The trigger is organized based on the catalog dimension, so there
1243            // is no need to check the permission of the database(schema).
1244        }
1245        Statement::ShowStatus(_stmt) => {}
1246        Statement::ShowSearchPath(_stmt) => {}
1247        Statement::DescribeTable(stmt) => {
1248            validate_param(stmt.name(), query_ctx)?;
1249        }
1250        Statement::Copy(sql::statements::copy::Copy::CopyTable(stmt)) => match stmt {
1251            CopyTable::To(copy_table_to) => validate_param(&copy_table_to.table_name, query_ctx)?,
1252            CopyTable::From(copy_table_from) => {
1253                validate_param(&copy_table_from.table_name, query_ctx)?
1254            }
1255        },
1256        Statement::Copy(sql::statements::copy::Copy::CopyDatabase(copy_database)) => {
1257            match copy_database {
1258                CopyDatabase::To(stmt) => validate_database(&stmt.database_name, query_ctx)?,
1259                CopyDatabase::From(stmt) => validate_database(&stmt.database_name, query_ctx)?,
1260            }
1261        }
1262        Statement::TruncateTable(stmt) => {
1263            validate_param(stmt.table_name(), query_ctx)?;
1264        }
1265        // cursor operations are always allowed once it's created
1266        Statement::FetchCursor(_) | Statement::CloseCursor(_) => {}
1267        // User can only kill process in their own catalog.
1268        Statement::Kill(_) => {}
1269        // SHOW PROCESSLIST
1270        Statement::ShowProcesslist(_) => {}
1271    }
1272    Ok(())
1273}
1274
1275fn validate_param(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
1276    let (catalog, schema, _) = table_idents_to_full_name(name, query_ctx)
1277        .map_err(BoxedError::new)
1278        .context(ExternalSnafu)?;
1279
1280    validate_catalog_and_schema(&catalog, &schema, query_ctx)
1281        .map_err(BoxedError::new)
1282        .context(SqlExecInterceptedSnafu)
1283}
1284
1285fn validate_flow(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
1286    let catalog = match &name.0[..] {
1287        [_flow] => query_ctx.current_catalog().to_string(),
1288        [catalog, _flow] => catalog.to_string_unquoted(),
1289        _ => {
1290            return InvalidSqlSnafu {
1291                err_msg: format!(
1292                    "expect flow name to be <catalog>.<flow_name> or <flow_name>, actual: {name}",
1293                ),
1294            }
1295            .fail();
1296        }
1297    };
1298
1299    let schema = query_ctx.current_schema();
1300
1301    validate_catalog_and_schema(&catalog, &schema, query_ctx)
1302        .map_err(BoxedError::new)
1303        .context(SqlExecInterceptedSnafu)
1304}
1305
1306fn validate_database(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
1307    let (catalog, schema) = match &name.0[..] {
1308        [schema] => (
1309            query_ctx.current_catalog().to_string(),
1310            schema.to_string_unquoted(),
1311        ),
1312        [catalog, schema] => (catalog.to_string_unquoted(), schema.to_string_unquoted()),
1313        _ => InvalidSqlSnafu {
1314            err_msg: format!(
1315                "expect database name to be <catalog>.<schema> or <schema>, actual: {name}",
1316            ),
1317        }
1318        .fail()?,
1319    };
1320
1321    validate_catalog_and_schema(&catalog, &schema, query_ctx)
1322        .map_err(BoxedError::new)
1323        .context(SqlExecInterceptedSnafu)
1324}
1325
1326fn is_readonly_plan(plan: &LogicalPlan) -> bool {
1327    !matches!(plan, LogicalPlan::Dml(_) | LogicalPlan::Ddl(_))
1328}
1329
1330fn should_track_statement_process(stmt: &Statement) -> bool {
1331    stmt.is_readonly()
1332        || matches!(stmt, Statement::Insert(insert) if insert.has_non_values_query_source())
1333}
1334
1335fn should_track_plan_process(stmt: Option<&Statement>, plan: &LogicalPlan) -> bool {
1336    is_readonly_plan(plan)
1337        || matches!(stmt, Some(Statement::Insert(insert)) if insert.has_non_values_query_source())
1338}
1339
1340#[cfg(test)]
1341mod tests {
1342    use std::collections::HashMap;
1343    use std::future::Future;
1344    use std::pin::Pin;
1345    use std::sync::atomic::{AtomicBool, Ordering};
1346    use std::sync::{Arc, Barrier};
1347    use std::task::{Context, Poll};
1348    use std::thread;
1349    use std::time::{Duration, Instant};
1350
1351    use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse};
1352    use catalog::process_manager::ProcessManager;
1353    use common_base::Plugins;
1354    use common_error::ext::{BoxedError, PlainError};
1355    use common_error::status_code::StatusCode;
1356    use common_meta::cache::LayeredCacheRegistryBuilder;
1357    use common_meta::kv_backend::memory::MemoryKvBackend;
1358    use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor};
1359    use common_meta::rpc::ddl::{SubmitDdlTaskRequest, SubmitDdlTaskResponse};
1360    use common_meta::rpc::procedure::{
1361        MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse,
1362    };
1363    use common_query::Output;
1364    use common_recordbatch::{
1365        OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream,
1366    };
1367    use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
1368    use datafusion_expr::dml::InsertOp;
1369    use datafusion_expr::{LogicalPlanBuilder, LogicalTableSource};
1370    use datatypes::prelude::ConcreteDataType;
1371    use datatypes::schema::{ColumnSchema, Schema as GtSchema, SchemaRef as GtSchemaRef};
1372    use query::query_engine::options::QueryOptions;
1373    use session::context::{Channel, ConnInfo, QueryContext, QueryContextBuilder};
1374    use snafu::{Location, Snafu};
1375    use sql::dialect::GreptimeDbDialect;
1376    use store_api::data_source::DataSource;
1377    use store_api::storage::ScanRequest;
1378    use strfmt::Format;
1379    use table::metadata::{FilterPushDownType, TableInfo, TableInfoBuilder, TableMetaBuilder};
1380    use table::test_util::EmptyTable;
1381    use table::{Table, TableRef};
1382    use tokio::sync::{mpsc, oneshot};
1383
1384    use super::*;
1385    use crate::frontend::FrontendOptions;
1386    use crate::instance::builder::FrontendBuilder;
1387
1388    fn parse_test_sql(sql: &str) -> Vec<Statement> {
1389        parse_stmt(sql, &GreptimeDbDialect {}).unwrap()
1390    }
1391
1392    #[test]
1393    fn test_validate_analyze_stream_statement_strictness() {
1394        for sql in [
1395            "select 1",
1396            "explain analyze select 1",
1397            "explain analyze verbose format text select 1",
1398            "explain analyze verbose format graphviz select 1",
1399        ] {
1400            let mut stmts = parse_test_sql(sql);
1401            assert!(
1402                validate_analyze_stream_statement(&mut stmts[0]).is_err(),
1403                "{sql}"
1404            );
1405        }
1406
1407        for sql in [
1408            "explain analyze verbose select 1",
1409            "explain analyze verbose format json select 1",
1410        ] {
1411            let mut stmts = parse_test_sql(sql);
1412            assert!(
1413                validate_analyze_stream_statement(&mut stmts[0]).is_ok(),
1414                "{sql}"
1415            );
1416            let Statement::Explain(explain) = &stmts[0] else {
1417                unreachable!();
1418            };
1419            assert!(explain.format.is_none());
1420        }
1421
1422        assert_eq!(
1423            parse_test_sql("explain analyze verbose select 1; select 2").len(),
1424            2
1425        );
1426    }
1427
1428    #[derive(Debug, Snafu)]
1429    enum TestError {
1430        #[snafu(display("Failed to build test cache registry"))]
1431        BuildCacheRegistry {
1432            source: cache::error::Error,
1433            #[snafu(implicit)]
1434            location: Location,
1435        },
1436
1437        #[snafu(display("Failed to build test table meta for table: {table_name}"))]
1438        BuildTableMeta {
1439            table_name: String,
1440            source: table::metadata::TableMetaBuilderError,
1441            #[snafu(implicit)]
1442            location: Location,
1443        },
1444
1445        #[snafu(display("Failed to build test table info for table: {table_name}"))]
1446        BuildTableInfo {
1447            table_name: String,
1448            source: table::metadata::TableInfoBuilderError,
1449            #[snafu(implicit)]
1450            location: Location,
1451        },
1452
1453        #[snafu(display("Failed to register test table: {table_name}"))]
1454        RegisterTable {
1455            table_name: String,
1456            source: catalog::error::Error,
1457            #[snafu(implicit)]
1458            location: Location,
1459        },
1460
1461        #[snafu(display("Failed to build test frontend instance"))]
1462        BuildFrontend {
1463            source: crate::error::Error,
1464            #[snafu(implicit)]
1465            location: Location,
1466        },
1467
1468        #[snafu(display("Expected exactly one output for SQL `{sql}`, got {actual}"))]
1469        UnexpectedOutputCount {
1470            sql: String,
1471            actual: usize,
1472            #[snafu(implicit)]
1473            location: Location,
1474        },
1475
1476        #[snafu(display("Failed to execute SQL `{sql}`"))]
1477        ExecuteSql {
1478            sql: String,
1479            source: crate::error::Error,
1480            #[snafu(implicit)]
1481            location: Location,
1482        },
1483
1484        #[snafu(display("Timed out waiting for insert-select start notification"))]
1485        InsertStartTimeout {
1486            source: tokio::time::error::Elapsed,
1487            #[snafu(implicit)]
1488            location: Location,
1489        },
1490
1491        #[snafu(display("Insert-select start notification channel closed"))]
1492        InsertStartChannelClosed {
1493            #[snafu(implicit)]
1494            location: Location,
1495        },
1496
1497        #[snafu(display("Failed to release blocking insert-select interceptor"))]
1498        ReleaseBlockedInsert {
1499            #[snafu(implicit)]
1500            location: Location,
1501        },
1502
1503        #[snafu(display("Timed out waiting for insert-select source to be polled"))]
1504        SourcePollTimeout {
1505            source: tokio::time::error::Elapsed,
1506            #[snafu(implicit)]
1507            location: Location,
1508        },
1509
1510        #[snafu(display("Insert-select source poll notification channel closed"))]
1511        SourcePollChannelClosed {
1512            source: oneshot::error::RecvError,
1513            #[snafu(implicit)]
1514            location: Location,
1515        },
1516
1517        #[snafu(display("Timed out waiting for insert task to finish"))]
1518        InsertTaskTimeout {
1519            source: tokio::time::error::Elapsed,
1520            #[snafu(implicit)]
1521            location: Location,
1522        },
1523
1524        #[snafu(display("Insert task panicked"))]
1525        InsertTaskPanic {
1526            source: tokio::task::JoinError,
1527            #[snafu(implicit)]
1528            location: Location,
1529        },
1530
1531        #[snafu(display("Expected insert-select to be cancelled"))]
1532        InsertSelectNotCancelled {
1533            #[snafu(implicit)]
1534            location: Location,
1535        },
1536    }
1537
1538    type TestResult<T> = std::result::Result<T, TestError>;
1539
1540    fn parse_one_sql(sql: &str) -> Statement {
1541        parse_stmt(sql, &GreptimeDbDialect {}).unwrap().remove(0)
1542    }
1543
1544    fn test_query_ctx(process_id: u32) -> QueryContextRef {
1545        Arc::new(
1546            QueryContextBuilder::default()
1547                .channel(Channel::Mysql)
1548                .conn_info(ConnInfo::new(None, Channel::Mysql))
1549                .process_id(process_id)
1550                .build(),
1551        )
1552    }
1553
1554    struct BlockingInsertSelectInterceptor {
1555        started_tx: mpsc::UnboundedSender<()>,
1556        finish_rx: std::sync::Mutex<Option<oneshot::Receiver<()>>>,
1557    }
1558
1559    impl BlockingInsertSelectInterceptor {
1560        fn new(started_tx: mpsc::UnboundedSender<()>, finish_rx: oneshot::Receiver<()>) -> Self {
1561            Self {
1562                started_tx,
1563                finish_rx: std::sync::Mutex::new(Some(finish_rx)),
1564            }
1565        }
1566    }
1567
1568    impl SqlQueryInterceptor for BlockingInsertSelectInterceptor {
1569        type Error = Error;
1570
1571        fn pre_execute(
1572            &self,
1573            statement: Option<&Statement>,
1574            _plan: Option<&LogicalPlan>,
1575            _query_ctx: QueryContextRef,
1576        ) -> Result<()> {
1577            let Some(Statement::Insert(insert)) = statement else {
1578                return Ok(());
1579            };
1580            if !insert.has_non_values_query_source() {
1581                return Ok(());
1582            }
1583
1584            let finish_rx = self.finish_rx.lock().unwrap().take().unwrap();
1585            let _ = self.started_tx.send(());
1586            tokio::task::block_in_place(|| {
1587                tokio::runtime::Handle::current()
1588                    .block_on(finish_rx)
1589                    .unwrap();
1590            });
1591            Ok(())
1592        }
1593    }
1594
1595    struct PendingRecordBatchStream {
1596        schema: GtSchemaRef,
1597        polled_tx: Option<oneshot::Sender<()>>,
1598        _finish_tx: oneshot::Sender<()>,
1599        finish_rx: Pin<Box<oneshot::Receiver<()>>>,
1600    }
1601
1602    impl RecordBatchStream for PendingRecordBatchStream {
1603        fn schema(&self) -> GtSchemaRef {
1604            self.schema.clone()
1605        }
1606
1607        fn output_ordering(&self) -> Option<&[OrderOption]> {
1608            None
1609        }
1610
1611        fn metrics(&self) -> Option<common_recordbatch::adapter::RecordBatchMetrics> {
1612            None
1613        }
1614    }
1615
1616    impl Stream for PendingRecordBatchStream {
1617        type Item = common_recordbatch::error::Result<RecordBatch>;
1618
1619        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
1620            if let Some(polled_tx) = self.polled_tx.take() {
1621                let _ = polled_tx.send(());
1622            }
1623
1624            match self.finish_rx.as_mut().poll(cx) {
1625                Poll::Ready(_) => Poll::Ready(None),
1626                Poll::Pending => Poll::Pending,
1627            }
1628        }
1629    }
1630
1631    impl Unpin for PendingRecordBatchStream {}
1632
1633    struct PendingDataSource {
1634        schema: GtSchemaRef,
1635        polled_tx: std::sync::Mutex<Option<oneshot::Sender<()>>>,
1636    }
1637
1638    impl DataSource for PendingDataSource {
1639        fn get_stream(
1640            &self,
1641            _request: ScanRequest,
1642        ) -> std::result::Result<SendableRecordBatchStream, BoxedError> {
1643            let (finish_tx, finish_rx) = oneshot::channel();
1644            let mut polled_tx = self.polled_tx.lock().map_err(|_| {
1645                BoxedError::new(PlainError::new(
1646                    "pending data source lock poisoned".to_string(),
1647                    StatusCode::Unexpected,
1648                ))
1649            })?;
1650            Ok(Box::pin(PendingRecordBatchStream {
1651                schema: self.schema.clone(),
1652                polled_tx: polled_tx.take(),
1653                _finish_tx: finish_tx,
1654                finish_rx: Box::pin(finish_rx),
1655            }))
1656        }
1657    }
1658
1659    struct NoopProcedureExecutor;
1660
1661    #[async_trait::async_trait]
1662    impl ProcedureExecutor for NoopProcedureExecutor {
1663        async fn submit_ddl_task(
1664            &self,
1665            _ctx: &ExecutorContext,
1666            _request: SubmitDdlTaskRequest,
1667        ) -> common_meta::error::Result<SubmitDdlTaskResponse> {
1668            common_meta::error::UnsupportedSnafu {
1669                operation: "submit_ddl_task",
1670            }
1671            .fail()
1672        }
1673
1674        async fn migrate_region(
1675            &self,
1676            _ctx: &ExecutorContext,
1677            _request: MigrateRegionRequest,
1678        ) -> common_meta::error::Result<MigrateRegionResponse> {
1679            common_meta::error::UnsupportedSnafu {
1680                operation: "migrate_region",
1681            }
1682            .fail()
1683        }
1684
1685        async fn reconcile(
1686            &self,
1687            _ctx: &ExecutorContext,
1688            _request: ReconcileRequest,
1689        ) -> common_meta::error::Result<ReconcileResponse> {
1690            common_meta::error::UnsupportedSnafu {
1691                operation: "reconcile",
1692            }
1693            .fail()
1694        }
1695
1696        async fn query_procedure_state(
1697            &self,
1698            _ctx: &ExecutorContext,
1699            _pid: &str,
1700        ) -> common_meta::error::Result<ProcedureStateResponse> {
1701            common_meta::error::UnsupportedSnafu {
1702                operation: "query_procedure_state",
1703            }
1704            .fail()
1705        }
1706
1707        async fn list_procedures(
1708            &self,
1709            _ctx: &ExecutorContext,
1710        ) -> common_meta::error::Result<ProcedureDetailResponse> {
1711            common_meta::error::UnsupportedSnafu {
1712                operation: "list_procedures",
1713            }
1714            .fail()
1715        }
1716    }
1717
1718    fn test_cache_registry(
1719        kv_backend: common_meta::kv_backend::KvBackendRef,
1720    ) -> TestResult<common_meta::cache::LayeredCacheRegistryRef> {
1721        Ok(Arc::new(
1722            cache::with_default_composite_cache_registry(
1723                LayeredCacheRegistryBuilder::default()
1724                    .add_cache_registry(cache::build_fundamental_cache_registry(kv_backend)),
1725            )
1726            .context(BuildCacheRegistrySnafu)?
1727            .build(),
1728        ))
1729    }
1730
1731    fn test_table_info(table_id: u32, table_name: &str) -> TestResult<TableInfo> {
1732        let schema = Arc::new(GtSchema::new(vec![
1733            ColumnSchema::new("id", ConcreteDataType::int32_datatype(), false),
1734            ColumnSchema::new(
1735                "ts",
1736                ConcreteDataType::timestamp_millisecond_datatype(),
1737                false,
1738            )
1739            .with_time_index(true),
1740        ]));
1741        let table_meta = TableMetaBuilder::empty()
1742            .schema(schema)
1743            .primary_key_indices(vec![0])
1744            .value_indices(vec![1])
1745            .next_column_id(1024)
1746            .build()
1747            .with_context(|_| BuildTableMetaSnafu {
1748                table_name: table_name.to_string(),
1749            })?;
1750
1751        TableInfoBuilder::new(table_name, table_meta)
1752            .table_id(table_id)
1753            .build()
1754            .with_context(|_| BuildTableInfoSnafu {
1755                table_name: table_name.to_string(),
1756            })
1757    }
1758
1759    fn test_table(table_id: u32, table_name: &str) -> TestResult<table::TableRef> {
1760        let table_info = test_table_info(table_id, table_name)?;
1761        Ok(EmptyTable::from_table_info(&table_info))
1762    }
1763
1764    fn pending_table(
1765        table_id: u32,
1766        table_name: &str,
1767        polled_tx: oneshot::Sender<()>,
1768    ) -> TestResult<table::TableRef> {
1769        let table_info = test_table_info(table_id, table_name)?;
1770        let data_source = Arc::new(PendingDataSource {
1771            schema: table_info.meta.schema.clone(),
1772            polled_tx: std::sync::Mutex::new(Some(polled_tx)),
1773        });
1774
1775        Ok(Arc::new(Table::new(
1776            Arc::new(table_info),
1777            FilterPushDownType::Unsupported,
1778            data_source,
1779        )))
1780    }
1781
1782    async fn test_instance_with_tables(
1783        source_table: TableRef,
1784        target_table: TableRef,
1785    ) -> TestResult<Instance> {
1786        test_instance_with_plugins(source_table, target_table, Plugins::new()).await
1787    }
1788
1789    async fn test_instance_with_insert_select_interceptor(
1790        interceptor: SqlQueryInterceptorRef<Error>,
1791    ) -> TestResult<Instance> {
1792        let plugins = Plugins::new();
1793        plugins.insert::<SqlQueryInterceptorRef<Error>>(interceptor);
1794
1795        test_instance_with_plugins(
1796            test_table(1024, "source")?,
1797            test_table(1025, "target")?,
1798            plugins,
1799        )
1800        .await
1801    }
1802
1803    async fn test_instance_with_plugins(
1804        source_table: TableRef,
1805        target_table: TableRef,
1806        plugins: Plugins,
1807    ) -> TestResult<Instance> {
1808        let kv_backend = Arc::new(MemoryKvBackend::new());
1809        let process_manager = Arc::new(ProcessManager::new("test-frontend".to_string(), None));
1810        let catalog_manager = catalog::memory::MemoryCatalogManager::new_with_table(source_table);
1811        let target_table_name = "target";
1812        catalog_manager
1813            .register_table_sync(catalog::RegisterTableRequest {
1814                catalog: "greptime".to_string(),
1815                schema: "public".to_string(),
1816                table_name: target_table_name.to_string(),
1817                table_id: 1025,
1818                table: target_table,
1819            })
1820            .with_context(|_| RegisterTableSnafu {
1821                table_name: target_table_name.to_string(),
1822            })?;
1823        catalog_manager.register_process_list_table(process_manager.clone());
1824
1825        let cache_registry = test_cache_registry(kv_backend.clone())?;
1826
1827        FrontendBuilder::new(
1828            FrontendOptions::default(),
1829            kv_backend,
1830            cache_registry,
1831            catalog_manager,
1832            Arc::new(client::client_manager::NodeClients::default()),
1833            Arc::new(NoopProcedureExecutor),
1834            process_manager,
1835        )
1836        .with_plugin(plugins)
1837        .try_build()
1838        .await
1839        .context(BuildFrontendSnafu)
1840    }
1841
1842    async fn execute_one_sql(
1843        instance: &Instance,
1844        sql: &str,
1845        query_ctx: QueryContextRef,
1846    ) -> TestResult<Output> {
1847        let mut results = instance.do_query_inner(sql, query_ctx).await;
1848        ensure!(
1849            results.len() == 1,
1850            UnexpectedOutputCountSnafu {
1851                sql: sql.to_string(),
1852                actual: results.len(),
1853            }
1854        );
1855        results.remove(0).with_context(|_| ExecuteSqlSnafu {
1856            sql: sql.to_string(),
1857        })
1858    }
1859
1860    #[test]
1861    fn test_fast_legacy_check_deadlock_prevention() {
1862        // Create a DashMap to simulate the cache
1863        let cache = DashMap::new();
1864
1865        // Pre-populate cache with some entries
1866        cache.insert("metric1".to_string(), true); // legacy mode
1867        cache.insert("metric2".to_string(), false); // prom mode
1868        cache.insert("metric3".to_string(), true); // legacy mode
1869
1870        // Test case 1: Normal operation with cache hits
1871        let metric1 = "metric1".to_string();
1872        let metric4 = "metric4".to_string();
1873        let names1 = vec![&metric1, &metric4];
1874        let result = fast_legacy_check(&cache, &names1);
1875        assert!(result.is_ok());
1876        assert_eq!(result.unwrap(), Some(true)); // should return legacy mode
1877
1878        // Verify that metric4 was added to cache
1879        assert!(cache.contains_key("metric4"));
1880        assert!(*cache.get("metric4").unwrap().value());
1881
1882        // Test case 2: No cache hits
1883        let metric5 = "metric5".to_string();
1884        let metric6 = "metric6".to_string();
1885        let names2 = vec![&metric5, &metric6];
1886        let result = fast_legacy_check(&cache, &names2);
1887        assert!(result.is_ok());
1888        assert_eq!(result.unwrap(), None); // should return None as no cache hits
1889
1890        // Test case 3: Incompatible modes should return error
1891        let cache_incompatible = DashMap::new();
1892        cache_incompatible.insert("metric1".to_string(), true); // legacy
1893        cache_incompatible.insert("metric2".to_string(), false); // prom
1894        let metric1_test = "metric1".to_string();
1895        let metric2_test = "metric2".to_string();
1896        let names3 = vec![&metric1_test, &metric2_test];
1897        let result = fast_legacy_check(&cache_incompatible, &names3);
1898        assert!(result.is_err()); // should error due to incompatible modes
1899
1900        // Test case 4: Intensive concurrent access to test deadlock prevention
1901        // This test specifically targets the scenario where multiple threads
1902        // access the same cache entries simultaneously
1903        let cache_concurrent = Arc::new(DashMap::new());
1904        cache_concurrent.insert("shared_metric".to_string(), true);
1905
1906        let num_threads = 8;
1907        let operations_per_thread = 100;
1908        let barrier = Arc::new(Barrier::new(num_threads));
1909        let success_flag = Arc::new(AtomicBool::new(true));
1910
1911        let handles: Vec<_> = (0..num_threads)
1912            .map(|thread_id| {
1913                let cache_clone = Arc::clone(&cache_concurrent);
1914                let barrier_clone = Arc::clone(&barrier);
1915                let success_flag_clone = Arc::clone(&success_flag);
1916
1917                thread::spawn(move || {
1918                    // Wait for all threads to be ready
1919                    barrier_clone.wait();
1920
1921                    let start_time = Instant::now();
1922                    for i in 0..operations_per_thread {
1923                        // Each operation references existing cache entry and adds new ones
1924                        let shared_metric = "shared_metric".to_string();
1925                        let new_metric = format!("thread_{}_metric_{}", thread_id, i);
1926                        let names = vec![&shared_metric, &new_metric];
1927
1928                        match fast_legacy_check(&cache_clone, &names) {
1929                            Ok(_) => {}
1930                            Err(_) => {
1931                                success_flag_clone.store(false, Ordering::Relaxed);
1932                                return;
1933                            }
1934                        }
1935
1936                        // If the test takes too long, it likely means deadlock
1937                        if start_time.elapsed() > Duration::from_secs(10) {
1938                            success_flag_clone.store(false, Ordering::Relaxed);
1939                            return;
1940                        }
1941                    }
1942                })
1943            })
1944            .collect();
1945
1946        // Join all threads with timeout
1947        let start_time = Instant::now();
1948        for (i, handle) in handles.into_iter().enumerate() {
1949            let join_result = handle.join();
1950
1951            // Check if we're taking too long (potential deadlock)
1952            if start_time.elapsed() > Duration::from_secs(30) {
1953                panic!("Test timed out - possible deadlock detected!");
1954            }
1955
1956            if join_result.is_err() {
1957                panic!("Thread {} panicked during execution", i);
1958            }
1959        }
1960
1961        // Verify all operations completed successfully
1962        assert!(
1963            success_flag.load(Ordering::Relaxed),
1964            "Some operations failed"
1965        );
1966
1967        // Verify that many new entries were added (proving operations completed)
1968        let final_count = cache_concurrent.len();
1969        assert!(
1970            final_count > 1 + num_threads * operations_per_thread / 2,
1971            "Expected more cache entries, got {}",
1972            final_count
1973        );
1974    }
1975
1976    #[test]
1977    fn test_should_track_statement_process() {
1978        assert!(should_track_statement_process(&parse_one_sql(
1979            "SELECT * FROM demo"
1980        )));
1981        assert!(should_track_statement_process(&parse_one_sql(
1982            "INSERT INTO demo SELECT * FROM source"
1983        )));
1984        assert!(!should_track_statement_process(&parse_one_sql(
1985            "INSERT INTO demo VALUES (1)"
1986        )));
1987        assert!(!should_track_statement_process(&parse_one_sql(
1988            "INSERT INTO demo VALUES (now())"
1989        )));
1990    }
1991
1992    #[test]
1993    fn test_should_track_plan_process() {
1994        let select_stmt = parse_one_sql("SELECT * FROM demo");
1995        let insert_select_stmt = parse_one_sql("INSERT INTO demo SELECT * FROM source");
1996        let insert_values_stmt = parse_one_sql("INSERT INTO demo VALUES (now())");
1997
1998        let empty_plan = LogicalPlanBuilder::empty(false).build().unwrap();
1999        assert!(should_track_plan_process(Some(&select_stmt), &empty_plan));
2000        assert!(should_track_plan_process(
2001            Some(&insert_select_stmt),
2002            &insert_dml_plan()
2003        ));
2004        assert!(!should_track_plan_process(
2005            Some(&insert_values_stmt),
2006            &insert_dml_plan()
2007        ));
2008        assert!(!should_track_plan_process(None, &insert_dml_plan()));
2009    }
2010
2011    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2012    async fn test_insert_select_is_visible_in_show_processlist() -> TestResult<()> {
2013        let insert_sql = "INSERT INTO target SELECT * FROM source";
2014        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
2015        let (finish_tx, finish_rx) = oneshot::channel();
2016        let interceptor = Arc::new(BlockingInsertSelectInterceptor::new(started_tx, finish_rx));
2017        let instance = Arc::new(test_instance_with_insert_select_interceptor(interceptor).await?);
2018
2019        let insert_task = tokio::spawn({
2020            let instance = instance.clone();
2021            async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
2022        });
2023
2024        tokio::time::timeout(Duration::from_secs(5), started_rx.recv())
2025            .await
2026            .context(InsertStartTimeoutSnafu)?
2027            .context(InsertStartChannelClosedSnafu)?;
2028
2029        let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
2030        let process_list = output.data.pretty_print().await;
2031        assert!(
2032            process_list.contains(insert_sql),
2033            "process list did not contain running insert:\n{process_list}"
2034        );
2035
2036        finish_tx
2037            .send(())
2038            .map_err(|_| ReleaseBlockedInsertSnafu.build())?;
2039        insert_task.await.context(InsertTaskPanicSnafu)??;
2040
2041        Ok(())
2042    }
2043
2044    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2045    async fn test_kill_query_cancels_insert_select() -> TestResult<()> {
2046        assert_kill_cancels_insert_select("KILL QUERY 4242").await
2047    }
2048
2049    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2050    async fn test_kill_process_id_cancels_insert_select() -> TestResult<()> {
2051        assert_kill_cancels_insert_select("KILL 'test-frontend/4242'").await
2052    }
2053
2054    async fn assert_kill_cancels_insert_select(kill_sql: &str) -> TestResult<()> {
2055        let insert_sql = "INSERT INTO target SELECT * FROM source";
2056        let (source_polled_tx, source_polled_rx) = oneshot::channel();
2057        let instance = Arc::new(
2058            test_instance_with_tables(
2059                pending_table(1024, "source", source_polled_tx)?,
2060                test_table(1025, "target")?,
2061            )
2062            .await?,
2063        );
2064
2065        let insert_task = tokio::spawn({
2066            let instance = instance.clone();
2067            async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
2068        });
2069
2070        tokio::time::timeout(Duration::from_secs(5), source_polled_rx)
2071            .await
2072            .context(SourcePollTimeoutSnafu)?
2073            .context(SourcePollChannelClosedSnafu)?;
2074
2075        let output = execute_one_sql(&instance, kill_sql, test_query_ctx(43)).await?;
2076        assert!(matches!(output.data, OutputData::AffectedRows(1)));
2077
2078        let insert_result = tokio::time::timeout(Duration::from_secs(5), insert_task)
2079            .await
2080            .context(InsertTaskTimeoutSnafu)?
2081            .context(InsertTaskPanicSnafu)?;
2082        let err = match insert_result {
2083            Ok(_) => return InsertSelectNotCancelledSnafu.fail(),
2084            Err(TestError::ExecuteSql { source, .. }) => source,
2085            Err(err) => return Err(err),
2086        };
2087        assert_eq!(StatusCode::Cancelled, err.status_code());
2088
2089        let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
2090        let process_list = output.data.pretty_print().await;
2091        assert!(
2092            !process_list.contains(insert_sql),
2093            "process list still contains killed insert:\n{process_list}"
2094        );
2095
2096        Ok(())
2097    }
2098
2099    fn insert_dml_plan() -> LogicalPlan {
2100        let schema = SchemaRef::new(Schema::new(vec![Field::new(
2101            "value",
2102            DataType::Int64,
2103            true,
2104        )]));
2105        let target = Arc::new(LogicalTableSource::new(schema));
2106        let input = LogicalPlanBuilder::empty(false).build().unwrap();
2107
2108        LogicalPlanBuilder::insert_into(input, "demo", target, InsertOp::Append)
2109            .unwrap()
2110            .build()
2111            .unwrap()
2112    }
2113
2114    #[test]
2115    fn test_exec_validation() {
2116        let query_ctx = QueryContext::arc();
2117        let plugins: Plugins = Plugins::new();
2118        plugins.insert(QueryOptions {
2119            disallow_cross_catalog_query: true,
2120        });
2121
2122        let sql = r#"
2123        SELECT * FROM demo;
2124        EXPLAIN SELECT * FROM demo;
2125        CREATE DATABASE test_database;
2126        SHOW DATABASES;
2127        "#;
2128        let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
2129        assert_eq!(stmts.len(), 4);
2130        for stmt in stmts {
2131            let re = check_permission(plugins.clone(), &stmt, &query_ctx);
2132            re.unwrap();
2133        }
2134
2135        let sql = r#"
2136        SHOW CREATE TABLE demo;
2137        ALTER TABLE demo ADD COLUMN new_col INT;
2138        "#;
2139        let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
2140        assert_eq!(stmts.len(), 2);
2141        for stmt in stmts {
2142            let re = check_permission(plugins.clone(), &stmt, &query_ctx);
2143            re.unwrap();
2144        }
2145
2146        fn replace_test(template_sql: &str, plugins: Plugins, query_ctx: &QueryContextRef) {
2147            // test right
2148            let right = vec![("", ""), ("", "public."), ("greptime.", "public.")];
2149            for (catalog, schema) in right {
2150                let sql = do_fmt(template_sql, catalog, schema);
2151                do_test(&sql, plugins.clone(), query_ctx, true);
2152            }
2153
2154            let wrong = vec![
2155                ("wrongcatalog.", "public."),
2156                ("wrongcatalog.", "wrongschema."),
2157            ];
2158            for (catalog, schema) in wrong {
2159                let sql = do_fmt(template_sql, catalog, schema);
2160                do_test(&sql, plugins.clone(), query_ctx, false);
2161            }
2162        }
2163
2164        fn do_fmt(template: &str, catalog: &str, schema: &str) -> String {
2165            let vars = HashMap::from([
2166                ("catalog".to_string(), catalog),
2167                ("schema".to_string(), schema),
2168            ]);
2169            template.format(&vars).unwrap()
2170        }
2171
2172        fn do_test(sql: &str, plugins: Plugins, query_ctx: &QueryContextRef, is_ok: bool) {
2173            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
2174            let re = check_permission(plugins, stmt, query_ctx);
2175            if is_ok {
2176                re.unwrap();
2177            } else {
2178                assert!(re.is_err());
2179            }
2180        }
2181
2182        // test insert
2183        let sql = "INSERT INTO {catalog}{schema}monitor(host) VALUES ('host1');";
2184        replace_test(sql, plugins.clone(), &query_ctx);
2185
2186        // test create table
2187        let sql = r#"CREATE TABLE {catalog}{schema}demo(
2188                            host STRING,
2189                            ts TIMESTAMP,
2190                            TIME INDEX (ts),
2191                            PRIMARY KEY(host)
2192                        ) engine=mito;"#;
2193        replace_test(sql, plugins.clone(), &query_ctx);
2194
2195        // test drop table
2196        let sql = "DROP TABLE {catalog}{schema}demo;";
2197        replace_test(sql, plugins.clone(), &query_ctx);
2198
2199        // test undrop table
2200        let sql = "UNDROP TABLE {catalog}{schema}demo;";
2201        replace_test(sql, plugins.clone(), &query_ctx);
2202
2203        // test show tables
2204        let sql = "SHOW TABLES FROM public";
2205        let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
2206        check_permission(plugins.clone(), &stmt[0], &query_ctx).unwrap();
2207
2208        let sql = "SHOW TABLES FROM private";
2209        let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
2210        let re = check_permission(plugins.clone(), &stmt[0], &query_ctx);
2211        assert!(re.is_ok());
2212
2213        // test describe table
2214        let sql = "DESC TABLE {catalog}{schema}demo;";
2215        replace_test(sql, plugins.clone(), &query_ctx);
2216
2217        let comment_flow_cases = [
2218            ("COMMENT ON FLOW my_flow IS 'comment';", true),
2219            ("COMMENT ON FLOW greptime.my_flow IS 'comment';", true),
2220            ("COMMENT ON FLOW wrongcatalog.my_flow IS 'comment';", false),
2221        ];
2222        for (sql, is_ok) in comment_flow_cases {
2223            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
2224            let result = check_permission(plugins.clone(), stmt, &query_ctx);
2225            assert_eq!(result.is_ok(), is_ok);
2226        }
2227
2228        let show_flow_cases = [
2229            ("SHOW CREATE FLOW my_flow;", true),
2230            ("SHOW CREATE FLOW greptime.my_flow;", true),
2231            ("SHOW CREATE FLOW wrongcatalog.my_flow;", false),
2232        ];
2233        for (sql, is_ok) in show_flow_cases {
2234            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
2235            let result = check_permission(plugins.clone(), stmt, &query_ctx);
2236            assert_eq!(result.is_ok(), is_ok);
2237        }
2238    }
2239}