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