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