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