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