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