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