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