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