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