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 ReleasedExportSource {
2767        schema: GtSchemaRef,
2768        channels: std::sync::Mutex<Option<(oneshot::Sender<()>, oneshot::Receiver<()>)>>,
2769    }
2770
2771    impl DataSource for ReleasedExportSource {
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_drains_ordinary_writer() {
2792        let destination = common_test_util::temp_dir::create_temp_dir("metric_export_timeout");
2793        let (started_tx, started_rx) = oneshot::channel();
2794        let (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(ReleasedExportSource {
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 sql = format!(
2838            "COPY DATABASE greptime.public TO '{}/' WITH (experimental_metric_export='true')",
2839            destination.path().display()
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        // A dropped ordinary stream would close this receiver before release.
2859        release_tx.send(()).unwrap();
2860        let file = destination.path().join("source.parquet");
2861        tokio::time::timeout(Duration::from_secs(5), async {
2862            loop {
2863                if std::fs::read(&file)
2864                    .is_ok_and(|bytes| bytes.len() > 8 && bytes.ends_with(b"PAR1"))
2865                {
2866                    break;
2867                }
2868                tokio::time::sleep(Duration::from_millis(10)).await;
2869            }
2870        })
2871        .await
2872        .unwrap();
2873        server_task.abort();
2874    }
2875
2876    fn pending_table(
2877        table_id: u32,
2878        table_name: &str,
2879        polled_tx: oneshot::Sender<()>,
2880    ) -> TestResult<table::TableRef> {
2881        let table_info = test_table_info(table_id, table_name)?;
2882        let data_source = Arc::new(PendingDataSource {
2883            schema: table_info.meta.schema.clone(),
2884            polled_tx: std::sync::Mutex::new(Some(polled_tx)),
2885        });
2886
2887        Ok(Arc::new(Table::new(
2888            Arc::new(table_info),
2889            FilterPushDownType::Unsupported,
2890            data_source,
2891        )))
2892    }
2893
2894    async fn test_instance_with_tables(
2895        source_table: TableRef,
2896        target_table: TableRef,
2897    ) -> TestResult<Instance> {
2898        test_instance_with_plugins(source_table, target_table, Plugins::new()).await
2899    }
2900
2901    async fn test_instance_with_insert_select_interceptor(
2902        interceptor: SqlQueryInterceptorRef<Error>,
2903    ) -> TestResult<Instance> {
2904        let plugins = Plugins::new();
2905        plugins.insert::<SqlQueryInterceptorRef<Error>>(interceptor);
2906
2907        test_instance_with_plugins(
2908            test_table(1024, "source")?,
2909            test_table(1025, "target")?,
2910            plugins,
2911        )
2912        .await
2913    }
2914
2915    async fn test_instance_with_plugins(
2916        source_table: TableRef,
2917        target_table: TableRef,
2918        plugins: Plugins,
2919    ) -> TestResult<Instance> {
2920        test_instance_with_plugins_and_metric_names(source_table, target_table, plugins, None).await
2921    }
2922
2923    async fn test_instance_with_plugins_and_metric_names(
2924        source_table: TableRef,
2925        target_table: TableRef,
2926        plugins: Plugins,
2927        metric_names_table: Option<TableRef>,
2928    ) -> TestResult<Instance> {
2929        let catalog_manager = catalog::memory::MemoryCatalogManager::new_with_table(source_table);
2930        test_instance_with_catalog_manager(
2931            catalog_manager,
2932            target_table,
2933            plugins,
2934            metric_names_table,
2935            Arc::new(NoopProcedureExecutor),
2936        )
2937        .await
2938    }
2939
2940    /// Builds a test frontend `Instance` over the given (already source-registered)
2941    /// catalog manager, completing DDL tasks through `procedure_executor`.
2942    async fn test_instance_with_catalog_manager(
2943        catalog_manager: Arc<catalog::memory::MemoryCatalogManager>,
2944        target_table: TableRef,
2945        plugins: Plugins,
2946        metric_names_table: Option<TableRef>,
2947        procedure_executor: ProcedureExecutorRef,
2948    ) -> TestResult<Instance> {
2949        let kv_backend = Arc::new(MemoryKvBackend::new());
2950        let process_manager = Arc::new(ProcessManager::new("test-frontend".to_string(), None));
2951        let target_table_name = "target";
2952        catalog_manager
2953            .register_table_sync(catalog::RegisterTableRequest {
2954                catalog: "greptime".to_string(),
2955                schema: "public".to_string(),
2956                table_name: target_table_name.to_string(),
2957                table_id: 1025,
2958                table: target_table,
2959            })
2960            .with_context(|_| RegisterTableSnafu {
2961                table_name: target_table_name.to_string(),
2962            })?;
2963        if let Some(table) = metric_names_table {
2964            catalog_manager
2965                .deregister_table_sync(catalog::DeregisterTableRequest {
2966                    catalog: "greptime".to_string(),
2967                    schema: "information_schema".to_string(),
2968                    table_name: "tables".to_string(),
2969                })
2970                .unwrap();
2971            catalog_manager
2972                .register_table_sync(catalog::RegisterTableRequest {
2973                    catalog: "greptime".to_string(),
2974                    schema: "information_schema".to_string(),
2975                    table_name: "tables".to_string(),
2976                    table_id: 2048,
2977                    table,
2978                })
2979                .unwrap();
2980        }
2981        catalog_manager.register_process_list_table(process_manager.clone());
2982
2983        let cache_registry = test_cache_registry(kv_backend.clone())?;
2984
2985        FrontendBuilder::new(
2986            FrontendOptions::default(),
2987            kv_backend,
2988            cache_registry,
2989            catalog_manager,
2990            Arc::new(client::client_manager::NodeClients::default()),
2991            procedure_executor,
2992            process_manager,
2993        )
2994        .with_plugin(plugins)
2995        .try_build()
2996        .await
2997        .context(BuildFrontendSnafu)
2998    }
2999
3000    async fn execute_one_sql(
3001        instance: &Instance,
3002        sql: &str,
3003        query_ctx: QueryContextRef,
3004    ) -> TestResult<Output> {
3005        let mut results = instance.do_query_inner(sql, query_ctx).await;
3006        ensure!(
3007            results.len() == 1,
3008            UnexpectedOutputCountSnafu {
3009                sql: sql.to_string(),
3010                actual: results.len(),
3011            }
3012        );
3013        results.remove(0).with_context(|_| ExecuteSqlSnafu {
3014            sql: sql.to_string(),
3015        })
3016    }
3017
3018    fn assert_permission_denied<T>(result: servers::error::Result<T>) {
3019        let err = match result {
3020            Ok(_) => panic!("request should be rejected"),
3021            Err(err) => err,
3022        };
3023        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3024    }
3025
3026    fn assert_action_checked(
3027        checker: &RejectEndpointPermissionChecker,
3028        action: PermissionAction,
3029        targets: Option<PermissionTableTargets>,
3030    ) {
3031        assert_eq!(CheckedAction { action, targets }, checker.take_check());
3032    }
3033
3034    #[tokio::test]
3035    async fn database_export_authorizes_all_tables_before_preparation() -> TestResult<()> {
3036        struct ExportAcl(std::sync::Mutex<Vec<PermissionTableTargets>>);
3037        impl PermissionChecker for ExportAcl {
3038            fn check_permission(
3039                &self,
3040                _: UserInfoRef,
3041                req: PermissionReq,
3042            ) -> auth::error::Result<PermissionResp> {
3043                assert!(matches!(
3044                    req,
3045                    PermissionReq::SqlStatement(Statement::Copy(
3046                        sql::statements::copy::Copy::CopyDatabase(
3047                            sql::statements::copy::CopyDatabase::To(_)
3048                        )
3049                    ))
3050                ));
3051                Ok(PermissionResp::Allow)
3052            }
3053            fn check_permission_with_table_targets(
3054                &self,
3055                _: UserInfoRef,
3056                req: PermissionReq,
3057                targets: PermissionTableTargets,
3058            ) -> auth::error::Result<PermissionResp> {
3059                self.0.lock().unwrap().push(targets.clone());
3060                let PermissionTableTargets::Resolved(tables) = targets else {
3061                    panic!("unresolved export")
3062                };
3063                if req.is_readonly() {
3064                    // The first table is readable; the later table has only write access.
3065                    Ok(if tables.iter().any(|t| t.table == "target") {
3066                        PermissionResp::Reject
3067                    } else {
3068                        PermissionResp::Allow
3069                    })
3070                } else {
3071                    self.check_permission(QueryContext::arc().current_user(), req)
3072                }
3073            }
3074        }
3075        let checker = Arc::new(ExportAcl(Default::default()));
3076        let plugins = Plugins::new();
3077        plugins.insert::<PermissionCheckerRef>(checker.clone());
3078        let instance = test_instance_with_plugins(
3079            test_logical_table(1024, "source")?,
3080            test_table(1025, "target")?,
3081            plugins,
3082        )
3083        .await?;
3084        let req = table::requests::CopyDatabaseRequest {
3085            catalog_name: "greptime".into(),
3086            schema_name: "public".into(),
3087            location: "invalid-destination".into(),
3088            with: Default::default(),
3089            connection: Default::default(),
3090            time_range: None,
3091        };
3092        let result = instance
3093            .export_database_for_test(
3094                req.clone(),
3095                None,
3096                &tokio_util::sync::CancellationToken::new(),
3097                QueryContext::arc(),
3098            )
3099            .await;
3100        assert!(matches!(result, Err(Error::Permission { .. })));
3101        let expected = PermissionTableTargets::resolved(vec![
3102            PermissionTableTarget::new("greptime", "public", "source"),
3103            PermissionTableTarget::new("greptime", "public", "target"),
3104        ]);
3105        assert_eq!(*checker.0.lock().unwrap(), vec![expected.clone(), expected]);
3106        // An empty selection must still check the operation privilege.
3107        instance
3108            .plugins
3109            .map_mut::<PermissionCheckerRef, _, _>(|checker| {
3110                *checker.unwrap() = Arc::new(WriteOnlyPermissionChecker)
3111            });
3112        let result = instance
3113            .export_database_for_test(
3114                req,
3115                Some(&[]),
3116                &tokio_util::sync::CancellationToken::new(),
3117                QueryContext::arc(),
3118            )
3119            .await;
3120        assert!(matches!(result, Err(Error::Permission { .. })));
3121        Ok(())
3122    }
3123
3124    #[tokio::test]
3125    async fn test_prom_remote_read_with_custom_timestamp_and_value_columns() -> TestResult<()> {
3126        let schema = Arc::new(GtSchema::new(vec![
3127            ColumnSchema::new(
3128                "custom_ts",
3129                ConcreteDataType::timestamp_millisecond_datatype(),
3130                false,
3131            )
3132            .with_time_index(true),
3133            ColumnSchema::new("custom_value", ConcreteDataType::float64_datatype(), false),
3134        ]));
3135        let recordbatch = RecordBatch::new(
3136            schema,
3137            vec![
3138                Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as VectorRef,
3139                Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as VectorRef,
3140            ],
3141        )
3142        .unwrap();
3143        let instance = test_instance_with_tables(
3144            MemTable::table("custom_metric", recordbatch),
3145            test_table(1025, "target")?,
3146        )
3147        .await?;
3148
3149        let response = PromStoreProtocolHandler::read(
3150            &instance,
3151            ReadRequest {
3152                queries: vec![RemoteQuery {
3153                    start_timestamp_ms: 1500,
3154                    end_timestamp_ms: 2500,
3155                    matchers: vec![LabelMatcher {
3156                        r#type: PromMatcherType::Eq as i32,
3157                        name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3158                        value: "custom_metric".to_string(),
3159                    }],
3160                    ..Default::default()
3161                }],
3162                ..Default::default()
3163            },
3164            test_query_ctx(1),
3165        )
3166        .await
3167        .unwrap();
3168        let body = servers::prom_store::snappy_decompress(&response.body).unwrap();
3169        let response = ReadResponse::decode(body.as_slice()).unwrap();
3170
3171        assert_eq!(1, response.results.len());
3172        assert_eq!(1, response.results[0].timeseries.len());
3173        let timeseries = &response.results[0].timeseries[0];
3174        assert_eq!(
3175            vec![Label {
3176                name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3177                value: "custom_metric".to_string(),
3178            }],
3179            timeseries.labels
3180        );
3181        assert_eq!(
3182            vec![Sample {
3183                value: 2.0,
3184                timestamp: 2000,
3185            }],
3186            timeseries.samples
3187        );
3188
3189        Ok(())
3190    }
3191
3192    #[tokio::test]
3193    async fn test_prom_remote_read_prefers_default_value_column() -> TestResult<()> {
3194        let schema = Arc::new(GtSchema::new(vec![
3195            ColumnSchema::new(
3196                "custom_ts",
3197                ConcreteDataType::timestamp_millisecond_datatype(),
3198                false,
3199            )
3200            .with_time_index(true),
3201            ColumnSchema::new("extra_field", ConcreteDataType::float64_datatype(), false),
3202            ColumnSchema::new(
3203                greptime_value(),
3204                ConcreteDataType::float64_datatype(),
3205                false,
3206            ),
3207        ]));
3208        let recordbatch = RecordBatch::new(
3209            schema,
3210            vec![
3211                Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as VectorRef,
3212                Arc::new(Float64Vector::from_vec(vec![99.0, 99.0, 99.0])) as VectorRef,
3213                Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as VectorRef,
3214            ],
3215        )
3216        .unwrap();
3217        let instance = test_instance_with_tables(
3218            MemTable::table("multi_field_metric", recordbatch),
3219            test_table(1025, "target")?,
3220        )
3221        .await?;
3222
3223        let response = PromStoreProtocolHandler::read(
3224            &instance,
3225            ReadRequest {
3226                queries: vec![RemoteQuery {
3227                    start_timestamp_ms: 1500,
3228                    end_timestamp_ms: 2500,
3229                    matchers: vec![LabelMatcher {
3230                        r#type: PromMatcherType::Eq as i32,
3231                        name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3232                        value: "multi_field_metric".to_string(),
3233                    }],
3234                    ..Default::default()
3235                }],
3236                ..Default::default()
3237            },
3238            test_query_ctx(1),
3239        )
3240        .await
3241        .unwrap();
3242        let body = servers::prom_store::snappy_decompress(&response.body).unwrap();
3243        let response = ReadResponse::decode(body.as_slice()).unwrap();
3244
3245        assert_eq!(1, response.results.len());
3246        assert_eq!(1, response.results[0].timeseries.len());
3247        let timeseries = &response.results[0].timeseries[0];
3248        assert_eq!(
3249            vec![
3250                Label {
3251                    name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3252                    value: "multi_field_metric".to_string(),
3253                },
3254                Label {
3255                    name: "extra_field".to_string(),
3256                    value: "99".to_string(),
3257                },
3258            ],
3259            timeseries.labels
3260        );
3261        assert_eq!(
3262            vec![Sample {
3263                value: 2.0,
3264                timestamp: 2000,
3265            }],
3266            timeseries.samples
3267        );
3268
3269        Ok(())
3270    }
3271
3272    #[tokio::test]
3273    async fn test_prom_remote_read_rejects_ambiguous_value_columns() -> TestResult<()> {
3274        let schema = Arc::new(GtSchema::new(vec![
3275            ColumnSchema::new(
3276                "custom_ts",
3277                ConcreteDataType::timestamp_millisecond_datatype(),
3278                false,
3279            )
3280            .with_time_index(true),
3281            ColumnSchema::new("field_a", ConcreteDataType::float64_datatype(), false),
3282            ColumnSchema::new("field_b", ConcreteDataType::float64_datatype(), false),
3283        ]));
3284        let recordbatch = RecordBatch::new(
3285            schema,
3286            vec![
3287                Arc::new(TimestampMillisecondVector::from_vec(vec![1000])) as VectorRef,
3288                Arc::new(Float64Vector::from_vec(vec![1.0])) as VectorRef,
3289                Arc::new(Float64Vector::from_vec(vec![2.0])) as VectorRef,
3290            ],
3291        )
3292        .unwrap();
3293        let instance = test_instance_with_tables(
3294            MemTable::table("ambiguous_metric", recordbatch),
3295            test_table(1025, "target")?,
3296        )
3297        .await?;
3298
3299        let err = PromStoreProtocolHandler::read(
3300            &instance,
3301            ReadRequest {
3302                queries: vec![RemoteQuery {
3303                    matchers: vec![LabelMatcher {
3304                        r#type: PromMatcherType::Eq as i32,
3305                        name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3306                        value: "ambiguous_metric".to_string(),
3307                    }],
3308                    ..Default::default()
3309                }],
3310                ..Default::default()
3311            },
3312            test_query_ctx(1),
3313        )
3314        .await
3315        .err()
3316        .expect("ambiguous value columns should fail remote read");
3317
3318        assert_eq!(StatusCode::InvalidArguments, err.status_code());
3319        assert!(format!("{err:?}").contains("Ambiguous value column"));
3320
3321        Ok(())
3322    }
3323
3324    #[tokio::test]
3325    async fn test_event_recorder_is_exposed() -> TestResult<()> {
3326        let instance =
3327            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
3328                .await?;
3329
3330        let _event_recorder = instance.event_recorder();
3331
3332        Ok(())
3333    }
3334
3335    #[tokio::test]
3336    async fn test_restricted_endpoint_handlers_check_permissions() -> TestResult<()> {
3337        let checker = Arc::new(RejectEndpointPermissionChecker::default());
3338        let plugins = Plugins::new();
3339        plugins.insert::<PermissionCheckerRef>(checker.clone());
3340        let instance = test_instance_with_plugins(
3341            test_table(1024, "denied")?,
3342            test_table(1025, "target")?,
3343            plugins,
3344        )
3345        .await?;
3346        let mut ctx = test_query_ctx(1);
3347        Arc::get_mut(&mut ctx).unwrap().set_extension(
3348            servers::http::jaeger::JAEGER_QUERY_TABLE_NAME_KEY,
3349            "denied".to_string(),
3350        );
3351        let jaeger_targets = Some(PermissionTableTargets::resolved(vec![
3352            PermissionTableTarget::new("greptime", "public", "denied"),
3353        ]));
3354
3355        assert_permission_denied(JaegerQueryHandler::get_services(&instance, ctx.clone()).await);
3356        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
3357        assert_permission_denied(
3358            JaegerQueryHandler::get_operations(&instance, ctx.clone(), "service", None).await,
3359        );
3360        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
3361        assert_permission_denied(
3362            JaegerQueryHandler::get_trace(&instance, ctx.clone(), "trace", None, None, None).await,
3363        );
3364        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
3365        assert_permission_denied(
3366            JaegerQueryHandler::find_traces(
3367                &instance,
3368                ctx.clone(),
3369                servers::http::jaeger::QueryTraceParams {
3370                    service_name: "service".to_string(),
3371                    ..Default::default()
3372                },
3373            )
3374            .await,
3375        );
3376        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets);
3377
3378        assert_permission_denied(
3379            PipelineHandler::get_pipeline_str(&instance, "pipeline", None, ctx.clone()).await,
3380        );
3381        assert_action_checked(&checker, PIPELINE_QUERY, None);
3382        assert_permission_denied(
3383            PipelineHandler::insert_pipeline(
3384                &instance,
3385                "pipeline",
3386                "application/yaml",
3387                "",
3388                ctx.clone(),
3389            )
3390            .await,
3391        );
3392        assert_action_checked(&checker, PIPELINE_INSERT, None);
3393        assert_permission_denied(
3394            PipelineHandler::delete_pipeline(&instance, "pipeline", None, ctx.clone()).await,
3395        );
3396        assert_action_checked(&checker, PIPELINE_DELETE, None);
3397        let app = axum::Router::new()
3398            .route(
3399                "/pipelines/_dryrun",
3400                axum::routing::post(servers::http::event::pipeline_dryrun),
3401            )
3402            .with_state(servers::http::event::LogState {
3403                log_handler: Arc::new(instance.clone()),
3404                log_validator: None,
3405                ingest_interceptor: None,
3406            })
3407            .layer(axum::Extension((*ctx).clone()));
3408        let response = app
3409            .oneshot(
3410                axum::http::Request::post("/pipelines/_dryrun")
3411                    .header("content-type", "application/json")
3412                    .body(axum::body::Body::from("{}"))
3413                    .unwrap(),
3414            )
3415            .await
3416            .unwrap();
3417        assert_eq!(axum::http::StatusCode::FORBIDDEN, response.status());
3418        assert_action_checked(&checker, PIPELINE_QUERY, None);
3419
3420        assert_permission_denied(
3421            DashboardHandler::save(&instance, "dashboard", "{}", ctx.clone()).await,
3422        );
3423        assert_action_checked(&checker, DASHBOARD_SAVE, None);
3424        assert_permission_denied(DashboardHandler::list(&instance, ctx.clone()).await);
3425        assert_action_checked(&checker, DASHBOARD_QUERY, None);
3426        assert_permission_denied(
3427            DashboardHandler::delete(&instance, "dashboard", ctx.clone()).await,
3428        );
3429        assert_action_checked(&checker, DASHBOARD_DELETE, None);
3430
3431        Ok(())
3432    }
3433
3434    #[tokio::test]
3435    async fn test_write_only_ingestion_loads_named_pipeline() -> TestResult<()> {
3436        let plugins = Plugins::new();
3437        plugins.insert::<PermissionCheckerRef>(Arc::new(WriteOnlyPermissionChecker));
3438        let instance = test_instance_with_plugins(
3439            test_table(1024, "source")?,
3440            test_table(1025, "target")?,
3441            plugins,
3442        )
3443        .await?;
3444        instance
3445            .catalog_manager()
3446            .as_any()
3447            .downcast_ref::<catalog::memory::MemoryCatalogManager>()
3448            .unwrap()
3449            .register_table_sync(catalog::RegisterTableRequest {
3450                catalog: "greptime".to_string(),
3451                schema: DEFAULT_PRIVATE_SCHEMA_NAME.to_string(),
3452                table_name: "pipelines".to_string(),
3453                table_id: 2049,
3454                table: test_pipeline_table(),
3455            })
3456            .with_context(|_| RegisterTableSnafu {
3457                table_name: "pipelines".to_string(),
3458            })?;
3459        let ctx = test_query_ctx(1);
3460        let handler: PipelineHandlerRef = Arc::new(instance.clone());
3461
3462        handler
3463            .get_pipeline("pipeline", None, ctx.clone())
3464            .await
3465            .unwrap();
3466        assert_permission_denied(
3467            PipelineHandler::get_pipeline_str(&instance, "pipeline", None, ctx.clone()).await,
3468        );
3469
3470        let app = axum::Router::new()
3471            .route(
3472                "/pipelines/_dryrun",
3473                axum::routing::post(servers::http::event::pipeline_dryrun),
3474            )
3475            .with_state(servers::http::event::LogState {
3476                log_handler: handler,
3477                log_validator: None,
3478                ingest_interceptor: None,
3479            })
3480            .layer(axum::Extension((*ctx).clone()));
3481        let response = app
3482            .oneshot(
3483                axum::http::Request::post("/pipelines/_dryrun")
3484                    .header("content-type", "application/json")
3485                    .body(axum::body::Body::from("{}"))
3486                    .unwrap(),
3487            )
3488            .await
3489            .unwrap();
3490        assert_eq!(axum::http::StatusCode::FORBIDDEN, response.status());
3491
3492        Ok(())
3493    }
3494
3495    #[tokio::test]
3496    async fn test_write_only_grpc_sql_is_checked_after_parsing() -> TestResult<()> {
3497        let plugins = Plugins::new();
3498        plugins.insert::<PermissionCheckerRef>(Arc::new(WriteOnlyPermissionChecker));
3499        let instance = test_instance_with_plugins(
3500            test_table(1024, "source")?,
3501            test_table(1025, "target")?,
3502            plugins,
3503        )
3504        .await?;
3505
3506        let insert = Request::Query(api::v1::QueryRequest {
3507            query: Some(Query::Sql(
3508                "INSERT INTO target SELECT * FROM source".to_string(),
3509            )),
3510        });
3511        servers::query_handler::grpc::GrpcQueryHandler::do_query(
3512            &instance,
3513            insert,
3514            QueryContext::arc(),
3515        )
3516        .await
3517        .unwrap();
3518
3519        let select = Request::Query(api::v1::QueryRequest {
3520            query: Some(Query::Sql("SELECT * FROM source".to_string())),
3521        });
3522        assert_permission_denied(
3523            servers::query_handler::grpc::GrpcQueryHandler::do_query(
3524                &instance,
3525                select,
3526                QueryContext::arc(),
3527            )
3528            .await,
3529        );
3530
3531        Ok(())
3532    }
3533
3534    #[tokio::test]
3535    async fn test_target_independent_checker_skips_target_resolution() -> TestResult<()> {
3536        let physical_table = "physical_metric";
3537        let checker = Arc::new(TargetIndependentPermissionChecker::default());
3538        let plugins = Plugins::new();
3539        plugins.insert::<PermissionCheckerRef>(checker.clone());
3540        let instance = test_instance_with_plugins(
3541            test_physical_table(1024, physical_table)?,
3542            test_table(1025, "target")?,
3543            plugins,
3544        )
3545        .await?;
3546
3547        let ctx = test_query_ctx(1);
3548        let physical_target = PermissionTableTarget::new("greptime", "public", physical_table);
3549        assert_eq!(
3550            PermissionTableTargets::Resolved(vec![physical_target.clone()]),
3551            instance
3552                .resolve_query_permission_targets(
3553                    PermissionTableTargets::resolved(vec![physical_target]),
3554                    &ctx,
3555                )
3556                .await
3557                .unwrap()
3558        );
3559        assert_eq!(
3560            vec![physical_table.to_string(), "target".to_string()],
3561            PrometheusHandler::filter_metadata_metric_names(
3562                &instance,
3563                vec![physical_table.to_string(), "target".to_string()],
3564                "public",
3565                &ctx,
3566            )
3567            .await
3568            .unwrap()
3569        );
3570        assert_eq!(1, checker.checks.load(atomic::Ordering::Relaxed));
3571
3572        Ok(())
3573    }
3574
3575    #[tokio::test]
3576    async fn test_query_permission_targets_are_deduplicated() -> TestResult<()> {
3577        let plugins = Plugins::new();
3578        plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
3579        let instance = test_instance_with_plugins(
3580            test_table(1024, "source")?,
3581            test_table(1025, "target")?,
3582            plugins,
3583        )
3584        .await?;
3585        let ctx = test_query_ctx(1);
3586        let target = PermissionTableTarget::new("greptime", "public", "target");
3587
3588        assert_eq!(
3589            PermissionTableTargets::Resolved(vec![target.clone()]),
3590            instance
3591                .resolve_query_permission_targets(
3592                    PermissionTableTargets::resolved(vec![target.clone(), target]),
3593                    &ctx,
3594                )
3595                .await
3596                .unwrap()
3597        );
3598
3599        Ok(())
3600    }
3601
3602    #[tokio::test]
3603    async fn test_physical_query_targets_fail_closed() -> TestResult<()> {
3604        let physical_table = "physical_metric";
3605        let plugins = Plugins::new();
3606        plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
3607        let instance = test_instance_with_plugins(
3608            test_physical_table(1024, physical_table)?,
3609            test_table(1025, "target")?,
3610            plugins,
3611        )
3612        .await?;
3613
3614        let ctx = test_query_ctx(1);
3615        let logical_target = PermissionTableTarget::new("greptime", "public", "target");
3616        assert_eq!(
3617            PermissionTableTargets::Resolved(vec![logical_target.clone()]),
3618            instance
3619                .resolve_query_permission_targets(
3620                    PermissionTableTargets::resolved(vec![logical_target.clone()]),
3621                    &ctx,
3622                )
3623                .await
3624                .unwrap()
3625        );
3626        let physical_target = PermissionTableTarget::new("greptime", "public", physical_table);
3627        assert_eq!(
3628            PermissionTableTargets::Unresolved,
3629            instance
3630                .resolve_query_permission_targets(
3631                    PermissionTableTargets::resolved(
3632                        vec![logical_target, physical_target.clone(),]
3633                    ),
3634                    &ctx,
3635                )
3636                .await
3637                .unwrap()
3638        );
3639        assert_eq!(
3640            vec!["target".to_string()],
3641            PrometheusHandler::filter_metadata_metric_names(
3642                &instance,
3643                vec!["target".to_string(), "denied".to_string()],
3644                "public",
3645                &ctx,
3646            )
3647            .await
3648            .unwrap()
3649        );
3650
3651        let query = PromQuery {
3652            query: physical_table.to_string(),
3653            ..Default::default()
3654        };
3655        let err = PrometheusHandler::check_query_target_permission(
3656            &instance,
3657            PermissionTableTargets::resolved(vec![physical_target]),
3658            &ctx,
3659        )
3660        .await
3661        .unwrap_err();
3662        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3663        let err = PrometheusHandler::check_query_permission(
3664            &instance,
3665            std::slice::from_ref(&query),
3666            &ctx,
3667        )
3668        .await
3669        .unwrap_err();
3670        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3671        let err = PrometheusHandler::do_query(&instance, &query, ctx.clone())
3672            .await
3673            .unwrap_err();
3674        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3675
3676        for sql in [
3677            "SELECT * FROM physical_metric",
3678            "TQL EVAL (0, 10, '5s') physical_metric",
3679            "INSERT INTO target SELECT * FROM physical_metric",
3680        ] {
3681            let mut results = instance.do_query_inner(sql, ctx.clone()).await;
3682            assert_eq!(1, results.len(), "{sql}");
3683            let err = results.remove(0).unwrap_err();
3684            assert_eq!(StatusCode::PermissionDenied, err.status_code(), "{sql}");
3685        }
3686        let err = LogQueryHandler::query(
3687            &instance,
3688            LogQuery {
3689                table: TableName::new("greptime", "public", physical_table),
3690                ..Default::default()
3691            },
3692            ctx.clone(),
3693        )
3694        .await
3695        .unwrap_err();
3696        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3697        let err = instance
3698            .do_describe_inner(parse_one_sql("SELECT * FROM physical_metric"), ctx.clone())
3699            .await
3700            .unwrap_err();
3701        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3702
3703        let request = ReadRequest {
3704            queries: vec![RemoteQuery {
3705                matchers: vec![LabelMatcher {
3706                    r#type: PromMatcherType::Eq as i32,
3707                    name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3708                    value: physical_table.to_string(),
3709                }],
3710                ..Default::default()
3711            }],
3712            ..Default::default()
3713        };
3714        let Err(err) = PromStoreProtocolHandler::read(&instance, request, ctx.clone()).await else {
3715            panic!("physical remote-read target must be rejected");
3716        };
3717        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3718
3719        let err = PrometheusHandler::query_label_values(
3720            &instance,
3721            physical_table.to_string(),
3722            "host".to_string(),
3723            vec![],
3724            SystemTime::UNIX_EPOCH,
3725            SystemTime::UNIX_EPOCH,
3726            &ctx,
3727        )
3728        .await
3729        .unwrap_err();
3730        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3731
3732        Ok(())
3733    }
3734
3735    #[tokio::test]
3736    async fn test_non_exact_query_discovery_keeps_denied_targets_for_batch_check() -> TestResult<()>
3737    {
3738        let plugins = Plugins::new();
3739        plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
3740        let instance = test_instance_with_plugins_and_metric_names(
3741            test_logical_table(1024, "denied")?,
3742            test_logical_table(1025, "target")?,
3743            plugins,
3744            Some(test_metric_names_table()),
3745        )
3746        .await?;
3747        let ctx = test_query_ctx(1);
3748
3749        let mut metric_names = PrometheusHandler::query_metric_names(
3750            &instance,
3751            vec![Matcher::new(
3752                promql_parser::label::MatchOp::NotEqual,
3753                "__name__",
3754                "",
3755            )],
3756            "public",
3757            &ctx,
3758        )
3759        .await
3760        .unwrap();
3761        metric_names.sort_unstable();
3762        assert_eq!(
3763            vec!["denied".to_string(), "target".to_string()],
3764            metric_names
3765        );
3766
3767        let queries = metric_names
3768            .into_iter()
3769            .map(|query| PromQuery {
3770                query,
3771                ..Default::default()
3772            })
3773            .collect::<Vec<_>>();
3774        let err = PrometheusHandler::check_query_permission(&instance, &queries, &ctx)
3775            .await
3776            .unwrap_err();
3777        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3778
3779        Ok(())
3780    }
3781
3782    #[test]
3783    fn test_fast_legacy_check_is_read_only() {
3784        let cache = DashMap::new();
3785        cache.insert("metric1".to_string(), true);
3786
3787        let names = vec!["metric1".to_string(), "metric2".to_string()];
3788        assert_eq!(Some(true), fast_legacy_check(&cache, &names).unwrap());
3789        assert!(!cache.contains_key("metric2"));
3790
3791        cache_legacy_mode(&cache, &names, true).unwrap();
3792        assert!(*cache.get("metric2").unwrap().value());
3793        assert!(cache_legacy_mode(&cache, &names, false).is_err());
3794        assert!(*cache.get("metric2").unwrap().value());
3795
3796        let cache_incompatible = DashMap::new();
3797        cache_incompatible.insert("metric1".to_string(), true);
3798        cache_incompatible.insert("metric2".to_string(), false);
3799        assert!(fast_legacy_check(&cache_incompatible, &names).is_err());
3800    }
3801
3802    #[test]
3803    fn test_should_track_statement_process() {
3804        assert!(should_track_statement_process(&parse_one_sql(
3805            "SELECT * FROM demo"
3806        )));
3807        assert!(should_track_statement_process(&parse_one_sql(
3808            "INSERT INTO demo SELECT * FROM source"
3809        )));
3810        assert!(!should_track_statement_process(&parse_one_sql(
3811            "INSERT INTO demo VALUES (1)"
3812        )));
3813        assert!(!should_track_statement_process(&parse_one_sql(
3814            "INSERT INTO demo VALUES (now())"
3815        )));
3816    }
3817
3818    #[test]
3819    fn test_should_track_plan_process() {
3820        let select_stmt = parse_one_sql("SELECT * FROM demo");
3821        let insert_select_stmt = parse_one_sql("INSERT INTO demo SELECT * FROM source");
3822        let insert_values_stmt = parse_one_sql("INSERT INTO demo VALUES (now())");
3823
3824        let empty_plan = LogicalPlanBuilder::empty(false).build().unwrap();
3825        assert!(should_track_plan_process(Some(&select_stmt), &empty_plan));
3826        assert!(should_track_plan_process(
3827            Some(&insert_select_stmt),
3828            &insert_dml_plan()
3829        ));
3830        assert!(!should_track_plan_process(
3831            Some(&insert_values_stmt),
3832            &insert_dml_plan()
3833        ));
3834        assert!(!should_track_plan_process(None, &insert_dml_plan()));
3835    }
3836
3837    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3838    async fn test_insert_select_is_visible_in_show_processlist() -> TestResult<()> {
3839        let insert_sql = "INSERT INTO target SELECT * FROM source";
3840        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
3841        let (finish_tx, finish_rx) = oneshot::channel();
3842        let interceptor = Arc::new(BlockingInsertSelectInterceptor::new(started_tx, finish_rx));
3843        let instance = Arc::new(test_instance_with_insert_select_interceptor(interceptor).await?);
3844
3845        let insert_task = tokio::spawn({
3846            let instance = instance.clone();
3847            async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
3848        });
3849
3850        tokio::time::timeout(Duration::from_secs(5), started_rx.recv())
3851            .await
3852            .context(InsertStartTimeoutSnafu)?
3853            .context(InsertStartChannelClosedSnafu)?;
3854
3855        let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
3856        let process_list = output.data.pretty_print().await;
3857        assert!(
3858            process_list.contains(insert_sql),
3859            "process list did not contain running insert:\n{process_list}"
3860        );
3861
3862        finish_tx
3863            .send(())
3864            .map_err(|_| ReleaseBlockedInsertSnafu.build())?;
3865        insert_task.await.context(InsertTaskPanicSnafu)??;
3866
3867        Ok(())
3868    }
3869
3870    #[tokio::test]
3871    async fn test_show_processlist_catalog_scope() -> TestResult<()> {
3872        let instance =
3873            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
3874                .await?;
3875        let _current_catalog = instance.process_manager().register_query(
3876            "greptime".to_string(),
3877            vec!["public".to_string()],
3878            "current_catalog_query".to_string(),
3879            String::new(),
3880            None,
3881            None,
3882        );
3883        let _other_catalog = instance.process_manager().register_query(
3884            "other".to_string(),
3885            vec!["public".to_string()],
3886            "other_catalog_query".to_string(),
3887            String::new(),
3888            None,
3889            None,
3890        );
3891
3892        for sql in ["SHOW PROCESSLIST", "SHOW FULL PROCESSLIST"] {
3893            let output = execute_one_sql(&instance, sql, test_query_ctx(43)).await?;
3894            let process_list = output.data.pretty_print().await;
3895            assert!(
3896                process_list.contains("current_catalog_query"),
3897                "{process_list}"
3898            );
3899            assert!(
3900                !process_list.contains("other_catalog_query"),
3901                "{process_list}"
3902            );
3903
3904            let admin_ctx = test_query_ctx(44);
3905            admin_ctx.set_current_user(Arc::new(AdminUserInfo));
3906            let output = execute_one_sql(&instance, sql, admin_ctx).await?;
3907            let process_list = output.data.pretty_print().await;
3908            assert!(
3909                process_list.contains("current_catalog_query"),
3910                "{process_list}"
3911            );
3912            assert!(
3913                process_list.contains("other_catalog_query"),
3914                "{process_list}"
3915            );
3916        }
3917
3918        Ok(())
3919    }
3920
3921    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3922    async fn test_kill_query_cancels_insert_select() -> TestResult<()> {
3923        assert_kill_cancels_insert_select("KILL QUERY 4242").await
3924    }
3925
3926    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3927    async fn test_kill_process_id_cancels_insert_select() -> TestResult<()> {
3928        assert_kill_cancels_insert_select("KILL 'test-frontend/4242'").await
3929    }
3930
3931    async fn assert_kill_cancels_insert_select(kill_sql: &str) -> TestResult<()> {
3932        let insert_sql = "INSERT INTO target SELECT * FROM source";
3933        let (source_polled_tx, source_polled_rx) = oneshot::channel();
3934        let instance = Arc::new(
3935            test_instance_with_tables(
3936                pending_table(1024, "source", source_polled_tx)?,
3937                test_table(1025, "target")?,
3938            )
3939            .await?,
3940        );
3941
3942        let insert_task = tokio::spawn({
3943            let instance = instance.clone();
3944            async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
3945        });
3946
3947        tokio::time::timeout(Duration::from_secs(5), source_polled_rx)
3948            .await
3949            .context(SourcePollTimeoutSnafu)?
3950            .context(SourcePollChannelClosedSnafu)?;
3951
3952        let output = execute_one_sql(&instance, kill_sql, test_query_ctx(43)).await?;
3953        assert!(matches!(output.data, OutputData::AffectedRows(1)));
3954
3955        let insert_result = tokio::time::timeout(Duration::from_secs(5), insert_task)
3956            .await
3957            .context(InsertTaskTimeoutSnafu)?
3958            .context(InsertTaskPanicSnafu)?;
3959        let err = match insert_result {
3960            Ok(_) => return InsertSelectNotCancelledSnafu.fail(),
3961            Err(TestError::ExecuteSql { source, .. }) => source,
3962            Err(err) => return Err(err),
3963        };
3964        assert_eq!(StatusCode::Cancelled, err.status_code());
3965
3966        let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
3967        let process_list = output.data.pretty_print().await;
3968        assert!(
3969            !process_list.contains(insert_sql),
3970            "process list still contains killed insert:\n{process_list}"
3971        );
3972
3973        Ok(())
3974    }
3975
3976    fn insert_dml_plan() -> LogicalPlan {
3977        let schema = SchemaRef::new(Schema::new(vec![Field::new(
3978            "value",
3979            DataType::Int64,
3980            true,
3981        )]));
3982        let target = Arc::new(LogicalTableSource::new(schema));
3983        let input = LogicalPlanBuilder::empty(false).build().unwrap();
3984
3985        LogicalPlanBuilder::insert_into(input, "demo", target, InsertOp::Append)
3986            .unwrap()
3987            .build()
3988            .unwrap()
3989    }
3990
3991    #[test]
3992    fn test_exec_validation() {
3993        let query_ctx = QueryContext::arc();
3994        let plugins: Plugins = Plugins::new();
3995        plugins.insert(QueryOptions {
3996            disallow_cross_catalog_query: true,
3997        });
3998
3999        let sql = r#"
4000        SELECT * FROM demo;
4001        EXPLAIN SELECT * FROM demo;
4002        CREATE DATABASE test_database;
4003        SHOW DATABASES;
4004        "#;
4005        let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
4006        assert_eq!(stmts.len(), 4);
4007        for stmt in stmts {
4008            let re = check_permission(plugins.clone(), &stmt, &query_ctx);
4009            re.unwrap();
4010        }
4011
4012        let sql = r#"
4013        SHOW CREATE TABLE demo;
4014        ALTER TABLE demo ADD COLUMN new_col INT;
4015        "#;
4016        let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
4017        assert_eq!(stmts.len(), 2);
4018        for stmt in stmts {
4019            let re = check_permission(plugins.clone(), &stmt, &query_ctx);
4020            re.unwrap();
4021        }
4022
4023        fn replace_test(template_sql: &str, plugins: Plugins, query_ctx: &QueryContextRef) {
4024            // test right
4025            let right = vec![("", ""), ("", "public."), ("greptime.", "public.")];
4026            for (catalog, schema) in right {
4027                let sql = do_fmt(template_sql, catalog, schema);
4028                do_test(&sql, plugins.clone(), query_ctx, true);
4029            }
4030
4031            let wrong = vec![
4032                ("wrongcatalog.", "public."),
4033                ("wrongcatalog.", "wrongschema."),
4034            ];
4035            for (catalog, schema) in wrong {
4036                let sql = do_fmt(template_sql, catalog, schema);
4037                do_test(&sql, plugins.clone(), query_ctx, false);
4038            }
4039        }
4040
4041        fn do_fmt(template: &str, catalog: &str, schema: &str) -> String {
4042            let vars = HashMap::from([
4043                ("catalog".to_string(), catalog),
4044                ("schema".to_string(), schema),
4045            ]);
4046            template.format(&vars).unwrap()
4047        }
4048
4049        fn do_test(sql: &str, plugins: Plugins, query_ctx: &QueryContextRef, is_ok: bool) {
4050            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
4051            let re = check_permission(plugins, stmt, query_ctx);
4052            if is_ok {
4053                re.unwrap();
4054            } else {
4055                assert!(re.is_err());
4056            }
4057        }
4058
4059        // test insert
4060        let sql = "INSERT INTO {catalog}{schema}monitor(host) VALUES ('host1');";
4061        replace_test(sql, plugins.clone(), &query_ctx);
4062
4063        // test create table
4064        let sql = r#"CREATE TABLE {catalog}{schema}demo(
4065                            host STRING,
4066                            ts TIMESTAMP,
4067                            TIME INDEX (ts),
4068                            PRIMARY KEY(host)
4069                        ) engine=mito;"#;
4070        replace_test(sql, plugins.clone(), &query_ctx);
4071
4072        // test drop table
4073        let sql = "DROP TABLE {catalog}{schema}demo;";
4074        replace_test(sql, plugins.clone(), &query_ctx);
4075
4076        // test undrop table
4077        #[cfg(feature = "enterprise")]
4078        {
4079            let sql = "UNDROP TABLE {catalog}{schema}demo;";
4080            replace_test(sql, plugins.clone(), &query_ctx);
4081        }
4082
4083        // test show tables
4084        let sql = "SHOW TABLES FROM public";
4085        let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
4086        check_permission(plugins.clone(), &stmt[0], &query_ctx).unwrap();
4087
4088        let sql = "SHOW TABLES FROM private";
4089        let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
4090        let re = check_permission(plugins.clone(), &stmt[0], &query_ctx);
4091        assert!(re.is_ok());
4092
4093        // test describe table
4094        let sql = "DESC TABLE {catalog}{schema}demo;";
4095        replace_test(sql, plugins.clone(), &query_ctx);
4096
4097        let comment_flow_cases = [
4098            ("COMMENT ON FLOW my_flow IS 'comment';", true),
4099            ("COMMENT ON FLOW greptime.my_flow IS 'comment';", true),
4100            ("COMMENT ON FLOW wrongcatalog.my_flow IS 'comment';", false),
4101        ];
4102        for (sql, is_ok) in comment_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        let show_flow_cases = [
4109            ("SHOW CREATE FLOW my_flow;", true),
4110            ("SHOW CREATE FLOW greptime.my_flow;", true),
4111            ("SHOW CREATE FLOW wrongcatalog.my_flow;", false),
4112        ];
4113        for (sql, is_ok) in show_flow_cases {
4114            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
4115            let result = check_permission(plugins.clone(), stmt, &query_ctx);
4116            assert_eq!(result.is_ok(), is_ok);
4117        }
4118    }
4119
4120    /// A `DropView` DDL sent through the direct gRPC ingress must return an error
4121    /// (e.g. table not found) instead of panicking on `todo!()`.
4122    #[tokio::test]
4123    async fn qx_152_drop_view_via_grpc_ddl_returns_error_not_panic() -> TestResult<()> {
4124        let instance =
4125            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
4126                .await?;
4127
4128        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4129            expr: Some(api::v1::ddl_request::Expr::DropView(
4130                api::v1::DropViewExpr {
4131                    catalog_name: String::new(),
4132                    schema_name: String::new(),
4133                    view_name: "non_existent_view".to_string(),
4134                    view_id: None,
4135                    drop_if_exists: false,
4136                },
4137            )),
4138        });
4139
4140        let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
4141            &instance,
4142            request,
4143            QueryContext::arc(),
4144        )
4145        .await;
4146
4147        let err = match result {
4148            Ok(_) => panic!("DropView DDL request must return an error instead of panicking"),
4149            Err(err) => err,
4150        };
4151        assert_eq!(
4152            err.status_code(),
4153            StatusCode::TableNotFound,
4154            "dropping a non-existent view without IF EXISTS must report TableNotFound, got {err}"
4155        );
4156        Ok(())
4157    }
4158
4159    /// `DROP VIEW IF EXISTS` on a missing view through the direct gRPC ingress must
4160    /// succeed with 0 affected rows (no error, no DDL task submitted), instead of
4161    /// returning `TableNotFound`.
4162    #[tokio::test]
4163    async fn qx_152_drop_view_if_exists_missing_view_via_grpc_ddl_succeeds() -> TestResult<()> {
4164        let catalog_manager =
4165            catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
4166        let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
4167        let instance = test_instance_with_catalog_manager(
4168            catalog_manager,
4169            test_table(1025, "target")?,
4170            Plugins::new(),
4171            None,
4172            procedure_executor.clone() as ProcedureExecutorRef,
4173        )
4174        .await?;
4175
4176        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4177            expr: Some(api::v1::ddl_request::Expr::DropView(
4178                api::v1::DropViewExpr {
4179                    catalog_name: String::new(),
4180                    schema_name: String::new(),
4181                    view_name: "non_existent_view".to_string(),
4182                    view_id: None,
4183                    drop_if_exists: true,
4184                },
4185            )),
4186        });
4187
4188        let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
4189            &instance,
4190            request,
4191            QueryContext::arc(),
4192        )
4193        .await;
4194
4195        let output = match result {
4196            Ok(output) => output,
4197            Err(err) => {
4198                panic!("DROP VIEW IF EXISTS on a missing view must succeed, got error: {err}")
4199            }
4200        };
4201        assert!(
4202            matches!(output.data, OutputData::AffectedRows(0)),
4203            "DROP VIEW IF EXISTS on a missing view must report 0 affected rows"
4204        );
4205        assert!(
4206            procedure_executor.submitted.lock().unwrap().is_empty(),
4207            "DROP VIEW IF EXISTS on a missing view must not submit a DDL task"
4208        );
4209        Ok(())
4210    }
4211
4212    /// A `CREATE VIEW` followed by `DROP VIEW` through the direct gRPC ingress must
4213    /// succeed end to end: the view is registered in the catalog and then removed.
4214    #[tokio::test]
4215    async fn qx_152_drop_existing_view_via_grpc_ddl_succeeds() -> TestResult<()> {
4216        let catalog_manager =
4217            catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
4218        let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
4219        let instance = test_instance_with_catalog_manager(
4220            catalog_manager,
4221            test_table(1025, "target")?,
4222            Plugins::new(),
4223            None,
4224            procedure_executor.clone() as ProcedureExecutorRef,
4225        )
4226        .await?;
4227
4228        // The default "greptime.public" schema must be visible to the kv-backed table
4229        // metadata manager for `CREATE VIEW`/`CREATE TABLE` to pass validation.
4230        instance
4231            .table_metadata_manager()
4232            .schema_manager()
4233            .create(
4234                common_meta::key::schema_name::SchemaNameKey::new("greptime", "public"),
4235                None,
4236                true,
4237            )
4238            .await
4239            .unwrap();
4240
4241        let create_view_request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4242            expr: Some(api::v1::ddl_request::Expr::CreateView(
4243                api::v1::CreateViewExpr {
4244                    catalog_name: String::new(),
4245                    schema_name: String::new(),
4246                    view_name: "my_view".to_string(),
4247                    logical_plan: vec![1, 2, 3],
4248                    create_if_not_exists: false,
4249                    or_replace: false,
4250                    table_names: vec![],
4251                    columns: vec![],
4252                    plan_columns: vec![],
4253                    definition: "CREATE VIEW my_view AS SELECT * FROM source".to_string(),
4254                },
4255            )),
4256        });
4257
4258        let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
4259            &instance,
4260            create_view_request,
4261            QueryContext::arc(),
4262        )
4263        .await
4264        {
4265            Ok(output) => output,
4266            Err(err) => panic!("CREATE VIEW via gRPC DDL must succeed, got error: {err}"),
4267        };
4268        assert!(
4269            matches!(output.data, OutputData::AffectedRows(0)),
4270            "CREATE VIEW via gRPC DDL must report 0 affected rows"
4271        );
4272
4273        // The view is registered in the catalog as a view.
4274        let view = instance
4275            .catalog_manager()
4276            .table("greptime", "public", "my_view", None)
4277            .await
4278            .unwrap()
4279            .expect("view should exist after CREATE VIEW");
4280        assert_eq!(view.table_info().table_type, TableType::View);
4281
4282        let drop_view_request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4283            expr: Some(api::v1::ddl_request::Expr::DropView(
4284                api::v1::DropViewExpr {
4285                    catalog_name: String::new(),
4286                    schema_name: String::new(),
4287                    view_name: "my_view".to_string(),
4288                    view_id: None,
4289                    drop_if_exists: false,
4290                },
4291            )),
4292        });
4293
4294        let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
4295            &instance,
4296            drop_view_request,
4297            QueryContext::arc(),
4298        )
4299        .await
4300        {
4301            Ok(output) => output,
4302            Err(err) => panic!("DROP VIEW via gRPC DDL must succeed, got error: {err}"),
4303        };
4304        assert!(
4305            matches!(output.data, OutputData::AffectedRows(0)),
4306            "DROP VIEW via gRPC DDL must report 0 affected rows"
4307        );
4308
4309        // The view is gone after the drop.
4310        assert!(
4311            instance
4312                .catalog_manager()
4313                .table("greptime", "public", "my_view", None)
4314                .await
4315                .unwrap()
4316                .is_none(),
4317            "view should be removed after DROP VIEW"
4318        );
4319
4320        let submitted = procedure_executor.submitted.lock().unwrap();
4321        assert_eq!(
4322            submitted.len(),
4323            2,
4324            "expected create and drop view tasks, got {submitted:?}"
4325        );
4326        assert!(matches!(&submitted[0], DdlTask::CreateView(_)));
4327        assert!(matches!(&submitted[1], DdlTask::DropView(_)));
4328        Ok(())
4329    }
4330
4331    /// A direct gRPC `CreateTable` whose time index column is not a timestamp must
4332    /// be rejected with `InvalidArguments` instead of panicking while building the schema.
4333    #[tokio::test]
4334    async fn qx_153_create_table_with_non_timestamp_time_index_via_grpc_returns_error()
4335    -> TestResult<()> {
4336        let instance =
4337            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
4338                .await?;
4339
4340        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4341            expr: Some(api::v1::ddl_request::Expr::CreateTable(
4342                api::v1::CreateTableExpr {
4343                    catalog_name: String::new(),
4344                    schema_name: String::new(),
4345                    table_name: "demo".to_string(),
4346                    desc: String::new(),
4347                    column_defs: vec![api::v1::ColumnDef {
4348                        name: "host".to_string(),
4349                        data_type: api::v1::ColumnDataType::String as i32,
4350                        is_nullable: true,
4351                        default_constraint: vec![],
4352                        semantic_type: 0,
4353                        comment: String::new(),
4354                        datatype_extension: None,
4355                        options: None,
4356                    }],
4357                    time_index: "host".to_string(),
4358                    primary_keys: vec![],
4359                    create_if_not_exists: false,
4360                    table_options: HashMap::new(),
4361                    table_id: None,
4362                    engine: "mito".to_string(),
4363                },
4364            )),
4365        });
4366
4367        let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
4368            &instance,
4369            request,
4370            QueryContext::arc(),
4371        )
4372        .await;
4373
4374        let err = match result {
4375            Ok(_) => panic!("CreateTable with a non-timestamp time index must be rejected"),
4376            Err(err) => err,
4377        };
4378        assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}");
4379        Ok(())
4380    }
4381
4382    /// A valid `CREATE TABLE` (timestamp time index) through the direct gRPC ingress
4383    /// must succeed, guarding that the `validate_create_expr` ingress check doesn't
4384    /// accidentally reject good requests.
4385    #[tokio::test]
4386    async fn qx_153_create_table_with_timestamp_time_index_via_grpc_succeeds() -> TestResult<()> {
4387        let catalog_manager =
4388            catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
4389        let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
4390        let instance = test_instance_with_catalog_manager(
4391            catalog_manager,
4392            test_table(1025, "target")?,
4393            Plugins::new(),
4394            None,
4395            procedure_executor.clone() as ProcedureExecutorRef,
4396        )
4397        .await?;
4398
4399        // The default "greptime.public" schema must be visible to the kv-backed table
4400        // metadata manager for `CREATE TABLE` to pass validation.
4401        instance
4402            .table_metadata_manager()
4403            .schema_manager()
4404            .create(
4405                common_meta::key::schema_name::SchemaNameKey::new("greptime", "public"),
4406                None,
4407                true,
4408            )
4409            .await
4410            .unwrap();
4411
4412        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4413            expr: Some(api::v1::ddl_request::Expr::CreateTable(
4414                api::v1::CreateTableExpr {
4415                    catalog_name: String::new(),
4416                    schema_name: String::new(),
4417                    table_name: "demo".to_string(),
4418                    desc: String::new(),
4419                    column_defs: vec![
4420                        api::v1::ColumnDef {
4421                            name: "host".to_string(),
4422                            data_type: api::v1::ColumnDataType::String as i32,
4423                            is_nullable: true,
4424                            default_constraint: vec![],
4425                            semantic_type: 0,
4426                            comment: String::new(),
4427                            datatype_extension: None,
4428                            options: None,
4429                        },
4430                        api::v1::ColumnDef {
4431                            name: "ts".to_string(),
4432                            data_type: api::v1::ColumnDataType::TimestampMillisecond as i32,
4433                            is_nullable: true,
4434                            default_constraint: vec![],
4435                            semantic_type: 0,
4436                            comment: String::new(),
4437                            datatype_extension: None,
4438                            options: None,
4439                        },
4440                    ],
4441                    time_index: "ts".to_string(),
4442                    primary_keys: vec![],
4443                    create_if_not_exists: false,
4444                    table_options: HashMap::new(),
4445                    table_id: None,
4446                    engine: "mito".to_string(),
4447                },
4448            )),
4449        });
4450
4451        let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
4452            &instance,
4453            request,
4454            QueryContext::arc(),
4455        )
4456        .await
4457        {
4458            Ok(output) => output,
4459            Err(err) => panic!("CREATE TABLE via gRPC DDL must succeed, got error: {err}"),
4460        };
4461        assert!(
4462            matches!(output.data, OutputData::AffectedRows(0)),
4463            "CREATE TABLE via gRPC DDL must report 0 affected rows"
4464        );
4465
4466        // The table is registered in the catalog.
4467        let table = instance
4468            .catalog_manager()
4469            .table("greptime", "public", "demo", None)
4470            .await
4471            .unwrap()
4472            .expect("table should exist after CREATE TABLE");
4473        assert_eq!(table.table_info().table_type, TableType::Base);
4474
4475        let submitted = procedure_executor.submitted.lock().unwrap();
4476        assert_eq!(
4477            submitted.len(),
4478            1,
4479            "expected one create table task, got {submitted:?}"
4480        );
4481        assert!(matches!(&submitted[0], DdlTask::CreateTable(_)));
4482        Ok(())
4483    }
4484}