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