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