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