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