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