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