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