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