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