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