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 requires_output_ordering = query.requires_output_ordering();
1333        let (query, stmt) = query.into_parts();
1334
1335        let QueryStatement::Promql(eval_stmt, _) = &stmt else {
1336            unreachable!("query is parsed from promql");
1337        };
1338
1339        let plan = self
1340            .statement_executor
1341            .plan(&stmt, query_ctx.clone())
1342            .await
1343            .map_err(BoxedError::new)
1344            .context(ExecuteQuerySnafu)?;
1345
1346        let plan = if requires_output_ordering {
1347            plan
1348        } else {
1349            promql::remove_output_sort(plan)
1350        };
1351
1352        interceptor.pre_execute(&query, &eval_stmt.expr, Some(&plan), query_ctx.clone())?;
1353
1354        // Take the EvalStmt from the original QueryStatement and use it to create the CatalogQueryStatement.
1355        let query_statement = if let QueryStatement::Promql(eval_stmt, alias) = stmt {
1356            CatalogQueryStatement::Promql(eval_stmt, alias)
1357        } else {
1358            // It should not happen since the query is already parsed successfully.
1359            return UnexpectedResultSnafu {
1360                reason: "The query should always be promql.".to_string(),
1361            }
1362            .fail();
1363        };
1364        let raw_query = query_statement.to_string();
1365
1366        let slow_query_timer = self.slow_query_options.enable.then(|| {
1367            SlowQueryTimer::new(
1368                query_statement,
1369                query_ctx.current_schema(),
1370                self.slow_query_options.threshold,
1371                self.slow_query_options.sample_ratio,
1372                self.slow_query_options.record_type,
1373                self.slow_query_recorder.clone(),
1374            )
1375        });
1376
1377        let ticket = self.process_manager.register_query(
1378            query_ctx.current_catalog().to_string(),
1379            vec![query_ctx.current_schema()],
1380            raw_query,
1381            query_ctx.conn_info().to_string(),
1382            Some(query_ctx.process_id()),
1383            slow_query_timer,
1384        );
1385
1386        let query_fut = self.statement_executor.exec_plan(plan, query_ctx.clone());
1387
1388        let output = CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
1389            .await
1390            .map_err(|_| servers::error::CancelledSnafu.build())?
1391            .map_err(BoxedError::new)
1392            .context(ExecuteQuerySnafu)?;
1393        let output = map_query_output(output)
1394            .map_err(BoxedError::new)
1395            .context(ExecuteQuerySnafu)?;
1396        let Output { meta, data } = output;
1397        let data = match data {
1398            OutputData::Stream(stream) => {
1399                OutputData::Stream(Box::pin(CancellableStreamWrapper::new(stream, ticket)))
1400            }
1401            other => other,
1402        };
1403        let output = Output { data, meta };
1404        Ok(interceptor.post_execute(output, query_ctx)?)
1405    }
1406
1407    async fn check_query_permission(
1408        &self,
1409        queries: &[PromQuery],
1410        query_ctx: &QueryContextRef,
1411    ) -> server_error::Result<()> {
1412        let queries = queries
1413            .iter()
1414            .cloned()
1415            .map(|query| ParsedPromQuery::parse(query, query_ctx))
1416            .collect::<server_error::Result<Vec<_>>>()?;
1417        self.check_query_permission_parsed(&queries, query_ctx)
1418            .await
1419    }
1420
1421    async fn check_query_permission_parsed(
1422        &self,
1423        queries: &[ParsedPromQuery],
1424        query_ctx: &QueryContextRef,
1425    ) -> server_error::Result<()> {
1426        self.check_prom_query_privilege(query_ctx)?;
1427        let targets = self.prom_queries_permission_targets(queries, query_ctx)?;
1428        self.check_query_target_permission(targets, query_ctx).await
1429    }
1430
1431    async fn check_query_target_permission(
1432        &self,
1433        targets: PermissionTableTargets,
1434        query_ctx: &QueryContextRef,
1435    ) -> server_error::Result<()> {
1436        let targets = self
1437            .resolve_query_permission_targets(targets, query_ctx)
1438            .await?;
1439        self.check_table_permission(query_ctx, PermissionReq::Action(PROMQL_QUERY), targets)
1440            .context(AuthSnafu)?;
1441        Ok(())
1442    }
1443
1444    async fn filter_metadata_metric_names(
1445        &self,
1446        metric_names: Vec<String>,
1447        schema: &str,
1448        query_ctx: &QueryContextRef,
1449    ) -> server_error::Result<Vec<String>> {
1450        let checker = self.plugins.get::<PermissionCheckerRef>();
1451        if !checker.as_ref().uses_table_targets() {
1452            let Some(metric) = metric_names.first() else {
1453                return Ok(metric_names);
1454            };
1455            let target =
1456                PermissionTableTarget::new(query_ctx.current_catalog(), schema, metric.as_str());
1457            let result = checker
1458                .as_ref()
1459                .check_permission_with_table_targets(
1460                    query_ctx.current_user(),
1461                    PermissionReq::Action(PROMQL_QUERY),
1462                    PermissionTableTargets::resolved(vec![target]),
1463                )
1464                .context(AuthSnafu);
1465            return match result {
1466                Ok(_) => Ok(metric_names),
1467                Err(error)
1468                    if error.status_code()
1469                        == common_error::status_code::StatusCode::PermissionDenied =>
1470                {
1471                    Ok(Vec::new())
1472                }
1473                Err(error) => Err(error),
1474            };
1475        }
1476
1477        let mut allowed = Vec::with_capacity(metric_names.len());
1478        for metric in metric_names {
1479            let target =
1480                PermissionTableTarget::new(query_ctx.current_catalog(), schema, metric.as_str());
1481            match checker
1482                .as_ref()
1483                .check_permission_with_table_targets(
1484                    query_ctx.current_user(),
1485                    PermissionReq::Action(PROMQL_QUERY),
1486                    PermissionTableTargets::resolved(vec![target]),
1487                )
1488                .context(AuthSnafu)
1489            {
1490                Ok(_) => allowed.push(metric),
1491                Err(error)
1492                    if error.status_code()
1493                        == common_error::status_code::StatusCode::PermissionDenied => {}
1494                Err(error) => return Err(error),
1495            }
1496        }
1497        Ok(allowed)
1498    }
1499
1500    async fn query_metric_names(
1501        &self,
1502        matchers: Vec<Matcher>,
1503        schema: &str,
1504        ctx: &QueryContextRef,
1505    ) -> server_error::Result<Vec<String>> {
1506        self.handle_query_metric_names(matchers, schema, ctx)
1507            .await
1508            .map_err(BoxedError::new)
1509            .context(ExecuteQuerySnafu)
1510    }
1511
1512    async fn query_label_values(
1513        &self,
1514        metric: String,
1515        label_name: String,
1516        matchers: Vec<Matcher>,
1517        start: SystemTime,
1518        end: SystemTime,
1519        ctx: &QueryContextRef,
1520    ) -> server_error::Result<Vec<String>> {
1521        let schema =
1522            resolve_schema_from_matchers(&matchers)?.unwrap_or_else(|| ctx.current_schema());
1523        let target = PermissionTableTarget::new(ctx.current_catalog(), schema.as_str(), &metric);
1524        self.check_query_target_permission(
1525            PermissionTableTargets::resolved(vec![target.clone()]),
1526            ctx,
1527        )
1528        .await?;
1529
1530        self.handle_query_label_values(target, label_name, matchers, start, end, ctx)
1531            .await
1532            .map_err(BoxedError::new)
1533            .context(ExecuteQuerySnafu)
1534    }
1535
1536    fn catalog_manager(&self) -> CatalogManagerRef {
1537        self.catalog_manager.clone()
1538    }
1539}
1540
1541/// Validate `stmt.database` permission if it's presented.
1542macro_rules! validate_db_permission {
1543    ($stmt: expr, $query_ctx: expr) => {
1544        if let Some(database) = &$stmt.database {
1545            validate_catalog_and_schema($query_ctx.current_catalog(), database, $query_ctx)
1546                .map_err(BoxedError::new)
1547                .context(SqlExecInterceptedSnafu)?;
1548        }
1549    };
1550}
1551
1552pub fn check_permission(
1553    plugins: Plugins,
1554    stmt: &Statement,
1555    query_ctx: &QueryContextRef,
1556) -> Result<()> {
1557    let need_validate = plugins
1558        .get::<QueryOptions>()
1559        .map(|opts| opts.disallow_cross_catalog_query)
1560        .unwrap_or_default();
1561
1562    if !need_validate {
1563        return Ok(());
1564    }
1565
1566    match stmt {
1567        // Will be checked in execution.
1568        // TODO(dennis): add a hook for admin commands.
1569        Statement::Admin(_) => {}
1570        // These are executed by query engine, and will be checked there.
1571        Statement::Query(_)
1572        | Statement::Explain(_)
1573        | Statement::Tql(_)
1574        | Statement::Delete(_)
1575        | Statement::DeclareCursor(_)
1576        | Statement::Copy(sql::statements::copy::Copy::CopyQueryTo(_)) => {}
1577        // database ops won't be checked
1578        Statement::CreateDatabase(_)
1579        | Statement::ShowDatabases(_)
1580        | Statement::DropDatabase(_)
1581        | Statement::AlterDatabase(_)
1582        | Statement::DropFlow(_)
1583        | Statement::Use(_) => {}
1584        #[cfg(feature = "enterprise")]
1585        Statement::DropTrigger(_) => {}
1586        Statement::ShowCreateDatabase(stmt) => {
1587            validate_database(&stmt.database_name, query_ctx)?;
1588        }
1589        Statement::ShowCreateTable(stmt) => {
1590            validate_param(&stmt.table_name, query_ctx)?;
1591        }
1592        Statement::ShowCreateFlow(stmt) => {
1593            validate_flow(&stmt.flow_name, query_ctx)?;
1594        }
1595        #[cfg(feature = "enterprise")]
1596        Statement::ShowCreateTrigger(stmt) => {
1597            validate_param(&stmt.trigger_name, query_ctx)?;
1598        }
1599        Statement::ShowCreateView(stmt) => {
1600            validate_param(&stmt.view_name, query_ctx)?;
1601        }
1602        Statement::CreateExternalTable(stmt) => {
1603            validate_param(&stmt.name, query_ctx)?;
1604        }
1605        Statement::CreateFlow(stmt) => {
1606            // TODO: should also validate source table name here?
1607            validate_param(&stmt.sink_table_name, query_ctx)?;
1608        }
1609        #[cfg(feature = "enterprise")]
1610        Statement::CreateTrigger(stmt) => {
1611            validate_param(&stmt.trigger_name, query_ctx)?;
1612        }
1613        Statement::CreateView(stmt) => {
1614            validate_param(&stmt.name, query_ctx)?;
1615        }
1616        Statement::AlterTable(stmt) => {
1617            validate_param(stmt.table_name(), query_ctx)?;
1618        }
1619        #[cfg(feature = "enterprise")]
1620        Statement::AlterTrigger(_) => {}
1621        // set/show variable now only alter/show variable in session
1622        Statement::SetVariables(_) | Statement::ShowVariables(_) => {}
1623        // show charset and show collation won't be checked
1624        Statement::ShowCharset(_) | Statement::ShowCollation(_) => {}
1625
1626        Statement::Comment(comment) => match &comment.object {
1627            CommentObject::Table(table) => validate_param(table, query_ctx)?,
1628            CommentObject::Column { table, .. } => validate_param(table, query_ctx)?,
1629            CommentObject::Flow(flow) => validate_flow(flow, query_ctx)?,
1630        },
1631
1632        Statement::Insert(insert) => {
1633            let name = insert.table_name().context(ParseSqlSnafu)?;
1634            validate_param(name, query_ctx)?;
1635        }
1636        Statement::CreateTable(stmt) => {
1637            validate_param(&stmt.name, query_ctx)?;
1638        }
1639        Statement::CreateTableLike(stmt) => {
1640            validate_param(&stmt.table_name, query_ctx)?;
1641            validate_param(&stmt.source_name, query_ctx)?;
1642        }
1643        Statement::DropTable(drop_stmt) => {
1644            for table_name in drop_stmt.table_names() {
1645                validate_param(table_name, query_ctx)?;
1646            }
1647        }
1648        #[cfg(feature = "enterprise")]
1649        Statement::UndropTable(stmt) => {
1650            validate_param(stmt.table_name(), query_ctx)?;
1651        }
1652        Statement::DropView(stmt) => {
1653            validate_param(&stmt.view_name, query_ctx)?;
1654        }
1655        Statement::ShowTables(stmt) => {
1656            validate_db_permission!(stmt, query_ctx);
1657        }
1658        Statement::ShowTableStatus(stmt) => {
1659            validate_db_permission!(stmt, query_ctx);
1660        }
1661        Statement::ShowColumns(stmt) => {
1662            validate_db_permission!(stmt, query_ctx);
1663        }
1664        Statement::ShowIndex(stmt) => {
1665            validate_db_permission!(stmt, query_ctx);
1666        }
1667        Statement::ShowRegion(stmt) => {
1668            validate_db_permission!(stmt, query_ctx);
1669        }
1670        Statement::ShowViews(stmt) => {
1671            validate_db_permission!(stmt, query_ctx);
1672        }
1673        Statement::ShowFlows(stmt) => {
1674            validate_db_permission!(stmt, query_ctx);
1675        }
1676        Statement::ShowFlowStatus(_stmt) => {
1677            // Flow statistics are organized based on the catalog dimension and
1678            // filtered by the current catalog, so there is no need to check the
1679            // permission of the database(schema).
1680        }
1681        #[cfg(feature = "enterprise")]
1682        Statement::ShowTriggers(_stmt) => {
1683            // The trigger is organized based on the catalog dimension, so there
1684            // is no need to check the permission of the database(schema).
1685        }
1686        Statement::ShowStatus(_stmt) => {}
1687        Statement::ShowSearchPath(_stmt) => {}
1688        Statement::DescribeTable(stmt) => {
1689            validate_param(stmt.name(), query_ctx)?;
1690        }
1691        Statement::Copy(sql::statements::copy::Copy::CopyTable(stmt)) => match stmt {
1692            CopyTable::To(copy_table_to) => validate_param(&copy_table_to.table_name, query_ctx)?,
1693            CopyTable::From(copy_table_from) => {
1694                validate_param(&copy_table_from.table_name, query_ctx)?
1695            }
1696        },
1697        Statement::Copy(sql::statements::copy::Copy::CopyDatabase(copy_database)) => {
1698            match copy_database {
1699                CopyDatabase::To(stmt) => validate_database(&stmt.database_name, query_ctx)?,
1700                CopyDatabase::From(stmt) => validate_database(&stmt.database_name, query_ctx)?,
1701            }
1702        }
1703        Statement::TruncateTable(stmt) => {
1704            validate_param(stmt.table_name(), query_ctx)?;
1705        }
1706        // cursor operations are always allowed once it's created
1707        Statement::FetchCursor(_) | Statement::CloseCursor(_) => {}
1708        // User can only kill process in their own catalog.
1709        Statement::Kill(_) => {}
1710        // SHOW PROCESSLIST
1711        Statement::ShowProcesslist(_) => {}
1712    }
1713    Ok(())
1714}
1715
1716fn validate_param(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
1717    let (catalog, schema, _) = table_idents_to_full_name(name, query_ctx)
1718        .map_err(BoxedError::new)
1719        .context(ExternalSnafu)?;
1720
1721    validate_catalog_and_schema(&catalog, &schema, query_ctx)
1722        .map_err(BoxedError::new)
1723        .context(SqlExecInterceptedSnafu)
1724}
1725
1726fn validate_flow(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
1727    let catalog = match &name.0[..] {
1728        [_flow] => query_ctx.current_catalog().to_string(),
1729        [catalog, _flow] => catalog.to_string_unquoted(),
1730        _ => {
1731            return InvalidSqlSnafu {
1732                err_msg: format!(
1733                    "expect flow name to be <catalog>.<flow_name> or <flow_name>, actual: {name}",
1734                ),
1735            }
1736            .fail();
1737        }
1738    };
1739
1740    let schema = query_ctx.current_schema();
1741
1742    validate_catalog_and_schema(&catalog, &schema, query_ctx)
1743        .map_err(BoxedError::new)
1744        .context(SqlExecInterceptedSnafu)
1745}
1746
1747fn validate_database(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
1748    let (catalog, schema) = match &name.0[..] {
1749        [schema] => (
1750            query_ctx.current_catalog().to_string(),
1751            schema.to_string_unquoted(),
1752        ),
1753        [catalog, schema] => (catalog.to_string_unquoted(), schema.to_string_unquoted()),
1754        _ => InvalidSqlSnafu {
1755            err_msg: format!(
1756                "expect database name to be <catalog>.<schema> or <schema>, actual: {name}",
1757            ),
1758        }
1759        .fail()?,
1760    };
1761
1762    validate_catalog_and_schema(&catalog, &schema, query_ctx)
1763        .map_err(BoxedError::new)
1764        .context(SqlExecInterceptedSnafu)
1765}
1766
1767fn is_readonly_plan(plan: &LogicalPlan) -> bool {
1768    !matches!(plan, LogicalPlan::Dml(_) | LogicalPlan::Ddl(_))
1769}
1770
1771fn should_track_statement_process(stmt: &Statement) -> bool {
1772    stmt.is_readonly()
1773        || matches!(stmt, Statement::Insert(insert) if insert.has_non_values_query_source())
1774}
1775
1776fn should_track_plan_process(stmt: Option<&Statement>, plan: &LogicalPlan) -> bool {
1777    is_readonly_plan(plan)
1778        || matches!(stmt, Some(Statement::Insert(insert)) if insert.has_non_values_query_source())
1779}
1780
1781#[cfg(test)]
1782mod tests {
1783    use std::any::Any;
1784    use std::collections::HashMap;
1785    use std::future::Future;
1786    use std::pin::Pin;
1787    use std::sync::Arc;
1788    use std::task::{Context, Poll};
1789    use std::time::Duration;
1790
1791    use api::prom_store::remote::label_matcher::Type as PromMatcherType;
1792    use api::prom_store::remote::{
1793        Label, LabelMatcher, Query as RemoteQuery, ReadRequest, ReadResponse, Sample,
1794    };
1795    use api::v1::greptime_request::Request;
1796    use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse};
1797    use api::v1::query_request::Query;
1798    use auth::{
1799        DASHBOARD_DELETE, DASHBOARD_QUERY, DASHBOARD_SAVE, JAEGER_QUERY, PIPELINE_DELETE,
1800        PIPELINE_INSERT, PIPELINE_QUERY, PermissionAction, PermissionResp, UserInfo, UserInfoRef,
1801    };
1802    use catalog::process_manager::{ProcessManager, QueryStatement, SlowQueryTimer};
1803    use common_base::Plugins;
1804    use common_catalog::consts::DEFAULT_PRIVATE_SCHEMA_NAME;
1805    use common_error::ext::{BoxedError, ErrorExt, PlainError};
1806    use common_error::status_code::StatusCode;
1807    use common_event_recorder::{Event, EventRecorder, EventTypeFilter, EventTypeFilterRef};
1808    use common_frontend::slow_query_event::SlowQueryEvent;
1809    use common_meta::cache::LayeredCacheRegistryBuilder;
1810    use common_meta::kv_backend::memory::MemoryKvBackend;
1811    use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor};
1812    use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse};
1813    use common_meta::rpc::procedure::{
1814        MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse,
1815    };
1816    use common_query::prelude::greptime_value;
1817    use common_query::{Output, OutputMeta};
1818    use common_recordbatch::{
1819        OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream,
1820    };
1821    use common_telemetry::logging::SlowQueriesRecordType;
1822    use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
1823    use datafusion::physical_plan::empty::EmptyExec;
1824    use datafusion_expr::dml::InsertOp;
1825    use datafusion_expr::{LogicalPlanBuilder, LogicalTableSource};
1826    use datatypes::prelude::ConcreteDataType;
1827    use datatypes::schema::{ColumnSchema, Schema as GtSchema, SchemaRef as GtSchemaRef};
1828    use datatypes::vectors::{
1829        Float64Vector, StringVector, TimestampMillisecondVector, TimestampNanosecondVector,
1830        VectorRef,
1831    };
1832    use log_query::LogQuery;
1833    use prost::Message;
1834    use query::query_engine::options::QueryOptions;
1835    use servers::query_handler::{
1836        DashboardHandler, JaegerQueryHandler, LogQueryHandler, PipelineHandler, PipelineHandlerRef,
1837        PromStoreProtocolHandler,
1838    };
1839    use session::context::{Channel, ConnInfo, QueryContext, QueryContextBuilder};
1840    use snafu::{Location, Snafu};
1841    use sql::dialect::GreptimeDbDialect;
1842    use store_api::data_source::DataSource;
1843    use store_api::metric_engine_consts::{
1844        LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY,
1845    };
1846    use store_api::storage::ScanRequest;
1847    use strfmt::Format;
1848    use table::metadata::{
1849        FilterPushDownType, TableInfo, TableInfoBuilder, TableMetaBuilder, TableType,
1850    };
1851    use table::table_name::TableName;
1852    use table::test_util::{EmptyTable, MemTable};
1853    use table::{Table, TableRef};
1854    use tokio::sync::{mpsc, oneshot};
1855    use tower::ServiceExt;
1856
1857    use super::*;
1858    use crate::frontend::FrontendOptions;
1859    use crate::instance::builder::FrontendBuilder;
1860
1861    fn parse_test_sql(sql: &str) -> Vec<Statement> {
1862        parse_stmt(sql, &GreptimeDbDialect {}).unwrap()
1863    }
1864
1865    #[derive(Debug, Default)]
1866    struct RecordingSlowQueryEventRecorder {
1867        payloads: std::sync::Mutex<Vec<serde_json::Value>>,
1868    }
1869
1870    impl EventRecorder for RecordingSlowQueryEventRecorder {
1871        fn record(&self, event: Box<dyn Event>) {
1872            let event = event
1873                .as_any()
1874                .downcast_ref::<SlowQueryEvent>()
1875                .expect("expected a slow query event");
1876            self.payloads.lock().unwrap().push(event.payload.clone());
1877        }
1878
1879        fn event_type_filter(&self) -> EventTypeFilterRef {
1880            Arc::new(EventTypeFilter::All)
1881        }
1882
1883        fn close(&self) {}
1884    }
1885
1886    #[test]
1887    fn test_validate_analyze_stream_statement_strictness() {
1888        for sql in [
1889            "select 1",
1890            "explain analyze select 1",
1891            "explain analyze verbose format text select 1",
1892            "explain analyze verbose format graphviz select 1",
1893            "TQL ANALYZE (0, 10, '5s') physical_metric",
1894            "TQL EXPLAIN VERBOSE (0, 10, '5s') physical_metric",
1895            "TQL ANALYZE VERBOSE FORMAT TEXT (0, 10, '5s') physical_metric",
1896        ] {
1897            let mut stmts = parse_test_sql(sql);
1898            assert!(
1899                validate_analyze_stream_statement(&mut stmts[0]).is_err(),
1900                "{sql}"
1901            );
1902        }
1903
1904        for sql in [
1905            "explain analyze verbose select 1",
1906            "explain analyze verbose format json select 1",
1907            "TQL ANALYZE VERBOSE (0, 10, '5s') physical_metric",
1908            "TQL ANALYZE VERBOSE FORMAT JSON (0, 10, '5s') physical_metric",
1909        ] {
1910            let mut stmts = parse_test_sql(sql);
1911            assert!(
1912                validate_analyze_stream_statement(&mut stmts[0]).is_ok(),
1913                "{sql}"
1914            );
1915            match &stmts[0] {
1916                Statement::Explain(explain) => assert!(explain.format.is_none()),
1917                Statement::Tql(Tql::Analyze(analyze)) => assert!(analyze.format.is_none()),
1918                _ => unreachable!(),
1919            }
1920        }
1921
1922        assert_eq!(
1923            parse_test_sql("explain analyze verbose select 1; select 2").len(),
1924            2
1925        );
1926
1927        assert!(is_explain_analyze_verbose(
1928            &parse_test_sql("explain analyze verbose select 1")[0]
1929        ));
1930        assert!(is_explain_analyze_verbose(
1931            &parse_test_sql("TQL ANALYZE VERBOSE (0, 10, '5s') physical_metric")[0]
1932        ));
1933        for sql in [
1934            "select 1",
1935            "explain select 1",
1936            "explain analyze select 1",
1937            "explain verbose select 1",
1938            "TQL ANALYZE (0, 10, '5s') physical_metric",
1939            "TQL EXPLAIN VERBOSE (0, 10, '5s') physical_metric",
1940        ] {
1941            assert!(
1942                !is_explain_analyze_verbose(&parse_test_sql(sql)[0]),
1943                "{sql}"
1944            );
1945        }
1946    }
1947
1948    #[derive(Debug, Snafu)]
1949    enum TestError {
1950        #[snafu(display("Failed to build test cache registry"))]
1951        BuildCacheRegistry {
1952            source: cache::error::Error,
1953            #[snafu(implicit)]
1954            location: Location,
1955        },
1956
1957        #[snafu(display("Failed to build test table meta for table: {table_name}"))]
1958        BuildTableMeta {
1959            table_name: String,
1960            source: table::metadata::TableMetaBuilderError,
1961            #[snafu(implicit)]
1962            location: Location,
1963        },
1964
1965        #[snafu(display("Failed to build test table info for table: {table_name}"))]
1966        BuildTableInfo {
1967            table_name: String,
1968            source: table::metadata::TableInfoBuilderError,
1969            #[snafu(implicit)]
1970            location: Location,
1971        },
1972
1973        #[snafu(display("Failed to register test table: {table_name}"))]
1974        RegisterTable {
1975            table_name: String,
1976            source: catalog::error::Error,
1977            #[snafu(implicit)]
1978            location: Location,
1979        },
1980
1981        #[snafu(display("Failed to build test frontend instance"))]
1982        BuildFrontend {
1983            source: crate::error::Error,
1984            #[snafu(implicit)]
1985            location: Location,
1986        },
1987
1988        #[snafu(display("Expected exactly one output for SQL `{sql}`, got {actual}"))]
1989        UnexpectedOutputCount {
1990            sql: String,
1991            actual: usize,
1992            #[snafu(implicit)]
1993            location: Location,
1994        },
1995
1996        #[snafu(display("Failed to execute SQL `{sql}`"))]
1997        ExecuteSql {
1998            sql: String,
1999            source: crate::error::Error,
2000            #[snafu(implicit)]
2001            location: Location,
2002        },
2003
2004        #[snafu(display("Timed out waiting for insert-select start notification"))]
2005        InsertStartTimeout {
2006            source: tokio::time::error::Elapsed,
2007            #[snafu(implicit)]
2008            location: Location,
2009        },
2010
2011        #[snafu(display("Insert-select start notification channel closed"))]
2012        InsertStartChannelClosed {
2013            #[snafu(implicit)]
2014            location: Location,
2015        },
2016
2017        #[snafu(display("Failed to release blocking insert-select interceptor"))]
2018        ReleaseBlockedInsert {
2019            #[snafu(implicit)]
2020            location: Location,
2021        },
2022
2023        #[snafu(display("Timed out waiting for insert-select source to be polled"))]
2024        SourcePollTimeout {
2025            source: tokio::time::error::Elapsed,
2026            #[snafu(implicit)]
2027            location: Location,
2028        },
2029
2030        #[snafu(display("Insert-select source poll notification channel closed"))]
2031        SourcePollChannelClosed {
2032            source: oneshot::error::RecvError,
2033            #[snafu(implicit)]
2034            location: Location,
2035        },
2036
2037        #[snafu(display("Timed out waiting for insert task to finish"))]
2038        InsertTaskTimeout {
2039            source: tokio::time::error::Elapsed,
2040            #[snafu(implicit)]
2041            location: Location,
2042        },
2043
2044        #[snafu(display("Insert task panicked"))]
2045        InsertTaskPanic {
2046            source: tokio::task::JoinError,
2047            #[snafu(implicit)]
2048            location: Location,
2049        },
2050
2051        #[snafu(display("Expected insert-select to be cancelled"))]
2052        InsertSelectNotCancelled {
2053            #[snafu(implicit)]
2054            location: Location,
2055        },
2056    }
2057
2058    type TestResult<T> = std::result::Result<T, TestError>;
2059
2060    fn parse_one_sql(sql: &str) -> Statement {
2061        parse_stmt(sql, &GreptimeDbDialect {}).unwrap().remove(0)
2062    }
2063
2064    fn test_query_ctx(process_id: u32) -> QueryContextRef {
2065        Arc::new(
2066            QueryContextBuilder::default()
2067                .channel(Channel::Mysql)
2068                .conn_info(ConnInfo::new(None, Channel::Mysql))
2069                .process_id(process_id)
2070                .build(),
2071        )
2072    }
2073
2074    #[derive(Debug)]
2075    struct AdminUserInfo;
2076
2077    impl UserInfo for AdminUserInfo {
2078        fn as_any(&self) -> &dyn Any {
2079            self
2080        }
2081
2082        fn username(&self) -> &str {
2083            "admin"
2084        }
2085
2086        fn is_admin(&self) -> bool {
2087            true
2088        }
2089    }
2090
2091    struct RejectUnresolvedPermissionChecker;
2092
2093    impl PermissionChecker for RejectUnresolvedPermissionChecker {
2094        fn check_permission(
2095            &self,
2096            _user_info: UserInfoRef,
2097            _req: PermissionReq,
2098        ) -> auth::error::Result<PermissionResp> {
2099            Ok(PermissionResp::Allow)
2100        }
2101
2102        fn check_permission_with_table_targets(
2103            &self,
2104            _user_info: UserInfoRef,
2105            _req: PermissionReq,
2106            targets: PermissionTableTargets,
2107        ) -> auth::error::Result<PermissionResp> {
2108            let reject = match targets {
2109                PermissionTableTargets::Unresolved => true,
2110                PermissionTableTargets::Resolved(targets) => {
2111                    targets.iter().any(|target| target.table == "denied")
2112                }
2113            };
2114            Ok(if reject {
2115                PermissionResp::Reject
2116            } else {
2117                PermissionResp::Allow
2118            })
2119        }
2120    }
2121
2122    #[derive(Debug, PartialEq, Eq)]
2123    struct CheckedAction {
2124        action: PermissionAction,
2125        targets: Option<PermissionTableTargets>,
2126    }
2127
2128    #[derive(Default)]
2129    struct RejectEndpointPermissionChecker {
2130        checks: std::sync::Mutex<Vec<CheckedAction>>,
2131    }
2132
2133    impl RejectEndpointPermissionChecker {
2134        fn reject(
2135            &self,
2136            action: PermissionAction,
2137            targets: Option<PermissionTableTargets>,
2138        ) -> PermissionResp {
2139            self.checks
2140                .lock()
2141                .unwrap()
2142                .push(CheckedAction { action, targets });
2143            PermissionResp::Reject
2144        }
2145
2146        fn take_check(&self) -> CheckedAction {
2147            let mut checks = self.checks.lock().unwrap();
2148            assert_eq!(1, checks.len());
2149            checks.pop().unwrap()
2150        }
2151    }
2152
2153    impl PermissionChecker for RejectEndpointPermissionChecker {
2154        fn check_permission(
2155            &self,
2156            _user_info: UserInfoRef,
2157            req: PermissionReq,
2158        ) -> auth::error::Result<PermissionResp> {
2159            Ok(match req {
2160                PermissionReq::Action(action) => self.reject(action, None),
2161                _ => PermissionResp::Allow,
2162            })
2163        }
2164
2165        fn check_permission_with_table_targets(
2166            &self,
2167            _user_info: UserInfoRef,
2168            req: PermissionReq,
2169            targets: PermissionTableTargets,
2170        ) -> auth::error::Result<PermissionResp> {
2171            Ok(match req {
2172                PermissionReq::Action(action) => self.reject(action, Some(targets)),
2173                _ => PermissionResp::Allow,
2174            })
2175        }
2176    }
2177
2178    struct WriteOnlyPermissionChecker;
2179
2180    impl PermissionChecker for WriteOnlyPermissionChecker {
2181        fn check_permission(
2182            &self,
2183            _user_info: UserInfoRef,
2184            req: PermissionReq,
2185        ) -> auth::error::Result<PermissionResp> {
2186            Ok(if req.is_readonly() {
2187                PermissionResp::Reject
2188            } else {
2189                PermissionResp::Allow
2190            })
2191        }
2192
2193        fn check_permission_with_table_targets(
2194            &self,
2195            user_info: UserInfoRef,
2196            req: PermissionReq,
2197            _targets: PermissionTableTargets,
2198        ) -> auth::error::Result<PermissionResp> {
2199            self.check_permission(user_info, req)
2200        }
2201    }
2202
2203    #[derive(Default)]
2204    struct TargetIndependentPermissionChecker {
2205        checks: atomic::AtomicUsize,
2206    }
2207
2208    impl PermissionChecker for TargetIndependentPermissionChecker {
2209        fn check_permission(
2210            &self,
2211            _user_info: UserInfoRef,
2212            _req: PermissionReq,
2213        ) -> auth::error::Result<PermissionResp> {
2214            self.checks.fetch_add(1, atomic::Ordering::Relaxed);
2215            Ok(PermissionResp::Allow)
2216        }
2217
2218        fn uses_table_targets(&self) -> bool {
2219            false
2220        }
2221
2222        fn check_permission_with_table_targets(
2223            &self,
2224            user_info: UserInfoRef,
2225            req: PermissionReq,
2226            _targets: PermissionTableTargets,
2227        ) -> auth::error::Result<PermissionResp> {
2228            self.check_permission(user_info, req)
2229        }
2230    }
2231
2232    struct BlockingInsertSelectInterceptor {
2233        started_tx: mpsc::UnboundedSender<()>,
2234        finish_rx: std::sync::Mutex<Option<oneshot::Receiver<()>>>,
2235    }
2236
2237    impl BlockingInsertSelectInterceptor {
2238        fn new(started_tx: mpsc::UnboundedSender<()>, finish_rx: oneshot::Receiver<()>) -> Self {
2239            Self {
2240                started_tx,
2241                finish_rx: std::sync::Mutex::new(Some(finish_rx)),
2242            }
2243        }
2244    }
2245
2246    impl SqlQueryInterceptor for BlockingInsertSelectInterceptor {
2247        type Error = Error;
2248
2249        fn pre_execute(
2250            &self,
2251            statement: Option<&Statement>,
2252            _plan: Option<&LogicalPlan>,
2253            _query_ctx: QueryContextRef,
2254        ) -> Result<()> {
2255            let Some(Statement::Insert(insert)) = statement else {
2256                return Ok(());
2257            };
2258            if !insert.has_non_values_query_source() {
2259                return Ok(());
2260            }
2261
2262            let finish_rx = self.finish_rx.lock().unwrap().take().unwrap();
2263            let _ = self.started_tx.send(());
2264            tokio::task::block_in_place(|| {
2265                tokio::runtime::Handle::current()
2266                    .block_on(finish_rx)
2267                    .unwrap();
2268            });
2269            Ok(())
2270        }
2271    }
2272
2273    struct PendingRecordBatchStream {
2274        schema: GtSchemaRef,
2275        polled_tx: Option<oneshot::Sender<()>>,
2276        _finish_tx: oneshot::Sender<()>,
2277        finish_rx: Pin<Box<oneshot::Receiver<()>>>,
2278    }
2279
2280    impl RecordBatchStream for PendingRecordBatchStream {
2281        fn schema(&self) -> GtSchemaRef {
2282            self.schema.clone()
2283        }
2284
2285        fn output_ordering(&self) -> Option<&[OrderOption]> {
2286            None
2287        }
2288
2289        fn metrics(&self) -> Option<common_recordbatch::adapter::RecordBatchMetrics> {
2290            None
2291        }
2292    }
2293
2294    impl Stream for PendingRecordBatchStream {
2295        type Item = common_recordbatch::error::Result<RecordBatch>;
2296
2297        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
2298            if let Some(polled_tx) = self.polled_tx.take() {
2299                let _ = polled_tx.send(());
2300            }
2301
2302            match self.finish_rx.as_mut().poll(cx) {
2303                Poll::Ready(_) => Poll::Ready(None),
2304                Poll::Pending => Poll::Pending,
2305            }
2306        }
2307    }
2308
2309    impl Unpin for PendingRecordBatchStream {}
2310
2311    #[test]
2312    fn test_record_explain_analyze_timeout_uses_empty_metrics_without_plan() {
2313        let event_recorder = Arc::new(RecordingSlowQueryEventRecorder::default());
2314        let timer = SlowQueryTimer::new(
2315            QueryStatement::Plan("EXPLAIN ANALYZE VERBOSE SELECT 1".to_string()),
2316            "public".to_string(),
2317            Duration::from_secs(3600),
2318            0.0,
2319            SlowQueriesRecordType::SystemTable,
2320            event_recorder.clone(),
2321        );
2322        let timeout_recorder = timer.recorder();
2323
2324        record_explain_analyze_timeout(Some(&timeout_recorder), None);
2325        drop(timer);
2326
2327        let payloads = event_recorder.payloads.lock().unwrap();
2328        assert_eq!(payloads.len(), 1);
2329        assert_eq!(payloads[0]["timed_out"], true);
2330        assert_eq!(payloads[0]["metrics"], serde_json::json!([]));
2331    }
2332
2333    #[tokio::test]
2334    async fn test_attach_timeout_records_explain_analyze_metrics() {
2335        let event_recorder = Arc::new(RecordingSlowQueryEventRecorder::default());
2336        let timer = SlowQueryTimer::new(
2337            QueryStatement::Plan("EXPLAIN ANALYZE VERBOSE SELECT 1".to_string()),
2338            "public".to_string(),
2339            Duration::from_secs(3600),
2340            0.0,
2341            SlowQueriesRecordType::SystemTable,
2342            event_recorder.clone(),
2343        );
2344        let timeout_recorder = timer.recorder();
2345        let plan: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(Arc::new(Schema::empty())));
2346        let (finish_tx, finish_rx) = oneshot::channel();
2347        let stream = PendingRecordBatchStream {
2348            schema: Arc::new(GtSchema::new(vec![])),
2349            polled_tx: None,
2350            _finish_tx: finish_tx,
2351            finish_rx: Box::pin(finish_rx),
2352        };
2353        let output = Output::new(
2354            OutputData::Stream(Box::pin(stream)),
2355            OutputMeta::new_with_plan(plan),
2356        );
2357        let output =
2358            attach_timeout(output, Duration::from_millis(10), Some(timeout_recorder)).unwrap();
2359        let OutputData::Stream(mut stream) = output.data else {
2360            unreachable!();
2361        };
2362
2363        let err = stream.next().await.unwrap().unwrap_err();
2364        assert_eq!(err.to_string(), "Stream timeout");
2365        drop(stream);
2366        drop(timer);
2367
2368        let payloads = event_recorder.payloads.lock().unwrap();
2369        assert_eq!(payloads.len(), 1);
2370        assert_eq!(payloads[0]["timed_out"], true);
2371        assert!(
2372            payloads[0]["metrics"]
2373                .as_array()
2374                .is_some_and(|metrics| !metrics.is_empty())
2375        );
2376    }
2377
2378    struct PendingDataSource {
2379        schema: GtSchemaRef,
2380        polled_tx: std::sync::Mutex<Option<oneshot::Sender<()>>>,
2381    }
2382
2383    impl DataSource for PendingDataSource {
2384        fn get_stream(
2385            &self,
2386            _request: ScanRequest,
2387        ) -> std::result::Result<SendableRecordBatchStream, BoxedError> {
2388            let (finish_tx, finish_rx) = oneshot::channel();
2389            let mut polled_tx = self.polled_tx.lock().map_err(|_| {
2390                BoxedError::new(PlainError::new(
2391                    "pending data source lock poisoned".to_string(),
2392                    StatusCode::Unexpected,
2393                ))
2394            })?;
2395            Ok(Box::pin(PendingRecordBatchStream {
2396                schema: self.schema.clone(),
2397                polled_tx: polled_tx.take(),
2398                _finish_tx: finish_tx,
2399                finish_rx: Box::pin(finish_rx),
2400            }))
2401        }
2402    }
2403
2404    struct NoopProcedureExecutor;
2405
2406    #[async_trait::async_trait]
2407    impl ProcedureExecutor for NoopProcedureExecutor {
2408        async fn submit_ddl_task(
2409            &self,
2410            _ctx: ExecutorContext,
2411            _request: SubmitDdlTaskRequest,
2412        ) -> common_meta::error::Result<SubmitDdlTaskResponse> {
2413            common_meta::error::UnsupportedSnafu {
2414                operation: "submit_ddl_task",
2415            }
2416            .fail()
2417        }
2418
2419        async fn migrate_region(
2420            &self,
2421            _ctx: &ExecutorContext,
2422            _request: MigrateRegionRequest,
2423        ) -> common_meta::error::Result<MigrateRegionResponse> {
2424            common_meta::error::UnsupportedSnafu {
2425                operation: "migrate_region",
2426            }
2427            .fail()
2428        }
2429
2430        async fn reconcile(
2431            &self,
2432            _ctx: &ExecutorContext,
2433            _request: ReconcileRequest,
2434        ) -> common_meta::error::Result<ReconcileResponse> {
2435            common_meta::error::UnsupportedSnafu {
2436                operation: "reconcile",
2437            }
2438            .fail()
2439        }
2440
2441        async fn query_procedure_state(
2442            &self,
2443            _ctx: &ExecutorContext,
2444            _pid: &str,
2445        ) -> common_meta::error::Result<ProcedureStateResponse> {
2446            common_meta::error::UnsupportedSnafu {
2447                operation: "query_procedure_state",
2448            }
2449            .fail()
2450        }
2451
2452        async fn list_procedures(
2453            &self,
2454            _ctx: &ExecutorContext,
2455        ) -> common_meta::error::Result<ProcedureDetailResponse> {
2456            common_meta::error::UnsupportedSnafu {
2457                operation: "list_procedures",
2458            }
2459            .fail()
2460        }
2461    }
2462
2463    /// A test [`ProcedureExecutor`] that completes create/drop DDL tasks against the
2464    /// in-memory catalog, mimicking what the meta DDL procedures do in production.
2465    /// This allows happy-path DDL requests (create/drop table/view) to be exercised
2466    /// end to end through the gRPC ingress.
2467    struct MockProcedureExecutor {
2468        catalog_manager: Arc<catalog::memory::MemoryCatalogManager>,
2469        next_table_id: std::sync::atomic::AtomicU32,
2470        submitted: std::sync::Mutex<Vec<DdlTask>>,
2471    }
2472
2473    impl MockProcedureExecutor {
2474        fn new(catalog_manager: Arc<catalog::memory::MemoryCatalogManager>) -> Self {
2475            Self {
2476                catalog_manager,
2477                next_table_id: std::sync::atomic::AtomicU32::new(1026),
2478                submitted: std::sync::Mutex::new(Vec::new()),
2479            }
2480        }
2481    }
2482
2483    #[async_trait::async_trait]
2484    impl ProcedureExecutor for MockProcedureExecutor {
2485        async fn submit_ddl_task(
2486            &self,
2487            _ctx: ExecutorContext,
2488            request: SubmitDdlTaskRequest,
2489        ) -> common_meta::error::Result<SubmitDdlTaskResponse> {
2490            self.submitted.lock().unwrap().push(request.task.clone());
2491            match request.task {
2492                DdlTask::CreateTable(task) => {
2493                    let table_id = self
2494                        .next_table_id
2495                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2496                    let mut table_info = task.table_info;
2497                    table_info.ident.table_id = table_id;
2498                    self.catalog_manager
2499                        .register_table_sync(catalog::RegisterTableRequest {
2500                            catalog: table_info.catalog_name.clone(),
2501                            schema: table_info.schema_name.clone(),
2502                            table_name: table_info.name.clone(),
2503                            table_id,
2504                            table: table::dist_table::DistTable::table(Arc::new(table_info)),
2505                        })
2506                        .map_err(BoxedError::new)
2507                        .context(common_meta::error::ExternalSnafu)?;
2508                    Ok(SubmitDdlTaskResponse {
2509                        key: Vec::new(),
2510                        table_ids: vec![table_id],
2511                    })
2512                }
2513                DdlTask::CreateView(task) => {
2514                    let view_id = self
2515                        .next_table_id
2516                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2517                    let mut view_info = task.view_info;
2518                    view_info.ident.table_id = view_id;
2519                    self.catalog_manager
2520                        .register_table_sync(catalog::RegisterTableRequest {
2521                            catalog: task.create_view.catalog_name.clone(),
2522                            schema: task.create_view.schema_name.clone(),
2523                            table_name: task.create_view.view_name.clone(),
2524                            table_id: view_id,
2525                            table: table::dist_table::DistTable::table(Arc::new(view_info)),
2526                        })
2527                        .map_err(BoxedError::new)
2528                        .context(common_meta::error::ExternalSnafu)?;
2529                    Ok(SubmitDdlTaskResponse {
2530                        key: Vec::new(),
2531                        table_ids: vec![view_id],
2532                    })
2533                }
2534                DdlTask::DropView(task) => {
2535                    self.catalog_manager
2536                        .deregister_table_sync(catalog::DeregisterTableRequest {
2537                            catalog: task.catalog.clone(),
2538                            schema: task.schema.clone(),
2539                            table_name: task.view.clone(),
2540                        })
2541                        .map_err(BoxedError::new)
2542                        .context(common_meta::error::ExternalSnafu)?;
2543                    Ok(SubmitDdlTaskResponse::default())
2544                }
2545                other => common_meta::error::UnsupportedSnafu {
2546                    operation: format!("mock submit_ddl_task: {other:?}"),
2547                }
2548                .fail(),
2549            }
2550        }
2551
2552        async fn migrate_region(
2553            &self,
2554            _ctx: &ExecutorContext,
2555            _request: MigrateRegionRequest,
2556        ) -> common_meta::error::Result<MigrateRegionResponse> {
2557            common_meta::error::UnsupportedSnafu {
2558                operation: "migrate_region",
2559            }
2560            .fail()
2561        }
2562
2563        async fn reconcile(
2564            &self,
2565            _ctx: &ExecutorContext,
2566            _request: ReconcileRequest,
2567        ) -> common_meta::error::Result<ReconcileResponse> {
2568            common_meta::error::UnsupportedSnafu {
2569                operation: "reconcile",
2570            }
2571            .fail()
2572        }
2573
2574        async fn query_procedure_state(
2575            &self,
2576            _ctx: &ExecutorContext,
2577            _pid: &str,
2578        ) -> common_meta::error::Result<ProcedureStateResponse> {
2579            common_meta::error::UnsupportedSnafu {
2580                operation: "query_procedure_state",
2581            }
2582            .fail()
2583        }
2584
2585        async fn list_procedures(
2586            &self,
2587            _ctx: &ExecutorContext,
2588        ) -> common_meta::error::Result<ProcedureDetailResponse> {
2589            common_meta::error::UnsupportedSnafu {
2590                operation: "list_procedures",
2591            }
2592            .fail()
2593        }
2594    }
2595
2596    fn test_cache_registry(
2597        kv_backend: common_meta::kv_backend::KvBackendRef,
2598    ) -> TestResult<common_meta::cache::LayeredCacheRegistryRef> {
2599        Ok(Arc::new(
2600            cache::with_default_composite_cache_registry(
2601                LayeredCacheRegistryBuilder::default()
2602                    .add_cache_registry(cache::build_fundamental_cache_registry(kv_backend)),
2603            )
2604            .context(BuildCacheRegistrySnafu)?
2605            .build(),
2606        ))
2607    }
2608
2609    fn test_table_info(table_id: u32, table_name: &str) -> TestResult<TableInfo> {
2610        let schema = Arc::new(GtSchema::new(vec![
2611            ColumnSchema::new("id", ConcreteDataType::int32_datatype(), false),
2612            ColumnSchema::new(
2613                "ts",
2614                ConcreteDataType::timestamp_millisecond_datatype(),
2615                false,
2616            )
2617            .with_time_index(true),
2618        ]));
2619        let table_meta = TableMetaBuilder::empty()
2620            .schema(schema)
2621            .primary_key_indices(vec![0])
2622            .value_indices(vec![1])
2623            .next_column_id(1024)
2624            .build()
2625            .with_context(|_| BuildTableMetaSnafu {
2626                table_name: table_name.to_string(),
2627            })?;
2628
2629        TableInfoBuilder::new(table_name, table_meta)
2630            .table_id(table_id)
2631            .build()
2632            .with_context(|_| BuildTableInfoSnafu {
2633                table_name: table_name.to_string(),
2634            })
2635    }
2636
2637    fn test_table(table_id: u32, table_name: &str) -> TestResult<table::TableRef> {
2638        let table_info = test_table_info(table_id, table_name)?;
2639        Ok(EmptyTable::from_table_info(&table_info))
2640    }
2641
2642    fn test_physical_table(table_id: u32, table_name: &str) -> TestResult<table::TableRef> {
2643        let mut table_info = test_table_info(table_id, table_name)?;
2644        table_info
2645            .meta
2646            .options
2647            .extra_options
2648            .insert(PHYSICAL_TABLE_METADATA_KEY.to_string(), String::new());
2649        Ok(EmptyTable::from_table_info(&table_info))
2650    }
2651
2652    fn test_logical_table(table_id: u32, table_name: &str) -> TestResult<table::TableRef> {
2653        let mut table_info = test_table_info(table_id, table_name)?;
2654        table_info.meta.engine = METRIC_ENGINE_NAME.to_string();
2655        table_info.meta.options.extra_options.insert(
2656            LOGICAL_TABLE_METADATA_KEY.to_string(),
2657            "physical_metric".to_string(),
2658        );
2659        Ok(EmptyTable::from_table_info(&table_info))
2660    }
2661
2662    fn test_metric_names_table() -> TableRef {
2663        let schema = Arc::new(GtSchema::new(vec![
2664            ColumnSchema::new("table_catalog", ConcreteDataType::string_datatype(), false),
2665            ColumnSchema::new("table_schema", ConcreteDataType::string_datatype(), false),
2666            ColumnSchema::new("table_name", ConcreteDataType::string_datatype(), false),
2667            ColumnSchema::new("engine", ConcreteDataType::string_datatype(), false),
2668            ColumnSchema::new("create_options", ConcreteDataType::string_datatype(), false),
2669        ]));
2670        let columns: Vec<VectorRef> = vec![
2671            Arc::new(StringVector::from(vec!["greptime", "greptime"])),
2672            Arc::new(StringVector::from(vec!["public", "public"])),
2673            Arc::new(StringVector::from(vec!["denied", "target"])),
2674            Arc::new(StringVector::from(vec!["metric", "metric"])),
2675            Arc::new(StringVector::from(vec![
2676                "on_physical_table=physical_metric",
2677                "on_physical_table=physical_metric",
2678            ])),
2679        ];
2680        let record_batch = RecordBatch::new(schema, columns).unwrap();
2681        MemTable::new_with_catalog(
2682            "tables",
2683            record_batch,
2684            2048,
2685            "greptime".to_string(),
2686            "information_schema".to_string(),
2687        )
2688    }
2689
2690    fn test_pipeline_table() -> TableRef {
2691        let schema = Arc::new(GtSchema::new(vec![
2692            ColumnSchema::new("name", ConcreteDataType::string_datatype(), false),
2693            ColumnSchema::new("schema", ConcreteDataType::string_datatype(), false),
2694            ColumnSchema::new("content_type", ConcreteDataType::string_datatype(), false),
2695            ColumnSchema::new("pipeline", ConcreteDataType::string_datatype(), false),
2696            ColumnSchema::new(
2697                "created_at",
2698                ConcreteDataType::timestamp_nanosecond_datatype(),
2699                false,
2700            )
2701            .with_time_index(true),
2702        ]));
2703        let columns: Vec<VectorRef> = vec![
2704            Arc::new(StringVector::from(vec!["pipeline"])),
2705            Arc::new(StringVector::from(vec!["public"])),
2706            Arc::new(StringVector::from(vec!["application/yaml"])),
2707            Arc::new(StringVector::from(vec![
2708                "transform:\n- field: ts\n  type: timestamp, ns\n  index: time\n",
2709            ])),
2710            Arc::new(TimestampNanosecondVector::from_values([1])),
2711        ];
2712        let record_batch = RecordBatch::new(schema, columns).unwrap();
2713        MemTable::new_with_catalog(
2714            "pipelines",
2715            record_batch,
2716            2049,
2717            "greptime".to_string(),
2718            DEFAULT_PRIVATE_SCHEMA_NAME.to_string(),
2719        )
2720    }
2721
2722    fn pending_table(
2723        table_id: u32,
2724        table_name: &str,
2725        polled_tx: oneshot::Sender<()>,
2726    ) -> TestResult<table::TableRef> {
2727        let table_info = test_table_info(table_id, table_name)?;
2728        let data_source = Arc::new(PendingDataSource {
2729            schema: table_info.meta.schema.clone(),
2730            polled_tx: std::sync::Mutex::new(Some(polled_tx)),
2731        });
2732
2733        Ok(Arc::new(Table::new(
2734            Arc::new(table_info),
2735            FilterPushDownType::Unsupported,
2736            data_source,
2737        )))
2738    }
2739
2740    async fn test_instance_with_tables(
2741        source_table: TableRef,
2742        target_table: TableRef,
2743    ) -> TestResult<Instance> {
2744        test_instance_with_plugins(source_table, target_table, Plugins::new()).await
2745    }
2746
2747    async fn test_instance_with_insert_select_interceptor(
2748        interceptor: SqlQueryInterceptorRef<Error>,
2749    ) -> TestResult<Instance> {
2750        let plugins = Plugins::new();
2751        plugins.insert::<SqlQueryInterceptorRef<Error>>(interceptor);
2752
2753        test_instance_with_plugins(
2754            test_table(1024, "source")?,
2755            test_table(1025, "target")?,
2756            plugins,
2757        )
2758        .await
2759    }
2760
2761    async fn test_instance_with_plugins(
2762        source_table: TableRef,
2763        target_table: TableRef,
2764        plugins: Plugins,
2765    ) -> TestResult<Instance> {
2766        test_instance_with_plugins_and_metric_names(source_table, target_table, plugins, None).await
2767    }
2768
2769    async fn test_instance_with_plugins_and_metric_names(
2770        source_table: TableRef,
2771        target_table: TableRef,
2772        plugins: Plugins,
2773        metric_names_table: Option<TableRef>,
2774    ) -> TestResult<Instance> {
2775        let catalog_manager = catalog::memory::MemoryCatalogManager::new_with_table(source_table);
2776        test_instance_with_catalog_manager(
2777            catalog_manager,
2778            target_table,
2779            plugins,
2780            metric_names_table,
2781            Arc::new(NoopProcedureExecutor),
2782        )
2783        .await
2784    }
2785
2786    /// Builds a test frontend `Instance` over the given (already source-registered)
2787    /// catalog manager, completing DDL tasks through `procedure_executor`.
2788    async fn test_instance_with_catalog_manager(
2789        catalog_manager: Arc<catalog::memory::MemoryCatalogManager>,
2790        target_table: TableRef,
2791        plugins: Plugins,
2792        metric_names_table: Option<TableRef>,
2793        procedure_executor: ProcedureExecutorRef,
2794    ) -> TestResult<Instance> {
2795        let kv_backend = Arc::new(MemoryKvBackend::new());
2796        let process_manager = Arc::new(ProcessManager::new("test-frontend".to_string(), None));
2797        let target_table_name = "target";
2798        catalog_manager
2799            .register_table_sync(catalog::RegisterTableRequest {
2800                catalog: "greptime".to_string(),
2801                schema: "public".to_string(),
2802                table_name: target_table_name.to_string(),
2803                table_id: 1025,
2804                table: target_table,
2805            })
2806            .with_context(|_| RegisterTableSnafu {
2807                table_name: target_table_name.to_string(),
2808            })?;
2809        if let Some(table) = metric_names_table {
2810            catalog_manager
2811                .deregister_table_sync(catalog::DeregisterTableRequest {
2812                    catalog: "greptime".to_string(),
2813                    schema: "information_schema".to_string(),
2814                    table_name: "tables".to_string(),
2815                })
2816                .unwrap();
2817            catalog_manager
2818                .register_table_sync(catalog::RegisterTableRequest {
2819                    catalog: "greptime".to_string(),
2820                    schema: "information_schema".to_string(),
2821                    table_name: "tables".to_string(),
2822                    table_id: 2048,
2823                    table,
2824                })
2825                .unwrap();
2826        }
2827        catalog_manager.register_process_list_table(process_manager.clone());
2828
2829        let cache_registry = test_cache_registry(kv_backend.clone())?;
2830
2831        FrontendBuilder::new(
2832            FrontendOptions::default(),
2833            kv_backend,
2834            cache_registry,
2835            catalog_manager,
2836            Arc::new(client::client_manager::NodeClients::default()),
2837            procedure_executor,
2838            process_manager,
2839        )
2840        .with_plugin(plugins)
2841        .try_build()
2842        .await
2843        .context(BuildFrontendSnafu)
2844    }
2845
2846    async fn execute_one_sql(
2847        instance: &Instance,
2848        sql: &str,
2849        query_ctx: QueryContextRef,
2850    ) -> TestResult<Output> {
2851        let mut results = instance.do_query_inner(sql, query_ctx).await;
2852        ensure!(
2853            results.len() == 1,
2854            UnexpectedOutputCountSnafu {
2855                sql: sql.to_string(),
2856                actual: results.len(),
2857            }
2858        );
2859        results.remove(0).with_context(|_| ExecuteSqlSnafu {
2860            sql: sql.to_string(),
2861        })
2862    }
2863
2864    fn assert_permission_denied<T>(result: servers::error::Result<T>) {
2865        let err = match result {
2866            Ok(_) => panic!("request should be rejected"),
2867            Err(err) => err,
2868        };
2869        assert_eq!(StatusCode::PermissionDenied, err.status_code());
2870    }
2871
2872    fn assert_action_checked(
2873        checker: &RejectEndpointPermissionChecker,
2874        action: PermissionAction,
2875        targets: Option<PermissionTableTargets>,
2876    ) {
2877        assert_eq!(CheckedAction { action, targets }, checker.take_check());
2878    }
2879
2880    #[tokio::test]
2881    async fn test_prom_remote_read_with_custom_timestamp_and_value_columns() -> TestResult<()> {
2882        let schema = Arc::new(GtSchema::new(vec![
2883            ColumnSchema::new(
2884                "custom_ts",
2885                ConcreteDataType::timestamp_millisecond_datatype(),
2886                false,
2887            )
2888            .with_time_index(true),
2889            ColumnSchema::new("custom_value", ConcreteDataType::float64_datatype(), false),
2890        ]));
2891        let recordbatch = RecordBatch::new(
2892            schema,
2893            vec![
2894                Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as VectorRef,
2895                Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as VectorRef,
2896            ],
2897        )
2898        .unwrap();
2899        let instance = test_instance_with_tables(
2900            MemTable::table("custom_metric", recordbatch),
2901            test_table(1025, "target")?,
2902        )
2903        .await?;
2904
2905        let response = PromStoreProtocolHandler::read(
2906            &instance,
2907            ReadRequest {
2908                queries: vec![RemoteQuery {
2909                    start_timestamp_ms: 1500,
2910                    end_timestamp_ms: 2500,
2911                    matchers: vec![LabelMatcher {
2912                        r#type: PromMatcherType::Eq as i32,
2913                        name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
2914                        value: "custom_metric".to_string(),
2915                    }],
2916                    ..Default::default()
2917                }],
2918                ..Default::default()
2919            },
2920            test_query_ctx(1),
2921        )
2922        .await
2923        .unwrap();
2924        let body = servers::prom_store::snappy_decompress(&response.body).unwrap();
2925        let response = ReadResponse::decode(body.as_slice()).unwrap();
2926
2927        assert_eq!(1, response.results.len());
2928        assert_eq!(1, response.results[0].timeseries.len());
2929        let timeseries = &response.results[0].timeseries[0];
2930        assert_eq!(
2931            vec![Label {
2932                name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
2933                value: "custom_metric".to_string(),
2934            }],
2935            timeseries.labels
2936        );
2937        assert_eq!(
2938            vec![Sample {
2939                value: 2.0,
2940                timestamp: 2000,
2941            }],
2942            timeseries.samples
2943        );
2944
2945        Ok(())
2946    }
2947
2948    #[tokio::test]
2949    async fn test_prom_remote_read_prefers_default_value_column() -> TestResult<()> {
2950        let schema = Arc::new(GtSchema::new(vec![
2951            ColumnSchema::new(
2952                "custom_ts",
2953                ConcreteDataType::timestamp_millisecond_datatype(),
2954                false,
2955            )
2956            .with_time_index(true),
2957            ColumnSchema::new("extra_field", ConcreteDataType::float64_datatype(), false),
2958            ColumnSchema::new(
2959                greptime_value(),
2960                ConcreteDataType::float64_datatype(),
2961                false,
2962            ),
2963        ]));
2964        let recordbatch = RecordBatch::new(
2965            schema,
2966            vec![
2967                Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as VectorRef,
2968                Arc::new(Float64Vector::from_vec(vec![99.0, 99.0, 99.0])) as VectorRef,
2969                Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as VectorRef,
2970            ],
2971        )
2972        .unwrap();
2973        let instance = test_instance_with_tables(
2974            MemTable::table("multi_field_metric", recordbatch),
2975            test_table(1025, "target")?,
2976        )
2977        .await?;
2978
2979        let response = PromStoreProtocolHandler::read(
2980            &instance,
2981            ReadRequest {
2982                queries: vec![RemoteQuery {
2983                    start_timestamp_ms: 1500,
2984                    end_timestamp_ms: 2500,
2985                    matchers: vec![LabelMatcher {
2986                        r#type: PromMatcherType::Eq as i32,
2987                        name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
2988                        value: "multi_field_metric".to_string(),
2989                    }],
2990                    ..Default::default()
2991                }],
2992                ..Default::default()
2993            },
2994            test_query_ctx(1),
2995        )
2996        .await
2997        .unwrap();
2998        let body = servers::prom_store::snappy_decompress(&response.body).unwrap();
2999        let response = ReadResponse::decode(body.as_slice()).unwrap();
3000
3001        assert_eq!(1, response.results.len());
3002        assert_eq!(1, response.results[0].timeseries.len());
3003        let timeseries = &response.results[0].timeseries[0];
3004        assert_eq!(
3005            vec![
3006                Label {
3007                    name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3008                    value: "multi_field_metric".to_string(),
3009                },
3010                Label {
3011                    name: "extra_field".to_string(),
3012                    value: "99".to_string(),
3013                },
3014            ],
3015            timeseries.labels
3016        );
3017        assert_eq!(
3018            vec![Sample {
3019                value: 2.0,
3020                timestamp: 2000,
3021            }],
3022            timeseries.samples
3023        );
3024
3025        Ok(())
3026    }
3027
3028    #[tokio::test]
3029    async fn test_prom_remote_read_rejects_ambiguous_value_columns() -> TestResult<()> {
3030        let schema = Arc::new(GtSchema::new(vec![
3031            ColumnSchema::new(
3032                "custom_ts",
3033                ConcreteDataType::timestamp_millisecond_datatype(),
3034                false,
3035            )
3036            .with_time_index(true),
3037            ColumnSchema::new("field_a", ConcreteDataType::float64_datatype(), false),
3038            ColumnSchema::new("field_b", ConcreteDataType::float64_datatype(), false),
3039        ]));
3040        let recordbatch = RecordBatch::new(
3041            schema,
3042            vec![
3043                Arc::new(TimestampMillisecondVector::from_vec(vec![1000])) as VectorRef,
3044                Arc::new(Float64Vector::from_vec(vec![1.0])) as VectorRef,
3045                Arc::new(Float64Vector::from_vec(vec![2.0])) as VectorRef,
3046            ],
3047        )
3048        .unwrap();
3049        let instance = test_instance_with_tables(
3050            MemTable::table("ambiguous_metric", recordbatch),
3051            test_table(1025, "target")?,
3052        )
3053        .await?;
3054
3055        let err = PromStoreProtocolHandler::read(
3056            &instance,
3057            ReadRequest {
3058                queries: vec![RemoteQuery {
3059                    matchers: vec![LabelMatcher {
3060                        r#type: PromMatcherType::Eq as i32,
3061                        name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3062                        value: "ambiguous_metric".to_string(),
3063                    }],
3064                    ..Default::default()
3065                }],
3066                ..Default::default()
3067            },
3068            test_query_ctx(1),
3069        )
3070        .await
3071        .err()
3072        .expect("ambiguous value columns should fail remote read");
3073
3074        assert_eq!(StatusCode::InvalidArguments, err.status_code());
3075        assert!(format!("{err:?}").contains("Ambiguous value column"));
3076
3077        Ok(())
3078    }
3079
3080    #[tokio::test]
3081    async fn test_event_recorder_is_exposed() -> TestResult<()> {
3082        let instance =
3083            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
3084                .await?;
3085
3086        let _event_recorder = instance.event_recorder();
3087
3088        Ok(())
3089    }
3090
3091    #[tokio::test]
3092    async fn test_restricted_endpoint_handlers_check_permissions() -> TestResult<()> {
3093        let checker = Arc::new(RejectEndpointPermissionChecker::default());
3094        let plugins = Plugins::new();
3095        plugins.insert::<PermissionCheckerRef>(checker.clone());
3096        let instance = test_instance_with_plugins(
3097            test_table(1024, "denied")?,
3098            test_table(1025, "target")?,
3099            plugins,
3100        )
3101        .await?;
3102        let mut ctx = test_query_ctx(1);
3103        Arc::get_mut(&mut ctx).unwrap().set_extension(
3104            servers::http::jaeger::JAEGER_QUERY_TABLE_NAME_KEY,
3105            "denied".to_string(),
3106        );
3107        let jaeger_targets = Some(PermissionTableTargets::resolved(vec![
3108            PermissionTableTarget::new("greptime", "public", "denied"),
3109        ]));
3110
3111        assert_permission_denied(JaegerQueryHandler::get_services(&instance, ctx.clone()).await);
3112        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
3113        assert_permission_denied(
3114            JaegerQueryHandler::get_operations(&instance, ctx.clone(), "service", None).await,
3115        );
3116        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
3117        assert_permission_denied(
3118            JaegerQueryHandler::get_trace(&instance, ctx.clone(), "trace", None, None, None).await,
3119        );
3120        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
3121        assert_permission_denied(
3122            JaegerQueryHandler::find_traces(
3123                &instance,
3124                ctx.clone(),
3125                servers::http::jaeger::QueryTraceParams {
3126                    service_name: "service".to_string(),
3127                    ..Default::default()
3128                },
3129            )
3130            .await,
3131        );
3132        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets);
3133
3134        assert_permission_denied(
3135            PipelineHandler::get_pipeline_str(&instance, "pipeline", None, ctx.clone()).await,
3136        );
3137        assert_action_checked(&checker, PIPELINE_QUERY, None);
3138        assert_permission_denied(
3139            PipelineHandler::insert_pipeline(
3140                &instance,
3141                "pipeline",
3142                "application/yaml",
3143                "",
3144                ctx.clone(),
3145            )
3146            .await,
3147        );
3148        assert_action_checked(&checker, PIPELINE_INSERT, None);
3149        assert_permission_denied(
3150            PipelineHandler::delete_pipeline(&instance, "pipeline", None, ctx.clone()).await,
3151        );
3152        assert_action_checked(&checker, PIPELINE_DELETE, None);
3153        let app = axum::Router::new()
3154            .route(
3155                "/pipelines/_dryrun",
3156                axum::routing::post(servers::http::event::pipeline_dryrun),
3157            )
3158            .with_state(servers::http::event::LogState {
3159                log_handler: Arc::new(instance.clone()),
3160                log_validator: None,
3161                ingest_interceptor: None,
3162            })
3163            .layer(axum::Extension((*ctx).clone()));
3164        let response = app
3165            .oneshot(
3166                axum::http::Request::post("/pipelines/_dryrun")
3167                    .header("content-type", "application/json")
3168                    .body(axum::body::Body::from("{}"))
3169                    .unwrap(),
3170            )
3171            .await
3172            .unwrap();
3173        assert_eq!(axum::http::StatusCode::FORBIDDEN, response.status());
3174        assert_action_checked(&checker, PIPELINE_QUERY, None);
3175
3176        assert_permission_denied(
3177            DashboardHandler::save(&instance, "dashboard", "{}", ctx.clone()).await,
3178        );
3179        assert_action_checked(&checker, DASHBOARD_SAVE, None);
3180        assert_permission_denied(DashboardHandler::list(&instance, ctx.clone()).await);
3181        assert_action_checked(&checker, DASHBOARD_QUERY, None);
3182        assert_permission_denied(
3183            DashboardHandler::delete(&instance, "dashboard", ctx.clone()).await,
3184        );
3185        assert_action_checked(&checker, DASHBOARD_DELETE, None);
3186
3187        Ok(())
3188    }
3189
3190    #[tokio::test]
3191    async fn test_write_only_ingestion_loads_named_pipeline() -> TestResult<()> {
3192        let plugins = Plugins::new();
3193        plugins.insert::<PermissionCheckerRef>(Arc::new(WriteOnlyPermissionChecker));
3194        let instance = test_instance_with_plugins(
3195            test_table(1024, "source")?,
3196            test_table(1025, "target")?,
3197            plugins,
3198        )
3199        .await?;
3200        instance
3201            .catalog_manager()
3202            .as_any()
3203            .downcast_ref::<catalog::memory::MemoryCatalogManager>()
3204            .unwrap()
3205            .register_table_sync(catalog::RegisterTableRequest {
3206                catalog: "greptime".to_string(),
3207                schema: DEFAULT_PRIVATE_SCHEMA_NAME.to_string(),
3208                table_name: "pipelines".to_string(),
3209                table_id: 2049,
3210                table: test_pipeline_table(),
3211            })
3212            .with_context(|_| RegisterTableSnafu {
3213                table_name: "pipelines".to_string(),
3214            })?;
3215        let ctx = test_query_ctx(1);
3216        let handler: PipelineHandlerRef = Arc::new(instance.clone());
3217
3218        handler
3219            .get_pipeline("pipeline", None, ctx.clone())
3220            .await
3221            .unwrap();
3222        assert_permission_denied(
3223            PipelineHandler::get_pipeline_str(&instance, "pipeline", None, ctx.clone()).await,
3224        );
3225
3226        let app = axum::Router::new()
3227            .route(
3228                "/pipelines/_dryrun",
3229                axum::routing::post(servers::http::event::pipeline_dryrun),
3230            )
3231            .with_state(servers::http::event::LogState {
3232                log_handler: handler,
3233                log_validator: None,
3234                ingest_interceptor: None,
3235            })
3236            .layer(axum::Extension((*ctx).clone()));
3237        let response = app
3238            .oneshot(
3239                axum::http::Request::post("/pipelines/_dryrun")
3240                    .header("content-type", "application/json")
3241                    .body(axum::body::Body::from("{}"))
3242                    .unwrap(),
3243            )
3244            .await
3245            .unwrap();
3246        assert_eq!(axum::http::StatusCode::FORBIDDEN, response.status());
3247
3248        Ok(())
3249    }
3250
3251    #[tokio::test]
3252    async fn test_write_only_grpc_sql_is_checked_after_parsing() -> TestResult<()> {
3253        let plugins = Plugins::new();
3254        plugins.insert::<PermissionCheckerRef>(Arc::new(WriteOnlyPermissionChecker));
3255        let instance = test_instance_with_plugins(
3256            test_table(1024, "source")?,
3257            test_table(1025, "target")?,
3258            plugins,
3259        )
3260        .await?;
3261
3262        let insert = Request::Query(api::v1::QueryRequest {
3263            query: Some(Query::Sql(
3264                "INSERT INTO target SELECT * FROM source".to_string(),
3265            )),
3266        });
3267        servers::query_handler::grpc::GrpcQueryHandler::do_query(
3268            &instance,
3269            insert,
3270            QueryContext::arc(),
3271        )
3272        .await
3273        .unwrap();
3274
3275        let select = Request::Query(api::v1::QueryRequest {
3276            query: Some(Query::Sql("SELECT * FROM source".to_string())),
3277        });
3278        assert_permission_denied(
3279            servers::query_handler::grpc::GrpcQueryHandler::do_query(
3280                &instance,
3281                select,
3282                QueryContext::arc(),
3283            )
3284            .await,
3285        );
3286
3287        Ok(())
3288    }
3289
3290    #[tokio::test]
3291    async fn test_target_independent_checker_skips_target_resolution() -> TestResult<()> {
3292        let physical_table = "physical_metric";
3293        let checker = Arc::new(TargetIndependentPermissionChecker::default());
3294        let plugins = Plugins::new();
3295        plugins.insert::<PermissionCheckerRef>(checker.clone());
3296        let instance = test_instance_with_plugins(
3297            test_physical_table(1024, physical_table)?,
3298            test_table(1025, "target")?,
3299            plugins,
3300        )
3301        .await?;
3302
3303        let ctx = test_query_ctx(1);
3304        let physical_target = PermissionTableTarget::new("greptime", "public", physical_table);
3305        assert_eq!(
3306            PermissionTableTargets::Resolved(vec![physical_target.clone()]),
3307            instance
3308                .resolve_query_permission_targets(
3309                    PermissionTableTargets::resolved(vec![physical_target]),
3310                    &ctx,
3311                )
3312                .await
3313                .unwrap()
3314        );
3315        assert_eq!(
3316            vec![physical_table.to_string(), "target".to_string()],
3317            PrometheusHandler::filter_metadata_metric_names(
3318                &instance,
3319                vec![physical_table.to_string(), "target".to_string()],
3320                "public",
3321                &ctx,
3322            )
3323            .await
3324            .unwrap()
3325        );
3326        assert_eq!(1, checker.checks.load(atomic::Ordering::Relaxed));
3327
3328        Ok(())
3329    }
3330
3331    #[tokio::test]
3332    async fn test_query_permission_targets_are_deduplicated() -> TestResult<()> {
3333        let plugins = Plugins::new();
3334        plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
3335        let instance = test_instance_with_plugins(
3336            test_table(1024, "source")?,
3337            test_table(1025, "target")?,
3338            plugins,
3339        )
3340        .await?;
3341        let ctx = test_query_ctx(1);
3342        let target = PermissionTableTarget::new("greptime", "public", "target");
3343
3344        assert_eq!(
3345            PermissionTableTargets::Resolved(vec![target.clone()]),
3346            instance
3347                .resolve_query_permission_targets(
3348                    PermissionTableTargets::resolved(vec![target.clone(), target]),
3349                    &ctx,
3350                )
3351                .await
3352                .unwrap()
3353        );
3354
3355        Ok(())
3356    }
3357
3358    #[tokio::test]
3359    async fn test_physical_query_targets_fail_closed() -> TestResult<()> {
3360        let physical_table = "physical_metric";
3361        let plugins = Plugins::new();
3362        plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
3363        let instance = test_instance_with_plugins(
3364            test_physical_table(1024, physical_table)?,
3365            test_table(1025, "target")?,
3366            plugins,
3367        )
3368        .await?;
3369
3370        let ctx = test_query_ctx(1);
3371        let logical_target = PermissionTableTarget::new("greptime", "public", "target");
3372        assert_eq!(
3373            PermissionTableTargets::Resolved(vec![logical_target.clone()]),
3374            instance
3375                .resolve_query_permission_targets(
3376                    PermissionTableTargets::resolved(vec![logical_target.clone()]),
3377                    &ctx,
3378                )
3379                .await
3380                .unwrap()
3381        );
3382        let physical_target = PermissionTableTarget::new("greptime", "public", physical_table);
3383        assert_eq!(
3384            PermissionTableTargets::Unresolved,
3385            instance
3386                .resolve_query_permission_targets(
3387                    PermissionTableTargets::resolved(
3388                        vec![logical_target, physical_target.clone(),]
3389                    ),
3390                    &ctx,
3391                )
3392                .await
3393                .unwrap()
3394        );
3395        assert_eq!(
3396            vec!["target".to_string()],
3397            PrometheusHandler::filter_metadata_metric_names(
3398                &instance,
3399                vec!["target".to_string(), "denied".to_string()],
3400                "public",
3401                &ctx,
3402            )
3403            .await
3404            .unwrap()
3405        );
3406
3407        let query = PromQuery {
3408            query: physical_table.to_string(),
3409            ..Default::default()
3410        };
3411        let err = PrometheusHandler::check_query_target_permission(
3412            &instance,
3413            PermissionTableTargets::resolved(vec![physical_target]),
3414            &ctx,
3415        )
3416        .await
3417        .unwrap_err();
3418        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3419        let err = PrometheusHandler::check_query_permission(
3420            &instance,
3421            std::slice::from_ref(&query),
3422            &ctx,
3423        )
3424        .await
3425        .unwrap_err();
3426        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3427        let err = PrometheusHandler::do_query(&instance, &query, ctx.clone())
3428            .await
3429            .unwrap_err();
3430        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3431
3432        for sql in [
3433            "SELECT * FROM physical_metric",
3434            "TQL EVAL (0, 10, '5s') physical_metric",
3435            "INSERT INTO target SELECT * FROM physical_metric",
3436        ] {
3437            let mut results = instance.do_query_inner(sql, ctx.clone()).await;
3438            assert_eq!(1, results.len(), "{sql}");
3439            let err = results.remove(0).unwrap_err();
3440            assert_eq!(StatusCode::PermissionDenied, err.status_code(), "{sql}");
3441        }
3442        let err = LogQueryHandler::query(
3443            &instance,
3444            LogQuery {
3445                table: TableName::new("greptime", "public", physical_table),
3446                ..Default::default()
3447            },
3448            ctx.clone(),
3449        )
3450        .await
3451        .unwrap_err();
3452        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3453        let err = instance
3454            .do_describe_inner(parse_one_sql("SELECT * FROM physical_metric"), ctx.clone())
3455            .await
3456            .unwrap_err();
3457        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3458
3459        let request = ReadRequest {
3460            queries: vec![RemoteQuery {
3461                matchers: vec![LabelMatcher {
3462                    r#type: PromMatcherType::Eq as i32,
3463                    name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3464                    value: physical_table.to_string(),
3465                }],
3466                ..Default::default()
3467            }],
3468            ..Default::default()
3469        };
3470        let Err(err) = PromStoreProtocolHandler::read(&instance, request, ctx.clone()).await else {
3471            panic!("physical remote-read target must be rejected");
3472        };
3473        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3474
3475        let err = PrometheusHandler::query_label_values(
3476            &instance,
3477            physical_table.to_string(),
3478            "host".to_string(),
3479            vec![],
3480            SystemTime::UNIX_EPOCH,
3481            SystemTime::UNIX_EPOCH,
3482            &ctx,
3483        )
3484        .await
3485        .unwrap_err();
3486        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3487
3488        Ok(())
3489    }
3490
3491    #[tokio::test]
3492    async fn test_non_exact_query_discovery_keeps_denied_targets_for_batch_check() -> TestResult<()>
3493    {
3494        let plugins = Plugins::new();
3495        plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
3496        let instance = test_instance_with_plugins_and_metric_names(
3497            test_logical_table(1024, "denied")?,
3498            test_logical_table(1025, "target")?,
3499            plugins,
3500            Some(test_metric_names_table()),
3501        )
3502        .await?;
3503        let ctx = test_query_ctx(1);
3504
3505        let mut metric_names = PrometheusHandler::query_metric_names(
3506            &instance,
3507            vec![Matcher::new(
3508                promql_parser::label::MatchOp::NotEqual,
3509                "__name__",
3510                "",
3511            )],
3512            "public",
3513            &ctx,
3514        )
3515        .await
3516        .unwrap();
3517        metric_names.sort_unstable();
3518        assert_eq!(
3519            vec!["denied".to_string(), "target".to_string()],
3520            metric_names
3521        );
3522
3523        let queries = metric_names
3524            .into_iter()
3525            .map(|query| PromQuery {
3526                query,
3527                ..Default::default()
3528            })
3529            .collect::<Vec<_>>();
3530        let err = PrometheusHandler::check_query_permission(&instance, &queries, &ctx)
3531            .await
3532            .unwrap_err();
3533        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3534
3535        Ok(())
3536    }
3537
3538    #[test]
3539    fn test_fast_legacy_check_is_read_only() {
3540        let cache = DashMap::new();
3541        cache.insert("metric1".to_string(), true);
3542
3543        let names = vec!["metric1".to_string(), "metric2".to_string()];
3544        assert_eq!(Some(true), fast_legacy_check(&cache, &names).unwrap());
3545        assert!(!cache.contains_key("metric2"));
3546
3547        cache_legacy_mode(&cache, &names, true).unwrap();
3548        assert!(*cache.get("metric2").unwrap().value());
3549        assert!(cache_legacy_mode(&cache, &names, false).is_err());
3550        assert!(*cache.get("metric2").unwrap().value());
3551
3552        let cache_incompatible = DashMap::new();
3553        cache_incompatible.insert("metric1".to_string(), true);
3554        cache_incompatible.insert("metric2".to_string(), false);
3555        assert!(fast_legacy_check(&cache_incompatible, &names).is_err());
3556    }
3557
3558    #[test]
3559    fn test_should_track_statement_process() {
3560        assert!(should_track_statement_process(&parse_one_sql(
3561            "SELECT * FROM demo"
3562        )));
3563        assert!(should_track_statement_process(&parse_one_sql(
3564            "INSERT INTO demo SELECT * FROM source"
3565        )));
3566        assert!(!should_track_statement_process(&parse_one_sql(
3567            "INSERT INTO demo VALUES (1)"
3568        )));
3569        assert!(!should_track_statement_process(&parse_one_sql(
3570            "INSERT INTO demo VALUES (now())"
3571        )));
3572    }
3573
3574    #[test]
3575    fn test_should_track_plan_process() {
3576        let select_stmt = parse_one_sql("SELECT * FROM demo");
3577        let insert_select_stmt = parse_one_sql("INSERT INTO demo SELECT * FROM source");
3578        let insert_values_stmt = parse_one_sql("INSERT INTO demo VALUES (now())");
3579
3580        let empty_plan = LogicalPlanBuilder::empty(false).build().unwrap();
3581        assert!(should_track_plan_process(Some(&select_stmt), &empty_plan));
3582        assert!(should_track_plan_process(
3583            Some(&insert_select_stmt),
3584            &insert_dml_plan()
3585        ));
3586        assert!(!should_track_plan_process(
3587            Some(&insert_values_stmt),
3588            &insert_dml_plan()
3589        ));
3590        assert!(!should_track_plan_process(None, &insert_dml_plan()));
3591    }
3592
3593    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3594    async fn test_insert_select_is_visible_in_show_processlist() -> TestResult<()> {
3595        let insert_sql = "INSERT INTO target SELECT * FROM source";
3596        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
3597        let (finish_tx, finish_rx) = oneshot::channel();
3598        let interceptor = Arc::new(BlockingInsertSelectInterceptor::new(started_tx, finish_rx));
3599        let instance = Arc::new(test_instance_with_insert_select_interceptor(interceptor).await?);
3600
3601        let insert_task = tokio::spawn({
3602            let instance = instance.clone();
3603            async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
3604        });
3605
3606        tokio::time::timeout(Duration::from_secs(5), started_rx.recv())
3607            .await
3608            .context(InsertStartTimeoutSnafu)?
3609            .context(InsertStartChannelClosedSnafu)?;
3610
3611        let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
3612        let process_list = output.data.pretty_print().await;
3613        assert!(
3614            process_list.contains(insert_sql),
3615            "process list did not contain running insert:\n{process_list}"
3616        );
3617
3618        finish_tx
3619            .send(())
3620            .map_err(|_| ReleaseBlockedInsertSnafu.build())?;
3621        insert_task.await.context(InsertTaskPanicSnafu)??;
3622
3623        Ok(())
3624    }
3625
3626    #[tokio::test]
3627    async fn test_show_processlist_catalog_scope() -> TestResult<()> {
3628        let instance =
3629            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
3630                .await?;
3631        let _current_catalog = instance.process_manager().register_query(
3632            "greptime".to_string(),
3633            vec!["public".to_string()],
3634            "current_catalog_query".to_string(),
3635            String::new(),
3636            None,
3637            None,
3638        );
3639        let _other_catalog = instance.process_manager().register_query(
3640            "other".to_string(),
3641            vec!["public".to_string()],
3642            "other_catalog_query".to_string(),
3643            String::new(),
3644            None,
3645            None,
3646        );
3647
3648        for sql in ["SHOW PROCESSLIST", "SHOW FULL PROCESSLIST"] {
3649            let output = execute_one_sql(&instance, sql, test_query_ctx(43)).await?;
3650            let process_list = output.data.pretty_print().await;
3651            assert!(
3652                process_list.contains("current_catalog_query"),
3653                "{process_list}"
3654            );
3655            assert!(
3656                !process_list.contains("other_catalog_query"),
3657                "{process_list}"
3658            );
3659
3660            let admin_ctx = test_query_ctx(44);
3661            admin_ctx.set_current_user(Arc::new(AdminUserInfo));
3662            let output = execute_one_sql(&instance, sql, admin_ctx).await?;
3663            let process_list = output.data.pretty_print().await;
3664            assert!(
3665                process_list.contains("current_catalog_query"),
3666                "{process_list}"
3667            );
3668            assert!(
3669                process_list.contains("other_catalog_query"),
3670                "{process_list}"
3671            );
3672        }
3673
3674        Ok(())
3675    }
3676
3677    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3678    async fn test_kill_query_cancels_insert_select() -> TestResult<()> {
3679        assert_kill_cancels_insert_select("KILL QUERY 4242").await
3680    }
3681
3682    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3683    async fn test_kill_process_id_cancels_insert_select() -> TestResult<()> {
3684        assert_kill_cancels_insert_select("KILL 'test-frontend/4242'").await
3685    }
3686
3687    async fn assert_kill_cancels_insert_select(kill_sql: &str) -> TestResult<()> {
3688        let insert_sql = "INSERT INTO target SELECT * FROM source";
3689        let (source_polled_tx, source_polled_rx) = oneshot::channel();
3690        let instance = Arc::new(
3691            test_instance_with_tables(
3692                pending_table(1024, "source", source_polled_tx)?,
3693                test_table(1025, "target")?,
3694            )
3695            .await?,
3696        );
3697
3698        let insert_task = tokio::spawn({
3699            let instance = instance.clone();
3700            async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
3701        });
3702
3703        tokio::time::timeout(Duration::from_secs(5), source_polled_rx)
3704            .await
3705            .context(SourcePollTimeoutSnafu)?
3706            .context(SourcePollChannelClosedSnafu)?;
3707
3708        let output = execute_one_sql(&instance, kill_sql, test_query_ctx(43)).await?;
3709        assert!(matches!(output.data, OutputData::AffectedRows(1)));
3710
3711        let insert_result = tokio::time::timeout(Duration::from_secs(5), insert_task)
3712            .await
3713            .context(InsertTaskTimeoutSnafu)?
3714            .context(InsertTaskPanicSnafu)?;
3715        let err = match insert_result {
3716            Ok(_) => return InsertSelectNotCancelledSnafu.fail(),
3717            Err(TestError::ExecuteSql { source, .. }) => source,
3718            Err(err) => return Err(err),
3719        };
3720        assert_eq!(StatusCode::Cancelled, err.status_code());
3721
3722        let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
3723        let process_list = output.data.pretty_print().await;
3724        assert!(
3725            !process_list.contains(insert_sql),
3726            "process list still contains killed insert:\n{process_list}"
3727        );
3728
3729        Ok(())
3730    }
3731
3732    fn insert_dml_plan() -> LogicalPlan {
3733        let schema = SchemaRef::new(Schema::new(vec![Field::new(
3734            "value",
3735            DataType::Int64,
3736            true,
3737        )]));
3738        let target = Arc::new(LogicalTableSource::new(schema));
3739        let input = LogicalPlanBuilder::empty(false).build().unwrap();
3740
3741        LogicalPlanBuilder::insert_into(input, "demo", target, InsertOp::Append)
3742            .unwrap()
3743            .build()
3744            .unwrap()
3745    }
3746
3747    #[test]
3748    fn test_exec_validation() {
3749        let query_ctx = QueryContext::arc();
3750        let plugins: Plugins = Plugins::new();
3751        plugins.insert(QueryOptions {
3752            disallow_cross_catalog_query: true,
3753        });
3754
3755        let sql = r#"
3756        SELECT * FROM demo;
3757        EXPLAIN SELECT * FROM demo;
3758        CREATE DATABASE test_database;
3759        SHOW DATABASES;
3760        "#;
3761        let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
3762        assert_eq!(stmts.len(), 4);
3763        for stmt in stmts {
3764            let re = check_permission(plugins.clone(), &stmt, &query_ctx);
3765            re.unwrap();
3766        }
3767
3768        let sql = r#"
3769        SHOW CREATE TABLE demo;
3770        ALTER TABLE demo ADD COLUMN new_col INT;
3771        "#;
3772        let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
3773        assert_eq!(stmts.len(), 2);
3774        for stmt in stmts {
3775            let re = check_permission(plugins.clone(), &stmt, &query_ctx);
3776            re.unwrap();
3777        }
3778
3779        fn replace_test(template_sql: &str, plugins: Plugins, query_ctx: &QueryContextRef) {
3780            // test right
3781            let right = vec![("", ""), ("", "public."), ("greptime.", "public.")];
3782            for (catalog, schema) in right {
3783                let sql = do_fmt(template_sql, catalog, schema);
3784                do_test(&sql, plugins.clone(), query_ctx, true);
3785            }
3786
3787            let wrong = vec![
3788                ("wrongcatalog.", "public."),
3789                ("wrongcatalog.", "wrongschema."),
3790            ];
3791            for (catalog, schema) in wrong {
3792                let sql = do_fmt(template_sql, catalog, schema);
3793                do_test(&sql, plugins.clone(), query_ctx, false);
3794            }
3795        }
3796
3797        fn do_fmt(template: &str, catalog: &str, schema: &str) -> String {
3798            let vars = HashMap::from([
3799                ("catalog".to_string(), catalog),
3800                ("schema".to_string(), schema),
3801            ]);
3802            template.format(&vars).unwrap()
3803        }
3804
3805        fn do_test(sql: &str, plugins: Plugins, query_ctx: &QueryContextRef, is_ok: bool) {
3806            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
3807            let re = check_permission(plugins, stmt, query_ctx);
3808            if is_ok {
3809                re.unwrap();
3810            } else {
3811                assert!(re.is_err());
3812            }
3813        }
3814
3815        // test insert
3816        let sql = "INSERT INTO {catalog}{schema}monitor(host) VALUES ('host1');";
3817        replace_test(sql, plugins.clone(), &query_ctx);
3818
3819        // test create table
3820        let sql = r#"CREATE TABLE {catalog}{schema}demo(
3821                            host STRING,
3822                            ts TIMESTAMP,
3823                            TIME INDEX (ts),
3824                            PRIMARY KEY(host)
3825                        ) engine=mito;"#;
3826        replace_test(sql, plugins.clone(), &query_ctx);
3827
3828        // test drop table
3829        let sql = "DROP TABLE {catalog}{schema}demo;";
3830        replace_test(sql, plugins.clone(), &query_ctx);
3831
3832        // test undrop table
3833        #[cfg(feature = "enterprise")]
3834        {
3835            let sql = "UNDROP TABLE {catalog}{schema}demo;";
3836            replace_test(sql, plugins.clone(), &query_ctx);
3837        }
3838
3839        // test show tables
3840        let sql = "SHOW TABLES FROM public";
3841        let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
3842        check_permission(plugins.clone(), &stmt[0], &query_ctx).unwrap();
3843
3844        let sql = "SHOW TABLES FROM private";
3845        let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
3846        let re = check_permission(plugins.clone(), &stmt[0], &query_ctx);
3847        assert!(re.is_ok());
3848
3849        // test describe table
3850        let sql = "DESC TABLE {catalog}{schema}demo;";
3851        replace_test(sql, plugins.clone(), &query_ctx);
3852
3853        let comment_flow_cases = [
3854            ("COMMENT ON FLOW my_flow IS 'comment';", true),
3855            ("COMMENT ON FLOW greptime.my_flow IS 'comment';", true),
3856            ("COMMENT ON FLOW wrongcatalog.my_flow IS 'comment';", false),
3857        ];
3858        for (sql, is_ok) in comment_flow_cases {
3859            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
3860            let result = check_permission(plugins.clone(), stmt, &query_ctx);
3861            assert_eq!(result.is_ok(), is_ok);
3862        }
3863
3864        let show_flow_cases = [
3865            ("SHOW CREATE FLOW my_flow;", true),
3866            ("SHOW CREATE FLOW greptime.my_flow;", true),
3867            ("SHOW CREATE FLOW wrongcatalog.my_flow;", false),
3868        ];
3869        for (sql, is_ok) in show_flow_cases {
3870            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
3871            let result = check_permission(plugins.clone(), stmt, &query_ctx);
3872            assert_eq!(result.is_ok(), is_ok);
3873        }
3874    }
3875
3876    /// A `DropView` DDL sent through the direct gRPC ingress must return an error
3877    /// (e.g. table not found) instead of panicking on `todo!()`.
3878    #[tokio::test]
3879    async fn qx_152_drop_view_via_grpc_ddl_returns_error_not_panic() -> TestResult<()> {
3880        let instance =
3881            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
3882                .await?;
3883
3884        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
3885            expr: Some(api::v1::ddl_request::Expr::DropView(
3886                api::v1::DropViewExpr {
3887                    catalog_name: String::new(),
3888                    schema_name: String::new(),
3889                    view_name: "non_existent_view".to_string(),
3890                    view_id: None,
3891                    drop_if_exists: false,
3892                },
3893            )),
3894        });
3895
3896        let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
3897            &instance,
3898            request,
3899            QueryContext::arc(),
3900        )
3901        .await;
3902
3903        let err = match result {
3904            Ok(_) => panic!("DropView DDL request must return an error instead of panicking"),
3905            Err(err) => err,
3906        };
3907        assert_eq!(
3908            err.status_code(),
3909            StatusCode::TableNotFound,
3910            "dropping a non-existent view without IF EXISTS must report TableNotFound, got {err}"
3911        );
3912        Ok(())
3913    }
3914
3915    /// `DROP VIEW IF EXISTS` on a missing view through the direct gRPC ingress must
3916    /// succeed with 0 affected rows (no error, no DDL task submitted), instead of
3917    /// returning `TableNotFound`.
3918    #[tokio::test]
3919    async fn qx_152_drop_view_if_exists_missing_view_via_grpc_ddl_succeeds() -> TestResult<()> {
3920        let catalog_manager =
3921            catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
3922        let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
3923        let instance = test_instance_with_catalog_manager(
3924            catalog_manager,
3925            test_table(1025, "target")?,
3926            Plugins::new(),
3927            None,
3928            procedure_executor.clone() as ProcedureExecutorRef,
3929        )
3930        .await?;
3931
3932        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
3933            expr: Some(api::v1::ddl_request::Expr::DropView(
3934                api::v1::DropViewExpr {
3935                    catalog_name: String::new(),
3936                    schema_name: String::new(),
3937                    view_name: "non_existent_view".to_string(),
3938                    view_id: None,
3939                    drop_if_exists: true,
3940                },
3941            )),
3942        });
3943
3944        let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
3945            &instance,
3946            request,
3947            QueryContext::arc(),
3948        )
3949        .await;
3950
3951        let output = match result {
3952            Ok(output) => output,
3953            Err(err) => {
3954                panic!("DROP VIEW IF EXISTS on a missing view must succeed, got error: {err}")
3955            }
3956        };
3957        assert!(
3958            matches!(output.data, OutputData::AffectedRows(0)),
3959            "DROP VIEW IF EXISTS on a missing view must report 0 affected rows"
3960        );
3961        assert!(
3962            procedure_executor.submitted.lock().unwrap().is_empty(),
3963            "DROP VIEW IF EXISTS on a missing view must not submit a DDL task"
3964        );
3965        Ok(())
3966    }
3967
3968    /// A `CREATE VIEW` followed by `DROP VIEW` through the direct gRPC ingress must
3969    /// succeed end to end: the view is registered in the catalog and then removed.
3970    #[tokio::test]
3971    async fn qx_152_drop_existing_view_via_grpc_ddl_succeeds() -> TestResult<()> {
3972        let catalog_manager =
3973            catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
3974        let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
3975        let instance = test_instance_with_catalog_manager(
3976            catalog_manager,
3977            test_table(1025, "target")?,
3978            Plugins::new(),
3979            None,
3980            procedure_executor.clone() as ProcedureExecutorRef,
3981        )
3982        .await?;
3983
3984        // The default "greptime.public" schema must be visible to the kv-backed table
3985        // metadata manager for `CREATE VIEW`/`CREATE TABLE` to pass validation.
3986        instance
3987            .table_metadata_manager()
3988            .schema_manager()
3989            .create(
3990                common_meta::key::schema_name::SchemaNameKey::new("greptime", "public"),
3991                None,
3992                true,
3993            )
3994            .await
3995            .unwrap();
3996
3997        let create_view_request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
3998            expr: Some(api::v1::ddl_request::Expr::CreateView(
3999                api::v1::CreateViewExpr {
4000                    catalog_name: String::new(),
4001                    schema_name: String::new(),
4002                    view_name: "my_view".to_string(),
4003                    logical_plan: vec![1, 2, 3],
4004                    create_if_not_exists: false,
4005                    or_replace: false,
4006                    table_names: vec![],
4007                    columns: vec![],
4008                    plan_columns: vec![],
4009                    definition: "CREATE VIEW my_view AS SELECT * FROM source".to_string(),
4010                },
4011            )),
4012        });
4013
4014        let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
4015            &instance,
4016            create_view_request,
4017            QueryContext::arc(),
4018        )
4019        .await
4020        {
4021            Ok(output) => output,
4022            Err(err) => panic!("CREATE VIEW via gRPC DDL must succeed, got error: {err}"),
4023        };
4024        assert!(
4025            matches!(output.data, OutputData::AffectedRows(0)),
4026            "CREATE VIEW via gRPC DDL must report 0 affected rows"
4027        );
4028
4029        // The view is registered in the catalog as a view.
4030        let view = instance
4031            .catalog_manager()
4032            .table("greptime", "public", "my_view", None)
4033            .await
4034            .unwrap()
4035            .expect("view should exist after CREATE VIEW");
4036        assert_eq!(view.table_info().table_type, TableType::View);
4037
4038        let drop_view_request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4039            expr: Some(api::v1::ddl_request::Expr::DropView(
4040                api::v1::DropViewExpr {
4041                    catalog_name: String::new(),
4042                    schema_name: String::new(),
4043                    view_name: "my_view".to_string(),
4044                    view_id: None,
4045                    drop_if_exists: false,
4046                },
4047            )),
4048        });
4049
4050        let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
4051            &instance,
4052            drop_view_request,
4053            QueryContext::arc(),
4054        )
4055        .await
4056        {
4057            Ok(output) => output,
4058            Err(err) => panic!("DROP VIEW via gRPC DDL must succeed, got error: {err}"),
4059        };
4060        assert!(
4061            matches!(output.data, OutputData::AffectedRows(0)),
4062            "DROP VIEW via gRPC DDL must report 0 affected rows"
4063        );
4064
4065        // The view is gone after the drop.
4066        assert!(
4067            instance
4068                .catalog_manager()
4069                .table("greptime", "public", "my_view", None)
4070                .await
4071                .unwrap()
4072                .is_none(),
4073            "view should be removed after DROP VIEW"
4074        );
4075
4076        let submitted = procedure_executor.submitted.lock().unwrap();
4077        assert_eq!(
4078            submitted.len(),
4079            2,
4080            "expected create and drop view tasks, got {submitted:?}"
4081        );
4082        assert!(matches!(&submitted[0], DdlTask::CreateView(_)));
4083        assert!(matches!(&submitted[1], DdlTask::DropView(_)));
4084        Ok(())
4085    }
4086
4087    /// A direct gRPC `CreateTable` whose time index column is not a timestamp must
4088    /// be rejected with `InvalidArguments` instead of panicking while building the schema.
4089    #[tokio::test]
4090    async fn qx_153_create_table_with_non_timestamp_time_index_via_grpc_returns_error()
4091    -> TestResult<()> {
4092        let instance =
4093            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
4094                .await?;
4095
4096        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4097            expr: Some(api::v1::ddl_request::Expr::CreateTable(
4098                api::v1::CreateTableExpr {
4099                    catalog_name: String::new(),
4100                    schema_name: String::new(),
4101                    table_name: "demo".to_string(),
4102                    desc: String::new(),
4103                    column_defs: vec![api::v1::ColumnDef {
4104                        name: "host".to_string(),
4105                        data_type: api::v1::ColumnDataType::String as i32,
4106                        is_nullable: true,
4107                        default_constraint: vec![],
4108                        semantic_type: 0,
4109                        comment: String::new(),
4110                        datatype_extension: None,
4111                        options: None,
4112                    }],
4113                    time_index: "host".to_string(),
4114                    primary_keys: vec![],
4115                    create_if_not_exists: false,
4116                    table_options: HashMap::new(),
4117                    table_id: None,
4118                    engine: "mito".to_string(),
4119                },
4120            )),
4121        });
4122
4123        let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
4124            &instance,
4125            request,
4126            QueryContext::arc(),
4127        )
4128        .await;
4129
4130        let err = match result {
4131            Ok(_) => panic!("CreateTable with a non-timestamp time index must be rejected"),
4132            Err(err) => err,
4133        };
4134        assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}");
4135        Ok(())
4136    }
4137
4138    /// A valid `CREATE TABLE` (timestamp time index) through the direct gRPC ingress
4139    /// must succeed, guarding that the `validate_create_expr` ingress check doesn't
4140    /// accidentally reject good requests.
4141    #[tokio::test]
4142    async fn qx_153_create_table_with_timestamp_time_index_via_grpc_succeeds() -> TestResult<()> {
4143        let catalog_manager =
4144            catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
4145        let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
4146        let instance = test_instance_with_catalog_manager(
4147            catalog_manager,
4148            test_table(1025, "target")?,
4149            Plugins::new(),
4150            None,
4151            procedure_executor.clone() as ProcedureExecutorRef,
4152        )
4153        .await?;
4154
4155        // The default "greptime.public" schema must be visible to the kv-backed table
4156        // metadata manager for `CREATE TABLE` to pass validation.
4157        instance
4158            .table_metadata_manager()
4159            .schema_manager()
4160            .create(
4161                common_meta::key::schema_name::SchemaNameKey::new("greptime", "public"),
4162                None,
4163                true,
4164            )
4165            .await
4166            .unwrap();
4167
4168        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4169            expr: Some(api::v1::ddl_request::Expr::CreateTable(
4170                api::v1::CreateTableExpr {
4171                    catalog_name: String::new(),
4172                    schema_name: String::new(),
4173                    table_name: "demo".to_string(),
4174                    desc: String::new(),
4175                    column_defs: vec![
4176                        api::v1::ColumnDef {
4177                            name: "host".to_string(),
4178                            data_type: api::v1::ColumnDataType::String as i32,
4179                            is_nullable: true,
4180                            default_constraint: vec![],
4181                            semantic_type: 0,
4182                            comment: String::new(),
4183                            datatype_extension: None,
4184                            options: None,
4185                        },
4186                        api::v1::ColumnDef {
4187                            name: "ts".to_string(),
4188                            data_type: api::v1::ColumnDataType::TimestampMillisecond as i32,
4189                            is_nullable: true,
4190                            default_constraint: vec![],
4191                            semantic_type: 0,
4192                            comment: String::new(),
4193                            datatype_extension: None,
4194                            options: None,
4195                        },
4196                    ],
4197                    time_index: "ts".to_string(),
4198                    primary_keys: vec![],
4199                    create_if_not_exists: false,
4200                    table_options: HashMap::new(),
4201                    table_id: None,
4202                    engine: "mito".to_string(),
4203                },
4204            )),
4205        });
4206
4207        let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
4208            &instance,
4209            request,
4210            QueryContext::arc(),
4211        )
4212        .await
4213        {
4214            Ok(output) => output,
4215            Err(err) => panic!("CREATE TABLE via gRPC DDL must succeed, got error: {err}"),
4216        };
4217        assert!(
4218            matches!(output.data, OutputData::AffectedRows(0)),
4219            "CREATE TABLE via gRPC DDL must report 0 affected rows"
4220        );
4221
4222        // The table is registered in the catalog.
4223        let table = instance
4224            .catalog_manager()
4225            .table("greptime", "public", "demo", None)
4226            .await
4227            .unwrap()
4228            .expect("table should exist after CREATE TABLE");
4229        assert_eq!(table.table_info().table_type, TableType::Base);
4230
4231        let submitted = procedure_executor.submitted.lock().unwrap();
4232        assert_eq!(
4233            submitted.len(),
4234            1,
4235            "expected one create table task, got {submitted:?}"
4236        );
4237        assert!(matches!(&submitted[0], DdlTask::CreateTable(_)));
4238        Ok(())
4239    }
4240}