Skip to main content

frontend/
instance.rs

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