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