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