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