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        if let [target] = targets.as_slice() {
1242            return if self
1243                .is_physical_query_permission_target(target, query_ctx)
1244                .await?
1245            {
1246                Ok(PermissionTableTargets::Unresolved)
1247            } else {
1248                Ok(PermissionTableTargets::resolved(targets))
1249            };
1250        }
1251
1252        // Bound catalog load and inspect results in target order to preserve serial semantics.
1253        for chunk in targets.chunks(CONCURRENCY) {
1254            let results = future::join_all(
1255                chunk
1256                    .iter()
1257                    .map(|target| self.is_physical_query_permission_target(target, query_ctx)),
1258            )
1259            .await;
1260            for result in results {
1261                if result? {
1262                    return Ok(PermissionTableTargets::Unresolved);
1263                }
1264            }
1265        }
1266
1267        Ok(PermissionTableTargets::resolved(targets))
1268    }
1269
1270    fn prom_queries_permission_targets(
1271        &self,
1272        queries: &[ParsedPromQuery],
1273        query_ctx: &QueryContextRef,
1274    ) -> server_error::Result<PermissionTableTargets> {
1275        let mut targets = Vec::new();
1276        let mut resolved = true;
1277
1278        for query in queries {
1279            let QueryStatement::Promql(eval_stmt, _) = query.statement() else {
1280                unreachable!("query is parsed from promql");
1281            };
1282
1283            if let Some(query_targets) =
1284                self.prom_expr_permission_targets(&eval_stmt.expr, query_ctx)?
1285            {
1286                targets.extend(query_targets);
1287            } else {
1288                resolved = false;
1289            }
1290        }
1291
1292        Ok(if resolved {
1293            PermissionTableTargets::resolved(targets)
1294        } else {
1295            PermissionTableTargets::Unresolved
1296        })
1297    }
1298}
1299
1300#[async_trait]
1301impl PrometheusHandler for Instance {
1302    #[tracing::instrument(skip_all)]
1303    async fn do_query(
1304        &self,
1305        query: &PromQuery,
1306        query_ctx: QueryContextRef,
1307    ) -> server_error::Result<Output> {
1308        let query = ParsedPromQuery::parse(query.clone(), &query_ctx)?;
1309        self.do_query_parsed(query, query_ctx).await
1310    }
1311
1312    #[tracing::instrument(skip_all)]
1313    async fn do_query_parsed(
1314        &self,
1315        query: ParsedPromQuery,
1316        query_ctx: QueryContextRef,
1317    ) -> server_error::Result<Output> {
1318        let interceptor = self
1319            .plugins
1320            .get::<PromQueryInterceptorRef<server_error::Error>>();
1321
1322        self.check_prom_query_privilege(&query_ctx)?;
1323
1324        let targets =
1325            self.prom_queries_permission_targets(std::slice::from_ref(&query), &query_ctx)?;
1326        self.check_query_target_permission(targets, &query_ctx)
1327            .await?;
1328
1329        let requires_output_ordering = query.requires_output_ordering();
1330        let (query, stmt) = query.into_parts();
1331
1332        let QueryStatement::Promql(eval_stmt, _) = &stmt else {
1333            unreachable!("query is parsed from promql");
1334        };
1335
1336        let plan = self
1337            .statement_executor
1338            .plan(&stmt, query_ctx.clone())
1339            .await
1340            .map_err(BoxedError::new)
1341            .context(ExecuteQuerySnafu)?;
1342
1343        let plan = if requires_output_ordering {
1344            plan
1345        } else {
1346            promql::remove_output_sort(plan)
1347        };
1348
1349        interceptor.pre_execute(&query, &eval_stmt.expr, Some(&plan), query_ctx.clone())?;
1350
1351        // Take the EvalStmt from the original QueryStatement and use it to create the CatalogQueryStatement.
1352        let query_statement = if let QueryStatement::Promql(eval_stmt, alias) = stmt {
1353            CatalogQueryStatement::Promql(eval_stmt, alias)
1354        } else {
1355            // It should not happen since the query is already parsed successfully.
1356            return UnexpectedResultSnafu {
1357                reason: "The query should always be promql.".to_string(),
1358            }
1359            .fail();
1360        };
1361        let raw_query = query_statement.to_string();
1362
1363        let slow_query_timer = self.slow_query_options.enable.then(|| {
1364            SlowQueryTimer::new(
1365                query_statement,
1366                query_ctx.current_schema(),
1367                self.slow_query_options.threshold,
1368                self.slow_query_options.sample_ratio,
1369                self.slow_query_options.record_type,
1370                self.slow_query_recorder.clone(),
1371            )
1372        });
1373
1374        let ticket = self.process_manager.register_query(
1375            query_ctx.current_catalog().to_string(),
1376            vec![query_ctx.current_schema()],
1377            raw_query,
1378            query_ctx.conn_info().to_string(),
1379            Some(query_ctx.process_id()),
1380            slow_query_timer,
1381        );
1382
1383        let query_fut = self.statement_executor.exec_plan(plan, query_ctx.clone());
1384
1385        let output = CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
1386            .await
1387            .map_err(|_| servers::error::CancelledSnafu.build())?
1388            .map_err(BoxedError::new)
1389            .context(ExecuteQuerySnafu)?;
1390        let output = map_query_output(output)
1391            .map_err(BoxedError::new)
1392            .context(ExecuteQuerySnafu)?;
1393        let Output { meta, data } = output;
1394        let data = match data {
1395            OutputData::Stream(stream) => {
1396                OutputData::Stream(Box::pin(CancellableStreamWrapper::new(stream, ticket)))
1397            }
1398            other => other,
1399        };
1400        let output = Output { data, meta };
1401        Ok(interceptor.post_execute(output, query_ctx)?)
1402    }
1403
1404    async fn check_query_permission(
1405        &self,
1406        queries: &[PromQuery],
1407        query_ctx: &QueryContextRef,
1408    ) -> server_error::Result<()> {
1409        let queries = queries
1410            .iter()
1411            .cloned()
1412            .map(|query| ParsedPromQuery::parse(query, query_ctx))
1413            .collect::<server_error::Result<Vec<_>>>()?;
1414        self.check_query_permission_parsed(&queries, query_ctx)
1415            .await
1416    }
1417
1418    async fn check_query_permission_parsed(
1419        &self,
1420        queries: &[ParsedPromQuery],
1421        query_ctx: &QueryContextRef,
1422    ) -> server_error::Result<()> {
1423        self.check_prom_query_privilege(query_ctx)?;
1424        let targets = self.prom_queries_permission_targets(queries, query_ctx)?;
1425        self.check_query_target_permission(targets, query_ctx).await
1426    }
1427
1428    async fn check_query_target_permission(
1429        &self,
1430        targets: PermissionTableTargets,
1431        query_ctx: &QueryContextRef,
1432    ) -> server_error::Result<()> {
1433        let targets = self
1434            .resolve_query_permission_targets(targets, query_ctx)
1435            .await?;
1436        self.check_table_permission(query_ctx, PermissionReq::Action(PROMQL_QUERY), targets)
1437            .context(AuthSnafu)?;
1438        Ok(())
1439    }
1440
1441    async fn filter_metadata_metric_names(
1442        &self,
1443        metric_names: Vec<String>,
1444        schema: &str,
1445        query_ctx: &QueryContextRef,
1446    ) -> server_error::Result<Vec<String>> {
1447        let checker = self.plugins.get::<PermissionCheckerRef>();
1448        if !checker.as_ref().uses_table_targets() {
1449            let Some(metric) = metric_names.first() else {
1450                return Ok(metric_names);
1451            };
1452            let target =
1453                PermissionTableTarget::new(query_ctx.current_catalog(), schema, metric.as_str());
1454            let result = checker
1455                .as_ref()
1456                .check_permission_with_table_targets(
1457                    query_ctx.current_user(),
1458                    PermissionReq::Action(PROMQL_QUERY),
1459                    PermissionTableTargets::resolved(vec![target]),
1460                )
1461                .context(AuthSnafu);
1462            return match result {
1463                Ok(_) => Ok(metric_names),
1464                Err(error)
1465                    if error.status_code()
1466                        == common_error::status_code::StatusCode::PermissionDenied =>
1467                {
1468                    Ok(Vec::new())
1469                }
1470                Err(error) => Err(error),
1471            };
1472        }
1473
1474        let mut allowed = Vec::with_capacity(metric_names.len());
1475        for metric in metric_names {
1476            let target =
1477                PermissionTableTarget::new(query_ctx.current_catalog(), schema, metric.as_str());
1478            match checker
1479                .as_ref()
1480                .check_permission_with_table_targets(
1481                    query_ctx.current_user(),
1482                    PermissionReq::Action(PROMQL_QUERY),
1483                    PermissionTableTargets::resolved(vec![target]),
1484                )
1485                .context(AuthSnafu)
1486            {
1487                Ok(_) => allowed.push(metric),
1488                Err(error)
1489                    if error.status_code()
1490                        == common_error::status_code::StatusCode::PermissionDenied => {}
1491                Err(error) => return Err(error),
1492            }
1493        }
1494        Ok(allowed)
1495    }
1496
1497    async fn query_metric_names(
1498        &self,
1499        matchers: Vec<Matcher>,
1500        schema: &str,
1501        ctx: &QueryContextRef,
1502    ) -> server_error::Result<Vec<String>> {
1503        self.handle_query_metric_names(matchers, schema, ctx)
1504            .await
1505            .map_err(BoxedError::new)
1506            .context(ExecuteQuerySnafu)
1507    }
1508
1509    async fn query_metric_names_by_labels(
1510        &self,
1511        matchers: Vec<Matcher>,
1512        schema: &str,
1513        start: SystemTime,
1514        end: SystemTime,
1515        ctx: &QueryContextRef,
1516    ) -> server_error::Result<Vec<String>> {
1517        self.handle_query_metric_names_by_labels(matchers, schema, start, end, ctx)
1518            .await
1519            .map_err(BoxedError::new)
1520            .context(ExecuteQuerySnafu)
1521    }
1522
1523    async fn query_label_values(
1524        &self,
1525        metric: String,
1526        label_name: String,
1527        matchers: Vec<Matcher>,
1528        start: SystemTime,
1529        end: SystemTime,
1530        ctx: &QueryContextRef,
1531    ) -> server_error::Result<Vec<String>> {
1532        let schema =
1533            resolve_schema_from_matchers(&matchers)?.unwrap_or_else(|| ctx.current_schema());
1534        let target = PermissionTableTarget::new(ctx.current_catalog(), schema.as_str(), &metric);
1535        self.check_query_target_permission(
1536            PermissionTableTargets::resolved(vec![target.clone()]),
1537            ctx,
1538        )
1539        .await?;
1540
1541        self.handle_query_label_values(target, label_name, matchers, start, end, ctx)
1542            .await
1543            .map_err(BoxedError::new)
1544            .context(ExecuteQuerySnafu)
1545    }
1546
1547    fn catalog_manager(&self) -> CatalogManagerRef {
1548        self.catalog_manager.clone()
1549    }
1550}
1551
1552/// Validate `stmt.database` permission if it's presented.
1553macro_rules! validate_db_permission {
1554    ($stmt: expr, $query_ctx: expr) => {
1555        if let Some(database) = &$stmt.database {
1556            validate_catalog_and_schema($query_ctx.current_catalog(), database, $query_ctx)
1557                .map_err(BoxedError::new)
1558                .context(SqlExecInterceptedSnafu)?;
1559        }
1560    };
1561}
1562
1563pub fn check_permission(
1564    plugins: Plugins,
1565    stmt: &Statement,
1566    query_ctx: &QueryContextRef,
1567) -> Result<()> {
1568    let need_validate = plugins
1569        .get::<QueryOptions>()
1570        .map(|opts| opts.disallow_cross_catalog_query)
1571        .unwrap_or_default();
1572
1573    if !need_validate {
1574        return Ok(());
1575    }
1576
1577    match stmt {
1578        // Will be checked in execution.
1579        // TODO(dennis): add a hook for admin commands.
1580        Statement::Admin(_) => {}
1581        // These are executed by query engine, and will be checked there.
1582        Statement::Query(_)
1583        | Statement::Explain(_)
1584        | Statement::Tql(_)
1585        | Statement::Delete(_)
1586        | Statement::DeclareCursor(_)
1587        | Statement::Copy(sql::statements::copy::Copy::CopyQueryTo(_)) => {}
1588        // database ops won't be checked
1589        Statement::CreateDatabase(_)
1590        | Statement::ShowDatabases(_)
1591        | Statement::DropDatabase(_)
1592        | Statement::AlterDatabase(_)
1593        | Statement::DropFlow(_)
1594        | Statement::Use(_) => {}
1595        #[cfg(feature = "enterprise")]
1596        Statement::DropTrigger(_) => {}
1597        Statement::ShowCreateDatabase(stmt) => {
1598            validate_database(&stmt.database_name, query_ctx)?;
1599        }
1600        Statement::ShowCreateTable(stmt) => {
1601            validate_param(&stmt.table_name, query_ctx)?;
1602        }
1603        Statement::ShowCreateFlow(stmt) => {
1604            validate_flow(&stmt.flow_name, query_ctx)?;
1605        }
1606        #[cfg(feature = "enterprise")]
1607        Statement::ShowCreateTrigger(stmt) => {
1608            validate_param(&stmt.trigger_name, query_ctx)?;
1609        }
1610        Statement::ShowCreateView(stmt) => {
1611            validate_param(&stmt.view_name, query_ctx)?;
1612        }
1613        Statement::CreateExternalTable(stmt) => {
1614            validate_param(&stmt.name, query_ctx)?;
1615        }
1616        Statement::CreateFlow(stmt) => {
1617            // TODO: should also validate source table name here?
1618            validate_param(&stmt.sink_table_name, query_ctx)?;
1619        }
1620        #[cfg(feature = "enterprise")]
1621        Statement::CreateTrigger(stmt) => {
1622            validate_param(&stmt.trigger_name, query_ctx)?;
1623        }
1624        Statement::CreateView(stmt) => {
1625            validate_param(&stmt.name, query_ctx)?;
1626        }
1627        Statement::AlterTable(stmt) => {
1628            validate_param(stmt.table_name(), query_ctx)?;
1629        }
1630        #[cfg(feature = "enterprise")]
1631        Statement::AlterTrigger(_) => {}
1632        // set/show variable now only alter/show variable in session
1633        Statement::SetVariables(_) | Statement::ShowVariables(_) => {}
1634        // show charset and show collation won't be checked
1635        Statement::ShowCharset(_) | Statement::ShowCollation(_) => {}
1636
1637        Statement::Comment(comment) => match &comment.object {
1638            CommentObject::Table(table) => validate_param(table, query_ctx)?,
1639            CommentObject::Column { table, .. } => validate_param(table, query_ctx)?,
1640            CommentObject::Flow(flow) => validate_flow(flow, query_ctx)?,
1641        },
1642
1643        Statement::Insert(insert) => {
1644            let name = insert.table_name().context(ParseSqlSnafu)?;
1645            validate_param(name, query_ctx)?;
1646        }
1647        Statement::CreateTable(stmt) => {
1648            validate_param(&stmt.name, query_ctx)?;
1649        }
1650        Statement::CreateTableLike(stmt) => {
1651            validate_param(&stmt.table_name, query_ctx)?;
1652            validate_param(&stmt.source_name, query_ctx)?;
1653        }
1654        Statement::DropTable(drop_stmt) => {
1655            for table_name in drop_stmt.table_names() {
1656                validate_param(table_name, query_ctx)?;
1657            }
1658        }
1659        #[cfg(feature = "enterprise")]
1660        Statement::UndropTable(stmt) => {
1661            validate_param(stmt.table_name(), query_ctx)?;
1662        }
1663        Statement::DropView(stmt) => {
1664            validate_param(&stmt.view_name, query_ctx)?;
1665        }
1666        Statement::ShowTables(stmt) => {
1667            validate_db_permission!(stmt, query_ctx);
1668        }
1669        Statement::ShowTableStatus(stmt) => {
1670            validate_db_permission!(stmt, query_ctx);
1671        }
1672        Statement::ShowColumns(stmt) => {
1673            validate_db_permission!(stmt, query_ctx);
1674        }
1675        Statement::ShowIndex(stmt) => {
1676            validate_db_permission!(stmt, query_ctx);
1677        }
1678        Statement::ShowRegion(stmt) => {
1679            validate_db_permission!(stmt, query_ctx);
1680        }
1681        Statement::ShowViews(stmt) => {
1682            validate_db_permission!(stmt, query_ctx);
1683        }
1684        Statement::ShowFlows(stmt) => {
1685            validate_db_permission!(stmt, query_ctx);
1686        }
1687        Statement::ShowFlowStatus(_stmt) => {
1688            // Flow statistics are organized based on the catalog dimension and
1689            // filtered by the current catalog, so there is no need to check the
1690            // permission of the database(schema).
1691        }
1692        #[cfg(feature = "enterprise")]
1693        Statement::ShowTriggers(_stmt) => {
1694            // The trigger is organized based on the catalog dimension, so there
1695            // is no need to check the permission of the database(schema).
1696        }
1697        Statement::ShowStatus(_stmt) => {}
1698        Statement::ShowSearchPath(_stmt) => {}
1699        Statement::DescribeTable(stmt) => {
1700            validate_param(stmt.name(), query_ctx)?;
1701        }
1702        Statement::Copy(sql::statements::copy::Copy::CopyTable(stmt)) => match stmt {
1703            CopyTable::To(copy_table_to) => validate_param(&copy_table_to.table_name, query_ctx)?,
1704            CopyTable::From(copy_table_from) => {
1705                validate_param(&copy_table_from.table_name, query_ctx)?
1706            }
1707        },
1708        Statement::Copy(sql::statements::copy::Copy::CopyDatabase(copy_database)) => {
1709            match copy_database {
1710                CopyDatabase::To(stmt) => validate_database(&stmt.database_name, query_ctx)?,
1711                CopyDatabase::From(stmt) => validate_database(&stmt.database_name, query_ctx)?,
1712            }
1713        }
1714        Statement::TruncateTable(stmt) => {
1715            validate_param(stmt.table_name(), query_ctx)?;
1716        }
1717        // cursor operations are always allowed once it's created
1718        Statement::FetchCursor(_) | Statement::CloseCursor(_) => {}
1719        // User can only kill process in their own catalog.
1720        Statement::Kill(_) => {}
1721        // SHOW PROCESSLIST
1722        Statement::ShowProcesslist(_) => {}
1723    }
1724    Ok(())
1725}
1726
1727fn validate_param(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
1728    let (catalog, schema, _) = table_idents_to_full_name(name, query_ctx)
1729        .map_err(BoxedError::new)
1730        .context(ExternalSnafu)?;
1731
1732    validate_catalog_and_schema(&catalog, &schema, query_ctx)
1733        .map_err(BoxedError::new)
1734        .context(SqlExecInterceptedSnafu)
1735}
1736
1737fn validate_flow(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
1738    let catalog = match &name.0[..] {
1739        [_flow] => query_ctx.current_catalog().to_string(),
1740        [catalog, _flow] => catalog.to_string_unquoted(),
1741        _ => {
1742            return InvalidSqlSnafu {
1743                err_msg: format!(
1744                    "expect flow name to be <catalog>.<flow_name> or <flow_name>, actual: {name}",
1745                ),
1746            }
1747            .fail();
1748        }
1749    };
1750
1751    let schema = query_ctx.current_schema();
1752
1753    validate_catalog_and_schema(&catalog, &schema, query_ctx)
1754        .map_err(BoxedError::new)
1755        .context(SqlExecInterceptedSnafu)
1756}
1757
1758fn validate_database(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
1759    let (catalog, schema) = match &name.0[..] {
1760        [schema] => (
1761            query_ctx.current_catalog().to_string(),
1762            schema.to_string_unquoted(),
1763        ),
1764        [catalog, schema] => (catalog.to_string_unquoted(), schema.to_string_unquoted()),
1765        _ => InvalidSqlSnafu {
1766            err_msg: format!(
1767                "expect database name to be <catalog>.<schema> or <schema>, actual: {name}",
1768            ),
1769        }
1770        .fail()?,
1771    };
1772
1773    validate_catalog_and_schema(&catalog, &schema, query_ctx)
1774        .map_err(BoxedError::new)
1775        .context(SqlExecInterceptedSnafu)
1776}
1777
1778fn is_readonly_plan(plan: &LogicalPlan) -> bool {
1779    !matches!(plan, LogicalPlan::Dml(_) | LogicalPlan::Ddl(_))
1780}
1781
1782fn should_track_statement_process(stmt: &Statement) -> bool {
1783    stmt.is_readonly()
1784        || matches!(stmt, Statement::Insert(insert) if insert.has_non_values_query_source())
1785}
1786
1787fn should_track_plan_process(stmt: Option<&Statement>, plan: &LogicalPlan) -> bool {
1788    is_readonly_plan(plan)
1789        || matches!(stmt, Some(Statement::Insert(insert)) if insert.has_non_values_query_source())
1790}
1791
1792#[cfg(test)]
1793mod tests {
1794    use std::any::Any;
1795    use std::collections::HashMap;
1796    use std::future::Future;
1797    use std::pin::Pin;
1798    use std::sync::Arc;
1799    use std::task::{Context, Poll};
1800    use std::time::Duration;
1801
1802    use api::prom_store::remote::label_matcher::Type as PromMatcherType;
1803    use api::prom_store::remote::{
1804        Label, LabelMatcher, Query as RemoteQuery, ReadRequest, ReadResponse, Sample,
1805    };
1806    use api::v1::greptime_request::Request;
1807    use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse};
1808    use api::v1::query_request::Query;
1809    use auth::{
1810        DASHBOARD_DELETE, DASHBOARD_QUERY, DASHBOARD_SAVE, JAEGER_QUERY, PIPELINE_DELETE,
1811        PIPELINE_INSERT, PIPELINE_QUERY, PermissionAction, PermissionResp, UserInfo, UserInfoRef,
1812    };
1813    use catalog::process_manager::{ProcessManager, QueryStatement, SlowQueryTimer};
1814    use common_base::Plugins;
1815    use common_catalog::consts::DEFAULT_PRIVATE_SCHEMA_NAME;
1816    use common_error::ext::{BoxedError, ErrorExt, PlainError};
1817    use common_error::status_code::StatusCode;
1818    use common_event_recorder::{Event, EventRecorder, EventTypeFilter, EventTypeFilterRef};
1819    use common_frontend::metrics::{CACHE_HIT, CACHE_MISS};
1820    use common_frontend::slow_query_event::SlowQueryEvent;
1821    use common_meta::cache::LayeredCacheRegistryBuilder;
1822    use common_meta::kv_backend::memory::MemoryKvBackend;
1823    use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor};
1824    use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse};
1825    use common_meta::rpc::procedure::{
1826        MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse,
1827    };
1828    use common_query::prelude::greptime_value;
1829    use common_query::{Output, OutputMeta};
1830    use common_recordbatch::{
1831        OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream,
1832    };
1833    use common_telemetry::logging::SlowQueriesRecordType;
1834    use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
1835    use datafusion::physical_plan::empty::EmptyExec;
1836    use datafusion_expr::dml::InsertOp;
1837    use datafusion_expr::{LogicalPlanBuilder, LogicalTableSource};
1838    use datatypes::prelude::ConcreteDataType;
1839    use datatypes::schema::{ColumnSchema, Schema as GtSchema, SchemaRef as GtSchemaRef};
1840    use datatypes::vectors::{
1841        Float64Vector, StringVector, TimestampMillisecondVector, TimestampNanosecondVector,
1842        VectorRef,
1843    };
1844    use log_query::LogQuery;
1845    use prost::Message;
1846    use query::query_engine::options::QueryOptions;
1847    use servers::query_handler::{
1848        DashboardHandler, JaegerQueryHandler, LogQueryHandler, PipelineHandler, PipelineHandlerRef,
1849        PromStoreProtocolHandler,
1850    };
1851    use session::context::{Channel, ConnInfo, QueryContext, QueryContextBuilder};
1852    use snafu::{Location, Snafu};
1853    use sql::dialect::GreptimeDbDialect;
1854    use store_api::data_source::DataSource;
1855    use store_api::metric_engine_consts::{
1856        LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY,
1857    };
1858    use store_api::storage::ScanRequest;
1859    use strfmt::Format;
1860    use table::metadata::{
1861        FilterPushDownType, TableInfo, TableInfoBuilder, TableMetaBuilder, TableType,
1862    };
1863    use table::table_name::TableName;
1864    use table::test_util::{EmptyTable, MemTable};
1865    use table::{Table, TableRef};
1866    use tokio::sync::{mpsc, oneshot};
1867    use tower::ServiceExt;
1868
1869    use crate::frontend::FrontendOptions;
1870    use crate::instance::builder::FrontendBuilder;
1871    use crate::instance::*;
1872
1873    fn parse_test_sql(sql: &str) -> Vec<Statement> {
1874        parse_stmt(sql, &GreptimeDbDialect {}).unwrap()
1875    }
1876
1877    #[derive(Debug, Default)]
1878    struct RecordingSlowQueryEventRecorder {
1879        queries: std::sync::Mutex<Vec<String>>,
1880        payloads: std::sync::Mutex<Vec<serde_json::Value>>,
1881    }
1882
1883    impl EventRecorder for RecordingSlowQueryEventRecorder {
1884        fn record(&self, event: Box<dyn Event>) {
1885            let event = event
1886                .as_any()
1887                .downcast_ref::<SlowQueryEvent>()
1888                .expect("expected a slow query event");
1889            self.queries.lock().unwrap().push(event.query.clone());
1890            self.payloads.lock().unwrap().push(event.payload.clone());
1891        }
1892
1893        fn event_type_filter(&self) -> EventTypeFilterRef {
1894            Arc::new(EventTypeFilter::All)
1895        }
1896
1897        fn close(&self) {}
1898    }
1899
1900    #[test]
1901    fn test_validate_analyze_stream_statement_strictness() {
1902        for sql in [
1903            "select 1",
1904            "explain analyze select 1",
1905            "explain analyze verbose format text select 1",
1906            "explain analyze verbose format graphviz select 1",
1907            "TQL ANALYZE (0, 10, '5s') physical_metric",
1908            "TQL EXPLAIN VERBOSE (0, 10, '5s') physical_metric",
1909            "TQL ANALYZE VERBOSE FORMAT TEXT (0, 10, '5s') physical_metric",
1910        ] {
1911            let mut stmts = parse_test_sql(sql);
1912            assert!(
1913                validate_analyze_stream_statement(&mut stmts[0]).is_err(),
1914                "{sql}"
1915            );
1916        }
1917
1918        for sql in [
1919            "explain analyze verbose select 1",
1920            "explain analyze verbose format json select 1",
1921            "TQL ANALYZE VERBOSE (0, 10, '5s') physical_metric",
1922            "TQL ANALYZE VERBOSE FORMAT JSON (0, 10, '5s') physical_metric",
1923        ] {
1924            let mut stmts = parse_test_sql(sql);
1925            assert!(
1926                validate_analyze_stream_statement(&mut stmts[0]).is_ok(),
1927                "{sql}"
1928            );
1929            match &stmts[0] {
1930                Statement::Explain(explain) => assert!(explain.format.is_none()),
1931                Statement::Tql(Tql::Analyze(analyze)) => assert!(analyze.format.is_none()),
1932                _ => unreachable!(),
1933            }
1934        }
1935
1936        assert_eq!(
1937            parse_test_sql("explain analyze verbose select 1; select 2").len(),
1938            2
1939        );
1940
1941        assert!(is_explain_analyze_verbose(
1942            &parse_test_sql("explain analyze verbose select 1")[0]
1943        ));
1944        assert!(is_explain_analyze_verbose(
1945            &parse_test_sql("TQL ANALYZE VERBOSE (0, 10, '5s') physical_metric")[0]
1946        ));
1947        for sql in [
1948            "select 1",
1949            "explain select 1",
1950            "explain analyze select 1",
1951            "explain verbose select 1",
1952            "TQL ANALYZE (0, 10, '5s') physical_metric",
1953            "TQL EXPLAIN VERBOSE (0, 10, '5s') physical_metric",
1954        ] {
1955            assert!(
1956                !is_explain_analyze_verbose(&parse_test_sql(sql)[0]),
1957                "{sql}"
1958            );
1959        }
1960    }
1961
1962    #[derive(Debug, Snafu)]
1963    enum TestError {
1964        #[snafu(display("Failed to build test cache registry"))]
1965        BuildCacheRegistry {
1966            source: cache::error::Error,
1967            #[snafu(implicit)]
1968            location: Location,
1969        },
1970
1971        #[snafu(display("Failed to build test table meta for table: {table_name}"))]
1972        BuildTableMeta {
1973            table_name: String,
1974            source: table::metadata::TableMetaBuilderError,
1975            #[snafu(implicit)]
1976            location: Location,
1977        },
1978
1979        #[snafu(display("Failed to build test table info for table: {table_name}"))]
1980        BuildTableInfo {
1981            table_name: String,
1982            source: table::metadata::TableInfoBuilderError,
1983            #[snafu(implicit)]
1984            location: Location,
1985        },
1986
1987        #[snafu(display("Failed to register test table: {table_name}"))]
1988        RegisterTable {
1989            table_name: String,
1990            source: catalog::error::Error,
1991            #[snafu(implicit)]
1992            location: Location,
1993        },
1994
1995        #[snafu(display("Failed to build test frontend instance"))]
1996        BuildFrontend {
1997            source: crate::error::Error,
1998            #[snafu(implicit)]
1999            location: Location,
2000        },
2001
2002        #[snafu(display("Expected exactly one output for SQL `{sql}`, got {actual}"))]
2003        UnexpectedOutputCount {
2004            sql: String,
2005            actual: usize,
2006            #[snafu(implicit)]
2007            location: Location,
2008        },
2009
2010        #[snafu(display("Failed to execute SQL `{sql}`"))]
2011        ExecuteSql {
2012            sql: String,
2013            source: crate::error::Error,
2014            #[snafu(implicit)]
2015            location: Location,
2016        },
2017
2018        #[snafu(display("Timed out waiting for insert-select start notification"))]
2019        InsertStartTimeout {
2020            source: tokio::time::error::Elapsed,
2021            #[snafu(implicit)]
2022            location: Location,
2023        },
2024
2025        #[snafu(display("Insert-select start notification channel closed"))]
2026        InsertStartChannelClosed {
2027            #[snafu(implicit)]
2028            location: Location,
2029        },
2030
2031        #[snafu(display("Failed to release blocking insert-select interceptor"))]
2032        ReleaseBlockedInsert {
2033            #[snafu(implicit)]
2034            location: Location,
2035        },
2036
2037        #[snafu(display("Timed out waiting for insert-select source to be polled"))]
2038        SourcePollTimeout {
2039            source: tokio::time::error::Elapsed,
2040            #[snafu(implicit)]
2041            location: Location,
2042        },
2043
2044        #[snafu(display("Insert-select source poll notification channel closed"))]
2045        SourcePollChannelClosed {
2046            source: oneshot::error::RecvError,
2047            #[snafu(implicit)]
2048            location: Location,
2049        },
2050
2051        #[snafu(display("Timed out waiting for insert task to finish"))]
2052        InsertTaskTimeout {
2053            source: tokio::time::error::Elapsed,
2054            #[snafu(implicit)]
2055            location: Location,
2056        },
2057
2058        #[snafu(display("Insert task panicked"))]
2059        InsertTaskPanic {
2060            source: tokio::task::JoinError,
2061            #[snafu(implicit)]
2062            location: Location,
2063        },
2064
2065        #[snafu(display("Expected insert-select to be cancelled"))]
2066        InsertSelectNotCancelled {
2067            #[snafu(implicit)]
2068            location: Location,
2069        },
2070    }
2071
2072    type TestResult<T> = std::result::Result<T, TestError>;
2073
2074    fn parse_one_sql(sql: &str) -> Statement {
2075        parse_stmt(sql, &GreptimeDbDialect {}).unwrap().remove(0)
2076    }
2077
2078    fn test_query_ctx(process_id: u32) -> QueryContextRef {
2079        Arc::new(
2080            QueryContextBuilder::default()
2081                .channel(Channel::Mysql)
2082                .conn_info(ConnInfo::new(None, Channel::Mysql))
2083                .process_id(process_id)
2084                .build(),
2085        )
2086    }
2087
2088    #[derive(Debug)]
2089    struct AdminUserInfo;
2090
2091    impl UserInfo for AdminUserInfo {
2092        fn as_any(&self) -> &dyn Any {
2093            self
2094        }
2095
2096        fn username(&self) -> &str {
2097            "admin"
2098        }
2099
2100        fn is_admin(&self) -> bool {
2101            true
2102        }
2103    }
2104
2105    struct RejectUnresolvedPermissionChecker;
2106
2107    impl PermissionChecker for RejectUnresolvedPermissionChecker {
2108        fn check_permission(
2109            &self,
2110            _user_info: UserInfoRef,
2111            _req: PermissionReq,
2112        ) -> auth::error::Result<PermissionResp> {
2113            Ok(PermissionResp::Allow)
2114        }
2115
2116        fn check_permission_with_table_targets(
2117            &self,
2118            _user_info: UserInfoRef,
2119            _req: PermissionReq,
2120            targets: PermissionTableTargets,
2121        ) -> auth::error::Result<PermissionResp> {
2122            let reject = match targets {
2123                PermissionTableTargets::Unresolved => true,
2124                PermissionTableTargets::Resolved(targets) => {
2125                    targets.iter().any(|target| target.table == "denied")
2126                }
2127            };
2128            Ok(if reject {
2129                PermissionResp::Reject
2130            } else {
2131                PermissionResp::Allow
2132            })
2133        }
2134    }
2135
2136    #[derive(Debug, PartialEq, Eq)]
2137    struct CheckedAction {
2138        action: PermissionAction,
2139        targets: Option<PermissionTableTargets>,
2140    }
2141
2142    #[derive(Default)]
2143    struct RejectEndpointPermissionChecker {
2144        checks: std::sync::Mutex<Vec<CheckedAction>>,
2145    }
2146
2147    impl RejectEndpointPermissionChecker {
2148        fn reject(
2149            &self,
2150            action: PermissionAction,
2151            targets: Option<PermissionTableTargets>,
2152        ) -> PermissionResp {
2153            self.checks
2154                .lock()
2155                .unwrap()
2156                .push(CheckedAction { action, targets });
2157            PermissionResp::Reject
2158        }
2159
2160        fn take_check(&self) -> CheckedAction {
2161            let mut checks = self.checks.lock().unwrap();
2162            assert_eq!(1, checks.len());
2163            checks.pop().unwrap()
2164        }
2165    }
2166
2167    impl PermissionChecker for RejectEndpointPermissionChecker {
2168        fn check_permission(
2169            &self,
2170            _user_info: UserInfoRef,
2171            req: PermissionReq,
2172        ) -> auth::error::Result<PermissionResp> {
2173            Ok(match req {
2174                PermissionReq::Action(action) => self.reject(action, None),
2175                _ => PermissionResp::Allow,
2176            })
2177        }
2178
2179        fn check_permission_with_table_targets(
2180            &self,
2181            _user_info: UserInfoRef,
2182            req: PermissionReq,
2183            targets: PermissionTableTargets,
2184        ) -> auth::error::Result<PermissionResp> {
2185            Ok(match req {
2186                PermissionReq::Action(action) => self.reject(action, Some(targets)),
2187                _ => PermissionResp::Allow,
2188            })
2189        }
2190    }
2191
2192    struct WriteOnlyPermissionChecker;
2193
2194    impl PermissionChecker for WriteOnlyPermissionChecker {
2195        fn check_permission(
2196            &self,
2197            _user_info: UserInfoRef,
2198            req: PermissionReq,
2199        ) -> auth::error::Result<PermissionResp> {
2200            Ok(if req.is_readonly() {
2201                PermissionResp::Reject
2202            } else {
2203                PermissionResp::Allow
2204            })
2205        }
2206
2207        fn check_permission_with_table_targets(
2208            &self,
2209            user_info: UserInfoRef,
2210            req: PermissionReq,
2211            _targets: PermissionTableTargets,
2212        ) -> auth::error::Result<PermissionResp> {
2213            self.check_permission(user_info, req)
2214        }
2215    }
2216
2217    #[derive(Default)]
2218    struct TargetIndependentPermissionChecker {
2219        checks: atomic::AtomicUsize,
2220    }
2221
2222    impl PermissionChecker for TargetIndependentPermissionChecker {
2223        fn check_permission(
2224            &self,
2225            _user_info: UserInfoRef,
2226            _req: PermissionReq,
2227        ) -> auth::error::Result<PermissionResp> {
2228            self.checks.fetch_add(1, atomic::Ordering::Relaxed);
2229            Ok(PermissionResp::Allow)
2230        }
2231
2232        fn uses_table_targets(&self) -> bool {
2233            false
2234        }
2235
2236        fn check_permission_with_table_targets(
2237            &self,
2238            user_info: UserInfoRef,
2239            req: PermissionReq,
2240            _targets: PermissionTableTargets,
2241        ) -> auth::error::Result<PermissionResp> {
2242            self.check_permission(user_info, req)
2243        }
2244    }
2245
2246    struct BlockingInsertSelectInterceptor {
2247        started_tx: mpsc::UnboundedSender<()>,
2248        finish_rx: std::sync::Mutex<Option<oneshot::Receiver<()>>>,
2249    }
2250
2251    impl BlockingInsertSelectInterceptor {
2252        fn new(started_tx: mpsc::UnboundedSender<()>, finish_rx: oneshot::Receiver<()>) -> Self {
2253            Self {
2254                started_tx,
2255                finish_rx: std::sync::Mutex::new(Some(finish_rx)),
2256            }
2257        }
2258    }
2259
2260    impl SqlQueryInterceptor for BlockingInsertSelectInterceptor {
2261        type Error = Error;
2262
2263        fn pre_execute(
2264            &self,
2265            statement: Option<&Statement>,
2266            _plan: Option<&LogicalPlan>,
2267            _query_ctx: QueryContextRef,
2268        ) -> Result<()> {
2269            let Some(Statement::Insert(insert)) = statement else {
2270                return Ok(());
2271            };
2272            if !insert.has_non_values_query_source() {
2273                return Ok(());
2274            }
2275
2276            let finish_rx = self.finish_rx.lock().unwrap().take().unwrap();
2277            let _ = self.started_tx.send(());
2278            tokio::task::block_in_place(|| {
2279                tokio::runtime::Handle::current()
2280                    .block_on(finish_rx)
2281                    .unwrap();
2282            });
2283            Ok(())
2284        }
2285    }
2286
2287    struct PendingRecordBatchStream {
2288        schema: GtSchemaRef,
2289        polled_tx: Option<oneshot::Sender<()>>,
2290        _finish_tx: oneshot::Sender<()>,
2291        finish_rx: Pin<Box<oneshot::Receiver<()>>>,
2292    }
2293
2294    impl RecordBatchStream for PendingRecordBatchStream {
2295        fn schema(&self) -> GtSchemaRef {
2296            self.schema.clone()
2297        }
2298
2299        fn output_ordering(&self) -> Option<&[OrderOption]> {
2300            None
2301        }
2302
2303        fn metrics(&self) -> Option<common_recordbatch::adapter::RecordBatchMetrics> {
2304            None
2305        }
2306    }
2307
2308    impl Stream for PendingRecordBatchStream {
2309        type Item = common_recordbatch::error::Result<RecordBatch>;
2310
2311        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
2312            if let Some(polled_tx) = self.polled_tx.take() {
2313                let _ = polled_tx.send(());
2314            }
2315
2316            match self.finish_rx.as_mut().poll(cx) {
2317                Poll::Ready(_) => Poll::Ready(None),
2318                Poll::Pending => Poll::Pending,
2319            }
2320        }
2321    }
2322
2323    impl Unpin for PendingRecordBatchStream {}
2324
2325    #[test]
2326    fn test_record_explain_analyze_timeout_uses_empty_metrics_without_plan() {
2327        let event_recorder = Arc::new(RecordingSlowQueryEventRecorder::default());
2328        let timer = SlowQueryTimer::new(
2329            QueryStatement::Plan("EXPLAIN ANALYZE VERBOSE SELECT 1".to_string()),
2330            "public".to_string(),
2331            Duration::from_secs(3600),
2332            0.0,
2333            SlowQueriesRecordType::SystemTable,
2334            event_recorder.clone(),
2335        );
2336        let timeout_recorder = timer.recorder();
2337
2338        record_explain_analyze_timeout(Some(&timeout_recorder), None);
2339        drop(timer);
2340
2341        let payloads = event_recorder.payloads.lock().unwrap();
2342        assert_eq!(payloads.len(), 1);
2343        assert_eq!(payloads[0]["timed_out"], true);
2344        assert_eq!(payloads[0]["metrics"], serde_json::json!([]));
2345    }
2346
2347    #[tokio::test]
2348    async fn test_attach_timeout_records_explain_analyze_metrics() {
2349        let event_recorder = Arc::new(RecordingSlowQueryEventRecorder::default());
2350        let timer = SlowQueryTimer::new(
2351            QueryStatement::Plan("EXPLAIN ANALYZE VERBOSE SELECT 1".to_string()),
2352            "public".to_string(),
2353            Duration::from_secs(3600),
2354            0.0,
2355            SlowQueriesRecordType::SystemTable,
2356            event_recorder.clone(),
2357        );
2358        let timeout_recorder = timer.recorder();
2359        let plan: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(Arc::new(Schema::empty())));
2360        let (finish_tx, finish_rx) = oneshot::channel();
2361        let stream = PendingRecordBatchStream {
2362            schema: Arc::new(GtSchema::new(vec![])),
2363            polled_tx: None,
2364            _finish_tx: finish_tx,
2365            finish_rx: Box::pin(finish_rx),
2366        };
2367        let output = Output::new(
2368            OutputData::Stream(Box::pin(stream)),
2369            OutputMeta::new_with_plan(plan),
2370        );
2371        let output =
2372            attach_timeout(output, Duration::from_millis(10), Some(timeout_recorder)).unwrap();
2373        let OutputData::Stream(mut stream) = output.data else {
2374            unreachable!();
2375        };
2376
2377        let err = stream.next().await.unwrap().unwrap_err();
2378        assert_eq!(err.to_string(), "Stream timeout");
2379        drop(stream);
2380        drop(timer);
2381
2382        let payloads = event_recorder.payloads.lock().unwrap();
2383        assert_eq!(payloads.len(), 1);
2384        assert_eq!(payloads[0]["timed_out"], true);
2385        assert!(
2386            payloads[0]["metrics"]
2387                .as_array()
2388                .is_some_and(|metrics| !metrics.is_empty())
2389        );
2390    }
2391
2392    struct PendingDataSource {
2393        schema: GtSchemaRef,
2394        polled_tx: std::sync::Mutex<Option<oneshot::Sender<()>>>,
2395    }
2396
2397    impl DataSource for PendingDataSource {
2398        fn get_stream(
2399            &self,
2400            _request: ScanRequest,
2401        ) -> std::result::Result<SendableRecordBatchStream, BoxedError> {
2402            let (finish_tx, finish_rx) = oneshot::channel();
2403            let mut polled_tx = self.polled_tx.lock().map_err(|_| {
2404                BoxedError::new(PlainError::new(
2405                    "pending data source lock poisoned".to_string(),
2406                    StatusCode::Unexpected,
2407                ))
2408            })?;
2409            Ok(Box::pin(PendingRecordBatchStream {
2410                schema: self.schema.clone(),
2411                polled_tx: polled_tx.take(),
2412                _finish_tx: finish_tx,
2413                finish_rx: Box::pin(finish_rx),
2414            }))
2415        }
2416    }
2417
2418    struct NoopProcedureExecutor;
2419
2420    #[async_trait::async_trait]
2421    impl ProcedureExecutor for NoopProcedureExecutor {
2422        async fn submit_ddl_task(
2423            &self,
2424            _ctx: ExecutorContext,
2425            _request: SubmitDdlTaskRequest,
2426        ) -> common_meta::error::Result<SubmitDdlTaskResponse> {
2427            common_meta::error::UnsupportedSnafu {
2428                operation: "submit_ddl_task",
2429            }
2430            .fail()
2431        }
2432
2433        async fn migrate_region(
2434            &self,
2435            _ctx: &ExecutorContext,
2436            _request: MigrateRegionRequest,
2437        ) -> common_meta::error::Result<MigrateRegionResponse> {
2438            common_meta::error::UnsupportedSnafu {
2439                operation: "migrate_region",
2440            }
2441            .fail()
2442        }
2443
2444        async fn reconcile(
2445            &self,
2446            _ctx: &ExecutorContext,
2447            _request: ReconcileRequest,
2448        ) -> common_meta::error::Result<ReconcileResponse> {
2449            common_meta::error::UnsupportedSnafu {
2450                operation: "reconcile",
2451            }
2452            .fail()
2453        }
2454
2455        async fn query_procedure_state(
2456            &self,
2457            _ctx: &ExecutorContext,
2458            _pid: &str,
2459        ) -> common_meta::error::Result<ProcedureStateResponse> {
2460            common_meta::error::UnsupportedSnafu {
2461                operation: "query_procedure_state",
2462            }
2463            .fail()
2464        }
2465
2466        async fn list_procedures(
2467            &self,
2468            _ctx: &ExecutorContext,
2469        ) -> common_meta::error::Result<ProcedureDetailResponse> {
2470            common_meta::error::UnsupportedSnafu {
2471                operation: "list_procedures",
2472            }
2473            .fail()
2474        }
2475    }
2476
2477    /// A test [`ProcedureExecutor`] that completes create/drop DDL tasks against the
2478    /// in-memory catalog, mimicking what the meta DDL procedures do in production.
2479    /// This allows happy-path DDL requests (create/drop table/view) to be exercised
2480    /// end to end through the gRPC ingress.
2481    struct MockProcedureExecutor {
2482        catalog_manager: Arc<catalog::memory::MemoryCatalogManager>,
2483        next_table_id: std::sync::atomic::AtomicU32,
2484        submitted: std::sync::Mutex<Vec<DdlTask>>,
2485    }
2486
2487    impl MockProcedureExecutor {
2488        fn new(catalog_manager: Arc<catalog::memory::MemoryCatalogManager>) -> Self {
2489            Self {
2490                catalog_manager,
2491                next_table_id: std::sync::atomic::AtomicU32::new(1026),
2492                submitted: std::sync::Mutex::new(Vec::new()),
2493            }
2494        }
2495    }
2496
2497    #[async_trait::async_trait]
2498    impl ProcedureExecutor for MockProcedureExecutor {
2499        async fn submit_ddl_task(
2500            &self,
2501            _ctx: ExecutorContext,
2502            request: SubmitDdlTaskRequest,
2503        ) -> common_meta::error::Result<SubmitDdlTaskResponse> {
2504            self.submitted.lock().unwrap().push(request.task.clone());
2505            match request.task {
2506                DdlTask::CreateTable(task) => {
2507                    let table_id = self
2508                        .next_table_id
2509                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2510                    let mut table_info = task.table_info;
2511                    table_info.ident.table_id = table_id;
2512                    self.catalog_manager
2513                        .register_table_sync(catalog::RegisterTableRequest {
2514                            catalog: table_info.catalog_name.clone(),
2515                            schema: table_info.schema_name.clone(),
2516                            table_name: table_info.name.clone(),
2517                            table_id,
2518                            table: table::dist_table::DistTable::table(Arc::new(table_info)),
2519                        })
2520                        .map_err(BoxedError::new)
2521                        .context(common_meta::error::ExternalSnafu)?;
2522                    Ok(SubmitDdlTaskResponse {
2523                        key: Vec::new(),
2524                        table_ids: vec![table_id],
2525                    })
2526                }
2527                DdlTask::CreateView(task) => {
2528                    let view_id = self
2529                        .next_table_id
2530                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2531                    let mut view_info = task.view_info;
2532                    view_info.ident.table_id = view_id;
2533                    self.catalog_manager
2534                        .register_table_sync(catalog::RegisterTableRequest {
2535                            catalog: task.create_view.catalog_name.clone(),
2536                            schema: task.create_view.schema_name.clone(),
2537                            table_name: task.create_view.view_name.clone(),
2538                            table_id: view_id,
2539                            table: table::dist_table::DistTable::table(Arc::new(view_info)),
2540                        })
2541                        .map_err(BoxedError::new)
2542                        .context(common_meta::error::ExternalSnafu)?;
2543                    Ok(SubmitDdlTaskResponse {
2544                        key: Vec::new(),
2545                        table_ids: vec![view_id],
2546                    })
2547                }
2548                DdlTask::DropView(task) => {
2549                    self.catalog_manager
2550                        .deregister_table_sync(catalog::DeregisterTableRequest {
2551                            catalog: task.catalog.clone(),
2552                            schema: task.schema.clone(),
2553                            table_name: task.view.clone(),
2554                        })
2555                        .map_err(BoxedError::new)
2556                        .context(common_meta::error::ExternalSnafu)?;
2557                    Ok(SubmitDdlTaskResponse::default())
2558                }
2559                other => common_meta::error::UnsupportedSnafu {
2560                    operation: format!("mock submit_ddl_task: {other:?}"),
2561                }
2562                .fail(),
2563            }
2564        }
2565
2566        async fn migrate_region(
2567            &self,
2568            _ctx: &ExecutorContext,
2569            _request: MigrateRegionRequest,
2570        ) -> common_meta::error::Result<MigrateRegionResponse> {
2571            common_meta::error::UnsupportedSnafu {
2572                operation: "migrate_region",
2573            }
2574            .fail()
2575        }
2576
2577        async fn reconcile(
2578            &self,
2579            _ctx: &ExecutorContext,
2580            _request: ReconcileRequest,
2581        ) -> common_meta::error::Result<ReconcileResponse> {
2582            common_meta::error::UnsupportedSnafu {
2583                operation: "reconcile",
2584            }
2585            .fail()
2586        }
2587
2588        async fn query_procedure_state(
2589            &self,
2590            _ctx: &ExecutorContext,
2591            _pid: &str,
2592        ) -> common_meta::error::Result<ProcedureStateResponse> {
2593            common_meta::error::UnsupportedSnafu {
2594                operation: "query_procedure_state",
2595            }
2596            .fail()
2597        }
2598
2599        async fn list_procedures(
2600            &self,
2601            _ctx: &ExecutorContext,
2602        ) -> common_meta::error::Result<ProcedureDetailResponse> {
2603            common_meta::error::UnsupportedSnafu {
2604                operation: "list_procedures",
2605            }
2606            .fail()
2607        }
2608    }
2609
2610    fn test_cache_registry(
2611        kv_backend: common_meta::kv_backend::KvBackendRef,
2612    ) -> TestResult<common_meta::cache::LayeredCacheRegistryRef> {
2613        Ok(Arc::new(
2614            cache::with_default_composite_cache_registry(
2615                LayeredCacheRegistryBuilder::default()
2616                    .add_cache_registry(cache::build_fundamental_cache_registry(kv_backend)),
2617            )
2618            .context(BuildCacheRegistrySnafu)?
2619            .build(),
2620        ))
2621    }
2622
2623    fn test_table_info(table_id: u32, table_name: &str) -> TestResult<TableInfo> {
2624        let schema = Arc::new(GtSchema::new(vec![
2625            ColumnSchema::new("id", ConcreteDataType::int32_datatype(), false),
2626            ColumnSchema::new(
2627                "ts",
2628                ConcreteDataType::timestamp_millisecond_datatype(),
2629                false,
2630            )
2631            .with_time_index(true),
2632        ]));
2633        let table_meta = TableMetaBuilder::empty()
2634            .schema(schema)
2635            .primary_key_indices(vec![0])
2636            .value_indices(vec![1])
2637            .next_column_id(1024)
2638            .build()
2639            .with_context(|_| BuildTableMetaSnafu {
2640                table_name: table_name.to_string(),
2641            })?;
2642
2643        TableInfoBuilder::new(table_name, table_meta)
2644            .table_id(table_id)
2645            .build()
2646            .with_context(|_| BuildTableInfoSnafu {
2647                table_name: table_name.to_string(),
2648            })
2649    }
2650
2651    fn test_table(table_id: u32, table_name: &str) -> TestResult<table::TableRef> {
2652        let table_info = test_table_info(table_id, table_name)?;
2653        Ok(EmptyTable::from_table_info(&table_info))
2654    }
2655
2656    fn test_physical_table(table_id: u32, table_name: &str) -> TestResult<table::TableRef> {
2657        let mut table_info = test_table_info(table_id, table_name)?;
2658        table_info
2659            .meta
2660            .options
2661            .extra_options
2662            .insert(PHYSICAL_TABLE_METADATA_KEY.to_string(), String::new());
2663        Ok(EmptyTable::from_table_info(&table_info))
2664    }
2665
2666    fn test_logical_table(table_id: u32, table_name: &str) -> TestResult<table::TableRef> {
2667        let mut table_info = test_table_info(table_id, table_name)?;
2668        table_info.meta.engine = METRIC_ENGINE_NAME.to_string();
2669        table_info.meta.options.extra_options.insert(
2670            LOGICAL_TABLE_METADATA_KEY.to_string(),
2671            "physical_metric".to_string(),
2672        );
2673        Ok(EmptyTable::from_table_info(&table_info))
2674    }
2675
2676    fn test_metric_names_table() -> TableRef {
2677        let schema = Arc::new(GtSchema::new(vec![
2678            ColumnSchema::new("table_catalog", ConcreteDataType::string_datatype(), false),
2679            ColumnSchema::new("table_schema", ConcreteDataType::string_datatype(), false),
2680            ColumnSchema::new("table_name", ConcreteDataType::string_datatype(), false),
2681            ColumnSchema::new("engine", ConcreteDataType::string_datatype(), false),
2682            ColumnSchema::new("create_options", ConcreteDataType::string_datatype(), false),
2683        ]));
2684        let columns: Vec<VectorRef> = vec![
2685            Arc::new(StringVector::from(vec!["greptime", "greptime"])),
2686            Arc::new(StringVector::from(vec!["public", "public"])),
2687            Arc::new(StringVector::from(vec!["denied", "target"])),
2688            Arc::new(StringVector::from(vec!["metric", "metric"])),
2689            Arc::new(StringVector::from(vec![
2690                "on_physical_table=physical_metric",
2691                "on_physical_table=physical_metric",
2692            ])),
2693        ];
2694        let record_batch = RecordBatch::new(schema, columns).unwrap();
2695        MemTable::new_with_catalog(
2696            "tables",
2697            record_batch,
2698            2048,
2699            "greptime".to_string(),
2700            "information_schema".to_string(),
2701        )
2702    }
2703
2704    fn test_pipeline_table() -> TableRef {
2705        let schema = Arc::new(GtSchema::new(vec![
2706            ColumnSchema::new("name", ConcreteDataType::string_datatype(), false),
2707            ColumnSchema::new("schema", ConcreteDataType::string_datatype(), false),
2708            ColumnSchema::new("content_type", ConcreteDataType::string_datatype(), false),
2709            ColumnSchema::new("pipeline", ConcreteDataType::string_datatype(), false),
2710            ColumnSchema::new(
2711                "created_at",
2712                ConcreteDataType::timestamp_nanosecond_datatype(),
2713                false,
2714            )
2715            .with_time_index(true),
2716        ]));
2717        let columns: Vec<VectorRef> = vec![
2718            Arc::new(StringVector::from(vec!["pipeline"])),
2719            Arc::new(StringVector::from(vec!["public"])),
2720            Arc::new(StringVector::from(vec!["application/yaml"])),
2721            Arc::new(StringVector::from(vec![
2722                "transform:\n- field: ts\n  type: timestamp, ns\n  index: time\n",
2723            ])),
2724            Arc::new(TimestampNanosecondVector::from_values([1])),
2725        ];
2726        let record_batch = RecordBatch::new(schema, columns).unwrap();
2727        MemTable::new_with_catalog(
2728            "pipelines",
2729            record_batch,
2730            2049,
2731            "greptime".to_string(),
2732            DEFAULT_PRIVATE_SCHEMA_NAME.to_string(),
2733        )
2734    }
2735
2736    struct CancellableExportSource {
2737        schema: GtSchemaRef,
2738        channels: std::sync::Mutex<Option<(oneshot::Sender<()>, oneshot::Receiver<()>)>>,
2739    }
2740
2741    impl DataSource for CancellableExportSource {
2742        fn get_stream(
2743            &self,
2744            _request: ScanRequest,
2745        ) -> std::result::Result<SendableRecordBatchStream, BoxedError> {
2746            let (started, release) = self.channels.lock().unwrap().take().unwrap();
2747            let schema = self.schema.clone();
2748            let stream = futures::stream::once(async move {
2749                started.send(()).unwrap();
2750                release.await.unwrap();
2751                Ok(RecordBatch::new_empty(schema))
2752            });
2753            Ok(Box::pin(RecordBatchStreamWrapper::new(
2754                self.schema.clone(),
2755                Box::pin(stream),
2756            )))
2757        }
2758    }
2759
2760    #[tokio::test]
2761    async fn test_metric_export_http_timeout_cancels_ordinary_source() {
2762        let destination = common_test_util::temp_dir::create_temp_dir("metric_export_timeout");
2763        let (started_tx, started_rx) = oneshot::channel();
2764        let (mut release_tx, release_rx) = oneshot::channel();
2765        let info = test_table_info(1024, "source").unwrap();
2766        let source = Arc::new(Table::new(
2767            Arc::new(info.clone()),
2768            FilterPushDownType::Unsupported,
2769            Arc::new(CancellableExportSource {
2770                schema: info.meta.schema.clone(),
2771                channels: std::sync::Mutex::new(Some((started_tx, release_rx))),
2772            }),
2773        ));
2774        let catalog = catalog::memory::MemoryCatalogManager::new_with_table(source);
2775        let kv = Arc::new(MemoryKvBackend::new());
2776        let instance = FrontendBuilder::new(
2777            FrontendOptions::default(),
2778            kv.clone(),
2779            test_cache_registry(kv).unwrap(),
2780            catalog,
2781            Arc::new(client::client_manager::NodeClients::default()),
2782            Arc::new(NoopProcedureExecutor),
2783            Arc::new(ProcessManager::new("export-timeout".into(), None)),
2784        )
2785        .with_local_file_access(
2786            common_datasource::object_store::LocalFileAccess::sandboxed(destination.path())
2787                .unwrap(),
2788        )
2789        .try_build()
2790        .await
2791        .unwrap();
2792        let server = servers::http::HttpServerBuilder::new(servers::http::HttpOptions {
2793            timeout: Duration::from_secs(2),
2794            ..Default::default()
2795        })
2796        .with_sql_handler(Arc::new(instance))
2797        .build();
2798        let app = server.build(server.make_app()).unwrap();
2799        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
2800        let addr = listener.local_addr().unwrap();
2801        let server_task = tokio::spawn(async move {
2802            axum::serve(listener, app).await.unwrap();
2803        });
2804        let uri = reqwest::Url::from_directory_path(destination.path()).unwrap();
2805        let sql = format!(
2806            "COPY DATABASE greptime.public TO '{uri}' WITH (experimental_metric_export='true')"
2807        );
2808        let request = tokio::spawn(async move {
2809            reqwest::Client::builder()
2810                .no_proxy()
2811                .build()
2812                .unwrap()
2813                .post(format!("http://{addr}/v1/sql"))
2814                .form(&[("sql", sql)])
2815                .send()
2816                .await
2817                .unwrap()
2818        });
2819        tokio::time::timeout(Duration::from_secs(5), started_rx)
2820            .await
2821            .unwrap()
2822            .unwrap();
2823        let response = request.await.unwrap();
2824        assert_eq!(response.status(), reqwest::StatusCode::REQUEST_TIMEOUT);
2825        tokio::time::timeout(Duration::from_secs(5), release_tx.closed())
2826            .await
2827            .unwrap();
2828        assert!(!destination.path().join("source.parquet").exists());
2829        server_task.abort();
2830    }
2831
2832    fn pending_table(
2833        table_id: u32,
2834        table_name: &str,
2835        polled_tx: oneshot::Sender<()>,
2836    ) -> TestResult<table::TableRef> {
2837        let table_info = test_table_info(table_id, table_name)?;
2838        let data_source = Arc::new(PendingDataSource {
2839            schema: table_info.meta.schema.clone(),
2840            polled_tx: std::sync::Mutex::new(Some(polled_tx)),
2841        });
2842
2843        Ok(Arc::new(Table::new(
2844            Arc::new(table_info),
2845            FilterPushDownType::Unsupported,
2846            data_source,
2847        )))
2848    }
2849
2850    async fn test_instance_with_tables(
2851        source_table: TableRef,
2852        target_table: TableRef,
2853    ) -> TestResult<Instance> {
2854        test_instance_with_plugins(source_table, target_table, Plugins::new()).await
2855    }
2856
2857    async fn test_instance_with_insert_select_interceptor(
2858        interceptor: SqlQueryInterceptorRef<Error>,
2859    ) -> TestResult<Instance> {
2860        let plugins = Plugins::new();
2861        plugins.insert::<SqlQueryInterceptorRef<Error>>(interceptor);
2862
2863        test_instance_with_plugins(
2864            test_table(1024, "source")?,
2865            test_table(1025, "target")?,
2866            plugins,
2867        )
2868        .await
2869    }
2870
2871    async fn test_instance_with_plugins(
2872        source_table: TableRef,
2873        target_table: TableRef,
2874        plugins: Plugins,
2875    ) -> TestResult<Instance> {
2876        test_instance_with_plugins_and_metric_names(source_table, target_table, plugins, None).await
2877    }
2878
2879    async fn test_instance_with_plugins_and_metric_names(
2880        source_table: TableRef,
2881        target_table: TableRef,
2882        plugins: Plugins,
2883        metric_names_table: Option<TableRef>,
2884    ) -> TestResult<Instance> {
2885        let catalog_manager = catalog::memory::MemoryCatalogManager::new_with_table(source_table);
2886        test_instance_with_catalog_manager(
2887            catalog_manager,
2888            target_table,
2889            plugins,
2890            metric_names_table,
2891            Arc::new(NoopProcedureExecutor),
2892        )
2893        .await
2894    }
2895
2896    /// Builds a test frontend `Instance` over the given (already source-registered)
2897    /// catalog manager, completing DDL tasks through `procedure_executor`.
2898    async fn test_instance_with_catalog_manager(
2899        catalog_manager: Arc<catalog::memory::MemoryCatalogManager>,
2900        target_table: TableRef,
2901        plugins: Plugins,
2902        metric_names_table: Option<TableRef>,
2903        procedure_executor: ProcedureExecutorRef,
2904    ) -> TestResult<Instance> {
2905        let kv_backend = Arc::new(MemoryKvBackend::new());
2906        let process_manager = Arc::new(ProcessManager::new("test-frontend".to_string(), None));
2907        let target_table_name = "target";
2908        catalog_manager
2909            .register_table_sync(catalog::RegisterTableRequest {
2910                catalog: "greptime".to_string(),
2911                schema: "public".to_string(),
2912                table_name: target_table_name.to_string(),
2913                table_id: 1025,
2914                table: target_table,
2915            })
2916            .with_context(|_| RegisterTableSnafu {
2917                table_name: target_table_name.to_string(),
2918            })?;
2919        if let Some(table) = metric_names_table {
2920            catalog_manager
2921                .deregister_table_sync(catalog::DeregisterTableRequest {
2922                    catalog: "greptime".to_string(),
2923                    schema: "information_schema".to_string(),
2924                    table_name: "tables".to_string(),
2925                })
2926                .unwrap();
2927            catalog_manager
2928                .register_table_sync(catalog::RegisterTableRequest {
2929                    catalog: "greptime".to_string(),
2930                    schema: "information_schema".to_string(),
2931                    table_name: "tables".to_string(),
2932                    table_id: 2048,
2933                    table,
2934                })
2935                .unwrap();
2936        }
2937        catalog_manager.register_process_list_table(process_manager.clone());
2938
2939        let cache_registry = test_cache_registry(kv_backend.clone())?;
2940
2941        FrontendBuilder::new(
2942            FrontendOptions::default(),
2943            kv_backend,
2944            cache_registry,
2945            catalog_manager,
2946            Arc::new(client::client_manager::NodeClients::default()),
2947            procedure_executor,
2948            process_manager,
2949        )
2950        .with_plugin(plugins)
2951        .try_build()
2952        .await
2953        .context(BuildFrontendSnafu)
2954    }
2955
2956    async fn execute_one_sql(
2957        instance: &Instance,
2958        sql: &str,
2959        query_ctx: QueryContextRef,
2960    ) -> TestResult<Output> {
2961        let mut results = instance.do_query_inner(sql, query_ctx).await;
2962        ensure!(
2963            results.len() == 1,
2964            UnexpectedOutputCountSnafu {
2965                sql: sql.to_string(),
2966                actual: results.len(),
2967            }
2968        );
2969        results.remove(0).with_context(|_| ExecuteSqlSnafu {
2970            sql: sql.to_string(),
2971        })
2972    }
2973
2974    fn assert_permission_denied<T>(result: servers::error::Result<T>) {
2975        let err = match result {
2976            Ok(_) => panic!("request should be rejected"),
2977            Err(err) => err,
2978        };
2979        assert_eq!(StatusCode::PermissionDenied, err.status_code());
2980    }
2981
2982    fn assert_action_checked(
2983        checker: &RejectEndpointPermissionChecker,
2984        action: PermissionAction,
2985        targets: Option<PermissionTableTargets>,
2986    ) {
2987        assert_eq!(CheckedAction { action, targets }, checker.take_check());
2988    }
2989
2990    #[tokio::test]
2991    async fn database_export_authorizes_all_tables_before_preparation() -> TestResult<()> {
2992        struct ExportAcl(std::sync::Mutex<Vec<PermissionTableTargets>>);
2993        impl PermissionChecker for ExportAcl {
2994            fn check_permission(
2995                &self,
2996                _: UserInfoRef,
2997                req: PermissionReq,
2998            ) -> auth::error::Result<PermissionResp> {
2999                assert!(matches!(
3000                    req,
3001                    PermissionReq::SqlStatement(Statement::Copy(
3002                        sql::statements::copy::Copy::CopyDatabase(
3003                            sql::statements::copy::CopyDatabase::To(_)
3004                        )
3005                    ))
3006                ));
3007                Ok(PermissionResp::Allow)
3008            }
3009            fn check_permission_with_table_targets(
3010                &self,
3011                _: UserInfoRef,
3012                req: PermissionReq,
3013                targets: PermissionTableTargets,
3014            ) -> auth::error::Result<PermissionResp> {
3015                self.0.lock().unwrap().push(targets.clone());
3016                let PermissionTableTargets::Resolved(tables) = targets else {
3017                    panic!("unresolved export")
3018                };
3019                if req.is_readonly() {
3020                    // The first table is readable; the later table has only write access.
3021                    Ok(if tables.iter().any(|t| t.table == "target") {
3022                        PermissionResp::Reject
3023                    } else {
3024                        PermissionResp::Allow
3025                    })
3026                } else {
3027                    self.check_permission(QueryContext::arc().current_user(), req)
3028                }
3029            }
3030        }
3031        let checker = Arc::new(ExportAcl(Default::default()));
3032        let plugins = Plugins::new();
3033        plugins.insert::<PermissionCheckerRef>(checker.clone());
3034        let instance = test_instance_with_plugins(
3035            test_logical_table(1024, "source")?,
3036            test_table(1025, "target")?,
3037            plugins,
3038        )
3039        .await?;
3040        let req = table::requests::CopyDatabaseRequest {
3041            catalog_name: "greptime".into(),
3042            schema_name: "public".into(),
3043            location: "invalid-destination".into(),
3044            with: Default::default(),
3045            connection: Default::default(),
3046            time_range: None,
3047        };
3048        let result = instance
3049            .export_database_for_test(
3050                req.clone(),
3051                None,
3052                &tokio_util::sync::CancellationToken::new(),
3053                QueryContext::arc(),
3054            )
3055            .await;
3056        assert!(matches!(result, Err(Error::Permission { .. })));
3057        let expected = PermissionTableTargets::resolved(vec![
3058            PermissionTableTarget::new("greptime", "public", "source"),
3059            PermissionTableTarget::new("greptime", "public", "target"),
3060        ]);
3061        assert_eq!(*checker.0.lock().unwrap(), vec![expected.clone(), expected]);
3062        // An empty selection must still check the operation privilege.
3063        instance
3064            .plugins
3065            .map_mut::<PermissionCheckerRef, _, _>(|checker| {
3066                *checker.unwrap() = Arc::new(WriteOnlyPermissionChecker)
3067            });
3068        let result = instance
3069            .export_database_for_test(
3070                req,
3071                Some(&[]),
3072                &tokio_util::sync::CancellationToken::new(),
3073                QueryContext::arc(),
3074            )
3075            .await;
3076        assert!(matches!(result, Err(Error::Permission { .. })));
3077        Ok(())
3078    }
3079
3080    #[tokio::test]
3081    async fn test_prom_remote_read_with_custom_timestamp_and_value_columns() -> TestResult<()> {
3082        let schema = Arc::new(GtSchema::new(vec![
3083            ColumnSchema::new(
3084                "custom_ts",
3085                ConcreteDataType::timestamp_millisecond_datatype(),
3086                false,
3087            )
3088            .with_time_index(true),
3089            ColumnSchema::new("custom_value", ConcreteDataType::float64_datatype(), false),
3090        ]));
3091        let recordbatch = RecordBatch::new(
3092            schema,
3093            vec![
3094                Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as VectorRef,
3095                Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as VectorRef,
3096            ],
3097        )
3098        .unwrap();
3099        let instance = test_instance_with_tables(
3100            MemTable::table("custom_metric", recordbatch),
3101            test_table(1025, "target")?,
3102        )
3103        .await?;
3104
3105        let response = PromStoreProtocolHandler::read(
3106            &instance,
3107            ReadRequest {
3108                queries: vec![RemoteQuery {
3109                    start_timestamp_ms: 1500,
3110                    end_timestamp_ms: 2500,
3111                    matchers: vec![LabelMatcher {
3112                        r#type: PromMatcherType::Eq as i32,
3113                        name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3114                        value: "custom_metric".to_string(),
3115                    }],
3116                    ..Default::default()
3117                }],
3118                ..Default::default()
3119            },
3120            test_query_ctx(1),
3121        )
3122        .await
3123        .unwrap();
3124        let body = servers::prom_store::snappy_decompress(&response.body).unwrap();
3125        let response = ReadResponse::decode(body.as_slice()).unwrap();
3126
3127        assert_eq!(1, response.results.len());
3128        assert_eq!(1, response.results[0].timeseries.len());
3129        let timeseries = &response.results[0].timeseries[0];
3130        assert_eq!(
3131            vec![Label {
3132                name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3133                value: "custom_metric".to_string(),
3134            }],
3135            timeseries.labels
3136        );
3137        assert_eq!(
3138            vec![Sample {
3139                value: 2.0,
3140                timestamp: 2000,
3141            }],
3142            timeseries.samples
3143        );
3144
3145        Ok(())
3146    }
3147
3148    #[tokio::test]
3149    async fn test_prom_remote_read_prefers_default_value_column() -> TestResult<()> {
3150        let schema = Arc::new(GtSchema::new(vec![
3151            ColumnSchema::new(
3152                "custom_ts",
3153                ConcreteDataType::timestamp_millisecond_datatype(),
3154                false,
3155            )
3156            .with_time_index(true),
3157            ColumnSchema::new("extra_field", ConcreteDataType::string_datatype(), false),
3158            ColumnSchema::new(
3159                greptime_value(),
3160                ConcreteDataType::float64_datatype(),
3161                false,
3162            ),
3163        ]));
3164        let recordbatch = RecordBatch::new(
3165            schema,
3166            vec![
3167                Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as VectorRef,
3168                Arc::new(StringVector::from(vec!["99", "99", "99"])) as VectorRef,
3169                Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as VectorRef,
3170            ],
3171        )
3172        .unwrap();
3173        let instance = test_instance_with_tables(
3174            MemTable::table("multi_field_metric", recordbatch),
3175            test_table(1025, "target")?,
3176        )
3177        .await?;
3178
3179        let response = PromStoreProtocolHandler::read(
3180            &instance,
3181            ReadRequest {
3182                queries: vec![RemoteQuery {
3183                    start_timestamp_ms: 1500,
3184                    end_timestamp_ms: 2500,
3185                    matchers: vec![LabelMatcher {
3186                        r#type: PromMatcherType::Eq as i32,
3187                        name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3188                        value: "multi_field_metric".to_string(),
3189                    }],
3190                    ..Default::default()
3191                }],
3192                ..Default::default()
3193            },
3194            test_query_ctx(1),
3195        )
3196        .await
3197        .unwrap();
3198        let body = servers::prom_store::snappy_decompress(&response.body).unwrap();
3199        let response = ReadResponse::decode(body.as_slice()).unwrap();
3200
3201        assert_eq!(1, response.results.len());
3202        assert_eq!(1, response.results[0].timeseries.len());
3203        let timeseries = &response.results[0].timeseries[0];
3204        assert_eq!(
3205            vec![
3206                Label {
3207                    name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3208                    value: "multi_field_metric".to_string(),
3209                },
3210                Label {
3211                    name: "extra_field".to_string(),
3212                    value: "99".to_string(),
3213                },
3214            ],
3215            timeseries.labels
3216        );
3217        assert_eq!(
3218            vec![Sample {
3219                value: 2.0,
3220                timestamp: 2000,
3221            }],
3222            timeseries.samples
3223        );
3224
3225        Ok(())
3226    }
3227
3228    #[tokio::test]
3229    async fn test_prom_remote_read_rejects_ambiguous_value_columns() -> TestResult<()> {
3230        let schema = Arc::new(GtSchema::new(vec![
3231            ColumnSchema::new(
3232                "custom_ts",
3233                ConcreteDataType::timestamp_millisecond_datatype(),
3234                false,
3235            )
3236            .with_time_index(true),
3237            ColumnSchema::new("field_a", ConcreteDataType::float64_datatype(), false),
3238            ColumnSchema::new("field_b", ConcreteDataType::float64_datatype(), false),
3239        ]));
3240        let recordbatch = RecordBatch::new(
3241            schema,
3242            vec![
3243                Arc::new(TimestampMillisecondVector::from_vec(vec![1000])) as VectorRef,
3244                Arc::new(Float64Vector::from_vec(vec![1.0])) as VectorRef,
3245                Arc::new(Float64Vector::from_vec(vec![2.0])) as VectorRef,
3246            ],
3247        )
3248        .unwrap();
3249        let instance = test_instance_with_tables(
3250            MemTable::table("ambiguous_metric", recordbatch),
3251            test_table(1025, "target")?,
3252        )
3253        .await?;
3254
3255        let err = PromStoreProtocolHandler::read(
3256            &instance,
3257            ReadRequest {
3258                queries: vec![RemoteQuery {
3259                    matchers: vec![LabelMatcher {
3260                        r#type: PromMatcherType::Eq as i32,
3261                        name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3262                        value: "ambiguous_metric".to_string(),
3263                    }],
3264                    ..Default::default()
3265                }],
3266                ..Default::default()
3267            },
3268            test_query_ctx(1),
3269        )
3270        .await
3271        .err()
3272        .expect("ambiguous value columns should fail remote read");
3273
3274        assert_eq!(StatusCode::InvalidArguments, err.status_code());
3275        assert!(format!("{err:?}").contains("Ambiguous value column"));
3276
3277        Ok(())
3278    }
3279
3280    #[tokio::test]
3281    async fn test_restricted_endpoint_handlers_check_permissions() -> TestResult<()> {
3282        let checker = Arc::new(RejectEndpointPermissionChecker::default());
3283        let plugins = Plugins::new();
3284        plugins.insert::<PermissionCheckerRef>(checker.clone());
3285        let instance = test_instance_with_plugins(
3286            test_table(1024, "denied")?,
3287            test_table(1025, "target")?,
3288            plugins,
3289        )
3290        .await?;
3291        let mut ctx = test_query_ctx(1);
3292        Arc::get_mut(&mut ctx).unwrap().set_extension(
3293            servers::http::jaeger::JAEGER_QUERY_TABLE_NAME_KEY,
3294            "denied".to_string(),
3295        );
3296        let jaeger_targets = Some(PermissionTableTargets::resolved(vec![
3297            PermissionTableTarget::new("greptime", "public", "denied"),
3298        ]));
3299
3300        assert_permission_denied(JaegerQueryHandler::get_services(&instance, ctx.clone()).await);
3301        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
3302        assert_permission_denied(
3303            JaegerQueryHandler::get_operations(&instance, ctx.clone(), "service", None).await,
3304        );
3305        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
3306        assert_permission_denied(
3307            JaegerQueryHandler::get_trace(&instance, ctx.clone(), "trace", None, None, None).await,
3308        );
3309        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
3310        assert_permission_denied(
3311            JaegerQueryHandler::find_traces(
3312                &instance,
3313                ctx.clone(),
3314                servers::http::jaeger::QueryTraceParams {
3315                    service_name: Some("service".to_string()),
3316                    ..Default::default()
3317                },
3318            )
3319            .await,
3320        );
3321        assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets);
3322
3323        assert_permission_denied(
3324            PipelineHandler::get_pipeline_str(&instance, "pipeline", None, ctx.clone()).await,
3325        );
3326        assert_action_checked(&checker, PIPELINE_QUERY, None);
3327        assert_permission_denied(
3328            PipelineHandler::insert_pipeline(
3329                &instance,
3330                "pipeline",
3331                "application/yaml",
3332                "",
3333                ctx.clone(),
3334            )
3335            .await,
3336        );
3337        assert_action_checked(&checker, PIPELINE_INSERT, None);
3338        assert_permission_denied(
3339            PipelineHandler::delete_pipeline(&instance, "pipeline", None, ctx.clone()).await,
3340        );
3341        assert_action_checked(&checker, PIPELINE_DELETE, None);
3342        let app = axum::Router::new()
3343            .route(
3344                "/pipelines/_dryrun",
3345                axum::routing::post(servers::http::event::pipeline_dryrun),
3346            )
3347            .with_state(servers::http::event::LogState {
3348                log_handler: Arc::new(instance.clone()),
3349                log_validator: None,
3350                ingest_interceptor: None,
3351            })
3352            .layer(axum::Extension((*ctx).clone()));
3353        let response = app
3354            .oneshot(
3355                axum::http::Request::post("/pipelines/_dryrun")
3356                    .header("content-type", "application/json")
3357                    .body(axum::body::Body::from("{}"))
3358                    .unwrap(),
3359            )
3360            .await
3361            .unwrap();
3362        assert_eq!(axum::http::StatusCode::FORBIDDEN, response.status());
3363        assert_action_checked(&checker, PIPELINE_QUERY, None);
3364
3365        assert_permission_denied(
3366            DashboardHandler::save(&instance, "dashboard", "{}", ctx.clone()).await,
3367        );
3368        assert_action_checked(&checker, DASHBOARD_SAVE, None);
3369        assert_permission_denied(DashboardHandler::list(&instance, ctx.clone()).await);
3370        assert_action_checked(&checker, DASHBOARD_QUERY, None);
3371        assert_permission_denied(
3372            DashboardHandler::delete(&instance, "dashboard", ctx.clone()).await,
3373        );
3374        assert_action_checked(&checker, DASHBOARD_DELETE, None);
3375
3376        Ok(())
3377    }
3378
3379    #[tokio::test]
3380    async fn test_write_only_ingestion_loads_named_pipeline() -> TestResult<()> {
3381        let plugins = Plugins::new();
3382        plugins.insert::<PermissionCheckerRef>(Arc::new(WriteOnlyPermissionChecker));
3383        let instance = test_instance_with_plugins(
3384            test_table(1024, "source")?,
3385            test_table(1025, "target")?,
3386            plugins,
3387        )
3388        .await?;
3389        instance
3390            .catalog_manager()
3391            .as_any()
3392            .downcast_ref::<catalog::memory::MemoryCatalogManager>()
3393            .unwrap()
3394            .register_table_sync(catalog::RegisterTableRequest {
3395                catalog: "greptime".to_string(),
3396                schema: DEFAULT_PRIVATE_SCHEMA_NAME.to_string(),
3397                table_name: "pipelines".to_string(),
3398                table_id: 2049,
3399                table: test_pipeline_table(),
3400            })
3401            .with_context(|_| RegisterTableSnafu {
3402                table_name: "pipelines".to_string(),
3403            })?;
3404        let ctx = test_query_ctx(1);
3405        let handler: PipelineHandlerRef = Arc::new(instance.clone());
3406        let hits = CACHE_HIT.with_label_values(&["pipeline_table"]);
3407        let misses = CACHE_MISS.with_label_values(&["pipeline_table"]);
3408        let before = (hits.get(), misses.get());
3409
3410        handler
3411            .get_pipeline("pipeline", None, ctx.clone())
3412            .await
3413            .unwrap();
3414        assert_eq!((hits.get(), misses.get()), (before.0 + 1, before.1 + 1));
3415        handler
3416            .get_pipeline("pipeline", None, ctx.clone())
3417            .await
3418            .unwrap();
3419        assert_eq!((hits.get(), misses.get()), (before.0 + 3, before.1 + 1));
3420        assert_permission_denied(
3421            PipelineHandler::get_pipeline_str(&instance, "pipeline", None, ctx.clone()).await,
3422        );
3423
3424        let app = axum::Router::new()
3425            .route(
3426                "/pipelines/_dryrun",
3427                axum::routing::post(servers::http::event::pipeline_dryrun),
3428            )
3429            .with_state(servers::http::event::LogState {
3430                log_handler: handler,
3431                log_validator: None,
3432                ingest_interceptor: None,
3433            })
3434            .layer(axum::Extension((*ctx).clone()));
3435        let response = app
3436            .oneshot(
3437                axum::http::Request::post("/pipelines/_dryrun")
3438                    .header("content-type", "application/json")
3439                    .body(axum::body::Body::from("{}"))
3440                    .unwrap(),
3441            )
3442            .await
3443            .unwrap();
3444        assert_eq!(axum::http::StatusCode::FORBIDDEN, response.status());
3445
3446        Ok(())
3447    }
3448
3449    #[tokio::test]
3450    async fn test_write_only_grpc_sql_is_checked_after_parsing() -> TestResult<()> {
3451        let plugins = Plugins::new();
3452        plugins.insert::<PermissionCheckerRef>(Arc::new(WriteOnlyPermissionChecker));
3453        let instance = test_instance_with_plugins(
3454            test_table(1024, "source")?,
3455            test_table(1025, "target")?,
3456            plugins,
3457        )
3458        .await?;
3459
3460        let insert = Request::Query(api::v1::QueryRequest {
3461            query: Some(Query::Sql(
3462                "INSERT INTO target SELECT * FROM source".to_string(),
3463            )),
3464        });
3465        servers::query_handler::grpc::GrpcQueryHandler::do_query(
3466            &instance,
3467            insert,
3468            QueryContext::arc(),
3469        )
3470        .await
3471        .unwrap();
3472
3473        let select = Request::Query(api::v1::QueryRequest {
3474            query: Some(Query::Sql("SELECT * FROM source".to_string())),
3475        });
3476        assert_permission_denied(
3477            servers::query_handler::grpc::GrpcQueryHandler::do_query(
3478                &instance,
3479                select,
3480                QueryContext::arc(),
3481            )
3482            .await,
3483        );
3484
3485        Ok(())
3486    }
3487
3488    #[tokio::test]
3489    async fn test_target_independent_checker_skips_target_resolution() -> TestResult<()> {
3490        let physical_table = "physical_metric";
3491        let checker = Arc::new(TargetIndependentPermissionChecker::default());
3492        let plugins = Plugins::new();
3493        plugins.insert::<PermissionCheckerRef>(checker.clone());
3494        let instance = test_instance_with_plugins(
3495            test_physical_table(1024, physical_table)?,
3496            test_table(1025, "target")?,
3497            plugins,
3498        )
3499        .await?;
3500
3501        let ctx = test_query_ctx(1);
3502        let physical_target = PermissionTableTarget::new("greptime", "public", physical_table);
3503        assert_eq!(
3504            PermissionTableTargets::Resolved(vec![physical_target.clone()]),
3505            instance
3506                .resolve_query_permission_targets(
3507                    PermissionTableTargets::resolved(vec![physical_target]),
3508                    &ctx,
3509                )
3510                .await
3511                .unwrap()
3512        );
3513        assert_eq!(
3514            vec![physical_table.to_string(), "target".to_string()],
3515            PrometheusHandler::filter_metadata_metric_names(
3516                &instance,
3517                vec![physical_table.to_string(), "target".to_string()],
3518                "public",
3519                &ctx,
3520            )
3521            .await
3522            .unwrap()
3523        );
3524        assert_eq!(1, checker.checks.load(atomic::Ordering::Relaxed));
3525
3526        Ok(())
3527    }
3528
3529    #[tokio::test]
3530    async fn test_query_permission_targets_are_deduplicated() -> TestResult<()> {
3531        let plugins = Plugins::new();
3532        plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
3533        let instance = test_instance_with_plugins(
3534            test_table(1024, "source")?,
3535            test_table(1025, "target")?,
3536            plugins,
3537        )
3538        .await?;
3539        let ctx = test_query_ctx(1);
3540        let target = PermissionTableTarget::new("greptime", "public", "target");
3541
3542        assert_eq!(
3543            PermissionTableTargets::Resolved(vec![target.clone()]),
3544            instance
3545                .resolve_query_permission_targets(
3546                    PermissionTableTargets::resolved(vec![target.clone(), target]),
3547                    &ctx,
3548                )
3549                .await
3550                .unwrap()
3551        );
3552
3553        Ok(())
3554    }
3555
3556    #[tokio::test]
3557    async fn test_physical_query_targets_fail_closed() -> TestResult<()> {
3558        let physical_table = "physical_metric";
3559        let plugins = Plugins::new();
3560        plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
3561        let instance = test_instance_with_plugins(
3562            test_physical_table(1024, physical_table)?,
3563            test_table(1025, "target")?,
3564            plugins,
3565        )
3566        .await?;
3567
3568        let ctx = test_query_ctx(1);
3569        let logical_target = PermissionTableTarget::new("greptime", "public", "target");
3570        assert_eq!(
3571            PermissionTableTargets::Resolved(vec![logical_target.clone()]),
3572            instance
3573                .resolve_query_permission_targets(
3574                    PermissionTableTargets::resolved(vec![logical_target.clone()]),
3575                    &ctx,
3576                )
3577                .await
3578                .unwrap()
3579        );
3580        let physical_target = PermissionTableTarget::new("greptime", "public", physical_table);
3581        assert_eq!(
3582            PermissionTableTargets::Unresolved,
3583            instance
3584                .resolve_query_permission_targets(
3585                    PermissionTableTargets::resolved(
3586                        vec![logical_target, physical_target.clone(),]
3587                    ),
3588                    &ctx,
3589                )
3590                .await
3591                .unwrap()
3592        );
3593        assert_eq!(
3594            vec!["target".to_string()],
3595            PrometheusHandler::filter_metadata_metric_names(
3596                &instance,
3597                vec!["target".to_string(), "denied".to_string()],
3598                "public",
3599                &ctx,
3600            )
3601            .await
3602            .unwrap()
3603        );
3604
3605        let query = PromQuery {
3606            query: physical_table.to_string(),
3607            ..Default::default()
3608        };
3609        let err = PrometheusHandler::check_query_target_permission(
3610            &instance,
3611            PermissionTableTargets::resolved(vec![physical_target]),
3612            &ctx,
3613        )
3614        .await
3615        .unwrap_err();
3616        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3617        let err = PrometheusHandler::check_query_permission(
3618            &instance,
3619            std::slice::from_ref(&query),
3620            &ctx,
3621        )
3622        .await
3623        .unwrap_err();
3624        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3625        let err = PrometheusHandler::do_query(&instance, &query, ctx.clone())
3626            .await
3627            .unwrap_err();
3628        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3629
3630        for sql in [
3631            "SELECT * FROM physical_metric",
3632            "TQL EVAL (0, 10, '5s') physical_metric",
3633            "INSERT INTO target SELECT * FROM physical_metric",
3634        ] {
3635            let mut results = instance.do_query_inner(sql, ctx.clone()).await;
3636            assert_eq!(1, results.len(), "{sql}");
3637            let err = results.remove(0).unwrap_err();
3638            assert_eq!(StatusCode::PermissionDenied, err.status_code(), "{sql}");
3639        }
3640        let err = LogQueryHandler::query(
3641            &instance,
3642            LogQuery {
3643                table: TableName::new("greptime", "public", physical_table),
3644                ..Default::default()
3645            },
3646            ctx.clone(),
3647        )
3648        .await
3649        .unwrap_err();
3650        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3651        let err = instance
3652            .do_describe_inner(parse_one_sql("SELECT * FROM physical_metric"), ctx.clone())
3653            .await
3654            .unwrap_err();
3655        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3656
3657        let request = ReadRequest {
3658            queries: vec![RemoteQuery {
3659                matchers: vec![LabelMatcher {
3660                    r#type: PromMatcherType::Eq as i32,
3661                    name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3662                    value: physical_table.to_string(),
3663                }],
3664                ..Default::default()
3665            }],
3666            ..Default::default()
3667        };
3668        let Err(err) = PromStoreProtocolHandler::read(&instance, request, ctx.clone()).await else {
3669            panic!("physical remote-read target must be rejected");
3670        };
3671        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3672
3673        let err = PrometheusHandler::query_label_values(
3674            &instance,
3675            physical_table.to_string(),
3676            "host".to_string(),
3677            vec![],
3678            SystemTime::UNIX_EPOCH,
3679            SystemTime::UNIX_EPOCH,
3680            &ctx,
3681        )
3682        .await
3683        .unwrap_err();
3684        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3685
3686        Ok(())
3687    }
3688
3689    #[tokio::test]
3690    async fn test_non_exact_query_discovery_keeps_denied_targets_for_batch_check() -> TestResult<()>
3691    {
3692        let plugins = Plugins::new();
3693        plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
3694        let instance = test_instance_with_plugins_and_metric_names(
3695            test_logical_table(1024, "denied")?,
3696            test_logical_table(1025, "target")?,
3697            plugins,
3698            Some(test_metric_names_table()),
3699        )
3700        .await?;
3701        let ctx = test_query_ctx(1);
3702
3703        let mut metric_names = PrometheusHandler::query_metric_names(
3704            &instance,
3705            vec![Matcher::new(
3706                promql_parser::label::MatchOp::NotEqual,
3707                "__name__",
3708                "",
3709            )],
3710            "public",
3711            &ctx,
3712        )
3713        .await
3714        .unwrap();
3715        metric_names.sort_unstable();
3716        assert_eq!(
3717            vec!["denied".to_string(), "target".to_string()],
3718            metric_names
3719        );
3720
3721        let queries = metric_names
3722            .into_iter()
3723            .map(|query| PromQuery {
3724                query,
3725                ..Default::default()
3726            })
3727            .collect::<Vec<_>>();
3728        let err = PrometheusHandler::check_query_permission(&instance, &queries, &ctx)
3729            .await
3730            .unwrap_err();
3731        assert_eq!(StatusCode::PermissionDenied, err.status_code());
3732
3733        Ok(())
3734    }
3735
3736    #[test]
3737    fn test_should_track_statement_process() {
3738        assert!(should_track_statement_process(&parse_one_sql(
3739            "SELECT * FROM demo"
3740        )));
3741        assert!(should_track_statement_process(&parse_one_sql(
3742            "INSERT INTO demo SELECT * FROM source"
3743        )));
3744        assert!(!should_track_statement_process(&parse_one_sql(
3745            "INSERT INTO demo VALUES (1)"
3746        )));
3747        assert!(!should_track_statement_process(&parse_one_sql(
3748            "INSERT INTO demo VALUES (now())"
3749        )));
3750    }
3751
3752    #[test]
3753    fn test_should_track_plan_process() {
3754        let select_stmt = parse_one_sql("SELECT * FROM demo");
3755        let insert_select_stmt = parse_one_sql("INSERT INTO demo SELECT * FROM source");
3756        let insert_values_stmt = parse_one_sql("INSERT INTO demo VALUES (now())");
3757
3758        let empty_plan = LogicalPlanBuilder::empty(false).build().unwrap();
3759        assert!(should_track_plan_process(Some(&select_stmt), &empty_plan));
3760        assert!(should_track_plan_process(
3761            Some(&insert_select_stmt),
3762            &insert_dml_plan()
3763        ));
3764        assert!(!should_track_plan_process(
3765            Some(&insert_values_stmt),
3766            &insert_dml_plan()
3767        ));
3768        assert!(!should_track_plan_process(None, &insert_dml_plan()));
3769    }
3770
3771    #[tokio::test]
3772    async fn test_insert_select_slow_query() -> TestResult<()> {
3773        let mut instance =
3774            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
3775                .await?;
3776        let recorder = Arc::new(RecordingSlowQueryEventRecorder::default());
3777        instance.slow_query_recorder = recorder.clone();
3778        let sql = "INSERT INTO target SELECT * FROM source";
3779
3780        for prepared in [false, true] {
3781            for (enabled, threshold, should_record) in [
3782                (true, Duration::ZERO, true),
3783                (false, Duration::ZERO, false),
3784                (true, Duration::from_secs(3600), false),
3785            ] {
3786                instance.slow_query_options.enable = enabled;
3787                instance.slow_query_options.threshold = threshold;
3788                let query_ctx = test_query_ctx(42);
3789                let output = if prepared {
3790                    let stmt = parse_one_sql(sql);
3791                    let described = instance
3792                        .do_describe_inner(stmt.clone(), query_ctx.clone())
3793                        .await
3794                        .unwrap()
3795                        .unwrap();
3796                    instance
3797                        .do_exec_plan_inner(described.logical_plan, Some(stmt), query_ctx)
3798                        .await
3799                        .unwrap()
3800                } else {
3801                    execute_one_sql(&instance, sql, query_ctx).await?
3802                };
3803                assert!(matches!(output.data, OutputData::AffectedRows(0)));
3804
3805                let queries = std::mem::take(&mut *recorder.queries.lock().unwrap());
3806                let expected = if should_record { vec![sql] } else { vec![] };
3807                assert_eq!(
3808                    queries, expected,
3809                    "prepared={prepared}, enabled={enabled}, threshold={threshold:?}"
3810                );
3811            }
3812        }
3813
3814        Ok(())
3815    }
3816
3817    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3818    async fn test_insert_select_is_visible_in_show_processlist() -> TestResult<()> {
3819        let insert_sql = "INSERT INTO target SELECT * FROM source";
3820        let (started_tx, mut started_rx) = mpsc::unbounded_channel();
3821        let (finish_tx, finish_rx) = oneshot::channel();
3822        let interceptor = Arc::new(BlockingInsertSelectInterceptor::new(started_tx, finish_rx));
3823        let instance = Arc::new(test_instance_with_insert_select_interceptor(interceptor).await?);
3824
3825        let insert_task = tokio::spawn({
3826            let instance = instance.clone();
3827            async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
3828        });
3829
3830        tokio::time::timeout(Duration::from_secs(5), started_rx.recv())
3831            .await
3832            .context(InsertStartTimeoutSnafu)?
3833            .context(InsertStartChannelClosedSnafu)?;
3834
3835        let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
3836        let process_list = output.data.pretty_print().await;
3837        assert!(
3838            process_list.contains(insert_sql),
3839            "process list did not contain running insert:\n{process_list}"
3840        );
3841
3842        finish_tx
3843            .send(())
3844            .map_err(|_| ReleaseBlockedInsertSnafu.build())?;
3845        insert_task.await.context(InsertTaskPanicSnafu)??;
3846
3847        Ok(())
3848    }
3849
3850    #[tokio::test]
3851    async fn test_show_processlist_catalog_scope() -> TestResult<()> {
3852        let instance =
3853            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
3854                .await?;
3855        let _current_catalog = instance.process_manager().register_query(
3856            "greptime".to_string(),
3857            vec!["public".to_string()],
3858            "current_catalog_query".to_string(),
3859            String::new(),
3860            None,
3861            None,
3862        );
3863        let _other_catalog = instance.process_manager().register_query(
3864            "other".to_string(),
3865            vec!["public".to_string()],
3866            "other_catalog_query".to_string(),
3867            String::new(),
3868            None,
3869            None,
3870        );
3871
3872        for sql in ["SHOW PROCESSLIST", "SHOW FULL PROCESSLIST"] {
3873            let output = execute_one_sql(&instance, sql, test_query_ctx(43)).await?;
3874            let process_list = output.data.pretty_print().await;
3875            assert!(
3876                process_list.contains("current_catalog_query"),
3877                "{process_list}"
3878            );
3879            assert!(
3880                !process_list.contains("other_catalog_query"),
3881                "{process_list}"
3882            );
3883
3884            let admin_ctx = test_query_ctx(44);
3885            admin_ctx.set_current_user(Arc::new(AdminUserInfo));
3886            let output = execute_one_sql(&instance, sql, admin_ctx).await?;
3887            let process_list = output.data.pretty_print().await;
3888            assert!(
3889                process_list.contains("current_catalog_query"),
3890                "{process_list}"
3891            );
3892            assert!(
3893                process_list.contains("other_catalog_query"),
3894                "{process_list}"
3895            );
3896        }
3897
3898        Ok(())
3899    }
3900
3901    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3902    async fn test_kill_query_cancels_insert_select() -> TestResult<()> {
3903        assert_kill_cancels_insert_select("KILL QUERY 4242").await
3904    }
3905
3906    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3907    async fn test_kill_process_id_cancels_insert_select() -> TestResult<()> {
3908        assert_kill_cancels_insert_select("KILL 'test-frontend/4242'").await
3909    }
3910
3911    async fn assert_kill_cancels_insert_select(kill_sql: &str) -> TestResult<()> {
3912        let insert_sql = "INSERT INTO target SELECT * FROM source";
3913        let (source_polled_tx, source_polled_rx) = oneshot::channel();
3914        let instance = Arc::new(
3915            test_instance_with_tables(
3916                pending_table(1024, "source", source_polled_tx)?,
3917                test_table(1025, "target")?,
3918            )
3919            .await?,
3920        );
3921
3922        let insert_task = tokio::spawn({
3923            let instance = instance.clone();
3924            async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
3925        });
3926
3927        tokio::time::timeout(Duration::from_secs(5), source_polled_rx)
3928            .await
3929            .context(SourcePollTimeoutSnafu)?
3930            .context(SourcePollChannelClosedSnafu)?;
3931
3932        let output = execute_one_sql(&instance, kill_sql, test_query_ctx(43)).await?;
3933        assert!(matches!(output.data, OutputData::AffectedRows(1)));
3934
3935        let insert_result = tokio::time::timeout(Duration::from_secs(5), insert_task)
3936            .await
3937            .context(InsertTaskTimeoutSnafu)?
3938            .context(InsertTaskPanicSnafu)?;
3939        let err = match insert_result {
3940            Ok(_) => return InsertSelectNotCancelledSnafu.fail(),
3941            Err(TestError::ExecuteSql { source, .. }) => source,
3942            Err(err) => return Err(err),
3943        };
3944        assert_eq!(StatusCode::Cancelled, err.status_code());
3945
3946        let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
3947        let process_list = output.data.pretty_print().await;
3948        assert!(
3949            !process_list.contains(insert_sql),
3950            "process list still contains killed insert:\n{process_list}"
3951        );
3952
3953        Ok(())
3954    }
3955
3956    fn insert_dml_plan() -> LogicalPlan {
3957        let schema = SchemaRef::new(Schema::new(vec![Field::new(
3958            "value",
3959            DataType::Int64,
3960            true,
3961        )]));
3962        let target = Arc::new(LogicalTableSource::new(schema));
3963        let input = LogicalPlanBuilder::empty(false).build().unwrap();
3964
3965        LogicalPlanBuilder::insert_into(input, "demo", target, InsertOp::Append)
3966            .unwrap()
3967            .build()
3968            .unwrap()
3969    }
3970
3971    #[test]
3972    fn test_exec_validation() {
3973        let query_ctx = QueryContext::arc();
3974        let plugins: Plugins = Plugins::new();
3975        plugins.insert(QueryOptions {
3976            disallow_cross_catalog_query: true,
3977        });
3978
3979        let sql = r#"
3980        SELECT * FROM demo;
3981        EXPLAIN SELECT * FROM demo;
3982        CREATE DATABASE test_database;
3983        SHOW DATABASES;
3984        "#;
3985        let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
3986        assert_eq!(stmts.len(), 4);
3987        for stmt in stmts {
3988            let re = check_permission(plugins.clone(), &stmt, &query_ctx);
3989            re.unwrap();
3990        }
3991
3992        let sql = r#"
3993        SHOW CREATE TABLE demo;
3994        ALTER TABLE demo ADD COLUMN new_col INT;
3995        "#;
3996        let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
3997        assert_eq!(stmts.len(), 2);
3998        for stmt in stmts {
3999            let re = check_permission(plugins.clone(), &stmt, &query_ctx);
4000            re.unwrap();
4001        }
4002
4003        fn replace_test(template_sql: &str, plugins: Plugins, query_ctx: &QueryContextRef) {
4004            // test right
4005            let right = vec![("", ""), ("", "public."), ("greptime.", "public.")];
4006            for (catalog, schema) in right {
4007                let sql = do_fmt(template_sql, catalog, schema);
4008                do_test(&sql, plugins.clone(), query_ctx, true);
4009            }
4010
4011            let wrong = vec![
4012                ("wrongcatalog.", "public."),
4013                ("wrongcatalog.", "wrongschema."),
4014            ];
4015            for (catalog, schema) in wrong {
4016                let sql = do_fmt(template_sql, catalog, schema);
4017                do_test(&sql, plugins.clone(), query_ctx, false);
4018            }
4019        }
4020
4021        fn do_fmt(template: &str, catalog: &str, schema: &str) -> String {
4022            let vars = HashMap::from([
4023                ("catalog".to_string(), catalog),
4024                ("schema".to_string(), schema),
4025            ]);
4026            template.format(&vars).unwrap()
4027        }
4028
4029        fn do_test(sql: &str, plugins: Plugins, query_ctx: &QueryContextRef, is_ok: bool) {
4030            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
4031            let re = check_permission(plugins, stmt, query_ctx);
4032            if is_ok {
4033                re.unwrap();
4034            } else {
4035                assert!(re.is_err());
4036            }
4037        }
4038
4039        // test insert
4040        let sql = "INSERT INTO {catalog}{schema}monitor(host) VALUES ('host1');";
4041        replace_test(sql, plugins.clone(), &query_ctx);
4042
4043        // test create table
4044        let sql = r#"CREATE TABLE {catalog}{schema}demo(
4045                            host STRING,
4046                            ts TIMESTAMP,
4047                            TIME INDEX (ts),
4048                            PRIMARY KEY(host)
4049                        ) engine=mito;"#;
4050        replace_test(sql, plugins.clone(), &query_ctx);
4051
4052        // test drop table
4053        let sql = "DROP TABLE {catalog}{schema}demo;";
4054        replace_test(sql, plugins.clone(), &query_ctx);
4055
4056        // test undrop table
4057        #[cfg(feature = "enterprise")]
4058        {
4059            let sql = "UNDROP TABLE {catalog}{schema}demo;";
4060            replace_test(sql, plugins.clone(), &query_ctx);
4061        }
4062
4063        // test show tables
4064        let sql = "SHOW TABLES FROM public";
4065        let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
4066        check_permission(plugins.clone(), &stmt[0], &query_ctx).unwrap();
4067
4068        let sql = "SHOW TABLES FROM private";
4069        let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
4070        let re = check_permission(plugins.clone(), &stmt[0], &query_ctx);
4071        assert!(re.is_ok());
4072
4073        // test describe table
4074        let sql = "DESC TABLE {catalog}{schema}demo;";
4075        replace_test(sql, plugins.clone(), &query_ctx);
4076
4077        let comment_flow_cases = [
4078            ("COMMENT ON FLOW my_flow IS 'comment';", true),
4079            ("COMMENT ON FLOW greptime.my_flow IS 'comment';", true),
4080            ("COMMENT ON FLOW wrongcatalog.my_flow IS 'comment';", false),
4081        ];
4082        for (sql, is_ok) in comment_flow_cases {
4083            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
4084            let result = check_permission(plugins.clone(), stmt, &query_ctx);
4085            assert_eq!(result.is_ok(), is_ok);
4086        }
4087
4088        let show_flow_cases = [
4089            ("SHOW CREATE FLOW my_flow;", true),
4090            ("SHOW CREATE FLOW greptime.my_flow;", true),
4091            ("SHOW CREATE FLOW wrongcatalog.my_flow;", false),
4092        ];
4093        for (sql, is_ok) in show_flow_cases {
4094            let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
4095            let result = check_permission(plugins.clone(), stmt, &query_ctx);
4096            assert_eq!(result.is_ok(), is_ok);
4097        }
4098    }
4099
4100    /// A `DropView` DDL sent through the direct gRPC ingress must return an error
4101    /// (e.g. table not found) instead of panicking on `todo!()`.
4102    #[tokio::test]
4103    async fn qx_152_drop_view_via_grpc_ddl_returns_error_not_panic() -> TestResult<()> {
4104        let instance =
4105            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
4106                .await?;
4107
4108        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4109            expr: Some(api::v1::ddl_request::Expr::DropView(
4110                api::v1::DropViewExpr {
4111                    catalog_name: String::new(),
4112                    schema_name: String::new(),
4113                    view_name: "non_existent_view".to_string(),
4114                    view_id: None,
4115                    drop_if_exists: false,
4116                },
4117            )),
4118        });
4119
4120        let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
4121            &instance,
4122            request,
4123            QueryContext::arc(),
4124        )
4125        .await;
4126
4127        let err = match result {
4128            Ok(_) => panic!("DropView DDL request must return an error instead of panicking"),
4129            Err(err) => err,
4130        };
4131        assert_eq!(
4132            err.status_code(),
4133            StatusCode::TableNotFound,
4134            "dropping a non-existent view without IF EXISTS must report TableNotFound, got {err}"
4135        );
4136        Ok(())
4137    }
4138
4139    #[tokio::test]
4140    async fn logical_ddl_validates_and_authorizes_rewritten_batch_before_submission()
4141    -> TestResult<()> {
4142        struct Rewrite(&'static str);
4143        impl SqlQueryInterceptor for Rewrite {
4144            type Error = Error;
4145            fn pre_parsing<'a>(
4146                &self,
4147                query: &'a str,
4148                _: QueryContextRef,
4149            ) -> Result<std::borrow::Cow<'a, str>> {
4150                Ok(if self.0 == "sql_bytes" {
4151                    std::borrow::Cow::Owned(format!("{query}{}", " ".repeat(1024 * 1024 + 1)))
4152                } else {
4153                    std::borrow::Cow::Borrowed(query)
4154                })
4155            }
4156            fn pre_execute(
4157                &self,
4158                statement: Option<&Statement>,
4159                _: Option<&LogicalPlan>,
4160                _: QueryContextRef,
4161            ) -> Result<()> {
4162                ensure!(
4163                    !(self.0 == "pre_execute"
4164                        && matches!(statement,
4165                    Some(Statement::CreateTable(create)) if create.name.to_string() == "allowed2")),
4166                    InvalidSqlSnafu {
4167                        err_msg: "second member rejected by interceptor"
4168                    }
4169                );
4170                Ok(())
4171            }
4172            fn post_parsing(
4173                &self,
4174                mut statements: Vec<Statement>,
4175                _: QueryContextRef,
4176            ) -> Result<Vec<Statement>> {
4177                match self.0 {
4178                    "target" => {
4179                        statements[1] = parse_one_sql(
4180                            "CREATE TABLE denied (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH (on_physical_table='phy')",
4181                        )
4182                    }
4183                    "count" => statements = (0..129).map(|i| parse_one_sql(&format!(
4184                        "CREATE TABLE allowed_{i} (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH (on_physical_table='phy')"
4185                    ))).collect(),
4186                    "type" => statements[1] = parse_one_sql("SELECT 1"),
4187                    "secret_bytes" | "array_bytes" | "small_option" => {
4188                        if let Statement::CreateTable(create) = &mut statements[1] {
4189                            let value = "x".repeat(if self.0 == "small_option" { 10 } else { 2 * 1024 * 1024 });
4190                            if self.0 == "array_bytes" {
4191                                create.options.insert_options("non_scalar", vec![value.as_str()].into());
4192                            } else {
4193                                create.options.insert("secret_access_key".to_string(), value);
4194                            }
4195                        }
4196                    }
4197                    "partition_shape" => {
4198                        if let Statement::CreateTable(create) = &mut statements[1] {
4199                            create.partitions = Some(sql::statements::create::Partitions {
4200                                column_list: vec![],
4201                                exprs: vec![sqlparser::ast::Expr::Identifier(sqlparser::ast::Ident::new("hidden"))],
4202                            });
4203                        }
4204                    }
4205                    "vector_hidden" => {
4206                        statements[1] = parse_one_sql("CREATE TABLE allowed2 (ts TIMESTAMP TIME INDEX, v VECTOR(3)) ENGINE=metric WITH (on_physical_table='phy')");
4207                        if let Statement::CreateTable(create) = &mut statements[1] {
4208                            create.columns[1].column_def.options.push(sqlparser::ast::ColumnOptionDef {
4209                                name: None,
4210                                option: sqlparser::ast::ColumnOption::Comment("private-value".repeat(200_000)),
4211                            });
4212                        }
4213                    }
4214                    mode if mode.starts_with("fulltext_") || mode.starts_with("skipping_")
4215                        || mode.starts_with("inverted_") || mode.starts_with("vector_") => {
4216                        if let Statement::CreateTable(create) = &mut statements[1] {
4217                            let (kind, value_type) = mode.split_once('_').unwrap();
4218                            let value = "x".repeat(if value_type == "small" { 10 } else { 2 * 1024 * 1024 });
4219                            let extensions = &mut create.columns[1].extensions;
4220                            let options = match kind {
4221                                "fulltext" => &mut extensions.fulltext_index_options,
4222                                "skipping" => &mut extensions.skipping_index_options,
4223                                "inverted" => &mut extensions.inverted_index_options,
4224                                _ => &mut extensions.vector_options,
4225                            }.get_or_insert_default();
4226                            if kind == "vector" {
4227                                options.insert("dim".to_string(), "3".to_string());
4228                            }
4229                            if value_type == "array" {
4230                                options.insert_options("non_scalar", vec![value.as_str()].into());
4231                            } else {
4232                                options.insert("secret_access_key".to_string(), value);
4233                            }
4234                        }
4235                    }
4236                    "physical_null" | "physical_array" => {
4237                        let value = if self.0 == "physical_null" { "NULL" } else { "['x']" };
4238                        statements[1] = parse_one_sql(&format!(
4239                            "CREATE TABLE allowed2 (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH (on_physical_table='phy', physical_metric_table={value})"
4240                        ));
4241                    }
4242                    "ast_bytes" => {
4243                        if let Statement::CreateTable(create) = &mut statements[1] {
4244                            create.name =
4245                                ObjectName(vec![sqlparser::ast::ObjectNamePart::Identifier(
4246                                    sqlparser::ast::Ident::new("t".repeat(1024 * 1024)),
4247                                )]);
4248                        }
4249                    }
4250                    _ => {}
4251                }
4252                Ok(statements)
4253            }
4254        }
4255        for mode in [
4256            "none",
4257            "target",
4258            "count",
4259            "type",
4260            "sql_bytes",
4261            "ast_bytes",
4262            "secret_bytes",
4263            "array_bytes",
4264            "small_option",
4265            "physical_null",
4266            "physical_array",
4267            "fulltext_secret",
4268            "fulltext_array",
4269            "fulltext_small",
4270            "skipping_secret",
4271            "skipping_array",
4272            "skipping_small",
4273            "inverted_secret",
4274            "inverted_array",
4275            "inverted_small",
4276            "vector_secret",
4277            "vector_array",
4278            "vector_small",
4279            "vector_hidden",
4280            "partition_shape",
4281            "pre_execute",
4282            "cross_catalog",
4283            "valid",
4284        ] {
4285            let catalog =
4286                catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
4287            let executor = Arc::new(MockProcedureExecutor::new(catalog.clone()));
4288            let plugins = Plugins::new();
4289            plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
4290            plugins.insert::<SqlQueryInterceptorRef<Error>>(Arc::new(Rewrite(mode)));
4291            plugins.insert(QueryOptions {
4292                disallow_cross_catalog_query: mode == "cross_catalog",
4293            });
4294            let instance = test_instance_with_catalog_manager(
4295                catalog,
4296                test_table(1025, "target")?,
4297                plugins,
4298                None,
4299                executor.clone(),
4300            )
4301            .await?;
4302            let create = |name| {
4303                format!(
4304                    "CREATE TABLE {name} (ts TIMESTAMP TIME INDEX, host STRING) ENGINE=metric WITH (on_physical_table='phy')"
4305                )
4306            };
4307            let sql = format!(
4308                "{};{}",
4309                create("allowed"),
4310                create(match mode {
4311                    "none" => "denied",
4312                    "cross_catalog" => "other.public.allowed2",
4313                    _ => "allowed2",
4314                })
4315            );
4316            let result =
4317                SqlQueryHandler::create_logical_tables(&instance, &sql, QueryContext::arc()).await;
4318            let submits = matches!(mode, "valid" | "small_option")
4319                || (mode.ends_with("_small") && !mode.starts_with("vector_"));
4320            if matches!(mode, "none" | "target") {
4321                let err = result.unwrap_err();
4322                assert_eq!(
4323                    err.status_code(),
4324                    StatusCode::PermissionDenied,
4325                    "{mode}: {err}"
4326                );
4327            } else {
4328                assert_eq!(
4329                    result.unwrap_err().status_code(),
4330                    match mode {
4331                        _ if submits => StatusCode::Unsupported,
4332                        "cross_catalog" => StatusCode::AccessDenied,
4333                        _ => StatusCode::InvalidArguments,
4334                    },
4335                    "{mode}"
4336                );
4337            }
4338            assert_eq!(
4339                executor.submitted.lock().unwrap().len(),
4340                usize::from(submits),
4341                "{mode}"
4342            );
4343        }
4344        Ok(())
4345    }
4346
4347    /// `DROP VIEW IF EXISTS` on a missing view through the direct gRPC ingress must
4348    /// succeed with 0 affected rows (no error, no DDL task submitted), instead of
4349    /// returning `TableNotFound`.
4350    #[tokio::test]
4351    async fn qx_152_drop_view_if_exists_missing_view_via_grpc_ddl_succeeds() -> TestResult<()> {
4352        let catalog_manager =
4353            catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
4354        let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
4355        let instance = test_instance_with_catalog_manager(
4356            catalog_manager,
4357            test_table(1025, "target")?,
4358            Plugins::new(),
4359            None,
4360            procedure_executor.clone() as ProcedureExecutorRef,
4361        )
4362        .await?;
4363
4364        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4365            expr: Some(api::v1::ddl_request::Expr::DropView(
4366                api::v1::DropViewExpr {
4367                    catalog_name: String::new(),
4368                    schema_name: String::new(),
4369                    view_name: "non_existent_view".to_string(),
4370                    view_id: None,
4371                    drop_if_exists: true,
4372                },
4373            )),
4374        });
4375
4376        let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
4377            &instance,
4378            request,
4379            QueryContext::arc(),
4380        )
4381        .await;
4382
4383        let output = match result {
4384            Ok(output) => output,
4385            Err(err) => {
4386                panic!("DROP VIEW IF EXISTS on a missing view must succeed, got error: {err}")
4387            }
4388        };
4389        assert!(
4390            matches!(output.data, OutputData::AffectedRows(0)),
4391            "DROP VIEW IF EXISTS on a missing view must report 0 affected rows"
4392        );
4393        assert!(
4394            procedure_executor.submitted.lock().unwrap().is_empty(),
4395            "DROP VIEW IF EXISTS on a missing view must not submit a DDL task"
4396        );
4397        Ok(())
4398    }
4399
4400    /// A `CREATE VIEW` followed by `DROP VIEW` through the direct gRPC ingress must
4401    /// succeed end to end: the view is registered in the catalog and then removed.
4402    #[tokio::test]
4403    async fn qx_152_drop_existing_view_via_grpc_ddl_succeeds() -> TestResult<()> {
4404        let catalog_manager =
4405            catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
4406        let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
4407        let instance = test_instance_with_catalog_manager(
4408            catalog_manager,
4409            test_table(1025, "target")?,
4410            Plugins::new(),
4411            None,
4412            procedure_executor.clone() as ProcedureExecutorRef,
4413        )
4414        .await?;
4415
4416        // The default "greptime.public" schema must be visible to the kv-backed table
4417        // metadata manager for `CREATE VIEW`/`CREATE TABLE` to pass validation.
4418        instance
4419            .table_metadata_manager()
4420            .schema_manager()
4421            .create(
4422                common_meta::key::schema_name::SchemaNameKey::new("greptime", "public"),
4423                None,
4424                true,
4425            )
4426            .await
4427            .unwrap();
4428
4429        let create_view_request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4430            expr: Some(api::v1::ddl_request::Expr::CreateView(
4431                api::v1::CreateViewExpr {
4432                    catalog_name: String::new(),
4433                    schema_name: String::new(),
4434                    view_name: "my_view".to_string(),
4435                    logical_plan: vec![1, 2, 3],
4436                    create_if_not_exists: false,
4437                    or_replace: false,
4438                    table_names: vec![],
4439                    columns: vec![],
4440                    plan_columns: vec![],
4441                    definition: "CREATE VIEW my_view AS SELECT * FROM source".to_string(),
4442                },
4443            )),
4444        });
4445
4446        let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
4447            &instance,
4448            create_view_request,
4449            QueryContext::arc(),
4450        )
4451        .await
4452        {
4453            Ok(output) => output,
4454            Err(err) => panic!("CREATE VIEW via gRPC DDL must succeed, got error: {err}"),
4455        };
4456        assert!(
4457            matches!(output.data, OutputData::AffectedRows(0)),
4458            "CREATE VIEW via gRPC DDL must report 0 affected rows"
4459        );
4460
4461        // The view is registered in the catalog as a view.
4462        let view = instance
4463            .catalog_manager()
4464            .table("greptime", "public", "my_view", None)
4465            .await
4466            .unwrap()
4467            .expect("view should exist after CREATE VIEW");
4468        assert_eq!(view.table_info().table_type, TableType::View);
4469
4470        let drop_view_request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4471            expr: Some(api::v1::ddl_request::Expr::DropView(
4472                api::v1::DropViewExpr {
4473                    catalog_name: String::new(),
4474                    schema_name: String::new(),
4475                    view_name: "my_view".to_string(),
4476                    view_id: None,
4477                    drop_if_exists: false,
4478                },
4479            )),
4480        });
4481
4482        let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
4483            &instance,
4484            drop_view_request,
4485            QueryContext::arc(),
4486        )
4487        .await
4488        {
4489            Ok(output) => output,
4490            Err(err) => panic!("DROP VIEW via gRPC DDL must succeed, got error: {err}"),
4491        };
4492        assert!(
4493            matches!(output.data, OutputData::AffectedRows(0)),
4494            "DROP VIEW via gRPC DDL must report 0 affected rows"
4495        );
4496
4497        // The view is gone after the drop.
4498        assert!(
4499            instance
4500                .catalog_manager()
4501                .table("greptime", "public", "my_view", None)
4502                .await
4503                .unwrap()
4504                .is_none(),
4505            "view should be removed after DROP VIEW"
4506        );
4507
4508        let submitted = procedure_executor.submitted.lock().unwrap();
4509        assert_eq!(
4510            submitted.len(),
4511            2,
4512            "expected create and drop view tasks, got {submitted:?}"
4513        );
4514        assert!(matches!(&submitted[0], DdlTask::CreateView(_)));
4515        assert!(matches!(&submitted[1], DdlTask::DropView(_)));
4516        Ok(())
4517    }
4518
4519    /// A direct gRPC `CreateTable` whose time index column is not a timestamp must
4520    /// be rejected with `InvalidArguments` instead of panicking while building the schema.
4521    #[tokio::test]
4522    async fn qx_153_create_table_with_non_timestamp_time_index_via_grpc_returns_error()
4523    -> TestResult<()> {
4524        let instance =
4525            test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
4526                .await?;
4527
4528        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4529            expr: Some(api::v1::ddl_request::Expr::CreateTable(
4530                api::v1::CreateTableExpr {
4531                    catalog_name: String::new(),
4532                    schema_name: String::new(),
4533                    table_name: "demo".to_string(),
4534                    desc: String::new(),
4535                    column_defs: vec![api::v1::ColumnDef {
4536                        name: "host".to_string(),
4537                        data_type: api::v1::ColumnDataType::String as i32,
4538                        is_nullable: true,
4539                        default_constraint: vec![],
4540                        semantic_type: 0,
4541                        comment: String::new(),
4542                        datatype_extension: None,
4543                        options: None,
4544                    }],
4545                    time_index: "host".to_string(),
4546                    primary_keys: vec![],
4547                    create_if_not_exists: false,
4548                    table_options: HashMap::new(),
4549                    table_id: None,
4550                    engine: "mito".to_string(),
4551                },
4552            )),
4553        });
4554
4555        let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
4556            &instance,
4557            request,
4558            QueryContext::arc(),
4559        )
4560        .await;
4561
4562        let err = match result {
4563            Ok(_) => panic!("CreateTable with a non-timestamp time index must be rejected"),
4564            Err(err) => err,
4565        };
4566        assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}");
4567        Ok(())
4568    }
4569
4570    /// A valid `CREATE TABLE` (timestamp time index) through the direct gRPC ingress
4571    /// must succeed, guarding that the `validate_create_expr` ingress check doesn't
4572    /// accidentally reject good requests.
4573    #[tokio::test]
4574    async fn qx_153_create_table_with_timestamp_time_index_via_grpc_succeeds() -> TestResult<()> {
4575        let catalog_manager =
4576            catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
4577        let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
4578        let instance = test_instance_with_catalog_manager(
4579            catalog_manager,
4580            test_table(1025, "target")?,
4581            Plugins::new(),
4582            None,
4583            procedure_executor.clone() as ProcedureExecutorRef,
4584        )
4585        .await?;
4586
4587        // The default "greptime.public" schema must be visible to the kv-backed table
4588        // metadata manager for `CREATE TABLE` to pass validation.
4589        instance
4590            .table_metadata_manager()
4591            .schema_manager()
4592            .create(
4593                common_meta::key::schema_name::SchemaNameKey::new("greptime", "public"),
4594                None,
4595                true,
4596            )
4597            .await
4598            .unwrap();
4599
4600        let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4601            expr: Some(api::v1::ddl_request::Expr::CreateTable(
4602                api::v1::CreateTableExpr {
4603                    catalog_name: String::new(),
4604                    schema_name: String::new(),
4605                    table_name: "demo".to_string(),
4606                    desc: String::new(),
4607                    column_defs: vec![
4608                        api::v1::ColumnDef {
4609                            name: "host".to_string(),
4610                            data_type: api::v1::ColumnDataType::String as i32,
4611                            is_nullable: true,
4612                            default_constraint: vec![],
4613                            semantic_type: 0,
4614                            comment: String::new(),
4615                            datatype_extension: None,
4616                            options: None,
4617                        },
4618                        api::v1::ColumnDef {
4619                            name: "ts".to_string(),
4620                            data_type: api::v1::ColumnDataType::TimestampMillisecond as i32,
4621                            is_nullable: true,
4622                            default_constraint: vec![],
4623                            semantic_type: 0,
4624                            comment: String::new(),
4625                            datatype_extension: None,
4626                            options: None,
4627                        },
4628                    ],
4629                    time_index: "ts".to_string(),
4630                    primary_keys: vec![],
4631                    create_if_not_exists: false,
4632                    table_options: HashMap::new(),
4633                    table_id: None,
4634                    engine: "mito".to_string(),
4635                },
4636            )),
4637        });
4638
4639        let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
4640            &instance,
4641            request,
4642            QueryContext::arc(),
4643        )
4644        .await
4645        {
4646            Ok(output) => output,
4647            Err(err) => panic!("CREATE TABLE via gRPC DDL must succeed, got error: {err}"),
4648        };
4649        assert!(
4650            matches!(output.data, OutputData::AffectedRows(0)),
4651            "CREATE TABLE via gRPC DDL must report 0 affected rows"
4652        );
4653
4654        // The table is registered in the catalog.
4655        let table = instance
4656            .catalog_manager()
4657            .table("greptime", "public", "demo", None)
4658            .await
4659            .unwrap()
4660            .expect("table should exist after CREATE TABLE");
4661        assert_eq!(table.table_info().table_type, TableType::Base);
4662
4663        let submitted = procedure_executor.submitted.lock().unwrap();
4664        assert_eq!(
4665            submitted.len(),
4666            1,
4667            "expected one create table task, got {submitted:?}"
4668        );
4669        assert!(matches!(&submitted[0], DdlTask::CreateTable(_)));
4670        Ok(())
4671    }
4672
4673    /// COPY FROM STDIN must enforce the same cross-catalog restriction as
4674    /// SQL INSERT when `disallow_cross_catalog_query` is configured.
4675    #[tokio::test]
4676    async fn test_copy_in_cross_catalog_restriction() -> TestResult<()> {
4677        use catalog::memory::MemoryCatalogManager;
4678        use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
4679        use servers::query_handler::CopyInHandler;
4680
4681        let catalog = MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
4682        // A table in another catalog.
4683        catalog.register_catalog_sync("other").unwrap();
4684        catalog
4685            .register_schema_sync(catalog::RegisterSchemaRequest {
4686                catalog: "other".to_string(),
4687                schema: DEFAULT_SCHEMA_NAME.to_string(),
4688            })
4689            .unwrap();
4690        catalog
4691            .register_table_sync(catalog::RegisterTableRequest {
4692                catalog: "other".to_string(),
4693                schema: DEFAULT_SCHEMA_NAME.to_string(),
4694                table_name: "copy_guard".to_string(),
4695                table_id: 2048,
4696                table: test_table(2048, "copy_guard")?,
4697            })
4698            .unwrap();
4699        let executor = Arc::new(MockProcedureExecutor::new(catalog.clone()));
4700        let plugins = Plugins::new();
4701        plugins.insert(QueryOptions {
4702            disallow_cross_catalog_query: true,
4703        });
4704        let instance = test_instance_with_catalog_manager(
4705            catalog,
4706            test_table(1025, "target")?,
4707            plugins,
4708            None,
4709            executor,
4710        )
4711        .await?;
4712
4713        let ctx = Arc::new(QueryContext::with(
4714            DEFAULT_CATALOG_NAME,
4715            DEFAULT_SCHEMA_NAME,
4716        ));
4717
4718        // Cross-catalog COPY is rejected before any data is streamed.
4719        let err = instance
4720            .copy_in_table("other", DEFAULT_SCHEMA_NAME, "copy_guard", ctx.clone())
4721            .await
4722            .unwrap_err();
4723        assert!(
4724            format!("{err:?}").contains("Illegal access to catalog"),
4725            "{err}"
4726        );
4727
4728        // Same-catalog COPY resolves normally.
4729        let found = instance
4730            .copy_in_table(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "target", ctx)
4731            .await
4732            .unwrap();
4733        assert!(found.is_some());
4734
4735        Ok(())
4736    }
4737}