Skip to main content

frontend/
instance.rs

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