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