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