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