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 entity_graph;
18mod export_database;
19mod grpc;
20mod import_packed;
21mod influxdb;
22mod jaeger;
23mod log_handler;
24mod logical_batcher;
25mod logs;
26mod opentsdb;
27mod otlp;
28pub mod prom_store;
29mod promql;
30mod region_query;
31
32use std::collections::HashSet;
33use std::pin::Pin;
34use std::sync::atomic::AtomicBool;
35use std::sync::{Arc, OnceLock, atomic};
36use std::time::{Duration, SystemTime};
37
38use async_stream::stream;
39use async_trait::async_trait;
40use auth::{
41    PROMQL_QUERY, PermissionChecker, PermissionCheckerRef, PermissionReq, PermissionTableTarget,
42    PermissionTableTargets,
43};
44use catalog::CatalogManagerRef;
45use catalog::process_manager::{
46    ProcessManagerRef, QueryStatement as CatalogQueryStatement, SlowQueryRecorder, SlowQueryTimer,
47};
48use client::OutputData;
49use common_base::Plugins;
50use common_base::cancellation::CancellableFuture;
51use common_error::ext::{BoxedError, ErrorExt};
52use common_event_recorder::EventRecorderRef;
53use common_meta::cache::TableFlownodeSetCacheRef;
54use common_meta::cache_invalidator::CacheInvalidatorRef;
55use common_meta::key::TableMetadataManagerRef;
56use common_meta::key::table_name::TableNameKey;
57use common_meta::node_manager::NodeManagerRef;
58use common_meta::procedure_executor::ProcedureExecutorRef;
59use common_query::Output;
60use common_recordbatch::RecordBatchStreamWrapper;
61use common_recordbatch::error::StreamTimeoutSnafu;
62use common_telemetry::logging::SlowQueryOptions;
63use common_telemetry::{debug, error, tracing};
64use dashmap::DashMap;
65use datafusion::dataframe::DataFrame;
66use datafusion::physical_plan::ExecutionPlan;
67use datafusion_expr::LogicalPlan;
68use futures::{Stream, StreamExt, future};
69use lazy_static::lazy_static;
70use operator::delete::DeleterRef;
71use operator::insert::InserterRef;
72use operator::statement::{StatementExecutor, StatementExecutorRef};
73use partition::manager::PartitionRuleManagerRef;
74use pipeline::pipeline_operator::PipelineOperator;
75use prometheus::HistogramTimer;
76use promql_parser::label::Matcher;
77use query::QueryEngineRef;
78use query::metrics::OnDone;
79use query::parser::{PromQuery, QueryStatement};
80use query::query_engine::DescribeResult;
81use query::query_engine::options::{QueryOptions, validate_catalog_and_schema};
82use servers::batcher::logical_table::LogicalTablePendingRowsBatcher;
83use servers::error::{
84    self as server_error, AuthSnafu, CommonMetaSnafu, ExecuteQuerySnafu,
85    OtlpMetricModeIncompatibleSnafu, UnexpectedResultSnafu,
86};
87use servers::interceptor::{
88    PromQueryInterceptor, PromQueryInterceptorRef, SqlQueryInterceptor, SqlQueryInterceptorRef,
89};
90use servers::otlp::metrics::legacy_normalize_otlp_name;
91use servers::prometheus_handler::{
92    ParsedPromQuery, PrometheusHandler, resolve_schema_from_matchers,
93};
94use servers::query_handler::sql::SqlQueryHandler;
95use session::context::{Channel, QueryContextRef};
96use session::table_name::table_idents_to_full_name;
97use snafu::prelude::*;
98use sql::ast::ObjectNamePartExt;
99use sql::dialect::Dialect;
100use sql::parser::{ParseOptions, ParserContext};
101use sql::statements::comment::CommentObject;
102use sql::statements::copy::{CopyDatabase, CopyTable};
103use sql::statements::statement::Statement;
104use sql::statements::tql::Tql;
105use sql::util::{extract_tables_from_prom_expr_checked, extract_tables_from_statement_checked};
106use sqlparser::ast::{AnalyzeFormat, ObjectName};
107use table::requests::{OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM};
108use tracing::Span;
109
110use crate::error::{
111    self, CollectRecordbatchSnafu, Error, ExecLogicalPlanSnafu, ExecutePromqlSnafu, ExternalSnafu,
112    InvalidSqlSnafu, ParseSqlSnafu, PermissionSnafu, PlanStatementSnafu, Result,
113    SqlExecInterceptedSnafu, StatementTimeoutSnafu, TableOperationSnafu,
114};
115use crate::service_config::InfluxdbMergeMode;
116use crate::stream_wrapper::CancellableStreamWrapper;
117
118lazy_static! {
119    static ref OTLP_LEGACY_DEFAULT_VALUE: String = "legacy".to_string();
120}
121
122/// The frontend instance contains necessary components, and implements many
123/// traits, like [`servers::query_handler::grpc::GrpcQueryHandler`],
124/// [`servers::query_handler::sql::SqlQueryHandler`], etc.
125#[derive(Clone)]
126pub struct Instance {
127    frontend_peer_addr: String,
128    experimental_metric_export: bool,
129    catalog_manager: CatalogManagerRef,
130    pipeline_operator: Arc<PipelineOperator>,
131    statement_executor: Arc<StatementExecutor>,
132    query_engine: QueryEngineRef,
133    plugins: Plugins,
134    inserter: InserterRef,
135    logical_batcher: Arc<OnceLock<Option<Arc<LogicalTablePendingRowsBatcher>>>>,
136    deleter: DeleterRef,
137    table_metadata_manager: TableMetadataManagerRef,
138    event_recorder: EventRecorderRef,
139    slow_query_recorder: EventRecorderRef,
140    process_manager: ProcessManagerRef,
141    slow_query_options: SlowQueryOptions,
142    influxdb_default_merge_mode: InfluxdbMergeMode,
143    trace_ingest_chunk_size: usize,
144    otlp_resource_info: bool,
145    suspend: Arc<AtomicBool>,
146
147    // cache for otlp metrics
148    // first layer key: db-string
149    // key: direct input metric name
150    // value: if runs in legacy mode
151    otlp_metrics_table_legacy_cache: DashMap<String, DashMap<String, bool>>,
152}
153
154impl Instance {
155    pub fn frontend_peer_addr(&self) -> &str {
156        &self.frontend_peer_addr
157    }
158
159    pub fn catalog_manager(&self) -> &CatalogManagerRef {
160        &self.catalog_manager
161    }
162
163    pub fn query_engine(&self) -> &QueryEngineRef {
164        &self.query_engine
165    }
166
167    pub fn plugins(&self) -> &Plugins {
168        &self.plugins
169    }
170
171    fn check_permission(
172        &self,
173        ctx: &QueryContextRef,
174        req: PermissionReq<'_>,
175    ) -> server_error::Result<()> {
176        self.plugins
177            .get::<PermissionCheckerRef>()
178            .as_ref()
179            .check_permission(ctx.current_user(), req)
180            .context(AuthSnafu)?;
181        Ok(())
182    }
183
184    pub fn statement_executor(&self) -> &StatementExecutorRef {
185        &self.statement_executor
186    }
187
188    pub fn table_metadata_manager(&self) -> &TableMetadataManagerRef {
189        &self.table_metadata_manager
190    }
191
192    pub fn inserter(&self) -> &InserterRef {
193        &self.inserter
194    }
195
196    pub fn process_manager(&self) -> &ProcessManagerRef {
197        &self.process_manager
198    }
199
200    /// Returns the event recorder configured for this frontend instance.
201    pub fn event_recorder(&self) -> EventRecorderRef {
202        self.event_recorder.clone()
203    }
204
205    pub fn node_manager(&self) -> &NodeManagerRef {
206        self.inserter.node_manager()
207    }
208
209    pub fn partition_manager(&self) -> &PartitionRuleManagerRef {
210        self.inserter.partition_manager()
211    }
212
213    pub fn table_flownode_set_cache(&self) -> &TableFlownodeSetCacheRef {
214        self.inserter.table_flownode_set_cache()
215    }
216
217    pub fn cache_invalidator(&self) -> &CacheInvalidatorRef {
218        self.statement_executor.cache_invalidator()
219    }
220
221    pub fn procedure_executor(&self) -> &ProcedureExecutorRef {
222        self.statement_executor.procedure_executor()
223    }
224
225    pub fn suspend_state(&self) -> Arc<AtomicBool> {
226        self.suspend.clone()
227    }
228
229    pub(crate) fn is_suspended(&self) -> bool {
230        self.suspend.load(atomic::Ordering::Relaxed)
231    }
232}
233
234fn parse_stmt(sql: &str, dialect: &(dyn Dialect + Send + Sync)) -> Result<Vec<Statement>> {
235    ParserContext::create_with_dialect(sql, dialect, ParseOptions::default()).context(ParseSqlSnafu)
236}
237
238fn is_explain_analyze_verbose(stmt: &Statement) -> bool {
239    matches!(stmt, Statement::Explain(explain) if explain.analyze && explain.verbose)
240        || matches!(stmt, Statement::Tql(Tql::Analyze(analyze)) if analyze.is_verbose)
241}
242
243fn validate_analyze_stream_statement(stmt: &mut Statement) -> Result<()> {
244    let (is_verbose, format) = match stmt {
245        Statement::Explain(explain) => (explain.analyze && explain.verbose, &mut explain.format),
246        Statement::Tql(Tql::Analyze(analyze)) => (analyze.is_verbose, &mut analyze.format),
247        _ => {
248            return InvalidSqlSnafu {
249                err_msg: "only EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE statement is supported",
250            }
251            .fail();
252        }
253    };
254
255    ensure!(
256        is_verbose,
257        InvalidSqlSnafu {
258            err_msg: "statement must be EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE"
259        }
260    );
261    match format {
262        None | Some(AnalyzeFormat::JSON) => {
263            // Keep explicit FORMAT JSON accepted, but pass JSON through
264            // QueryContext.explain_format instead of the statement to avoid
265            // the planner's current `EXPLAIN VERBOSE with FORMAT` limitation.
266            *format = None;
267            Ok(())
268        }
269        Some(_) => InvalidSqlSnafu {
270            err_msg: "only FORMAT JSON is supported for EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE",
271        }
272        .fail(),
273    }
274}
275
276impl Instance {
277    fn statement_slow_query_timer(
278        &self,
279        stmt: &Statement,
280        schema_name: String,
281    ) -> Option<SlowQueryTimer> {
282        if !stmt.is_readonly() || !self.slow_query_options.enable {
283            return None;
284        }
285
286        Some(SlowQueryTimer::new(
287            CatalogQueryStatement::Sql(stmt.clone()),
288            schema_name,
289            self.slow_query_options.threshold,
290            self.slow_query_options.sample_ratio,
291            self.slow_query_options.record_type,
292            self.slow_query_recorder.clone(),
293        ))
294    }
295
296    async fn query_statement(&self, stmt: Statement, query_ctx: QueryContextRef) -> Result<Output> {
297        check_permission(self.plugins.clone(), &stmt, &query_ctx)?;
298
299        let query_interceptor = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
300        let query_interceptor = query_interceptor.as_ref();
301
302        if should_track_statement_process(&stmt) {
303            let catalog_name = query_ctx.current_catalog().to_string();
304            let schema_name = query_ctx.current_schema();
305            let slow_query_timer = self.statement_slow_query_timer(&stmt, schema_name.clone());
306            let timeout_recorder = is_explain_analyze_verbose(&stmt)
307                .then(|| slow_query_timer.as_ref().map(SlowQueryTimer::recorder))
308                .flatten();
309
310            let ticket = self.process_manager.register_query(
311                catalog_name,
312                vec![schema_name],
313                stmt.to_string(),
314                query_ctx.conn_info().to_string(),
315                Some(query_ctx.process_id()),
316                slow_query_timer,
317            );
318
319            let query_fut = self.exec_statement_with_timeout(
320                stmt,
321                query_ctx,
322                query_interceptor,
323                timeout_recorder,
324            );
325
326            CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
327                .await
328                .map_err(|_| error::CancelledSnafu.build())?
329                .map(|output| {
330                    let Output { meta, data } = output;
331
332                    let data = match data {
333                        OutputData::Stream(stream) => OutputData::Stream(Box::pin(
334                            CancellableStreamWrapper::new(stream, ticket),
335                        )),
336                        other => other,
337                    };
338                    Output { data, meta }
339                })
340        } else {
341            self.exec_statement_with_timeout(stmt, query_ctx, query_interceptor, None)
342                .await
343        }
344    }
345
346    async fn exec_statement_with_timeout(
347        &self,
348        stmt: Statement,
349        query_ctx: QueryContextRef,
350        query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
351        timeout_recorder: Option<SlowQueryRecorder>,
352    ) -> Result<Output> {
353        let timeout = derive_timeout(&stmt, &query_ctx);
354        match timeout {
355            Some(timeout) => {
356                let start = tokio::time::Instant::now();
357                let output = tokio::time::timeout(
358                    timeout,
359                    self.exec_statement(stmt, query_ctx, query_interceptor),
360                )
361                .await
362                .map_err(|_| StatementTimeoutSnafu.build())??;
363                let output = map_query_output(output)?;
364                // compute remaining timeout
365                let remaining_timeout = timeout.checked_sub(start.elapsed()).unwrap_or_default();
366                attach_timeout(output, remaining_timeout, timeout_recorder)
367            }
368            None => self
369                .exec_statement(stmt, query_ctx, query_interceptor)
370                .await
371                .and_then(map_query_output),
372        }
373    }
374
375    async fn exec_statement(
376        &self,
377        stmt: Statement,
378        query_ctx: QueryContextRef,
379        query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
380    ) -> Result<Output> {
381        match stmt {
382            Statement::Query(_) | Statement::Explain(_) | Statement::Delete(_) => {
383                // TODO: remove this when format is supported in datafusion
384                if let Statement::Explain(explain) = &stmt
385                    && let Some(format) = explain.format()
386                {
387                    query_ctx.set_explain_format(format.to_string());
388                }
389
390                self.plan_and_exec_sql(stmt, &query_ctx, query_interceptor)
391                    .await
392            }
393            Statement::Tql(tql) => {
394                self.plan_and_exec_tql(&query_ctx, query_interceptor, tql)
395                    .await
396            }
397            _ => {
398                query_interceptor.pre_execute(Some(&stmt), None, query_ctx.clone())?;
399                if let Statement::Copy(sql::statements::copy::Copy::CopyDatabase(
400                    CopyDatabase::From(arg),
401                )) = &stmt
402                    && arg.with.get("metric_data_layout").is_some()
403                {
404                    return self
405                        .copy_packed_database(arg.clone(), &stmt, query_ctx)
406                        .await;
407                }
408                if let Statement::ShowVariables(show) = &stmt
409                    && show
410                        .variable
411                        .to_string()
412                        .eq_ignore_ascii_case("experimental_metric_export")
413                {
414                    return self.show_metric_export_capability();
415                }
416                if let Statement::Copy(sql::statements::copy::Copy::CopyDatabase(CopyDatabase::To(
417                    arg,
418                ))) = &stmt
419                    && export_database::parse_metric_export_requested(&arg.with)?
420                {
421                    return self.copy_metric_database(arg.clone(), query_ctx).await;
422                }
423                self.statement_executor
424                    .execute_sql(stmt, query_ctx)
425                    .await
426                    .context(TableOperationSnafu)
427            }
428        }
429    }
430
431    async fn plan_and_exec_sql(
432        &self,
433        stmt: Statement,
434        query_ctx: &QueryContextRef,
435        query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
436    ) -> Result<Output> {
437        let stmt = QueryStatement::Sql(stmt);
438        let plan = self
439            .statement_executor
440            .plan(&stmt, query_ctx.clone())
441            .await?;
442        let QueryStatement::Sql(stmt) = stmt else {
443            unreachable!()
444        };
445        query_interceptor.pre_execute(Some(&stmt), Some(&plan), query_ctx.clone())?;
446
447        self.statement_executor
448            .exec_plan(plan, query_ctx.clone())
449            .await
450            .context(TableOperationSnafu)
451    }
452
453    async fn plan_and_exec_tql(
454        &self,
455        query_ctx: &QueryContextRef,
456        query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
457        tql: Tql,
458    ) -> Result<Output> {
459        let plan = self
460            .statement_executor
461            .plan_tql(tql.clone(), query_ctx)
462            .await?;
463        query_interceptor.pre_execute(
464            Some(&Statement::Tql(tql)),
465            Some(&plan),
466            query_ctx.clone(),
467        )?;
468        self.statement_executor
469            .exec_plan(plan, query_ctx.clone())
470            .await
471            .context(TableOperationSnafu)
472    }
473
474    async fn check_otlp_legacy(
475        &self,
476        names: &[String],
477        ctx: &QueryContextRef,
478    ) -> server_error::Result<bool> {
479        let db_string = ctx.get_db_string();
480        // fast cache check
481        let cache = self
482            .otlp_metrics_table_legacy_cache
483            .entry(db_string.clone())
484            .or_default();
485        if let Some(flag) = fast_legacy_check(&cache, names)? {
486            return Ok(flag);
487        }
488        // release cache reference to avoid lock contention
489        drop(cache);
490
491        let catalog = ctx.current_catalog();
492        let schema = ctx.current_schema();
493
494        // query legacy table names
495        let normalized_names = names
496            .iter()
497            .map(|n| legacy_normalize_otlp_name(n))
498            .collect::<Vec<_>>();
499        let table_names = normalized_names
500            .iter()
501            .map(|n| TableNameKey::new(catalog, &schema, n))
502            .collect::<Vec<_>>();
503        let table_values = self
504            .table_metadata_manager()
505            .table_name_manager()
506            .batch_get(table_names)
507            .await
508            .context(CommonMetaSnafu)?;
509        let table_ids = table_values
510            .into_iter()
511            .filter_map(|v| v.map(|vi| vi.table_id()))
512            .collect::<Vec<_>>();
513
514        // means no existing table is found, use new mode
515        if table_ids.is_empty() {
516            return Ok(false);
517        }
518
519        // has existing table, check table options
520        let table_infos = self
521            .table_metadata_manager()
522            .table_info_manager()
523            .batch_get(&table_ids)
524            .await
525            .context(CommonMetaSnafu)?;
526        let options = table_infos
527            .values()
528            .map(|info| {
529                info.table_info
530                    .meta
531                    .options
532                    .extra_options
533                    .get(OTLP_METRIC_COMPAT_KEY)
534                    .unwrap_or(&OTLP_LEGACY_DEFAULT_VALUE)
535            })
536            .collect::<Vec<_>>();
537        if !options.is_empty() {
538            // check value consistency
539            let has_prom = options.iter().any(|opt| *opt == OTLP_METRIC_COMPAT_PROM);
540            let has_legacy = options
541                .iter()
542                .any(|opt| *opt == OTLP_LEGACY_DEFAULT_VALUE.as_str());
543            ensure!(!(has_prom && has_legacy), OtlpMetricModeIncompatibleSnafu);
544            Ok(has_legacy)
545        } else {
546            // no table info, use new mode
547            Ok(false)
548        }
549    }
550
551    fn cache_otlp_legacy(
552        &self,
553        names: &[String],
554        ctx: &QueryContextRef,
555        is_legacy: bool,
556    ) -> server_error::Result<()> {
557        let cache = self
558            .otlp_metrics_table_legacy_cache
559            .entry(ctx.get_db_string())
560            .or_default();
561        cache_legacy_mode(&cache, names, is_legacy)
562    }
563}
564
565fn fast_legacy_check(
566    cache: &DashMap<String, bool>,
567    names: &[String],
568) -> server_error::Result<Option<bool>> {
569    let hit_cache = names
570        .iter()
571        .filter_map(|name| cache.get(name))
572        .collect::<Vec<_>>();
573    if !hit_cache.is_empty() {
574        let hit_legacy = hit_cache.iter().any(|en| *en.value());
575        let hit_prom = hit_cache.iter().any(|en| !*en.value());
576
577        // hit but have true and false, means both legacy and new mode are used
578        // we cannot handle this case, so return error
579        // add doc links in err msg later
580        ensure!(!(hit_legacy && hit_prom), OtlpMetricModeIncompatibleSnafu);
581
582        Ok(Some(hit_legacy))
583    } else {
584        Ok(None)
585    }
586}
587
588fn cache_legacy_mode(
589    cache: &DashMap<String, bool>,
590    names: &[String],
591    is_legacy: bool,
592) -> server_error::Result<()> {
593    for name in names {
594        let cached = cache.entry(name.clone()).or_insert(is_legacy);
595        ensure!(*cached == is_legacy, OtlpMetricModeIncompatibleSnafu);
596    }
597    Ok(())
598}
599
600/// If the relevant variables are set, the timeout is enforced for all PostgreSQL statements.
601/// For MySQL, it applies only to read-only statements.
602fn derive_timeout(stmt: &Statement, query_ctx: &QueryContextRef) -> Option<Duration> {
603    let query_timeout = query_ctx.query_timeout()?;
604    if query_timeout.is_zero() {
605        return None;
606    }
607    match query_ctx.channel() {
608        Channel::Mysql if stmt.is_readonly() => Some(query_timeout),
609        Channel::Postgres => Some(query_timeout),
610        _ => None,
611    }
612}
613
614/// Derives timeout for plan execution.
615fn derive_timeout_for_plan(plan: &LogicalPlan, query_ctx: &QueryContextRef) -> Option<Duration> {
616    let query_timeout = query_ctx.query_timeout()?;
617    if query_timeout.is_zero() {
618        return None;
619    }
620    match query_ctx.channel() {
621        Channel::Mysql if is_readonly_plan(plan) => Some(query_timeout),
622        Channel::Postgres => Some(query_timeout),
623        _ => None,
624    }
625}
626
627fn record_explain_analyze_timeout(
628    recorder: Option<&SlowQueryRecorder>,
629    plan: Option<&Arc<dyn ExecutionPlan>>,
630) {
631    let Some(recorder) = recorder else {
632        return;
633    };
634    let metrics = plan
635        .and_then(|plan| query::analyze_plan_metrics_to_json_value(plan, true).ok())
636        .unwrap_or_else(|| serde_json::json!([]));
637    recorder.force_record_with_payload(serde_json::json!({
638        "timed_out": true,
639        "metrics": metrics,
640    }));
641}
642
643fn attach_timeout(
644    output: Output,
645    mut timeout: Duration,
646    timeout_recorder: Option<SlowQueryRecorder>,
647) -> Result<Output> {
648    if timeout.is_zero() {
649        return StatementTimeoutSnafu.fail();
650    }
651
652    let plan = timeout_recorder
653        .as_ref()
654        .and_then(|_| output.meta.plan.clone());
655    let output = match output.data {
656        OutputData::AffectedRows(_) | OutputData::RecordBatches(_) => output,
657        OutputData::Stream(mut stream) => {
658            let schema = stream.schema();
659            let s = Box::pin(stream! {
660                let mut start = tokio::time::Instant::now();
661                while let Some(item) = tokio::time::timeout(timeout, stream.next()).await.map_err(|_| {
662                    record_explain_analyze_timeout(timeout_recorder.as_ref(), plan.as_ref());
663                    StreamTimeoutSnafu.build()
664                })? {
665                    yield item;
666
667                    let now = tokio::time::Instant::now();
668                    timeout = timeout.checked_sub(now - start).unwrap_or(Duration::ZERO);
669                    start = now;
670                    // tokio::time::timeout may not return an error immediately when timeout is 0.
671                    if timeout.is_zero() {
672                        record_explain_analyze_timeout(timeout_recorder.as_ref(), plan.as_ref());
673                        StreamTimeoutSnafu.fail()?;
674                    }
675                }
676            }) as Pin<Box<dyn Stream<Item = _> + Send>>;
677            let stream = RecordBatchStreamWrapper {
678                schema,
679                stream: s,
680                output_ordering: None,
681                metrics: Default::default(),
682                span: Span::current(),
683            };
684            Output::new(OutputData::Stream(Box::pin(stream)), output.meta)
685        }
686    };
687
688    Ok(output)
689}
690
691impl Instance {
692    async fn check_sql_permission(
693        &self,
694        stmt: &Statement,
695        query_ctx: &QueryContextRef,
696    ) -> Result<()> {
697        self.plugins
698            .get::<PermissionCheckerRef>()
699            .as_ref()
700            .check_permission_with_context(
701                query_ctx.current_user(),
702                PermissionReq::SqlStatement(stmt),
703                Some(&query_ctx.current_schema()),
704            )
705            .context(PermissionSnafu)?;
706
707        let targets = match extract_tables_from_statement_checked(stmt) {
708            Some(tables) => PermissionTableTargets::resolved(
709                tables
710                    .map(|name| {
711                        table_idents_to_full_name(&name, query_ctx).map(
712                            |(catalog, schema, table)| {
713                                PermissionTableTarget::new(catalog, schema, table)
714                            },
715                        )
716                    })
717                    .collect::<std::result::Result<Vec<_>, _>>()
718                    .map_err(BoxedError::new)
719                    .context(ExternalSnafu)?,
720            ),
721            None => PermissionTableTargets::Unresolved,
722        };
723        let targets = self
724            .resolve_query_permission_targets(targets, query_ctx)
725            .await
726            .map_err(BoxedError::new)
727            .context(ExternalSnafu)?;
728        self.check_table_permission(query_ctx, PermissionReq::SqlStatement(stmt), targets)
729            .context(PermissionSnafu)?;
730        Ok(())
731    }
732
733    #[tracing::instrument(skip_all, name = "SqlQueryHandler::do_analyze_stream_query")]
734    async fn do_analyze_stream_query_inner(
735        &self,
736        query: &str,
737        query_ctx: QueryContextRef,
738    ) -> Result<Output> {
739        ensure!(!self.is_suspended(), error::SuspendedSnafu);
740
741        let query_interceptor_opt = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
742        let query_interceptor = query_interceptor_opt.as_ref();
743        let query = query_interceptor.pre_parsing(query, query_ctx.clone())?;
744        let mut stmts = parse_stmt(query.as_ref(), query_ctx.sql_dialect())
745            .and_then(|stmts| query_interceptor.post_parsing(stmts, query_ctx.clone()))?;
746
747        ensure!(
748            stmts.len() == 1,
749            InvalidSqlSnafu {
750                err_msg: "only a single EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE statement is supported"
751            }
752        );
753        let mut stmt = stmts.remove(0);
754        validate_analyze_stream_statement(&mut stmt)?;
755        query_ctx.set_explain_format(AnalyzeFormat::JSON.to_string());
756
757        self.check_sql_permission(&stmt, &query_ctx).await?;
758        check_permission(self.plugins.clone(), &stmt, &query_ctx)?;
759        let catalog_name = query_ctx.current_catalog().to_string();
760        let schema_name = query_ctx.current_schema();
761        let slow_query_timer = self.statement_slow_query_timer(&stmt, schema_name.clone());
762        let ticket = self.process_manager.register_query(
763            catalog_name,
764            vec![schema_name],
765            stmt.to_string(),
766            query_ctx.conn_info().to_string(),
767            Some(query_ctx.process_id()),
768            slow_query_timer,
769        );
770        let query_fut =
771            self.exec_statement_with_timeout(stmt, query_ctx.clone(), query_interceptor, None);
772        let output = CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
773            .await
774            .map_err(|_| error::CancelledSnafu.build())??;
775        let Output { meta, data } = output;
776        let data = match data {
777            OutputData::Stream(stream) => OutputData::Stream(Box::pin(
778                CancellableStreamWrapper::new_cancel_on_drop(stream, ticket),
779            )),
780            other => other,
781        };
782        query_interceptor.post_execute(Output { data, meta }, query_ctx)
783    }
784
785    #[tracing::instrument(skip_all, name = "SqlQueryHandler::do_query")]
786    async fn do_query_inner(&self, query: &str, query_ctx: QueryContextRef) -> Vec<Result<Output>> {
787        if self.is_suspended() {
788            return vec![error::SuspendedSnafu {}.fail()];
789        }
790
791        let query_interceptor_opt = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
792        let query_interceptor = query_interceptor_opt.as_ref();
793        let query = match query_interceptor.pre_parsing(query, query_ctx.clone()) {
794            Ok(q) => q,
795            Err(e) => return vec![Err(e)],
796        };
797
798        match parse_stmt(query.as_ref(), query_ctx.sql_dialect())
799            .and_then(|stmts| query_interceptor.post_parsing(stmts, query_ctx.clone()))
800        {
801            Ok(stmts) => {
802                if stmts.is_empty() {
803                    return vec![
804                        InvalidSqlSnafu {
805                            err_msg: "empty statements",
806                        }
807                        .fail(),
808                    ];
809                }
810
811                let mut results = Vec::with_capacity(stmts.len());
812                for stmt in stmts {
813                    if let Err(e) = self.check_sql_permission(&stmt, &query_ctx).await {
814                        results.push(Err(e));
815                        break;
816                    }
817
818                    match self.query_statement(stmt.clone(), query_ctx.clone()).await {
819                        Ok(output) => {
820                            let output_result =
821                                query_interceptor.post_execute(output, query_ctx.clone());
822                            results.push(output_result);
823                        }
824                        Err(e) => {
825                            if e.status_code().should_log_error() {
826                                error!(e; "Failed to execute query: {stmt}");
827                            } else {
828                                debug!("Failed to execute query: {stmt}, {e}");
829                            }
830                            results.push(Err(e));
831                            break;
832                        }
833                    }
834                }
835                results
836            }
837            Err(e) => {
838                vec![Err(e)]
839            }
840        }
841    }
842
843    async fn exec_plan(&self, plan: LogicalPlan, query_ctx: QueryContextRef) -> Result<Output> {
844        self.query_engine
845            .execute(plan, query_ctx)
846            .await
847            .context(ExecLogicalPlanSnafu)
848    }
849
850    async fn exec_plan_with_timeout(
851        &self,
852        plan: LogicalPlan,
853        query_ctx: QueryContextRef,
854        timeout_recorder: Option<SlowQueryRecorder>,
855    ) -> Result<Output> {
856        let timeout = derive_timeout_for_plan(&plan, &query_ctx);
857        match timeout {
858            Some(timeout) => {
859                let start = tokio::time::Instant::now();
860                let output = tokio::time::timeout(timeout, self.exec_plan(plan, query_ctx))
861                    .await
862                    .map_err(|_| StatementTimeoutSnafu.build())??;
863                let output = map_query_output(output)?;
864                let remaining_timeout = timeout.checked_sub(start.elapsed()).unwrap_or_default();
865                attach_timeout(output, remaining_timeout, timeout_recorder)
866            }
867            None => self
868                .exec_plan(plan, query_ctx)
869                .await
870                .and_then(map_query_output),
871        }
872    }
873
874    async fn do_exec_plan_inner(
875        &self,
876        plan: LogicalPlan,
877        stmt: Option<Statement>,
878        query_ctx: QueryContextRef,
879    ) -> Result<Output> {
880        ensure!(!self.is_suspended(), error::SuspendedSnafu);
881
882        let query_interceptor_opt = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
883        let query_interceptor = query_interceptor_opt.as_ref();
884
885        query_interceptor.pre_execute(stmt.as_ref(), Some(&plan), query_ctx.clone())?;
886
887        // TQL EXPLAIN/ANALYZE formats are consumed from the query context at
888        // execution time (see `optimize_physical_plan`); re-apply the side
889        // effect of `plan_tql` that was lost when the plan was built during
890        // Describe. `explain_format` is per-query state, so this never
891        // overwrites anything.
892        if let Some(Statement::Tql(tql)) = &stmt {
893            let format = match tql {
894                Tql::Explain(explain) => explain.format.as_ref(),
895                Tql::Analyze(analyze) => analyze.format.as_ref(),
896                Tql::Eval(_) => None,
897            };
898            if let Some(format) = format {
899                query_ctx.set_explain_format(format.to_string());
900            }
901        }
902
903        let query = stmt
904            .as_ref()
905            .map(|s| s.to_string())
906            .unwrap_or_else(|| plan.display_indent().to_string());
907
908        let plan_is_readonly = is_readonly_plan(&plan);
909        let result = if should_track_plan_process(stmt.as_ref(), &plan) {
910            let catalog_name = query_ctx.current_catalog().to_string();
911            let schema_name = query_ctx.current_schema();
912            let slow_query_timer = if plan_is_readonly {
913                self.slow_query_options.enable.then(|| {
914                    SlowQueryTimer::new(
915                        CatalogQueryStatement::Plan(query.clone()),
916                        schema_name.clone(),
917                        self.slow_query_options.threshold,
918                        self.slow_query_options.sample_ratio,
919                        self.slow_query_options.record_type,
920                        self.slow_query_recorder.clone(),
921                    )
922                })
923            } else {
924                None
925            };
926
927            let timeout_recorder = stmt
928                .as_ref()
929                .is_some_and(is_explain_analyze_verbose)
930                .then(|| slow_query_timer.as_ref().map(SlowQueryTimer::recorder))
931                .flatten();
932            let ticket = self.process_manager.register_query(
933                catalog_name,
934                vec![schema_name],
935                query,
936                query_ctx.conn_info().to_string(),
937                Some(query_ctx.process_id()),
938                slow_query_timer,
939            );
940
941            let query_fut = self.exec_plan_with_timeout(plan, query_ctx.clone(), timeout_recorder);
942
943            CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
944                .await
945                .map_err(|_| error::CancelledSnafu.build())?
946                .map(|output| {
947                    let Output { meta, data } = output;
948
949                    let data = match data {
950                        OutputData::Stream(stream) => OutputData::Stream(Box::pin(
951                            CancellableStreamWrapper::new(stream, ticket),
952                        )),
953                        other => other,
954                    };
955                    Output { data, meta }
956                })
957        } else {
958            self.exec_plan_with_timeout(plan, query_ctx.clone(), None)
959                .await
960        };
961
962        result.and_then(|output| query_interceptor.post_execute(output, query_ctx))
963    }
964
965    #[tracing::instrument(skip_all, name = "SqlQueryHandler::do_promql_query")]
966    async fn do_promql_query_inner(
967        &self,
968        query: &PromQuery,
969        query_ctx: QueryContextRef,
970    ) -> Vec<Result<Output>> {
971        if self.is_suspended() {
972            return vec![error::SuspendedSnafu {}.fail()];
973        }
974
975        // check will be done in prometheus handler's do_query
976        let result = PrometheusHandler::do_query(self, query, query_ctx)
977            .await
978            .with_context(|_| ExecutePromqlSnafu {
979                query: format!("{query:?}"),
980            });
981        vec![result]
982    }
983
984    /// Builds the [`DataFrame`] for an information-schema-backed `SHOW`
985    /// statement; `None` for other statements. The future is boxed to keep
986    /// `do_describe_inner`'s state machine small.
987    fn show_statement_dataframe<'a>(
988        &'a self,
989        stmt: &'a Statement,
990        query_ctx: &'a QueryContextRef,
991    ) -> Pin<Box<dyn Future<Output = Option<query::error::Result<DataFrame>>> + Send + 'a>> {
992        Box::pin(async move {
993            let engine = &self.query_engine;
994            let catalog_manager = self.catalog_manager();
995            let ctx = query_ctx.clone();
996            let dataframe = match stmt {
997                Statement::ShowDatabases(show) => {
998                    query::sql::show_databases_dataframe(show, engine, catalog_manager, ctx).await
999                }
1000                Statement::ShowTables(show) => {
1001                    query::sql::show_tables_dataframe(show, engine, catalog_manager, ctx).await
1002                }
1003                Statement::ShowViews(show) => {
1004                    query::sql::show_views_dataframe(show, engine, catalog_manager, ctx).await
1005                }
1006                Statement::ShowFlows(show) => {
1007                    query::sql::show_flows_dataframe(show, engine, catalog_manager, ctx).await
1008                }
1009                Statement::ShowColumns(show) => {
1010                    query::sql::show_columns_dataframe(show, engine, catalog_manager, ctx).await
1011                }
1012                Statement::ShowTableStatus(show) => {
1013                    query::sql::show_table_status_dataframe(show, engine, catalog_manager, ctx)
1014                        .await
1015                }
1016                Statement::ShowCharset(kind) => {
1017                    query::sql::show_charsets_dataframe(kind, engine, catalog_manager, ctx).await
1018                }
1019                Statement::ShowCollation(kind) => {
1020                    query::sql::show_collations_dataframe(kind, engine, catalog_manager, ctx).await
1021                }
1022                Statement::ShowIndex(show) => {
1023                    query::sql::show_index_dataframe(show, engine, catalog_manager, ctx).await
1024                }
1025                Statement::ShowRegion(show) => {
1026                    query::sql::show_region_dataframe(show, engine, catalog_manager, ctx).await
1027                }
1028                Statement::ShowProcesslist(show) => {
1029                    query::sql::show_processlist_dataframe(show, engine, catalog_manager, ctx).await
1030                }
1031                _ => return None,
1032            };
1033            Some(dataframe)
1034        })
1035    }
1036
1037    async fn do_describe_inner(
1038        &self,
1039        stmt: Statement,
1040        query_ctx: QueryContextRef,
1041    ) -> Result<Option<DescribeResult>> {
1042        ensure!(!self.is_suspended(), error::SuspendedSnafu);
1043
1044        // EXPLAIN / EXPLAIN ANALYZE wrap an inner statement; describe them when the
1045        // wrapped statement is something we already plan (so that bind parameters
1046        // in the inner query get their types inferred). See #8029.
1047        let is_inner_plannable = |s: &Statement| {
1048            matches!(
1049                s,
1050                Statement::Insert(_) | Statement::Query(_) | Statement::Delete(_)
1051            )
1052        };
1053        let plannable = is_inner_plannable(&stmt)
1054            || matches!(&stmt, Statement::Explain(explain) if is_inner_plannable(explain.statement.as_ref()));
1055
1056        if let Statement::Tql(tql) = stmt {
1057            // TQL produces a logical plan; describe it from the plan so the
1058            // extended-protocol RowDescription matches the executed DataRows.
1059            self.check_sql_permission(&Statement::Tql(tql.clone()), &query_ctx)
1060                .await?;
1061            let plan = self.statement_executor.plan_tql(tql, &query_ctx).await?;
1062            return self
1063                .query_engine
1064                .describe(plan, query_ctx)
1065                .await
1066                .map(Some)
1067                .context(error::DescribeStatementSnafu);
1068        }
1069
1070        // Describe SHOW statements from the same projection the executor builds.
1071        if let Some(dataframe) = self
1072            .show_statement_dataframe(&stmt, &query_ctx)
1073            .await
1074            .transpose()
1075            .context(PlanStatementSnafu)?
1076        {
1077            self.check_sql_permission(&stmt, &query_ctx).await?;
1078            let plan = dataframe.into_unoptimized_plan();
1079            return self
1080                .query_engine
1081                .describe(plan, query_ctx)
1082                .await
1083                .map(Some)
1084                .context(error::DescribeStatementSnafu);
1085        }
1086
1087        if plannable {
1088            self.check_sql_permission(&stmt, &query_ctx).await?;
1089
1090            let plan = self
1091                .query_engine
1092                .planner()
1093                .plan(&QueryStatement::Sql(stmt), query_ctx.clone())
1094                .await
1095                .context(PlanStatementSnafu)?;
1096            self.query_engine
1097                .describe(plan, query_ctx)
1098                .await
1099                .map(Some)
1100                .context(error::DescribeStatementSnafu)
1101        } else {
1102            Ok(None)
1103        }
1104    }
1105
1106    async fn is_valid_schema_inner(&self, catalog: &str, schema: &str) -> Result<bool> {
1107        self.catalog_manager
1108            .schema_exists(catalog, schema, None)
1109            .await
1110            .context(error::CatalogSnafu)
1111    }
1112}
1113
1114#[async_trait]
1115impl SqlQueryHandler for Instance {
1116    async fn do_query(
1117        &self,
1118        query: &str,
1119        query_ctx: QueryContextRef,
1120    ) -> Vec<server_error::Result<Output>> {
1121        self.do_query_inner(query, query_ctx)
1122            .await
1123            .into_iter()
1124            .map(|result| result.map_err(BoxedError::new).context(ExecuteQuerySnafu))
1125            .collect()
1126    }
1127
1128    async fn do_analyze_stream_query(
1129        &self,
1130        query: &str,
1131        query_ctx: QueryContextRef,
1132    ) -> server_error::Result<Output> {
1133        self.do_analyze_stream_query_inner(query, query_ctx)
1134            .await
1135            .map_err(BoxedError::new)
1136            .context(ExecuteQuerySnafu)
1137    }
1138
1139    async fn do_exec_plan(
1140        &self,
1141        plan: LogicalPlan,
1142        stmt: Option<Statement>,
1143        query_ctx: QueryContextRef,
1144    ) -> server_error::Result<Output> {
1145        self.do_exec_plan_inner(plan, stmt, query_ctx)
1146            .await
1147            .map_err(BoxedError::new)
1148            .context(server_error::ExecutePlanSnafu)
1149    }
1150
1151    async fn do_promql_query(
1152        &self,
1153        query: &PromQuery,
1154        query_ctx: QueryContextRef,
1155    ) -> Vec<server_error::Result<Output>> {
1156        self.do_promql_query_inner(query, query_ctx)
1157            .await
1158            .into_iter()
1159            .map(|result| result.map_err(BoxedError::new).context(ExecuteQuerySnafu))
1160            .collect()
1161    }
1162
1163    async fn do_describe(
1164        &self,
1165        stmt: Statement,
1166        query_ctx: QueryContextRef,
1167    ) -> server_error::Result<Option<DescribeResult>> {
1168        self.do_describe_inner(stmt, query_ctx)
1169            .await
1170            .map_err(BoxedError::new)
1171            .context(server_error::DescribeStatementSnafu)
1172    }
1173
1174    async fn is_valid_schema(&self, catalog: &str, schema: &str) -> server_error::Result<bool> {
1175        self.is_valid_schema_inner(catalog, schema)
1176            .await
1177            .map_err(BoxedError::new)
1178            .context(server_error::CheckDatabaseValiditySnafu)
1179    }
1180}
1181
1182/// Expands scan-time dictionaries only when a query result leaves the frontend.
1183pub(crate) fn map_query_output(output: Output) -> Result<Output> {
1184    output
1185        .map_dictionary_to_values()
1186        .context(CollectRecordbatchSnafu)
1187}
1188
1189/// Attaches a timer to the output and observes it once the output is exhausted.
1190pub fn attach_timer(output: Output, timer: HistogramTimer) -> Output {
1191    match output.data {
1192        OutputData::AffectedRows(_) | OutputData::RecordBatches(_) => output,
1193        OutputData::Stream(stream) => {
1194            let stream = OnDone::new(stream, move || {
1195                timer.observe_duration();
1196            });
1197            Output::new(OutputData::Stream(Box::pin(stream)), output.meta)
1198        }
1199    }
1200}
1201
1202impl Instance {
1203    fn check_prom_query_privilege(&self, query_ctx: &QueryContextRef) -> server_error::Result<()> {
1204        self.plugins
1205            .get::<PermissionCheckerRef>()
1206            .as_ref()
1207            .check_permission(
1208                query_ctx.current_user(),
1209                PermissionReq::Action(PROMQL_QUERY),
1210            )
1211            .context(AuthSnafu)?;
1212        Ok(())
1213    }
1214
1215    fn prom_expr_permission_targets(
1216        &self,
1217        expr: &promql_parser::parser::Expr,
1218        query_ctx: &QueryContextRef,
1219    ) -> server_error::Result<Option<Vec<PermissionTableTarget>>> {
1220        extract_tables_from_prom_expr_checked(expr)
1221            .map(|tables| {
1222                tables
1223                    .map(|name| {
1224                        table_idents_to_full_name(&name, query_ctx).map(
1225                            |(catalog, schema, table)| {
1226                                PermissionTableTarget::new(catalog, schema, table)
1227                            },
1228                        )
1229                    })
1230                    .collect::<std::result::Result<Vec<_>, _>>()
1231                    .map_err(BoxedError::new)
1232                    .context(ExecuteQuerySnafu)
1233            })
1234            .transpose()
1235    }
1236
1237    async fn is_physical_query_permission_target(
1238        &self,
1239        target: &PermissionTableTarget,
1240        query_ctx: &QueryContextRef,
1241    ) -> server_error::Result<bool> {
1242        self.catalog_manager
1243            .table(
1244                &target.catalog,
1245                &target.schema,
1246                &target.table,
1247                Some(query_ctx),
1248            )
1249            .await
1250            .map(|table| table.is_some_and(|table| table.table_info().is_physical_table()))
1251            .map_err(BoxedError::new)
1252            .context(ExecuteQuerySnafu)
1253    }
1254
1255    async fn resolve_query_permission_targets(
1256        &self,
1257        targets: PermissionTableTargets,
1258        query_ctx: &QueryContextRef,
1259    ) -> server_error::Result<PermissionTableTargets> {
1260        const CONCURRENCY: usize = 8;
1261
1262        let checker = self.plugins.get::<PermissionCheckerRef>();
1263        if !checker.as_ref().uses_table_targets() {
1264            return Ok(targets);
1265        }
1266
1267        let PermissionTableTargets::Resolved(mut targets) = targets else {
1268            return Ok(PermissionTableTargets::Unresolved);
1269        };
1270        if targets.len() > 1 {
1271            let mut seen = HashSet::with_capacity(targets.len());
1272            targets.retain(|target| seen.insert(target.clone()));
1273        }
1274        if let [target] = targets.as_slice() {
1275            return if self
1276                .is_physical_query_permission_target(target, query_ctx)
1277                .await?
1278            {
1279                Ok(PermissionTableTargets::Unresolved)
1280            } else {
1281                Ok(PermissionTableTargets::resolved(targets))
1282            };
1283        }
1284
1285        // Bound catalog load and inspect results in target order to preserve serial semantics.
1286        for chunk in targets.chunks(CONCURRENCY) {
1287            let results = future::join_all(
1288                chunk
1289                    .iter()
1290                    .map(|target| self.is_physical_query_permission_target(target, query_ctx)),
1291            )
1292            .await;
1293            for result in results {
1294                if result? {
1295                    return Ok(PermissionTableTargets::Unresolved);
1296                }
1297            }
1298        }
1299
1300        Ok(PermissionTableTargets::resolved(targets))
1301    }
1302
1303    fn prom_queries_permission_targets(
1304        &self,
1305        queries: &[ParsedPromQuery],
1306        query_ctx: &QueryContextRef,
1307    ) -> server_error::Result<PermissionTableTargets> {
1308        let mut targets = Vec::new();
1309        let mut resolved = true;
1310
1311        for query in queries {
1312            let QueryStatement::Promql(eval_stmt, _) = query.statement() else {
1313                unreachable!("query is parsed from promql");
1314            };
1315
1316            if let Some(query_targets) =
1317                self.prom_expr_permission_targets(&eval_stmt.expr, query_ctx)?
1318            {
1319                targets.extend(query_targets);
1320            } else {
1321                resolved = false;
1322            }
1323        }
1324
1325        Ok(if resolved {
1326            PermissionTableTargets::resolved(targets)
1327        } else {
1328            PermissionTableTargets::Unresolved
1329        })
1330    }
1331}
1332
1333#[async_trait]
1334impl PrometheusHandler for Instance {
1335    #[tracing::instrument(skip_all)]
1336    async fn do_query(
1337        &self,
1338        query: &PromQuery,
1339        query_ctx: QueryContextRef,
1340    ) -> server_error::Result<Output> {
1341        let query = ParsedPromQuery::parse(query.clone(), &query_ctx)?;
1342        self.do_query_parsed(query, query_ctx).await
1343    }
1344
1345    #[tracing::instrument(skip_all)]
1346    async fn do_query_parsed(
1347        &self,
1348        query: ParsedPromQuery,
1349        query_ctx: QueryContextRef,
1350    ) -> server_error::Result<Output> {
1351        let interceptor = self
1352            .plugins
1353            .get::<PromQueryInterceptorRef<server_error::Error>>();
1354
1355        self.check_prom_query_privilege(&query_ctx)?;
1356
1357        let targets =
1358            self.prom_queries_permission_targets(std::slice::from_ref(&query), &query_ctx)?;
1359        self.check_query_target_permission(targets, &query_ctx)
1360            .await?;
1361
1362        let requires_output_ordering = query.requires_output_ordering();
1363        let (query, stmt) = query.into_parts();
1364
1365        let QueryStatement::Promql(eval_stmt, _) = &stmt else {
1366            unreachable!("query is parsed from promql");
1367        };
1368
1369        let plan = self
1370            .statement_executor
1371            .plan(&stmt, query_ctx.clone())
1372            .await
1373            .map_err(BoxedError::new)
1374            .context(ExecuteQuerySnafu)?;
1375
1376        let plan = if requires_output_ordering {
1377            plan
1378        } else {
1379            promql::remove_output_sort(plan)
1380        };
1381
1382        interceptor.pre_execute(&query, &eval_stmt.expr, Some(&plan), query_ctx.clone())?;
1383
1384        // Take the EvalStmt from the original QueryStatement and use it to create the CatalogQueryStatement.
1385        let query_statement = if let QueryStatement::Promql(eval_stmt, alias) = stmt {
1386            CatalogQueryStatement::Promql(eval_stmt, alias)
1387        } else {
1388            // It should not happen since the query is already parsed successfully.
1389            return UnexpectedResultSnafu {
1390                reason: "The query should always be promql.".to_string(),
1391            }
1392            .fail();
1393        };
1394        let raw_query = query_statement.to_string();
1395
1396        let slow_query_timer = self.slow_query_options.enable.then(|| {
1397            SlowQueryTimer::new(
1398                query_statement,
1399                query_ctx.current_schema(),
1400                self.slow_query_options.threshold,
1401                self.slow_query_options.sample_ratio,
1402                self.slow_query_options.record_type,
1403                self.slow_query_recorder.clone(),
1404            )
1405        });
1406
1407        let ticket = self.process_manager.register_query(
1408            query_ctx.current_catalog().to_string(),
1409            vec![query_ctx.current_schema()],
1410            raw_query,
1411            query_ctx.conn_info().to_string(),
1412            Some(query_ctx.process_id()),
1413            slow_query_timer,
1414        );
1415
1416        let query_fut = self.statement_executor.exec_plan(plan, query_ctx.clone());
1417
1418        let output = CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
1419            .await
1420            .map_err(|_| servers::error::CancelledSnafu.build())?
1421            .map_err(BoxedError::new)
1422            .context(ExecuteQuerySnafu)?;
1423        let output = map_query_output(output)
1424            .map_err(BoxedError::new)
1425            .context(ExecuteQuerySnafu)?;
1426        let Output { meta, data } = output;
1427        let data = match data {
1428            OutputData::Stream(stream) => {
1429                OutputData::Stream(Box::pin(CancellableStreamWrapper::new(stream, ticket)))
1430            }
1431            other => other,
1432        };
1433        let output = Output { data, meta };
1434        Ok(interceptor.post_execute(output, query_ctx)?)
1435    }
1436
1437    async fn check_query_permission(
1438        &self,
1439        queries: &[PromQuery],
1440        query_ctx: &QueryContextRef,
1441    ) -> server_error::Result<()> {
1442        let queries = queries
1443            .iter()
1444            .cloned()
1445            .map(|query| ParsedPromQuery::parse(query, query_ctx))
1446            .collect::<server_error::Result<Vec<_>>>()?;
1447        self.check_query_permission_parsed(&queries, query_ctx)
1448            .await
1449    }
1450
1451    async fn check_query_permission_parsed(
1452        &self,
1453        queries: &[ParsedPromQuery],
1454        query_ctx: &QueryContextRef,
1455    ) -> server_error::Result<()> {
1456        self.check_prom_query_privilege(query_ctx)?;
1457        let targets = self.prom_queries_permission_targets(queries, query_ctx)?;
1458        self.check_query_target_permission(targets, query_ctx).await
1459    }
1460
1461    async fn check_query_target_permission(
1462        &self,
1463        targets: PermissionTableTargets,
1464        query_ctx: &QueryContextRef,
1465    ) -> server_error::Result<()> {
1466        let targets = self
1467            .resolve_query_permission_targets(targets, query_ctx)
1468            .await?;
1469        self.check_table_permission(query_ctx, PermissionReq::Action(PROMQL_QUERY), targets)
1470            .context(AuthSnafu)?;
1471        Ok(())
1472    }
1473
1474    async fn filter_metadata_metric_names(
1475        &self,
1476        metric_names: Vec<String>,
1477        schema: &str,
1478        query_ctx: &QueryContextRef,
1479    ) -> server_error::Result<Vec<String>> {
1480        let checker = self.plugins.get::<PermissionCheckerRef>();
1481        if !checker.as_ref().uses_table_targets() {
1482            let Some(metric) = metric_names.first() else {
1483                return Ok(metric_names);
1484            };
1485            let target =
1486                PermissionTableTarget::new(query_ctx.current_catalog(), schema, metric.as_str());
1487            let result = checker
1488                .as_ref()
1489                .check_permission_with_table_targets(
1490                    query_ctx.current_user(),
1491                    PermissionReq::Action(PROMQL_QUERY),
1492                    PermissionTableTargets::resolved(vec![target]),
1493                )
1494                .context(AuthSnafu);
1495            return match result {
1496                Ok(_) => Ok(metric_names),
1497                Err(error)
1498                    if error.status_code()
1499                        == common_error::status_code::StatusCode::PermissionDenied =>
1500                {
1501                    Ok(Vec::new())
1502                }
1503                Err(error) => Err(error),
1504            };
1505        }
1506
1507        let mut allowed = Vec::with_capacity(metric_names.len());
1508        for metric in metric_names {
1509            let target =
1510                PermissionTableTarget::new(query_ctx.current_catalog(), schema, metric.as_str());
1511            match checker
1512                .as_ref()
1513                .check_permission_with_table_targets(
1514                    query_ctx.current_user(),
1515                    PermissionReq::Action(PROMQL_QUERY),
1516                    PermissionTableTargets::resolved(vec![target]),
1517                )
1518                .context(AuthSnafu)
1519            {
1520                Ok(_) => allowed.push(metric),
1521                Err(error)
1522                    if error.status_code()
1523                        == common_error::status_code::StatusCode::PermissionDenied => {}
1524                Err(error) => return Err(error),
1525            }
1526        }
1527        Ok(allowed)
1528    }
1529
1530    async fn query_metric_names(
1531        &self,
1532        matchers: Vec<Matcher>,
1533        schema: &str,
1534        ctx: &QueryContextRef,
1535    ) -> server_error::Result<Vec<String>> {
1536        self.handle_query_metric_names(matchers, schema, ctx)
1537            .await
1538            .map_err(BoxedError::new)
1539            .context(ExecuteQuerySnafu)
1540    }
1541
1542    async fn query_metric_names_by_labels(
1543        &self,
1544        matchers: Vec<Matcher>,
1545        schema: &str,
1546        start: SystemTime,
1547        end: SystemTime,
1548        ctx: &QueryContextRef,
1549    ) -> server_error::Result<Vec<String>> {
1550        self.handle_query_metric_names_by_labels(matchers, schema, start, end, ctx)
1551            .await
1552            .map_err(BoxedError::new)
1553            .context(ExecuteQuerySnafu)
1554    }
1555
1556    async fn query_label_values(
1557        &self,
1558        metric: String,
1559        label_name: String,
1560        matchers: Vec<Matcher>,
1561        start: SystemTime,
1562        end: SystemTime,
1563        ctx: &QueryContextRef,
1564    ) -> server_error::Result<Vec<String>> {
1565        let schema =
1566            resolve_schema_from_matchers(&matchers)?.unwrap_or_else(|| ctx.current_schema());
1567        let target = PermissionTableTarget::new(ctx.current_catalog(), schema.as_str(), &metric);
1568        self.check_query_target_permission(
1569            PermissionTableTargets::resolved(vec![target.clone()]),
1570            ctx,
1571        )
1572        .await?;
1573
1574        self.handle_query_label_values(target, label_name, matchers, start, end, ctx)
1575            .await
1576            .map_err(BoxedError::new)
1577            .context(ExecuteQuerySnafu)
1578    }
1579
1580    fn catalog_manager(&self) -> CatalogManagerRef {
1581        self.catalog_manager.clone()
1582    }
1583}
1584
1585/// Validate `stmt.database` permission if it's presented.
1586macro_rules! validate_db_permission {
1587    ($stmt: expr, $query_ctx: expr) => {
1588        if let Some(database) = &$stmt.database {
1589            validate_catalog_and_schema($query_ctx.current_catalog(), database, $query_ctx)
1590                .map_err(BoxedError::new)
1591                .context(SqlExecInterceptedSnafu)?;
1592        }
1593    };
1594}
1595
1596pub fn check_permission(
1597    plugins: Plugins,
1598    stmt: &Statement,
1599    query_ctx: &QueryContextRef,
1600) -> Result<()> {
1601    let need_validate = plugins
1602        .get::<QueryOptions>()
1603        .map(|opts| opts.disallow_cross_catalog_query)
1604        .unwrap_or_default();
1605
1606    if !need_validate {
1607        return Ok(());
1608    }
1609
1610    match stmt {
1611        // Will be checked in execution.
1612        // TODO(dennis): add a hook for admin commands.
1613        Statement::Admin(_) => {}
1614        // These are executed by query engine, and will be checked there.
1615        Statement::Query(_)
1616        | Statement::Explain(_)
1617        | Statement::Tql(_)
1618        | Statement::Delete(_)
1619        | Statement::DeclareCursor(_)
1620        | Statement::Copy(sql::statements::copy::Copy::CopyQueryTo(_)) => {}
1621        // database ops won't be checked
1622        Statement::CreateDatabase(_)
1623        | Statement::ShowDatabases(_)
1624        | Statement::DropDatabase(_)
1625        | Statement::AlterDatabase(_)
1626        | Statement::DropFlow(_)
1627        | Statement::Use(_) => {}
1628        #[cfg(feature = "enterprise")]
1629        Statement::DropTrigger(_) => {}
1630        Statement::ShowCreateDatabase(stmt) => {
1631            validate_database(&stmt.database_name, query_ctx)?;
1632        }
1633        Statement::ShowCreateTable(stmt) => {
1634            validate_param(&stmt.table_name, query_ctx)?;
1635        }
1636        Statement::ShowCreateFlow(stmt) => {
1637            validate_flow(&stmt.flow_name, query_ctx)?;
1638        }
1639        #[cfg(feature = "enterprise")]
1640        Statement::ShowCreateTrigger(stmt) => {
1641            validate_param(&stmt.trigger_name, query_ctx)?;
1642        }
1643        Statement::ShowCreateView(stmt) => {
1644            validate_param(&stmt.view_name, query_ctx)?;
1645        }
1646        Statement::CreateExternalTable(stmt) => {
1647            validate_param(&stmt.name, query_ctx)?;
1648        }
1649        Statement::CreateFlow(stmt) => {
1650            // TODO: should also validate source table name here?
1651            validate_param(&stmt.sink_table_name, query_ctx)?;
1652        }
1653        #[cfg(feature = "enterprise")]
1654        Statement::CreateTrigger(stmt) => {
1655            validate_param(&stmt.trigger_name, query_ctx)?;
1656        }
1657        Statement::CreateView(stmt) => {
1658            validate_param(&stmt.name, query_ctx)?;
1659        }
1660        Statement::AlterTable(stmt) => {
1661            validate_param(stmt.table_name(), query_ctx)?;
1662        }
1663        #[cfg(feature = "enterprise")]
1664        Statement::AlterTrigger(_) => {}
1665        // set/show variable now only alter/show variable in session
1666        Statement::SetVariables(_) | Statement::ShowVariables(_) => {}
1667        // show charset and show collation won't be checked
1668        Statement::ShowCharset(_) | Statement::ShowCollation(_) => {}
1669
1670        Statement::Comment(comment) => match &comment.object {
1671            CommentObject::Table(table) => validate_param(table, query_ctx)?,
1672            CommentObject::Column { table, .. } => validate_param(table, query_ctx)?,
1673            CommentObject::Flow(flow) => validate_flow(flow, query_ctx)?,
1674        },
1675
1676        Statement::Insert(insert) => {
1677            let name = insert.table_name().context(ParseSqlSnafu)?;
1678            validate_param(name, query_ctx)?;
1679        }
1680        Statement::CreateTable(stmt) => {
1681            validate_param(&stmt.name, query_ctx)?;
1682        }
1683        Statement::CreateTableLike(stmt) => {
1684            validate_param(&stmt.table_name, query_ctx)?;
1685            validate_param(&stmt.source_name, query_ctx)?;
1686        }
1687        Statement::DropTable(drop_stmt) => {
1688            for table_name in drop_stmt.table_names() {
1689                validate_param(table_name, query_ctx)?;
1690            }
1691        }
1692        #[cfg(feature = "enterprise")]
1693        Statement::UndropTable(stmt) => {
1694            validate_param(stmt.table_name(), query_ctx)?;
1695        }
1696        Statement::DropView(stmt) => {
1697            validate_param(&stmt.view_name, query_ctx)?;
1698        }
1699        Statement::ShowTables(stmt) => {
1700            validate_db_permission!(stmt, query_ctx);
1701        }
1702        Statement::ShowTableStatus(stmt) => {
1703            validate_db_permission!(stmt, query_ctx);
1704        }
1705        Statement::ShowColumns(stmt) => {
1706            validate_db_permission!(stmt, query_ctx);
1707        }
1708        Statement::ShowIndex(stmt) => {
1709            validate_db_permission!(stmt, query_ctx);
1710        }
1711        Statement::ShowRegion(stmt) => {
1712            validate_db_permission!(stmt, query_ctx);
1713        }
1714        Statement::ShowViews(stmt) => {
1715            validate_db_permission!(stmt, query_ctx);
1716        }
1717        Statement::ShowFlows(stmt) => {
1718            validate_db_permission!(stmt, query_ctx);
1719        }
1720        Statement::ShowFlowStatus(_stmt) => {
1721            // Flow statistics are organized based on the catalog dimension and
1722            // filtered by the current catalog, so there is no need to check the
1723            // permission of the database(schema).
1724        }
1725        #[cfg(feature = "enterprise")]
1726        Statement::ShowTriggers(_stmt) => {
1727            // The trigger is organized based on the catalog dimension, so there
1728            // is no need to check the permission of the database(schema).
1729        }
1730        Statement::ShowStatus(_stmt) => {}
1731        Statement::ShowSearchPath(_stmt) => {}
1732        Statement::DescribeTable(stmt) => {
1733            validate_param(stmt.name(), query_ctx)?;
1734        }
1735        Statement::Copy(sql::statements::copy::Copy::CopyTable(stmt)) => match stmt {
1736            CopyTable::To(copy_table_to) => validate_param(&copy_table_to.table_name, query_ctx)?,
1737            CopyTable::From(copy_table_from) => {
1738                validate_param(&copy_table_from.table_name, query_ctx)?
1739            }
1740        },
1741        Statement::Copy(sql::statements::copy::Copy::CopyDatabase(copy_database)) => {
1742            match copy_database {
1743                CopyDatabase::To(stmt) => validate_database(&stmt.database_name, query_ctx)?,
1744                CopyDatabase::From(stmt) => validate_database(&stmt.database_name, query_ctx)?,
1745            }
1746        }
1747        Statement::TruncateTable(stmt) => {
1748            validate_param(stmt.table_name(), query_ctx)?;
1749        }
1750        // cursor operations are always allowed once it's created
1751        Statement::FetchCursor(_) | Statement::CloseCursor(_) => {}
1752        // User can only kill process in their own catalog.
1753        Statement::Kill(_) => {}
1754        // SHOW PROCESSLIST
1755        Statement::ShowProcesslist(_) => {}
1756    }
1757    Ok(())
1758}
1759
1760fn validate_param(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
1761    let (catalog, schema, _) = table_idents_to_full_name(name, query_ctx)
1762        .map_err(BoxedError::new)
1763        .context(ExternalSnafu)?;
1764
1765    validate_catalog_and_schema(&catalog, &schema, query_ctx)
1766        .map_err(BoxedError::new)
1767        .context(SqlExecInterceptedSnafu)
1768}
1769
1770fn validate_flow(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
1771    let catalog = match &name.0[..] {
1772        [_flow] => query_ctx.current_catalog().to_string(),
1773        [catalog, _flow] => catalog.to_string_unquoted(),
1774        _ => {
1775            return InvalidSqlSnafu {
1776                err_msg: format!(
1777                    "expect flow name to be <catalog>.<flow_name> or <flow_name>, actual: {name}",
1778                ),
1779            }
1780            .fail();
1781        }
1782    };
1783
1784    let schema = query_ctx.current_schema();
1785
1786    validate_catalog_and_schema(&catalog, &schema, query_ctx)
1787        .map_err(BoxedError::new)
1788        .context(SqlExecInterceptedSnafu)
1789}
1790
1791fn validate_database(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
1792    let (catalog, schema) = match &name.0[..] {
1793        [schema] => (
1794            query_ctx.current_catalog().to_string(),
1795            schema.to_string_unquoted(),
1796        ),
1797        [catalog, schema] => (catalog.to_string_unquoted(), schema.to_string_unquoted()),
1798        _ => InvalidSqlSnafu {
1799            err_msg: format!(
1800                "expect database name to be <catalog>.<schema> or <schema>, actual: {name}",
1801            ),
1802        }
1803        .fail()?,
1804    };
1805
1806    validate_catalog_and_schema(&catalog, &schema, query_ctx)
1807        .map_err(BoxedError::new)
1808        .context(SqlExecInterceptedSnafu)
1809}
1810
1811fn is_readonly_plan(plan: &LogicalPlan) -> bool {
1812    !matches!(plan, LogicalPlan::Dml(_) | LogicalPlan::Ddl(_))
1813}
1814
1815fn should_track_statement_process(stmt: &Statement) -> bool {
1816    stmt.is_readonly()
1817        || matches!(stmt, Statement::Insert(insert) if insert.has_non_values_query_source())
1818}
1819
1820fn should_track_plan_process(stmt: Option<&Statement>, plan: &LogicalPlan) -> bool {
1821    is_readonly_plan(plan)
1822        || matches!(stmt, Some(Statement::Insert(insert)) if insert.has_non_values_query_source())
1823}
1824
1825#[cfg(test)]
1826mod tests {
1827    use std::any::Any;
1828    use std::collections::HashMap;
1829    use std::future::Future;
1830    use std::pin::Pin;
1831    use std::sync::Arc;
1832    use std::task::{Context, Poll};
1833    use std::time::Duration;
1834
1835    use api::prom_store::remote::label_matcher::Type as PromMatcherType;
1836    use api::prom_store::remote::{
1837        Label, LabelMatcher, Query as RemoteQuery, ReadRequest, ReadResponse, Sample,
1838    };
1839    use api::v1::greptime_request::Request;
1840    use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse};
1841    use api::v1::query_request::Query;
1842    use auth::{
1843        DASHBOARD_DELETE, DASHBOARD_QUERY, DASHBOARD_SAVE, JAEGER_QUERY, PIPELINE_DELETE,
1844        PIPELINE_INSERT, PIPELINE_QUERY, PermissionAction, PermissionResp, UserInfo, UserInfoRef,
1845    };
1846    use catalog::process_manager::{ProcessManager, QueryStatement, SlowQueryTimer};
1847    use common_base::Plugins;
1848    use common_catalog::consts::DEFAULT_PRIVATE_SCHEMA_NAME;
1849    use common_error::ext::{BoxedError, ErrorExt, PlainError};
1850    use common_error::status_code::StatusCode;
1851    use common_event_recorder::{Event, EventRecorder, EventTypeFilter, EventTypeFilterRef};
1852    use common_frontend::slow_query_event::SlowQueryEvent;
1853    use common_meta::cache::LayeredCacheRegistryBuilder;
1854    use common_meta::kv_backend::memory::MemoryKvBackend;
1855    use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor};
1856    use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse};
1857    use common_meta::rpc::procedure::{
1858        MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse,
1859    };
1860    use common_query::prelude::greptime_value;
1861    use common_query::{Output, OutputMeta};
1862    use common_recordbatch::{
1863        OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream,
1864    };
1865    use common_telemetry::logging::SlowQueriesRecordType;
1866    use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
1867    use datafusion::physical_plan::empty::EmptyExec;
1868    use datafusion_expr::dml::InsertOp;
1869    use datafusion_expr::{LogicalPlanBuilder, LogicalTableSource};
1870    use datatypes::prelude::ConcreteDataType;
1871    use datatypes::schema::{ColumnSchema, Schema as GtSchema, SchemaRef as GtSchemaRef};
1872    use datatypes::vectors::{
1873        Float64Vector, StringVector, TimestampMillisecondVector, TimestampNanosecondVector,
1874        VectorRef,
1875    };
1876    use log_query::LogQuery;
1877    use prost::Message;
1878    use query::query_engine::options::QueryOptions;
1879    use servers::query_handler::{
1880        DashboardHandler, JaegerQueryHandler, LogQueryHandler, PipelineHandler, PipelineHandlerRef,
1881        PromStoreProtocolHandler,
1882    };
1883    use session::context::{Channel, ConnInfo, QueryContext, QueryContextBuilder};
1884    use snafu::{Location, Snafu};
1885    use sql::dialect::GreptimeDbDialect;
1886    use store_api::data_source::DataSource;
1887    use store_api::metric_engine_consts::{
1888        LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY,
1889    };
1890    use store_api::storage::ScanRequest;
1891    use strfmt::Format;
1892    use table::metadata::{
1893        FilterPushDownType, TableInfo, TableInfoBuilder, TableMetaBuilder, TableType,
1894    };
1895    use table::table_name::TableName;
1896    use table::test_util::{EmptyTable, MemTable};
1897    use table::{Table, TableRef};
1898    use tokio::sync::{mpsc, oneshot};
1899    use tower::ServiceExt;
1900
1901    use crate::frontend::FrontendOptions;
1902    use crate::instance::builder::FrontendBuilder;
1903    use crate::instance::*;
1904
1905    fn parse_test_sql(sql: &str) -> Vec<Statement> {
1906        parse_stmt(sql, &GreptimeDbDialect {}).unwrap()
1907    }
1908
1909    #[derive(Debug, Default)]
1910    struct RecordingSlowQueryEventRecorder {
1911        payloads: std::sync::Mutex<Vec<serde_json::Value>>,
1912    }
1913
1914    impl EventRecorder for RecordingSlowQueryEventRecorder {
1915        fn record(&self, event: Box<dyn Event>) {
1916            let event = event
1917                .as_any()
1918                .downcast_ref::<SlowQueryEvent>()
1919                .expect("expected a slow query event");
1920            self.payloads.lock().unwrap().push(event.payload.clone());
1921        }
1922
1923        fn event_type_filter(&self) -> EventTypeFilterRef {
1924            Arc::new(EventTypeFilter::All)
1925        }
1926
1927        fn close(&self) {}
1928    }
1929
1930    #[test]
1931    fn test_validate_analyze_stream_statement_strictness() {
1932        for sql in [
1933            "select 1",
1934            "explain analyze select 1",
1935            "explain analyze verbose format text select 1",
1936            "explain analyze verbose format graphviz select 1",
1937            "TQL ANALYZE (0, 10, '5s') physical_metric",
1938            "TQL EXPLAIN VERBOSE (0, 10, '5s') physical_metric",
1939            "TQL ANALYZE VERBOSE FORMAT TEXT (0, 10, '5s') physical_metric",
1940        ] {
1941            let mut stmts = parse_test_sql(sql);
1942            assert!(
1943                validate_analyze_stream_statement(&mut stmts[0]).is_err(),
1944                "{sql}"
1945            );
1946        }
1947
1948        for sql in [
1949            "explain analyze verbose select 1",
1950            "explain analyze verbose format json select 1",
1951            "TQL ANALYZE VERBOSE (0, 10, '5s') physical_metric",
1952            "TQL ANALYZE VERBOSE FORMAT JSON (0, 10, '5s') physical_metric",
1953        ] {
1954            let mut stmts = parse_test_sql(sql);
1955            assert!(
1956                validate_analyze_stream_statement(&mut stmts[0]).is_ok(),
1957                "{sql}"
1958            );
1959            match &stmts[0] {
1960                Statement::Explain(explain) => assert!(explain.format.is_none()),
1961                Statement::Tql(Tql::Analyze(analyze)) => assert!(analyze.format.is_none()),
1962                _ => unreachable!(),
1963            }
1964        }
1965
1966        assert_eq!(
1967            parse_test_sql("explain analyze verbose select 1; select 2").len(),
1968            2
1969        );
1970
1971        assert!(is_explain_analyze_verbose(
1972            &parse_test_sql("explain analyze verbose select 1")[0]
1973        ));
1974        assert!(is_explain_analyze_verbose(
1975            &parse_test_sql("TQL ANALYZE VERBOSE (0, 10, '5s') physical_metric")[0]
1976        ));
1977        for sql in [
1978            "select 1",
1979            "explain select 1",
1980            "explain analyze select 1",
1981            "explain verbose select 1",
1982            "TQL ANALYZE (0, 10, '5s') physical_metric",
1983            "TQL EXPLAIN VERBOSE (0, 10, '5s') physical_metric",
1984        ] {
1985            assert!(
1986                !is_explain_analyze_verbose(&parse_test_sql(sql)[0]),
1987                "{sql}"
1988            );
1989        }
1990    }
1991
1992    #[derive(Debug, Snafu)]
1993    enum TestError {
1994        #[snafu(display("Failed to build test cache registry"))]
1995        BuildCacheRegistry {
1996            source: cache::error::Error,
1997            #[snafu(implicit)]
1998            location: Location,
1999        },
2000
2001        #[snafu(display("Failed to build test table meta for table: {table_name}"))]
2002        BuildTableMeta {
2003            table_name: String,
2004            source: table::metadata::TableMetaBuilderError,
2005            #[snafu(implicit)]
2006            location: Location,
2007        },
2008
2009        #[snafu(display("Failed to build test table info for table: {table_name}"))]
2010        BuildTableInfo {
2011            table_name: String,
2012            source: table::metadata::TableInfoBuilderError,
2013            #[snafu(implicit)]
2014            location: Location,
2015        },
2016
2017        #[snafu(display("Failed to register test table: {table_name}"))]
2018        RegisterTable {
2019            table_name: String,
2020            source: catalog::error::Error,
2021            #[snafu(implicit)]
2022            location: Location,
2023        },
2024
2025        #[snafu(display("Failed to build test frontend instance"))]
2026        BuildFrontend {
2027            source: crate::error::Error,
2028            #[snafu(implicit)]
2029            location: Location,
2030        },
2031
2032        #[snafu(display("Expected exactly one output for SQL `{sql}`, got {actual}"))]
2033        UnexpectedOutputCount {
2034            sql: String,
2035            actual: usize,
2036            #[snafu(implicit)]
2037            location: Location,
2038        },
2039
2040        #[snafu(display("Failed to execute SQL `{sql}`"))]
2041        ExecuteSql {
2042            sql: String,
2043            source: crate::error::Error,
2044            #[snafu(implicit)]
2045            location: Location,
2046        },
2047
2048        #[snafu(display("Timed out waiting for insert-select start notification"))]
2049        InsertStartTimeout {
2050            source: tokio::time::error::Elapsed,
2051            #[snafu(implicit)]
2052            location: Location,
2053        },
2054
2055        #[snafu(display("Insert-select start notification channel closed"))]
2056        InsertStartChannelClosed {
2057            #[snafu(implicit)]
2058            location: Location,
2059        },
2060
2061        #[snafu(display("Failed to release blocking insert-select interceptor"))]
2062        ReleaseBlockedInsert {
2063            #[snafu(implicit)]
2064            location: Location,
2065        },
2066
2067        #[snafu(display("Timed out waiting for insert-select source to be polled"))]
2068        SourcePollTimeout {
2069            source: tokio::time::error::Elapsed,
2070            #[snafu(implicit)]
2071            location: Location,
2072        },
2073
2074        #[snafu(display("Insert-select source poll notification channel closed"))]
2075        SourcePollChannelClosed {
2076            source: oneshot::error::RecvError,
2077            #[snafu(implicit)]
2078            location: Location,
2079        },
2080
2081        #[snafu(display("Timed out waiting for insert task to finish"))]
2082        InsertTaskTimeout {
2083            source: tokio::time::error::Elapsed,
2084            #[snafu(implicit)]
2085            location: Location,
2086        },
2087
2088        #[snafu(display("Insert task panicked"))]
2089        InsertTaskPanic {
2090            source: tokio::task::JoinError,
2091            #[snafu(implicit)]
2092            location: Location,
2093        },
2094
2095        #[snafu(display("Expected insert-select to be cancelled"))]
2096        InsertSelectNotCancelled {
2097            #[snafu(implicit)]
2098            location: Location,
2099        },
2100    }
2101
2102    type TestResult<T> = std::result::Result<T, TestError>;
2103
2104    fn parse_one_sql(sql: &str) -> Statement {
2105        parse_stmt(sql, &GreptimeDbDialect {}).unwrap().remove(0)
2106    }
2107
2108    fn test_query_ctx(process_id: u32) -> QueryContextRef {
2109        Arc::new(
2110            QueryContextBuilder::default()
2111                .channel(Channel::Mysql)
2112                .conn_info(ConnInfo::new(None, Channel::Mysql))
2113                .process_id(process_id)
2114                .build(),
2115        )
2116    }
2117
2118    #[derive(Debug)]
2119    struct AdminUserInfo;
2120
2121    impl UserInfo for AdminUserInfo {
2122        fn as_any(&self) -> &dyn Any {
2123            self
2124        }
2125
2126        fn username(&self) -> &str {
2127            "admin"
2128        }
2129
2130        fn is_admin(&self) -> bool {
2131            true
2132        }
2133    }
2134
2135    struct RejectUnresolvedPermissionChecker;
2136
2137    impl PermissionChecker for RejectUnresolvedPermissionChecker {
2138        fn check_permission(
2139            &self,
2140            _user_info: UserInfoRef,
2141            _req: PermissionReq,
2142        ) -> auth::error::Result<PermissionResp> {
2143            Ok(PermissionResp::Allow)
2144        }
2145
2146        fn check_permission_with_table_targets(
2147            &self,
2148            _user_info: UserInfoRef,
2149            _req: PermissionReq,
2150            targets: PermissionTableTargets,
2151        ) -> auth::error::Result<PermissionResp> {
2152            let reject = match targets {
2153                PermissionTableTargets::Unresolved => true,
2154                PermissionTableTargets::Resolved(targets) => {
2155                    targets.iter().any(|target| target.table == "denied")
2156                }
2157            };
2158            Ok(if reject {
2159                PermissionResp::Reject
2160            } else {
2161                PermissionResp::Allow
2162            })
2163        }
2164    }
2165
2166    #[derive(Debug, PartialEq, Eq)]
2167    struct CheckedAction {
2168        action: PermissionAction,
2169        targets: Option<PermissionTableTargets>,
2170    }
2171
2172    #[derive(Default)]
2173    struct RejectEndpointPermissionChecker {
2174        checks: std::sync::Mutex<Vec<CheckedAction>>,
2175    }
2176
2177    impl RejectEndpointPermissionChecker {
2178        fn reject(
2179            &self,
2180            action: PermissionAction,
2181            targets: Option<PermissionTableTargets>,
2182        ) -> PermissionResp {
2183            self.checks
2184                .lock()
2185                .unwrap()
2186                .push(CheckedAction { action, targets });
2187            PermissionResp::Reject
2188        }
2189
2190        fn take_check(&self) -> CheckedAction {
2191            let mut checks = self.checks.lock().unwrap();
2192            assert_eq!(1, checks.len());
2193            checks.pop().unwrap()
2194        }
2195    }
2196
2197    impl PermissionChecker for RejectEndpointPermissionChecker {
2198        fn check_permission(
2199            &self,
2200            _user_info: UserInfoRef,
2201            req: PermissionReq,
2202        ) -> auth::error::Result<PermissionResp> {
2203            Ok(match req {
2204                PermissionReq::Action(action) => self.reject(action, None),
2205                _ => PermissionResp::Allow,
2206            })
2207        }
2208
2209        fn check_permission_with_table_targets(
2210            &self,
2211            _user_info: UserInfoRef,
2212            req: PermissionReq,
2213            targets: PermissionTableTargets,
2214        ) -> auth::error::Result<PermissionResp> {
2215            Ok(match req {
2216                PermissionReq::Action(action) => self.reject(action, Some(targets)),
2217                _ => PermissionResp::Allow,
2218            })
2219        }
2220    }
2221
2222    struct WriteOnlyPermissionChecker;
2223
2224    impl PermissionChecker for WriteOnlyPermissionChecker {
2225        fn check_permission(
2226            &self,
2227            _user_info: UserInfoRef,
2228            req: PermissionReq,
2229        ) -> auth::error::Result<PermissionResp> {
2230            Ok(if req.is_readonly() {
2231                PermissionResp::Reject
2232            } else {
2233                PermissionResp::Allow
2234            })
2235        }
2236
2237        fn check_permission_with_table_targets(
2238            &self,
2239            user_info: UserInfoRef,
2240            req: PermissionReq,
2241            _targets: PermissionTableTargets,
2242        ) -> auth::error::Result<PermissionResp> {
2243            self.check_permission(user_info, req)
2244        }
2245    }
2246
2247    #[derive(Default)]
2248    struct TargetIndependentPermissionChecker {
2249        checks: atomic::AtomicUsize,
2250    }
2251
2252    impl PermissionChecker for TargetIndependentPermissionChecker {
2253        fn check_permission(
2254            &self,
2255            _user_info: UserInfoRef,
2256            _req: PermissionReq,
2257        ) -> auth::error::Result<PermissionResp> {
2258            self.checks.fetch_add(1, atomic::Ordering::Relaxed);
2259            Ok(PermissionResp::Allow)
2260        }
2261
2262        fn uses_table_targets(&self) -> bool {
2263            false
2264        }
2265
2266        fn check_permission_with_table_targets(
2267            &self,
2268            user_info: UserInfoRef,
2269            req: PermissionReq,
2270            _targets: PermissionTableTargets,
2271        ) -> auth::error::Result<PermissionResp> {
2272            self.check_permission(user_info, req)
2273        }
2274    }
2275
2276    struct BlockingInsertSelectInterceptor {
2277        started_tx: mpsc::UnboundedSender<()>,
2278        finish_rx: std::sync::Mutex<Option<oneshot::Receiver<()>>>,
2279    }
2280
2281    impl BlockingInsertSelectInterceptor {
2282        fn new(started_tx: mpsc::UnboundedSender<()>, finish_rx: oneshot::Receiver<()>) -> Self {
2283            Self {
2284                started_tx,
2285                finish_rx: std::sync::Mutex::new(Some(finish_rx)),
2286            }
2287        }
2288    }
2289
2290    impl SqlQueryInterceptor for BlockingInsertSelectInterceptor {
2291        type Error = Error;
2292
2293        fn pre_execute(
2294            &self,
2295            statement: Option<&Statement>,
2296            _plan: Option<&LogicalPlan>,
2297            _query_ctx: QueryContextRef,
2298        ) -> Result<()> {
2299            let Some(Statement::Insert(insert)) = statement else {
2300                return Ok(());
2301            };
2302            if !insert.has_non_values_query_source() {
2303                return Ok(());
2304            }
2305
2306            let finish_rx = self.finish_rx.lock().unwrap().take().unwrap();
2307            let _ = self.started_tx.send(());
2308            tokio::task::block_in_place(|| {
2309                tokio::runtime::Handle::current()
2310                    .block_on(finish_rx)
2311                    .unwrap();
2312            });
2313            Ok(())
2314        }
2315    }
2316
2317    struct PendingRecordBatchStream {
2318        schema: GtSchemaRef,
2319        polled_tx: Option<oneshot::Sender<()>>,
2320        _finish_tx: oneshot::Sender<()>,
2321        finish_rx: Pin<Box<oneshot::Receiver<()>>>,
2322    }
2323
2324    impl RecordBatchStream for PendingRecordBatchStream {
2325        fn schema(&self) -> GtSchemaRef {
2326            self.schema.clone()
2327        }
2328
2329        fn output_ordering(&self) -> Option<&[OrderOption]> {
2330            None
2331        }
2332
2333        fn metrics(&self) -> Option<common_recordbatch::adapter::RecordBatchMetrics> {
2334            None
2335        }
2336    }
2337
2338    impl Stream for PendingRecordBatchStream {
2339        type Item = common_recordbatch::error::Result<RecordBatch>;
2340
2341        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
2342            if let Some(polled_tx) = self.polled_tx.take() {
2343                let _ = polled_tx.send(());
2344            }
2345
2346            match self.finish_rx.as_mut().poll(cx) {
2347                Poll::Ready(_) => Poll::Ready(None),
2348                Poll::Pending => Poll::Pending,
2349            }
2350        }
2351    }
2352
2353    impl Unpin for PendingRecordBatchStream {}
2354
2355    #[test]
2356    fn test_record_explain_analyze_timeout_uses_empty_metrics_without_plan() {
2357        let event_recorder = Arc::new(RecordingSlowQueryEventRecorder::default());
2358        let timer = SlowQueryTimer::new(
2359            QueryStatement::Plan("EXPLAIN ANALYZE VERBOSE SELECT 1".to_string()),
2360            "public".to_string(),
2361            Duration::from_secs(3600),
2362            0.0,
2363            SlowQueriesRecordType::SystemTable,
2364            event_recorder.clone(),
2365        );
2366        let timeout_recorder = timer.recorder();
2367
2368        record_explain_analyze_timeout(Some(&timeout_recorder), None);
2369        drop(timer);
2370
2371        let payloads = event_recorder.payloads.lock().unwrap();
2372        assert_eq!(payloads.len(), 1);
2373        assert_eq!(payloads[0]["timed_out"], true);
2374        assert_eq!(payloads[0]["metrics"], serde_json::json!([]));
2375    }
2376
2377    #[tokio::test]
2378    async fn test_attach_timeout_records_explain_analyze_metrics() {
2379        let event_recorder = Arc::new(RecordingSlowQueryEventRecorder::default());
2380        let timer = SlowQueryTimer::new(
2381            QueryStatement::Plan("EXPLAIN ANALYZE VERBOSE SELECT 1".to_string()),
2382            "public".to_string(),
2383            Duration::from_secs(3600),
2384            0.0,
2385            SlowQueriesRecordType::SystemTable,
2386            event_recorder.clone(),
2387        );
2388        let timeout_recorder = timer.recorder();
2389        let plan: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(Arc::new(Schema::empty())));
2390        let (finish_tx, finish_rx) = oneshot::channel();
2391        let stream = PendingRecordBatchStream {
2392            schema: Arc::new(GtSchema::new(vec![])),
2393            polled_tx: None,
2394            _finish_tx: finish_tx,
2395            finish_rx: Box::pin(finish_rx),
2396        };
2397        let output = Output::new(
2398            OutputData::Stream(Box::pin(stream)),
2399            OutputMeta::new_with_plan(plan),
2400        );
2401        let output =
2402            attach_timeout(output, Duration::from_millis(10), Some(timeout_recorder)).unwrap();
2403        let OutputData::Stream(mut stream) = output.data else {
2404            unreachable!();
2405        };
2406
2407        let err = stream.next().await.unwrap().unwrap_err();
2408        assert_eq!(err.to_string(), "Stream timeout");
2409        drop(stream);
2410        drop(timer);
2411
2412        let payloads = event_recorder.payloads.lock().unwrap();
2413        assert_eq!(payloads.len(), 1);
2414        assert_eq!(payloads[0]["timed_out"], true);
2415        assert!(
2416            payloads[0]["metrics"]
2417                .as_array()
2418                .is_some_and(|metrics| !metrics.is_empty())
2419        );
2420    }
2421
2422    struct PendingDataSource {
2423        schema: GtSchemaRef,
2424        polled_tx: std::sync::Mutex<Option<oneshot::Sender<()>>>,
2425    }
2426
2427    impl DataSource for PendingDataSource {
2428        fn get_stream(
2429            &self,
2430            _request: ScanRequest,
2431        ) -> std::result::Result<SendableRecordBatchStream, BoxedError> {
2432            let (finish_tx, finish_rx) = oneshot::channel();
2433            let mut polled_tx = self.polled_tx.lock().map_err(|_| {
2434                BoxedError::new(PlainError::new(
2435                    "pending data source lock poisoned".to_string(),
2436                    StatusCode::Unexpected,
2437                ))
2438            })?;
2439            Ok(Box::pin(PendingRecordBatchStream {
2440                schema: self.schema.clone(),
2441                polled_tx: polled_tx.take(),
2442                _finish_tx: finish_tx,
2443                finish_rx: Box::pin(finish_rx),
2444            }))
2445        }
2446    }
2447
2448    struct NoopProcedureExecutor;
2449
2450    #[async_trait::async_trait]
2451    impl ProcedureExecutor for NoopProcedureExecutor {
2452        async fn submit_ddl_task(
2453            &self,
2454            _ctx: ExecutorContext,
2455            _request: SubmitDdlTaskRequest,
2456        ) -> common_meta::error::Result<SubmitDdlTaskResponse> {
2457            common_meta::error::UnsupportedSnafu {
2458                operation: "submit_ddl_task",
2459            }
2460            .fail()
2461        }
2462
2463        async fn migrate_region(
2464            &self,
2465            _ctx: &ExecutorContext,
2466            _request: MigrateRegionRequest,
2467        ) -> common_meta::error::Result<MigrateRegionResponse> {
2468            common_meta::error::UnsupportedSnafu {
2469                operation: "migrate_region",
2470            }
2471            .fail()
2472        }
2473
2474        async fn reconcile(
2475            &self,
2476            _ctx: &ExecutorContext,
2477            _request: ReconcileRequest,
2478        ) -> common_meta::error::Result<ReconcileResponse> {
2479            common_meta::error::UnsupportedSnafu {
2480                operation: "reconcile",
2481            }
2482            .fail()
2483        }
2484
2485        async fn query_procedure_state(
2486            &self,
2487            _ctx: &ExecutorContext,
2488            _pid: &str,
2489        ) -> common_meta::error::Result<ProcedureStateResponse> {
2490            common_meta::error::UnsupportedSnafu {
2491                operation: "query_procedure_state",
2492            }
2493            .fail()
2494        }
2495
2496        async fn list_procedures(
2497            &self,
2498            _ctx: &ExecutorContext,
2499        ) -> common_meta::error::Result<ProcedureDetailResponse> {
2500            common_meta::error::UnsupportedSnafu {
2501                operation: "list_procedures",
2502            }
2503            .fail()
2504        }
2505    }
2506
2507    /// A test [`ProcedureExecutor`] that completes create/drop DDL tasks against the
2508    /// in-memory catalog, mimicking what the meta DDL procedures do in production.
2509    /// This allows happy-path DDL requests (create/drop table/view) to be exercised
2510    /// end to end through the gRPC ingress.
2511    struct MockProcedureExecutor {
2512        catalog_manager: Arc<catalog::memory::MemoryCatalogManager>,
2513        next_table_id: std::sync::atomic::AtomicU32,
2514        submitted: std::sync::Mutex<Vec<DdlTask>>,
2515    }
2516
2517    impl MockProcedureExecutor {
2518        fn new(catalog_manager: Arc<catalog::memory::MemoryCatalogManager>) -> Self {
2519            Self {
2520                catalog_manager,
2521                next_table_id: std::sync::atomic::AtomicU32::new(1026),
2522                submitted: std::sync::Mutex::new(Vec::new()),
2523            }
2524        }
2525    }
2526
2527    #[async_trait::async_trait]
2528    impl ProcedureExecutor for MockProcedureExecutor {
2529        async fn submit_ddl_task(
2530            &self,
2531            _ctx: ExecutorContext,
2532            request: SubmitDdlTaskRequest,
2533        ) -> common_meta::error::Result<SubmitDdlTaskResponse> {
2534            self.submitted.lock().unwrap().push(request.task.clone());
2535            match request.task {
2536                DdlTask::CreateTable(task) => {
2537                    let table_id = self
2538                        .next_table_id
2539                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2540                    let mut table_info = task.table_info;
2541                    table_info.ident.table_id = table_id;
2542                    self.catalog_manager
2543                        .register_table_sync(catalog::RegisterTableRequest {
2544                            catalog: table_info.catalog_name.clone(),
2545                            schema: table_info.schema_name.clone(),
2546                            table_name: table_info.name.clone(),
2547                            table_id,
2548                            table: table::dist_table::DistTable::table(Arc::new(table_info)),
2549                        })
2550                        .map_err(BoxedError::new)
2551                        .context(common_meta::error::ExternalSnafu)?;
2552                    Ok(SubmitDdlTaskResponse {
2553                        key: Vec::new(),
2554                        table_ids: vec![table_id],
2555                    })
2556                }
2557                DdlTask::CreateView(task) => {
2558                    let view_id = self
2559                        .next_table_id
2560                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2561                    let mut view_info = task.view_info;
2562                    view_info.ident.table_id = view_id;
2563                    self.catalog_manager
2564                        .register_table_sync(catalog::RegisterTableRequest {
2565                            catalog: task.create_view.catalog_name.clone(),
2566                            schema: task.create_view.schema_name.clone(),
2567                            table_name: task.create_view.view_name.clone(),
2568                            table_id: view_id,
2569                            table: table::dist_table::DistTable::table(Arc::new(view_info)),
2570                        })
2571                        .map_err(BoxedError::new)
2572                        .context(common_meta::error::ExternalSnafu)?;
2573                    Ok(SubmitDdlTaskResponse {
2574                        key: Vec::new(),
2575                        table_ids: vec![view_id],
2576                    })
2577                }
2578                DdlTask::DropView(task) => {
2579                    self.catalog_manager
2580                        .deregister_table_sync(catalog::DeregisterTableRequest {
2581                            catalog: task.catalog.clone(),
2582                            schema: task.schema.clone(),
2583                            table_name: task.view.clone(),
2584                        })
2585                        .map_err(BoxedError::new)
2586                        .context(common_meta::error::ExternalSnafu)?;
2587                    Ok(SubmitDdlTaskResponse::default())
2588                }
2589                other => common_meta::error::UnsupportedSnafu {
2590                    operation: format!("mock submit_ddl_task: {other:?}"),
2591                }
2592                .fail(),
2593            }
2594        }
2595
2596        async fn migrate_region(
2597            &self,
2598            _ctx: &ExecutorContext,
2599            _request: MigrateRegionRequest,
2600        ) -> common_meta::error::Result<MigrateRegionResponse> {
2601            common_meta::error::UnsupportedSnafu {
2602                operation: "migrate_region",
2603            }
2604            .fail()
2605        }
2606
2607        async fn reconcile(
2608            &self,
2609            _ctx: &ExecutorContext,
2610            _request: ReconcileRequest,
2611        ) -> common_meta::error::Result<ReconcileResponse> {
2612            common_meta::error::UnsupportedSnafu {
2613                operation: "reconcile",
2614            }
2615            .fail()
2616        }
2617
2618        async fn query_procedure_state(
2619            &self,
2620            _ctx: &ExecutorContext,
2621            _pid: &str,
2622        ) -> common_meta::error::Result<ProcedureStateResponse> {
2623            common_meta::error::UnsupportedSnafu {
2624                operation: "query_procedure_state",
2625            }
2626            .fail()
2627        }
2628
2629        async fn list_procedures(
2630            &self,
2631            _ctx: &ExecutorContext,
2632        ) -> common_meta::error::Result<ProcedureDetailResponse> {
2633            common_meta::error::UnsupportedSnafu {
2634                operation: "list_procedures",
2635            }
2636            .fail()
2637        }
2638    }
2639
2640    fn test_cache_registry(
2641        kv_backend: common_meta::kv_backend::KvBackendRef,
2642    ) -> TestResult<common_meta::cache::LayeredCacheRegistryRef> {
2643        Ok(Arc::new(
2644            cache::with_default_composite_cache_registry(
2645                LayeredCacheRegistryBuilder::default()
2646                    .add_cache_registry(cache::build_fundamental_cache_registry(kv_backend)),
2647            )
2648            .context(BuildCacheRegistrySnafu)?
2649            .build(),
2650        ))
2651    }
2652
2653    fn test_table_info(table_id: u32, table_name: &str) -> TestResult<TableInfo> {
2654        let schema = Arc::new(GtSchema::new(vec![
2655            ColumnSchema::new("id", ConcreteDataType::int32_datatype(), false),
2656            ColumnSchema::new(
2657                "ts",
2658                ConcreteDataType::timestamp_millisecond_datatype(),
2659                false,
2660            )
2661            .with_time_index(true),
2662        ]));
2663        let table_meta = TableMetaBuilder::empty()
2664            .schema(schema)
2665            .primary_key_indices(vec![0])
2666            .value_indices(vec![1])
2667            .next_column_id(1024)
2668            .build()
2669            .with_context(|_| BuildTableMetaSnafu {
2670                table_name: table_name.to_string(),
2671            })?;
2672
2673        TableInfoBuilder::new(table_name, table_meta)
2674            .table_id(table_id)
2675            .build()
2676            .with_context(|_| BuildTableInfoSnafu {
2677                table_name: table_name.to_string(),
2678            })
2679    }
2680
2681    fn test_table(table_id: u32, table_name: &str) -> TestResult<table::TableRef> {
2682        let table_info = test_table_info(table_id, table_name)?;
2683        Ok(EmptyTable::from_table_info(&table_info))
2684    }
2685
2686    fn test_physical_table(table_id: u32, table_name: &str) -> TestResult<table::TableRef> {
2687        let mut table_info = test_table_info(table_id, table_name)?;
2688        table_info
2689            .meta
2690            .options
2691            .extra_options
2692            .insert(PHYSICAL_TABLE_METADATA_KEY.to_string(), String::new());
2693        Ok(EmptyTable::from_table_info(&table_info))
2694    }
2695
2696    fn test_logical_table(table_id: u32, table_name: &str) -> TestResult<table::TableRef> {
2697        let mut table_info = test_table_info(table_id, table_name)?;
2698        table_info.meta.engine = METRIC_ENGINE_NAME.to_string();
2699        table_info.meta.options.extra_options.insert(
2700            LOGICAL_TABLE_METADATA_KEY.to_string(),
2701            "physical_metric".to_string(),
2702        );
2703        Ok(EmptyTable::from_table_info(&table_info))
2704    }
2705
2706    fn test_metric_names_table() -> TableRef {
2707        let schema = Arc::new(GtSchema::new(vec![
2708            ColumnSchema::new("table_catalog", ConcreteDataType::string_datatype(), false),
2709            ColumnSchema::new("table_schema", ConcreteDataType::string_datatype(), false),
2710            ColumnSchema::new("table_name", ConcreteDataType::string_datatype(), false),
2711            ColumnSchema::new("engine", ConcreteDataType::string_datatype(), false),
2712            ColumnSchema::new("create_options", ConcreteDataType::string_datatype(), false),
2713        ]));
2714        let columns: Vec<VectorRef> = vec![
2715            Arc::new(StringVector::from(vec!["greptime", "greptime"])),
2716            Arc::new(StringVector::from(vec!["public", "public"])),
2717            Arc::new(StringVector::from(vec!["denied", "target"])),
2718            Arc::new(StringVector::from(vec!["metric", "metric"])),
2719            Arc::new(StringVector::from(vec![
2720                "on_physical_table=physical_metric",
2721                "on_physical_table=physical_metric",
2722            ])),
2723        ];
2724        let record_batch = RecordBatch::new(schema, columns).unwrap();
2725        MemTable::new_with_catalog(
2726            "tables",
2727            record_batch,
2728            2048,
2729            "greptime".to_string(),
2730            "information_schema".to_string(),
2731        )
2732    }
2733
2734    fn test_pipeline_table() -> TableRef {
2735        let schema = Arc::new(GtSchema::new(vec![
2736            ColumnSchema::new("name", ConcreteDataType::string_datatype(), false),
2737            ColumnSchema::new("schema", ConcreteDataType::string_datatype(), false),
2738            ColumnSchema::new("content_type", ConcreteDataType::string_datatype(), false),
2739            ColumnSchema::new("pipeline", ConcreteDataType::string_datatype(), false),
2740            ColumnSchema::new(
2741                "created_at",
2742                ConcreteDataType::timestamp_nanosecond_datatype(),
2743                false,
2744            )
2745            .with_time_index(true),
2746        ]));
2747        let columns: Vec<VectorRef> = vec![
2748            Arc::new(StringVector::from(vec!["pipeline"])),
2749            Arc::new(StringVector::from(vec!["public"])),
2750            Arc::new(StringVector::from(vec!["application/yaml"])),
2751            Arc::new(StringVector::from(vec![
2752                "transform:\n- field: ts\n  type: timestamp, ns\n  index: time\n",
2753            ])),
2754            Arc::new(TimestampNanosecondVector::from_values([1])),
2755        ];
2756        let record_batch = RecordBatch::new(schema, columns).unwrap();
2757        MemTable::new_with_catalog(
2758            "pipelines",
2759            record_batch,
2760            2049,
2761            "greptime".to_string(),
2762            DEFAULT_PRIVATE_SCHEMA_NAME.to_string(),
2763        )
2764    }
2765
2766    struct CancellableExportSource {
2767        schema: GtSchemaRef,
2768        channels: std::sync::Mutex<Option<(oneshot::Sender<()>, oneshot::Receiver<()>)>>,
2769    }
2770
2771    impl DataSource for CancellableExportSource {
2772        fn get_stream(
2773            &self,
2774            _request: ScanRequest,
2775        ) -> std::result::Result<SendableRecordBatchStream, BoxedError> {
2776            let (started, release) = self.channels.lock().unwrap().take().unwrap();
2777            let schema = self.schema.clone();
2778            let stream = futures::stream::once(async move {
2779                started.send(()).unwrap();
2780                release.await.unwrap();
2781                Ok(RecordBatch::new_empty(schema))
2782            });
2783            Ok(Box::pin(RecordBatchStreamWrapper::new(
2784                self.schema.clone(),
2785                Box::pin(stream),
2786            )))
2787        }
2788    }
2789
2790    #[tokio::test]
2791    async fn test_metric_export_http_timeout_cancels_ordinary_source() {
2792        let destination = common_test_util::temp_dir::create_temp_dir("metric_export_timeout");
2793        let (started_tx, started_rx) = oneshot::channel();
2794        let (mut release_tx, release_rx) = oneshot::channel();
2795        let info = test_table_info(1024, "source").unwrap();
2796        let source = Arc::new(Table::new(
2797            Arc::new(info.clone()),
2798            FilterPushDownType::Unsupported,
2799            Arc::new(CancellableExportSource {
2800                schema: info.meta.schema.clone(),
2801                channels: std::sync::Mutex::new(Some((started_tx, release_rx))),
2802            }),
2803        ));
2804        let catalog = catalog::memory::MemoryCatalogManager::new_with_table(source);
2805        let kv = Arc::new(MemoryKvBackend::new());
2806        let instance = FrontendBuilder::new(
2807            FrontendOptions {
2808                experimental_metric_export: true,
2809                ..Default::default()
2810            },
2811            kv.clone(),
2812            test_cache_registry(kv).unwrap(),
2813            catalog,
2814            Arc::new(client::client_manager::NodeClients::default()),
2815            Arc::new(NoopProcedureExecutor),
2816            Arc::new(ProcessManager::new("export-timeout".into(), None)),
2817        )
2818        .with_local_file_access(
2819            common_datasource::object_store::LocalFileAccess::sandboxed(destination.path())
2820                .unwrap(),
2821        )
2822        .try_build()
2823        .await
2824        .unwrap();
2825        let server = servers::http::HttpServerBuilder::new(servers::http::HttpOptions {
2826            timeout: Duration::from_secs(2),
2827            ..Default::default()
2828        })
2829        .with_sql_handler(Arc::new(instance))
2830        .build();
2831        let app = server.build(server.make_app()).unwrap();
2832        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
2833        let addr = listener.local_addr().unwrap();
2834        let server_task = tokio::spawn(async move {
2835            axum::serve(listener, app).await.unwrap();
2836        });
2837        let uri = reqwest::Url::from_directory_path(destination.path()).unwrap();
2838        let sql = format!(
2839            "COPY DATABASE greptime.public TO '{uri}' WITH (experimental_metric_export='true')"
2840        );
2841        let request = tokio::spawn(async move {
2842            reqwest::Client::builder()
2843                .no_proxy()
2844                .build()
2845                .unwrap()
2846                .post(format!("http://{addr}/v1/sql"))
2847                .form(&[("sql", sql)])
2848                .send()
2849                .await
2850                .unwrap()
2851        });
2852        tokio::time::timeout(Duration::from_secs(5), started_rx)
2853            .await
2854            .unwrap()
2855            .unwrap();
2856        let response = request.await.unwrap();
2857        assert_eq!(response.status(), reqwest::StatusCode::REQUEST_TIMEOUT);
2858        tokio::time::timeout(Duration::from_secs(5), release_tx.closed())
2859            .await
2860            .unwrap();
2861        assert!(!destination.path().join("source.parquet").exists());
2862        server_task.abort();
2863    }
2864
2865    fn pending_table(
2866        table_id: u32,
2867        table_name: &str,
2868        polled_tx: oneshot::Sender<()>,
2869    ) -> TestResult<table::TableRef> {
2870        let table_info = test_table_info(table_id, table_name)?;
2871        let data_source = Arc::new(PendingDataSource {
2872            schema: table_info.meta.schema.clone(),
2873            polled_tx: std::sync::Mutex::new(Some(polled_tx)),
2874        });
2875
2876        Ok(Arc::new(Table::new(
2877            Arc::new(table_info),
2878            FilterPushDownType::Unsupported,
2879            data_source,
2880        )))
2881    }
2882
2883    async fn test_instance_with_tables(
2884        source_table: TableRef,
2885        target_table: TableRef,
2886    ) -> TestResult<Instance> {
2887        test_instance_with_plugins(source_table, target_table, Plugins::new()).await
2888    }
2889
2890    async fn test_instance_with_insert_select_interceptor(
2891        interceptor: SqlQueryInterceptorRef<Error>,
2892    ) -> TestResult<Instance> {
2893        let plugins = Plugins::new();
2894        plugins.insert::<SqlQueryInterceptorRef<Error>>(interceptor);
2895
2896        test_instance_with_plugins(
2897            test_table(1024, "source")?,
2898            test_table(1025, "target")?,
2899            plugins,
2900        )
2901        .await
2902    }
2903
2904    async fn test_instance_with_plugins(
2905        source_table: TableRef,
2906        target_table: TableRef,
2907        plugins: Plugins,
2908    ) -> TestResult<Instance> {
2909        test_instance_with_plugins_and_metric_names(source_table, target_table, plugins, None).await
2910    }
2911
2912    async fn test_instance_with_plugins_and_metric_names(
2913        source_table: TableRef,
2914        target_table: TableRef,
2915        plugins: Plugins,
2916        metric_names_table: Option<TableRef>,
2917    ) -> TestResult<Instance> {
2918        let catalog_manager = catalog::memory::MemoryCatalogManager::new_with_table(source_table);
2919        test_instance_with_catalog_manager(
2920            catalog_manager,
2921            target_table,
2922            plugins,
2923            metric_names_table,
2924            Arc::new(NoopProcedureExecutor),
2925        )
2926        .await
2927    }
2928
2929    /// Builds a test frontend `Instance` over the given (already source-registered)
2930    /// catalog manager, completing DDL tasks through `procedure_executor`.
2931    async fn test_instance_with_catalog_manager(
2932        catalog_manager: Arc<catalog::memory::MemoryCatalogManager>,
2933        target_table: TableRef,
2934        plugins: Plugins,
2935        metric_names_table: Option<TableRef>,
2936        procedure_executor: ProcedureExecutorRef,
2937    ) -> TestResult<Instance> {
2938        let kv_backend = Arc::new(MemoryKvBackend::new());
2939        let process_manager = Arc::new(ProcessManager::new("test-frontend".to_string(), None));
2940        let target_table_name = "target";
2941        catalog_manager
2942            .register_table_sync(catalog::RegisterTableRequest {
2943                catalog: "greptime".to_string(),
2944                schema: "public".to_string(),
2945                table_name: target_table_name.to_string(),
2946                table_id: 1025,
2947                table: target_table,
2948            })
2949            .with_context(|_| RegisterTableSnafu {
2950                table_name: target_table_name.to_string(),
2951            })?;
2952        if let Some(table) = metric_names_table {
2953            catalog_manager
2954                .deregister_table_sync(catalog::DeregisterTableRequest {
2955                    catalog: "greptime".to_string(),
2956                    schema: "information_schema".to_string(),
2957                    table_name: "tables".to_string(),
2958                })
2959                .unwrap();
2960            catalog_manager
2961                .register_table_sync(catalog::RegisterTableRequest {
2962                    catalog: "greptime".to_string(),
2963                    schema: "information_schema".to_string(),
2964                    table_name: "tables".to_string(),
2965                    table_id: 2048,
2966                    table,
2967                })
2968                .unwrap();
2969        }
2970        catalog_manager.register_process_list_table(process_manager.clone());
2971
2972        let cache_registry = test_cache_registry(kv_backend.clone())?;
2973
2974        FrontendBuilder::new(
2975            FrontendOptions::default(),
2976            kv_backend,
2977            cache_registry,
2978            catalog_manager,
2979            Arc::new(client::client_manager::NodeClients::default()),
2980            procedure_executor,
2981            process_manager,
2982        )
2983        .with_plugin(plugins)
2984        .try_build()
2985        .await
2986        .context(BuildFrontendSnafu)
2987    }
2988
2989    async fn execute_one_sql(
2990        instance: &Instance,
2991        sql: &str,
2992        query_ctx: QueryContextRef,
2993    ) -> TestResult<Output> {
2994        let mut results = instance.do_query_inner(sql, query_ctx).await;
2995        ensure!(
2996            results.len() == 1,
2997            UnexpectedOutputCountSnafu {
2998                sql: sql.to_string(),
2999                actual: results.len(),
3000            }
3001        );
3002        results.remove(0).with_context(|_| ExecuteSqlSnafu {
3003            sql: sql.to_string(),
3004        })
3005    }
3006
3007    fn assert_permission_denied<T>(result: servers::error::Result<T>) {
3008        let err = match result {
3009            Ok(_) => panic!("request should be rejected"),
3010            Err(err) => err,
3011        };
3012        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3013    }
3014
3015    fn assert_action_checked(
3016        checker: &RejectEndpointPermissionChecker,
3017        action: PermissionAction,
3018        targets: Option<PermissionTableTargets>,
3019    ) {
3020        assert_eq!(CheckedAction { action, targets }, checker.take_check());
3021    }
3022
3023    #[tokio::test]
3024    async fn database_export_authorizes_all_tables_before_preparation() -> TestResult<()> {
3025        struct ExportAcl(std::sync::Mutex<Vec<PermissionTableTargets>>);
3026        impl PermissionChecker for ExportAcl {
3027            fn check_permission(
3028                &self,
3029                _: UserInfoRef,
3030                req: PermissionReq,
3031            ) -> auth::error::Result<PermissionResp> {
3032                assert!(matches!(
3033                    req,
3034                    PermissionReq::SqlStatement(Statement::Copy(
3035                        sql::statements::copy::Copy::CopyDatabase(
3036                            sql::statements::copy::CopyDatabase::To(_)
3037                        )
3038                    ))
3039                ));
3040                Ok(PermissionResp::Allow)
3041            }
3042            fn check_permission_with_table_targets(
3043                &self,
3044                _: UserInfoRef,
3045                req: PermissionReq,
3046                targets: PermissionTableTargets,
3047            ) -> auth::error::Result<PermissionResp> {
3048                self.0.lock().unwrap().push(targets.clone());
3049                let PermissionTableTargets::Resolved(tables) = targets else {
3050                    panic!("unresolved export")
3051                };
3052                if req.is_readonly() {
3053                    // The first table is readable; the later table has only write access.
3054                    Ok(if tables.iter().any(|t| t.table == "target") {
3055                        PermissionResp::Reject
3056                    } else {
3057                        PermissionResp::Allow
3058                    })
3059                } else {
3060                    self.check_permission(QueryContext::arc().current_user(), req)
3061                }
3062            }
3063        }
3064        let checker = Arc::new(ExportAcl(Default::default()));
3065        let plugins = Plugins::new();
3066        plugins.insert::<PermissionCheckerRef>(checker.clone());
3067        let instance = test_instance_with_plugins(
3068            test_logical_table(1024, "source")?,
3069            test_table(1025, "target")?,
3070            plugins,
3071        )
3072        .await?;
3073        let req = table::requests::CopyDatabaseRequest {
3074            catalog_name: "greptime".into(),
3075            schema_name: "public".into(),
3076            location: "invalid-destination".into(),
3077            with: Default::default(),
3078            connection: Default::default(),
3079            time_range: None,
3080        };
3081        let result = instance
3082            .export_database_for_test(
3083                req.clone(),
3084                None,
3085                &tokio_util::sync::CancellationToken::new(),
3086                QueryContext::arc(),
3087            )
3088            .await;
3089        assert!(matches!(result, Err(Error::Permission { .. })));
3090        let expected = PermissionTableTargets::resolved(vec![
3091            PermissionTableTarget::new("greptime", "public", "source"),
3092            PermissionTableTarget::new("greptime", "public", "target"),
3093        ]);
3094        assert_eq!(*checker.0.lock().unwrap(), vec![expected.clone(), expected]);
3095        // An empty selection must still check the operation privilege.
3096        instance
3097            .plugins
3098            .map_mut::<PermissionCheckerRef, _, _>(|checker| {
3099                *checker.unwrap() = Arc::new(WriteOnlyPermissionChecker)
3100            });
3101        let result = instance
3102            .export_database_for_test(
3103                req,
3104                Some(&[]),
3105                &tokio_util::sync::CancellationToken::new(),
3106                QueryContext::arc(),
3107            )
3108            .await;
3109        assert!(matches!(result, Err(Error::Permission { .. })));
3110        Ok(())
3111    }
3112
3113    #[tokio::test]
3114    async fn test_prom_remote_read_with_custom_timestamp_and_value_columns() -> TestResult<()> {
3115        let schema = Arc::new(GtSchema::new(vec![
3116            ColumnSchema::new(
3117                "custom_ts",
3118                ConcreteDataType::timestamp_millisecond_datatype(),
3119                false,
3120            )
3121            .with_time_index(true),
3122            ColumnSchema::new("custom_value", ConcreteDataType::float64_datatype(), false),
3123        ]));
3124        let recordbatch = RecordBatch::new(
3125            schema,
3126            vec![
3127                Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as VectorRef,
3128                Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as VectorRef,
3129            ],
3130        )
3131        .unwrap();
3132        let instance = test_instance_with_tables(
3133            MemTable::table("custom_metric", recordbatch),
3134            test_table(1025, "target")?,
3135        )
3136        .await?;
3137
3138        let response = PromStoreProtocolHandler::read(
3139            &instance,
3140            ReadRequest {
3141                queries: vec![RemoteQuery {
3142                    start_timestamp_ms: 1500,
3143                    end_timestamp_ms: 2500,
3144                    matchers: vec![LabelMatcher {
3145                        r#type: PromMatcherType::Eq as i32,
3146                        name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3147                        value: "custom_metric".to_string(),
3148                    }],
3149                    ..Default::default()
3150                }],
3151                ..Default::default()
3152            },
3153            test_query_ctx(1),
3154        )
3155        .await
3156        .unwrap();
3157        let body = servers::prom_store::snappy_decompress(&response.body).unwrap();
3158        let response = ReadResponse::decode(body.as_slice()).unwrap();
3159
3160        assert_eq!(1, response.results.len());
3161        assert_eq!(1, response.results[0].timeseries.len());
3162        let timeseries = &response.results[0].timeseries[0];
3163        assert_eq!(
3164            vec![Label {
3165                name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3166                value: "custom_metric".to_string(),
3167            }],
3168            timeseries.labels
3169        );
3170        assert_eq!(
3171            vec![Sample {
3172                value: 2.0,
3173                timestamp: 2000,
3174            }],
3175            timeseries.samples
3176        );
3177
3178        Ok(())
3179    }
3180
3181    #[tokio::test]
3182    async fn test_prom_remote_read_prefers_default_value_column() -> TestResult<()> {
3183        let schema = Arc::new(GtSchema::new(vec![
3184            ColumnSchema::new(
3185                "custom_ts",
3186                ConcreteDataType::timestamp_millisecond_datatype(),
3187                false,
3188            )
3189            .with_time_index(true),
3190            ColumnSchema::new("extra_field", ConcreteDataType::float64_datatype(), false),
3191            ColumnSchema::new(
3192                greptime_value(),
3193                ConcreteDataType::float64_datatype(),
3194                false,
3195            ),
3196        ]));
3197        let recordbatch = RecordBatch::new(
3198            schema,
3199            vec![
3200                Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as VectorRef,
3201                Arc::new(Float64Vector::from_vec(vec![99.0, 99.0, 99.0])) as VectorRef,
3202                Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as VectorRef,
3203            ],
3204        )
3205        .unwrap();
3206        let instance = test_instance_with_tables(
3207            MemTable::table("multi_field_metric", recordbatch),
3208            test_table(1025, "target")?,
3209        )
3210        .await?;
3211
3212        let response = PromStoreProtocolHandler::read(
3213            &instance,
3214            ReadRequest {
3215                queries: vec![RemoteQuery {
3216                    start_timestamp_ms: 1500,
3217                    end_timestamp_ms: 2500,
3218                    matchers: vec![LabelMatcher {
3219                        r#type: PromMatcherType::Eq as i32,
3220                        name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3221                        value: "multi_field_metric".to_string(),
3222                    }],
3223                    ..Default::default()
3224                }],
3225                ..Default::default()
3226            },
3227            test_query_ctx(1),
3228        )
3229        .await
3230        .unwrap();
3231        let body = servers::prom_store::snappy_decompress(&response.body).unwrap();
3232        let response = ReadResponse::decode(body.as_slice()).unwrap();
3233
3234        assert_eq!(1, response.results.len());
3235        assert_eq!(1, response.results[0].timeseries.len());
3236        let timeseries = &response.results[0].timeseries[0];
3237        assert_eq!(
3238            vec![
3239                Label {
3240                    name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3241                    value: "multi_field_metric".to_string(),
3242                },
3243                Label {
3244                    name: "extra_field".to_string(),
3245                    value: "99".to_string(),
3246                },
3247            ],
3248            timeseries.labels
3249        );
3250        assert_eq!(
3251            vec![Sample {
3252                value: 2.0,
3253                timestamp: 2000,
3254            }],
3255            timeseries.samples
3256        );
3257
3258        Ok(())
3259    }
3260
3261    #[tokio::test]
3262    async fn test_prom_remote_read_rejects_ambiguous_value_columns() -> TestResult<()> {
3263        let schema = Arc::new(GtSchema::new(vec![
3264            ColumnSchema::new(
3265                "custom_ts",
3266                ConcreteDataType::timestamp_millisecond_datatype(),
3267                false,
3268            )
3269            .with_time_index(true),
3270            ColumnSchema::new("field_a", ConcreteDataType::float64_datatype(), false),
3271            ColumnSchema::new("field_b", ConcreteDataType::float64_datatype(), false),
3272        ]));
3273        let recordbatch = RecordBatch::new(
3274            schema,
3275            vec![
3276                Arc::new(TimestampMillisecondVector::from_vec(vec![1000])) as VectorRef,
3277                Arc::new(Float64Vector::from_vec(vec![1.0])) as VectorRef,
3278                Arc::new(Float64Vector::from_vec(vec![2.0])) as VectorRef,
3279            ],
3280        )
3281        .unwrap();
3282        let instance = test_instance_with_tables(
3283            MemTable::table("ambiguous_metric", recordbatch),
3284            test_table(1025, "target")?,
3285        )
3286        .await?;
3287
3288        let err = PromStoreProtocolHandler::read(
3289            &instance,
3290            ReadRequest {
3291                queries: vec![RemoteQuery {
3292                    matchers: vec![LabelMatcher {
3293                        r#type: PromMatcherType::Eq as i32,
3294                        name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3295                        value: "ambiguous_metric".to_string(),
3296                    }],
3297                    ..Default::default()
3298                }],
3299                ..Default::default()
3300            },
3301            test_query_ctx(1),
3302        )
3303        .await
3304        .err()
3305        .expect("ambiguous value columns should fail remote read");
3306
3307        assert_eq!(StatusCode::InvalidArguments, err.status_code());
3308        assert!(format!("{err:?}").contains("Ambiguous value column"));
3309
3310        Ok(())
3311    }
3312
3313    #[tokio::test]
3314    async fn test_event_recorder_is_exposed() -> TestResult<()> {
3315        let instance =
3316            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
3317                .await?;
3318
3319        let _event_recorder = instance.event_recorder();
3320
3321        Ok(())
3322    }
3323
3324    #[tokio::test]
3325    async fn test_restricted_endpoint_handlers_check_permissions() -> TestResult<()> {
3326        let checker = Arc::new(RejectEndpointPermissionChecker::default());
3327        let plugins = Plugins::new();
3328        plugins.insert::<PermissionCheckerRef>(checker.clone());
3329        let instance = test_instance_with_plugins(
3330            test_table(1024, "denied")?,
3331            test_table(1025, "target")?,
3332            plugins,
3333        )
3334        .await?;
3335        let mut ctx = test_query_ctx(1);
3336        Arc::get_mut(&mut ctx).unwrap().set_extension(
3337            servers::http::jaeger::JAEGER_QUERY_TABLE_NAME_KEY,
3338            "denied".to_string(),
3339        );
3340        let jaeger_targets = Some(PermissionTableTargets::resolved(vec![
3341            PermissionTableTarget::new("greptime", "public", "denied"),
3342        ]));
3343
3344        assert_permission_denied(JaegerQueryHandler::get_services(&instance, ctx.clone()).await);
3345        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
3346        assert_permission_denied(
3347            JaegerQueryHandler::get_operations(&instance, ctx.clone(), "service", None).await,
3348        );
3349        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
3350        assert_permission_denied(
3351            JaegerQueryHandler::get_trace(&instance, ctx.clone(), "trace", None, None, None).await,
3352        );
3353        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
3354        assert_permission_denied(
3355            JaegerQueryHandler::find_traces(
3356                &instance,
3357                ctx.clone(),
3358                servers::http::jaeger::QueryTraceParams {
3359                    service_name: "service".to_string(),
3360                    ..Default::default()
3361                },
3362            )
3363            .await,
3364        );
3365        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets);
3366
3367        assert_permission_denied(
3368            PipelineHandler::get_pipeline_str(&instance, "pipeline", None, ctx.clone()).await,
3369        );
3370        assert_action_checked(&checker, PIPELINE_QUERY, None);
3371        assert_permission_denied(
3372            PipelineHandler::insert_pipeline(
3373                &instance,
3374                "pipeline",
3375                "application/yaml",
3376                "",
3377                ctx.clone(),
3378            )
3379            .await,
3380        );
3381        assert_action_checked(&checker, PIPELINE_INSERT, None);
3382        assert_permission_denied(
3383            PipelineHandler::delete_pipeline(&instance, "pipeline", None, ctx.clone()).await,
3384        );
3385        assert_action_checked(&checker, PIPELINE_DELETE, None);
3386        let app = axum::Router::new()
3387            .route(
3388                "/pipelines/_dryrun",
3389                axum::routing::post(servers::http::event::pipeline_dryrun),
3390            )
3391            .with_state(servers::http::event::LogState {
3392                log_handler: Arc::new(instance.clone()),
3393                log_validator: None,
3394                ingest_interceptor: None,
3395            })
3396            .layer(axum::Extension((*ctx).clone()));
3397        let response = app
3398            .oneshot(
3399                axum::http::Request::post("/pipelines/_dryrun")
3400                    .header("content-type", "application/json")
3401                    .body(axum::body::Body::from("{}"))
3402                    .unwrap(),
3403            )
3404            .await
3405            .unwrap();
3406        assert_eq!(axum::http::StatusCode::FORBIDDEN, response.status());
3407        assert_action_checked(&checker, PIPELINE_QUERY, None);
3408
3409        assert_permission_denied(
3410            DashboardHandler::save(&instance, "dashboard", "{}", ctx.clone()).await,
3411        );
3412        assert_action_checked(&checker, DASHBOARD_SAVE, None);
3413        assert_permission_denied(DashboardHandler::list(&instance, ctx.clone()).await);
3414        assert_action_checked(&checker, DASHBOARD_QUERY, None);
3415        assert_permission_denied(
3416            DashboardHandler::delete(&instance, "dashboard", ctx.clone()).await,
3417        );
3418        assert_action_checked(&checker, DASHBOARD_DELETE, None);
3419
3420        Ok(())
3421    }
3422
3423    #[tokio::test]
3424    async fn test_write_only_ingestion_loads_named_pipeline() -> TestResult<()> {
3425        let plugins = Plugins::new();
3426        plugins.insert::<PermissionCheckerRef>(Arc::new(WriteOnlyPermissionChecker));
3427        let instance = test_instance_with_plugins(
3428            test_table(1024, "source")?,
3429            test_table(1025, "target")?,
3430            plugins,
3431        )
3432        .await?;
3433        instance
3434            .catalog_manager()
3435            .as_any()
3436            .downcast_ref::<catalog::memory::MemoryCatalogManager>()
3437            .unwrap()
3438            .register_table_sync(catalog::RegisterTableRequest {
3439                catalog: "greptime".to_string(),
3440                schema: DEFAULT_PRIVATE_SCHEMA_NAME.to_string(),
3441                table_name: "pipelines".to_string(),
3442                table_id: 2049,
3443                table: test_pipeline_table(),
3444            })
3445            .with_context(|_| RegisterTableSnafu {
3446                table_name: "pipelines".to_string(),
3447            })?;
3448        let ctx = test_query_ctx(1);
3449        let handler: PipelineHandlerRef = Arc::new(instance.clone());
3450
3451        handler
3452            .get_pipeline("pipeline", None, ctx.clone())
3453            .await
3454            .unwrap();
3455        assert_permission_denied(
3456            PipelineHandler::get_pipeline_str(&instance, "pipeline", None, ctx.clone()).await,
3457        );
3458
3459        let app = axum::Router::new()
3460            .route(
3461                "/pipelines/_dryrun",
3462                axum::routing::post(servers::http::event::pipeline_dryrun),
3463            )
3464            .with_state(servers::http::event::LogState {
3465                log_handler: handler,
3466                log_validator: None,
3467                ingest_interceptor: None,
3468            })
3469            .layer(axum::Extension((*ctx).clone()));
3470        let response = app
3471            .oneshot(
3472                axum::http::Request::post("/pipelines/_dryrun")
3473                    .header("content-type", "application/json")
3474                    .body(axum::body::Body::from("{}"))
3475                    .unwrap(),
3476            )
3477            .await
3478            .unwrap();
3479        assert_eq!(axum::http::StatusCode::FORBIDDEN, response.status());
3480
3481        Ok(())
3482    }
3483
3484    #[tokio::test]
3485    async fn test_write_only_grpc_sql_is_checked_after_parsing() -> TestResult<()> {
3486        let plugins = Plugins::new();
3487        plugins.insert::<PermissionCheckerRef>(Arc::new(WriteOnlyPermissionChecker));
3488        let instance = test_instance_with_plugins(
3489            test_table(1024, "source")?,
3490            test_table(1025, "target")?,
3491            plugins,
3492        )
3493        .await?;
3494
3495        let insert = Request::Query(api::v1::QueryRequest {
3496            query: Some(Query::Sql(
3497                "INSERT INTO target SELECT * FROM source".to_string(),
3498            )),
3499        });
3500        servers::query_handler::grpc::GrpcQueryHandler::do_query(
3501            &instance,
3502            insert,
3503            QueryContext::arc(),
3504        )
3505        .await
3506        .unwrap();
3507
3508        let select = Request::Query(api::v1::QueryRequest {
3509            query: Some(Query::Sql("SELECT * FROM source".to_string())),
3510        });
3511        assert_permission_denied(
3512            servers::query_handler::grpc::GrpcQueryHandler::do_query(
3513                &instance,
3514                select,
3515                QueryContext::arc(),
3516            )
3517            .await,
3518        );
3519
3520        Ok(())
3521    }
3522
3523    #[tokio::test]
3524    async fn test_target_independent_checker_skips_target_resolution() -> TestResult<()> {
3525        let physical_table = "physical_metric";
3526        let checker = Arc::new(TargetIndependentPermissionChecker::default());
3527        let plugins = Plugins::new();
3528        plugins.insert::<PermissionCheckerRef>(checker.clone());
3529        let instance = test_instance_with_plugins(
3530            test_physical_table(1024, physical_table)?,
3531            test_table(1025, "target")?,
3532            plugins,
3533        )
3534        .await?;
3535
3536        let ctx = test_query_ctx(1);
3537        let physical_target = PermissionTableTarget::new("greptime", "public", physical_table);
3538        assert_eq!(
3539            PermissionTableTargets::Resolved(vec![physical_target.clone()]),
3540            instance
3541                .resolve_query_permission_targets(
3542                    PermissionTableTargets::resolved(vec![physical_target]),
3543                    &ctx,
3544                )
3545                .await
3546                .unwrap()
3547        );
3548        assert_eq!(
3549            vec![physical_table.to_string(), "target".to_string()],
3550            PrometheusHandler::filter_metadata_metric_names(
3551                &instance,
3552                vec![physical_table.to_string(), "target".to_string()],
3553                "public",
3554                &ctx,
3555            )
3556            .await
3557            .unwrap()
3558        );
3559        assert_eq!(1, checker.checks.load(atomic::Ordering::Relaxed));
3560
3561        Ok(())
3562    }
3563
3564    #[tokio::test]
3565    async fn test_query_permission_targets_are_deduplicated() -> TestResult<()> {
3566        let plugins = Plugins::new();
3567        plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
3568        let instance = test_instance_with_plugins(
3569            test_table(1024, "source")?,
3570            test_table(1025, "target")?,
3571            plugins,
3572        )
3573        .await?;
3574        let ctx = test_query_ctx(1);
3575        let target = PermissionTableTarget::new("greptime", "public", "target");
3576
3577        assert_eq!(
3578            PermissionTableTargets::Resolved(vec![target.clone()]),
3579            instance
3580                .resolve_query_permission_targets(
3581                    PermissionTableTargets::resolved(vec![target.clone(), target]),
3582                    &ctx,
3583                )
3584                .await
3585                .unwrap()
3586        );
3587
3588        Ok(())
3589    }
3590
3591    #[tokio::test]
3592    async fn test_physical_query_targets_fail_closed() -> TestResult<()> {
3593        let physical_table = "physical_metric";
3594        let plugins = Plugins::new();
3595        plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
3596        let instance = test_instance_with_plugins(
3597            test_physical_table(1024, physical_table)?,
3598            test_table(1025, "target")?,
3599            plugins,
3600        )
3601        .await?;
3602
3603        let ctx = test_query_ctx(1);
3604        let logical_target = PermissionTableTarget::new("greptime", "public", "target");
3605        assert_eq!(
3606            PermissionTableTargets::Resolved(vec![logical_target.clone()]),
3607            instance
3608                .resolve_query_permission_targets(
3609                    PermissionTableTargets::resolved(vec![logical_target.clone()]),
3610                    &ctx,
3611                )
3612                .await
3613                .unwrap()
3614        );
3615        let physical_target = PermissionTableTarget::new("greptime", "public", physical_table);
3616        assert_eq!(
3617            PermissionTableTargets::Unresolved,
3618            instance
3619                .resolve_query_permission_targets(
3620                    PermissionTableTargets::resolved(
3621                        vec![logical_target, physical_target.clone(),]
3622                    ),
3623                    &ctx,
3624                )
3625                .await
3626                .unwrap()
3627        );
3628        assert_eq!(
3629            vec!["target".to_string()],
3630            PrometheusHandler::filter_metadata_metric_names(
3631                &instance,
3632                vec!["target".to_string(), "denied".to_string()],
3633                "public",
3634                &ctx,
3635            )
3636            .await
3637            .unwrap()
3638        );
3639
3640        let query = PromQuery {
3641            query: physical_table.to_string(),
3642            ..Default::default()
3643        };
3644        let err = PrometheusHandler::check_query_target_permission(
3645            &instance,
3646            PermissionTableTargets::resolved(vec![physical_target]),
3647            &ctx,
3648        )
3649        .await
3650        .unwrap_err();
3651        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3652        let err = PrometheusHandler::check_query_permission(
3653            &instance,
3654            std::slice::from_ref(&query),
3655            &ctx,
3656        )
3657        .await
3658        .unwrap_err();
3659        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3660        let err = PrometheusHandler::do_query(&instance, &query, ctx.clone())
3661            .await
3662            .unwrap_err();
3663        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3664
3665        for sql in [
3666            "SELECT * FROM physical_metric",
3667            "TQL EVAL (0, 10, '5s') physical_metric",
3668            "INSERT INTO target SELECT * FROM physical_metric",
3669        ] {
3670            let mut results = instance.do_query_inner(sql, ctx.clone()).await;
3671            assert_eq!(1, results.len(), "{sql}");
3672            let err = results.remove(0).unwrap_err();
3673            assert_eq!(StatusCode::PermissionDenied, err.status_code(), "{sql}");
3674        }
3675        let err = LogQueryHandler::query(
3676            &instance,
3677            LogQuery {
3678                table: TableName::new("greptime", "public", physical_table),
3679                ..Default::default()
3680            },
3681            ctx.clone(),
3682        )
3683        .await
3684        .unwrap_err();
3685        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3686        let err = instance
3687            .do_describe_inner(parse_one_sql("SELECT * FROM physical_metric"), ctx.clone())
3688            .await
3689            .unwrap_err();
3690        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3691
3692        let request = ReadRequest {
3693            queries: vec![RemoteQuery {
3694                matchers: vec![LabelMatcher {
3695                    r#type: PromMatcherType::Eq as i32,
3696                    name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3697                    value: physical_table.to_string(),
3698                }],
3699                ..Default::default()
3700            }],
3701            ..Default::default()
3702        };
3703        let Err(err) = PromStoreProtocolHandler::read(&instance, request, ctx.clone()).await else {
3704            panic!("physical remote-read target must be rejected");
3705        };
3706        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3707
3708        let err = PrometheusHandler::query_label_values(
3709            &instance,
3710            physical_table.to_string(),
3711            "host".to_string(),
3712            vec![],
3713            SystemTime::UNIX_EPOCH,
3714            SystemTime::UNIX_EPOCH,
3715            &ctx,
3716        )
3717        .await
3718        .unwrap_err();
3719        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3720
3721        Ok(())
3722    }
3723
3724    #[tokio::test]
3725    async fn test_non_exact_query_discovery_keeps_denied_targets_for_batch_check() -> TestResult<()>
3726    {
3727        let plugins = Plugins::new();
3728        plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
3729        let instance = test_instance_with_plugins_and_metric_names(
3730            test_logical_table(1024, "denied")?,
3731            test_logical_table(1025, "target")?,
3732            plugins,
3733            Some(test_metric_names_table()),
3734        )
3735        .await?;
3736        let ctx = test_query_ctx(1);
3737
3738        let mut metric_names = PrometheusHandler::query_metric_names(
3739            &instance,
3740            vec![Matcher::new(
3741                promql_parser::label::MatchOp::NotEqual,
3742                "__name__",
3743                "",
3744            )],
3745            "public",
3746            &ctx,
3747        )
3748        .await
3749        .unwrap();
3750        metric_names.sort_unstable();
3751        assert_eq!(
3752            vec!["denied".to_string(), "target".to_string()],
3753            metric_names
3754        );
3755
3756        let queries = metric_names
3757            .into_iter()
3758            .map(|query| PromQuery {
3759                query,
3760                ..Default::default()
3761            })
3762            .collect::<Vec<_>>();
3763        let err = PrometheusHandler::check_query_permission(&instance, &queries, &ctx)
3764            .await
3765            .unwrap_err();
3766        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3767
3768        Ok(())
3769    }
3770
3771    #[test]
3772    fn test_fast_legacy_check_is_read_only() {
3773        let cache = DashMap::new();
3774        cache.insert("metric1".to_string(), true);
3775
3776        let names = vec!["metric1".to_string(), "metric2".to_string()];
3777        assert_eq!(Some(true), fast_legacy_check(&cache, &names).unwrap());
3778        assert!(!cache.contains_key("metric2"));
3779
3780        cache_legacy_mode(&cache, &names, true).unwrap();
3781        assert!(*cache.get("metric2").unwrap().value());
3782        assert!(cache_legacy_mode(&cache, &names, false).is_err());
3783        assert!(*cache.get("metric2").unwrap().value());
3784
3785        let cache_incompatible = DashMap::new();
3786        cache_incompatible.insert("metric1".to_string(), true);
3787        cache_incompatible.insert("metric2".to_string(), false);
3788        assert!(fast_legacy_check(&cache_incompatible, &names).is_err());
3789    }
3790
3791    #[test]
3792    fn test_should_track_statement_process() {
3793        assert!(should_track_statement_process(&parse_one_sql(
3794            "SELECT * FROM demo"
3795        )));
3796        assert!(should_track_statement_process(&parse_one_sql(
3797            "INSERT INTO demo SELECT * FROM source"
3798        )));
3799        assert!(!should_track_statement_process(&parse_one_sql(
3800            "INSERT INTO demo VALUES (1)"
3801        )));
3802        assert!(!should_track_statement_process(&parse_one_sql(
3803            "INSERT INTO demo VALUES (now())"
3804        )));
3805    }
3806
3807    #[test]
3808    fn test_should_track_plan_process() {
3809        let select_stmt = parse_one_sql("SELECT * FROM demo");
3810        let insert_select_stmt = parse_one_sql("INSERT INTO demo SELECT * FROM source");
3811        let insert_values_stmt = parse_one_sql("INSERT INTO demo VALUES (now())");
3812
3813        let empty_plan = LogicalPlanBuilder::empty(false).build().unwrap();
3814        assert!(should_track_plan_process(Some(&select_stmt), &empty_plan));
3815        assert!(should_track_plan_process(
3816            Some(&insert_select_stmt),
3817            &insert_dml_plan()
3818        ));
3819        assert!(!should_track_plan_process(
3820            Some(&insert_values_stmt),
3821            &insert_dml_plan()
3822        ));
3823        assert!(!should_track_plan_process(None, &insert_dml_plan()));
3824    }
3825
3826    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3827    async fn test_insert_select_is_visible_in_show_processlist() -> TestResult<()> {
3828        let insert_sql = "INSERT INTO target SELECT * FROM source";
3829        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
3830        let (finish_tx, finish_rx) = oneshot::channel();
3831        let interceptor = Arc::new(BlockingInsertSelectInterceptor::new(started_tx, finish_rx));
3832        let instance = Arc::new(test_instance_with_insert_select_interceptor(interceptor).await?);
3833
3834        let insert_task = tokio::spawn({
3835            let instance = instance.clone();
3836            async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
3837        });
3838
3839        tokio::time::timeout(Duration::from_secs(5), started_rx.recv())
3840            .await
3841            .context(InsertStartTimeoutSnafu)?
3842            .context(InsertStartChannelClosedSnafu)?;
3843
3844        let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
3845        let process_list = output.data.pretty_print().await;
3846        assert!(
3847            process_list.contains(insert_sql),
3848            "process list did not contain running insert:\n{process_list}"
3849        );
3850
3851        finish_tx
3852            .send(())
3853            .map_err(|_| ReleaseBlockedInsertSnafu.build())?;
3854        insert_task.await.context(InsertTaskPanicSnafu)??;
3855
3856        Ok(())
3857    }
3858
3859    #[tokio::test]
3860    async fn test_show_processlist_catalog_scope() -> TestResult<()> {
3861        let instance =
3862            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
3863                .await?;
3864        let _current_catalog = instance.process_manager().register_query(
3865            "greptime".to_string(),
3866            vec!["public".to_string()],
3867            "current_catalog_query".to_string(),
3868            String::new(),
3869            None,
3870            None,
3871        );
3872        let _other_catalog = instance.process_manager().register_query(
3873            "other".to_string(),
3874            vec!["public".to_string()],
3875            "other_catalog_query".to_string(),
3876            String::new(),
3877            None,
3878            None,
3879        );
3880
3881        for sql in ["SHOW PROCESSLIST", "SHOW FULL PROCESSLIST"] {
3882            let output = execute_one_sql(&instance, sql, test_query_ctx(43)).await?;
3883            let process_list = output.data.pretty_print().await;
3884            assert!(
3885                process_list.contains("current_catalog_query"),
3886                "{process_list}"
3887            );
3888            assert!(
3889                !process_list.contains("other_catalog_query"),
3890                "{process_list}"
3891            );
3892
3893            let admin_ctx = test_query_ctx(44);
3894            admin_ctx.set_current_user(Arc::new(AdminUserInfo));
3895            let output = execute_one_sql(&instance, sql, admin_ctx).await?;
3896            let process_list = output.data.pretty_print().await;
3897            assert!(
3898                process_list.contains("current_catalog_query"),
3899                "{process_list}"
3900            );
3901            assert!(
3902                process_list.contains("other_catalog_query"),
3903                "{process_list}"
3904            );
3905        }
3906
3907        Ok(())
3908    }
3909
3910    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3911    async fn test_kill_query_cancels_insert_select() -> TestResult<()> {
3912        assert_kill_cancels_insert_select("KILL QUERY 4242").await
3913    }
3914
3915    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3916    async fn test_kill_process_id_cancels_insert_select() -> TestResult<()> {
3917        assert_kill_cancels_insert_select("KILL 'test-frontend/4242'").await
3918    }
3919
3920    async fn assert_kill_cancels_insert_select(kill_sql: &str) -> TestResult<()> {
3921        let insert_sql = "INSERT INTO target SELECT * FROM source";
3922        let (source_polled_tx, source_polled_rx) = oneshot::channel();
3923        let instance = Arc::new(
3924            test_instance_with_tables(
3925                pending_table(1024, "source", source_polled_tx)?,
3926                test_table(1025, "target")?,
3927            )
3928            .await?,
3929        );
3930
3931        let insert_task = tokio::spawn({
3932            let instance = instance.clone();
3933            async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
3934        });
3935
3936        tokio::time::timeout(Duration::from_secs(5), source_polled_rx)
3937            .await
3938            .context(SourcePollTimeoutSnafu)?
3939            .context(SourcePollChannelClosedSnafu)?;
3940
3941        let output = execute_one_sql(&instance, kill_sql, test_query_ctx(43)).await?;
3942        assert!(matches!(output.data, OutputData::AffectedRows(1)));
3943
3944        let insert_result = tokio::time::timeout(Duration::from_secs(5), insert_task)
3945            .await
3946            .context(InsertTaskTimeoutSnafu)?
3947            .context(InsertTaskPanicSnafu)?;
3948        let err = match insert_result {
3949            Ok(_) => return InsertSelectNotCancelledSnafu.fail(),
3950            Err(TestError::ExecuteSql { source, .. }) => source,
3951            Err(err) => return Err(err),
3952        };
3953        assert_eq!(StatusCode::Cancelled, err.status_code());
3954
3955        let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
3956        let process_list = output.data.pretty_print().await;
3957        assert!(
3958            !process_list.contains(insert_sql),
3959            "process list still contains killed insert:\n{process_list}"
3960        );
3961
3962        Ok(())
3963    }
3964
3965    fn insert_dml_plan() -> LogicalPlan {
3966        let schema = SchemaRef::new(Schema::new(vec![Field::new(
3967            "value",
3968            DataType::Int64,
3969            true,
3970        )]));
3971        let target = Arc::new(LogicalTableSource::new(schema));
3972        let input = LogicalPlanBuilder::empty(false).build().unwrap();
3973
3974        LogicalPlanBuilder::insert_into(input, "demo", target, InsertOp::Append)
3975            .unwrap()
3976            .build()
3977            .unwrap()
3978    }
3979
3980    #[test]
3981    fn test_exec_validation() {
3982        let query_ctx = QueryContext::arc();
3983        let plugins: Plugins = Plugins::new();
3984        plugins.insert(QueryOptions {
3985            disallow_cross_catalog_query: true,
3986        });
3987
3988        let sql = r#"
3989        SELECT * FROM demo;
3990        EXPLAIN SELECT * FROM demo;
3991        CREATE DATABASE test_database;
3992        SHOW DATABASES;
3993        "#;
3994        let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
3995        assert_eq!(stmts.len(), 4);
3996        for stmt in stmts {
3997            let re = check_permission(plugins.clone(), &stmt, &query_ctx);
3998            re.unwrap();
3999        }
4000
4001        let sql = r#"
4002        SHOW CREATE TABLE demo;
4003        ALTER TABLE demo ADD COLUMN new_col INT;
4004        "#;
4005        let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
4006        assert_eq!(stmts.len(), 2);
4007        for stmt in stmts {
4008            let re = check_permission(plugins.clone(), &stmt, &query_ctx);
4009            re.unwrap();
4010        }
4011
4012        fn replace_test(template_sql: &str, plugins: Plugins, query_ctx: &QueryContextRef) {
4013            // test right
4014            let right = vec![("", ""), ("", "public."), ("greptime.", "public.")];
4015            for (catalog, schema) in right {
4016                let sql = do_fmt(template_sql, catalog, schema);
4017                do_test(&sql, plugins.clone(), query_ctx, true);
4018            }
4019
4020            let wrong = vec![
4021                ("wrongcatalog.", "public."),
4022                ("wrongcatalog.", "wrongschema."),
4023            ];
4024            for (catalog, schema) in wrong {
4025                let sql = do_fmt(template_sql, catalog, schema);
4026                do_test(&sql, plugins.clone(), query_ctx, false);
4027            }
4028        }
4029
4030        fn do_fmt(template: &str, catalog: &str, schema: &str) -> String {
4031            let vars = HashMap::from([
4032                ("catalog".to_string(), catalog),
4033                ("schema".to_string(), schema),
4034            ]);
4035            template.format(&vars).unwrap()
4036        }
4037
4038        fn do_test(sql: &str, plugins: Plugins, query_ctx: &QueryContextRef, is_ok: bool) {
4039            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
4040            let re = check_permission(plugins, stmt, query_ctx);
4041            if is_ok {
4042                re.unwrap();
4043            } else {
4044                assert!(re.is_err());
4045            }
4046        }
4047
4048        // test insert
4049        let sql = "INSERT INTO {catalog}{schema}monitor(host) VALUES ('host1');";
4050        replace_test(sql, plugins.clone(), &query_ctx);
4051
4052        // test create table
4053        let sql = r#"CREATE TABLE {catalog}{schema}demo(
4054                            host STRING,
4055                            ts TIMESTAMP,
4056                            TIME INDEX (ts),
4057                            PRIMARY KEY(host)
4058                        ) engine=mito;"#;
4059        replace_test(sql, plugins.clone(), &query_ctx);
4060
4061        // test drop table
4062        let sql = "DROP TABLE {catalog}{schema}demo;";
4063        replace_test(sql, plugins.clone(), &query_ctx);
4064
4065        // test undrop table
4066        #[cfg(feature = "enterprise")]
4067        {
4068            let sql = "UNDROP TABLE {catalog}{schema}demo;";
4069            replace_test(sql, plugins.clone(), &query_ctx);
4070        }
4071
4072        // test show tables
4073        let sql = "SHOW TABLES FROM public";
4074        let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
4075        check_permission(plugins.clone(), &stmt[0], &query_ctx).unwrap();
4076
4077        let sql = "SHOW TABLES FROM private";
4078        let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
4079        let re = check_permission(plugins.clone(), &stmt[0], &query_ctx);
4080        assert!(re.is_ok());
4081
4082        // test describe table
4083        let sql = "DESC TABLE {catalog}{schema}demo;";
4084        replace_test(sql, plugins.clone(), &query_ctx);
4085
4086        let comment_flow_cases = [
4087            ("COMMENT ON FLOW my_flow IS 'comment';", true),
4088            ("COMMENT ON FLOW greptime.my_flow IS 'comment';", true),
4089            ("COMMENT ON FLOW wrongcatalog.my_flow IS 'comment';", false),
4090        ];
4091        for (sql, is_ok) in comment_flow_cases {
4092            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
4093            let result = check_permission(plugins.clone(), stmt, &query_ctx);
4094            assert_eq!(result.is_ok(), is_ok);
4095        }
4096
4097        let show_flow_cases = [
4098            ("SHOW CREATE FLOW my_flow;", true),
4099            ("SHOW CREATE FLOW greptime.my_flow;", true),
4100            ("SHOW CREATE FLOW wrongcatalog.my_flow;", false),
4101        ];
4102        for (sql, is_ok) in show_flow_cases {
4103            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
4104            let result = check_permission(plugins.clone(), stmt, &query_ctx);
4105            assert_eq!(result.is_ok(), is_ok);
4106        }
4107    }
4108
4109    /// A `DropView` DDL sent through the direct gRPC ingress must return an error
4110    /// (e.g. table not found) instead of panicking on `todo!()`.
4111    #[tokio::test]
4112    async fn qx_152_drop_view_via_grpc_ddl_returns_error_not_panic() -> TestResult<()> {
4113        let instance =
4114            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
4115                .await?;
4116
4117        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4118            expr: Some(api::v1::ddl_request::Expr::DropView(
4119                api::v1::DropViewExpr {
4120                    catalog_name: String::new(),
4121                    schema_name: String::new(),
4122                    view_name: "non_existent_view".to_string(),
4123                    view_id: None,
4124                    drop_if_exists: false,
4125                },
4126            )),
4127        });
4128
4129        let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
4130            &instance,
4131            request,
4132            QueryContext::arc(),
4133        )
4134        .await;
4135
4136        let err = match result {
4137            Ok(_) => panic!("DropView DDL request must return an error instead of panicking"),
4138            Err(err) => err,
4139        };
4140        assert_eq!(
4141            err.status_code(),
4142            StatusCode::TableNotFound,
4143            "dropping a non-existent view without IF EXISTS must report TableNotFound, got {err}"
4144        );
4145        Ok(())
4146    }
4147
4148    /// `DROP VIEW IF EXISTS` on a missing view through the direct gRPC ingress must
4149    /// succeed with 0 affected rows (no error, no DDL task submitted), instead of
4150    /// returning `TableNotFound`.
4151    #[tokio::test]
4152    async fn qx_152_drop_view_if_exists_missing_view_via_grpc_ddl_succeeds() -> TestResult<()> {
4153        let catalog_manager =
4154            catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
4155        let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
4156        let instance = test_instance_with_catalog_manager(
4157            catalog_manager,
4158            test_table(1025, "target")?,
4159            Plugins::new(),
4160            None,
4161            procedure_executor.clone() as ProcedureExecutorRef,
4162        )
4163        .await?;
4164
4165        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4166            expr: Some(api::v1::ddl_request::Expr::DropView(
4167                api::v1::DropViewExpr {
4168                    catalog_name: String::new(),
4169                    schema_name: String::new(),
4170                    view_name: "non_existent_view".to_string(),
4171                    view_id: None,
4172                    drop_if_exists: true,
4173                },
4174            )),
4175        });
4176
4177        let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
4178            &instance,
4179            request,
4180            QueryContext::arc(),
4181        )
4182        .await;
4183
4184        let output = match result {
4185            Ok(output) => output,
4186            Err(err) => {
4187                panic!("DROP VIEW IF EXISTS on a missing view must succeed, got error: {err}")
4188            }
4189        };
4190        assert!(
4191            matches!(output.data, OutputData::AffectedRows(0)),
4192            "DROP VIEW IF EXISTS on a missing view must report 0 affected rows"
4193        );
4194        assert!(
4195            procedure_executor.submitted.lock().unwrap().is_empty(),
4196            "DROP VIEW IF EXISTS on a missing view must not submit a DDL task"
4197        );
4198        Ok(())
4199    }
4200
4201    /// A `CREATE VIEW` followed by `DROP VIEW` through the direct gRPC ingress must
4202    /// succeed end to end: the view is registered in the catalog and then removed.
4203    #[tokio::test]
4204    async fn qx_152_drop_existing_view_via_grpc_ddl_succeeds() -> TestResult<()> {
4205        let catalog_manager =
4206            catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
4207        let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
4208        let instance = test_instance_with_catalog_manager(
4209            catalog_manager,
4210            test_table(1025, "target")?,
4211            Plugins::new(),
4212            None,
4213            procedure_executor.clone() as ProcedureExecutorRef,
4214        )
4215        .await?;
4216
4217        // The default "greptime.public" schema must be visible to the kv-backed table
4218        // metadata manager for `CREATE VIEW`/`CREATE TABLE` to pass validation.
4219        instance
4220            .table_metadata_manager()
4221            .schema_manager()
4222            .create(
4223                common_meta::key::schema_name::SchemaNameKey::new("greptime", "public"),
4224                None,
4225                true,
4226            )
4227            .await
4228            .unwrap();
4229
4230        let create_view_request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4231            expr: Some(api::v1::ddl_request::Expr::CreateView(
4232                api::v1::CreateViewExpr {
4233                    catalog_name: String::new(),
4234                    schema_name: String::new(),
4235                    view_name: "my_view".to_string(),
4236                    logical_plan: vec![1, 2, 3],
4237                    create_if_not_exists: false,
4238                    or_replace: false,
4239                    table_names: vec![],
4240                    columns: vec![],
4241                    plan_columns: vec![],
4242                    definition: "CREATE VIEW my_view AS SELECT * FROM source".to_string(),
4243                },
4244            )),
4245        });
4246
4247        let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
4248            &instance,
4249            create_view_request,
4250            QueryContext::arc(),
4251        )
4252        .await
4253        {
4254            Ok(output) => output,
4255            Err(err) => panic!("CREATE VIEW via gRPC DDL must succeed, got error: {err}"),
4256        };
4257        assert!(
4258            matches!(output.data, OutputData::AffectedRows(0)),
4259            "CREATE VIEW via gRPC DDL must report 0 affected rows"
4260        );
4261
4262        // The view is registered in the catalog as a view.
4263        let view = instance
4264            .catalog_manager()
4265            .table("greptime", "public", "my_view", None)
4266            .await
4267            .unwrap()
4268            .expect("view should exist after CREATE VIEW");
4269        assert_eq!(view.table_info().table_type, TableType::View);
4270
4271        let drop_view_request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4272            expr: Some(api::v1::ddl_request::Expr::DropView(
4273                api::v1::DropViewExpr {
4274                    catalog_name: String::new(),
4275                    schema_name: String::new(),
4276                    view_name: "my_view".to_string(),
4277                    view_id: None,
4278                    drop_if_exists: false,
4279                },
4280            )),
4281        });
4282
4283        let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
4284            &instance,
4285            drop_view_request,
4286            QueryContext::arc(),
4287        )
4288        .await
4289        {
4290            Ok(output) => output,
4291            Err(err) => panic!("DROP VIEW via gRPC DDL must succeed, got error: {err}"),
4292        };
4293        assert!(
4294            matches!(output.data, OutputData::AffectedRows(0)),
4295            "DROP VIEW via gRPC DDL must report 0 affected rows"
4296        );
4297
4298        // The view is gone after the drop.
4299        assert!(
4300            instance
4301                .catalog_manager()
4302                .table("greptime", "public", "my_view", None)
4303                .await
4304                .unwrap()
4305                .is_none(),
4306            "view should be removed after DROP VIEW"
4307        );
4308
4309        let submitted = procedure_executor.submitted.lock().unwrap();
4310        assert_eq!(
4311            submitted.len(),
4312            2,
4313            "expected create and drop view tasks, got {submitted:?}"
4314        );
4315        assert!(matches!(&submitted[0], DdlTask::CreateView(_)));
4316        assert!(matches!(&submitted[1], DdlTask::DropView(_)));
4317        Ok(())
4318    }
4319
4320    /// A direct gRPC `CreateTable` whose time index column is not a timestamp must
4321    /// be rejected with `InvalidArguments` instead of panicking while building the schema.
4322    #[tokio::test]
4323    async fn qx_153_create_table_with_non_timestamp_time_index_via_grpc_returns_error()
4324    -> TestResult<()> {
4325        let instance =
4326            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
4327                .await?;
4328
4329        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4330            expr: Some(api::v1::ddl_request::Expr::CreateTable(
4331                api::v1::CreateTableExpr {
4332                    catalog_name: String::new(),
4333                    schema_name: String::new(),
4334                    table_name: "demo".to_string(),
4335                    desc: String::new(),
4336                    column_defs: vec![api::v1::ColumnDef {
4337                        name: "host".to_string(),
4338                        data_type: api::v1::ColumnDataType::String as i32,
4339                        is_nullable: true,
4340                        default_constraint: vec![],
4341                        semantic_type: 0,
4342                        comment: String::new(),
4343                        datatype_extension: None,
4344                        options: None,
4345                    }],
4346                    time_index: "host".to_string(),
4347                    primary_keys: vec![],
4348                    create_if_not_exists: false,
4349                    table_options: HashMap::new(),
4350                    table_id: None,
4351                    engine: "mito".to_string(),
4352                },
4353            )),
4354        });
4355
4356        let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
4357            &instance,
4358            request,
4359            QueryContext::arc(),
4360        )
4361        .await;
4362
4363        let err = match result {
4364            Ok(_) => panic!("CreateTable with a non-timestamp time index must be rejected"),
4365            Err(err) => err,
4366        };
4367        assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}");
4368        Ok(())
4369    }
4370
4371    /// A valid `CREATE TABLE` (timestamp time index) through the direct gRPC ingress
4372    /// must succeed, guarding that the `validate_create_expr` ingress check doesn't
4373    /// accidentally reject good requests.
4374    #[tokio::test]
4375    async fn qx_153_create_table_with_timestamp_time_index_via_grpc_succeeds() -> TestResult<()> {
4376        let catalog_manager =
4377            catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
4378        let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
4379        let instance = test_instance_with_catalog_manager(
4380            catalog_manager,
4381            test_table(1025, "target")?,
4382            Plugins::new(),
4383            None,
4384            procedure_executor.clone() as ProcedureExecutorRef,
4385        )
4386        .await?;
4387
4388        // The default "greptime.public" schema must be visible to the kv-backed table
4389        // metadata manager for `CREATE TABLE` to pass validation.
4390        instance
4391            .table_metadata_manager()
4392            .schema_manager()
4393            .create(
4394                common_meta::key::schema_name::SchemaNameKey::new("greptime", "public"),
4395                None,
4396                true,
4397            )
4398            .await
4399            .unwrap();
4400
4401        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4402            expr: Some(api::v1::ddl_request::Expr::CreateTable(
4403                api::v1::CreateTableExpr {
4404                    catalog_name: String::new(),
4405                    schema_name: String::new(),
4406                    table_name: "demo".to_string(),
4407                    desc: String::new(),
4408                    column_defs: vec![
4409                        api::v1::ColumnDef {
4410                            name: "host".to_string(),
4411                            data_type: api::v1::ColumnDataType::String as i32,
4412                            is_nullable: true,
4413                            default_constraint: vec![],
4414                            semantic_type: 0,
4415                            comment: String::new(),
4416                            datatype_extension: None,
4417                            options: None,
4418                        },
4419                        api::v1::ColumnDef {
4420                            name: "ts".to_string(),
4421                            data_type: api::v1::ColumnDataType::TimestampMillisecond as i32,
4422                            is_nullable: true,
4423                            default_constraint: vec![],
4424                            semantic_type: 0,
4425                            comment: String::new(),
4426                            datatype_extension: None,
4427                            options: None,
4428                        },
4429                    ],
4430                    time_index: "ts".to_string(),
4431                    primary_keys: vec![],
4432                    create_if_not_exists: false,
4433                    table_options: HashMap::new(),
4434                    table_id: None,
4435                    engine: "mito".to_string(),
4436                },
4437            )),
4438        });
4439
4440        let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
4441            &instance,
4442            request,
4443            QueryContext::arc(),
4444        )
4445        .await
4446        {
4447            Ok(output) => output,
4448            Err(err) => panic!("CREATE TABLE via gRPC DDL must succeed, got error: {err}"),
4449        };
4450        assert!(
4451            matches!(output.data, OutputData::AffectedRows(0)),
4452            "CREATE TABLE via gRPC DDL must report 0 affected rows"
4453        );
4454
4455        // The table is registered in the catalog.
4456        let table = instance
4457            .catalog_manager()
4458            .table("greptime", "public", "demo", None)
4459            .await
4460            .unwrap()
4461            .expect("table should exist after CREATE TABLE");
4462        assert_eq!(table.table_info().table_type, TableType::Base);
4463
4464        let submitted = procedure_executor.submitted.lock().unwrap();
4465        assert_eq!(
4466            submitted.len(),
4467            1,
4468            "expected one create table task, got {submitted:?}"
4469        );
4470        assert!(matches!(&submitted[0], DdlTask::CreateTable(_)));
4471        Ok(())
4472    }
4473}