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 CancellableExportSource {
2767 schema: GtSchemaRef,
2768 channels: std::sync::Mutex<Option<(oneshot::Sender<()>, oneshot::Receiver<()>)>>,
2769 }
2770
2771 impl DataSource for CancellableExportSource {
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_cancels_ordinary_source() {
2792 let destination = common_test_util::temp_dir::create_temp_dir("metric_export_timeout");
2793 let (started_tx, started_rx) = oneshot::channel();
2794 let (mut 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(CancellableExportSource {
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 uri = reqwest::Url::from_directory_path(destination.path()).unwrap();
2838 let sql = format!(
2839 "COPY DATABASE greptime.public TO '{uri}' WITH (experimental_metric_export='true')"
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 tokio::time::timeout(Duration::from_secs(5), release_tx.closed())
2859 .await
2860 .unwrap();
2861 assert!(!destination.path().join("source.parquet").exists());
2862 server_task.abort();
2863 }
2864
2865 fn pending_table(
2866 table_id: u32,
2867 table_name: &str,
2868 polled_tx: oneshot::Sender<()>,
2869 ) -> TestResult<table::TableRef> {
2870 let table_info = test_table_info(table_id, table_name)?;
2871 let data_source = Arc::new(PendingDataSource {
2872 schema: table_info.meta.schema.clone(),
2873 polled_tx: std::sync::Mutex::new(Some(polled_tx)),
2874 });
2875
2876 Ok(Arc::new(Table::new(
2877 Arc::new(table_info),
2878 FilterPushDownType::Unsupported,
2879 data_source,
2880 )))
2881 }
2882
2883 async fn test_instance_with_tables(
2884 source_table: TableRef,
2885 target_table: TableRef,
2886 ) -> TestResult<Instance> {
2887 test_instance_with_plugins(source_table, target_table, Plugins::new()).await
2888 }
2889
2890 async fn test_instance_with_insert_select_interceptor(
2891 interceptor: SqlQueryInterceptorRef<Error>,
2892 ) -> TestResult<Instance> {
2893 let plugins = Plugins::new();
2894 plugins.insert::<SqlQueryInterceptorRef<Error>>(interceptor);
2895
2896 test_instance_with_plugins(
2897 test_table(1024, "source")?,
2898 test_table(1025, "target")?,
2899 plugins,
2900 )
2901 .await
2902 }
2903
2904 async fn test_instance_with_plugins(
2905 source_table: TableRef,
2906 target_table: TableRef,
2907 plugins: Plugins,
2908 ) -> TestResult<Instance> {
2909 test_instance_with_plugins_and_metric_names(source_table, target_table, plugins, None).await
2910 }
2911
2912 async fn test_instance_with_plugins_and_metric_names(
2913 source_table: TableRef,
2914 target_table: TableRef,
2915 plugins: Plugins,
2916 metric_names_table: Option<TableRef>,
2917 ) -> TestResult<Instance> {
2918 let catalog_manager = catalog::memory::MemoryCatalogManager::new_with_table(source_table);
2919 test_instance_with_catalog_manager(
2920 catalog_manager,
2921 target_table,
2922 plugins,
2923 metric_names_table,
2924 Arc::new(NoopProcedureExecutor),
2925 )
2926 .await
2927 }
2928
2929 async fn test_instance_with_catalog_manager(
2932 catalog_manager: Arc<catalog::memory::MemoryCatalogManager>,
2933 target_table: TableRef,
2934 plugins: Plugins,
2935 metric_names_table: Option<TableRef>,
2936 procedure_executor: ProcedureExecutorRef,
2937 ) -> TestResult<Instance> {
2938 let kv_backend = Arc::new(MemoryKvBackend::new());
2939 let process_manager = Arc::new(ProcessManager::new("test-frontend".to_string(), None));
2940 let target_table_name = "target";
2941 catalog_manager
2942 .register_table_sync(catalog::RegisterTableRequest {
2943 catalog: "greptime".to_string(),
2944 schema: "public".to_string(),
2945 table_name: target_table_name.to_string(),
2946 table_id: 1025,
2947 table: target_table,
2948 })
2949 .with_context(|_| RegisterTableSnafu {
2950 table_name: target_table_name.to_string(),
2951 })?;
2952 if let Some(table) = metric_names_table {
2953 catalog_manager
2954 .deregister_table_sync(catalog::DeregisterTableRequest {
2955 catalog: "greptime".to_string(),
2956 schema: "information_schema".to_string(),
2957 table_name: "tables".to_string(),
2958 })
2959 .unwrap();
2960 catalog_manager
2961 .register_table_sync(catalog::RegisterTableRequest {
2962 catalog: "greptime".to_string(),
2963 schema: "information_schema".to_string(),
2964 table_name: "tables".to_string(),
2965 table_id: 2048,
2966 table,
2967 })
2968 .unwrap();
2969 }
2970 catalog_manager.register_process_list_table(process_manager.clone());
2971
2972 let cache_registry = test_cache_registry(kv_backend.clone())?;
2973
2974 FrontendBuilder::new(
2975 FrontendOptions::default(),
2976 kv_backend,
2977 cache_registry,
2978 catalog_manager,
2979 Arc::new(client::client_manager::NodeClients::default()),
2980 procedure_executor,
2981 process_manager,
2982 )
2983 .with_plugin(plugins)
2984 .try_build()
2985 .await
2986 .context(BuildFrontendSnafu)
2987 }
2988
2989 async fn execute_one_sql(
2990 instance: &Instance,
2991 sql: &str,
2992 query_ctx: QueryContextRef,
2993 ) -> TestResult<Output> {
2994 let mut results = instance.do_query_inner(sql, query_ctx).await;
2995 ensure!(
2996 results.len() == 1,
2997 UnexpectedOutputCountSnafu {
2998 sql: sql.to_string(),
2999 actual: results.len(),
3000 }
3001 );
3002 results.remove(0).with_context(|_| ExecuteSqlSnafu {
3003 sql: sql.to_string(),
3004 })
3005 }
3006
3007 fn assert_permission_denied<T>(result: servers::error::Result<T>) {
3008 let err = match result {
3009 Ok(_) => panic!("request should be rejected"),
3010 Err(err) => err,
3011 };
3012 assert_eq!(StatusCode::PermissionDenied, err.status_code());
3013 }
3014
3015 fn assert_action_checked(
3016 checker: &RejectEndpointPermissionChecker,
3017 action: PermissionAction,
3018 targets: Option<PermissionTableTargets>,
3019 ) {
3020 assert_eq!(CheckedAction { action, targets }, checker.take_check());
3021 }
3022
3023 #[tokio::test]
3024 async fn database_export_authorizes_all_tables_before_preparation() -> TestResult<()> {
3025 struct ExportAcl(std::sync::Mutex<Vec<PermissionTableTargets>>);
3026 impl PermissionChecker for ExportAcl {
3027 fn check_permission(
3028 &self,
3029 _: UserInfoRef,
3030 req: PermissionReq,
3031 ) -> auth::error::Result<PermissionResp> {
3032 assert!(matches!(
3033 req,
3034 PermissionReq::SqlStatement(Statement::Copy(
3035 sql::statements::copy::Copy::CopyDatabase(
3036 sql::statements::copy::CopyDatabase::To(_)
3037 )
3038 ))
3039 ));
3040 Ok(PermissionResp::Allow)
3041 }
3042 fn check_permission_with_table_targets(
3043 &self,
3044 _: UserInfoRef,
3045 req: PermissionReq,
3046 targets: PermissionTableTargets,
3047 ) -> auth::error::Result<PermissionResp> {
3048 self.0.lock().unwrap().push(targets.clone());
3049 let PermissionTableTargets::Resolved(tables) = targets else {
3050 panic!("unresolved export")
3051 };
3052 if req.is_readonly() {
3053 Ok(if tables.iter().any(|t| t.table == "target") {
3055 PermissionResp::Reject
3056 } else {
3057 PermissionResp::Allow
3058 })
3059 } else {
3060 self.check_permission(QueryContext::arc().current_user(), req)
3061 }
3062 }
3063 }
3064 let checker = Arc::new(ExportAcl(Default::default()));
3065 let plugins = Plugins::new();
3066 plugins.insert::<PermissionCheckerRef>(checker.clone());
3067 let instance = test_instance_with_plugins(
3068 test_logical_table(1024, "source")?,
3069 test_table(1025, "target")?,
3070 plugins,
3071 )
3072 .await?;
3073 let req = table::requests::CopyDatabaseRequest {
3074 catalog_name: "greptime".into(),
3075 schema_name: "public".into(),
3076 location: "invalid-destination".into(),
3077 with: Default::default(),
3078 connection: Default::default(),
3079 time_range: None,
3080 };
3081 let result = instance
3082 .export_database_for_test(
3083 req.clone(),
3084 None,
3085 &tokio_util::sync::CancellationToken::new(),
3086 QueryContext::arc(),
3087 )
3088 .await;
3089 assert!(matches!(result, Err(Error::Permission { .. })));
3090 let expected = PermissionTableTargets::resolved(vec![
3091 PermissionTableTarget::new("greptime", "public", "source"),
3092 PermissionTableTarget::new("greptime", "public", "target"),
3093 ]);
3094 assert_eq!(*checker.0.lock().unwrap(), vec![expected.clone(), expected]);
3095 instance
3097 .plugins
3098 .map_mut::<PermissionCheckerRef, _, _>(|checker| {
3099 *checker.unwrap() = Arc::new(WriteOnlyPermissionChecker)
3100 });
3101 let result = instance
3102 .export_database_for_test(
3103 req,
3104 Some(&[]),
3105 &tokio_util::sync::CancellationToken::new(),
3106 QueryContext::arc(),
3107 )
3108 .await;
3109 assert!(matches!(result, Err(Error::Permission { .. })));
3110 Ok(())
3111 }
3112
3113 #[tokio::test]
3114 async fn test_prom_remote_read_with_custom_timestamp_and_value_columns() -> TestResult<()> {
3115 let schema = Arc::new(GtSchema::new(vec![
3116 ColumnSchema::new(
3117 "custom_ts",
3118 ConcreteDataType::timestamp_millisecond_datatype(),
3119 false,
3120 )
3121 .with_time_index(true),
3122 ColumnSchema::new("custom_value", ConcreteDataType::float64_datatype(), false),
3123 ]));
3124 let recordbatch = RecordBatch::new(
3125 schema,
3126 vec![
3127 Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as VectorRef,
3128 Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as VectorRef,
3129 ],
3130 )
3131 .unwrap();
3132 let instance = test_instance_with_tables(
3133 MemTable::table("custom_metric", recordbatch),
3134 test_table(1025, "target")?,
3135 )
3136 .await?;
3137
3138 let response = PromStoreProtocolHandler::read(
3139 &instance,
3140 ReadRequest {
3141 queries: vec![RemoteQuery {
3142 start_timestamp_ms: 1500,
3143 end_timestamp_ms: 2500,
3144 matchers: vec![LabelMatcher {
3145 r#type: PromMatcherType::Eq as i32,
3146 name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3147 value: "custom_metric".to_string(),
3148 }],
3149 ..Default::default()
3150 }],
3151 ..Default::default()
3152 },
3153 test_query_ctx(1),
3154 )
3155 .await
3156 .unwrap();
3157 let body = servers::prom_store::snappy_decompress(&response.body).unwrap();
3158 let response = ReadResponse::decode(body.as_slice()).unwrap();
3159
3160 assert_eq!(1, response.results.len());
3161 assert_eq!(1, response.results[0].timeseries.len());
3162 let timeseries = &response.results[0].timeseries[0];
3163 assert_eq!(
3164 vec![Label {
3165 name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3166 value: "custom_metric".to_string(),
3167 }],
3168 timeseries.labels
3169 );
3170 assert_eq!(
3171 vec![Sample {
3172 value: 2.0,
3173 timestamp: 2000,
3174 }],
3175 timeseries.samples
3176 );
3177
3178 Ok(())
3179 }
3180
3181 #[tokio::test]
3182 async fn test_prom_remote_read_prefers_default_value_column() -> TestResult<()> {
3183 let schema = Arc::new(GtSchema::new(vec![
3184 ColumnSchema::new(
3185 "custom_ts",
3186 ConcreteDataType::timestamp_millisecond_datatype(),
3187 false,
3188 )
3189 .with_time_index(true),
3190 ColumnSchema::new("extra_field", ConcreteDataType::float64_datatype(), false),
3191 ColumnSchema::new(
3192 greptime_value(),
3193 ConcreteDataType::float64_datatype(),
3194 false,
3195 ),
3196 ]));
3197 let recordbatch = RecordBatch::new(
3198 schema,
3199 vec![
3200 Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as VectorRef,
3201 Arc::new(Float64Vector::from_vec(vec![99.0, 99.0, 99.0])) as VectorRef,
3202 Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as VectorRef,
3203 ],
3204 )
3205 .unwrap();
3206 let instance = test_instance_with_tables(
3207 MemTable::table("multi_field_metric", recordbatch),
3208 test_table(1025, "target")?,
3209 )
3210 .await?;
3211
3212 let response = PromStoreProtocolHandler::read(
3213 &instance,
3214 ReadRequest {
3215 queries: vec![RemoteQuery {
3216 start_timestamp_ms: 1500,
3217 end_timestamp_ms: 2500,
3218 matchers: vec![LabelMatcher {
3219 r#type: PromMatcherType::Eq as i32,
3220 name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3221 value: "multi_field_metric".to_string(),
3222 }],
3223 ..Default::default()
3224 }],
3225 ..Default::default()
3226 },
3227 test_query_ctx(1),
3228 )
3229 .await
3230 .unwrap();
3231 let body = servers::prom_store::snappy_decompress(&response.body).unwrap();
3232 let response = ReadResponse::decode(body.as_slice()).unwrap();
3233
3234 assert_eq!(1, response.results.len());
3235 assert_eq!(1, response.results[0].timeseries.len());
3236 let timeseries = &response.results[0].timeseries[0];
3237 assert_eq!(
3238 vec![
3239 Label {
3240 name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3241 value: "multi_field_metric".to_string(),
3242 },
3243 Label {
3244 name: "extra_field".to_string(),
3245 value: "99".to_string(),
3246 },
3247 ],
3248 timeseries.labels
3249 );
3250 assert_eq!(
3251 vec![Sample {
3252 value: 2.0,
3253 timestamp: 2000,
3254 }],
3255 timeseries.samples
3256 );
3257
3258 Ok(())
3259 }
3260
3261 #[tokio::test]
3262 async fn test_prom_remote_read_rejects_ambiguous_value_columns() -> TestResult<()> {
3263 let schema = Arc::new(GtSchema::new(vec![
3264 ColumnSchema::new(
3265 "custom_ts",
3266 ConcreteDataType::timestamp_millisecond_datatype(),
3267 false,
3268 )
3269 .with_time_index(true),
3270 ColumnSchema::new("field_a", ConcreteDataType::float64_datatype(), false),
3271 ColumnSchema::new("field_b", ConcreteDataType::float64_datatype(), false),
3272 ]));
3273 let recordbatch = RecordBatch::new(
3274 schema,
3275 vec![
3276 Arc::new(TimestampMillisecondVector::from_vec(vec![1000])) as VectorRef,
3277 Arc::new(Float64Vector::from_vec(vec![1.0])) as VectorRef,
3278 Arc::new(Float64Vector::from_vec(vec![2.0])) as VectorRef,
3279 ],
3280 )
3281 .unwrap();
3282 let instance = test_instance_with_tables(
3283 MemTable::table("ambiguous_metric", recordbatch),
3284 test_table(1025, "target")?,
3285 )
3286 .await?;
3287
3288 let err = PromStoreProtocolHandler::read(
3289 &instance,
3290 ReadRequest {
3291 queries: vec![RemoteQuery {
3292 matchers: vec![LabelMatcher {
3293 r#type: PromMatcherType::Eq as i32,
3294 name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3295 value: "ambiguous_metric".to_string(),
3296 }],
3297 ..Default::default()
3298 }],
3299 ..Default::default()
3300 },
3301 test_query_ctx(1),
3302 )
3303 .await
3304 .err()
3305 .expect("ambiguous value columns should fail remote read");
3306
3307 assert_eq!(StatusCode::InvalidArguments, err.status_code());
3308 assert!(format!("{err:?}").contains("Ambiguous value column"));
3309
3310 Ok(())
3311 }
3312
3313 #[tokio::test]
3314 async fn test_event_recorder_is_exposed() -> TestResult<()> {
3315 let instance =
3316 test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
3317 .await?;
3318
3319 let _event_recorder = instance.event_recorder();
3320
3321 Ok(())
3322 }
3323
3324 #[tokio::test]
3325 async fn test_restricted_endpoint_handlers_check_permissions() -> TestResult<()> {
3326 let checker = Arc::new(RejectEndpointPermissionChecker::default());
3327 let plugins = Plugins::new();
3328 plugins.insert::<PermissionCheckerRef>(checker.clone());
3329 let instance = test_instance_with_plugins(
3330 test_table(1024, "denied")?,
3331 test_table(1025, "target")?,
3332 plugins,
3333 )
3334 .await?;
3335 let mut ctx = test_query_ctx(1);
3336 Arc::get_mut(&mut ctx).unwrap().set_extension(
3337 servers::http::jaeger::JAEGER_QUERY_TABLE_NAME_KEY,
3338 "denied".to_string(),
3339 );
3340 let jaeger_targets = Some(PermissionTableTargets::resolved(vec![
3341 PermissionTableTarget::new("greptime", "public", "denied"),
3342 ]));
3343
3344 assert_permission_denied(JaegerQueryHandler::get_services(&instance, ctx.clone()).await);
3345 assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
3346 assert_permission_denied(
3347 JaegerQueryHandler::get_operations(&instance, ctx.clone(), "service", None).await,
3348 );
3349 assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
3350 assert_permission_denied(
3351 JaegerQueryHandler::get_trace(&instance, ctx.clone(), "trace", None, None, None).await,
3352 );
3353 assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
3354 assert_permission_denied(
3355 JaegerQueryHandler::find_traces(
3356 &instance,
3357 ctx.clone(),
3358 servers::http::jaeger::QueryTraceParams {
3359 service_name: "service".to_string(),
3360 ..Default::default()
3361 },
3362 )
3363 .await,
3364 );
3365 assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets);
3366
3367 assert_permission_denied(
3368 PipelineHandler::get_pipeline_str(&instance, "pipeline", None, ctx.clone()).await,
3369 );
3370 assert_action_checked(&checker, PIPELINE_QUERY, None);
3371 assert_permission_denied(
3372 PipelineHandler::insert_pipeline(
3373 &instance,
3374 "pipeline",
3375 "application/yaml",
3376 "",
3377 ctx.clone(),
3378 )
3379 .await,
3380 );
3381 assert_action_checked(&checker, PIPELINE_INSERT, None);
3382 assert_permission_denied(
3383 PipelineHandler::delete_pipeline(&instance, "pipeline", None, ctx.clone()).await,
3384 );
3385 assert_action_checked(&checker, PIPELINE_DELETE, None);
3386 let app = axum::Router::new()
3387 .route(
3388 "/pipelines/_dryrun",
3389 axum::routing::post(servers::http::event::pipeline_dryrun),
3390 )
3391 .with_state(servers::http::event::LogState {
3392 log_handler: Arc::new(instance.clone()),
3393 log_validator: None,
3394 ingest_interceptor: None,
3395 })
3396 .layer(axum::Extension((*ctx).clone()));
3397 let response = app
3398 .oneshot(
3399 axum::http::Request::post("/pipelines/_dryrun")
3400 .header("content-type", "application/json")
3401 .body(axum::body::Body::from("{}"))
3402 .unwrap(),
3403 )
3404 .await
3405 .unwrap();
3406 assert_eq!(axum::http::StatusCode::FORBIDDEN, response.status());
3407 assert_action_checked(&checker, PIPELINE_QUERY, None);
3408
3409 assert_permission_denied(
3410 DashboardHandler::save(&instance, "dashboard", "{}", ctx.clone()).await,
3411 );
3412 assert_action_checked(&checker, DASHBOARD_SAVE, None);
3413 assert_permission_denied(DashboardHandler::list(&instance, ctx.clone()).await);
3414 assert_action_checked(&checker, DASHBOARD_QUERY, None);
3415 assert_permission_denied(
3416 DashboardHandler::delete(&instance, "dashboard", ctx.clone()).await,
3417 );
3418 assert_action_checked(&checker, DASHBOARD_DELETE, None);
3419
3420 Ok(())
3421 }
3422
3423 #[tokio::test]
3424 async fn test_write_only_ingestion_loads_named_pipeline() -> TestResult<()> {
3425 let plugins = Plugins::new();
3426 plugins.insert::<PermissionCheckerRef>(Arc::new(WriteOnlyPermissionChecker));
3427 let instance = test_instance_with_plugins(
3428 test_table(1024, "source")?,
3429 test_table(1025, "target")?,
3430 plugins,
3431 )
3432 .await?;
3433 instance
3434 .catalog_manager()
3435 .as_any()
3436 .downcast_ref::<catalog::memory::MemoryCatalogManager>()
3437 .unwrap()
3438 .register_table_sync(catalog::RegisterTableRequest {
3439 catalog: "greptime".to_string(),
3440 schema: DEFAULT_PRIVATE_SCHEMA_NAME.to_string(),
3441 table_name: "pipelines".to_string(),
3442 table_id: 2049,
3443 table: test_pipeline_table(),
3444 })
3445 .with_context(|_| RegisterTableSnafu {
3446 table_name: "pipelines".to_string(),
3447 })?;
3448 let ctx = test_query_ctx(1);
3449 let handler: PipelineHandlerRef = Arc::new(instance.clone());
3450
3451 handler
3452 .get_pipeline("pipeline", None, ctx.clone())
3453 .await
3454 .unwrap();
3455 assert_permission_denied(
3456 PipelineHandler::get_pipeline_str(&instance, "pipeline", None, ctx.clone()).await,
3457 );
3458
3459 let app = axum::Router::new()
3460 .route(
3461 "/pipelines/_dryrun",
3462 axum::routing::post(servers::http::event::pipeline_dryrun),
3463 )
3464 .with_state(servers::http::event::LogState {
3465 log_handler: handler,
3466 log_validator: None,
3467 ingest_interceptor: None,
3468 })
3469 .layer(axum::Extension((*ctx).clone()));
3470 let response = app
3471 .oneshot(
3472 axum::http::Request::post("/pipelines/_dryrun")
3473 .header("content-type", "application/json")
3474 .body(axum::body::Body::from("{}"))
3475 .unwrap(),
3476 )
3477 .await
3478 .unwrap();
3479 assert_eq!(axum::http::StatusCode::FORBIDDEN, response.status());
3480
3481 Ok(())
3482 }
3483
3484 #[tokio::test]
3485 async fn test_write_only_grpc_sql_is_checked_after_parsing() -> TestResult<()> {
3486 let plugins = Plugins::new();
3487 plugins.insert::<PermissionCheckerRef>(Arc::new(WriteOnlyPermissionChecker));
3488 let instance = test_instance_with_plugins(
3489 test_table(1024, "source")?,
3490 test_table(1025, "target")?,
3491 plugins,
3492 )
3493 .await?;
3494
3495 let insert = Request::Query(api::v1::QueryRequest {
3496 query: Some(Query::Sql(
3497 "INSERT INTO target SELECT * FROM source".to_string(),
3498 )),
3499 });
3500 servers::query_handler::grpc::GrpcQueryHandler::do_query(
3501 &instance,
3502 insert,
3503 QueryContext::arc(),
3504 )
3505 .await
3506 .unwrap();
3507
3508 let select = Request::Query(api::v1::QueryRequest {
3509 query: Some(Query::Sql("SELECT * FROM source".to_string())),
3510 });
3511 assert_permission_denied(
3512 servers::query_handler::grpc::GrpcQueryHandler::do_query(
3513 &instance,
3514 select,
3515 QueryContext::arc(),
3516 )
3517 .await,
3518 );
3519
3520 Ok(())
3521 }
3522
3523 #[tokio::test]
3524 async fn test_target_independent_checker_skips_target_resolution() -> TestResult<()> {
3525 let physical_table = "physical_metric";
3526 let checker = Arc::new(TargetIndependentPermissionChecker::default());
3527 let plugins = Plugins::new();
3528 plugins.insert::<PermissionCheckerRef>(checker.clone());
3529 let instance = test_instance_with_plugins(
3530 test_physical_table(1024, physical_table)?,
3531 test_table(1025, "target")?,
3532 plugins,
3533 )
3534 .await?;
3535
3536 let ctx = test_query_ctx(1);
3537 let physical_target = PermissionTableTarget::new("greptime", "public", physical_table);
3538 assert_eq!(
3539 PermissionTableTargets::Resolved(vec![physical_target.clone()]),
3540 instance
3541 .resolve_query_permission_targets(
3542 PermissionTableTargets::resolved(vec![physical_target]),
3543 &ctx,
3544 )
3545 .await
3546 .unwrap()
3547 );
3548 assert_eq!(
3549 vec![physical_table.to_string(), "target".to_string()],
3550 PrometheusHandler::filter_metadata_metric_names(
3551 &instance,
3552 vec![physical_table.to_string(), "target".to_string()],
3553 "public",
3554 &ctx,
3555 )
3556 .await
3557 .unwrap()
3558 );
3559 assert_eq!(1, checker.checks.load(atomic::Ordering::Relaxed));
3560
3561 Ok(())
3562 }
3563
3564 #[tokio::test]
3565 async fn test_query_permission_targets_are_deduplicated() -> TestResult<()> {
3566 let plugins = Plugins::new();
3567 plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
3568 let instance = test_instance_with_plugins(
3569 test_table(1024, "source")?,
3570 test_table(1025, "target")?,
3571 plugins,
3572 )
3573 .await?;
3574 let ctx = test_query_ctx(1);
3575 let target = PermissionTableTarget::new("greptime", "public", "target");
3576
3577 assert_eq!(
3578 PermissionTableTargets::Resolved(vec![target.clone()]),
3579 instance
3580 .resolve_query_permission_targets(
3581 PermissionTableTargets::resolved(vec![target.clone(), target]),
3582 &ctx,
3583 )
3584 .await
3585 .unwrap()
3586 );
3587
3588 Ok(())
3589 }
3590
3591 #[tokio::test]
3592 async fn test_physical_query_targets_fail_closed() -> TestResult<()> {
3593 let physical_table = "physical_metric";
3594 let plugins = Plugins::new();
3595 plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
3596 let instance = test_instance_with_plugins(
3597 test_physical_table(1024, physical_table)?,
3598 test_table(1025, "target")?,
3599 plugins,
3600 )
3601 .await?;
3602
3603 let ctx = test_query_ctx(1);
3604 let logical_target = PermissionTableTarget::new("greptime", "public", "target");
3605 assert_eq!(
3606 PermissionTableTargets::Resolved(vec![logical_target.clone()]),
3607 instance
3608 .resolve_query_permission_targets(
3609 PermissionTableTargets::resolved(vec![logical_target.clone()]),
3610 &ctx,
3611 )
3612 .await
3613 .unwrap()
3614 );
3615 let physical_target = PermissionTableTarget::new("greptime", "public", physical_table);
3616 assert_eq!(
3617 PermissionTableTargets::Unresolved,
3618 instance
3619 .resolve_query_permission_targets(
3620 PermissionTableTargets::resolved(
3621 vec![logical_target, physical_target.clone(),]
3622 ),
3623 &ctx,
3624 )
3625 .await
3626 .unwrap()
3627 );
3628 assert_eq!(
3629 vec!["target".to_string()],
3630 PrometheusHandler::filter_metadata_metric_names(
3631 &instance,
3632 vec!["target".to_string(), "denied".to_string()],
3633 "public",
3634 &ctx,
3635 )
3636 .await
3637 .unwrap()
3638 );
3639
3640 let query = PromQuery {
3641 query: physical_table.to_string(),
3642 ..Default::default()
3643 };
3644 let err = PrometheusHandler::check_query_target_permission(
3645 &instance,
3646 PermissionTableTargets::resolved(vec![physical_target]),
3647 &ctx,
3648 )
3649 .await
3650 .unwrap_err();
3651 assert_eq!(StatusCode::PermissionDenied, err.status_code());
3652 let err = PrometheusHandler::check_query_permission(
3653 &instance,
3654 std::slice::from_ref(&query),
3655 &ctx,
3656 )
3657 .await
3658 .unwrap_err();
3659 assert_eq!(StatusCode::PermissionDenied, err.status_code());
3660 let err = PrometheusHandler::do_query(&instance, &query, ctx.clone())
3661 .await
3662 .unwrap_err();
3663 assert_eq!(StatusCode::PermissionDenied, err.status_code());
3664
3665 for sql in [
3666 "SELECT * FROM physical_metric",
3667 "TQL EVAL (0, 10, '5s') physical_metric",
3668 "INSERT INTO target SELECT * FROM physical_metric",
3669 ] {
3670 let mut results = instance.do_query_inner(sql, ctx.clone()).await;
3671 assert_eq!(1, results.len(), "{sql}");
3672 let err = results.remove(0).unwrap_err();
3673 assert_eq!(StatusCode::PermissionDenied, err.status_code(), "{sql}");
3674 }
3675 let err = LogQueryHandler::query(
3676 &instance,
3677 LogQuery {
3678 table: TableName::new("greptime", "public", physical_table),
3679 ..Default::default()
3680 },
3681 ctx.clone(),
3682 )
3683 .await
3684 .unwrap_err();
3685 assert_eq!(StatusCode::PermissionDenied, err.status_code());
3686 let err = instance
3687 .do_describe_inner(parse_one_sql("SELECT * FROM physical_metric"), ctx.clone())
3688 .await
3689 .unwrap_err();
3690 assert_eq!(StatusCode::PermissionDenied, err.status_code());
3691
3692 let request = ReadRequest {
3693 queries: vec![RemoteQuery {
3694 matchers: vec![LabelMatcher {
3695 r#type: PromMatcherType::Eq as i32,
3696 name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
3697 value: physical_table.to_string(),
3698 }],
3699 ..Default::default()
3700 }],
3701 ..Default::default()
3702 };
3703 let Err(err) = PromStoreProtocolHandler::read(&instance, request, ctx.clone()).await else {
3704 panic!("physical remote-read target must be rejected");
3705 };
3706 assert_eq!(StatusCode::PermissionDenied, err.status_code());
3707
3708 let err = PrometheusHandler::query_label_values(
3709 &instance,
3710 physical_table.to_string(),
3711 "host".to_string(),
3712 vec![],
3713 SystemTime::UNIX_EPOCH,
3714 SystemTime::UNIX_EPOCH,
3715 &ctx,
3716 )
3717 .await
3718 .unwrap_err();
3719 assert_eq!(StatusCode::PermissionDenied, err.status_code());
3720
3721 Ok(())
3722 }
3723
3724 #[tokio::test]
3725 async fn test_non_exact_query_discovery_keeps_denied_targets_for_batch_check() -> TestResult<()>
3726 {
3727 let plugins = Plugins::new();
3728 plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
3729 let instance = test_instance_with_plugins_and_metric_names(
3730 test_logical_table(1024, "denied")?,
3731 test_logical_table(1025, "target")?,
3732 plugins,
3733 Some(test_metric_names_table()),
3734 )
3735 .await?;
3736 let ctx = test_query_ctx(1);
3737
3738 let mut metric_names = PrometheusHandler::query_metric_names(
3739 &instance,
3740 vec![Matcher::new(
3741 promql_parser::label::MatchOp::NotEqual,
3742 "__name__",
3743 "",
3744 )],
3745 "public",
3746 &ctx,
3747 )
3748 .await
3749 .unwrap();
3750 metric_names.sort_unstable();
3751 assert_eq!(
3752 vec!["denied".to_string(), "target".to_string()],
3753 metric_names
3754 );
3755
3756 let queries = metric_names
3757 .into_iter()
3758 .map(|query| PromQuery {
3759 query,
3760 ..Default::default()
3761 })
3762 .collect::<Vec<_>>();
3763 let err = PrometheusHandler::check_query_permission(&instance, &queries, &ctx)
3764 .await
3765 .unwrap_err();
3766 assert_eq!(StatusCode::PermissionDenied, err.status_code());
3767
3768 Ok(())
3769 }
3770
3771 #[test]
3772 fn test_fast_legacy_check_is_read_only() {
3773 let cache = DashMap::new();
3774 cache.insert("metric1".to_string(), true);
3775
3776 let names = vec!["metric1".to_string(), "metric2".to_string()];
3777 assert_eq!(Some(true), fast_legacy_check(&cache, &names).unwrap());
3778 assert!(!cache.contains_key("metric2"));
3779
3780 cache_legacy_mode(&cache, &names, true).unwrap();
3781 assert!(*cache.get("metric2").unwrap().value());
3782 assert!(cache_legacy_mode(&cache, &names, false).is_err());
3783 assert!(*cache.get("metric2").unwrap().value());
3784
3785 let cache_incompatible = DashMap::new();
3786 cache_incompatible.insert("metric1".to_string(), true);
3787 cache_incompatible.insert("metric2".to_string(), false);
3788 assert!(fast_legacy_check(&cache_incompatible, &names).is_err());
3789 }
3790
3791 #[test]
3792 fn test_should_track_statement_process() {
3793 assert!(should_track_statement_process(&parse_one_sql(
3794 "SELECT * FROM demo"
3795 )));
3796 assert!(should_track_statement_process(&parse_one_sql(
3797 "INSERT INTO demo SELECT * FROM source"
3798 )));
3799 assert!(!should_track_statement_process(&parse_one_sql(
3800 "INSERT INTO demo VALUES (1)"
3801 )));
3802 assert!(!should_track_statement_process(&parse_one_sql(
3803 "INSERT INTO demo VALUES (now())"
3804 )));
3805 }
3806
3807 #[test]
3808 fn test_should_track_plan_process() {
3809 let select_stmt = parse_one_sql("SELECT * FROM demo");
3810 let insert_select_stmt = parse_one_sql("INSERT INTO demo SELECT * FROM source");
3811 let insert_values_stmt = parse_one_sql("INSERT INTO demo VALUES (now())");
3812
3813 let empty_plan = LogicalPlanBuilder::empty(false).build().unwrap();
3814 assert!(should_track_plan_process(Some(&select_stmt), &empty_plan));
3815 assert!(should_track_plan_process(
3816 Some(&insert_select_stmt),
3817 &insert_dml_plan()
3818 ));
3819 assert!(!should_track_plan_process(
3820 Some(&insert_values_stmt),
3821 &insert_dml_plan()
3822 ));
3823 assert!(!should_track_plan_process(None, &insert_dml_plan()));
3824 }
3825
3826 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3827 async fn test_insert_select_is_visible_in_show_processlist() -> TestResult<()> {
3828 let insert_sql = "INSERT INTO target SELECT * FROM source";
3829 let (started_tx, mut started_rx) = mpsc::unbounded_channel();
3830 let (finish_tx, finish_rx) = oneshot::channel();
3831 let interceptor = Arc::new(BlockingInsertSelectInterceptor::new(started_tx, finish_rx));
3832 let instance = Arc::new(test_instance_with_insert_select_interceptor(interceptor).await?);
3833
3834 let insert_task = tokio::spawn({
3835 let instance = instance.clone();
3836 async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
3837 });
3838
3839 tokio::time::timeout(Duration::from_secs(5), started_rx.recv())
3840 .await
3841 .context(InsertStartTimeoutSnafu)?
3842 .context(InsertStartChannelClosedSnafu)?;
3843
3844 let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
3845 let process_list = output.data.pretty_print().await;
3846 assert!(
3847 process_list.contains(insert_sql),
3848 "process list did not contain running insert:\n{process_list}"
3849 );
3850
3851 finish_tx
3852 .send(())
3853 .map_err(|_| ReleaseBlockedInsertSnafu.build())?;
3854 insert_task.await.context(InsertTaskPanicSnafu)??;
3855
3856 Ok(())
3857 }
3858
3859 #[tokio::test]
3860 async fn test_show_processlist_catalog_scope() -> TestResult<()> {
3861 let instance =
3862 test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
3863 .await?;
3864 let _current_catalog = instance.process_manager().register_query(
3865 "greptime".to_string(),
3866 vec!["public".to_string()],
3867 "current_catalog_query".to_string(),
3868 String::new(),
3869 None,
3870 None,
3871 );
3872 let _other_catalog = instance.process_manager().register_query(
3873 "other".to_string(),
3874 vec!["public".to_string()],
3875 "other_catalog_query".to_string(),
3876 String::new(),
3877 None,
3878 None,
3879 );
3880
3881 for sql in ["SHOW PROCESSLIST", "SHOW FULL PROCESSLIST"] {
3882 let output = execute_one_sql(&instance, sql, test_query_ctx(43)).await?;
3883 let process_list = output.data.pretty_print().await;
3884 assert!(
3885 process_list.contains("current_catalog_query"),
3886 "{process_list}"
3887 );
3888 assert!(
3889 !process_list.contains("other_catalog_query"),
3890 "{process_list}"
3891 );
3892
3893 let admin_ctx = test_query_ctx(44);
3894 admin_ctx.set_current_user(Arc::new(AdminUserInfo));
3895 let output = execute_one_sql(&instance, sql, admin_ctx).await?;
3896 let process_list = output.data.pretty_print().await;
3897 assert!(
3898 process_list.contains("current_catalog_query"),
3899 "{process_list}"
3900 );
3901 assert!(
3902 process_list.contains("other_catalog_query"),
3903 "{process_list}"
3904 );
3905 }
3906
3907 Ok(())
3908 }
3909
3910 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3911 async fn test_kill_query_cancels_insert_select() -> TestResult<()> {
3912 assert_kill_cancels_insert_select("KILL QUERY 4242").await
3913 }
3914
3915 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3916 async fn test_kill_process_id_cancels_insert_select() -> TestResult<()> {
3917 assert_kill_cancels_insert_select("KILL 'test-frontend/4242'").await
3918 }
3919
3920 async fn assert_kill_cancels_insert_select(kill_sql: &str) -> TestResult<()> {
3921 let insert_sql = "INSERT INTO target SELECT * FROM source";
3922 let (source_polled_tx, source_polled_rx) = oneshot::channel();
3923 let instance = Arc::new(
3924 test_instance_with_tables(
3925 pending_table(1024, "source", source_polled_tx)?,
3926 test_table(1025, "target")?,
3927 )
3928 .await?,
3929 );
3930
3931 let insert_task = tokio::spawn({
3932 let instance = instance.clone();
3933 async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
3934 });
3935
3936 tokio::time::timeout(Duration::from_secs(5), source_polled_rx)
3937 .await
3938 .context(SourcePollTimeoutSnafu)?
3939 .context(SourcePollChannelClosedSnafu)?;
3940
3941 let output = execute_one_sql(&instance, kill_sql, test_query_ctx(43)).await?;
3942 assert!(matches!(output.data, OutputData::AffectedRows(1)));
3943
3944 let insert_result = tokio::time::timeout(Duration::from_secs(5), insert_task)
3945 .await
3946 .context(InsertTaskTimeoutSnafu)?
3947 .context(InsertTaskPanicSnafu)?;
3948 let err = match insert_result {
3949 Ok(_) => return InsertSelectNotCancelledSnafu.fail(),
3950 Err(TestError::ExecuteSql { source, .. }) => source,
3951 Err(err) => return Err(err),
3952 };
3953 assert_eq!(StatusCode::Cancelled, err.status_code());
3954
3955 let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
3956 let process_list = output.data.pretty_print().await;
3957 assert!(
3958 !process_list.contains(insert_sql),
3959 "process list still contains killed insert:\n{process_list}"
3960 );
3961
3962 Ok(())
3963 }
3964
3965 fn insert_dml_plan() -> LogicalPlan {
3966 let schema = SchemaRef::new(Schema::new(vec![Field::new(
3967 "value",
3968 DataType::Int64,
3969 true,
3970 )]));
3971 let target = Arc::new(LogicalTableSource::new(schema));
3972 let input = LogicalPlanBuilder::empty(false).build().unwrap();
3973
3974 LogicalPlanBuilder::insert_into(input, "demo", target, InsertOp::Append)
3975 .unwrap()
3976 .build()
3977 .unwrap()
3978 }
3979
3980 #[test]
3981 fn test_exec_validation() {
3982 let query_ctx = QueryContext::arc();
3983 let plugins: Plugins = Plugins::new();
3984 plugins.insert(QueryOptions {
3985 disallow_cross_catalog_query: true,
3986 });
3987
3988 let sql = r#"
3989 SELECT * FROM demo;
3990 EXPLAIN SELECT * FROM demo;
3991 CREATE DATABASE test_database;
3992 SHOW DATABASES;
3993 "#;
3994 let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
3995 assert_eq!(stmts.len(), 4);
3996 for stmt in stmts {
3997 let re = check_permission(plugins.clone(), &stmt, &query_ctx);
3998 re.unwrap();
3999 }
4000
4001 let sql = r#"
4002 SHOW CREATE TABLE demo;
4003 ALTER TABLE demo ADD COLUMN new_col INT;
4004 "#;
4005 let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
4006 assert_eq!(stmts.len(), 2);
4007 for stmt in stmts {
4008 let re = check_permission(plugins.clone(), &stmt, &query_ctx);
4009 re.unwrap();
4010 }
4011
4012 fn replace_test(template_sql: &str, plugins: Plugins, query_ctx: &QueryContextRef) {
4013 let right = vec![("", ""), ("", "public."), ("greptime.", "public.")];
4015 for (catalog, schema) in right {
4016 let sql = do_fmt(template_sql, catalog, schema);
4017 do_test(&sql, plugins.clone(), query_ctx, true);
4018 }
4019
4020 let wrong = vec![
4021 ("wrongcatalog.", "public."),
4022 ("wrongcatalog.", "wrongschema."),
4023 ];
4024 for (catalog, schema) in wrong {
4025 let sql = do_fmt(template_sql, catalog, schema);
4026 do_test(&sql, plugins.clone(), query_ctx, false);
4027 }
4028 }
4029
4030 fn do_fmt(template: &str, catalog: &str, schema: &str) -> String {
4031 let vars = HashMap::from([
4032 ("catalog".to_string(), catalog),
4033 ("schema".to_string(), schema),
4034 ]);
4035 template.format(&vars).unwrap()
4036 }
4037
4038 fn do_test(sql: &str, plugins: Plugins, query_ctx: &QueryContextRef, is_ok: bool) {
4039 let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
4040 let re = check_permission(plugins, stmt, query_ctx);
4041 if is_ok {
4042 re.unwrap();
4043 } else {
4044 assert!(re.is_err());
4045 }
4046 }
4047
4048 let sql = "INSERT INTO {catalog}{schema}monitor(host) VALUES ('host1');";
4050 replace_test(sql, plugins.clone(), &query_ctx);
4051
4052 let sql = r#"CREATE TABLE {catalog}{schema}demo(
4054 host STRING,
4055 ts TIMESTAMP,
4056 TIME INDEX (ts),
4057 PRIMARY KEY(host)
4058 ) engine=mito;"#;
4059 replace_test(sql, plugins.clone(), &query_ctx);
4060
4061 let sql = "DROP TABLE {catalog}{schema}demo;";
4063 replace_test(sql, plugins.clone(), &query_ctx);
4064
4065 #[cfg(feature = "enterprise")]
4067 {
4068 let sql = "UNDROP TABLE {catalog}{schema}demo;";
4069 replace_test(sql, plugins.clone(), &query_ctx);
4070 }
4071
4072 let sql = "SHOW TABLES FROM public";
4074 let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
4075 check_permission(plugins.clone(), &stmt[0], &query_ctx).unwrap();
4076
4077 let sql = "SHOW TABLES FROM private";
4078 let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
4079 let re = check_permission(plugins.clone(), &stmt[0], &query_ctx);
4080 assert!(re.is_ok());
4081
4082 let sql = "DESC TABLE {catalog}{schema}demo;";
4084 replace_test(sql, plugins.clone(), &query_ctx);
4085
4086 let comment_flow_cases = [
4087 ("COMMENT ON FLOW my_flow IS 'comment';", true),
4088 ("COMMENT ON FLOW greptime.my_flow IS 'comment';", true),
4089 ("COMMENT ON FLOW wrongcatalog.my_flow IS 'comment';", false),
4090 ];
4091 for (sql, is_ok) in comment_flow_cases {
4092 let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
4093 let result = check_permission(plugins.clone(), stmt, &query_ctx);
4094 assert_eq!(result.is_ok(), is_ok);
4095 }
4096
4097 let show_flow_cases = [
4098 ("SHOW CREATE FLOW my_flow;", true),
4099 ("SHOW CREATE FLOW greptime.my_flow;", true),
4100 ("SHOW CREATE FLOW wrongcatalog.my_flow;", false),
4101 ];
4102 for (sql, is_ok) in show_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
4109 #[tokio::test]
4112 async fn qx_152_drop_view_via_grpc_ddl_returns_error_not_panic() -> TestResult<()> {
4113 let instance =
4114 test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
4115 .await?;
4116
4117 let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4118 expr: Some(api::v1::ddl_request::Expr::DropView(
4119 api::v1::DropViewExpr {
4120 catalog_name: String::new(),
4121 schema_name: String::new(),
4122 view_name: "non_existent_view".to_string(),
4123 view_id: None,
4124 drop_if_exists: false,
4125 },
4126 )),
4127 });
4128
4129 let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
4130 &instance,
4131 request,
4132 QueryContext::arc(),
4133 )
4134 .await;
4135
4136 let err = match result {
4137 Ok(_) => panic!("DropView DDL request must return an error instead of panicking"),
4138 Err(err) => err,
4139 };
4140 assert_eq!(
4141 err.status_code(),
4142 StatusCode::TableNotFound,
4143 "dropping a non-existent view without IF EXISTS must report TableNotFound, got {err}"
4144 );
4145 Ok(())
4146 }
4147
4148 #[tokio::test]
4152 async fn qx_152_drop_view_if_exists_missing_view_via_grpc_ddl_succeeds() -> TestResult<()> {
4153 let catalog_manager =
4154 catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
4155 let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
4156 let instance = test_instance_with_catalog_manager(
4157 catalog_manager,
4158 test_table(1025, "target")?,
4159 Plugins::new(),
4160 None,
4161 procedure_executor.clone() as ProcedureExecutorRef,
4162 )
4163 .await?;
4164
4165 let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4166 expr: Some(api::v1::ddl_request::Expr::DropView(
4167 api::v1::DropViewExpr {
4168 catalog_name: String::new(),
4169 schema_name: String::new(),
4170 view_name: "non_existent_view".to_string(),
4171 view_id: None,
4172 drop_if_exists: true,
4173 },
4174 )),
4175 });
4176
4177 let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
4178 &instance,
4179 request,
4180 QueryContext::arc(),
4181 )
4182 .await;
4183
4184 let output = match result {
4185 Ok(output) => output,
4186 Err(err) => {
4187 panic!("DROP VIEW IF EXISTS on a missing view must succeed, got error: {err}")
4188 }
4189 };
4190 assert!(
4191 matches!(output.data, OutputData::AffectedRows(0)),
4192 "DROP VIEW IF EXISTS on a missing view must report 0 affected rows"
4193 );
4194 assert!(
4195 procedure_executor.submitted.lock().unwrap().is_empty(),
4196 "DROP VIEW IF EXISTS on a missing view must not submit a DDL task"
4197 );
4198 Ok(())
4199 }
4200
4201 #[tokio::test]
4204 async fn qx_152_drop_existing_view_via_grpc_ddl_succeeds() -> TestResult<()> {
4205 let catalog_manager =
4206 catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
4207 let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
4208 let instance = test_instance_with_catalog_manager(
4209 catalog_manager,
4210 test_table(1025, "target")?,
4211 Plugins::new(),
4212 None,
4213 procedure_executor.clone() as ProcedureExecutorRef,
4214 )
4215 .await?;
4216
4217 instance
4220 .table_metadata_manager()
4221 .schema_manager()
4222 .create(
4223 common_meta::key::schema_name::SchemaNameKey::new("greptime", "public"),
4224 None,
4225 true,
4226 )
4227 .await
4228 .unwrap();
4229
4230 let create_view_request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4231 expr: Some(api::v1::ddl_request::Expr::CreateView(
4232 api::v1::CreateViewExpr {
4233 catalog_name: String::new(),
4234 schema_name: String::new(),
4235 view_name: "my_view".to_string(),
4236 logical_plan: vec![1, 2, 3],
4237 create_if_not_exists: false,
4238 or_replace: false,
4239 table_names: vec![],
4240 columns: vec![],
4241 plan_columns: vec![],
4242 definition: "CREATE VIEW my_view AS SELECT * FROM source".to_string(),
4243 },
4244 )),
4245 });
4246
4247 let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
4248 &instance,
4249 create_view_request,
4250 QueryContext::arc(),
4251 )
4252 .await
4253 {
4254 Ok(output) => output,
4255 Err(err) => panic!("CREATE VIEW via gRPC DDL must succeed, got error: {err}"),
4256 };
4257 assert!(
4258 matches!(output.data, OutputData::AffectedRows(0)),
4259 "CREATE VIEW via gRPC DDL must report 0 affected rows"
4260 );
4261
4262 let view = instance
4264 .catalog_manager()
4265 .table("greptime", "public", "my_view", None)
4266 .await
4267 .unwrap()
4268 .expect("view should exist after CREATE VIEW");
4269 assert_eq!(view.table_info().table_type, TableType::View);
4270
4271 let drop_view_request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4272 expr: Some(api::v1::ddl_request::Expr::DropView(
4273 api::v1::DropViewExpr {
4274 catalog_name: String::new(),
4275 schema_name: String::new(),
4276 view_name: "my_view".to_string(),
4277 view_id: None,
4278 drop_if_exists: false,
4279 },
4280 )),
4281 });
4282
4283 let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
4284 &instance,
4285 drop_view_request,
4286 QueryContext::arc(),
4287 )
4288 .await
4289 {
4290 Ok(output) => output,
4291 Err(err) => panic!("DROP VIEW via gRPC DDL must succeed, got error: {err}"),
4292 };
4293 assert!(
4294 matches!(output.data, OutputData::AffectedRows(0)),
4295 "DROP VIEW via gRPC DDL must report 0 affected rows"
4296 );
4297
4298 assert!(
4300 instance
4301 .catalog_manager()
4302 .table("greptime", "public", "my_view", None)
4303 .await
4304 .unwrap()
4305 .is_none(),
4306 "view should be removed after DROP VIEW"
4307 );
4308
4309 let submitted = procedure_executor.submitted.lock().unwrap();
4310 assert_eq!(
4311 submitted.len(),
4312 2,
4313 "expected create and drop view tasks, got {submitted:?}"
4314 );
4315 assert!(matches!(&submitted[0], DdlTask::CreateView(_)));
4316 assert!(matches!(&submitted[1], DdlTask::DropView(_)));
4317 Ok(())
4318 }
4319
4320 #[tokio::test]
4323 async fn qx_153_create_table_with_non_timestamp_time_index_via_grpc_returns_error()
4324 -> TestResult<()> {
4325 let instance =
4326 test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
4327 .await?;
4328
4329 let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4330 expr: Some(api::v1::ddl_request::Expr::CreateTable(
4331 api::v1::CreateTableExpr {
4332 catalog_name: String::new(),
4333 schema_name: String::new(),
4334 table_name: "demo".to_string(),
4335 desc: String::new(),
4336 column_defs: vec![api::v1::ColumnDef {
4337 name: "host".to_string(),
4338 data_type: api::v1::ColumnDataType::String as i32,
4339 is_nullable: true,
4340 default_constraint: vec![],
4341 semantic_type: 0,
4342 comment: String::new(),
4343 datatype_extension: None,
4344 options: None,
4345 }],
4346 time_index: "host".to_string(),
4347 primary_keys: vec![],
4348 create_if_not_exists: false,
4349 table_options: HashMap::new(),
4350 table_id: None,
4351 engine: "mito".to_string(),
4352 },
4353 )),
4354 });
4355
4356 let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
4357 &instance,
4358 request,
4359 QueryContext::arc(),
4360 )
4361 .await;
4362
4363 let err = match result {
4364 Ok(_) => panic!("CreateTable with a non-timestamp time index must be rejected"),
4365 Err(err) => err,
4366 };
4367 assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}");
4368 Ok(())
4369 }
4370
4371 #[tokio::test]
4375 async fn qx_153_create_table_with_timestamp_time_index_via_grpc_succeeds() -> TestResult<()> {
4376 let catalog_manager =
4377 catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
4378 let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
4379 let instance = test_instance_with_catalog_manager(
4380 catalog_manager,
4381 test_table(1025, "target")?,
4382 Plugins::new(),
4383 None,
4384 procedure_executor.clone() as ProcedureExecutorRef,
4385 )
4386 .await?;
4387
4388 instance
4391 .table_metadata_manager()
4392 .schema_manager()
4393 .create(
4394 common_meta::key::schema_name::SchemaNameKey::new("greptime", "public"),
4395 None,
4396 true,
4397 )
4398 .await
4399 .unwrap();
4400
4401 let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
4402 expr: Some(api::v1::ddl_request::Expr::CreateTable(
4403 api::v1::CreateTableExpr {
4404 catalog_name: String::new(),
4405 schema_name: String::new(),
4406 table_name: "demo".to_string(),
4407 desc: String::new(),
4408 column_defs: vec![
4409 api::v1::ColumnDef {
4410 name: "host".to_string(),
4411 data_type: api::v1::ColumnDataType::String as i32,
4412 is_nullable: true,
4413 default_constraint: vec![],
4414 semantic_type: 0,
4415 comment: String::new(),
4416 datatype_extension: None,
4417 options: None,
4418 },
4419 api::v1::ColumnDef {
4420 name: "ts".to_string(),
4421 data_type: api::v1::ColumnDataType::TimestampMillisecond as i32,
4422 is_nullable: true,
4423 default_constraint: vec![],
4424 semantic_type: 0,
4425 comment: String::new(),
4426 datatype_extension: None,
4427 options: None,
4428 },
4429 ],
4430 time_index: "ts".to_string(),
4431 primary_keys: vec![],
4432 create_if_not_exists: false,
4433 table_options: HashMap::new(),
4434 table_id: None,
4435 engine: "mito".to_string(),
4436 },
4437 )),
4438 });
4439
4440 let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
4441 &instance,
4442 request,
4443 QueryContext::arc(),
4444 )
4445 .await
4446 {
4447 Ok(output) => output,
4448 Err(err) => panic!("CREATE TABLE via gRPC DDL must succeed, got error: {err}"),
4449 };
4450 assert!(
4451 matches!(output.data, OutputData::AffectedRows(0)),
4452 "CREATE TABLE via gRPC DDL must report 0 affected rows"
4453 );
4454
4455 let table = instance
4457 .catalog_manager()
4458 .table("greptime", "public", "demo", None)
4459 .await
4460 .unwrap()
4461 .expect("table should exist after CREATE TABLE");
4462 assert_eq!(table.table_info().table_type, TableType::Base);
4463
4464 let submitted = procedure_executor.submitted.lock().unwrap();
4465 assert_eq!(
4466 submitted.len(),
4467 1,
4468 "expected one create table task, got {submitted:?}"
4469 );
4470 assert!(matches!(&submitted[0], DdlTask::CreateTable(_)));
4471 Ok(())
4472 }
4473}