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