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