Skip to main content

frontend/
instance.rs

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