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 !stmt.is_readonly() || !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 plan_is_readonly = is_readonly_plan(&plan);
1036        let result = if should_track_plan_process(stmt.as_ref(), &plan) {
1037            let catalog_name = query_ctx.current_catalog().to_string();
1038            let schema_name = query_ctx.current_schema();
1039            let slow_query_timer = if plan_is_readonly {
1040                self.slow_query_options.enable.then(|| {
1041                    SlowQueryTimer::new(
1042                        CatalogQueryStatement::Plan(query.clone()),
1043                        schema_name.clone(),
1044                        self.slow_query_options.threshold,
1045                        self.slow_query_options.sample_ratio,
1046                        self.slow_query_options.record_type,
1047                        self.slow_query_recorder.clone(),
1048                    )
1049                })
1050            } else {
1051                None
1052            };
1053
1054            let timeout_recorder = stmt
1055                .as_ref()
1056                .is_some_and(is_explain_analyze_verbose)
1057                .then(|| slow_query_timer.as_ref().map(SlowQueryTimer::recorder))
1058                .flatten();
1059            let ticket = self.process_manager.register_query(
1060                catalog_name,
1061                vec![schema_name],
1062                query,
1063                query_ctx.conn_info().to_string(),
1064                Some(query_ctx.process_id()),
1065                slow_query_timer,
1066            );
1067
1068            let query_fut = self.exec_plan_with_timeout(plan, query_ctx.clone(), timeout_recorder);
1069
1070            #[cfg(feature = "enterprise")]
1071            let result = if query_ctx.channel() == Channel::Internal
1072                && query_ctx.extension(common_grpc::flow_write::CHILD_KEY) == Some("1")
1073            {
1074                query_fut.await
1075            } else {
1076                CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
1077                    .await
1078                    .map_err(|_| error::CancelledSnafu.build())?
1079            };
1080            #[cfg(not(feature = "enterprise"))]
1081            let result = CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
1082                .await
1083                .map_err(|_| error::CancelledSnafu.build())?;
1084            result.map(|output| {
1085                let Output { meta, data } = output;
1086                let data = match data {
1087                    OutputData::Stream(stream) => {
1088                        OutputData::Stream(Box::pin(CancellableStreamWrapper::new(stream, ticket)))
1089                    }
1090                    other => other,
1091                };
1092                Output { data, meta }
1093            })
1094        } else {
1095            self.exec_plan_with_timeout(plan, query_ctx.clone(), None)
1096                .await
1097        };
1098
1099        result.and_then(|output| query_interceptor.post_execute(output, query_ctx))
1100    }
1101
1102    #[tracing::instrument(skip_all, name = "SqlQueryHandler::do_promql_query")]
1103    async fn do_promql_query_inner(
1104        &self,
1105        query: &PromQuery,
1106        query_ctx: QueryContextRef,
1107    ) -> Vec<Result<Output>> {
1108        if self.is_suspended() {
1109            return vec![error::SuspendedSnafu {}.fail()];
1110        }
1111
1112        // check will be done in prometheus handler's do_query
1113        let result = PrometheusHandler::do_query(self, query, query_ctx)
1114            .await
1115            .with_context(|_| ExecutePromqlSnafu {
1116                query: format!("{query:?}"),
1117            });
1118        vec![result]
1119    }
1120
1121    /// Builds the [`DataFrame`] for an information-schema-backed `SHOW`
1122    /// statement; `None` for other statements. The future is boxed to keep
1123    /// `do_describe_inner`'s state machine small.
1124    fn show_statement_dataframe<'a>(
1125        &'a self,
1126        stmt: &'a Statement,
1127        query_ctx: &'a QueryContextRef,
1128    ) -> Pin<Box<dyn Future<Output = Option<query::error::Result<DataFrame>>> + Send + 'a>> {
1129        Box::pin(async move {
1130            let engine = &self.query_engine;
1131            let catalog_manager = self.catalog_manager();
1132            let ctx = query_ctx.clone();
1133            let dataframe = match stmt {
1134                Statement::ShowDatabases(show) => {
1135                    query::sql::show_databases_dataframe(show, engine, catalog_manager, ctx).await
1136                }
1137                Statement::ShowTables(show) => {
1138                    query::sql::show_tables_dataframe(show, engine, catalog_manager, ctx).await
1139                }
1140                Statement::ShowViews(show) => {
1141                    query::sql::show_views_dataframe(show, engine, catalog_manager, ctx).await
1142                }
1143                Statement::ShowFlows(show) => {
1144                    query::sql::show_flows_dataframe(show, engine, catalog_manager, ctx).await
1145                }
1146                Statement::ShowColumns(show) => {
1147                    query::sql::show_columns_dataframe(show, engine, catalog_manager, ctx).await
1148                }
1149                Statement::ShowTableStatus(show) => {
1150                    query::sql::show_table_status_dataframe(show, engine, catalog_manager, ctx)
1151                        .await
1152                }
1153                Statement::ShowCharset(kind) => {
1154                    query::sql::show_charsets_dataframe(kind, engine, catalog_manager, ctx).await
1155                }
1156                Statement::ShowCollation(kind) => {
1157                    query::sql::show_collations_dataframe(kind, engine, catalog_manager, ctx).await
1158                }
1159                Statement::ShowIndex(show) => {
1160                    query::sql::show_index_dataframe(show, engine, catalog_manager, ctx).await
1161                }
1162                Statement::ShowRegion(show) => {
1163                    query::sql::show_region_dataframe(show, engine, catalog_manager, ctx).await
1164                }
1165                Statement::ShowProcesslist(show) => {
1166                    query::sql::show_processlist_dataframe(show, engine, catalog_manager, ctx).await
1167                }
1168                _ => return None,
1169            };
1170            Some(dataframe)
1171        })
1172    }
1173
1174    async fn do_describe_inner(
1175        &self,
1176        stmt: Statement,
1177        query_ctx: QueryContextRef,
1178    ) -> Result<Option<DescribeResult>> {
1179        ensure!(!self.is_suspended(), error::SuspendedSnafu);
1180
1181        // EXPLAIN / EXPLAIN ANALYZE wrap an inner statement; describe them when the
1182        // wrapped statement is something we already plan (so that bind parameters
1183        // in the inner query get their types inferred). See #8029.
1184        let is_inner_plannable = |s: &Statement| {
1185            matches!(
1186                s,
1187                Statement::Insert(_) | Statement::Query(_) | Statement::Delete(_)
1188            )
1189        };
1190        let plannable = is_inner_plannable(&stmt)
1191            || matches!(&stmt, Statement::Explain(explain) if is_inner_plannable(explain.statement.as_ref()));
1192
1193        if let Statement::Tql(tql) = stmt {
1194            // TQL produces a logical plan; describe it from the plan so the
1195            // extended-protocol RowDescription matches the executed DataRows.
1196            self.check_sql_permission(&Statement::Tql(tql.clone()), &query_ctx)
1197                .await?;
1198            let plan = self.statement_executor.plan_tql(tql, &query_ctx).await?;
1199            return self
1200                .query_engine
1201                .describe(plan, query_ctx)
1202                .await
1203                .map(Some)
1204                .context(error::DescribeStatementSnafu);
1205        }
1206
1207        // Describe SHOW statements from the same projection the executor builds.
1208        if let Some(dataframe) = self
1209            .show_statement_dataframe(&stmt, &query_ctx)
1210            .await
1211            .transpose()
1212            .context(PlanStatementSnafu)?
1213        {
1214            self.check_sql_permission(&stmt, &query_ctx).await?;
1215            let plan = dataframe.into_unoptimized_plan();
1216            return self
1217                .query_engine
1218                .describe(plan, query_ctx)
1219                .await
1220                .map(Some)
1221                .context(error::DescribeStatementSnafu);
1222        }
1223
1224        if plannable {
1225            self.check_sql_permission(&stmt, &query_ctx).await?;
1226
1227            let plan = self
1228                .query_engine
1229                .planner()
1230                .plan(&QueryStatement::Sql(stmt), query_ctx.clone())
1231                .await
1232                .context(PlanStatementSnafu)?;
1233            self.query_engine
1234                .describe(plan, query_ctx)
1235                .await
1236                .map(Some)
1237                .context(error::DescribeStatementSnafu)
1238        } else {
1239            Ok(None)
1240        }
1241    }
1242
1243    async fn is_valid_schema_inner(&self, catalog: &str, schema: &str) -> Result<bool> {
1244        self.catalog_manager
1245            .schema_exists(catalog, schema, None)
1246            .await
1247            .context(error::CatalogSnafu)
1248    }
1249}
1250
1251#[async_trait]
1252impl SqlQueryHandler for Instance {
1253    fn supports_metric_batch_ddl(&self) -> bool {
1254        true
1255    }
1256
1257    async fn create_logical_tables(
1258        &self,
1259        query: &str,
1260        query_ctx: QueryContextRef,
1261    ) -> server_error::Result<Vec<Output>> {
1262        self.batch_create_logical_tables(query, query_ctx)
1263            .await
1264            .map_err(BoxedError::new)
1265            .context(ExecuteQuerySnafu)
1266    }
1267
1268    async fn do_query(
1269        &self,
1270        query: &str,
1271        query_ctx: QueryContextRef,
1272    ) -> Vec<server_error::Result<Output>> {
1273        self.do_query_inner(query, query_ctx)
1274            .await
1275            .into_iter()
1276            .map(|result| result.map_err(BoxedError::new).context(ExecuteQuerySnafu))
1277            .collect()
1278    }
1279
1280    async fn do_analyze_stream_query(
1281        &self,
1282        query: &str,
1283        query_ctx: QueryContextRef,
1284    ) -> server_error::Result<Output> {
1285        self.do_analyze_stream_query_inner(query, query_ctx)
1286            .await
1287            .map_err(BoxedError::new)
1288            .context(ExecuteQuerySnafu)
1289    }
1290
1291    async fn do_exec_plan(
1292        &self,
1293        plan: LogicalPlan,
1294        stmt: Option<Statement>,
1295        query_ctx: QueryContextRef,
1296    ) -> server_error::Result<Output> {
1297        self.do_exec_plan_inner(plan, stmt, query_ctx)
1298            .await
1299            .map_err(BoxedError::new)
1300            .context(server_error::ExecutePlanSnafu)
1301    }
1302
1303    async fn do_promql_query(
1304        &self,
1305        query: &PromQuery,
1306        query_ctx: QueryContextRef,
1307    ) -> Vec<server_error::Result<Output>> {
1308        self.do_promql_query_inner(query, query_ctx)
1309            .await
1310            .into_iter()
1311            .map(|result| result.map_err(BoxedError::new).context(ExecuteQuerySnafu))
1312            .collect()
1313    }
1314
1315    async fn do_describe(
1316        &self,
1317        stmt: Statement,
1318        query_ctx: QueryContextRef,
1319    ) -> server_error::Result<Option<DescribeResult>> {
1320        self.do_describe_inner(stmt, query_ctx)
1321            .await
1322            .map_err(BoxedError::new)
1323            .context(server_error::DescribeStatementSnafu)
1324    }
1325
1326    async fn is_valid_schema(&self, catalog: &str, schema: &str) -> server_error::Result<bool> {
1327        self.is_valid_schema_inner(catalog, schema)
1328            .await
1329            .map_err(BoxedError::new)
1330            .context(server_error::CheckDatabaseValiditySnafu)
1331    }
1332}
1333
1334/// Expands scan-time dictionaries only when a query result leaves the frontend.
1335pub(crate) fn map_query_output(output: Output) -> Result<Output> {
1336    output
1337        .map_dictionary_to_values()
1338        .context(CollectRecordbatchSnafu)
1339}
1340
1341/// Attaches a timer to the output and observes it once the output is exhausted.
1342pub fn attach_timer(output: Output, timer: HistogramTimer) -> Output {
1343    match output.data {
1344        OutputData::AffectedRows(_) | OutputData::RecordBatches(_) => output,
1345        OutputData::Stream(stream) => {
1346            let stream = OnDone::new(stream, move || {
1347                timer.observe_duration();
1348            });
1349            Output::new(OutputData::Stream(Box::pin(stream)), output.meta)
1350        }
1351    }
1352}
1353
1354impl Instance {
1355    fn check_prom_query_privilege(&self, query_ctx: &QueryContextRef) -> server_error::Result<()> {
1356        self.plugins
1357            .get::<PermissionCheckerRef>()
1358            .as_ref()
1359            .check_permission(
1360                query_ctx.current_user(),
1361                PermissionReq::Action(PROMQL_QUERY),
1362            )
1363            .context(AuthSnafu)?;
1364        Ok(())
1365    }
1366
1367    fn prom_expr_permission_targets(
1368        &self,
1369        expr: &promql_parser::parser::Expr,
1370        query_ctx: &QueryContextRef,
1371    ) -> server_error::Result<Option<Vec<PermissionTableTarget>>> {
1372        extract_tables_from_prom_expr_checked(expr)
1373            .map(|tables| {
1374                tables
1375                    .map(|name| {
1376                        table_idents_to_full_name(&name, query_ctx).map(
1377                            |(catalog, schema, table)| {
1378                                PermissionTableTarget::new(catalog, schema, table)
1379                            },
1380                        )
1381                    })
1382                    .collect::<std::result::Result<Vec<_>, _>>()
1383                    .map_err(BoxedError::new)
1384                    .context(ExecuteQuerySnafu)
1385            })
1386            .transpose()
1387    }
1388
1389    async fn is_physical_query_permission_target(
1390        &self,
1391        target: &PermissionTableTarget,
1392        query_ctx: &QueryContextRef,
1393    ) -> server_error::Result<bool> {
1394        self.catalog_manager
1395            .table(
1396                &target.catalog,
1397                &target.schema,
1398                &target.table,
1399                Some(query_ctx),
1400            )
1401            .await
1402            .map(|table| table.is_some_and(|table| table.table_info().is_physical_table()))
1403            .map_err(BoxedError::new)
1404            .context(ExecuteQuerySnafu)
1405    }
1406
1407    async fn resolve_query_permission_targets(
1408        &self,
1409        targets: PermissionTableTargets,
1410        query_ctx: &QueryContextRef,
1411    ) -> server_error::Result<PermissionTableTargets> {
1412        const CONCURRENCY: usize = 8;
1413
1414        let checker = self.plugins.get::<PermissionCheckerRef>();
1415        if !checker.as_ref().uses_table_targets() {
1416            return Ok(targets);
1417        }
1418
1419        let PermissionTableTargets::Resolved(mut targets) = targets else {
1420            return Ok(PermissionTableTargets::Unresolved);
1421        };
1422        if targets.len() > 1 {
1423            let mut seen = HashSet::with_capacity(targets.len());
1424            targets.retain(|target| seen.insert(target.clone()));
1425        }
1426        if let [target] = targets.as_slice() {
1427            return if self
1428                .is_physical_query_permission_target(target, query_ctx)
1429                .await?
1430            {
1431                Ok(PermissionTableTargets::Unresolved)
1432            } else {
1433                Ok(PermissionTableTargets::resolved(targets))
1434            };
1435        }
1436
1437        // Bound catalog load and inspect results in target order to preserve serial semantics.
1438        for chunk in targets.chunks(CONCURRENCY) {
1439            let results = future::join_all(
1440                chunk
1441                    .iter()
1442                    .map(|target| self.is_physical_query_permission_target(target, query_ctx)),
1443            )
1444            .await;
1445            for result in results {
1446                if result? {
1447                    return Ok(PermissionTableTargets::Unresolved);
1448                }
1449            }
1450        }
1451
1452        Ok(PermissionTableTargets::resolved(targets))
1453    }
1454
1455    fn prom_queries_permission_targets(
1456        &self,
1457        queries: &[ParsedPromQuery],
1458        query_ctx: &QueryContextRef,
1459    ) -> server_error::Result<PermissionTableTargets> {
1460        let mut targets = Vec::new();
1461        let mut resolved = true;
1462
1463        for query in queries {
1464            let QueryStatement::Promql(eval_stmt, _) = query.statement() else {
1465                unreachable!("query is parsed from promql");
1466            };
1467
1468            if let Some(query_targets) =
1469                self.prom_expr_permission_targets(&eval_stmt.expr, query_ctx)?
1470            {
1471                targets.extend(query_targets);
1472            } else {
1473                resolved = false;
1474            }
1475        }
1476
1477        Ok(if resolved {
1478            PermissionTableTargets::resolved(targets)
1479        } else {
1480            PermissionTableTargets::Unresolved
1481        })
1482    }
1483}
1484
1485#[async_trait]
1486impl PrometheusHandler for Instance {
1487    #[tracing::instrument(skip_all)]
1488    async fn do_query(
1489        &self,
1490        query: &PromQuery,
1491        query_ctx: QueryContextRef,
1492    ) -> server_error::Result<Output> {
1493        let query = ParsedPromQuery::parse(query.clone(), &query_ctx)?;
1494        self.do_query_parsed(query, query_ctx).await
1495    }
1496
1497    #[tracing::instrument(skip_all)]
1498    async fn do_query_parsed(
1499        &self,
1500        query: ParsedPromQuery,
1501        query_ctx: QueryContextRef,
1502    ) -> server_error::Result<Output> {
1503        let interceptor = self
1504            .plugins
1505            .get::<PromQueryInterceptorRef<server_error::Error>>();
1506
1507        self.check_prom_query_privilege(&query_ctx)?;
1508
1509        let targets =
1510            self.prom_queries_permission_targets(std::slice::from_ref(&query), &query_ctx)?;
1511        self.check_query_target_permission(targets, &query_ctx)
1512            .await?;
1513
1514        let requires_output_ordering = query.requires_output_ordering();
1515        let (query, stmt) = query.into_parts();
1516
1517        let QueryStatement::Promql(eval_stmt, _) = &stmt else {
1518            unreachable!("query is parsed from promql");
1519        };
1520
1521        let plan = self
1522            .statement_executor
1523            .plan(&stmt, query_ctx.clone())
1524            .await
1525            .map_err(BoxedError::new)
1526            .context(ExecuteQuerySnafu)?;
1527
1528        let plan = if requires_output_ordering {
1529            plan
1530        } else {
1531            promql::remove_output_sort(plan)
1532        };
1533
1534        interceptor.pre_execute(&query, &eval_stmt.expr, Some(&plan), query_ctx.clone())?;
1535
1536        // Take the EvalStmt from the original QueryStatement and use it to create the CatalogQueryStatement.
1537        let query_statement = if let QueryStatement::Promql(eval_stmt, alias) = stmt {
1538            CatalogQueryStatement::Promql(eval_stmt, alias)
1539        } else {
1540            // It should not happen since the query is already parsed successfully.
1541            return UnexpectedResultSnafu {
1542                reason: "The query should always be promql.".to_string(),
1543            }
1544            .fail();
1545        };
1546        let raw_query = query_statement.to_string();
1547
1548        let slow_query_timer = self.slow_query_options.enable.then(|| {
1549            SlowQueryTimer::new(
1550                query_statement,
1551                query_ctx.current_schema(),
1552                self.slow_query_options.threshold,
1553                self.slow_query_options.sample_ratio,
1554                self.slow_query_options.record_type,
1555                self.slow_query_recorder.clone(),
1556            )
1557        });
1558
1559        let ticket = self.process_manager.register_query(
1560            query_ctx.current_catalog().to_string(),
1561            vec![query_ctx.current_schema()],
1562            raw_query,
1563            query_ctx.conn_info().to_string(),
1564            Some(query_ctx.process_id()),
1565            slow_query_timer,
1566        );
1567
1568        let query_fut = self.statement_executor.exec_plan(plan, query_ctx.clone());
1569
1570        let output = CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
1571            .await
1572            .map_err(|_| servers::error::CancelledSnafu.build())?
1573            .map_err(BoxedError::new)
1574            .context(ExecuteQuerySnafu)?;
1575        let output = map_query_output(output)
1576            .map_err(BoxedError::new)
1577            .context(ExecuteQuerySnafu)?;
1578        let Output { meta, data } = output;
1579        let data = match data {
1580            OutputData::Stream(stream) => {
1581                OutputData::Stream(Box::pin(CancellableStreamWrapper::new(stream, ticket)))
1582            }
1583            other => other,
1584        };
1585        let output = Output { data, meta };
1586        Ok(interceptor.post_execute(output, query_ctx)?)
1587    }
1588
1589    async fn check_query_permission(
1590        &self,
1591        queries: &[PromQuery],
1592        query_ctx: &QueryContextRef,
1593    ) -> server_error::Result<()> {
1594        let queries = queries
1595            .iter()
1596            .cloned()
1597            .map(|query| ParsedPromQuery::parse(query, query_ctx))
1598            .collect::<server_error::Result<Vec<_>>>()?;
1599        self.check_query_permission_parsed(&queries, query_ctx)
1600            .await
1601    }
1602
1603    async fn check_query_permission_parsed(
1604        &self,
1605        queries: &[ParsedPromQuery],
1606        query_ctx: &QueryContextRef,
1607    ) -> server_error::Result<()> {
1608        self.check_prom_query_privilege(query_ctx)?;
1609        let targets = self.prom_queries_permission_targets(queries, query_ctx)?;
1610        self.check_query_target_permission(targets, query_ctx).await
1611    }
1612
1613    async fn check_query_target_permission(
1614        &self,
1615        targets: PermissionTableTargets,
1616        query_ctx: &QueryContextRef,
1617    ) -> server_error::Result<()> {
1618        let targets = self
1619            .resolve_query_permission_targets(targets, query_ctx)
1620            .await?;
1621        self.check_table_permission(query_ctx, PermissionReq::Action(PROMQL_QUERY), targets)
1622            .context(AuthSnafu)?;
1623        Ok(())
1624    }
1625
1626    async fn filter_metadata_metric_names(
1627        &self,
1628        metric_names: Vec<String>,
1629        schema: &str,
1630        query_ctx: &QueryContextRef,
1631    ) -> server_error::Result<Vec<String>> {
1632        let checker = self.plugins.get::<PermissionCheckerRef>();
1633        if !checker.as_ref().uses_table_targets() {
1634            let Some(metric) = metric_names.first() else {
1635                return Ok(metric_names);
1636            };
1637            let target =
1638                PermissionTableTarget::new(query_ctx.current_catalog(), schema, metric.as_str());
1639            let result = checker
1640                .as_ref()
1641                .check_permission_with_table_targets(
1642                    query_ctx.current_user(),
1643                    PermissionReq::Action(PROMQL_QUERY),
1644                    PermissionTableTargets::resolved(vec![target]),
1645                )
1646                .context(AuthSnafu);
1647            return match result {
1648                Ok(_) => Ok(metric_names),
1649                Err(error)
1650                    if error.status_code()
1651                        == common_error::status_code::StatusCode::PermissionDenied =>
1652                {
1653                    Ok(Vec::new())
1654                }
1655                Err(error) => Err(error),
1656            };
1657        }
1658
1659        let mut allowed = Vec::with_capacity(metric_names.len());
1660        for metric in metric_names {
1661            let target =
1662                PermissionTableTarget::new(query_ctx.current_catalog(), schema, metric.as_str());
1663            match checker
1664                .as_ref()
1665                .check_permission_with_table_targets(
1666                    query_ctx.current_user(),
1667                    PermissionReq::Action(PROMQL_QUERY),
1668                    PermissionTableTargets::resolved(vec![target]),
1669                )
1670                .context(AuthSnafu)
1671            {
1672                Ok(_) => allowed.push(metric),
1673                Err(error)
1674                    if error.status_code()
1675                        == common_error::status_code::StatusCode::PermissionDenied => {}
1676                Err(error) => return Err(error),
1677            }
1678        }
1679        Ok(allowed)
1680    }
1681
1682    async fn query_metric_names(
1683        &self,
1684        matchers: Vec<Matcher>,
1685        schema: &str,
1686        ctx: &QueryContextRef,
1687    ) -> server_error::Result<Vec<String>> {
1688        self.handle_query_metric_names(matchers, schema, ctx)
1689            .await
1690            .map_err(BoxedError::new)
1691            .context(ExecuteQuerySnafu)
1692    }
1693
1694    async fn query_metric_names_by_labels(
1695        &self,
1696        matchers: Vec<Matcher>,
1697        schema: &str,
1698        start: SystemTime,
1699        end: SystemTime,
1700        ctx: &QueryContextRef,
1701    ) -> server_error::Result<Vec<String>> {
1702        self.handle_query_metric_names_by_labels(matchers, schema, start, end, ctx)
1703            .await
1704            .map_err(BoxedError::new)
1705            .context(ExecuteQuerySnafu)
1706    }
1707
1708    async fn query_label_values(
1709        &self,
1710        metric: String,
1711        label_name: String,
1712        matchers: Vec<Matcher>,
1713        start: SystemTime,
1714        end: SystemTime,
1715        ctx: &QueryContextRef,
1716    ) -> server_error::Result<Vec<String>> {
1717        let schema =
1718            resolve_schema_from_matchers(&matchers)?.unwrap_or_else(|| ctx.current_schema());
1719        let target = PermissionTableTarget::new(ctx.current_catalog(), schema.as_str(), &metric);
1720        self.check_query_target_permission(
1721            PermissionTableTargets::resolved(vec![target.clone()]),
1722            ctx,
1723        )
1724        .await?;
1725
1726        self.handle_query_label_values(target, label_name, matchers, start, end, ctx)
1727            .await
1728            .map_err(BoxedError::new)
1729            .context(ExecuteQuerySnafu)
1730    }
1731
1732    fn catalog_manager(&self) -> CatalogManagerRef {
1733        self.catalog_manager.clone()
1734    }
1735}
1736
1737/// Validate `stmt.database` permission if it's presented.
1738macro_rules! validate_db_permission {
1739    ($stmt: expr, $query_ctx: expr) => {
1740        if let Some(database) = &$stmt.database {
1741            validate_catalog_and_schema($query_ctx.current_catalog(), database, $query_ctx)
1742                .map_err(BoxedError::new)
1743                .context(SqlExecInterceptedSnafu)?;
1744        }
1745    };
1746}
1747
1748pub fn check_permission(
1749    plugins: Plugins,
1750    stmt: &Statement,
1751    query_ctx: &QueryContextRef,
1752) -> Result<()> {
1753    let need_validate = plugins
1754        .get::<QueryOptions>()
1755        .map(|opts| opts.disallow_cross_catalog_query)
1756        .unwrap_or_default();
1757
1758    if !need_validate {
1759        return Ok(());
1760    }
1761
1762    match stmt {
1763        // Will be checked in execution.
1764        // TODO(dennis): add a hook for admin commands.
1765        Statement::Admin(_) => {}
1766        // These are executed by query engine, and will be checked there.
1767        Statement::Query(_)
1768        | Statement::Explain(_)
1769        | Statement::Tql(_)
1770        | Statement::Delete(_)
1771        | Statement::DeclareCursor(_)
1772        | Statement::Copy(sql::statements::copy::Copy::CopyQueryTo(_)) => {}
1773        // database ops won't be checked
1774        Statement::CreateDatabase(_)
1775        | Statement::ShowDatabases(_)
1776        | Statement::DropDatabase(_)
1777        | Statement::AlterDatabase(_)
1778        | Statement::DropFlow(_)
1779        | Statement::Use(_) => {}
1780        #[cfg(feature = "enterprise")]
1781        Statement::DropTrigger(_) => {}
1782        Statement::ShowCreateDatabase(stmt) => {
1783            validate_database(&stmt.database_name, query_ctx)?;
1784        }
1785        Statement::ShowCreateTable(stmt) => {
1786            validate_param(&stmt.table_name, query_ctx)?;
1787        }
1788        Statement::ShowCreateFlow(stmt) => {
1789            validate_flow(&stmt.flow_name, query_ctx)?;
1790        }
1791        #[cfg(feature = "enterprise")]
1792        Statement::ShowCreateTrigger(stmt) => {
1793            validate_param(&stmt.trigger_name, query_ctx)?;
1794        }
1795        Statement::ShowCreateView(stmt) => {
1796            validate_param(&stmt.view_name, query_ctx)?;
1797        }
1798        Statement::CreateExternalTable(stmt) => {
1799            validate_param(&stmt.name, query_ctx)?;
1800        }
1801        Statement::CreateFlow(stmt) => {
1802            // TODO: should also validate source table name here?
1803            validate_param(&stmt.sink_table_name, query_ctx)?;
1804        }
1805        #[cfg(feature = "enterprise")]
1806        Statement::CreateTrigger(stmt) => {
1807            validate_param(&stmt.trigger_name, query_ctx)?;
1808        }
1809        Statement::CreateView(stmt) => {
1810            validate_param(&stmt.name, query_ctx)?;
1811        }
1812        Statement::AlterTable(stmt) => {
1813            validate_param(stmt.table_name(), query_ctx)?;
1814        }
1815        #[cfg(feature = "enterprise")]
1816        Statement::AlterTrigger(_) => {}
1817        // set/show variable now only alter/show variable in session
1818        Statement::SetVariables(_) | Statement::ShowVariables(_) => {}
1819        // show charset and show collation won't be checked
1820        Statement::ShowCharset(_) | Statement::ShowCollation(_) => {}
1821
1822        Statement::Comment(comment) => match &comment.object {
1823            CommentObject::Table(table) => validate_param(table, query_ctx)?,
1824            CommentObject::Column { table, .. } => validate_param(table, query_ctx)?,
1825            CommentObject::Flow(flow) => validate_flow(flow, query_ctx)?,
1826        },
1827
1828        Statement::Insert(insert) => {
1829            let name = insert.table_name().context(ParseSqlSnafu)?;
1830            validate_param(name, query_ctx)?;
1831        }
1832        Statement::CreateTable(stmt) => {
1833            validate_param(&stmt.name, query_ctx)?;
1834        }
1835        Statement::CreateTableLike(stmt) => {
1836            validate_param(&stmt.table_name, query_ctx)?;
1837            validate_param(&stmt.source_name, query_ctx)?;
1838        }
1839        Statement::DropTable(drop_stmt) => {
1840            for table_name in drop_stmt.table_names() {
1841                validate_param(table_name, query_ctx)?;
1842            }
1843        }
1844        #[cfg(feature = "enterprise")]
1845        Statement::UndropTable(stmt) => {
1846            validate_param(stmt.table_name(), query_ctx)?;
1847        }
1848        Statement::DropView(stmt) => {
1849            validate_param(&stmt.view_name, query_ctx)?;
1850        }
1851        Statement::ShowTables(stmt) => {
1852            validate_db_permission!(stmt, query_ctx);
1853        }
1854        Statement::ShowTableStatus(stmt) => {
1855            validate_db_permission!(stmt, query_ctx);
1856        }
1857        Statement::ShowColumns(stmt) => {
1858            validate_db_permission!(stmt, query_ctx);
1859        }
1860        Statement::ShowIndex(stmt) => {
1861            validate_db_permission!(stmt, query_ctx);
1862        }
1863        Statement::ShowRegion(stmt) => {
1864            validate_db_permission!(stmt, query_ctx);
1865        }
1866        Statement::ShowViews(stmt) => {
1867            validate_db_permission!(stmt, query_ctx);
1868        }
1869        Statement::ShowFlows(stmt) => {
1870            validate_db_permission!(stmt, query_ctx);
1871        }
1872        Statement::ShowFlowStatus(_stmt) => {
1873            // Flow statistics are organized based on the catalog dimension and
1874            // filtered by the current catalog, so there is no need to check the
1875            // permission of the database(schema).
1876        }
1877        #[cfg(feature = "enterprise")]
1878        Statement::ShowTriggers(_stmt) => {
1879            // The trigger is organized based on the catalog dimension, so there
1880            // is no need to check the permission of the database(schema).
1881        }
1882        Statement::ShowStatus(_stmt) => {}
1883        Statement::ShowSearchPath(_stmt) => {}
1884        Statement::DescribeTable(stmt) => {
1885            validate_param(stmt.name(), query_ctx)?;
1886        }
1887        Statement::Copy(sql::statements::copy::Copy::CopyTable(stmt)) => match stmt {
1888            CopyTable::To(copy_table_to) => validate_param(&copy_table_to.table_name, query_ctx)?,
1889            CopyTable::From(copy_table_from) => {
1890                validate_param(&copy_table_from.table_name, query_ctx)?
1891            }
1892        },
1893        Statement::Copy(sql::statements::copy::Copy::CopyDatabase(copy_database)) => {
1894            match copy_database {
1895                CopyDatabase::To(stmt) => validate_database(&stmt.database_name, query_ctx)?,
1896                CopyDatabase::From(stmt) => validate_database(&stmt.database_name, query_ctx)?,
1897            }
1898        }
1899        Statement::TruncateTable(stmt) => {
1900            validate_param(stmt.table_name(), query_ctx)?;
1901        }
1902        // cursor operations are always allowed once it's created
1903        Statement::FetchCursor(_) | Statement::CloseCursor(_) => {}
1904        // User can only kill process in their own catalog.
1905        Statement::Kill(_) => {}
1906        // SHOW PROCESSLIST
1907        Statement::ShowProcesslist(_) => {}
1908    }
1909    Ok(())
1910}
1911
1912fn validate_param(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
1913    let (catalog, schema, _) = table_idents_to_full_name(name, query_ctx)
1914        .map_err(BoxedError::new)
1915        .context(ExternalSnafu)?;
1916
1917    validate_catalog_and_schema(&catalog, &schema, query_ctx)
1918        .map_err(BoxedError::new)
1919        .context(SqlExecInterceptedSnafu)
1920}
1921
1922fn validate_flow(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
1923    let catalog = match &name.0[..] {
1924        [_flow] => query_ctx.current_catalog().to_string(),
1925        [catalog, _flow] => catalog.to_string_unquoted(),
1926        _ => {
1927            return InvalidSqlSnafu {
1928                err_msg: format!(
1929                    "expect flow name to be <catalog>.<flow_name> or <flow_name>, actual: {name}",
1930                ),
1931            }
1932            .fail();
1933        }
1934    };
1935
1936    let schema = query_ctx.current_schema();
1937
1938    validate_catalog_and_schema(&catalog, &schema, query_ctx)
1939        .map_err(BoxedError::new)
1940        .context(SqlExecInterceptedSnafu)
1941}
1942
1943fn validate_database(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
1944    let (catalog, schema) = match &name.0[..] {
1945        [schema] => (
1946            query_ctx.current_catalog().to_string(),
1947            schema.to_string_unquoted(),
1948        ),
1949        [catalog, schema] => (catalog.to_string_unquoted(), schema.to_string_unquoted()),
1950        _ => InvalidSqlSnafu {
1951            err_msg: format!(
1952                "expect database name to be <catalog>.<schema> or <schema>, actual: {name}",
1953            ),
1954        }
1955        .fail()?,
1956    };
1957
1958    validate_catalog_and_schema(&catalog, &schema, query_ctx)
1959        .map_err(BoxedError::new)
1960        .context(SqlExecInterceptedSnafu)
1961}
1962
1963fn is_readonly_plan(plan: &LogicalPlan) -> bool {
1964    !matches!(plan, LogicalPlan::Dml(_) | LogicalPlan::Ddl(_))
1965}
1966
1967fn should_track_statement_process(stmt: &Statement) -> bool {
1968    stmt.is_readonly()
1969        || matches!(stmt, Statement::Insert(insert) if insert.has_non_values_query_source())
1970}
1971
1972fn should_track_plan_process(stmt: Option<&Statement>, plan: &LogicalPlan) -> bool {
1973    is_readonly_plan(plan)
1974        || matches!(stmt, Some(Statement::Insert(insert)) if insert.has_non_values_query_source())
1975}
1976
1977#[cfg(test)]
1978mod tests {
1979    use std::any::Any;
1980    use std::collections::HashMap;
1981    use std::future::Future;
1982    use std::pin::Pin;
1983    use std::sync::Arc;
1984    use std::task::{Context, Poll};
1985    use std::time::Duration;
1986
1987    use api::prom_store::remote::label_matcher::Type as PromMatcherType;
1988    use api::prom_store::remote::{
1989        Label, LabelMatcher, Query as RemoteQuery, ReadRequest, ReadResponse, Sample,
1990    };
1991    use api::v1::greptime_request::Request;
1992    use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse};
1993    use api::v1::query_request::Query;
1994    use auth::{
1995        DASHBOARD_DELETE, DASHBOARD_QUERY, DASHBOARD_SAVE, JAEGER_QUERY, PIPELINE_DELETE,
1996        PIPELINE_INSERT, PIPELINE_QUERY, PermissionAction, PermissionResp, UserInfo, UserInfoRef,
1997    };
1998    use catalog::process_manager::{ProcessManager, QueryStatement, SlowQueryTimer};
1999    use common_base::Plugins;
2000    use common_catalog::consts::DEFAULT_PRIVATE_SCHEMA_NAME;
2001    use common_error::ext::{BoxedError, ErrorExt, PlainError};
2002    use common_error::status_code::StatusCode;
2003    use common_event_recorder::{Event, EventRecorder, EventTypeFilter, EventTypeFilterRef};
2004    use common_frontend::slow_query_event::SlowQueryEvent;
2005    use common_meta::cache::LayeredCacheRegistryBuilder;
2006    use common_meta::kv_backend::memory::MemoryKvBackend;
2007    use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor};
2008    use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse};
2009    use common_meta::rpc::procedure::{
2010        MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse,
2011    };
2012    use common_query::prelude::greptime_value;
2013    use common_query::{Output, OutputMeta};
2014    use common_recordbatch::{
2015        OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream,
2016    };
2017    use common_telemetry::logging::SlowQueriesRecordType;
2018    use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
2019    use datafusion::physical_plan::empty::EmptyExec;
2020    use datafusion_expr::dml::InsertOp;
2021    use datafusion_expr::{LogicalPlanBuilder, LogicalTableSource};
2022    use datatypes::prelude::ConcreteDataType;
2023    use datatypes::schema::{ColumnSchema, Schema as GtSchema, SchemaRef as GtSchemaRef};
2024    use datatypes::vectors::{
2025        Float64Vector, StringVector, TimestampMillisecondVector, TimestampNanosecondVector,
2026        VectorRef,
2027    };
2028    use log_query::LogQuery;
2029    use prost::Message;
2030    use query::query_engine::options::QueryOptions;
2031    use servers::query_handler::{
2032        DashboardHandler, JaegerQueryHandler, LogQueryHandler, PipelineHandler, PipelineHandlerRef,
2033        PromStoreProtocolHandler,
2034    };
2035    use session::context::{Channel, ConnInfo, QueryContext, QueryContextBuilder};
2036    use snafu::{Location, Snafu};
2037    use sql::dialect::GreptimeDbDialect;
2038    use store_api::data_source::DataSource;
2039    use store_api::metric_engine_consts::{
2040        LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY,
2041    };
2042    use store_api::storage::ScanRequest;
2043    use strfmt::Format;
2044    use table::metadata::{
2045        FilterPushDownType, TableInfo, TableInfoBuilder, TableMetaBuilder, TableType,
2046    };
2047    use table::table_name::TableName;
2048    use table::test_util::{EmptyTable, MemTable};
2049    use table::{Table, TableRef};
2050    use tokio::sync::{mpsc, oneshot};
2051    use tower::ServiceExt;
2052
2053    use crate::frontend::FrontendOptions;
2054    use crate::instance::builder::FrontendBuilder;
2055    use crate::instance::*;
2056
2057    fn parse_test_sql(sql: &str) -> Vec<Statement> {
2058        parse_stmt(sql, &GreptimeDbDialect {}).unwrap()
2059    }
2060
2061    #[derive(Debug, Default)]
2062    struct RecordingSlowQueryEventRecorder {
2063        payloads: std::sync::Mutex<Vec<serde_json::Value>>,
2064    }
2065
2066    impl EventRecorder for RecordingSlowQueryEventRecorder {
2067        fn record(&self, event: Box<dyn Event>) {
2068            let event = event
2069                .as_any()
2070                .downcast_ref::<SlowQueryEvent>()
2071                .expect("expected a slow query event");
2072            self.payloads.lock().unwrap().push(event.payload.clone());
2073        }
2074
2075        fn event_type_filter(&self) -> EventTypeFilterRef {
2076            Arc::new(EventTypeFilter::All)
2077        }
2078
2079        fn close(&self) {}
2080    }
2081
2082    #[test]
2083    fn test_validate_analyze_stream_statement_strictness() {
2084        for sql in [
2085            "select 1",
2086            "explain analyze select 1",
2087            "explain analyze verbose format text select 1",
2088            "explain analyze verbose format graphviz select 1",
2089            "TQL ANALYZE (0, 10, '5s') physical_metric",
2090            "TQL EXPLAIN VERBOSE (0, 10, '5s') physical_metric",
2091            "TQL ANALYZE VERBOSE FORMAT TEXT (0, 10, '5s') physical_metric",
2092        ] {
2093            let mut stmts = parse_test_sql(sql);
2094            assert!(
2095                validate_analyze_stream_statement(&mut stmts[0]).is_err(),
2096                "{sql}"
2097            );
2098        }
2099
2100        for sql in [
2101            "explain analyze verbose select 1",
2102            "explain analyze verbose format json select 1",
2103            "TQL ANALYZE VERBOSE (0, 10, '5s') physical_metric",
2104            "TQL ANALYZE VERBOSE FORMAT JSON (0, 10, '5s') physical_metric",
2105        ] {
2106            let mut stmts = parse_test_sql(sql);
2107            assert!(
2108                validate_analyze_stream_statement(&mut stmts[0]).is_ok(),
2109                "{sql}"
2110            );
2111            match &stmts[0] {
2112                Statement::Explain(explain) => assert!(explain.format.is_none()),
2113                Statement::Tql(Tql::Analyze(analyze)) => assert!(analyze.format.is_none()),
2114                _ => unreachable!(),
2115            }
2116        }
2117
2118        assert_eq!(
2119            parse_test_sql("explain analyze verbose select 1; select 2").len(),
2120            2
2121        );
2122
2123        assert!(is_explain_analyze_verbose(
2124            &parse_test_sql("explain analyze verbose select 1")[0]
2125        ));
2126        assert!(is_explain_analyze_verbose(
2127            &parse_test_sql("TQL ANALYZE VERBOSE (0, 10, '5s') physical_metric")[0]
2128        ));
2129        for sql in [
2130            "select 1",
2131            "explain select 1",
2132            "explain analyze select 1",
2133            "explain verbose select 1",
2134            "TQL ANALYZE (0, 10, '5s') physical_metric",
2135            "TQL EXPLAIN VERBOSE (0, 10, '5s') physical_metric",
2136        ] {
2137            assert!(
2138                !is_explain_analyze_verbose(&parse_test_sql(sql)[0]),
2139                "{sql}"
2140            );
2141        }
2142    }
2143
2144    #[derive(Debug, Snafu)]
2145    enum TestError {
2146        #[snafu(display("Failed to build test cache registry"))]
2147        BuildCacheRegistry {
2148            source: cache::error::Error,
2149            #[snafu(implicit)]
2150            location: Location,
2151        },
2152
2153        #[snafu(display("Failed to build test table meta for table: {table_name}"))]
2154        BuildTableMeta {
2155            table_name: String,
2156            source: table::metadata::TableMetaBuilderError,
2157            #[snafu(implicit)]
2158            location: Location,
2159        },
2160
2161        #[snafu(display("Failed to build test table info for table: {table_name}"))]
2162        BuildTableInfo {
2163            table_name: String,
2164            source: table::metadata::TableInfoBuilderError,
2165            #[snafu(implicit)]
2166            location: Location,
2167        },
2168
2169        #[snafu(display("Failed to register test table: {table_name}"))]
2170        RegisterTable {
2171            table_name: String,
2172            source: catalog::error::Error,
2173            #[snafu(implicit)]
2174            location: Location,
2175        },
2176
2177        #[snafu(display("Failed to build test frontend instance"))]
2178        BuildFrontend {
2179            source: crate::error::Error,
2180            #[snafu(implicit)]
2181            location: Location,
2182        },
2183
2184        #[snafu(display("Expected exactly one output for SQL `{sql}`, got {actual}"))]
2185        UnexpectedOutputCount {
2186            sql: String,
2187            actual: usize,
2188            #[snafu(implicit)]
2189            location: Location,
2190        },
2191
2192        #[snafu(display("Failed to execute SQL `{sql}`"))]
2193        ExecuteSql {
2194            sql: String,
2195            source: crate::error::Error,
2196            #[snafu(implicit)]
2197            location: Location,
2198        },
2199
2200        #[snafu(display("Timed out waiting for insert-select start notification"))]
2201        InsertStartTimeout {
2202            source: tokio::time::error::Elapsed,
2203            #[snafu(implicit)]
2204            location: Location,
2205        },
2206
2207        #[snafu(display("Insert-select start notification channel closed"))]
2208        InsertStartChannelClosed {
2209            #[snafu(implicit)]
2210            location: Location,
2211        },
2212
2213        #[snafu(display("Failed to release blocking insert-select interceptor"))]
2214        ReleaseBlockedInsert {
2215            #[snafu(implicit)]
2216            location: Location,
2217        },
2218
2219        #[snafu(display("Timed out waiting for insert-select source to be polled"))]
2220        SourcePollTimeout {
2221            source: tokio::time::error::Elapsed,
2222            #[snafu(implicit)]
2223            location: Location,
2224        },
2225
2226        #[snafu(display("Insert-select source poll notification channel closed"))]
2227        SourcePollChannelClosed {
2228            source: oneshot::error::RecvError,
2229            #[snafu(implicit)]
2230            location: Location,
2231        },
2232
2233        #[snafu(display("Timed out waiting for insert task to finish"))]
2234        InsertTaskTimeout {
2235            source: tokio::time::error::Elapsed,
2236            #[snafu(implicit)]
2237            location: Location,
2238        },
2239
2240        #[snafu(display("Insert task panicked"))]
2241        InsertTaskPanic {
2242            source: tokio::task::JoinError,
2243            #[snafu(implicit)]
2244            location: Location,
2245        },
2246
2247        #[snafu(display("Expected insert-select to be cancelled"))]
2248        InsertSelectNotCancelled {
2249            #[snafu(implicit)]
2250            location: Location,
2251        },
2252    }
2253
2254    type TestResult<T> = std::result::Result<T, TestError>;
2255
2256    fn parse_one_sql(sql: &str) -> Statement {
2257        parse_stmt(sql, &GreptimeDbDialect {}).unwrap().remove(0)
2258    }
2259
2260    fn test_query_ctx(process_id: u32) -> QueryContextRef {
2261        Arc::new(
2262            QueryContextBuilder::default()
2263                .channel(Channel::Mysql)
2264                .conn_info(ConnInfo::new(None, Channel::Mysql))
2265                .process_id(process_id)
2266                .build(),
2267        )
2268    }
2269
2270    #[derive(Debug)]
2271    struct AdminUserInfo;
2272
2273    impl UserInfo for AdminUserInfo {
2274        fn as_any(&self) -> &dyn Any {
2275            self
2276        }
2277
2278        fn username(&self) -> &str {
2279            "admin"
2280        }
2281
2282        fn is_admin(&self) -> bool {
2283            true
2284        }
2285    }
2286
2287    struct RejectUnresolvedPermissionChecker;
2288
2289    impl PermissionChecker for RejectUnresolvedPermissionChecker {
2290        fn check_permission(
2291            &self,
2292            _user_info: UserInfoRef,
2293            _req: PermissionReq,
2294        ) -> auth::error::Result<PermissionResp> {
2295            Ok(PermissionResp::Allow)
2296        }
2297
2298        fn check_permission_with_table_targets(
2299            &self,
2300            _user_info: UserInfoRef,
2301            _req: PermissionReq,
2302            targets: PermissionTableTargets,
2303        ) -> auth::error::Result<PermissionResp> {
2304            let reject = match targets {
2305                PermissionTableTargets::Unresolved => true,
2306                PermissionTableTargets::Resolved(targets) => {
2307                    targets.iter().any(|target| target.table == "denied")
2308                }
2309            };
2310            Ok(if reject {
2311                PermissionResp::Reject
2312            } else {
2313                PermissionResp::Allow
2314            })
2315        }
2316    }
2317
2318    #[derive(Debug, PartialEq, Eq)]
2319    struct CheckedAction {
2320        action: PermissionAction,
2321        targets: Option<PermissionTableTargets>,
2322    }
2323
2324    #[derive(Default)]
2325    struct RejectEndpointPermissionChecker {
2326        checks: std::sync::Mutex<Vec<CheckedAction>>,
2327    }
2328
2329    impl RejectEndpointPermissionChecker {
2330        fn reject(
2331            &self,
2332            action: PermissionAction,
2333            targets: Option<PermissionTableTargets>,
2334        ) -> PermissionResp {
2335            self.checks
2336                .lock()
2337                .unwrap()
2338                .push(CheckedAction { action, targets });
2339            PermissionResp::Reject
2340        }
2341
2342        fn take_check(&self) -> CheckedAction {
2343            let mut checks = self.checks.lock().unwrap();
2344            assert_eq!(1, checks.len());
2345            checks.pop().unwrap()
2346        }
2347    }
2348
2349    impl PermissionChecker for RejectEndpointPermissionChecker {
2350        fn check_permission(
2351            &self,
2352            _user_info: UserInfoRef,
2353            req: PermissionReq,
2354        ) -> auth::error::Result<PermissionResp> {
2355            Ok(match req {
2356                PermissionReq::Action(action) => self.reject(action, None),
2357                _ => PermissionResp::Allow,
2358            })
2359        }
2360
2361        fn check_permission_with_table_targets(
2362            &self,
2363            _user_info: UserInfoRef,
2364            req: PermissionReq,
2365            targets: PermissionTableTargets,
2366        ) -> auth::error::Result<PermissionResp> {
2367            Ok(match req {
2368                PermissionReq::Action(action) => self.reject(action, Some(targets)),
2369                _ => PermissionResp::Allow,
2370            })
2371        }
2372    }
2373
2374    struct WriteOnlyPermissionChecker;
2375
2376    impl PermissionChecker for WriteOnlyPermissionChecker {
2377        fn check_permission(
2378            &self,
2379            _user_info: UserInfoRef,
2380            req: PermissionReq,
2381        ) -> auth::error::Result<PermissionResp> {
2382            Ok(if req.is_readonly() {
2383                PermissionResp::Reject
2384            } else {
2385                PermissionResp::Allow
2386            })
2387        }
2388
2389        fn check_permission_with_table_targets(
2390            &self,
2391            user_info: UserInfoRef,
2392            req: PermissionReq,
2393            _targets: PermissionTableTargets,
2394        ) -> auth::error::Result<PermissionResp> {
2395            self.check_permission(user_info, req)
2396        }
2397    }
2398
2399    #[derive(Default)]
2400    struct TargetIndependentPermissionChecker {
2401        checks: atomic::AtomicUsize,
2402    }
2403
2404    impl PermissionChecker for TargetIndependentPermissionChecker {
2405        fn check_permission(
2406            &self,
2407            _user_info: UserInfoRef,
2408            _req: PermissionReq,
2409        ) -> auth::error::Result<PermissionResp> {
2410            self.checks.fetch_add(1, atomic::Ordering::Relaxed);
2411            Ok(PermissionResp::Allow)
2412        }
2413
2414        fn uses_table_targets(&self) -> bool {
2415            false
2416        }
2417
2418        fn check_permission_with_table_targets(
2419            &self,
2420            user_info: UserInfoRef,
2421            req: PermissionReq,
2422            _targets: PermissionTableTargets,
2423        ) -> auth::error::Result<PermissionResp> {
2424            self.check_permission(user_info, req)
2425        }
2426    }
2427
2428    struct BlockingInsertSelectInterceptor {
2429        started_tx: mpsc::UnboundedSender<()>,
2430        finish_rx: std::sync::Mutex<Option<oneshot::Receiver<()>>>,
2431    }
2432
2433    impl BlockingInsertSelectInterceptor {
2434        fn new(started_tx: mpsc::UnboundedSender<()>, finish_rx: oneshot::Receiver<()>) -> Self {
2435            Self {
2436                started_tx,
2437                finish_rx: std::sync::Mutex::new(Some(finish_rx)),
2438            }
2439        }
2440    }
2441
2442    impl SqlQueryInterceptor for BlockingInsertSelectInterceptor {
2443        type Error = Error;
2444
2445        fn pre_execute(
2446            &self,
2447            statement: Option<&Statement>,
2448            _plan: Option<&LogicalPlan>,
2449            _query_ctx: QueryContextRef,
2450        ) -> Result<()> {
2451            let Some(Statement::Insert(insert)) = statement else {
2452                return Ok(());
2453            };
2454            if !insert.has_non_values_query_source() {
2455                return Ok(());
2456            }
2457
2458            let finish_rx = self.finish_rx.lock().unwrap().take().unwrap();
2459            let _ = self.started_tx.send(());
2460            tokio::task::block_in_place(|| {
2461                tokio::runtime::Handle::current()
2462                    .block_on(finish_rx)
2463                    .unwrap();
2464            });
2465            Ok(())
2466        }
2467    }
2468
2469    struct PendingRecordBatchStream {
2470        schema: GtSchemaRef,
2471        polled_tx: Option<oneshot::Sender<()>>,
2472        _finish_tx: oneshot::Sender<()>,
2473        finish_rx: Pin<Box<oneshot::Receiver<()>>>,
2474    }
2475
2476    impl RecordBatchStream for PendingRecordBatchStream {
2477        fn schema(&self) -> GtSchemaRef {
2478            self.schema.clone()
2479        }
2480
2481        fn output_ordering(&self) -> Option<&[OrderOption]> {
2482            None
2483        }
2484
2485        fn metrics(&self) -> Option<common_recordbatch::adapter::RecordBatchMetrics> {
2486            None
2487        }
2488    }
2489
2490    impl Stream for PendingRecordBatchStream {
2491        type Item = common_recordbatch::error::Result<RecordBatch>;
2492
2493        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
2494            if let Some(polled_tx) = self.polled_tx.take() {
2495                let _ = polled_tx.send(());
2496            }
2497
2498            match self.finish_rx.as_mut().poll(cx) {
2499                Poll::Ready(_) => Poll::Ready(None),
2500                Poll::Pending => Poll::Pending,
2501            }
2502        }
2503    }
2504
2505    impl Unpin for PendingRecordBatchStream {}
2506
2507    #[test]
2508    fn test_record_explain_analyze_timeout_uses_empty_metrics_without_plan() {
2509        let event_recorder = Arc::new(RecordingSlowQueryEventRecorder::default());
2510        let timer = SlowQueryTimer::new(
2511            QueryStatement::Plan("EXPLAIN ANALYZE VERBOSE SELECT 1".to_string()),
2512            "public".to_string(),
2513            Duration::from_secs(3600),
2514            0.0,
2515            SlowQueriesRecordType::SystemTable,
2516            event_recorder.clone(),
2517        );
2518        let timeout_recorder = timer.recorder();
2519
2520        record_explain_analyze_timeout(Some(&timeout_recorder), None);
2521        drop(timer);
2522
2523        let payloads = event_recorder.payloads.lock().unwrap();
2524        assert_eq!(payloads.len(), 1);
2525        assert_eq!(payloads[0]["timed_out"], true);
2526        assert_eq!(payloads[0]["metrics"], serde_json::json!([]));
2527    }
2528
2529    #[tokio::test]
2530    async fn test_attach_timeout_records_explain_analyze_metrics() {
2531        let event_recorder = Arc::new(RecordingSlowQueryEventRecorder::default());
2532        let timer = SlowQueryTimer::new(
2533            QueryStatement::Plan("EXPLAIN ANALYZE VERBOSE SELECT 1".to_string()),
2534            "public".to_string(),
2535            Duration::from_secs(3600),
2536            0.0,
2537            SlowQueriesRecordType::SystemTable,
2538            event_recorder.clone(),
2539        );
2540        let timeout_recorder = timer.recorder();
2541        let plan: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(Arc::new(Schema::empty())));
2542        let (finish_tx, finish_rx) = oneshot::channel();
2543        let stream = PendingRecordBatchStream {
2544            schema: Arc::new(GtSchema::new(vec![])),
2545            polled_tx: None,
2546            _finish_tx: finish_tx,
2547            finish_rx: Box::pin(finish_rx),
2548        };
2549        let output = Output::new(
2550            OutputData::Stream(Box::pin(stream)),
2551            OutputMeta::new_with_plan(plan),
2552        );
2553        let output =
2554            attach_timeout(output, Duration::from_millis(10), Some(timeout_recorder)).unwrap();
2555        let OutputData::Stream(mut stream) = output.data else {
2556            unreachable!();
2557        };
2558
2559        let err = stream.next().await.unwrap().unwrap_err();
2560        assert_eq!(err.to_string(), "Stream timeout");
2561        drop(stream);
2562        drop(timer);
2563
2564        let payloads = event_recorder.payloads.lock().unwrap();
2565        assert_eq!(payloads.len(), 1);
2566        assert_eq!(payloads[0]["timed_out"], true);
2567        assert!(
2568            payloads[0]["metrics"]
2569                .as_array()
2570                .is_some_and(|metrics| !metrics.is_empty())
2571        );
2572    }
2573
2574    struct PendingDataSource {
2575        schema: GtSchemaRef,
2576        polled_tx: std::sync::Mutex<Option<oneshot::Sender<()>>>,
2577    }
2578
2579    impl DataSource for PendingDataSource {
2580        fn get_stream(
2581            &self,
2582            _request: ScanRequest,
2583        ) -> std::result::Result<SendableRecordBatchStream, BoxedError> {
2584            let (finish_tx, finish_rx) = oneshot::channel();
2585            let mut polled_tx = self.polled_tx.lock().map_err(|_| {
2586                BoxedError::new(PlainError::new(
2587                    "pending data source lock poisoned".to_string(),
2588                    StatusCode::Unexpected,
2589                ))
2590            })?;
2591            Ok(Box::pin(PendingRecordBatchStream {
2592                schema: self.schema.clone(),
2593                polled_tx: polled_tx.take(),
2594                _finish_tx: finish_tx,
2595                finish_rx: Box::pin(finish_rx),
2596            }))
2597        }
2598    }
2599
2600    struct NoopProcedureExecutor;
2601
2602    #[async_trait::async_trait]
2603    impl ProcedureExecutor for NoopProcedureExecutor {
2604        async fn submit_ddl_task(
2605            &self,
2606            _ctx: ExecutorContext,
2607            _request: SubmitDdlTaskRequest,
2608        ) -> common_meta::error::Result<SubmitDdlTaskResponse> {
2609            common_meta::error::UnsupportedSnafu {
2610                operation: "submit_ddl_task",
2611            }
2612            .fail()
2613        }
2614
2615        async fn migrate_region(
2616            &self,
2617            _ctx: &ExecutorContext,
2618            _request: MigrateRegionRequest,
2619        ) -> common_meta::error::Result<MigrateRegionResponse> {
2620            common_meta::error::UnsupportedSnafu {
2621                operation: "migrate_region",
2622            }
2623            .fail()
2624        }
2625
2626        async fn reconcile(
2627            &self,
2628            _ctx: &ExecutorContext,
2629            _request: ReconcileRequest,
2630        ) -> common_meta::error::Result<ReconcileResponse> {
2631            common_meta::error::UnsupportedSnafu {
2632                operation: "reconcile",
2633            }
2634            .fail()
2635        }
2636
2637        async fn query_procedure_state(
2638            &self,
2639            _ctx: &ExecutorContext,
2640            _pid: &str,
2641        ) -> common_meta::error::Result<ProcedureStateResponse> {
2642            common_meta::error::UnsupportedSnafu {
2643                operation: "query_procedure_state",
2644            }
2645            .fail()
2646        }
2647
2648        async fn list_procedures(
2649            &self,
2650            _ctx: &ExecutorContext,
2651        ) -> common_meta::error::Result<ProcedureDetailResponse> {
2652            common_meta::error::UnsupportedSnafu {
2653                operation: "list_procedures",
2654            }
2655            .fail()
2656        }
2657    }
2658
2659    /// A test [`ProcedureExecutor`] that completes create/drop DDL tasks against the
2660    /// in-memory catalog, mimicking what the meta DDL procedures do in production.
2661    /// This allows happy-path DDL requests (create/drop table/view) to be exercised
2662    /// end to end through the gRPC ingress.
2663    struct MockProcedureExecutor {
2664        catalog_manager: Arc<catalog::memory::MemoryCatalogManager>,
2665        next_table_id: std::sync::atomic::AtomicU32,
2666        submitted: std::sync::Mutex<Vec<DdlTask>>,
2667    }
2668
2669    impl MockProcedureExecutor {
2670        fn new(catalog_manager: Arc<catalog::memory::MemoryCatalogManager>) -> Self {
2671            Self {
2672                catalog_manager,
2673                next_table_id: std::sync::atomic::AtomicU32::new(1026),
2674                submitted: std::sync::Mutex::new(Vec::new()),
2675            }
2676        }
2677    }
2678
2679    #[async_trait::async_trait]
2680    impl ProcedureExecutor for MockProcedureExecutor {
2681        async fn submit_ddl_task(
2682            &self,
2683            _ctx: ExecutorContext,
2684            request: SubmitDdlTaskRequest,
2685        ) -> common_meta::error::Result<SubmitDdlTaskResponse> {
2686            self.submitted.lock().unwrap().push(request.task.clone());
2687            match request.task {
2688                DdlTask::CreateTable(task) => {
2689                    let table_id = self
2690                        .next_table_id
2691                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2692                    let mut table_info = task.table_info;
2693                    table_info.ident.table_id = table_id;
2694                    self.catalog_manager
2695                        .register_table_sync(catalog::RegisterTableRequest {
2696                            catalog: table_info.catalog_name.clone(),
2697                            schema: table_info.schema_name.clone(),
2698                            table_name: table_info.name.clone(),
2699                            table_id,
2700                            table: table::dist_table::DistTable::table(Arc::new(table_info)),
2701                        })
2702                        .map_err(BoxedError::new)
2703                        .context(common_meta::error::ExternalSnafu)?;
2704                    Ok(SubmitDdlTaskResponse {
2705                        key: Vec::new(),
2706                        table_ids: vec![table_id],
2707                    })
2708                }
2709                DdlTask::CreateView(task) => {
2710                    let view_id = self
2711                        .next_table_id
2712                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2713                    let mut view_info = task.view_info;
2714                    view_info.ident.table_id = view_id;
2715                    self.catalog_manager
2716                        .register_table_sync(catalog::RegisterTableRequest {
2717                            catalog: task.create_view.catalog_name.clone(),
2718                            schema: task.create_view.schema_name.clone(),
2719                            table_name: task.create_view.view_name.clone(),
2720                            table_id: view_id,
2721                            table: table::dist_table::DistTable::table(Arc::new(view_info)),
2722                        })
2723                        .map_err(BoxedError::new)
2724                        .context(common_meta::error::ExternalSnafu)?;
2725                    Ok(SubmitDdlTaskResponse {
2726                        key: Vec::new(),
2727                        table_ids: vec![view_id],
2728                    })
2729                }
2730                DdlTask::DropView(task) => {
2731                    self.catalog_manager
2732                        .deregister_table_sync(catalog::DeregisterTableRequest {
2733                            catalog: task.catalog.clone(),
2734                            schema: task.schema.clone(),
2735                            table_name: task.view.clone(),
2736                        })
2737                        .map_err(BoxedError::new)
2738                        .context(common_meta::error::ExternalSnafu)?;
2739                    Ok(SubmitDdlTaskResponse::default())
2740                }
2741                other => common_meta::error::UnsupportedSnafu {
2742                    operation: format!("mock submit_ddl_task: {other:?}"),
2743                }
2744                .fail(),
2745            }
2746        }
2747
2748        async fn migrate_region(
2749            &self,
2750            _ctx: &ExecutorContext,
2751            _request: MigrateRegionRequest,
2752        ) -> common_meta::error::Result<MigrateRegionResponse> {
2753            common_meta::error::UnsupportedSnafu {
2754                operation: "migrate_region",
2755            }
2756            .fail()
2757        }
2758
2759        async fn reconcile(
2760            &self,
2761            _ctx: &ExecutorContext,
2762            _request: ReconcileRequest,
2763        ) -> common_meta::error::Result<ReconcileResponse> {
2764            common_meta::error::UnsupportedSnafu {
2765                operation: "reconcile",
2766            }
2767            .fail()
2768        }
2769
2770        async fn query_procedure_state(
2771            &self,
2772            _ctx: &ExecutorContext,
2773            _pid: &str,
2774        ) -> common_meta::error::Result<ProcedureStateResponse> {
2775            common_meta::error::UnsupportedSnafu {
2776                operation: "query_procedure_state",
2777            }
2778            .fail()
2779        }
2780
2781        async fn list_procedures(
2782            &self,
2783            _ctx: &ExecutorContext,
2784        ) -> common_meta::error::Result<ProcedureDetailResponse> {
2785            common_meta::error::UnsupportedSnafu {
2786                operation: "list_procedures",
2787            }
2788            .fail()
2789        }
2790    }
2791
2792    fn test_cache_registry(
2793        kv_backend: common_meta::kv_backend::KvBackendRef,
2794    ) -> TestResult<common_meta::cache::LayeredCacheRegistryRef> {
2795        Ok(Arc::new(
2796            cache::with_default_composite_cache_registry(
2797                LayeredCacheRegistryBuilder::default()
2798                    .add_cache_registry(cache::build_fundamental_cache_registry(kv_backend)),
2799            )
2800            .context(BuildCacheRegistrySnafu)?
2801            .build(),
2802        ))
2803    }
2804
2805    fn test_table_info(table_id: u32, table_name: &str) -> TestResult<TableInfo> {
2806        let schema = Arc::new(GtSchema::new(vec![
2807            ColumnSchema::new("id", ConcreteDataType::int32_datatype(), false),
2808            ColumnSchema::new(
2809                "ts",
2810                ConcreteDataType::timestamp_millisecond_datatype(),
2811                false,
2812            )
2813            .with_time_index(true),
2814        ]));
2815        let table_meta = TableMetaBuilder::empty()
2816            .schema(schema)
2817            .primary_key_indices(vec![0])
2818            .value_indices(vec![1])
2819            .next_column_id(1024)
2820            .build()
2821            .with_context(|_| BuildTableMetaSnafu {
2822                table_name: table_name.to_string(),
2823            })?;
2824
2825        TableInfoBuilder::new(table_name, table_meta)
2826            .table_id(table_id)
2827            .build()
2828            .with_context(|_| BuildTableInfoSnafu {
2829                table_name: table_name.to_string(),
2830            })
2831    }
2832
2833    fn test_table(table_id: u32, table_name: &str) -> TestResult<table::TableRef> {
2834        let table_info = test_table_info(table_id, table_name)?;
2835        Ok(EmptyTable::from_table_info(&table_info))
2836    }
2837
2838    fn test_physical_table(table_id: u32, table_name: &str) -> TestResult<table::TableRef> {
2839        let mut table_info = test_table_info(table_id, table_name)?;
2840        table_info
2841            .meta
2842            .options
2843            .extra_options
2844            .insert(PHYSICAL_TABLE_METADATA_KEY.to_string(), String::new());
2845        Ok(EmptyTable::from_table_info(&table_info))
2846    }
2847
2848    fn test_logical_table(table_id: u32, table_name: &str) -> TestResult<table::TableRef> {
2849        let mut table_info = test_table_info(table_id, table_name)?;
2850        table_info.meta.engine = METRIC_ENGINE_NAME.to_string();
2851        table_info.meta.options.extra_options.insert(
2852            LOGICAL_TABLE_METADATA_KEY.to_string(),
2853            "physical_metric".to_string(),
2854        );
2855        Ok(EmptyTable::from_table_info(&table_info))
2856    }
2857
2858    fn test_metric_names_table() -> TableRef {
2859        let schema = Arc::new(GtSchema::new(vec![
2860            ColumnSchema::new("table_catalog", ConcreteDataType::string_datatype(), false),
2861            ColumnSchema::new("table_schema", ConcreteDataType::string_datatype(), false),
2862            ColumnSchema::new("table_name", ConcreteDataType::string_datatype(), false),
2863            ColumnSchema::new("engine", ConcreteDataType::string_datatype(), false),
2864            ColumnSchema::new("create_options", ConcreteDataType::string_datatype(), false),
2865        ]));
2866        let columns: Vec<VectorRef> = vec![
2867            Arc::new(StringVector::from(vec!["greptime", "greptime"])),
2868            Arc::new(StringVector::from(vec!["public", "public"])),
2869            Arc::new(StringVector::from(vec!["denied", "target"])),
2870            Arc::new(StringVector::from(vec!["metric", "metric"])),
2871            Arc::new(StringVector::from(vec![
2872                "on_physical_table=physical_metric",
2873                "on_physical_table=physical_metric",
2874            ])),
2875        ];
2876        let record_batch = RecordBatch::new(schema, columns).unwrap();
2877        MemTable::new_with_catalog(
2878            "tables",
2879            record_batch,
2880            2048,
2881            "greptime".to_string(),
2882            "information_schema".to_string(),
2883        )
2884    }
2885
2886    fn test_pipeline_table() -> TableRef {
2887        let schema = Arc::new(GtSchema::new(vec![
2888            ColumnSchema::new("name", ConcreteDataType::string_datatype(), false),
2889            ColumnSchema::new("schema", ConcreteDataType::string_datatype(), false),
2890            ColumnSchema::new("content_type", ConcreteDataType::string_datatype(), false),
2891            ColumnSchema::new("pipeline", ConcreteDataType::string_datatype(), false),
2892            ColumnSchema::new(
2893                "created_at",
2894                ConcreteDataType::timestamp_nanosecond_datatype(),
2895                false,
2896            )
2897            .with_time_index(true),
2898        ]));
2899        let columns: Vec<VectorRef> = vec![
2900            Arc::new(StringVector::from(vec!["pipeline"])),
2901            Arc::new(StringVector::from(vec!["public"])),
2902            Arc::new(StringVector::from(vec!["application/yaml"])),
2903            Arc::new(StringVector::from(vec![
2904                "transform:\n- field: ts\n  type: timestamp, ns\n  index: time\n",
2905            ])),
2906            Arc::new(TimestampNanosecondVector::from_values([1])),
2907        ];
2908        let record_batch = RecordBatch::new(schema, columns).unwrap();
2909        MemTable::new_with_catalog(
2910            "pipelines",
2911            record_batch,
2912            2049,
2913            "greptime".to_string(),
2914            DEFAULT_PRIVATE_SCHEMA_NAME.to_string(),
2915        )
2916    }
2917
2918    struct CancellableExportSource {
2919        schema: GtSchemaRef,
2920        channels: std::sync::Mutex<Option<(oneshot::Sender<()>, oneshot::Receiver<()>)>>,
2921    }
2922
2923    impl DataSource for CancellableExportSource {
2924        fn get_stream(
2925            &self,
2926            _request: ScanRequest,
2927        ) -> std::result::Result<SendableRecordBatchStream, BoxedError> {
2928            let (started, release) = self.channels.lock().unwrap().take().unwrap();
2929            let schema = self.schema.clone();
2930            let stream = futures::stream::once(async move {
2931                started.send(()).unwrap();
2932                release.await.unwrap();
2933                Ok(RecordBatch::new_empty(schema))
2934            });
2935            Ok(Box::pin(RecordBatchStreamWrapper::new(
2936                self.schema.clone(),
2937                Box::pin(stream),
2938            )))
2939        }
2940    }
2941
2942    #[tokio::test]
2943    async fn test_metric_export_http_timeout_cancels_ordinary_source() {
2944        let destination = common_test_util::temp_dir::create_temp_dir("metric_export_timeout");
2945        let (started_tx, started_rx) = oneshot::channel();
2946        let (mut release_tx, release_rx) = oneshot::channel();
2947        let info = test_table_info(1024, "source").unwrap();
2948        let source = Arc::new(Table::new(
2949            Arc::new(info.clone()),
2950            FilterPushDownType::Unsupported,
2951            Arc::new(CancellableExportSource {
2952                schema: info.meta.schema.clone(),
2953                channels: std::sync::Mutex::new(Some((started_tx, release_rx))),
2954            }),
2955        ));
2956        let catalog = catalog::memory::MemoryCatalogManager::new_with_table(source);
2957        let kv = Arc::new(MemoryKvBackend::new());
2958        let instance = FrontendBuilder::new(
2959            FrontendOptions::default(),
2960            kv.clone(),
2961            test_cache_registry(kv).unwrap(),
2962            catalog,
2963            Arc::new(client::client_manager::NodeClients::default()),
2964            Arc::new(NoopProcedureExecutor),
2965            Arc::new(ProcessManager::new("export-timeout".into(), None)),
2966        )
2967        .with_local_file_access(
2968            common_datasource::object_store::LocalFileAccess::sandboxed(destination.path())
2969                .unwrap(),
2970        )
2971        .try_build()
2972        .await
2973        .unwrap();
2974        let server = servers::http::HttpServerBuilder::new(servers::http::HttpOptions {
2975            timeout: Duration::from_secs(2),
2976            ..Default::default()
2977        })
2978        .with_sql_handler(Arc::new(instance))
2979        .build();
2980        let app = server.build(server.make_app()).unwrap();
2981        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
2982        let addr = listener.local_addr().unwrap();
2983        let server_task = tokio::spawn(async move {
2984            axum::serve(listener, app).await.unwrap();
2985        });
2986        let uri = reqwest::Url::from_directory_path(destination.path()).unwrap();
2987        let sql = format!(
2988            "COPY DATABASE greptime.public TO '{uri}' WITH (experimental_metric_export='true')"
2989        );
2990        let request = tokio::spawn(async move {
2991            reqwest::Client::builder()
2992                .no_proxy()
2993                .build()
2994                .unwrap()
2995                .post(format!("http://{addr}/v1/sql"))
2996                .form(&[("sql", sql)])
2997                .send()
2998                .await
2999                .unwrap()
3000        });
3001        tokio::time::timeout(Duration::from_secs(5), started_rx)
3002            .await
3003            .unwrap()
3004            .unwrap();
3005        let response = request.await.unwrap();
3006        assert_eq!(response.status(), reqwest::StatusCode::REQUEST_TIMEOUT);
3007        tokio::time::timeout(Duration::from_secs(5), release_tx.closed())
3008            .await
3009            .unwrap();
3010        assert!(!destination.path().join("source.parquet").exists());
3011        server_task.abort();
3012    }
3013
3014    fn pending_table(
3015        table_id: u32,
3016        table_name: &str,
3017        polled_tx: oneshot::Sender<()>,
3018    ) -> TestResult<table::TableRef> {
3019        let table_info = test_table_info(table_id, table_name)?;
3020        let data_source = Arc::new(PendingDataSource {
3021            schema: table_info.meta.schema.clone(),
3022            polled_tx: std::sync::Mutex::new(Some(polled_tx)),
3023        });
3024
3025        Ok(Arc::new(Table::new(
3026            Arc::new(table_info),
3027            FilterPushDownType::Unsupported,
3028            data_source,
3029        )))
3030    }
3031
3032    async fn test_instance_with_tables(
3033        source_table: TableRef,
3034        target_table: TableRef,
3035    ) -> TestResult<Instance> {
3036        test_instance_with_plugins(source_table, target_table, Plugins::new()).await
3037    }
3038
3039    async fn test_instance_with_insert_select_interceptor(
3040        interceptor: SqlQueryInterceptorRef<Error>,
3041    ) -> TestResult<Instance> {
3042        let plugins = Plugins::new();
3043        plugins.insert::<SqlQueryInterceptorRef<Error>>(interceptor);
3044
3045        test_instance_with_plugins(
3046            test_table(1024, "source")?,
3047            test_table(1025, "target")?,
3048            plugins,
3049        )
3050        .await
3051    }
3052
3053    async fn test_instance_with_plugins(
3054        source_table: TableRef,
3055        target_table: TableRef,
3056        plugins: Plugins,
3057    ) -> TestResult<Instance> {
3058        test_instance_with_plugins_and_metric_names(source_table, target_table, plugins, None).await
3059    }
3060
3061    async fn test_instance_with_plugins_and_metric_names(
3062        source_table: TableRef,
3063        target_table: TableRef,
3064        plugins: Plugins,
3065        metric_names_table: Option<TableRef>,
3066    ) -> TestResult<Instance> {
3067        let catalog_manager = catalog::memory::MemoryCatalogManager::new_with_table(source_table);
3068        test_instance_with_catalog_manager(
3069            catalog_manager,
3070            target_table,
3071            plugins,
3072            metric_names_table,
3073            Arc::new(NoopProcedureExecutor),
3074        )
3075        .await
3076    }
3077
3078    /// Builds a test frontend `Instance` over the given (already source-registered)
3079    /// catalog manager, completing DDL tasks through `procedure_executor`.
3080    async fn test_instance_with_catalog_manager(
3081        catalog_manager: Arc<catalog::memory::MemoryCatalogManager>,
3082        target_table: TableRef,
3083        plugins: Plugins,
3084        metric_names_table: Option<TableRef>,
3085        procedure_executor: ProcedureExecutorRef,
3086    ) -> TestResult<Instance> {
3087        let kv_backend = Arc::new(MemoryKvBackend::new());
3088        let process_manager = Arc::new(ProcessManager::new("test-frontend".to_string(), None));
3089        let target_table_name = "target";
3090        catalog_manager
3091            .register_table_sync(catalog::RegisterTableRequest {
3092                catalog: "greptime".to_string(),
3093                schema: "public".to_string(),
3094                table_name: target_table_name.to_string(),
3095                table_id: 1025,
3096                table: target_table,
3097            })
3098            .with_context(|_| RegisterTableSnafu {
3099                table_name: target_table_name.to_string(),
3100            })?;
3101        if let Some(table) = metric_names_table {
3102            catalog_manager
3103                .deregister_table_sync(catalog::DeregisterTableRequest {
3104                    catalog: "greptime".to_string(),
3105                    schema: "information_schema".to_string(),
3106                    table_name: "tables".to_string(),
3107                })
3108                .unwrap();
3109            catalog_manager
3110                .register_table_sync(catalog::RegisterTableRequest {
3111                    catalog: "greptime".to_string(),
3112                    schema: "information_schema".to_string(),
3113                    table_name: "tables".to_string(),
3114                    table_id: 2048,
3115                    table,
3116                })
3117                .unwrap();
3118        }
3119        catalog_manager.register_process_list_table(process_manager.clone());
3120
3121        let cache_registry = test_cache_registry(kv_backend.clone())?;
3122
3123        FrontendBuilder::new(
3124            FrontendOptions::default(),
3125            kv_backend,
3126            cache_registry,
3127            catalog_manager,
3128            Arc::new(client::client_manager::NodeClients::default()),
3129            procedure_executor,
3130            process_manager,
3131        )
3132        .with_plugin(plugins)
3133        .try_build()
3134        .await
3135        .context(BuildFrontendSnafu)
3136    }
3137
3138    async fn execute_one_sql(
3139        instance: &Instance,
3140        sql: &str,
3141        query_ctx: QueryContextRef,
3142    ) -> TestResult<Output> {
3143        let mut results = instance.do_query_inner(sql, query_ctx).await;
3144        ensure!(
3145            results.len() == 1,
3146            UnexpectedOutputCountSnafu {
3147                sql: sql.to_string(),
3148                actual: results.len(),
3149            }
3150        );
3151        results.remove(0).with_context(|_| ExecuteSqlSnafu {
3152            sql: sql.to_string(),
3153        })
3154    }
3155
3156    fn assert_permission_denied<T>(result: servers::error::Result<T>) {
3157        let err = match result {
3158            Ok(_) => panic!("request should be rejected"),
3159            Err(err) => err,
3160        };
3161        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3162    }
3163
3164    fn assert_action_checked(
3165        checker: &RejectEndpointPermissionChecker,
3166        action: PermissionAction,
3167        targets: Option<PermissionTableTargets>,
3168    ) {
3169        assert_eq!(CheckedAction { action, targets }, checker.take_check());
3170    }
3171
3172    #[tokio::test]
3173    async fn database_export_authorizes_all_tables_before_preparation() -> TestResult<()> {
3174        struct ExportAcl(std::sync::Mutex<Vec<PermissionTableTargets>>);
3175        impl PermissionChecker for ExportAcl {
3176            fn check_permission(
3177                &self,
3178                _: UserInfoRef,
3179                req: PermissionReq,
3180            ) -> auth::error::Result<PermissionResp> {
3181                assert!(matches!(
3182                    req,
3183                    PermissionReq::SqlStatement(Statement::Copy(
3184                        sql::statements::copy::Copy::CopyDatabase(
3185                            sql::statements::copy::CopyDatabase::To(_)
3186                        )
3187                    ))
3188                ));
3189                Ok(PermissionResp::Allow)
3190            }
3191            fn check_permission_with_table_targets(
3192                &self,
3193                _: UserInfoRef,
3194                req: PermissionReq,
3195                targets: PermissionTableTargets,
3196            ) -> auth::error::Result<PermissionResp> {
3197                self.0.lock().unwrap().push(targets.clone());
3198                let PermissionTableTargets::Resolved(tables) = targets else {
3199                    panic!("unresolved export")
3200                };
3201                if req.is_readonly() {
3202                    // The first table is readable; the later table has only write access.
3203                    Ok(if tables.iter().any(|t| t.table == "target") {
3204                        PermissionResp::Reject
3205                    } else {
3206                        PermissionResp::Allow
3207                    })
3208                } else {
3209                    self.check_permission(QueryContext::arc().current_user(), req)
3210                }
3211            }
3212        }
3213        let checker = Arc::new(ExportAcl(Default::default()));
3214        let plugins = Plugins::new();
3215        plugins.insert::<PermissionCheckerRef>(checker.clone());
3216        let instance = test_instance_with_plugins(
3217            test_logical_table(1024, "source")?,
3218            test_table(1025, "target")?,
3219            plugins,
3220        )
3221        .await?;
3222        let req = table::requests::CopyDatabaseRequest {
3223            catalog_name: "greptime".into(),
3224            schema_name: "public".into(),
3225            location: "invalid-destination".into(),
3226            with: Default::default(),
3227            connection: Default::default(),
3228            time_range: None,
3229        };
3230        let result = instance
3231            .export_database_for_test(
3232                req.clone(),
3233                None,
3234                &tokio_util::sync::CancellationToken::new(),
3235                QueryContext::arc(),
3236            )
3237            .await;
3238        assert!(matches!(result, Err(Error::Permission { .. })));
3239        let expected = PermissionTableTargets::resolved(vec![
3240            PermissionTableTarget::new("greptime", "public", "source"),
3241            PermissionTableTarget::new("greptime", "public", "target"),
3242        ]);
3243        assert_eq!(*checker.0.lock().unwrap(), vec![expected.clone(), expected]);
3244        // An empty selection must still check the operation privilege.
3245        instance
3246            .plugins
3247            .map_mut::<PermissionCheckerRef, _, _>(|checker| {
3248                *checker.unwrap() = Arc::new(WriteOnlyPermissionChecker)
3249            });
3250        let result = instance
3251            .export_database_for_test(
3252                req,
3253                Some(&[]),
3254                &tokio_util::sync::CancellationToken::new(),
3255                QueryContext::arc(),
3256            )
3257            .await;
3258        assert!(matches!(result, Err(Error::Permission { .. })));
3259        Ok(())
3260    }
3261
3262    #[tokio::test]
3263    async fn test_prom_remote_read_with_custom_timestamp_and_value_columns() -> TestResult<()> {
3264        let schema = Arc::new(GtSchema::new(vec![
3265            ColumnSchema::new(
3266                "custom_ts",
3267                ConcreteDataType::timestamp_millisecond_datatype(),
3268                false,
3269            )
3270            .with_time_index(true),
3271            ColumnSchema::new("custom_value", ConcreteDataType::float64_datatype(), false),
3272        ]));
3273        let recordbatch = RecordBatch::new(
3274            schema,
3275            vec![
3276                Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as VectorRef,
3277                Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as VectorRef,
3278            ],
3279        )
3280        .unwrap();
3281        let instance = test_instance_with_tables(
3282            MemTable::table("custom_metric", recordbatch),
3283            test_table(1025, "target")?,
3284        )
3285        .await?;
3286
3287        let response = PromStoreProtocolHandler::read(
3288            &instance,
3289            ReadRequest {
3290                queries: vec![RemoteQuery {
3291                    start_timestamp_ms: 1500,
3292                    end_timestamp_ms: 2500,
3293                    matchers: vec![LabelMatcher {
3294                        r#type: PromMatcherType::Eq as i32,
3295                        name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3296                        value: "custom_metric".to_string(),
3297                    }],
3298                    ..Default::default()
3299                }],
3300                ..Default::default()
3301            },
3302            test_query_ctx(1),
3303        )
3304        .await
3305        .unwrap();
3306        let body = servers::prom_store::snappy_decompress(&response.body).unwrap();
3307        let response = ReadResponse::decode(body.as_slice()).unwrap();
3308
3309        assert_eq!(1, response.results.len());
3310        assert_eq!(1, response.results[0].timeseries.len());
3311        let timeseries = &response.results[0].timeseries[0];
3312        assert_eq!(
3313            vec![Label {
3314                name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3315                value: "custom_metric".to_string(),
3316            }],
3317            timeseries.labels
3318        );
3319        assert_eq!(
3320            vec![Sample {
3321                value: 2.0,
3322                timestamp: 2000,
3323            }],
3324            timeseries.samples
3325        );
3326
3327        Ok(())
3328    }
3329
3330    #[tokio::test]
3331    async fn test_prom_remote_read_prefers_default_value_column() -> TestResult<()> {
3332        let schema = Arc::new(GtSchema::new(vec![
3333            ColumnSchema::new(
3334                "custom_ts",
3335                ConcreteDataType::timestamp_millisecond_datatype(),
3336                false,
3337            )
3338            .with_time_index(true),
3339            ColumnSchema::new("extra_field", ConcreteDataType::float64_datatype(), false),
3340            ColumnSchema::new(
3341                greptime_value(),
3342                ConcreteDataType::float64_datatype(),
3343                false,
3344            ),
3345        ]));
3346        let recordbatch = RecordBatch::new(
3347            schema,
3348            vec![
3349                Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as VectorRef,
3350                Arc::new(Float64Vector::from_vec(vec![99.0, 99.0, 99.0])) as VectorRef,
3351                Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as VectorRef,
3352            ],
3353        )
3354        .unwrap();
3355        let instance = test_instance_with_tables(
3356            MemTable::table("multi_field_metric", recordbatch),
3357            test_table(1025, "target")?,
3358        )
3359        .await?;
3360
3361        let response = PromStoreProtocolHandler::read(
3362            &instance,
3363            ReadRequest {
3364                queries: vec![RemoteQuery {
3365                    start_timestamp_ms: 1500,
3366                    end_timestamp_ms: 2500,
3367                    matchers: vec![LabelMatcher {
3368                        r#type: PromMatcherType::Eq as i32,
3369                        name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3370                        value: "multi_field_metric".to_string(),
3371                    }],
3372                    ..Default::default()
3373                }],
3374                ..Default::default()
3375            },
3376            test_query_ctx(1),
3377        )
3378        .await
3379        .unwrap();
3380        let body = servers::prom_store::snappy_decompress(&response.body).unwrap();
3381        let response = ReadResponse::decode(body.as_slice()).unwrap();
3382
3383        assert_eq!(1, response.results.len());
3384        assert_eq!(1, response.results[0].timeseries.len());
3385        let timeseries = &response.results[0].timeseries[0];
3386        assert_eq!(
3387            vec![
3388                Label {
3389                    name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3390                    value: "multi_field_metric".to_string(),
3391                },
3392                Label {
3393                    name: "extra_field".to_string(),
3394                    value: "99".to_string(),
3395                },
3396            ],
3397            timeseries.labels
3398        );
3399        assert_eq!(
3400            vec![Sample {
3401                value: 2.0,
3402                timestamp: 2000,
3403            }],
3404            timeseries.samples
3405        );
3406
3407        Ok(())
3408    }
3409
3410    #[tokio::test]
3411    async fn test_prom_remote_read_rejects_ambiguous_value_columns() -> TestResult<()> {
3412        let schema = Arc::new(GtSchema::new(vec![
3413            ColumnSchema::new(
3414                "custom_ts",
3415                ConcreteDataType::timestamp_millisecond_datatype(),
3416                false,
3417            )
3418            .with_time_index(true),
3419            ColumnSchema::new("field_a", ConcreteDataType::float64_datatype(), false),
3420            ColumnSchema::new("field_b", ConcreteDataType::float64_datatype(), false),
3421        ]));
3422        let recordbatch = RecordBatch::new(
3423            schema,
3424            vec![
3425                Arc::new(TimestampMillisecondVector::from_vec(vec![1000])) as VectorRef,
3426                Arc::new(Float64Vector::from_vec(vec![1.0])) as VectorRef,
3427                Arc::new(Float64Vector::from_vec(vec![2.0])) as VectorRef,
3428            ],
3429        )
3430        .unwrap();
3431        let instance = test_instance_with_tables(
3432            MemTable::table("ambiguous_metric", recordbatch),
3433            test_table(1025, "target")?,
3434        )
3435        .await?;
3436
3437        let err = PromStoreProtocolHandler::read(
3438            &instance,
3439            ReadRequest {
3440                queries: vec![RemoteQuery {
3441                    matchers: vec![LabelMatcher {
3442                        r#type: PromMatcherType::Eq as i32,
3443                        name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3444                        value: "ambiguous_metric".to_string(),
3445                    }],
3446                    ..Default::default()
3447                }],
3448                ..Default::default()
3449            },
3450            test_query_ctx(1),
3451        )
3452        .await
3453        .err()
3454        .expect("ambiguous value columns should fail remote read");
3455
3456        assert_eq!(StatusCode::InvalidArguments, err.status_code());
3457        assert!(format!("{err:?}").contains("Ambiguous value column"));
3458
3459        Ok(())
3460    }
3461
3462    #[tokio::test]
3463    async fn test_instance_shared_state() -> TestResult<()> {
3464        let instance =
3465            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
3466                .await?;
3467
3468        let _event_recorder = instance.event_recorder();
3469        assert!(Arc::ptr_eq(
3470            &instance.otlp_metrics_table_legacy_cache,
3471            &instance.clone().otlp_metrics_table_legacy_cache,
3472        ));
3473
3474        Ok(())
3475    }
3476
3477    #[tokio::test]
3478    async fn test_restricted_endpoint_handlers_check_permissions() -> TestResult<()> {
3479        let checker = Arc::new(RejectEndpointPermissionChecker::default());
3480        let plugins = Plugins::new();
3481        plugins.insert::<PermissionCheckerRef>(checker.clone());
3482        let instance = test_instance_with_plugins(
3483            test_table(1024, "denied")?,
3484            test_table(1025, "target")?,
3485            plugins,
3486        )
3487        .await?;
3488        let mut ctx = test_query_ctx(1);
3489        Arc::get_mut(&mut ctx).unwrap().set_extension(
3490            servers::http::jaeger::JAEGER_QUERY_TABLE_NAME_KEY,
3491            "denied".to_string(),
3492        );
3493        let jaeger_targets = Some(PermissionTableTargets::resolved(vec![
3494            PermissionTableTarget::new("greptime", "public", "denied"),
3495        ]));
3496
3497        assert_permission_denied(JaegerQueryHandler::get_services(&instance, ctx.clone()).await);
3498        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
3499        assert_permission_denied(
3500            JaegerQueryHandler::get_operations(&instance, ctx.clone(), "service", None).await,
3501        );
3502        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
3503        assert_permission_denied(
3504            JaegerQueryHandler::get_trace(&instance, ctx.clone(), "trace", None, None, None).await,
3505        );
3506        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
3507        assert_permission_denied(
3508            JaegerQueryHandler::find_traces(
3509                &instance,
3510                ctx.clone(),
3511                servers::http::jaeger::QueryTraceParams {
3512                    service_name: Some("service".to_string()),
3513                    ..Default::default()
3514                },
3515            )
3516            .await,
3517        );
3518        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets);
3519
3520        assert_permission_denied(
3521            PipelineHandler::get_pipeline_str(&instance, "pipeline", None, ctx.clone()).await,
3522        );
3523        assert_action_checked(&checker, PIPELINE_QUERY, None);
3524        assert_permission_denied(
3525            PipelineHandler::insert_pipeline(
3526                &instance,
3527                "pipeline",
3528                "application/yaml",
3529                "",
3530                ctx.clone(),
3531            )
3532            .await,
3533        );
3534        assert_action_checked(&checker, PIPELINE_INSERT, None);
3535        assert_permission_denied(
3536            PipelineHandler::delete_pipeline(&instance, "pipeline", None, ctx.clone()).await,
3537        );
3538        assert_action_checked(&checker, PIPELINE_DELETE, None);
3539        let app = axum::Router::new()
3540            .route(
3541                "/pipelines/_dryrun",
3542                axum::routing::post(servers::http::event::pipeline_dryrun),
3543            )
3544            .with_state(servers::http::event::LogState {
3545                log_handler: Arc::new(instance.clone()),
3546                log_validator: None,
3547                ingest_interceptor: None,
3548            })
3549            .layer(axum::Extension((*ctx).clone()));
3550        let response = app
3551            .oneshot(
3552                axum::http::Request::post("/pipelines/_dryrun")
3553                    .header("content-type", "application/json")
3554                    .body(axum::body::Body::from("{}"))
3555                    .unwrap(),
3556            )
3557            .await
3558            .unwrap();
3559        assert_eq!(axum::http::StatusCode::FORBIDDEN, response.status());
3560        assert_action_checked(&checker, PIPELINE_QUERY, None);
3561
3562        assert_permission_denied(
3563            DashboardHandler::save(&instance, "dashboard", "{}", ctx.clone()).await,
3564        );
3565        assert_action_checked(&checker, DASHBOARD_SAVE, None);
3566        assert_permission_denied(DashboardHandler::list(&instance, ctx.clone()).await);
3567        assert_action_checked(&checker, DASHBOARD_QUERY, None);
3568        assert_permission_denied(
3569            DashboardHandler::delete(&instance, "dashboard", ctx.clone()).await,
3570        );
3571        assert_action_checked(&checker, DASHBOARD_DELETE, None);
3572
3573        Ok(())
3574    }
3575
3576    #[tokio::test]
3577    async fn test_write_only_ingestion_loads_named_pipeline() -> TestResult<()> {
3578        let plugins = Plugins::new();
3579        plugins.insert::<PermissionCheckerRef>(Arc::new(WriteOnlyPermissionChecker));
3580        let instance = test_instance_with_plugins(
3581            test_table(1024, "source")?,
3582            test_table(1025, "target")?,
3583            plugins,
3584        )
3585        .await?;
3586        instance
3587            .catalog_manager()
3588            .as_any()
3589            .downcast_ref::<catalog::memory::MemoryCatalogManager>()
3590            .unwrap()
3591            .register_table_sync(catalog::RegisterTableRequest {
3592                catalog: "greptime".to_string(),
3593                schema: DEFAULT_PRIVATE_SCHEMA_NAME.to_string(),
3594                table_name: "pipelines".to_string(),
3595                table_id: 2049,
3596                table: test_pipeline_table(),
3597            })
3598            .with_context(|_| RegisterTableSnafu {
3599                table_name: "pipelines".to_string(),
3600            })?;
3601        let ctx = test_query_ctx(1);
3602        let handler: PipelineHandlerRef = Arc::new(instance.clone());
3603        let hits = CACHE_HIT.with_label_values(&["pipeline_table"]);
3604        let misses = CACHE_MISS.with_label_values(&["pipeline_table"]);
3605        let before = (hits.get(), misses.get());
3606
3607        handler
3608            .get_pipeline("pipeline", None, ctx.clone())
3609            .await
3610            .unwrap();
3611        assert_eq!((hits.get(), misses.get()), (before.0 + 1, before.1 + 1));
3612        handler
3613            .get_pipeline("pipeline", None, ctx.clone())
3614            .await
3615            .unwrap();
3616        assert_eq!((hits.get(), misses.get()), (before.0 + 3, before.1 + 1));
3617        assert_permission_denied(
3618            PipelineHandler::get_pipeline_str(&instance, "pipeline", None, ctx.clone()).await,
3619        );
3620
3621        let app = axum::Router::new()
3622            .route(
3623                "/pipelines/_dryrun",
3624                axum::routing::post(servers::http::event::pipeline_dryrun),
3625            )
3626            .with_state(servers::http::event::LogState {
3627                log_handler: handler,
3628                log_validator: None,
3629                ingest_interceptor: None,
3630            })
3631            .layer(axum::Extension((*ctx).clone()));
3632        let response = app
3633            .oneshot(
3634                axum::http::Request::post("/pipelines/_dryrun")
3635                    .header("content-type", "application/json")
3636                    .body(axum::body::Body::from("{}"))
3637                    .unwrap(),
3638            )
3639            .await
3640            .unwrap();
3641        assert_eq!(axum::http::StatusCode::FORBIDDEN, response.status());
3642
3643        Ok(())
3644    }
3645
3646    #[tokio::test]
3647    async fn test_write_only_grpc_sql_is_checked_after_parsing() -> TestResult<()> {
3648        let plugins = Plugins::new();
3649        plugins.insert::<PermissionCheckerRef>(Arc::new(WriteOnlyPermissionChecker));
3650        let instance = test_instance_with_plugins(
3651            test_table(1024, "source")?,
3652            test_table(1025, "target")?,
3653            plugins,
3654        )
3655        .await?;
3656
3657        let insert = Request::Query(api::v1::QueryRequest {
3658            query: Some(Query::Sql(
3659                "INSERT INTO target SELECT * FROM source".to_string(),
3660            )),
3661        });
3662        servers::query_handler::grpc::GrpcQueryHandler::do_query(
3663            &instance,
3664            insert,
3665            QueryContext::arc(),
3666        )
3667        .await
3668        .unwrap();
3669
3670        let select = Request::Query(api::v1::QueryRequest {
3671            query: Some(Query::Sql("SELECT * FROM source".to_string())),
3672        });
3673        assert_permission_denied(
3674            servers::query_handler::grpc::GrpcQueryHandler::do_query(
3675                &instance,
3676                select,
3677                QueryContext::arc(),
3678            )
3679            .await,
3680        );
3681
3682        Ok(())
3683    }
3684
3685    #[tokio::test]
3686    async fn test_target_independent_checker_skips_target_resolution() -> TestResult<()> {
3687        let physical_table = "physical_metric";
3688        let checker = Arc::new(TargetIndependentPermissionChecker::default());
3689        let plugins = Plugins::new();
3690        plugins.insert::<PermissionCheckerRef>(checker.clone());
3691        let instance = test_instance_with_plugins(
3692            test_physical_table(1024, physical_table)?,
3693            test_table(1025, "target")?,
3694            plugins,
3695        )
3696        .await?;
3697
3698        let ctx = test_query_ctx(1);
3699        let physical_target = PermissionTableTarget::new("greptime", "public", physical_table);
3700        assert_eq!(
3701            PermissionTableTargets::Resolved(vec![physical_target.clone()]),
3702            instance
3703                .resolve_query_permission_targets(
3704                    PermissionTableTargets::resolved(vec![physical_target]),
3705                    &ctx,
3706                )
3707                .await
3708                .unwrap()
3709        );
3710        assert_eq!(
3711            vec![physical_table.to_string(), "target".to_string()],
3712            PrometheusHandler::filter_metadata_metric_names(
3713                &instance,
3714                vec![physical_table.to_string(), "target".to_string()],
3715                "public",
3716                &ctx,
3717            )
3718            .await
3719            .unwrap()
3720        );
3721        assert_eq!(1, checker.checks.load(atomic::Ordering::Relaxed));
3722
3723        Ok(())
3724    }
3725
3726    #[tokio::test]
3727    async fn test_query_permission_targets_are_deduplicated() -> TestResult<()> {
3728        let plugins = Plugins::new();
3729        plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
3730        let instance = test_instance_with_plugins(
3731            test_table(1024, "source")?,
3732            test_table(1025, "target")?,
3733            plugins,
3734        )
3735        .await?;
3736        let ctx = test_query_ctx(1);
3737        let target = PermissionTableTarget::new("greptime", "public", "target");
3738
3739        assert_eq!(
3740            PermissionTableTargets::Resolved(vec![target.clone()]),
3741            instance
3742                .resolve_query_permission_targets(
3743                    PermissionTableTargets::resolved(vec![target.clone(), target]),
3744                    &ctx,
3745                )
3746                .await
3747                .unwrap()
3748        );
3749
3750        Ok(())
3751    }
3752
3753    #[tokio::test]
3754    async fn test_physical_query_targets_fail_closed() -> TestResult<()> {
3755        let physical_table = "physical_metric";
3756        let plugins = Plugins::new();
3757        plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
3758        let instance = test_instance_with_plugins(
3759            test_physical_table(1024, physical_table)?,
3760            test_table(1025, "target")?,
3761            plugins,
3762        )
3763        .await?;
3764
3765        let ctx = test_query_ctx(1);
3766        let logical_target = PermissionTableTarget::new("greptime", "public", "target");
3767        assert_eq!(
3768            PermissionTableTargets::Resolved(vec![logical_target.clone()]),
3769            instance
3770                .resolve_query_permission_targets(
3771                    PermissionTableTargets::resolved(vec![logical_target.clone()]),
3772                    &ctx,
3773                )
3774                .await
3775                .unwrap()
3776        );
3777        let physical_target = PermissionTableTarget::new("greptime", "public", physical_table);
3778        assert_eq!(
3779            PermissionTableTargets::Unresolved,
3780            instance
3781                .resolve_query_permission_targets(
3782                    PermissionTableTargets::resolved(
3783                        vec![logical_target, physical_target.clone(),]
3784                    ),
3785                    &ctx,
3786                )
3787                .await
3788                .unwrap()
3789        );
3790        assert_eq!(
3791            vec!["target".to_string()],
3792            PrometheusHandler::filter_metadata_metric_names(
3793                &instance,
3794                vec!["target".to_string(), "denied".to_string()],
3795                "public",
3796                &ctx,
3797            )
3798            .await
3799            .unwrap()
3800        );
3801
3802        let query = PromQuery {
3803            query: physical_table.to_string(),
3804            ..Default::default()
3805        };
3806        let err = PrometheusHandler::check_query_target_permission(
3807            &instance,
3808            PermissionTableTargets::resolved(vec![physical_target]),
3809            &ctx,
3810        )
3811        .await
3812        .unwrap_err();
3813        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3814        let err = PrometheusHandler::check_query_permission(
3815            &instance,
3816            std::slice::from_ref(&query),
3817            &ctx,
3818        )
3819        .await
3820        .unwrap_err();
3821        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3822        let err = PrometheusHandler::do_query(&instance, &query, ctx.clone())
3823            .await
3824            .unwrap_err();
3825        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3826
3827        for sql in [
3828            "SELECT * FROM physical_metric",
3829            "TQL EVAL (0, 10, '5s') physical_metric",
3830            "INSERT INTO target SELECT * FROM physical_metric",
3831        ] {
3832            let mut results = instance.do_query_inner(sql, ctx.clone()).await;
3833            assert_eq!(1, results.len(), "{sql}");
3834            let err = results.remove(0).unwrap_err();
3835            assert_eq!(StatusCode::PermissionDenied, err.status_code(), "{sql}");
3836        }
3837        let err = LogQueryHandler::query(
3838            &instance,
3839            LogQuery {
3840                table: TableName::new("greptime", "public", physical_table),
3841                ..Default::default()
3842            },
3843            ctx.clone(),
3844        )
3845        .await
3846        .unwrap_err();
3847        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3848        let err = instance
3849            .do_describe_inner(parse_one_sql("SELECT * FROM physical_metric"), ctx.clone())
3850            .await
3851            .unwrap_err();
3852        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3853
3854        let request = ReadRequest {
3855            queries: vec![RemoteQuery {
3856                matchers: vec![LabelMatcher {
3857                    r#type: PromMatcherType::Eq as i32,
3858                    name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3859                    value: physical_table.to_string(),
3860                }],
3861                ..Default::default()
3862            }],
3863            ..Default::default()
3864        };
3865        let Err(err) = PromStoreProtocolHandler::read(&instance, request, ctx.clone()).await else {
3866            panic!("physical remote-read target must be rejected");
3867        };
3868        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3869
3870        let err = PrometheusHandler::query_label_values(
3871            &instance,
3872            physical_table.to_string(),
3873            "host".to_string(),
3874            vec![],
3875            SystemTime::UNIX_EPOCH,
3876            SystemTime::UNIX_EPOCH,
3877            &ctx,
3878        )
3879        .await
3880        .unwrap_err();
3881        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3882
3883        Ok(())
3884    }
3885
3886    #[tokio::test]
3887    async fn test_non_exact_query_discovery_keeps_denied_targets_for_batch_check() -> TestResult<()>
3888    {
3889        let plugins = Plugins::new();
3890        plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
3891        let instance = test_instance_with_plugins_and_metric_names(
3892            test_logical_table(1024, "denied")?,
3893            test_logical_table(1025, "target")?,
3894            plugins,
3895            Some(test_metric_names_table()),
3896        )
3897        .await?;
3898        let ctx = test_query_ctx(1);
3899
3900        let mut metric_names = PrometheusHandler::query_metric_names(
3901            &instance,
3902            vec![Matcher::new(
3903                promql_parser::label::MatchOp::NotEqual,
3904                "__name__",
3905                "",
3906            )],
3907            "public",
3908            &ctx,
3909        )
3910        .await
3911        .unwrap();
3912        metric_names.sort_unstable();
3913        assert_eq!(
3914            vec!["denied".to_string(), "target".to_string()],
3915            metric_names
3916        );
3917
3918        let queries = metric_names
3919            .into_iter()
3920            .map(|query| PromQuery {
3921                query,
3922                ..Default::default()
3923            })
3924            .collect::<Vec<_>>();
3925        let err = PrometheusHandler::check_query_permission(&instance, &queries, &ctx)
3926            .await
3927            .unwrap_err();
3928        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3929
3930        Ok(())
3931    }
3932
3933    #[test]
3934    fn test_fast_legacy_check_is_read_only() {
3935        let hits = CACHE_HIT.with_label_values(&["otlp_metrics_legacy"]);
3936        let misses = CACHE_MISS.with_label_values(&["otlp_metrics_legacy"]);
3937        let before = (hits.get(), misses.get());
3938        let cache = DashMap::new();
3939        cache.insert("metric1".to_string(), true);
3940
3941        let names = vec!["metric1".to_string(), "metric2".to_string()];
3942        assert_eq!(Some(true), fast_legacy_check(&cache, &names).unwrap());
3943        assert_eq!((hits.get(), misses.get()), (before.0 + 1, before.1 + 1));
3944        assert!(!cache.contains_key("metric2"));
3945
3946        cache_legacy_mode(&cache, &names, true).unwrap();
3947        assert!(*cache.get("metric2").unwrap().value());
3948        assert!(cache_legacy_mode(&cache, &names, false).is_err());
3949        assert!(*cache.get("metric2").unwrap().value());
3950
3951        let cache_incompatible = DashMap::new();
3952        cache_incompatible.insert("metric1".to_string(), true);
3953        cache_incompatible.insert("metric2".to_string(), false);
3954        assert!(fast_legacy_check(&cache_incompatible, &names).is_err());
3955        assert_eq!((hits.get(), misses.get()), (before.0 + 3, before.1 + 1));
3956        assert_eq!(None, fast_legacy_check(&cache, &[]).unwrap());
3957        assert_eq!((hits.get(), misses.get()), (before.0 + 3, before.1 + 1));
3958    }
3959
3960    #[test]
3961    fn test_should_track_statement_process() {
3962        assert!(should_track_statement_process(&parse_one_sql(
3963            "SELECT * FROM demo"
3964        )));
3965        assert!(should_track_statement_process(&parse_one_sql(
3966            "INSERT INTO demo SELECT * FROM source"
3967        )));
3968        assert!(!should_track_statement_process(&parse_one_sql(
3969            "INSERT INTO demo VALUES (1)"
3970        )));
3971        assert!(!should_track_statement_process(&parse_one_sql(
3972            "INSERT INTO demo VALUES (now())"
3973        )));
3974    }
3975
3976    #[test]
3977    fn test_should_track_plan_process() {
3978        let select_stmt = parse_one_sql("SELECT * FROM demo");
3979        let insert_select_stmt = parse_one_sql("INSERT INTO demo SELECT * FROM source");
3980        let insert_values_stmt = parse_one_sql("INSERT INTO demo VALUES (now())");
3981
3982        let empty_plan = LogicalPlanBuilder::empty(false).build().unwrap();
3983        assert!(should_track_plan_process(Some(&select_stmt), &empty_plan));
3984        assert!(should_track_plan_process(
3985            Some(&insert_select_stmt),
3986            &insert_dml_plan()
3987        ));
3988        assert!(!should_track_plan_process(
3989            Some(&insert_values_stmt),
3990            &insert_dml_plan()
3991        ));
3992        assert!(!should_track_plan_process(None, &insert_dml_plan()));
3993    }
3994
3995    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3996    async fn test_insert_select_is_visible_in_show_processlist() -> TestResult<()> {
3997        let insert_sql = "INSERT INTO target SELECT * FROM source";
3998        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
3999        let (finish_tx, finish_rx) = oneshot::channel();
4000        let interceptor = Arc::new(BlockingInsertSelectInterceptor::new(started_tx, finish_rx));
4001        let instance = Arc::new(test_instance_with_insert_select_interceptor(interceptor).await?);
4002
4003        let insert_task = tokio::spawn({
4004            let instance = instance.clone();
4005            async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
4006        });
4007
4008        tokio::time::timeout(Duration::from_secs(5), started_rx.recv())
4009            .await
4010            .context(InsertStartTimeoutSnafu)?
4011            .context(InsertStartChannelClosedSnafu)?;
4012
4013        let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
4014        let process_list = output.data.pretty_print().await;
4015        assert!(
4016            process_list.contains(insert_sql),
4017            "process list did not contain running insert:\n{process_list}"
4018        );
4019
4020        finish_tx
4021            .send(())
4022            .map_err(|_| ReleaseBlockedInsertSnafu.build())?;
4023        insert_task.await.context(InsertTaskPanicSnafu)??;
4024
4025        Ok(())
4026    }
4027
4028    #[tokio::test]
4029    async fn test_show_processlist_catalog_scope() -> TestResult<()> {
4030        let instance =
4031            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
4032                .await?;
4033        let _current_catalog = instance.process_manager().register_query(
4034            "greptime".to_string(),
4035            vec!["public".to_string()],
4036            "current_catalog_query".to_string(),
4037            String::new(),
4038            None,
4039            None,
4040        );
4041        let _other_catalog = instance.process_manager().register_query(
4042            "other".to_string(),
4043            vec!["public".to_string()],
4044            "other_catalog_query".to_string(),
4045            String::new(),
4046            None,
4047            None,
4048        );
4049
4050        for sql in ["SHOW PROCESSLIST", "SHOW FULL PROCESSLIST"] {
4051            let output = execute_one_sql(&instance, sql, test_query_ctx(43)).await?;
4052            let process_list = output.data.pretty_print().await;
4053            assert!(
4054                process_list.contains("current_catalog_query"),
4055                "{process_list}"
4056            );
4057            assert!(
4058                !process_list.contains("other_catalog_query"),
4059                "{process_list}"
4060            );
4061
4062            let admin_ctx = test_query_ctx(44);
4063            admin_ctx.set_current_user(Arc::new(AdminUserInfo));
4064            let output = execute_one_sql(&instance, sql, admin_ctx).await?;
4065            let process_list = output.data.pretty_print().await;
4066            assert!(
4067                process_list.contains("current_catalog_query"),
4068                "{process_list}"
4069            );
4070            assert!(
4071                process_list.contains("other_catalog_query"),
4072                "{process_list}"
4073            );
4074        }
4075
4076        Ok(())
4077    }
4078
4079    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4080    async fn test_kill_query_cancels_insert_select() -> TestResult<()> {
4081        assert_kill_cancels_insert_select("KILL QUERY 4242").await
4082    }
4083
4084    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4085    async fn test_kill_process_id_cancels_insert_select() -> TestResult<()> {
4086        assert_kill_cancels_insert_select("KILL 'test-frontend/4242'").await
4087    }
4088
4089    async fn assert_kill_cancels_insert_select(kill_sql: &str) -> TestResult<()> {
4090        let insert_sql = "INSERT INTO target SELECT * FROM source";
4091        let (source_polled_tx, source_polled_rx) = oneshot::channel();
4092        let instance = Arc::new(
4093            test_instance_with_tables(
4094                pending_table(1024, "source", source_polled_tx)?,
4095                test_table(1025, "target")?,
4096            )
4097            .await?,
4098        );
4099
4100        let insert_task = tokio::spawn({
4101            let instance = instance.clone();
4102            async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
4103        });
4104
4105        tokio::time::timeout(Duration::from_secs(5), source_polled_rx)
4106            .await
4107            .context(SourcePollTimeoutSnafu)?
4108            .context(SourcePollChannelClosedSnafu)?;
4109
4110        let output = execute_one_sql(&instance, kill_sql, test_query_ctx(43)).await?;
4111        assert!(matches!(output.data, OutputData::AffectedRows(1)));
4112
4113        let insert_result = tokio::time::timeout(Duration::from_secs(5), insert_task)
4114            .await
4115            .context(InsertTaskTimeoutSnafu)?
4116            .context(InsertTaskPanicSnafu)?;
4117        let err = match insert_result {
4118            Ok(_) => return InsertSelectNotCancelledSnafu.fail(),
4119            Err(TestError::ExecuteSql { source, .. }) => source,
4120            Err(err) => return Err(err),
4121        };
4122        assert_eq!(StatusCode::Cancelled, err.status_code());
4123
4124        let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
4125        let process_list = output.data.pretty_print().await;
4126        assert!(
4127            !process_list.contains(insert_sql),
4128            "process list still contains killed insert:\n{process_list}"
4129        );
4130
4131        Ok(())
4132    }
4133
4134    fn insert_dml_plan() -> LogicalPlan {
4135        let schema = SchemaRef::new(Schema::new(vec![Field::new(
4136            "value",
4137            DataType::Int64,
4138            true,
4139        )]));
4140        let target = Arc::new(LogicalTableSource::new(schema));
4141        let input = LogicalPlanBuilder::empty(false).build().unwrap();
4142
4143        LogicalPlanBuilder::insert_into(input, "demo", target, InsertOp::Append)
4144            .unwrap()
4145            .build()
4146            .unwrap()
4147    }
4148
4149    #[test]
4150    fn test_exec_validation() {
4151        let query_ctx = QueryContext::arc();
4152        let plugins: Plugins = Plugins::new();
4153        plugins.insert(QueryOptions {
4154            disallow_cross_catalog_query: true,
4155        });
4156
4157        let sql = r#"
4158        SELECT * FROM demo;
4159        EXPLAIN SELECT * FROM demo;
4160        CREATE DATABASE test_database;
4161        SHOW DATABASES;
4162        "#;
4163        let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
4164        assert_eq!(stmts.len(), 4);
4165        for stmt in stmts {
4166            let re = check_permission(plugins.clone(), &stmt, &query_ctx);
4167            re.unwrap();
4168        }
4169
4170        let sql = r#"
4171        SHOW CREATE TABLE demo;
4172        ALTER TABLE demo ADD COLUMN new_col INT;
4173        "#;
4174        let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
4175        assert_eq!(stmts.len(), 2);
4176        for stmt in stmts {
4177            let re = check_permission(plugins.clone(), &stmt, &query_ctx);
4178            re.unwrap();
4179        }
4180
4181        fn replace_test(template_sql: &str, plugins: Plugins, query_ctx: &QueryContextRef) {
4182            // test right
4183            let right = vec![("", ""), ("", "public."), ("greptime.", "public.")];
4184            for (catalog, schema) in right {
4185                let sql = do_fmt(template_sql, catalog, schema);
4186                do_test(&sql, plugins.clone(), query_ctx, true);
4187            }
4188
4189            let wrong = vec![
4190                ("wrongcatalog.", "public."),
4191                ("wrongcatalog.", "wrongschema."),
4192            ];
4193            for (catalog, schema) in wrong {
4194                let sql = do_fmt(template_sql, catalog, schema);
4195                do_test(&sql, plugins.clone(), query_ctx, false);
4196            }
4197        }
4198
4199        fn do_fmt(template: &str, catalog: &str, schema: &str) -> String {
4200            let vars = HashMap::from([
4201                ("catalog".to_string(), catalog),
4202                ("schema".to_string(), schema),
4203            ]);
4204            template.format(&vars).unwrap()
4205        }
4206
4207        fn do_test(sql: &str, plugins: Plugins, query_ctx: &QueryContextRef, is_ok: bool) {
4208            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
4209            let re = check_permission(plugins, stmt, query_ctx);
4210            if is_ok {
4211                re.unwrap();
4212            } else {
4213                assert!(re.is_err());
4214            }
4215        }
4216
4217        // test insert
4218        let sql = "INSERT INTO {catalog}{schema}monitor(host) VALUES ('host1');";
4219        replace_test(sql, plugins.clone(), &query_ctx);
4220
4221        // test create table
4222        let sql = r#"CREATE TABLE {catalog}{schema}demo(
4223                            host STRING,
4224                            ts TIMESTAMP,
4225                            TIME INDEX (ts),
4226                            PRIMARY KEY(host)
4227                        ) engine=mito;"#;
4228        replace_test(sql, plugins.clone(), &query_ctx);
4229
4230        // test drop table
4231        let sql = "DROP TABLE {catalog}{schema}demo;";
4232        replace_test(sql, plugins.clone(), &query_ctx);
4233
4234        // test undrop table
4235        #[cfg(feature = "enterprise")]
4236        {
4237            let sql = "UNDROP TABLE {catalog}{schema}demo;";
4238            replace_test(sql, plugins.clone(), &query_ctx);
4239        }
4240
4241        // test show tables
4242        let sql = "SHOW TABLES FROM public";
4243        let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
4244        check_permission(plugins.clone(), &stmt[0], &query_ctx).unwrap();
4245
4246        let sql = "SHOW TABLES FROM private";
4247        let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
4248        let re = check_permission(plugins.clone(), &stmt[0], &query_ctx);
4249        assert!(re.is_ok());
4250
4251        // test describe table
4252        let sql = "DESC TABLE {catalog}{schema}demo;";
4253        replace_test(sql, plugins.clone(), &query_ctx);
4254
4255        let comment_flow_cases = [
4256            ("COMMENT ON FLOW my_flow IS 'comment';", true),
4257            ("COMMENT ON FLOW greptime.my_flow IS 'comment';", true),
4258            ("COMMENT ON FLOW wrongcatalog.my_flow IS 'comment';", false),
4259        ];
4260        for (sql, is_ok) in comment_flow_cases {
4261            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
4262            let result = check_permission(plugins.clone(), stmt, &query_ctx);
4263            assert_eq!(result.is_ok(), is_ok);
4264        }
4265
4266        let show_flow_cases = [
4267            ("SHOW CREATE FLOW my_flow;", true),
4268            ("SHOW CREATE FLOW greptime.my_flow;", true),
4269            ("SHOW CREATE FLOW wrongcatalog.my_flow;", false),
4270        ];
4271        for (sql, is_ok) in show_flow_cases {
4272            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
4273            let result = check_permission(plugins.clone(), stmt, &query_ctx);
4274            assert_eq!(result.is_ok(), is_ok);
4275        }
4276    }
4277
4278    /// A `DropView` DDL sent through the direct gRPC ingress must return an error
4279    /// (e.g. table not found) instead of panicking on `todo!()`.
4280    #[tokio::test]
4281    async fn qx_152_drop_view_via_grpc_ddl_returns_error_not_panic() -> TestResult<()> {
4282        let instance =
4283            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
4284                .await?;
4285
4286        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4287            expr: Some(api::v1::ddl_request::Expr::DropView(
4288                api::v1::DropViewExpr {
4289                    catalog_name: String::new(),
4290                    schema_name: String::new(),
4291                    view_name: "non_existent_view".to_string(),
4292                    view_id: None,
4293                    drop_if_exists: false,
4294                },
4295            )),
4296        });
4297
4298        let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
4299            &instance,
4300            request,
4301            QueryContext::arc(),
4302        )
4303        .await;
4304
4305        let err = match result {
4306            Ok(_) => panic!("DropView DDL request must return an error instead of panicking"),
4307            Err(err) => err,
4308        };
4309        assert_eq!(
4310            err.status_code(),
4311            StatusCode::TableNotFound,
4312            "dropping a non-existent view without IF EXISTS must report TableNotFound, got {err}"
4313        );
4314        Ok(())
4315    }
4316
4317    #[tokio::test]
4318    async fn logical_ddl_validates_and_authorizes_rewritten_batch_before_submission()
4319    -> TestResult<()> {
4320        struct Rewrite(&'static str);
4321        impl SqlQueryInterceptor for Rewrite {
4322            type Error = Error;
4323            fn pre_parsing<'a>(
4324                &self,
4325                query: &'a str,
4326                _: QueryContextRef,
4327            ) -> Result<std::borrow::Cow<'a, str>> {
4328                Ok(if self.0 == "sql_bytes" {
4329                    std::borrow::Cow::Owned(format!("{query}{}", " ".repeat(1024 * 1024 + 1)))
4330                } else {
4331                    std::borrow::Cow::Borrowed(query)
4332                })
4333            }
4334            fn pre_execute(
4335                &self,
4336                statement: Option<&Statement>,
4337                _: Option<&LogicalPlan>,
4338                _: QueryContextRef,
4339            ) -> Result<()> {
4340                ensure!(
4341                    !(self.0 == "pre_execute"
4342                        && matches!(statement,
4343                    Some(Statement::CreateTable(create)) if create.name.to_string() == "allowed2")),
4344                    InvalidSqlSnafu {
4345                        err_msg: "second member rejected by interceptor"
4346                    }
4347                );
4348                Ok(())
4349            }
4350            fn post_parsing(
4351                &self,
4352                mut statements: Vec<Statement>,
4353                _: QueryContextRef,
4354            ) -> Result<Vec<Statement>> {
4355                match self.0 {
4356                    "target" => {
4357                        statements[1] = parse_one_sql(
4358                            "CREATE TABLE denied (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH (on_physical_table='phy')",
4359                        )
4360                    }
4361                    "count" => statements = (0..129).map(|i| parse_one_sql(&format!(
4362                        "CREATE TABLE allowed_{i} (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH (on_physical_table='phy')"
4363                    ))).collect(),
4364                    "type" => statements[1] = parse_one_sql("SELECT 1"),
4365                    "secret_bytes" | "array_bytes" | "small_option" => {
4366                        if let Statement::CreateTable(create) = &mut statements[1] {
4367                            let value = "x".repeat(if self.0 == "small_option" { 10 } else { 2 * 1024 * 1024 });
4368                            if self.0 == "array_bytes" {
4369                                create.options.insert_options("non_scalar", vec![value.as_str()].into());
4370                            } else {
4371                                create.options.insert("secret_access_key".to_string(), value);
4372                            }
4373                        }
4374                    }
4375                    "partition_shape" => {
4376                        if let Statement::CreateTable(create) = &mut statements[1] {
4377                            create.partitions = Some(sql::statements::create::Partitions {
4378                                column_list: vec![],
4379                                exprs: vec![sqlparser::ast::Expr::Identifier(sqlparser::ast::Ident::new("hidden"))],
4380                            });
4381                        }
4382                    }
4383                    "vector_hidden" => {
4384                        statements[1] = parse_one_sql("CREATE TABLE allowed2 (ts TIMESTAMP TIME INDEX, v VECTOR(3)) ENGINE=metric WITH (on_physical_table='phy')");
4385                        if let Statement::CreateTable(create) = &mut statements[1] {
4386                            create.columns[1].column_def.options.push(sqlparser::ast::ColumnOptionDef {
4387                                name: None,
4388                                option: sqlparser::ast::ColumnOption::Comment("private-value".repeat(200_000)),
4389                            });
4390                        }
4391                    }
4392                    mode if mode.starts_with("fulltext_") || mode.starts_with("skipping_")
4393                        || mode.starts_with("inverted_") || mode.starts_with("vector_") => {
4394                        if let Statement::CreateTable(create) = &mut statements[1] {
4395                            let (kind, value_type) = mode.split_once('_').unwrap();
4396                            let value = "x".repeat(if value_type == "small" { 10 } else { 2 * 1024 * 1024 });
4397                            let extensions = &mut create.columns[1].extensions;
4398                            let options = match kind {
4399                                "fulltext" => &mut extensions.fulltext_index_options,
4400                                "skipping" => &mut extensions.skipping_index_options,
4401                                "inverted" => &mut extensions.inverted_index_options,
4402                                _ => &mut extensions.vector_options,
4403                            }.get_or_insert_default();
4404                            if kind == "vector" {
4405                                options.insert("dim".to_string(), "3".to_string());
4406                            }
4407                            if value_type == "array" {
4408                                options.insert_options("non_scalar", vec![value.as_str()].into());
4409                            } else {
4410                                options.insert("secret_access_key".to_string(), value);
4411                            }
4412                        }
4413                    }
4414                    "physical_null" | "physical_array" => {
4415                        let value = if self.0 == "physical_null" { "NULL" } else { "['x']" };
4416                        statements[1] = parse_one_sql(&format!(
4417                            "CREATE TABLE allowed2 (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH (on_physical_table='phy', physical_metric_table={value})"
4418                        ));
4419                    }
4420                    "ast_bytes" => {
4421                        if let Statement::CreateTable(create) = &mut statements[1] {
4422                            create.name =
4423                                ObjectName(vec![sqlparser::ast::ObjectNamePart::Identifier(
4424                                    sqlparser::ast::Ident::new("t".repeat(1024 * 1024)),
4425                                )]);
4426                        }
4427                    }
4428                    _ => {}
4429                }
4430                Ok(statements)
4431            }
4432        }
4433        for mode in [
4434            "none",
4435            "target",
4436            "count",
4437            "type",
4438            "sql_bytes",
4439            "ast_bytes",
4440            "secret_bytes",
4441            "array_bytes",
4442            "small_option",
4443            "physical_null",
4444            "physical_array",
4445            "fulltext_secret",
4446            "fulltext_array",
4447            "fulltext_small",
4448            "skipping_secret",
4449            "skipping_array",
4450            "skipping_small",
4451            "inverted_secret",
4452            "inverted_array",
4453            "inverted_small",
4454            "vector_secret",
4455            "vector_array",
4456            "vector_small",
4457            "vector_hidden",
4458            "partition_shape",
4459            "pre_execute",
4460            "cross_catalog",
4461            "valid",
4462        ] {
4463            let catalog =
4464                catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
4465            let executor = Arc::new(MockProcedureExecutor::new(catalog.clone()));
4466            let plugins = Plugins::new();
4467            plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
4468            plugins.insert::<SqlQueryInterceptorRef<Error>>(Arc::new(Rewrite(mode)));
4469            plugins.insert(QueryOptions {
4470                disallow_cross_catalog_query: mode == "cross_catalog",
4471            });
4472            let instance = test_instance_with_catalog_manager(
4473                catalog,
4474                test_table(1025, "target")?,
4475                plugins,
4476                None,
4477                executor.clone(),
4478            )
4479            .await?;
4480            let create = |name| {
4481                format!(
4482                    "CREATE TABLE {name} (ts TIMESTAMP TIME INDEX, host STRING) ENGINE=metric WITH (on_physical_table='phy')"
4483                )
4484            };
4485            let sql = format!(
4486                "{};{}",
4487                create("allowed"),
4488                create(match mode {
4489                    "none" => "denied",
4490                    "cross_catalog" => "other.public.allowed2",
4491                    _ => "allowed2",
4492                })
4493            );
4494            let result =
4495                SqlQueryHandler::create_logical_tables(&instance, &sql, QueryContext::arc()).await;
4496            let submits = matches!(mode, "valid" | "small_option")
4497                || (mode.ends_with("_small") && !mode.starts_with("vector_"));
4498            if matches!(mode, "none" | "target") {
4499                let err = result.unwrap_err();
4500                assert_eq!(
4501                    err.status_code(),
4502                    StatusCode::PermissionDenied,
4503                    "{mode}: {err}"
4504                );
4505            } else {
4506                assert_eq!(
4507                    result.unwrap_err().status_code(),
4508                    match mode {
4509                        _ if submits => StatusCode::Unsupported,
4510                        "cross_catalog" => StatusCode::AccessDenied,
4511                        _ => StatusCode::InvalidArguments,
4512                    },
4513                    "{mode}"
4514                );
4515            }
4516            assert_eq!(
4517                executor.submitted.lock().unwrap().len(),
4518                usize::from(submits),
4519                "{mode}"
4520            );
4521        }
4522        Ok(())
4523    }
4524
4525    /// `DROP VIEW IF EXISTS` on a missing view through the direct gRPC ingress must
4526    /// succeed with 0 affected rows (no error, no DDL task submitted), instead of
4527    /// returning `TableNotFound`.
4528    #[tokio::test]
4529    async fn qx_152_drop_view_if_exists_missing_view_via_grpc_ddl_succeeds() -> TestResult<()> {
4530        let catalog_manager =
4531            catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
4532        let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
4533        let instance = test_instance_with_catalog_manager(
4534            catalog_manager,
4535            test_table(1025, "target")?,
4536            Plugins::new(),
4537            None,
4538            procedure_executor.clone() as ProcedureExecutorRef,
4539        )
4540        .await?;
4541
4542        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4543            expr: Some(api::v1::ddl_request::Expr::DropView(
4544                api::v1::DropViewExpr {
4545                    catalog_name: String::new(),
4546                    schema_name: String::new(),
4547                    view_name: "non_existent_view".to_string(),
4548                    view_id: None,
4549                    drop_if_exists: true,
4550                },
4551            )),
4552        });
4553
4554        let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
4555            &instance,
4556            request,
4557            QueryContext::arc(),
4558        )
4559        .await;
4560
4561        let output = match result {
4562            Ok(output) => output,
4563            Err(err) => {
4564                panic!("DROP VIEW IF EXISTS on a missing view must succeed, got error: {err}")
4565            }
4566        };
4567        assert!(
4568            matches!(output.data, OutputData::AffectedRows(0)),
4569            "DROP VIEW IF EXISTS on a missing view must report 0 affected rows"
4570        );
4571        assert!(
4572            procedure_executor.submitted.lock().unwrap().is_empty(),
4573            "DROP VIEW IF EXISTS on a missing view must not submit a DDL task"
4574        );
4575        Ok(())
4576    }
4577
4578    /// A `CREATE VIEW` followed by `DROP VIEW` through the direct gRPC ingress must
4579    /// succeed end to end: the view is registered in the catalog and then removed.
4580    #[tokio::test]
4581    async fn qx_152_drop_existing_view_via_grpc_ddl_succeeds() -> TestResult<()> {
4582        let catalog_manager =
4583            catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
4584        let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
4585        let instance = test_instance_with_catalog_manager(
4586            catalog_manager,
4587            test_table(1025, "target")?,
4588            Plugins::new(),
4589            None,
4590            procedure_executor.clone() as ProcedureExecutorRef,
4591        )
4592        .await?;
4593
4594        // The default "greptime.public" schema must be visible to the kv-backed table
4595        // metadata manager for `CREATE VIEW`/`CREATE TABLE` to pass validation.
4596        instance
4597            .table_metadata_manager()
4598            .schema_manager()
4599            .create(
4600                common_meta::key::schema_name::SchemaNameKey::new("greptime", "public"),
4601                None,
4602                true,
4603            )
4604            .await
4605            .unwrap();
4606
4607        let create_view_request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4608            expr: Some(api::v1::ddl_request::Expr::CreateView(
4609                api::v1::CreateViewExpr {
4610                    catalog_name: String::new(),
4611                    schema_name: String::new(),
4612                    view_name: "my_view".to_string(),
4613                    logical_plan: vec![1, 2, 3],
4614                    create_if_not_exists: false,
4615                    or_replace: false,
4616                    table_names: vec![],
4617                    columns: vec![],
4618                    plan_columns: vec![],
4619                    definition: "CREATE VIEW my_view AS SELECT * FROM source".to_string(),
4620                },
4621            )),
4622        });
4623
4624        let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
4625            &instance,
4626            create_view_request,
4627            QueryContext::arc(),
4628        )
4629        .await
4630        {
4631            Ok(output) => output,
4632            Err(err) => panic!("CREATE VIEW via gRPC DDL must succeed, got error: {err}"),
4633        };
4634        assert!(
4635            matches!(output.data, OutputData::AffectedRows(0)),
4636            "CREATE VIEW via gRPC DDL must report 0 affected rows"
4637        );
4638
4639        // The view is registered in the catalog as a view.
4640        let view = instance
4641            .catalog_manager()
4642            .table("greptime", "public", "my_view", None)
4643            .await
4644            .unwrap()
4645            .expect("view should exist after CREATE VIEW");
4646        assert_eq!(view.table_info().table_type, TableType::View);
4647
4648        let drop_view_request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4649            expr: Some(api::v1::ddl_request::Expr::DropView(
4650                api::v1::DropViewExpr {
4651                    catalog_name: String::new(),
4652                    schema_name: String::new(),
4653                    view_name: "my_view".to_string(),
4654                    view_id: None,
4655                    drop_if_exists: false,
4656                },
4657            )),
4658        });
4659
4660        let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
4661            &instance,
4662            drop_view_request,
4663            QueryContext::arc(),
4664        )
4665        .await
4666        {
4667            Ok(output) => output,
4668            Err(err) => panic!("DROP VIEW via gRPC DDL must succeed, got error: {err}"),
4669        };
4670        assert!(
4671            matches!(output.data, OutputData::AffectedRows(0)),
4672            "DROP VIEW via gRPC DDL must report 0 affected rows"
4673        );
4674
4675        // The view is gone after the drop.
4676        assert!(
4677            instance
4678                .catalog_manager()
4679                .table("greptime", "public", "my_view", None)
4680                .await
4681                .unwrap()
4682                .is_none(),
4683            "view should be removed after DROP VIEW"
4684        );
4685
4686        let submitted = procedure_executor.submitted.lock().unwrap();
4687        assert_eq!(
4688            submitted.len(),
4689            2,
4690            "expected create and drop view tasks, got {submitted:?}"
4691        );
4692        assert!(matches!(&submitted[0], DdlTask::CreateView(_)));
4693        assert!(matches!(&submitted[1], DdlTask::DropView(_)));
4694        Ok(())
4695    }
4696
4697    /// A direct gRPC `CreateTable` whose time index column is not a timestamp must
4698    /// be rejected with `InvalidArguments` instead of panicking while building the schema.
4699    #[tokio::test]
4700    async fn qx_153_create_table_with_non_timestamp_time_index_via_grpc_returns_error()
4701    -> TestResult<()> {
4702        let instance =
4703            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
4704                .await?;
4705
4706        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4707            expr: Some(api::v1::ddl_request::Expr::CreateTable(
4708                api::v1::CreateTableExpr {
4709                    catalog_name: String::new(),
4710                    schema_name: String::new(),
4711                    table_name: "demo".to_string(),
4712                    desc: String::new(),
4713                    column_defs: vec![api::v1::ColumnDef {
4714                        name: "host".to_string(),
4715                        data_type: api::v1::ColumnDataType::String as i32,
4716                        is_nullable: true,
4717                        default_constraint: vec![],
4718                        semantic_type: 0,
4719                        comment: String::new(),
4720                        datatype_extension: None,
4721                        options: None,
4722                    }],
4723                    time_index: "host".to_string(),
4724                    primary_keys: vec![],
4725                    create_if_not_exists: false,
4726                    table_options: HashMap::new(),
4727                    table_id: None,
4728                    engine: "mito".to_string(),
4729                },
4730            )),
4731        });
4732
4733        let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
4734            &instance,
4735            request,
4736            QueryContext::arc(),
4737        )
4738        .await;
4739
4740        let err = match result {
4741            Ok(_) => panic!("CreateTable with a non-timestamp time index must be rejected"),
4742            Err(err) => err,
4743        };
4744        assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}");
4745        Ok(())
4746    }
4747
4748    /// A valid `CREATE TABLE` (timestamp time index) through the direct gRPC ingress
4749    /// must succeed, guarding that the `validate_create_expr` ingress check doesn't
4750    /// accidentally reject good requests.
4751    #[tokio::test]
4752    async fn qx_153_create_table_with_timestamp_time_index_via_grpc_succeeds() -> TestResult<()> {
4753        let catalog_manager =
4754            catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
4755        let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
4756        let instance = test_instance_with_catalog_manager(
4757            catalog_manager,
4758            test_table(1025, "target")?,
4759            Plugins::new(),
4760            None,
4761            procedure_executor.clone() as ProcedureExecutorRef,
4762        )
4763        .await?;
4764
4765        // The default "greptime.public" schema must be visible to the kv-backed table
4766        // metadata manager for `CREATE TABLE` to pass validation.
4767        instance
4768            .table_metadata_manager()
4769            .schema_manager()
4770            .create(
4771                common_meta::key::schema_name::SchemaNameKey::new("greptime", "public"),
4772                None,
4773                true,
4774            )
4775            .await
4776            .unwrap();
4777
4778        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4779            expr: Some(api::v1::ddl_request::Expr::CreateTable(
4780                api::v1::CreateTableExpr {
4781                    catalog_name: String::new(),
4782                    schema_name: String::new(),
4783                    table_name: "demo".to_string(),
4784                    desc: String::new(),
4785                    column_defs: vec![
4786                        api::v1::ColumnDef {
4787                            name: "host".to_string(),
4788                            data_type: api::v1::ColumnDataType::String as i32,
4789                            is_nullable: true,
4790                            default_constraint: vec![],
4791                            semantic_type: 0,
4792                            comment: String::new(),
4793                            datatype_extension: None,
4794                            options: None,
4795                        },
4796                        api::v1::ColumnDef {
4797                            name: "ts".to_string(),
4798                            data_type: api::v1::ColumnDataType::TimestampMillisecond as i32,
4799                            is_nullable: true,
4800                            default_constraint: vec![],
4801                            semantic_type: 0,
4802                            comment: String::new(),
4803                            datatype_extension: None,
4804                            options: None,
4805                        },
4806                    ],
4807                    time_index: "ts".to_string(),
4808                    primary_keys: vec![],
4809                    create_if_not_exists: false,
4810                    table_options: HashMap::new(),
4811                    table_id: None,
4812                    engine: "mito".to_string(),
4813                },
4814            )),
4815        });
4816
4817        let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
4818            &instance,
4819            request,
4820            QueryContext::arc(),
4821        )
4822        .await
4823        {
4824            Ok(output) => output,
4825            Err(err) => panic!("CREATE TABLE via gRPC DDL must succeed, got error: {err}"),
4826        };
4827        assert!(
4828            matches!(output.data, OutputData::AffectedRows(0)),
4829            "CREATE TABLE via gRPC DDL must report 0 affected rows"
4830        );
4831
4832        // The table is registered in the catalog.
4833        let table = instance
4834            .catalog_manager()
4835            .table("greptime", "public", "demo", None)
4836            .await
4837            .unwrap()
4838            .expect("table should exist after CREATE TABLE");
4839        assert_eq!(table.table_info().table_type, TableType::Base);
4840
4841        let submitted = procedure_executor.submitted.lock().unwrap();
4842        assert_eq!(
4843            submitted.len(),
4844            1,
4845            "expected one create table task, got {submitted:?}"
4846        );
4847        assert!(matches!(&submitted[0], DdlTask::CreateTable(_)));
4848        Ok(())
4849    }
4850}