Skip to main content

frontend/
instance.rs

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