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