1pub mod builder;
16mod dashboard;
17mod grpc;
18mod influxdb;
19mod jaeger;
20mod log_handler;
21mod logs;
22mod opentsdb;
23mod otlp;
24pub mod prom_store;
25mod promql;
26mod region_query;
27pub mod standalone;
28
29use std::pin::Pin;
30use std::sync::atomic::AtomicBool;
31use std::sync::{Arc, atomic};
32use std::time::{Duration, SystemTime};
33
34use async_stream::stream;
35use async_trait::async_trait;
36use auth::{PermissionChecker, PermissionCheckerRef, PermissionReq};
37use catalog::CatalogManagerRef;
38use catalog::process_manager::{
39 ProcessManagerRef, QueryStatement as CatalogQueryStatement, SlowQueryTimer,
40};
41use client::OutputData;
42use common_base::Plugins;
43use common_base::cancellation::CancellableFuture;
44use common_error::ext::{BoxedError, ErrorExt};
45use common_event_recorder::EventRecorderRef;
46use common_meta::cache_invalidator::CacheInvalidatorRef;
47use common_meta::key::TableMetadataManagerRef;
48use common_meta::key::table_name::TableNameKey;
49use common_meta::node_manager::NodeManagerRef;
50use common_meta::procedure_executor::ProcedureExecutorRef;
51use common_query::Output;
52use common_recordbatch::RecordBatchStreamWrapper;
53use common_recordbatch::error::StreamTimeoutSnafu;
54use common_telemetry::logging::SlowQueryOptions;
55use common_telemetry::{debug, error, tracing};
56use dashmap::DashMap;
57use datafusion_expr::LogicalPlan;
58use futures::{Stream, StreamExt};
59use lazy_static::lazy_static;
60use operator::delete::DeleterRef;
61use operator::insert::InserterRef;
62use operator::statement::{StatementExecutor, StatementExecutorRef};
63use partition::manager::PartitionRuleManagerRef;
64use pipeline::pipeline_operator::PipelineOperator;
65use prometheus::HistogramTimer;
66use promql_parser::label::Matcher;
67use query::QueryEngineRef;
68use query::metrics::OnDone;
69use query::parser::{PromQuery, QueryLanguageParser, QueryStatement};
70use query::query_engine::DescribeResult;
71use query::query_engine::options::{QueryOptions, validate_catalog_and_schema};
72use servers::error::{
73 self as server_error, AuthSnafu, CommonMetaSnafu, ExecuteQuerySnafu,
74 OtlpMetricModeIncompatibleSnafu, ParsePromQLSnafu, UnexpectedResultSnafu,
75};
76use servers::interceptor::{
77 PromQueryInterceptor, PromQueryInterceptorRef, SqlQueryInterceptor, SqlQueryInterceptorRef,
78};
79use servers::otlp::metrics::legacy_normalize_otlp_name;
80use servers::prometheus_handler::PrometheusHandler;
81use servers::query_handler::sql::SqlQueryHandler;
82use session::context::{Channel, QueryContextRef};
83use session::table_name::table_idents_to_full_name;
84use snafu::prelude::*;
85use sql::ast::ObjectNamePartExt;
86use sql::dialect::Dialect;
87use sql::parser::{ParseOptions, ParserContext};
88use sql::statements::comment::CommentObject;
89use sql::statements::copy::{CopyDatabase, CopyTable};
90use sql::statements::statement::Statement;
91use sql::statements::tql::Tql;
92use sqlparser::ast::{AnalyzeFormat, ObjectName};
93pub use standalone::StandaloneDatanodeManager;
94use table::requests::{OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM};
95use tracing::Span;
96
97use crate::error::{
98 self, CollectRecordbatchSnafu, Error, ExecLogicalPlanSnafu, ExecutePromqlSnafu, ExternalSnafu,
99 InvalidSqlSnafu, ParseSqlSnafu, PermissionSnafu, PlanStatementSnafu, Result,
100 SqlExecInterceptedSnafu, StatementTimeoutSnafu, TableOperationSnafu,
101};
102use crate::service_config::InfluxdbMergeMode;
103use crate::stream_wrapper::CancellableStreamWrapper;
104
105lazy_static! {
106 static ref OTLP_LEGACY_DEFAULT_VALUE: String = "legacy".to_string();
107}
108
109#[derive(Clone)]
113pub struct Instance {
114 frontend_peer_addr: String,
115 catalog_manager: CatalogManagerRef,
116 pipeline_operator: Arc<PipelineOperator>,
117 statement_executor: Arc<StatementExecutor>,
118 query_engine: QueryEngineRef,
119 plugins: Plugins,
120 inserter: InserterRef,
121 deleter: DeleterRef,
122 table_metadata_manager: TableMetadataManagerRef,
123 event_recorder: Option<EventRecorderRef>,
124 process_manager: ProcessManagerRef,
125 slow_query_options: SlowQueryOptions,
126 influxdb_default_merge_mode: InfluxdbMergeMode,
127 trace_ingest_chunk_size: usize,
128 suspend: Arc<AtomicBool>,
129
130 otlp_metrics_table_legacy_cache: DashMap<String, DashMap<String, bool>>,
135}
136
137impl Instance {
138 pub fn frontend_peer_addr(&self) -> &str {
139 &self.frontend_peer_addr
140 }
141
142 pub fn catalog_manager(&self) -> &CatalogManagerRef {
143 &self.catalog_manager
144 }
145
146 pub fn query_engine(&self) -> &QueryEngineRef {
147 &self.query_engine
148 }
149
150 pub fn plugins(&self) -> &Plugins {
151 &self.plugins
152 }
153
154 pub fn statement_executor(&self) -> &StatementExecutorRef {
155 &self.statement_executor
156 }
157
158 pub fn table_metadata_manager(&self) -> &TableMetadataManagerRef {
159 &self.table_metadata_manager
160 }
161
162 pub fn inserter(&self) -> &InserterRef {
163 &self.inserter
164 }
165
166 pub fn process_manager(&self) -> &ProcessManagerRef {
167 &self.process_manager
168 }
169
170 pub fn node_manager(&self) -> &NodeManagerRef {
171 self.inserter.node_manager()
172 }
173
174 pub fn partition_manager(&self) -> &PartitionRuleManagerRef {
175 self.inserter.partition_manager()
176 }
177
178 pub fn cache_invalidator(&self) -> &CacheInvalidatorRef {
179 self.statement_executor.cache_invalidator()
180 }
181
182 pub fn procedure_executor(&self) -> &ProcedureExecutorRef {
183 self.statement_executor.procedure_executor()
184 }
185
186 pub fn suspend_state(&self) -> Arc<AtomicBool> {
187 self.suspend.clone()
188 }
189
190 pub(crate) fn is_suspended(&self) -> bool {
191 self.suspend.load(atomic::Ordering::Relaxed)
192 }
193}
194
195fn parse_stmt(sql: &str, dialect: &(dyn Dialect + Send + Sync)) -> Result<Vec<Statement>> {
196 ParserContext::create_with_dialect(sql, dialect, ParseOptions::default()).context(ParseSqlSnafu)
197}
198
199fn validate_analyze_stream_statement(stmt: &mut Statement) -> Result<()> {
200 let Statement::Explain(explain) = stmt else {
201 return InvalidSqlSnafu {
202 err_msg: "only EXPLAIN ANALYZE VERBOSE statement is supported",
203 }
204 .fail();
205 };
206 ensure!(
207 explain.analyze && explain.verbose,
208 InvalidSqlSnafu {
209 err_msg: "statement must be EXPLAIN ANALYZE VERBOSE"
210 }
211 );
212 match explain.format {
213 None | Some(AnalyzeFormat::JSON) => {
214 explain.format = None;
218 Ok(())
219 }
220 Some(_) => InvalidSqlSnafu {
221 err_msg: "only FORMAT JSON is supported for analyze stream",
222 }
223 .fail(),
224 }
225}
226
227impl Instance {
228 fn statement_slow_query_timer(
229 &self,
230 stmt: &Statement,
231 schema_name: String,
232 ) -> Option<SlowQueryTimer> {
233 if !stmt.is_readonly() || !self.slow_query_options.enable {
234 return None;
235 }
236
237 self.event_recorder.clone().map(|event_recorder| {
238 SlowQueryTimer::new(
239 CatalogQueryStatement::Sql(stmt.clone()),
240 schema_name,
241 self.slow_query_options.threshold,
242 self.slow_query_options.sample_ratio,
243 self.slow_query_options.record_type,
244 event_recorder,
245 )
246 })
247 }
248
249 async fn query_statement(&self, stmt: Statement, query_ctx: QueryContextRef) -> Result<Output> {
250 check_permission(self.plugins.clone(), &stmt, &query_ctx)?;
251
252 let query_interceptor = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
253 let query_interceptor = query_interceptor.as_ref();
254
255 if should_track_statement_process(&stmt) {
256 let catalog_name = query_ctx.current_catalog().to_string();
257 let schema_name = query_ctx.current_schema();
258 let slow_query_timer = self.statement_slow_query_timer(&stmt, schema_name.clone());
259
260 let ticket = self.process_manager.register_query(
261 catalog_name,
262 vec![schema_name],
263 stmt.to_string(),
264 query_ctx.conn_info().to_string(),
265 Some(query_ctx.process_id()),
266 slow_query_timer,
267 );
268
269 let query_fut = self.exec_statement_with_timeout(stmt, query_ctx, query_interceptor);
270
271 CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
272 .await
273 .map_err(|_| error::CancelledSnafu.build())?
274 .map(|output| {
275 let Output { meta, data } = output;
276
277 let data = match data {
278 OutputData::Stream(stream) => OutputData::Stream(Box::pin(
279 CancellableStreamWrapper::new(stream, ticket),
280 )),
281 other => other,
282 };
283 Output { data, meta }
284 })
285 } else {
286 self.exec_statement_with_timeout(stmt, query_ctx, query_interceptor)
287 .await
288 }
289 }
290
291 async fn exec_statement_with_timeout(
292 &self,
293 stmt: Statement,
294 query_ctx: QueryContextRef,
295 query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
296 ) -> Result<Output> {
297 let timeout = derive_timeout(&stmt, &query_ctx);
298 match timeout {
299 Some(timeout) => {
300 let start = tokio::time::Instant::now();
301 let output = tokio::time::timeout(
302 timeout,
303 self.exec_statement(stmt, query_ctx, query_interceptor),
304 )
305 .await
306 .map_err(|_| StatementTimeoutSnafu.build())??;
307 let output = map_query_output(output)?;
308 let remaining_timeout = timeout.checked_sub(start.elapsed()).unwrap_or_default();
310 attach_timeout(output, remaining_timeout)
311 }
312 None => self
313 .exec_statement(stmt, query_ctx, query_interceptor)
314 .await
315 .and_then(map_query_output),
316 }
317 }
318
319 async fn exec_statement(
320 &self,
321 stmt: Statement,
322 query_ctx: QueryContextRef,
323 query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
324 ) -> Result<Output> {
325 match stmt {
326 Statement::Query(_) | Statement::Explain(_) | Statement::Delete(_) => {
327 if let Statement::Explain(explain) = &stmt
329 && let Some(format) = explain.format()
330 {
331 query_ctx.set_explain_format(format.to_string());
332 }
333
334 self.plan_and_exec_sql(stmt, &query_ctx, query_interceptor)
335 .await
336 }
337 Statement::Tql(tql) => {
338 self.plan_and_exec_tql(&query_ctx, query_interceptor, tql)
339 .await
340 }
341 _ => {
342 query_interceptor.pre_execute(Some(&stmt), None, query_ctx.clone())?;
343 self.statement_executor
344 .execute_sql(stmt, query_ctx)
345 .await
346 .context(TableOperationSnafu)
347 }
348 }
349 }
350
351 async fn plan_and_exec_sql(
352 &self,
353 stmt: Statement,
354 query_ctx: &QueryContextRef,
355 query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
356 ) -> Result<Output> {
357 let stmt = QueryStatement::Sql(stmt);
358 let plan = self
359 .statement_executor
360 .plan(&stmt, query_ctx.clone())
361 .await?;
362 let QueryStatement::Sql(stmt) = stmt else {
363 unreachable!()
364 };
365 query_interceptor.pre_execute(Some(&stmt), Some(&plan), query_ctx.clone())?;
366
367 self.statement_executor
368 .exec_plan(plan, query_ctx.clone())
369 .await
370 .context(TableOperationSnafu)
371 }
372
373 async fn plan_and_exec_tql(
374 &self,
375 query_ctx: &QueryContextRef,
376 query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
377 tql: Tql,
378 ) -> Result<Output> {
379 let plan = self
380 .statement_executor
381 .plan_tql(tql.clone(), query_ctx)
382 .await?;
383 query_interceptor.pre_execute(
384 Some(&Statement::Tql(tql)),
385 Some(&plan),
386 query_ctx.clone(),
387 )?;
388 self.statement_executor
389 .exec_plan(plan, query_ctx.clone())
390 .await
391 .context(TableOperationSnafu)
392 }
393
394 async fn check_otlp_legacy(
395 &self,
396 names: &[&String],
397 ctx: QueryContextRef,
398 ) -> server_error::Result<bool> {
399 let db_string = ctx.get_db_string();
400 let cache = self
402 .otlp_metrics_table_legacy_cache
403 .entry(db_string.clone())
404 .or_default();
405 if let Some(flag) = fast_legacy_check(&cache, names)? {
406 return Ok(flag);
407 }
408 drop(cache);
410
411 let catalog = ctx.current_catalog();
412 let schema = ctx.current_schema();
413
414 let normalized_names = names
416 .iter()
417 .map(|n| legacy_normalize_otlp_name(n))
418 .collect::<Vec<_>>();
419 let table_names = normalized_names
420 .iter()
421 .map(|n| TableNameKey::new(catalog, &schema, n))
422 .collect::<Vec<_>>();
423 let table_values = self
424 .table_metadata_manager()
425 .table_name_manager()
426 .batch_get(table_names)
427 .await
428 .context(CommonMetaSnafu)?;
429 let table_ids = table_values
430 .into_iter()
431 .filter_map(|v| v.map(|vi| vi.table_id()))
432 .collect::<Vec<_>>();
433
434 if table_ids.is_empty() {
436 let cache = self
437 .otlp_metrics_table_legacy_cache
438 .entry(db_string)
439 .or_default();
440 names.iter().for_each(|name| {
441 cache.insert((*name).clone(), false);
442 });
443 return Ok(false);
444 }
445
446 let table_infos = self
448 .table_metadata_manager()
449 .table_info_manager()
450 .batch_get(&table_ids)
451 .await
452 .context(CommonMetaSnafu)?;
453 let options = table_infos
454 .values()
455 .map(|info| {
456 info.table_info
457 .meta
458 .options
459 .extra_options
460 .get(OTLP_METRIC_COMPAT_KEY)
461 .unwrap_or(&OTLP_LEGACY_DEFAULT_VALUE)
462 })
463 .collect::<Vec<_>>();
464 let cache = self
465 .otlp_metrics_table_legacy_cache
466 .entry(db_string)
467 .or_default();
468 if !options.is_empty() {
469 let has_prom = options.iter().any(|opt| *opt == OTLP_METRIC_COMPAT_PROM);
471 let has_legacy = options
472 .iter()
473 .any(|opt| *opt == OTLP_LEGACY_DEFAULT_VALUE.as_str());
474 ensure!(!(has_prom && has_legacy), OtlpMetricModeIncompatibleSnafu);
475 let flag = has_legacy;
476 names.iter().for_each(|name| {
477 cache.insert((*name).clone(), flag);
478 });
479 Ok(flag)
480 } else {
481 names.iter().for_each(|name| {
483 cache.insert((*name).clone(), false);
484 });
485 Ok(false)
486 }
487 }
488}
489
490fn fast_legacy_check(
491 cache: &DashMap<String, bool>,
492 names: &[&String],
493) -> server_error::Result<Option<bool>> {
494 let hit_cache = names
495 .iter()
496 .filter_map(|name| cache.get(*name))
497 .collect::<Vec<_>>();
498 if !hit_cache.is_empty() {
499 let hit_legacy = hit_cache.iter().any(|en| *en.value());
500 let hit_prom = hit_cache.iter().any(|en| !*en.value());
501
502 ensure!(!(hit_legacy && hit_prom), OtlpMetricModeIncompatibleSnafu);
506
507 let flag = hit_legacy;
508 drop(hit_cache);
510
511 names.iter().for_each(|name| {
513 if !cache.contains_key(*name) {
514 cache.insert((*name).clone(), flag);
515 }
516 });
517 Ok(Some(flag))
518 } else {
519 Ok(None)
520 }
521}
522
523fn derive_timeout(stmt: &Statement, query_ctx: &QueryContextRef) -> Option<Duration> {
526 let query_timeout = query_ctx.query_timeout()?;
527 if query_timeout.is_zero() {
528 return None;
529 }
530 match query_ctx.channel() {
531 Channel::Mysql if stmt.is_readonly() => Some(query_timeout),
532 Channel::Postgres => Some(query_timeout),
533 _ => None,
534 }
535}
536
537fn derive_timeout_for_plan(plan: &LogicalPlan, query_ctx: &QueryContextRef) -> Option<Duration> {
539 let query_timeout = query_ctx.query_timeout()?;
540 if query_timeout.is_zero() {
541 return None;
542 }
543 match query_ctx.channel() {
544 Channel::Mysql if is_readonly_plan(plan) => Some(query_timeout),
545 Channel::Postgres => Some(query_timeout),
546 _ => None,
547 }
548}
549
550fn attach_timeout(output: Output, mut timeout: Duration) -> Result<Output> {
551 if timeout.is_zero() {
552 return StatementTimeoutSnafu.fail();
553 }
554
555 let output = match output.data {
556 OutputData::AffectedRows(_) | OutputData::RecordBatches(_) => output,
557 OutputData::Stream(mut stream) => {
558 let schema = stream.schema();
559 let s = Box::pin(stream! {
560 let mut start = tokio::time::Instant::now();
561 while let Some(item) = tokio::time::timeout(timeout, stream.next()).await.map_err(|_| StreamTimeoutSnafu.build())? {
562 yield item;
563
564 let now = tokio::time::Instant::now();
565 timeout = timeout.checked_sub(now - start).unwrap_or(Duration::ZERO);
566 start = now;
567 if timeout.is_zero() {
569 StreamTimeoutSnafu.fail()?;
570 }
571 }
572 }) as Pin<Box<dyn Stream<Item = _> + Send>>;
573 let stream = RecordBatchStreamWrapper {
574 schema,
575 stream: s,
576 output_ordering: None,
577 metrics: Default::default(),
578 span: Span::current(),
579 };
580 Output::new(OutputData::Stream(Box::pin(stream)), output.meta)
581 }
582 };
583
584 Ok(output)
585}
586
587impl Instance {
588 #[tracing::instrument(skip_all, name = "SqlQueryHandler::do_analyze_stream_query")]
589 async fn do_analyze_stream_query_inner(
590 &self,
591 query: &str,
592 query_ctx: QueryContextRef,
593 ) -> Result<Output> {
594 ensure!(!self.is_suspended(), error::SuspendedSnafu);
595
596 let query_interceptor_opt = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
597 let query_interceptor = query_interceptor_opt.as_ref();
598 let query = query_interceptor.pre_parsing(query, query_ctx.clone())?;
599 let mut stmts = parse_stmt(query.as_ref(), query_ctx.sql_dialect())
600 .and_then(|stmts| query_interceptor.post_parsing(stmts, query_ctx.clone()))?;
601
602 ensure!(
603 stmts.len() == 1,
604 InvalidSqlSnafu {
605 err_msg: "only single EXPLAIN ANALYZE VERBOSE statement is supported"
606 }
607 );
608 let mut stmt = stmts.remove(0);
609 validate_analyze_stream_statement(&mut stmt)?;
610 query_ctx.set_explain_format(AnalyzeFormat::JSON.to_string());
611
612 let checker_ref = self.plugins.get::<PermissionCheckerRef>();
613 checker_ref
614 .as_ref()
615 .check_permission_with_context(
616 query_ctx.current_user(),
617 PermissionReq::SqlStatement(&stmt),
618 Some(&query_ctx.current_schema()),
619 )
620 .context(PermissionSnafu)?;
621 check_permission(self.plugins.clone(), &stmt, &query_ctx)?;
622 let catalog_name = query_ctx.current_catalog().to_string();
623 let schema_name = query_ctx.current_schema();
624 let slow_query_timer = self.statement_slow_query_timer(&stmt, schema_name.clone());
625 let ticket = self.process_manager.register_query(
626 catalog_name,
627 vec![schema_name],
628 stmt.to_string(),
629 query_ctx.conn_info().to_string(),
630 Some(query_ctx.process_id()),
631 slow_query_timer,
632 );
633 let query_fut =
634 self.exec_statement_with_timeout(stmt, query_ctx.clone(), query_interceptor);
635 let output = CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
636 .await
637 .map_err(|_| error::CancelledSnafu.build())??;
638 let Output { meta, data } = output;
639 let data = match data {
640 OutputData::Stream(stream) => OutputData::Stream(Box::pin(
641 CancellableStreamWrapper::new_cancel_on_drop(stream, ticket),
642 )),
643 other => other,
644 };
645 query_interceptor.post_execute(Output { data, meta }, query_ctx)
646 }
647
648 #[tracing::instrument(skip_all, name = "SqlQueryHandler::do_query")]
649 async fn do_query_inner(&self, query: &str, query_ctx: QueryContextRef) -> Vec<Result<Output>> {
650 if self.is_suspended() {
651 return vec![error::SuspendedSnafu {}.fail()];
652 }
653
654 let query_interceptor_opt = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
655 let query_interceptor = query_interceptor_opt.as_ref();
656 let query = match query_interceptor.pre_parsing(query, query_ctx.clone()) {
657 Ok(q) => q,
658 Err(e) => return vec![Err(e)],
659 };
660
661 let checker_ref = self.plugins.get::<PermissionCheckerRef>();
662 let checker = checker_ref.as_ref();
663
664 match parse_stmt(query.as_ref(), query_ctx.sql_dialect())
665 .and_then(|stmts| query_interceptor.post_parsing(stmts, query_ctx.clone()))
666 {
667 Ok(stmts) => {
668 if stmts.is_empty() {
669 return vec![
670 InvalidSqlSnafu {
671 err_msg: "empty statements",
672 }
673 .fail(),
674 ];
675 }
676
677 let mut results = Vec::with_capacity(stmts.len());
678 for stmt in stmts {
679 if let Err(e) = checker
680 .check_permission_with_context(
681 query_ctx.current_user(),
682 PermissionReq::SqlStatement(&stmt),
683 Some(&query_ctx.current_schema()),
684 )
685 .context(PermissionSnafu)
686 {
687 results.push(Err(e));
688 break;
689 }
690
691 match self.query_statement(stmt.clone(), query_ctx.clone()).await {
692 Ok(output) => {
693 let output_result =
694 query_interceptor.post_execute(output, query_ctx.clone());
695 results.push(output_result);
696 }
697 Err(e) => {
698 if e.status_code().should_log_error() {
699 error!(e; "Failed to execute query: {stmt}");
700 } else {
701 debug!("Failed to execute query: {stmt}, {e}");
702 }
703 results.push(Err(e));
704 break;
705 }
706 }
707 }
708 results
709 }
710 Err(e) => {
711 vec![Err(e)]
712 }
713 }
714 }
715
716 async fn exec_plan(&self, plan: LogicalPlan, query_ctx: QueryContextRef) -> Result<Output> {
717 self.query_engine
718 .execute(plan, query_ctx)
719 .await
720 .context(ExecLogicalPlanSnafu)
721 }
722
723 async fn exec_plan_with_timeout(
724 &self,
725 plan: LogicalPlan,
726 query_ctx: QueryContextRef,
727 ) -> Result<Output> {
728 let timeout = derive_timeout_for_plan(&plan, &query_ctx);
729 match timeout {
730 Some(timeout) => {
731 let start = tokio::time::Instant::now();
732 let output = tokio::time::timeout(timeout, self.exec_plan(plan, query_ctx))
733 .await
734 .map_err(|_| StatementTimeoutSnafu.build())??;
735 let output = map_query_output(output)?;
736 let remaining_timeout = timeout.checked_sub(start.elapsed()).unwrap_or_default();
737 attach_timeout(output, remaining_timeout)
738 }
739 None => self
740 .exec_plan(plan, query_ctx)
741 .await
742 .and_then(map_query_output),
743 }
744 }
745
746 async fn do_exec_plan_inner(
747 &self,
748 plan: LogicalPlan,
749 stmt: Option<Statement>,
750 query_ctx: QueryContextRef,
751 ) -> Result<Output> {
752 ensure!(!self.is_suspended(), error::SuspendedSnafu);
753
754 let query_interceptor_opt = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
755 let query_interceptor = query_interceptor_opt.as_ref();
756
757 query_interceptor.pre_execute(stmt.as_ref(), Some(&plan), query_ctx.clone())?;
758
759 let query = stmt
760 .as_ref()
761 .map(|s| s.to_string())
762 .unwrap_or_else(|| plan.display_indent().to_string());
763
764 let plan_is_readonly = is_readonly_plan(&plan);
765 let result = if should_track_plan_process(stmt.as_ref(), &plan) {
766 let catalog_name = query_ctx.current_catalog().to_string();
767 let schema_name = query_ctx.current_schema();
768 let slow_query_timer = if plan_is_readonly {
769 self.slow_query_options
770 .enable
771 .then(|| self.event_recorder.clone())
772 .flatten()
773 .map(|event_recorder| {
774 SlowQueryTimer::new(
775 CatalogQueryStatement::Plan(query.clone()),
776 schema_name.clone(),
777 self.slow_query_options.threshold,
778 self.slow_query_options.sample_ratio,
779 self.slow_query_options.record_type,
780 event_recorder,
781 )
782 })
783 } else {
784 None
785 };
786
787 let ticket = self.process_manager.register_query(
788 catalog_name,
789 vec![schema_name],
790 query,
791 query_ctx.conn_info().to_string(),
792 Some(query_ctx.process_id()),
793 slow_query_timer,
794 );
795
796 let query_fut = self.exec_plan_with_timeout(plan, query_ctx.clone());
797
798 CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
799 .await
800 .map_err(|_| error::CancelledSnafu.build())?
801 .map(|output| {
802 let Output { meta, data } = output;
803
804 let data = match data {
805 OutputData::Stream(stream) => OutputData::Stream(Box::pin(
806 CancellableStreamWrapper::new(stream, ticket),
807 )),
808 other => other,
809 };
810 Output { data, meta }
811 })
812 } else {
813 self.exec_plan_with_timeout(plan, query_ctx.clone()).await
814 };
815
816 result.and_then(|output| query_interceptor.post_execute(output, query_ctx))
817 }
818
819 #[tracing::instrument(skip_all, name = "SqlQueryHandler::do_promql_query")]
820 async fn do_promql_query_inner(
821 &self,
822 query: &PromQuery,
823 query_ctx: QueryContextRef,
824 ) -> Vec<Result<Output>> {
825 if self.is_suspended() {
826 return vec![error::SuspendedSnafu {}.fail()];
827 }
828
829 let result = PrometheusHandler::do_query(self, query, query_ctx)
831 .await
832 .with_context(|_| ExecutePromqlSnafu {
833 query: format!("{query:?}"),
834 });
835 vec![result]
836 }
837
838 async fn do_describe_inner(
839 &self,
840 stmt: Statement,
841 query_ctx: QueryContextRef,
842 ) -> Result<Option<DescribeResult>> {
843 ensure!(!self.is_suspended(), error::SuspendedSnafu);
844
845 let is_inner_plannable = |s: &Statement| {
849 matches!(
850 s,
851 Statement::Insert(_) | Statement::Query(_) | Statement::Delete(_)
852 )
853 };
854 let plannable = is_inner_plannable(&stmt)
855 || matches!(&stmt, Statement::Explain(explain) if is_inner_plannable(explain.statement.as_ref()));
856
857 if plannable {
858 self.plugins
859 .get::<PermissionCheckerRef>()
860 .as_ref()
861 .check_permission_with_context(
862 query_ctx.current_user(),
863 PermissionReq::SqlStatement(&stmt),
864 Some(&query_ctx.current_schema()),
865 )
866 .context(PermissionSnafu)?;
867
868 let plan = self
869 .query_engine
870 .planner()
871 .plan(&QueryStatement::Sql(stmt), query_ctx.clone())
872 .await
873 .context(PlanStatementSnafu)?;
874 self.query_engine
875 .describe(plan, query_ctx)
876 .await
877 .map(Some)
878 .context(error::DescribeStatementSnafu)
879 } else {
880 Ok(None)
881 }
882 }
883
884 async fn is_valid_schema_inner(&self, catalog: &str, schema: &str) -> Result<bool> {
885 self.catalog_manager
886 .schema_exists(catalog, schema, None)
887 .await
888 .context(error::CatalogSnafu)
889 }
890}
891
892#[async_trait]
893impl SqlQueryHandler for Instance {
894 async fn do_query(
895 &self,
896 query: &str,
897 query_ctx: QueryContextRef,
898 ) -> Vec<server_error::Result<Output>> {
899 self.do_query_inner(query, query_ctx)
900 .await
901 .into_iter()
902 .map(|result| result.map_err(BoxedError::new).context(ExecuteQuerySnafu))
903 .collect()
904 }
905
906 async fn do_analyze_stream_query(
907 &self,
908 query: &str,
909 query_ctx: QueryContextRef,
910 ) -> server_error::Result<Output> {
911 self.do_analyze_stream_query_inner(query, query_ctx)
912 .await
913 .map_err(BoxedError::new)
914 .context(ExecuteQuerySnafu)
915 }
916
917 async fn do_exec_plan(
918 &self,
919 plan: LogicalPlan,
920 stmt: Option<Statement>,
921 query_ctx: QueryContextRef,
922 ) -> server_error::Result<Output> {
923 self.do_exec_plan_inner(plan, stmt, query_ctx)
924 .await
925 .map_err(BoxedError::new)
926 .context(server_error::ExecutePlanSnafu)
927 }
928
929 async fn do_promql_query(
930 &self,
931 query: &PromQuery,
932 query_ctx: QueryContextRef,
933 ) -> Vec<server_error::Result<Output>> {
934 self.do_promql_query_inner(query, query_ctx)
935 .await
936 .into_iter()
937 .map(|result| result.map_err(BoxedError::new).context(ExecuteQuerySnafu))
938 .collect()
939 }
940
941 async fn do_describe(
942 &self,
943 stmt: Statement,
944 query_ctx: QueryContextRef,
945 ) -> server_error::Result<Option<DescribeResult>> {
946 self.do_describe_inner(stmt, query_ctx)
947 .await
948 .map_err(BoxedError::new)
949 .context(server_error::DescribeStatementSnafu)
950 }
951
952 async fn is_valid_schema(&self, catalog: &str, schema: &str) -> server_error::Result<bool> {
953 self.is_valid_schema_inner(catalog, schema)
954 .await
955 .map_err(BoxedError::new)
956 .context(server_error::CheckDatabaseValiditySnafu)
957 }
958}
959
960pub(crate) fn map_query_output(output: Output) -> Result<Output> {
962 output
963 .map_dictionary_to_values()
964 .context(CollectRecordbatchSnafu)
965}
966
967pub fn attach_timer(output: Output, timer: HistogramTimer) -> Output {
969 match output.data {
970 OutputData::AffectedRows(_) | OutputData::RecordBatches(_) => output,
971 OutputData::Stream(stream) => {
972 let stream = OnDone::new(stream, move || {
973 timer.observe_duration();
974 });
975 Output::new(OutputData::Stream(Box::pin(stream)), output.meta)
976 }
977 }
978}
979
980#[async_trait]
981impl PrometheusHandler for Instance {
982 #[tracing::instrument(skip_all)]
983 async fn do_query(
984 &self,
985 query: &PromQuery,
986 query_ctx: QueryContextRef,
987 ) -> server_error::Result<Output> {
988 let interceptor = self
989 .plugins
990 .get::<PromQueryInterceptorRef<server_error::Error>>();
991
992 self.plugins
993 .get::<PermissionCheckerRef>()
994 .as_ref()
995 .check_permission(query_ctx.current_user(), PermissionReq::PromQuery)
996 .context(AuthSnafu)?;
997
998 let stmt = QueryLanguageParser::parse_promql(query, &query_ctx).with_context(|_| {
999 ParsePromQLSnafu {
1000 query: query.clone(),
1001 }
1002 })?;
1003
1004 let plan = self
1005 .statement_executor
1006 .plan(&stmt, query_ctx.clone())
1007 .await
1008 .map_err(BoxedError::new)
1009 .context(ExecuteQuerySnafu)?;
1010
1011 let QueryStatement::Promql(eval_stmt, _) = &stmt else {
1012 unreachable!("query is parsed from promql");
1013 };
1014
1015 interceptor.pre_execute(query, &eval_stmt.expr, Some(&plan), query_ctx.clone())?;
1016
1017 let query_statement = if let QueryStatement::Promql(eval_stmt, alias) = stmt {
1019 CatalogQueryStatement::Promql(eval_stmt, alias)
1020 } else {
1021 return UnexpectedResultSnafu {
1023 reason: "The query should always be promql.".to_string(),
1024 }
1025 .fail();
1026 };
1027 let raw_query = query_statement.to_string();
1028
1029 let slow_query_timer = self
1030 .slow_query_options
1031 .enable
1032 .then(|| self.event_recorder.clone())
1033 .flatten()
1034 .map(|event_recorder| {
1035 SlowQueryTimer::new(
1036 query_statement,
1037 query_ctx.current_schema(),
1038 self.slow_query_options.threshold,
1039 self.slow_query_options.sample_ratio,
1040 self.slow_query_options.record_type,
1041 event_recorder,
1042 )
1043 });
1044
1045 let ticket = self.process_manager.register_query(
1046 query_ctx.current_catalog().to_string(),
1047 vec![query_ctx.current_schema()],
1048 raw_query,
1049 query_ctx.conn_info().to_string(),
1050 Some(query_ctx.process_id()),
1051 slow_query_timer,
1052 );
1053
1054 let query_fut = self.statement_executor.exec_plan(plan, query_ctx.clone());
1055
1056 let output = CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
1057 .await
1058 .map_err(|_| servers::error::CancelledSnafu.build())?
1059 .map_err(BoxedError::new)
1060 .context(ExecuteQuerySnafu)?;
1061 let output = map_query_output(output)
1062 .map_err(BoxedError::new)
1063 .context(ExecuteQuerySnafu)?;
1064 let Output { meta, data } = output;
1065 let data = match data {
1066 OutputData::Stream(stream) => {
1067 OutputData::Stream(Box::pin(CancellableStreamWrapper::new(stream, ticket)))
1068 }
1069 other => other,
1070 };
1071 let output = Output { data, meta };
1072 Ok(interceptor.post_execute(output, query_ctx)?)
1073 }
1074
1075 async fn query_metric_names(
1076 &self,
1077 matchers: Vec<Matcher>,
1078 ctx: &QueryContextRef,
1079 ) -> server_error::Result<Vec<String>> {
1080 self.handle_query_metric_names(matchers, ctx)
1081 .await
1082 .map_err(BoxedError::new)
1083 .context(ExecuteQuerySnafu)
1084 }
1085
1086 async fn query_label_values(
1087 &self,
1088 metric: String,
1089 label_name: String,
1090 matchers: Vec<Matcher>,
1091 start: SystemTime,
1092 end: SystemTime,
1093 ctx: &QueryContextRef,
1094 ) -> server_error::Result<Vec<String>> {
1095 self.handle_query_label_values(metric, label_name, matchers, start, end, ctx)
1096 .await
1097 .map_err(BoxedError::new)
1098 .context(ExecuteQuerySnafu)
1099 }
1100
1101 fn catalog_manager(&self) -> CatalogManagerRef {
1102 self.catalog_manager.clone()
1103 }
1104}
1105
1106macro_rules! validate_db_permission {
1108 ($stmt: expr, $query_ctx: expr) => {
1109 if let Some(database) = &$stmt.database {
1110 validate_catalog_and_schema($query_ctx.current_catalog(), database, $query_ctx)
1111 .map_err(BoxedError::new)
1112 .context(SqlExecInterceptedSnafu)?;
1113 }
1114 };
1115}
1116
1117pub fn check_permission(
1118 plugins: Plugins,
1119 stmt: &Statement,
1120 query_ctx: &QueryContextRef,
1121) -> Result<()> {
1122 let need_validate = plugins
1123 .get::<QueryOptions>()
1124 .map(|opts| opts.disallow_cross_catalog_query)
1125 .unwrap_or_default();
1126
1127 if !need_validate {
1128 return Ok(());
1129 }
1130
1131 match stmt {
1132 Statement::Admin(_) => {}
1135 Statement::Query(_)
1137 | Statement::Explain(_)
1138 | Statement::Tql(_)
1139 | Statement::Delete(_)
1140 | Statement::DeclareCursor(_)
1141 | Statement::Copy(sql::statements::copy::Copy::CopyQueryTo(_)) => {}
1142 Statement::CreateDatabase(_)
1144 | Statement::ShowDatabases(_)
1145 | Statement::DropDatabase(_)
1146 | Statement::AlterDatabase(_)
1147 | Statement::DropFlow(_)
1148 | Statement::Use(_) => {}
1149 #[cfg(feature = "enterprise")]
1150 Statement::DropTrigger(_) => {}
1151 Statement::ShowCreateDatabase(stmt) => {
1152 validate_database(&stmt.database_name, query_ctx)?;
1153 }
1154 Statement::ShowCreateTable(stmt) => {
1155 validate_param(&stmt.table_name, query_ctx)?;
1156 }
1157 Statement::ShowCreateFlow(stmt) => {
1158 validate_flow(&stmt.flow_name, query_ctx)?;
1159 }
1160 #[cfg(feature = "enterprise")]
1161 Statement::ShowCreateTrigger(stmt) => {
1162 validate_param(&stmt.trigger_name, query_ctx)?;
1163 }
1164 Statement::ShowCreateView(stmt) => {
1165 validate_param(&stmt.view_name, query_ctx)?;
1166 }
1167 Statement::CreateExternalTable(stmt) => {
1168 validate_param(&stmt.name, query_ctx)?;
1169 }
1170 Statement::CreateFlow(stmt) => {
1171 validate_param(&stmt.sink_table_name, query_ctx)?;
1173 }
1174 #[cfg(feature = "enterprise")]
1175 Statement::CreateTrigger(stmt) => {
1176 validate_param(&stmt.trigger_name, query_ctx)?;
1177 }
1178 Statement::CreateView(stmt) => {
1179 validate_param(&stmt.name, query_ctx)?;
1180 }
1181 Statement::AlterTable(stmt) => {
1182 validate_param(stmt.table_name(), query_ctx)?;
1183 }
1184 #[cfg(feature = "enterprise")]
1185 Statement::AlterTrigger(_) => {}
1186 Statement::SetVariables(_) | Statement::ShowVariables(_) => {}
1188 Statement::ShowCharset(_) | Statement::ShowCollation(_) => {}
1190
1191 Statement::Comment(comment) => match &comment.object {
1192 CommentObject::Table(table) => validate_param(table, query_ctx)?,
1193 CommentObject::Column { table, .. } => validate_param(table, query_ctx)?,
1194 CommentObject::Flow(flow) => validate_flow(flow, query_ctx)?,
1195 },
1196
1197 Statement::Insert(insert) => {
1198 let name = insert.table_name().context(ParseSqlSnafu)?;
1199 validate_param(name, query_ctx)?;
1200 }
1201 Statement::CreateTable(stmt) => {
1202 validate_param(&stmt.name, query_ctx)?;
1203 }
1204 Statement::CreateTableLike(stmt) => {
1205 validate_param(&stmt.table_name, query_ctx)?;
1206 validate_param(&stmt.source_name, query_ctx)?;
1207 }
1208 Statement::DropTable(drop_stmt) => {
1209 for table_name in drop_stmt.table_names() {
1210 validate_param(table_name, query_ctx)?;
1211 }
1212 }
1213 Statement::UndropTable(stmt) => {
1214 validate_param(stmt.table_name(), query_ctx)?;
1215 }
1216 Statement::DropView(stmt) => {
1217 validate_param(&stmt.view_name, query_ctx)?;
1218 }
1219 Statement::ShowTables(stmt) => {
1220 validate_db_permission!(stmt, query_ctx);
1221 }
1222 Statement::ShowTableStatus(stmt) => {
1223 validate_db_permission!(stmt, query_ctx);
1224 }
1225 Statement::ShowColumns(stmt) => {
1226 validate_db_permission!(stmt, query_ctx);
1227 }
1228 Statement::ShowIndex(stmt) => {
1229 validate_db_permission!(stmt, query_ctx);
1230 }
1231 Statement::ShowRegion(stmt) => {
1232 validate_db_permission!(stmt, query_ctx);
1233 }
1234 Statement::ShowViews(stmt) => {
1235 validate_db_permission!(stmt, query_ctx);
1236 }
1237 Statement::ShowFlows(stmt) => {
1238 validate_db_permission!(stmt, query_ctx);
1239 }
1240 #[cfg(feature = "enterprise")]
1241 Statement::ShowTriggers(_stmt) => {
1242 }
1245 Statement::ShowStatus(_stmt) => {}
1246 Statement::ShowSearchPath(_stmt) => {}
1247 Statement::DescribeTable(stmt) => {
1248 validate_param(stmt.name(), query_ctx)?;
1249 }
1250 Statement::Copy(sql::statements::copy::Copy::CopyTable(stmt)) => match stmt {
1251 CopyTable::To(copy_table_to) => validate_param(©_table_to.table_name, query_ctx)?,
1252 CopyTable::From(copy_table_from) => {
1253 validate_param(©_table_from.table_name, query_ctx)?
1254 }
1255 },
1256 Statement::Copy(sql::statements::copy::Copy::CopyDatabase(copy_database)) => {
1257 match copy_database {
1258 CopyDatabase::To(stmt) => validate_database(&stmt.database_name, query_ctx)?,
1259 CopyDatabase::From(stmt) => validate_database(&stmt.database_name, query_ctx)?,
1260 }
1261 }
1262 Statement::TruncateTable(stmt) => {
1263 validate_param(stmt.table_name(), query_ctx)?;
1264 }
1265 Statement::FetchCursor(_) | Statement::CloseCursor(_) => {}
1267 Statement::Kill(_) => {}
1269 Statement::ShowProcesslist(_) => {}
1271 }
1272 Ok(())
1273}
1274
1275fn validate_param(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
1276 let (catalog, schema, _) = table_idents_to_full_name(name, query_ctx)
1277 .map_err(BoxedError::new)
1278 .context(ExternalSnafu)?;
1279
1280 validate_catalog_and_schema(&catalog, &schema, query_ctx)
1281 .map_err(BoxedError::new)
1282 .context(SqlExecInterceptedSnafu)
1283}
1284
1285fn validate_flow(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
1286 let catalog = match &name.0[..] {
1287 [_flow] => query_ctx.current_catalog().to_string(),
1288 [catalog, _flow] => catalog.to_string_unquoted(),
1289 _ => {
1290 return InvalidSqlSnafu {
1291 err_msg: format!(
1292 "expect flow name to be <catalog>.<flow_name> or <flow_name>, actual: {name}",
1293 ),
1294 }
1295 .fail();
1296 }
1297 };
1298
1299 let schema = query_ctx.current_schema();
1300
1301 validate_catalog_and_schema(&catalog, &schema, query_ctx)
1302 .map_err(BoxedError::new)
1303 .context(SqlExecInterceptedSnafu)
1304}
1305
1306fn validate_database(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
1307 let (catalog, schema) = match &name.0[..] {
1308 [schema] => (
1309 query_ctx.current_catalog().to_string(),
1310 schema.to_string_unquoted(),
1311 ),
1312 [catalog, schema] => (catalog.to_string_unquoted(), schema.to_string_unquoted()),
1313 _ => InvalidSqlSnafu {
1314 err_msg: format!(
1315 "expect database name to be <catalog>.<schema> or <schema>, actual: {name}",
1316 ),
1317 }
1318 .fail()?,
1319 };
1320
1321 validate_catalog_and_schema(&catalog, &schema, query_ctx)
1322 .map_err(BoxedError::new)
1323 .context(SqlExecInterceptedSnafu)
1324}
1325
1326fn is_readonly_plan(plan: &LogicalPlan) -> bool {
1327 !matches!(plan, LogicalPlan::Dml(_) | LogicalPlan::Ddl(_))
1328}
1329
1330fn should_track_statement_process(stmt: &Statement) -> bool {
1331 stmt.is_readonly()
1332 || matches!(stmt, Statement::Insert(insert) if insert.has_non_values_query_source())
1333}
1334
1335fn should_track_plan_process(stmt: Option<&Statement>, plan: &LogicalPlan) -> bool {
1336 is_readonly_plan(plan)
1337 || matches!(stmt, Some(Statement::Insert(insert)) if insert.has_non_values_query_source())
1338}
1339
1340#[cfg(test)]
1341mod tests {
1342 use std::collections::HashMap;
1343 use std::future::Future;
1344 use std::pin::Pin;
1345 use std::sync::atomic::{AtomicBool, Ordering};
1346 use std::sync::{Arc, Barrier};
1347 use std::task::{Context, Poll};
1348 use std::thread;
1349 use std::time::{Duration, Instant};
1350
1351 use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse};
1352 use catalog::process_manager::ProcessManager;
1353 use common_base::Plugins;
1354 use common_error::ext::{BoxedError, PlainError};
1355 use common_error::status_code::StatusCode;
1356 use common_meta::cache::LayeredCacheRegistryBuilder;
1357 use common_meta::kv_backend::memory::MemoryKvBackend;
1358 use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor};
1359 use common_meta::rpc::ddl::{SubmitDdlTaskRequest, SubmitDdlTaskResponse};
1360 use common_meta::rpc::procedure::{
1361 MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse,
1362 };
1363 use common_query::Output;
1364 use common_recordbatch::{
1365 OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream,
1366 };
1367 use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
1368 use datafusion_expr::dml::InsertOp;
1369 use datafusion_expr::{LogicalPlanBuilder, LogicalTableSource};
1370 use datatypes::prelude::ConcreteDataType;
1371 use datatypes::schema::{ColumnSchema, Schema as GtSchema, SchemaRef as GtSchemaRef};
1372 use query::query_engine::options::QueryOptions;
1373 use session::context::{Channel, ConnInfo, QueryContext, QueryContextBuilder};
1374 use snafu::{Location, Snafu};
1375 use sql::dialect::GreptimeDbDialect;
1376 use store_api::data_source::DataSource;
1377 use store_api::storage::ScanRequest;
1378 use strfmt::Format;
1379 use table::metadata::{FilterPushDownType, TableInfo, TableInfoBuilder, TableMetaBuilder};
1380 use table::test_util::EmptyTable;
1381 use table::{Table, TableRef};
1382 use tokio::sync::{mpsc, oneshot};
1383
1384 use super::*;
1385 use crate::frontend::FrontendOptions;
1386 use crate::instance::builder::FrontendBuilder;
1387
1388 fn parse_test_sql(sql: &str) -> Vec<Statement> {
1389 parse_stmt(sql, &GreptimeDbDialect {}).unwrap()
1390 }
1391
1392 #[test]
1393 fn test_validate_analyze_stream_statement_strictness() {
1394 for sql in [
1395 "select 1",
1396 "explain analyze select 1",
1397 "explain analyze verbose format text select 1",
1398 "explain analyze verbose format graphviz select 1",
1399 ] {
1400 let mut stmts = parse_test_sql(sql);
1401 assert!(
1402 validate_analyze_stream_statement(&mut stmts[0]).is_err(),
1403 "{sql}"
1404 );
1405 }
1406
1407 for sql in [
1408 "explain analyze verbose select 1",
1409 "explain analyze verbose format json select 1",
1410 ] {
1411 let mut stmts = parse_test_sql(sql);
1412 assert!(
1413 validate_analyze_stream_statement(&mut stmts[0]).is_ok(),
1414 "{sql}"
1415 );
1416 let Statement::Explain(explain) = &stmts[0] else {
1417 unreachable!();
1418 };
1419 assert!(explain.format.is_none());
1420 }
1421
1422 assert_eq!(
1423 parse_test_sql("explain analyze verbose select 1; select 2").len(),
1424 2
1425 );
1426 }
1427
1428 #[derive(Debug, Snafu)]
1429 enum TestError {
1430 #[snafu(display("Failed to build test cache registry"))]
1431 BuildCacheRegistry {
1432 source: cache::error::Error,
1433 #[snafu(implicit)]
1434 location: Location,
1435 },
1436
1437 #[snafu(display("Failed to build test table meta for table: {table_name}"))]
1438 BuildTableMeta {
1439 table_name: String,
1440 source: table::metadata::TableMetaBuilderError,
1441 #[snafu(implicit)]
1442 location: Location,
1443 },
1444
1445 #[snafu(display("Failed to build test table info for table: {table_name}"))]
1446 BuildTableInfo {
1447 table_name: String,
1448 source: table::metadata::TableInfoBuilderError,
1449 #[snafu(implicit)]
1450 location: Location,
1451 },
1452
1453 #[snafu(display("Failed to register test table: {table_name}"))]
1454 RegisterTable {
1455 table_name: String,
1456 source: catalog::error::Error,
1457 #[snafu(implicit)]
1458 location: Location,
1459 },
1460
1461 #[snafu(display("Failed to build test frontend instance"))]
1462 BuildFrontend {
1463 source: crate::error::Error,
1464 #[snafu(implicit)]
1465 location: Location,
1466 },
1467
1468 #[snafu(display("Expected exactly one output for SQL `{sql}`, got {actual}"))]
1469 UnexpectedOutputCount {
1470 sql: String,
1471 actual: usize,
1472 #[snafu(implicit)]
1473 location: Location,
1474 },
1475
1476 #[snafu(display("Failed to execute SQL `{sql}`"))]
1477 ExecuteSql {
1478 sql: String,
1479 source: crate::error::Error,
1480 #[snafu(implicit)]
1481 location: Location,
1482 },
1483
1484 #[snafu(display("Timed out waiting for insert-select start notification"))]
1485 InsertStartTimeout {
1486 source: tokio::time::error::Elapsed,
1487 #[snafu(implicit)]
1488 location: Location,
1489 },
1490
1491 #[snafu(display("Insert-select start notification channel closed"))]
1492 InsertStartChannelClosed {
1493 #[snafu(implicit)]
1494 location: Location,
1495 },
1496
1497 #[snafu(display("Failed to release blocking insert-select interceptor"))]
1498 ReleaseBlockedInsert {
1499 #[snafu(implicit)]
1500 location: Location,
1501 },
1502
1503 #[snafu(display("Timed out waiting for insert-select source to be polled"))]
1504 SourcePollTimeout {
1505 source: tokio::time::error::Elapsed,
1506 #[snafu(implicit)]
1507 location: Location,
1508 },
1509
1510 #[snafu(display("Insert-select source poll notification channel closed"))]
1511 SourcePollChannelClosed {
1512 source: oneshot::error::RecvError,
1513 #[snafu(implicit)]
1514 location: Location,
1515 },
1516
1517 #[snafu(display("Timed out waiting for insert task to finish"))]
1518 InsertTaskTimeout {
1519 source: tokio::time::error::Elapsed,
1520 #[snafu(implicit)]
1521 location: Location,
1522 },
1523
1524 #[snafu(display("Insert task panicked"))]
1525 InsertTaskPanic {
1526 source: tokio::task::JoinError,
1527 #[snafu(implicit)]
1528 location: Location,
1529 },
1530
1531 #[snafu(display("Expected insert-select to be cancelled"))]
1532 InsertSelectNotCancelled {
1533 #[snafu(implicit)]
1534 location: Location,
1535 },
1536 }
1537
1538 type TestResult<T> = std::result::Result<T, TestError>;
1539
1540 fn parse_one_sql(sql: &str) -> Statement {
1541 parse_stmt(sql, &GreptimeDbDialect {}).unwrap().remove(0)
1542 }
1543
1544 fn test_query_ctx(process_id: u32) -> QueryContextRef {
1545 Arc::new(
1546 QueryContextBuilder::default()
1547 .channel(Channel::Mysql)
1548 .conn_info(ConnInfo::new(None, Channel::Mysql))
1549 .process_id(process_id)
1550 .build(),
1551 )
1552 }
1553
1554 struct BlockingInsertSelectInterceptor {
1555 started_tx: mpsc::UnboundedSender<()>,
1556 finish_rx: std::sync::Mutex<Option<oneshot::Receiver<()>>>,
1557 }
1558
1559 impl BlockingInsertSelectInterceptor {
1560 fn new(started_tx: mpsc::UnboundedSender<()>, finish_rx: oneshot::Receiver<()>) -> Self {
1561 Self {
1562 started_tx,
1563 finish_rx: std::sync::Mutex::new(Some(finish_rx)),
1564 }
1565 }
1566 }
1567
1568 impl SqlQueryInterceptor for BlockingInsertSelectInterceptor {
1569 type Error = Error;
1570
1571 fn pre_execute(
1572 &self,
1573 statement: Option<&Statement>,
1574 _plan: Option<&LogicalPlan>,
1575 _query_ctx: QueryContextRef,
1576 ) -> Result<()> {
1577 let Some(Statement::Insert(insert)) = statement else {
1578 return Ok(());
1579 };
1580 if !insert.has_non_values_query_source() {
1581 return Ok(());
1582 }
1583
1584 let finish_rx = self.finish_rx.lock().unwrap().take().unwrap();
1585 let _ = self.started_tx.send(());
1586 tokio::task::block_in_place(|| {
1587 tokio::runtime::Handle::current()
1588 .block_on(finish_rx)
1589 .unwrap();
1590 });
1591 Ok(())
1592 }
1593 }
1594
1595 struct PendingRecordBatchStream {
1596 schema: GtSchemaRef,
1597 polled_tx: Option<oneshot::Sender<()>>,
1598 _finish_tx: oneshot::Sender<()>,
1599 finish_rx: Pin<Box<oneshot::Receiver<()>>>,
1600 }
1601
1602 impl RecordBatchStream for PendingRecordBatchStream {
1603 fn schema(&self) -> GtSchemaRef {
1604 self.schema.clone()
1605 }
1606
1607 fn output_ordering(&self) -> Option<&[OrderOption]> {
1608 None
1609 }
1610
1611 fn metrics(&self) -> Option<common_recordbatch::adapter::RecordBatchMetrics> {
1612 None
1613 }
1614 }
1615
1616 impl Stream for PendingRecordBatchStream {
1617 type Item = common_recordbatch::error::Result<RecordBatch>;
1618
1619 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
1620 if let Some(polled_tx) = self.polled_tx.take() {
1621 let _ = polled_tx.send(());
1622 }
1623
1624 match self.finish_rx.as_mut().poll(cx) {
1625 Poll::Ready(_) => Poll::Ready(None),
1626 Poll::Pending => Poll::Pending,
1627 }
1628 }
1629 }
1630
1631 impl Unpin for PendingRecordBatchStream {}
1632
1633 struct PendingDataSource {
1634 schema: GtSchemaRef,
1635 polled_tx: std::sync::Mutex<Option<oneshot::Sender<()>>>,
1636 }
1637
1638 impl DataSource for PendingDataSource {
1639 fn get_stream(
1640 &self,
1641 _request: ScanRequest,
1642 ) -> std::result::Result<SendableRecordBatchStream, BoxedError> {
1643 let (finish_tx, finish_rx) = oneshot::channel();
1644 let mut polled_tx = self.polled_tx.lock().map_err(|_| {
1645 BoxedError::new(PlainError::new(
1646 "pending data source lock poisoned".to_string(),
1647 StatusCode::Unexpected,
1648 ))
1649 })?;
1650 Ok(Box::pin(PendingRecordBatchStream {
1651 schema: self.schema.clone(),
1652 polled_tx: polled_tx.take(),
1653 _finish_tx: finish_tx,
1654 finish_rx: Box::pin(finish_rx),
1655 }))
1656 }
1657 }
1658
1659 struct NoopProcedureExecutor;
1660
1661 #[async_trait::async_trait]
1662 impl ProcedureExecutor for NoopProcedureExecutor {
1663 async fn submit_ddl_task(
1664 &self,
1665 _ctx: &ExecutorContext,
1666 _request: SubmitDdlTaskRequest,
1667 ) -> common_meta::error::Result<SubmitDdlTaskResponse> {
1668 common_meta::error::UnsupportedSnafu {
1669 operation: "submit_ddl_task",
1670 }
1671 .fail()
1672 }
1673
1674 async fn migrate_region(
1675 &self,
1676 _ctx: &ExecutorContext,
1677 _request: MigrateRegionRequest,
1678 ) -> common_meta::error::Result<MigrateRegionResponse> {
1679 common_meta::error::UnsupportedSnafu {
1680 operation: "migrate_region",
1681 }
1682 .fail()
1683 }
1684
1685 async fn reconcile(
1686 &self,
1687 _ctx: &ExecutorContext,
1688 _request: ReconcileRequest,
1689 ) -> common_meta::error::Result<ReconcileResponse> {
1690 common_meta::error::UnsupportedSnafu {
1691 operation: "reconcile",
1692 }
1693 .fail()
1694 }
1695
1696 async fn query_procedure_state(
1697 &self,
1698 _ctx: &ExecutorContext,
1699 _pid: &str,
1700 ) -> common_meta::error::Result<ProcedureStateResponse> {
1701 common_meta::error::UnsupportedSnafu {
1702 operation: "query_procedure_state",
1703 }
1704 .fail()
1705 }
1706
1707 async fn list_procedures(
1708 &self,
1709 _ctx: &ExecutorContext,
1710 ) -> common_meta::error::Result<ProcedureDetailResponse> {
1711 common_meta::error::UnsupportedSnafu {
1712 operation: "list_procedures",
1713 }
1714 .fail()
1715 }
1716 }
1717
1718 fn test_cache_registry(
1719 kv_backend: common_meta::kv_backend::KvBackendRef,
1720 ) -> TestResult<common_meta::cache::LayeredCacheRegistryRef> {
1721 Ok(Arc::new(
1722 cache::with_default_composite_cache_registry(
1723 LayeredCacheRegistryBuilder::default()
1724 .add_cache_registry(cache::build_fundamental_cache_registry(kv_backend)),
1725 )
1726 .context(BuildCacheRegistrySnafu)?
1727 .build(),
1728 ))
1729 }
1730
1731 fn test_table_info(table_id: u32, table_name: &str) -> TestResult<TableInfo> {
1732 let schema = Arc::new(GtSchema::new(vec![
1733 ColumnSchema::new("id", ConcreteDataType::int32_datatype(), false),
1734 ColumnSchema::new(
1735 "ts",
1736 ConcreteDataType::timestamp_millisecond_datatype(),
1737 false,
1738 )
1739 .with_time_index(true),
1740 ]));
1741 let table_meta = TableMetaBuilder::empty()
1742 .schema(schema)
1743 .primary_key_indices(vec![0])
1744 .value_indices(vec![1])
1745 .next_column_id(1024)
1746 .build()
1747 .with_context(|_| BuildTableMetaSnafu {
1748 table_name: table_name.to_string(),
1749 })?;
1750
1751 TableInfoBuilder::new(table_name, table_meta)
1752 .table_id(table_id)
1753 .build()
1754 .with_context(|_| BuildTableInfoSnafu {
1755 table_name: table_name.to_string(),
1756 })
1757 }
1758
1759 fn test_table(table_id: u32, table_name: &str) -> TestResult<table::TableRef> {
1760 let table_info = test_table_info(table_id, table_name)?;
1761 Ok(EmptyTable::from_table_info(&table_info))
1762 }
1763
1764 fn pending_table(
1765 table_id: u32,
1766 table_name: &str,
1767 polled_tx: oneshot::Sender<()>,
1768 ) -> TestResult<table::TableRef> {
1769 let table_info = test_table_info(table_id, table_name)?;
1770 let data_source = Arc::new(PendingDataSource {
1771 schema: table_info.meta.schema.clone(),
1772 polled_tx: std::sync::Mutex::new(Some(polled_tx)),
1773 });
1774
1775 Ok(Arc::new(Table::new(
1776 Arc::new(table_info),
1777 FilterPushDownType::Unsupported,
1778 data_source,
1779 )))
1780 }
1781
1782 async fn test_instance_with_tables(
1783 source_table: TableRef,
1784 target_table: TableRef,
1785 ) -> TestResult<Instance> {
1786 test_instance_with_plugins(source_table, target_table, Plugins::new()).await
1787 }
1788
1789 async fn test_instance_with_insert_select_interceptor(
1790 interceptor: SqlQueryInterceptorRef<Error>,
1791 ) -> TestResult<Instance> {
1792 let plugins = Plugins::new();
1793 plugins.insert::<SqlQueryInterceptorRef<Error>>(interceptor);
1794
1795 test_instance_with_plugins(
1796 test_table(1024, "source")?,
1797 test_table(1025, "target")?,
1798 plugins,
1799 )
1800 .await
1801 }
1802
1803 async fn test_instance_with_plugins(
1804 source_table: TableRef,
1805 target_table: TableRef,
1806 plugins: Plugins,
1807 ) -> TestResult<Instance> {
1808 let kv_backend = Arc::new(MemoryKvBackend::new());
1809 let process_manager = Arc::new(ProcessManager::new("test-frontend".to_string(), None));
1810 let catalog_manager = catalog::memory::MemoryCatalogManager::new_with_table(source_table);
1811 let target_table_name = "target";
1812 catalog_manager
1813 .register_table_sync(catalog::RegisterTableRequest {
1814 catalog: "greptime".to_string(),
1815 schema: "public".to_string(),
1816 table_name: target_table_name.to_string(),
1817 table_id: 1025,
1818 table: target_table,
1819 })
1820 .with_context(|_| RegisterTableSnafu {
1821 table_name: target_table_name.to_string(),
1822 })?;
1823 catalog_manager.register_process_list_table(process_manager.clone());
1824
1825 let cache_registry = test_cache_registry(kv_backend.clone())?;
1826
1827 FrontendBuilder::new(
1828 FrontendOptions::default(),
1829 kv_backend,
1830 cache_registry,
1831 catalog_manager,
1832 Arc::new(client::client_manager::NodeClients::default()),
1833 Arc::new(NoopProcedureExecutor),
1834 process_manager,
1835 )
1836 .with_plugin(plugins)
1837 .try_build()
1838 .await
1839 .context(BuildFrontendSnafu)
1840 }
1841
1842 async fn execute_one_sql(
1843 instance: &Instance,
1844 sql: &str,
1845 query_ctx: QueryContextRef,
1846 ) -> TestResult<Output> {
1847 let mut results = instance.do_query_inner(sql, query_ctx).await;
1848 ensure!(
1849 results.len() == 1,
1850 UnexpectedOutputCountSnafu {
1851 sql: sql.to_string(),
1852 actual: results.len(),
1853 }
1854 );
1855 results.remove(0).with_context(|_| ExecuteSqlSnafu {
1856 sql: sql.to_string(),
1857 })
1858 }
1859
1860 #[test]
1861 fn test_fast_legacy_check_deadlock_prevention() {
1862 let cache = DashMap::new();
1864
1865 cache.insert("metric1".to_string(), true); cache.insert("metric2".to_string(), false); cache.insert("metric3".to_string(), true); let metric1 = "metric1".to_string();
1872 let metric4 = "metric4".to_string();
1873 let names1 = vec![&metric1, &metric4];
1874 let result = fast_legacy_check(&cache, &names1);
1875 assert!(result.is_ok());
1876 assert_eq!(result.unwrap(), Some(true)); assert!(cache.contains_key("metric4"));
1880 assert!(*cache.get("metric4").unwrap().value());
1881
1882 let metric5 = "metric5".to_string();
1884 let metric6 = "metric6".to_string();
1885 let names2 = vec![&metric5, &metric6];
1886 let result = fast_legacy_check(&cache, &names2);
1887 assert!(result.is_ok());
1888 assert_eq!(result.unwrap(), None); let cache_incompatible = DashMap::new();
1892 cache_incompatible.insert("metric1".to_string(), true); cache_incompatible.insert("metric2".to_string(), false); let metric1_test = "metric1".to_string();
1895 let metric2_test = "metric2".to_string();
1896 let names3 = vec![&metric1_test, &metric2_test];
1897 let result = fast_legacy_check(&cache_incompatible, &names3);
1898 assert!(result.is_err()); let cache_concurrent = Arc::new(DashMap::new());
1904 cache_concurrent.insert("shared_metric".to_string(), true);
1905
1906 let num_threads = 8;
1907 let operations_per_thread = 100;
1908 let barrier = Arc::new(Barrier::new(num_threads));
1909 let success_flag = Arc::new(AtomicBool::new(true));
1910
1911 let handles: Vec<_> = (0..num_threads)
1912 .map(|thread_id| {
1913 let cache_clone = Arc::clone(&cache_concurrent);
1914 let barrier_clone = Arc::clone(&barrier);
1915 let success_flag_clone = Arc::clone(&success_flag);
1916
1917 thread::spawn(move || {
1918 barrier_clone.wait();
1920
1921 let start_time = Instant::now();
1922 for i in 0..operations_per_thread {
1923 let shared_metric = "shared_metric".to_string();
1925 let new_metric = format!("thread_{}_metric_{}", thread_id, i);
1926 let names = vec![&shared_metric, &new_metric];
1927
1928 match fast_legacy_check(&cache_clone, &names) {
1929 Ok(_) => {}
1930 Err(_) => {
1931 success_flag_clone.store(false, Ordering::Relaxed);
1932 return;
1933 }
1934 }
1935
1936 if start_time.elapsed() > Duration::from_secs(10) {
1938 success_flag_clone.store(false, Ordering::Relaxed);
1939 return;
1940 }
1941 }
1942 })
1943 })
1944 .collect();
1945
1946 let start_time = Instant::now();
1948 for (i, handle) in handles.into_iter().enumerate() {
1949 let join_result = handle.join();
1950
1951 if start_time.elapsed() > Duration::from_secs(30) {
1953 panic!("Test timed out - possible deadlock detected!");
1954 }
1955
1956 if join_result.is_err() {
1957 panic!("Thread {} panicked during execution", i);
1958 }
1959 }
1960
1961 assert!(
1963 success_flag.load(Ordering::Relaxed),
1964 "Some operations failed"
1965 );
1966
1967 let final_count = cache_concurrent.len();
1969 assert!(
1970 final_count > 1 + num_threads * operations_per_thread / 2,
1971 "Expected more cache entries, got {}",
1972 final_count
1973 );
1974 }
1975
1976 #[test]
1977 fn test_should_track_statement_process() {
1978 assert!(should_track_statement_process(&parse_one_sql(
1979 "SELECT * FROM demo"
1980 )));
1981 assert!(should_track_statement_process(&parse_one_sql(
1982 "INSERT INTO demo SELECT * FROM source"
1983 )));
1984 assert!(!should_track_statement_process(&parse_one_sql(
1985 "INSERT INTO demo VALUES (1)"
1986 )));
1987 assert!(!should_track_statement_process(&parse_one_sql(
1988 "INSERT INTO demo VALUES (now())"
1989 )));
1990 }
1991
1992 #[test]
1993 fn test_should_track_plan_process() {
1994 let select_stmt = parse_one_sql("SELECT * FROM demo");
1995 let insert_select_stmt = parse_one_sql("INSERT INTO demo SELECT * FROM source");
1996 let insert_values_stmt = parse_one_sql("INSERT INTO demo VALUES (now())");
1997
1998 let empty_plan = LogicalPlanBuilder::empty(false).build().unwrap();
1999 assert!(should_track_plan_process(Some(&select_stmt), &empty_plan));
2000 assert!(should_track_plan_process(
2001 Some(&insert_select_stmt),
2002 &insert_dml_plan()
2003 ));
2004 assert!(!should_track_plan_process(
2005 Some(&insert_values_stmt),
2006 &insert_dml_plan()
2007 ));
2008 assert!(!should_track_plan_process(None, &insert_dml_plan()));
2009 }
2010
2011 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2012 async fn test_insert_select_is_visible_in_show_processlist() -> TestResult<()> {
2013 let insert_sql = "INSERT INTO target SELECT * FROM source";
2014 let (started_tx, mut started_rx) = mpsc::unbounded_channel();
2015 let (finish_tx, finish_rx) = oneshot::channel();
2016 let interceptor = Arc::new(BlockingInsertSelectInterceptor::new(started_tx, finish_rx));
2017 let instance = Arc::new(test_instance_with_insert_select_interceptor(interceptor).await?);
2018
2019 let insert_task = tokio::spawn({
2020 let instance = instance.clone();
2021 async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
2022 });
2023
2024 tokio::time::timeout(Duration::from_secs(5), started_rx.recv())
2025 .await
2026 .context(InsertStartTimeoutSnafu)?
2027 .context(InsertStartChannelClosedSnafu)?;
2028
2029 let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
2030 let process_list = output.data.pretty_print().await;
2031 assert!(
2032 process_list.contains(insert_sql),
2033 "process list did not contain running insert:\n{process_list}"
2034 );
2035
2036 finish_tx
2037 .send(())
2038 .map_err(|_| ReleaseBlockedInsertSnafu.build())?;
2039 insert_task.await.context(InsertTaskPanicSnafu)??;
2040
2041 Ok(())
2042 }
2043
2044 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2045 async fn test_kill_query_cancels_insert_select() -> TestResult<()> {
2046 assert_kill_cancels_insert_select("KILL QUERY 4242").await
2047 }
2048
2049 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2050 async fn test_kill_process_id_cancels_insert_select() -> TestResult<()> {
2051 assert_kill_cancels_insert_select("KILL 'test-frontend/4242'").await
2052 }
2053
2054 async fn assert_kill_cancels_insert_select(kill_sql: &str) -> TestResult<()> {
2055 let insert_sql = "INSERT INTO target SELECT * FROM source";
2056 let (source_polled_tx, source_polled_rx) = oneshot::channel();
2057 let instance = Arc::new(
2058 test_instance_with_tables(
2059 pending_table(1024, "source", source_polled_tx)?,
2060 test_table(1025, "target")?,
2061 )
2062 .await?,
2063 );
2064
2065 let insert_task = tokio::spawn({
2066 let instance = instance.clone();
2067 async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
2068 });
2069
2070 tokio::time::timeout(Duration::from_secs(5), source_polled_rx)
2071 .await
2072 .context(SourcePollTimeoutSnafu)?
2073 .context(SourcePollChannelClosedSnafu)?;
2074
2075 let output = execute_one_sql(&instance, kill_sql, test_query_ctx(43)).await?;
2076 assert!(matches!(output.data, OutputData::AffectedRows(1)));
2077
2078 let insert_result = tokio::time::timeout(Duration::from_secs(5), insert_task)
2079 .await
2080 .context(InsertTaskTimeoutSnafu)?
2081 .context(InsertTaskPanicSnafu)?;
2082 let err = match insert_result {
2083 Ok(_) => return InsertSelectNotCancelledSnafu.fail(),
2084 Err(TestError::ExecuteSql { source, .. }) => source,
2085 Err(err) => return Err(err),
2086 };
2087 assert_eq!(StatusCode::Cancelled, err.status_code());
2088
2089 let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
2090 let process_list = output.data.pretty_print().await;
2091 assert!(
2092 !process_list.contains(insert_sql),
2093 "process list still contains killed insert:\n{process_list}"
2094 );
2095
2096 Ok(())
2097 }
2098
2099 fn insert_dml_plan() -> LogicalPlan {
2100 let schema = SchemaRef::new(Schema::new(vec![Field::new(
2101 "value",
2102 DataType::Int64,
2103 true,
2104 )]));
2105 let target = Arc::new(LogicalTableSource::new(schema));
2106 let input = LogicalPlanBuilder::empty(false).build().unwrap();
2107
2108 LogicalPlanBuilder::insert_into(input, "demo", target, InsertOp::Append)
2109 .unwrap()
2110 .build()
2111 .unwrap()
2112 }
2113
2114 #[test]
2115 fn test_exec_validation() {
2116 let query_ctx = QueryContext::arc();
2117 let plugins: Plugins = Plugins::new();
2118 plugins.insert(QueryOptions {
2119 disallow_cross_catalog_query: true,
2120 });
2121
2122 let sql = r#"
2123 SELECT * FROM demo;
2124 EXPLAIN SELECT * FROM demo;
2125 CREATE DATABASE test_database;
2126 SHOW DATABASES;
2127 "#;
2128 let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
2129 assert_eq!(stmts.len(), 4);
2130 for stmt in stmts {
2131 let re = check_permission(plugins.clone(), &stmt, &query_ctx);
2132 re.unwrap();
2133 }
2134
2135 let sql = r#"
2136 SHOW CREATE TABLE demo;
2137 ALTER TABLE demo ADD COLUMN new_col INT;
2138 "#;
2139 let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
2140 assert_eq!(stmts.len(), 2);
2141 for stmt in stmts {
2142 let re = check_permission(plugins.clone(), &stmt, &query_ctx);
2143 re.unwrap();
2144 }
2145
2146 fn replace_test(template_sql: &str, plugins: Plugins, query_ctx: &QueryContextRef) {
2147 let right = vec![("", ""), ("", "public."), ("greptime.", "public.")];
2149 for (catalog, schema) in right {
2150 let sql = do_fmt(template_sql, catalog, schema);
2151 do_test(&sql, plugins.clone(), query_ctx, true);
2152 }
2153
2154 let wrong = vec![
2155 ("wrongcatalog.", "public."),
2156 ("wrongcatalog.", "wrongschema."),
2157 ];
2158 for (catalog, schema) in wrong {
2159 let sql = do_fmt(template_sql, catalog, schema);
2160 do_test(&sql, plugins.clone(), query_ctx, false);
2161 }
2162 }
2163
2164 fn do_fmt(template: &str, catalog: &str, schema: &str) -> String {
2165 let vars = HashMap::from([
2166 ("catalog".to_string(), catalog),
2167 ("schema".to_string(), schema),
2168 ]);
2169 template.format(&vars).unwrap()
2170 }
2171
2172 fn do_test(sql: &str, plugins: Plugins, query_ctx: &QueryContextRef, is_ok: bool) {
2173 let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
2174 let re = check_permission(plugins, stmt, query_ctx);
2175 if is_ok {
2176 re.unwrap();
2177 } else {
2178 assert!(re.is_err());
2179 }
2180 }
2181
2182 let sql = "INSERT INTO {catalog}{schema}monitor(host) VALUES ('host1');";
2184 replace_test(sql, plugins.clone(), &query_ctx);
2185
2186 let sql = r#"CREATE TABLE {catalog}{schema}demo(
2188 host STRING,
2189 ts TIMESTAMP,
2190 TIME INDEX (ts),
2191 PRIMARY KEY(host)
2192 ) engine=mito;"#;
2193 replace_test(sql, plugins.clone(), &query_ctx);
2194
2195 let sql = "DROP TABLE {catalog}{schema}demo;";
2197 replace_test(sql, plugins.clone(), &query_ctx);
2198
2199 let sql = "UNDROP TABLE {catalog}{schema}demo;";
2201 replace_test(sql, plugins.clone(), &query_ctx);
2202
2203 let sql = "SHOW TABLES FROM public";
2205 let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
2206 check_permission(plugins.clone(), &stmt[0], &query_ctx).unwrap();
2207
2208 let sql = "SHOW TABLES FROM private";
2209 let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
2210 let re = check_permission(plugins.clone(), &stmt[0], &query_ctx);
2211 assert!(re.is_ok());
2212
2213 let sql = "DESC TABLE {catalog}{schema}demo;";
2215 replace_test(sql, plugins.clone(), &query_ctx);
2216
2217 let comment_flow_cases = [
2218 ("COMMENT ON FLOW my_flow IS 'comment';", true),
2219 ("COMMENT ON FLOW greptime.my_flow IS 'comment';", true),
2220 ("COMMENT ON FLOW wrongcatalog.my_flow IS 'comment';", false),
2221 ];
2222 for (sql, is_ok) in comment_flow_cases {
2223 let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
2224 let result = check_permission(plugins.clone(), stmt, &query_ctx);
2225 assert_eq!(result.is_ok(), is_ok);
2226 }
2227
2228 let show_flow_cases = [
2229 ("SHOW CREATE FLOW my_flow;", true),
2230 ("SHOW CREATE FLOW greptime.my_flow;", true),
2231 ("SHOW CREATE FLOW wrongcatalog.my_flow;", false),
2232 ];
2233 for (sql, is_ok) in show_flow_cases {
2234 let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
2235 let result = check_permission(plugins.clone(), stmt, &query_ctx);
2236 assert_eq!(result.is_ok(), is_ok);
2237 }
2238 }
2239}