mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-10-12 23:12:28 +00:00
4474 lines
160 KiB
Rust
4474 lines
160 KiB
Rust
// Copyright 2023 Greptime Team
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
pub mod builder;
|
|
mod dashboard;
|
|
mod entity_graph;
|
|
mod export_database;
|
|
mod grpc;
|
|
mod import_packed;
|
|
mod influxdb;
|
|
mod jaeger;
|
|
mod log_handler;
|
|
mod logical_batcher;
|
|
mod logs;
|
|
mod opentsdb;
|
|
mod otlp;
|
|
pub mod prom_store;
|
|
mod promql;
|
|
mod region_query;
|
|
|
|
use std::collections::HashSet;
|
|
use std::pin::Pin;
|
|
use std::sync::atomic::AtomicBool;
|
|
use std::sync::{Arc, OnceLock, atomic};
|
|
use std::time::{Duration, SystemTime};
|
|
|
|
use async_stream::stream;
|
|
use async_trait::async_trait;
|
|
use auth::{
|
|
PROMQL_QUERY, PermissionChecker, PermissionCheckerRef, PermissionReq, PermissionTableTarget,
|
|
PermissionTableTargets,
|
|
};
|
|
use catalog::CatalogManagerRef;
|
|
use catalog::process_manager::{
|
|
ProcessManagerRef, QueryStatement as CatalogQueryStatement, SlowQueryRecorder, SlowQueryTimer,
|
|
};
|
|
use client::OutputData;
|
|
use common_base::Plugins;
|
|
use common_base::cancellation::CancellableFuture;
|
|
use common_error::ext::{BoxedError, ErrorExt};
|
|
use common_event_recorder::EventRecorderRef;
|
|
use common_meta::cache::TableFlownodeSetCacheRef;
|
|
use common_meta::cache_invalidator::CacheInvalidatorRef;
|
|
use common_meta::key::TableMetadataManagerRef;
|
|
use common_meta::key::table_name::TableNameKey;
|
|
use common_meta::node_manager::NodeManagerRef;
|
|
use common_meta::procedure_executor::ProcedureExecutorRef;
|
|
use common_query::Output;
|
|
use common_recordbatch::RecordBatchStreamWrapper;
|
|
use common_recordbatch::error::StreamTimeoutSnafu;
|
|
use common_telemetry::logging::SlowQueryOptions;
|
|
use common_telemetry::{debug, error, tracing};
|
|
use dashmap::DashMap;
|
|
use datafusion::dataframe::DataFrame;
|
|
use datafusion::physical_plan::ExecutionPlan;
|
|
use datafusion_expr::LogicalPlan;
|
|
use futures::{Stream, StreamExt, future};
|
|
use lazy_static::lazy_static;
|
|
use operator::delete::DeleterRef;
|
|
use operator::insert::InserterRef;
|
|
use operator::statement::{StatementExecutor, StatementExecutorRef};
|
|
use partition::manager::PartitionRuleManagerRef;
|
|
use pipeline::pipeline_operator::PipelineOperator;
|
|
use prometheus::HistogramTimer;
|
|
use promql_parser::label::Matcher;
|
|
use query::QueryEngineRef;
|
|
use query::metrics::OnDone;
|
|
use query::parser::{PromQuery, QueryStatement};
|
|
use query::query_engine::DescribeResult;
|
|
use query::query_engine::options::{QueryOptions, validate_catalog_and_schema};
|
|
use servers::batcher::logical_table::LogicalTablePendingRowsBatcher;
|
|
use servers::error::{
|
|
self as server_error, AuthSnafu, CommonMetaSnafu, ExecuteQuerySnafu,
|
|
OtlpMetricModeIncompatibleSnafu, UnexpectedResultSnafu,
|
|
};
|
|
use servers::interceptor::{
|
|
PromQueryInterceptor, PromQueryInterceptorRef, SqlQueryInterceptor, SqlQueryInterceptorRef,
|
|
};
|
|
use servers::otlp::metrics::legacy_normalize_otlp_name;
|
|
use servers::prometheus_handler::{
|
|
ParsedPromQuery, PrometheusHandler, resolve_schema_from_matchers,
|
|
};
|
|
use servers::query_handler::sql::SqlQueryHandler;
|
|
use session::context::{Channel, QueryContextRef};
|
|
use session::table_name::table_idents_to_full_name;
|
|
use snafu::prelude::*;
|
|
use sql::ast::ObjectNamePartExt;
|
|
use sql::dialect::Dialect;
|
|
use sql::parser::{ParseOptions, ParserContext};
|
|
use sql::statements::comment::CommentObject;
|
|
use sql::statements::copy::{CopyDatabase, CopyTable};
|
|
use sql::statements::statement::Statement;
|
|
use sql::statements::tql::Tql;
|
|
use sql::util::{extract_tables_from_prom_expr_checked, extract_tables_from_statement_checked};
|
|
use sqlparser::ast::{AnalyzeFormat, ObjectName};
|
|
use table::requests::{OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM};
|
|
use tracing::Span;
|
|
|
|
use crate::error::{
|
|
self, CollectRecordbatchSnafu, Error, ExecLogicalPlanSnafu, ExecutePromqlSnafu, ExternalSnafu,
|
|
InvalidSqlSnafu, ParseSqlSnafu, PermissionSnafu, PlanStatementSnafu, Result,
|
|
SqlExecInterceptedSnafu, StatementTimeoutSnafu, TableOperationSnafu,
|
|
};
|
|
use crate::service_config::InfluxdbMergeMode;
|
|
use crate::stream_wrapper::CancellableStreamWrapper;
|
|
|
|
lazy_static! {
|
|
static ref OTLP_LEGACY_DEFAULT_VALUE: String = "legacy".to_string();
|
|
}
|
|
|
|
/// The frontend instance contains necessary components, and implements many
|
|
/// traits, like [`servers::query_handler::grpc::GrpcQueryHandler`],
|
|
/// [`servers::query_handler::sql::SqlQueryHandler`], etc.
|
|
#[derive(Clone)]
|
|
pub struct Instance {
|
|
frontend_peer_addr: String,
|
|
experimental_metric_export: bool,
|
|
catalog_manager: CatalogManagerRef,
|
|
pipeline_operator: Arc<PipelineOperator>,
|
|
statement_executor: Arc<StatementExecutor>,
|
|
query_engine: QueryEngineRef,
|
|
plugins: Plugins,
|
|
inserter: InserterRef,
|
|
logical_batcher: Arc<OnceLock<Option<Arc<LogicalTablePendingRowsBatcher>>>>,
|
|
deleter: DeleterRef,
|
|
table_metadata_manager: TableMetadataManagerRef,
|
|
event_recorder: EventRecorderRef,
|
|
slow_query_recorder: EventRecorderRef,
|
|
process_manager: ProcessManagerRef,
|
|
slow_query_options: SlowQueryOptions,
|
|
influxdb_default_merge_mode: InfluxdbMergeMode,
|
|
trace_ingest_chunk_size: usize,
|
|
otlp_resource_info: bool,
|
|
suspend: Arc<AtomicBool>,
|
|
|
|
// cache for otlp metrics
|
|
// first layer key: db-string
|
|
// key: direct input metric name
|
|
// value: if runs in legacy mode
|
|
otlp_metrics_table_legacy_cache: DashMap<String, DashMap<String, bool>>,
|
|
}
|
|
|
|
impl Instance {
|
|
pub fn frontend_peer_addr(&self) -> &str {
|
|
&self.frontend_peer_addr
|
|
}
|
|
|
|
pub fn catalog_manager(&self) -> &CatalogManagerRef {
|
|
&self.catalog_manager
|
|
}
|
|
|
|
pub fn query_engine(&self) -> &QueryEngineRef {
|
|
&self.query_engine
|
|
}
|
|
|
|
pub fn plugins(&self) -> &Plugins {
|
|
&self.plugins
|
|
}
|
|
|
|
fn check_permission(
|
|
&self,
|
|
ctx: &QueryContextRef,
|
|
req: PermissionReq<'_>,
|
|
) -> server_error::Result<()> {
|
|
self.plugins
|
|
.get::<PermissionCheckerRef>()
|
|
.as_ref()
|
|
.check_permission(ctx.current_user(), req)
|
|
.context(AuthSnafu)?;
|
|
Ok(())
|
|
}
|
|
|
|
pub fn statement_executor(&self) -> &StatementExecutorRef {
|
|
&self.statement_executor
|
|
}
|
|
|
|
pub fn table_metadata_manager(&self) -> &TableMetadataManagerRef {
|
|
&self.table_metadata_manager
|
|
}
|
|
|
|
pub fn inserter(&self) -> &InserterRef {
|
|
&self.inserter
|
|
}
|
|
|
|
pub fn process_manager(&self) -> &ProcessManagerRef {
|
|
&self.process_manager
|
|
}
|
|
|
|
/// Returns the event recorder configured for this frontend instance.
|
|
pub fn event_recorder(&self) -> EventRecorderRef {
|
|
self.event_recorder.clone()
|
|
}
|
|
|
|
pub fn node_manager(&self) -> &NodeManagerRef {
|
|
self.inserter.node_manager()
|
|
}
|
|
|
|
pub fn partition_manager(&self) -> &PartitionRuleManagerRef {
|
|
self.inserter.partition_manager()
|
|
}
|
|
|
|
pub fn table_flownode_set_cache(&self) -> &TableFlownodeSetCacheRef {
|
|
self.inserter.table_flownode_set_cache()
|
|
}
|
|
|
|
pub fn cache_invalidator(&self) -> &CacheInvalidatorRef {
|
|
self.statement_executor.cache_invalidator()
|
|
}
|
|
|
|
pub fn procedure_executor(&self) -> &ProcedureExecutorRef {
|
|
self.statement_executor.procedure_executor()
|
|
}
|
|
|
|
pub fn suspend_state(&self) -> Arc<AtomicBool> {
|
|
self.suspend.clone()
|
|
}
|
|
|
|
pub(crate) fn is_suspended(&self) -> bool {
|
|
self.suspend.load(atomic::Ordering::Relaxed)
|
|
}
|
|
}
|
|
|
|
fn parse_stmt(sql: &str, dialect: &(dyn Dialect + Send + Sync)) -> Result<Vec<Statement>> {
|
|
ParserContext::create_with_dialect(sql, dialect, ParseOptions::default()).context(ParseSqlSnafu)
|
|
}
|
|
|
|
fn is_explain_analyze_verbose(stmt: &Statement) -> bool {
|
|
matches!(stmt, Statement::Explain(explain) if explain.analyze && explain.verbose)
|
|
|| matches!(stmt, Statement::Tql(Tql::Analyze(analyze)) if analyze.is_verbose)
|
|
}
|
|
|
|
fn validate_analyze_stream_statement(stmt: &mut Statement) -> Result<()> {
|
|
let (is_verbose, format) = match stmt {
|
|
Statement::Explain(explain) => (explain.analyze && explain.verbose, &mut explain.format),
|
|
Statement::Tql(Tql::Analyze(analyze)) => (analyze.is_verbose, &mut analyze.format),
|
|
_ => {
|
|
return InvalidSqlSnafu {
|
|
err_msg: "only EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE statement is supported",
|
|
}
|
|
.fail();
|
|
}
|
|
};
|
|
|
|
ensure!(
|
|
is_verbose,
|
|
InvalidSqlSnafu {
|
|
err_msg: "statement must be EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE"
|
|
}
|
|
);
|
|
match format {
|
|
None | Some(AnalyzeFormat::JSON) => {
|
|
// Keep explicit FORMAT JSON accepted, but pass JSON through
|
|
// QueryContext.explain_format instead of the statement to avoid
|
|
// the planner's current `EXPLAIN VERBOSE with FORMAT` limitation.
|
|
*format = None;
|
|
Ok(())
|
|
}
|
|
Some(_) => InvalidSqlSnafu {
|
|
err_msg: "only FORMAT JSON is supported for EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE",
|
|
}
|
|
.fail(),
|
|
}
|
|
}
|
|
|
|
impl Instance {
|
|
fn statement_slow_query_timer(
|
|
&self,
|
|
stmt: &Statement,
|
|
schema_name: String,
|
|
) -> Option<SlowQueryTimer> {
|
|
if !stmt.is_readonly() || !self.slow_query_options.enable {
|
|
return None;
|
|
}
|
|
|
|
Some(SlowQueryTimer::new(
|
|
CatalogQueryStatement::Sql(stmt.clone()),
|
|
schema_name,
|
|
self.slow_query_options.threshold,
|
|
self.slow_query_options.sample_ratio,
|
|
self.slow_query_options.record_type,
|
|
self.slow_query_recorder.clone(),
|
|
))
|
|
}
|
|
|
|
async fn query_statement(&self, stmt: Statement, query_ctx: QueryContextRef) -> Result<Output> {
|
|
check_permission(self.plugins.clone(), &stmt, &query_ctx)?;
|
|
|
|
let query_interceptor = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
|
|
let query_interceptor = query_interceptor.as_ref();
|
|
|
|
if should_track_statement_process(&stmt) {
|
|
let catalog_name = query_ctx.current_catalog().to_string();
|
|
let schema_name = query_ctx.current_schema();
|
|
let slow_query_timer = self.statement_slow_query_timer(&stmt, schema_name.clone());
|
|
let timeout_recorder = is_explain_analyze_verbose(&stmt)
|
|
.then(|| slow_query_timer.as_ref().map(SlowQueryTimer::recorder))
|
|
.flatten();
|
|
|
|
let ticket = self.process_manager.register_query(
|
|
catalog_name,
|
|
vec![schema_name],
|
|
stmt.to_string(),
|
|
query_ctx.conn_info().to_string(),
|
|
Some(query_ctx.process_id()),
|
|
slow_query_timer,
|
|
);
|
|
|
|
let query_fut = self.exec_statement_with_timeout(
|
|
stmt,
|
|
query_ctx,
|
|
query_interceptor,
|
|
timeout_recorder,
|
|
);
|
|
|
|
CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
|
|
.await
|
|
.map_err(|_| error::CancelledSnafu.build())?
|
|
.map(|output| {
|
|
let Output { meta, data } = output;
|
|
|
|
let data = match data {
|
|
OutputData::Stream(stream) => OutputData::Stream(Box::pin(
|
|
CancellableStreamWrapper::new(stream, ticket),
|
|
)),
|
|
other => other,
|
|
};
|
|
Output { data, meta }
|
|
})
|
|
} else {
|
|
self.exec_statement_with_timeout(stmt, query_ctx, query_interceptor, None)
|
|
.await
|
|
}
|
|
}
|
|
|
|
async fn exec_statement_with_timeout(
|
|
&self,
|
|
stmt: Statement,
|
|
query_ctx: QueryContextRef,
|
|
query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
|
|
timeout_recorder: Option<SlowQueryRecorder>,
|
|
) -> Result<Output> {
|
|
let timeout = derive_timeout(&stmt, &query_ctx);
|
|
match timeout {
|
|
Some(timeout) => {
|
|
let start = tokio::time::Instant::now();
|
|
let output = tokio::time::timeout(
|
|
timeout,
|
|
self.exec_statement(stmt, query_ctx, query_interceptor),
|
|
)
|
|
.await
|
|
.map_err(|_| StatementTimeoutSnafu.build())??;
|
|
let output = map_query_output(output)?;
|
|
// compute remaining timeout
|
|
let remaining_timeout = timeout.checked_sub(start.elapsed()).unwrap_or_default();
|
|
attach_timeout(output, remaining_timeout, timeout_recorder)
|
|
}
|
|
None => self
|
|
.exec_statement(stmt, query_ctx, query_interceptor)
|
|
.await
|
|
.and_then(map_query_output),
|
|
}
|
|
}
|
|
|
|
async fn exec_statement(
|
|
&self,
|
|
stmt: Statement,
|
|
query_ctx: QueryContextRef,
|
|
query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
|
|
) -> Result<Output> {
|
|
match stmt {
|
|
Statement::Query(_) | Statement::Explain(_) | Statement::Delete(_) => {
|
|
// TODO: remove this when format is supported in datafusion
|
|
if let Statement::Explain(explain) = &stmt
|
|
&& let Some(format) = explain.format()
|
|
{
|
|
query_ctx.set_explain_format(format.to_string());
|
|
}
|
|
|
|
self.plan_and_exec_sql(stmt, &query_ctx, query_interceptor)
|
|
.await
|
|
}
|
|
Statement::Tql(tql) => {
|
|
self.plan_and_exec_tql(&query_ctx, query_interceptor, tql)
|
|
.await
|
|
}
|
|
_ => {
|
|
query_interceptor.pre_execute(Some(&stmt), None, query_ctx.clone())?;
|
|
if let Statement::Copy(sql::statements::copy::Copy::CopyDatabase(
|
|
CopyDatabase::From(arg),
|
|
)) = &stmt
|
|
&& arg.with.get("metric_data_layout").is_some()
|
|
{
|
|
return self
|
|
.copy_packed_database(arg.clone(), &stmt, query_ctx)
|
|
.await;
|
|
}
|
|
if let Statement::ShowVariables(show) = &stmt
|
|
&& show
|
|
.variable
|
|
.to_string()
|
|
.eq_ignore_ascii_case("experimental_metric_export")
|
|
{
|
|
return self.show_metric_export_capability();
|
|
}
|
|
if let Statement::Copy(sql::statements::copy::Copy::CopyDatabase(CopyDatabase::To(
|
|
arg,
|
|
))) = &stmt
|
|
&& export_database::parse_metric_export_requested(&arg.with)?
|
|
{
|
|
return self.copy_metric_database(arg.clone(), query_ctx).await;
|
|
}
|
|
self.statement_executor
|
|
.execute_sql(stmt, query_ctx)
|
|
.await
|
|
.context(TableOperationSnafu)
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn plan_and_exec_sql(
|
|
&self,
|
|
stmt: Statement,
|
|
query_ctx: &QueryContextRef,
|
|
query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
|
|
) -> Result<Output> {
|
|
let stmt = QueryStatement::Sql(stmt);
|
|
let plan = self
|
|
.statement_executor
|
|
.plan(&stmt, query_ctx.clone())
|
|
.await?;
|
|
let QueryStatement::Sql(stmt) = stmt else {
|
|
unreachable!()
|
|
};
|
|
query_interceptor.pre_execute(Some(&stmt), Some(&plan), query_ctx.clone())?;
|
|
|
|
self.statement_executor
|
|
.exec_plan(plan, query_ctx.clone())
|
|
.await
|
|
.context(TableOperationSnafu)
|
|
}
|
|
|
|
async fn plan_and_exec_tql(
|
|
&self,
|
|
query_ctx: &QueryContextRef,
|
|
query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
|
|
tql: Tql,
|
|
) -> Result<Output> {
|
|
let plan = self
|
|
.statement_executor
|
|
.plan_tql(tql.clone(), query_ctx)
|
|
.await?;
|
|
query_interceptor.pre_execute(
|
|
Some(&Statement::Tql(tql)),
|
|
Some(&plan),
|
|
query_ctx.clone(),
|
|
)?;
|
|
self.statement_executor
|
|
.exec_plan(plan, query_ctx.clone())
|
|
.await
|
|
.context(TableOperationSnafu)
|
|
}
|
|
|
|
async fn check_otlp_legacy(
|
|
&self,
|
|
names: &[String],
|
|
ctx: &QueryContextRef,
|
|
) -> server_error::Result<bool> {
|
|
let db_string = ctx.get_db_string();
|
|
// fast cache check
|
|
let cache = self
|
|
.otlp_metrics_table_legacy_cache
|
|
.entry(db_string.clone())
|
|
.or_default();
|
|
if let Some(flag) = fast_legacy_check(&cache, names)? {
|
|
return Ok(flag);
|
|
}
|
|
// release cache reference to avoid lock contention
|
|
drop(cache);
|
|
|
|
let catalog = ctx.current_catalog();
|
|
let schema = ctx.current_schema();
|
|
|
|
// query legacy table names
|
|
let normalized_names = names
|
|
.iter()
|
|
.map(|n| legacy_normalize_otlp_name(n))
|
|
.collect::<Vec<_>>();
|
|
let table_names = normalized_names
|
|
.iter()
|
|
.map(|n| TableNameKey::new(catalog, &schema, n))
|
|
.collect::<Vec<_>>();
|
|
let table_values = self
|
|
.table_metadata_manager()
|
|
.table_name_manager()
|
|
.batch_get(table_names)
|
|
.await
|
|
.context(CommonMetaSnafu)?;
|
|
let table_ids = table_values
|
|
.into_iter()
|
|
.filter_map(|v| v.map(|vi| vi.table_id()))
|
|
.collect::<Vec<_>>();
|
|
|
|
// means no existing table is found, use new mode
|
|
if table_ids.is_empty() {
|
|
return Ok(false);
|
|
}
|
|
|
|
// has existing table, check table options
|
|
let table_infos = self
|
|
.table_metadata_manager()
|
|
.table_info_manager()
|
|
.batch_get(&table_ids)
|
|
.await
|
|
.context(CommonMetaSnafu)?;
|
|
let options = table_infos
|
|
.values()
|
|
.map(|info| {
|
|
info.table_info
|
|
.meta
|
|
.options
|
|
.extra_options
|
|
.get(OTLP_METRIC_COMPAT_KEY)
|
|
.unwrap_or(&OTLP_LEGACY_DEFAULT_VALUE)
|
|
})
|
|
.collect::<Vec<_>>();
|
|
if !options.is_empty() {
|
|
// check value consistency
|
|
let has_prom = options.iter().any(|opt| *opt == OTLP_METRIC_COMPAT_PROM);
|
|
let has_legacy = options
|
|
.iter()
|
|
.any(|opt| *opt == OTLP_LEGACY_DEFAULT_VALUE.as_str());
|
|
ensure!(!(has_prom && has_legacy), OtlpMetricModeIncompatibleSnafu);
|
|
Ok(has_legacy)
|
|
} else {
|
|
// no table info, use new mode
|
|
Ok(false)
|
|
}
|
|
}
|
|
|
|
fn cache_otlp_legacy(
|
|
&self,
|
|
names: &[String],
|
|
ctx: &QueryContextRef,
|
|
is_legacy: bool,
|
|
) -> server_error::Result<()> {
|
|
let cache = self
|
|
.otlp_metrics_table_legacy_cache
|
|
.entry(ctx.get_db_string())
|
|
.or_default();
|
|
cache_legacy_mode(&cache, names, is_legacy)
|
|
}
|
|
}
|
|
|
|
fn fast_legacy_check(
|
|
cache: &DashMap<String, bool>,
|
|
names: &[String],
|
|
) -> server_error::Result<Option<bool>> {
|
|
let hit_cache = names
|
|
.iter()
|
|
.filter_map(|name| cache.get(name))
|
|
.collect::<Vec<_>>();
|
|
if !hit_cache.is_empty() {
|
|
let hit_legacy = hit_cache.iter().any(|en| *en.value());
|
|
let hit_prom = hit_cache.iter().any(|en| !*en.value());
|
|
|
|
// hit but have true and false, means both legacy and new mode are used
|
|
// we cannot handle this case, so return error
|
|
// add doc links in err msg later
|
|
ensure!(!(hit_legacy && hit_prom), OtlpMetricModeIncompatibleSnafu);
|
|
|
|
Ok(Some(hit_legacy))
|
|
} else {
|
|
Ok(None)
|
|
}
|
|
}
|
|
|
|
fn cache_legacy_mode(
|
|
cache: &DashMap<String, bool>,
|
|
names: &[String],
|
|
is_legacy: bool,
|
|
) -> server_error::Result<()> {
|
|
for name in names {
|
|
let cached = cache.entry(name.clone()).or_insert(is_legacy);
|
|
ensure!(*cached == is_legacy, OtlpMetricModeIncompatibleSnafu);
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// If the relevant variables are set, the timeout is enforced for all PostgreSQL statements.
|
|
/// For MySQL, it applies only to read-only statements.
|
|
fn derive_timeout(stmt: &Statement, query_ctx: &QueryContextRef) -> Option<Duration> {
|
|
let query_timeout = query_ctx.query_timeout()?;
|
|
if query_timeout.is_zero() {
|
|
return None;
|
|
}
|
|
match query_ctx.channel() {
|
|
Channel::Mysql if stmt.is_readonly() => Some(query_timeout),
|
|
Channel::Postgres => Some(query_timeout),
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
/// Derives timeout for plan execution.
|
|
fn derive_timeout_for_plan(plan: &LogicalPlan, query_ctx: &QueryContextRef) -> Option<Duration> {
|
|
let query_timeout = query_ctx.query_timeout()?;
|
|
if query_timeout.is_zero() {
|
|
return None;
|
|
}
|
|
match query_ctx.channel() {
|
|
Channel::Mysql if is_readonly_plan(plan) => Some(query_timeout),
|
|
Channel::Postgres => Some(query_timeout),
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
fn record_explain_analyze_timeout(
|
|
recorder: Option<&SlowQueryRecorder>,
|
|
plan: Option<&Arc<dyn ExecutionPlan>>,
|
|
) {
|
|
let Some(recorder) = recorder else {
|
|
return;
|
|
};
|
|
let metrics = plan
|
|
.and_then(|plan| query::analyze_plan_metrics_to_json_value(plan, true).ok())
|
|
.unwrap_or_else(|| serde_json::json!([]));
|
|
recorder.force_record_with_payload(serde_json::json!({
|
|
"timed_out": true,
|
|
"metrics": metrics,
|
|
}));
|
|
}
|
|
|
|
fn attach_timeout(
|
|
output: Output,
|
|
mut timeout: Duration,
|
|
timeout_recorder: Option<SlowQueryRecorder>,
|
|
) -> Result<Output> {
|
|
if timeout.is_zero() {
|
|
return StatementTimeoutSnafu.fail();
|
|
}
|
|
|
|
let plan = timeout_recorder
|
|
.as_ref()
|
|
.and_then(|_| output.meta.plan.clone());
|
|
let output = match output.data {
|
|
OutputData::AffectedRows(_) | OutputData::RecordBatches(_) => output,
|
|
OutputData::Stream(mut stream) => {
|
|
let schema = stream.schema();
|
|
let s = Box::pin(stream! {
|
|
let mut start = tokio::time::Instant::now();
|
|
while let Some(item) = tokio::time::timeout(timeout, stream.next()).await.map_err(|_| {
|
|
record_explain_analyze_timeout(timeout_recorder.as_ref(), plan.as_ref());
|
|
StreamTimeoutSnafu.build()
|
|
})? {
|
|
yield item;
|
|
|
|
let now = tokio::time::Instant::now();
|
|
timeout = timeout.checked_sub(now - start).unwrap_or(Duration::ZERO);
|
|
start = now;
|
|
// tokio::time::timeout may not return an error immediately when timeout is 0.
|
|
if timeout.is_zero() {
|
|
record_explain_analyze_timeout(timeout_recorder.as_ref(), plan.as_ref());
|
|
StreamTimeoutSnafu.fail()?;
|
|
}
|
|
}
|
|
}) as Pin<Box<dyn Stream<Item = _> + Send>>;
|
|
let stream = RecordBatchStreamWrapper {
|
|
schema,
|
|
stream: s,
|
|
output_ordering: None,
|
|
metrics: Default::default(),
|
|
span: Span::current(),
|
|
};
|
|
Output::new(OutputData::Stream(Box::pin(stream)), output.meta)
|
|
}
|
|
};
|
|
|
|
Ok(output)
|
|
}
|
|
|
|
impl Instance {
|
|
async fn check_sql_permission(
|
|
&self,
|
|
stmt: &Statement,
|
|
query_ctx: &QueryContextRef,
|
|
) -> Result<()> {
|
|
self.plugins
|
|
.get::<PermissionCheckerRef>()
|
|
.as_ref()
|
|
.check_permission_with_context(
|
|
query_ctx.current_user(),
|
|
PermissionReq::SqlStatement(stmt),
|
|
Some(&query_ctx.current_schema()),
|
|
)
|
|
.context(PermissionSnafu)?;
|
|
|
|
let targets = match extract_tables_from_statement_checked(stmt) {
|
|
Some(tables) => PermissionTableTargets::resolved(
|
|
tables
|
|
.map(|name| {
|
|
table_idents_to_full_name(&name, query_ctx).map(
|
|
|(catalog, schema, table)| {
|
|
PermissionTableTarget::new(catalog, schema, table)
|
|
},
|
|
)
|
|
})
|
|
.collect::<std::result::Result<Vec<_>, _>>()
|
|
.map_err(BoxedError::new)
|
|
.context(ExternalSnafu)?,
|
|
),
|
|
None => PermissionTableTargets::Unresolved,
|
|
};
|
|
let targets = self
|
|
.resolve_query_permission_targets(targets, query_ctx)
|
|
.await
|
|
.map_err(BoxedError::new)
|
|
.context(ExternalSnafu)?;
|
|
self.check_table_permission(query_ctx, PermissionReq::SqlStatement(stmt), targets)
|
|
.context(PermissionSnafu)?;
|
|
Ok(())
|
|
}
|
|
|
|
#[tracing::instrument(skip_all, name = "SqlQueryHandler::do_analyze_stream_query")]
|
|
async fn do_analyze_stream_query_inner(
|
|
&self,
|
|
query: &str,
|
|
query_ctx: QueryContextRef,
|
|
) -> Result<Output> {
|
|
ensure!(!self.is_suspended(), error::SuspendedSnafu);
|
|
|
|
let query_interceptor_opt = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
|
|
let query_interceptor = query_interceptor_opt.as_ref();
|
|
let query = query_interceptor.pre_parsing(query, query_ctx.clone())?;
|
|
let mut stmts = parse_stmt(query.as_ref(), query_ctx.sql_dialect())
|
|
.and_then(|stmts| query_interceptor.post_parsing(stmts, query_ctx.clone()))?;
|
|
|
|
ensure!(
|
|
stmts.len() == 1,
|
|
InvalidSqlSnafu {
|
|
err_msg: "only a single EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE statement is supported"
|
|
}
|
|
);
|
|
let mut stmt = stmts.remove(0);
|
|
validate_analyze_stream_statement(&mut stmt)?;
|
|
query_ctx.set_explain_format(AnalyzeFormat::JSON.to_string());
|
|
|
|
self.check_sql_permission(&stmt, &query_ctx).await?;
|
|
check_permission(self.plugins.clone(), &stmt, &query_ctx)?;
|
|
let catalog_name = query_ctx.current_catalog().to_string();
|
|
let schema_name = query_ctx.current_schema();
|
|
let slow_query_timer = self.statement_slow_query_timer(&stmt, schema_name.clone());
|
|
let ticket = self.process_manager.register_query(
|
|
catalog_name,
|
|
vec![schema_name],
|
|
stmt.to_string(),
|
|
query_ctx.conn_info().to_string(),
|
|
Some(query_ctx.process_id()),
|
|
slow_query_timer,
|
|
);
|
|
let query_fut =
|
|
self.exec_statement_with_timeout(stmt, query_ctx.clone(), query_interceptor, None);
|
|
let output = CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
|
|
.await
|
|
.map_err(|_| error::CancelledSnafu.build())??;
|
|
let Output { meta, data } = output;
|
|
let data = match data {
|
|
OutputData::Stream(stream) => OutputData::Stream(Box::pin(
|
|
CancellableStreamWrapper::new_cancel_on_drop(stream, ticket),
|
|
)),
|
|
other => other,
|
|
};
|
|
query_interceptor.post_execute(Output { data, meta }, query_ctx)
|
|
}
|
|
|
|
#[tracing::instrument(skip_all, name = "SqlQueryHandler::do_query")]
|
|
async fn do_query_inner(&self, query: &str, query_ctx: QueryContextRef) -> Vec<Result<Output>> {
|
|
if self.is_suspended() {
|
|
return vec![error::SuspendedSnafu {}.fail()];
|
|
}
|
|
|
|
let query_interceptor_opt = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
|
|
let query_interceptor = query_interceptor_opt.as_ref();
|
|
let query = match query_interceptor.pre_parsing(query, query_ctx.clone()) {
|
|
Ok(q) => q,
|
|
Err(e) => return vec![Err(e)],
|
|
};
|
|
|
|
match parse_stmt(query.as_ref(), query_ctx.sql_dialect())
|
|
.and_then(|stmts| query_interceptor.post_parsing(stmts, query_ctx.clone()))
|
|
{
|
|
Ok(stmts) => {
|
|
if stmts.is_empty() {
|
|
return vec![
|
|
InvalidSqlSnafu {
|
|
err_msg: "empty statements",
|
|
}
|
|
.fail(),
|
|
];
|
|
}
|
|
|
|
let mut results = Vec::with_capacity(stmts.len());
|
|
for stmt in stmts {
|
|
if let Err(e) = self.check_sql_permission(&stmt, &query_ctx).await {
|
|
results.push(Err(e));
|
|
break;
|
|
}
|
|
|
|
match self.query_statement(stmt.clone(), query_ctx.clone()).await {
|
|
Ok(output) => {
|
|
let output_result =
|
|
query_interceptor.post_execute(output, query_ctx.clone());
|
|
results.push(output_result);
|
|
}
|
|
Err(e) => {
|
|
if e.status_code().should_log_error() {
|
|
error!(e; "Failed to execute query: {stmt}");
|
|
} else {
|
|
debug!("Failed to execute query: {stmt}, {e}");
|
|
}
|
|
results.push(Err(e));
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
results
|
|
}
|
|
Err(e) => {
|
|
vec![Err(e)]
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn exec_plan(&self, plan: LogicalPlan, query_ctx: QueryContextRef) -> Result<Output> {
|
|
self.query_engine
|
|
.execute(plan, query_ctx)
|
|
.await
|
|
.context(ExecLogicalPlanSnafu)
|
|
}
|
|
|
|
async fn exec_plan_with_timeout(
|
|
&self,
|
|
plan: LogicalPlan,
|
|
query_ctx: QueryContextRef,
|
|
timeout_recorder: Option<SlowQueryRecorder>,
|
|
) -> Result<Output> {
|
|
let timeout = derive_timeout_for_plan(&plan, &query_ctx);
|
|
match timeout {
|
|
Some(timeout) => {
|
|
let start = tokio::time::Instant::now();
|
|
let output = tokio::time::timeout(timeout, self.exec_plan(plan, query_ctx))
|
|
.await
|
|
.map_err(|_| StatementTimeoutSnafu.build())??;
|
|
let output = map_query_output(output)?;
|
|
let remaining_timeout = timeout.checked_sub(start.elapsed()).unwrap_or_default();
|
|
attach_timeout(output, remaining_timeout, timeout_recorder)
|
|
}
|
|
None => self
|
|
.exec_plan(plan, query_ctx)
|
|
.await
|
|
.and_then(map_query_output),
|
|
}
|
|
}
|
|
|
|
async fn do_exec_plan_inner(
|
|
&self,
|
|
plan: LogicalPlan,
|
|
stmt: Option<Statement>,
|
|
query_ctx: QueryContextRef,
|
|
) -> Result<Output> {
|
|
ensure!(!self.is_suspended(), error::SuspendedSnafu);
|
|
|
|
let query_interceptor_opt = self.plugins.get::<SqlQueryInterceptorRef<Error>>();
|
|
let query_interceptor = query_interceptor_opt.as_ref();
|
|
|
|
query_interceptor.pre_execute(stmt.as_ref(), Some(&plan), query_ctx.clone())?;
|
|
|
|
// TQL EXPLAIN/ANALYZE formats are consumed from the query context at
|
|
// execution time (see `optimize_physical_plan`); re-apply the side
|
|
// effect of `plan_tql` that was lost when the plan was built during
|
|
// Describe. `explain_format` is per-query state, so this never
|
|
// overwrites anything.
|
|
if let Some(Statement::Tql(tql)) = &stmt {
|
|
let format = match tql {
|
|
Tql::Explain(explain) => explain.format.as_ref(),
|
|
Tql::Analyze(analyze) => analyze.format.as_ref(),
|
|
Tql::Eval(_) => None,
|
|
};
|
|
if let Some(format) = format {
|
|
query_ctx.set_explain_format(format.to_string());
|
|
}
|
|
}
|
|
|
|
let query = stmt
|
|
.as_ref()
|
|
.map(|s| s.to_string())
|
|
.unwrap_or_else(|| plan.display_indent().to_string());
|
|
|
|
let plan_is_readonly = is_readonly_plan(&plan);
|
|
let result = if should_track_plan_process(stmt.as_ref(), &plan) {
|
|
let catalog_name = query_ctx.current_catalog().to_string();
|
|
let schema_name = query_ctx.current_schema();
|
|
let slow_query_timer = if plan_is_readonly {
|
|
self.slow_query_options.enable.then(|| {
|
|
SlowQueryTimer::new(
|
|
CatalogQueryStatement::Plan(query.clone()),
|
|
schema_name.clone(),
|
|
self.slow_query_options.threshold,
|
|
self.slow_query_options.sample_ratio,
|
|
self.slow_query_options.record_type,
|
|
self.slow_query_recorder.clone(),
|
|
)
|
|
})
|
|
} else {
|
|
None
|
|
};
|
|
|
|
let timeout_recorder = stmt
|
|
.as_ref()
|
|
.is_some_and(is_explain_analyze_verbose)
|
|
.then(|| slow_query_timer.as_ref().map(SlowQueryTimer::recorder))
|
|
.flatten();
|
|
let ticket = self.process_manager.register_query(
|
|
catalog_name,
|
|
vec![schema_name],
|
|
query,
|
|
query_ctx.conn_info().to_string(),
|
|
Some(query_ctx.process_id()),
|
|
slow_query_timer,
|
|
);
|
|
|
|
let query_fut = self.exec_plan_with_timeout(plan, query_ctx.clone(), timeout_recorder);
|
|
|
|
CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
|
|
.await
|
|
.map_err(|_| error::CancelledSnafu.build())?
|
|
.map(|output| {
|
|
let Output { meta, data } = output;
|
|
|
|
let data = match data {
|
|
OutputData::Stream(stream) => OutputData::Stream(Box::pin(
|
|
CancellableStreamWrapper::new(stream, ticket),
|
|
)),
|
|
other => other,
|
|
};
|
|
Output { data, meta }
|
|
})
|
|
} else {
|
|
self.exec_plan_with_timeout(plan, query_ctx.clone(), None)
|
|
.await
|
|
};
|
|
|
|
result.and_then(|output| query_interceptor.post_execute(output, query_ctx))
|
|
}
|
|
|
|
#[tracing::instrument(skip_all, name = "SqlQueryHandler::do_promql_query")]
|
|
async fn do_promql_query_inner(
|
|
&self,
|
|
query: &PromQuery,
|
|
query_ctx: QueryContextRef,
|
|
) -> Vec<Result<Output>> {
|
|
if self.is_suspended() {
|
|
return vec![error::SuspendedSnafu {}.fail()];
|
|
}
|
|
|
|
// check will be done in prometheus handler's do_query
|
|
let result = PrometheusHandler::do_query(self, query, query_ctx)
|
|
.await
|
|
.with_context(|_| ExecutePromqlSnafu {
|
|
query: format!("{query:?}"),
|
|
});
|
|
vec![result]
|
|
}
|
|
|
|
/// Builds the [`DataFrame`] for an information-schema-backed `SHOW`
|
|
/// statement; `None` for other statements. The future is boxed to keep
|
|
/// `do_describe_inner`'s state machine small.
|
|
fn show_statement_dataframe<'a>(
|
|
&'a self,
|
|
stmt: &'a Statement,
|
|
query_ctx: &'a QueryContextRef,
|
|
) -> Pin<Box<dyn Future<Output = Option<query::error::Result<DataFrame>>> + Send + 'a>> {
|
|
Box::pin(async move {
|
|
let engine = &self.query_engine;
|
|
let catalog_manager = self.catalog_manager();
|
|
let ctx = query_ctx.clone();
|
|
let dataframe = match stmt {
|
|
Statement::ShowDatabases(show) => {
|
|
query::sql::show_databases_dataframe(show, engine, catalog_manager, ctx).await
|
|
}
|
|
Statement::ShowTables(show) => {
|
|
query::sql::show_tables_dataframe(show, engine, catalog_manager, ctx).await
|
|
}
|
|
Statement::ShowViews(show) => {
|
|
query::sql::show_views_dataframe(show, engine, catalog_manager, ctx).await
|
|
}
|
|
Statement::ShowFlows(show) => {
|
|
query::sql::show_flows_dataframe(show, engine, catalog_manager, ctx).await
|
|
}
|
|
Statement::ShowColumns(show) => {
|
|
query::sql::show_columns_dataframe(show, engine, catalog_manager, ctx).await
|
|
}
|
|
Statement::ShowTableStatus(show) => {
|
|
query::sql::show_table_status_dataframe(show, engine, catalog_manager, ctx)
|
|
.await
|
|
}
|
|
Statement::ShowCharset(kind) => {
|
|
query::sql::show_charsets_dataframe(kind, engine, catalog_manager, ctx).await
|
|
}
|
|
Statement::ShowCollation(kind) => {
|
|
query::sql::show_collations_dataframe(kind, engine, catalog_manager, ctx).await
|
|
}
|
|
Statement::ShowIndex(show) => {
|
|
query::sql::show_index_dataframe(show, engine, catalog_manager, ctx).await
|
|
}
|
|
Statement::ShowRegion(show) => {
|
|
query::sql::show_region_dataframe(show, engine, catalog_manager, ctx).await
|
|
}
|
|
Statement::ShowProcesslist(show) => {
|
|
query::sql::show_processlist_dataframe(show, engine, catalog_manager, ctx).await
|
|
}
|
|
_ => return None,
|
|
};
|
|
Some(dataframe)
|
|
})
|
|
}
|
|
|
|
async fn do_describe_inner(
|
|
&self,
|
|
stmt: Statement,
|
|
query_ctx: QueryContextRef,
|
|
) -> Result<Option<DescribeResult>> {
|
|
ensure!(!self.is_suspended(), error::SuspendedSnafu);
|
|
|
|
// EXPLAIN / EXPLAIN ANALYZE wrap an inner statement; describe them when the
|
|
// wrapped statement is something we already plan (so that bind parameters
|
|
// in the inner query get their types inferred). See #8029.
|
|
let is_inner_plannable = |s: &Statement| {
|
|
matches!(
|
|
s,
|
|
Statement::Insert(_) | Statement::Query(_) | Statement::Delete(_)
|
|
)
|
|
};
|
|
let plannable = is_inner_plannable(&stmt)
|
|
|| matches!(&stmt, Statement::Explain(explain) if is_inner_plannable(explain.statement.as_ref()));
|
|
|
|
if let Statement::Tql(tql) = stmt {
|
|
// TQL produces a logical plan; describe it from the plan so the
|
|
// extended-protocol RowDescription matches the executed DataRows.
|
|
self.check_sql_permission(&Statement::Tql(tql.clone()), &query_ctx)
|
|
.await?;
|
|
let plan = self.statement_executor.plan_tql(tql, &query_ctx).await?;
|
|
return self
|
|
.query_engine
|
|
.describe(plan, query_ctx)
|
|
.await
|
|
.map(Some)
|
|
.context(error::DescribeStatementSnafu);
|
|
}
|
|
|
|
// Describe SHOW statements from the same projection the executor builds.
|
|
if let Some(dataframe) = self
|
|
.show_statement_dataframe(&stmt, &query_ctx)
|
|
.await
|
|
.transpose()
|
|
.context(PlanStatementSnafu)?
|
|
{
|
|
self.check_sql_permission(&stmt, &query_ctx).await?;
|
|
let plan = dataframe.into_unoptimized_plan();
|
|
return self
|
|
.query_engine
|
|
.describe(plan, query_ctx)
|
|
.await
|
|
.map(Some)
|
|
.context(error::DescribeStatementSnafu);
|
|
}
|
|
|
|
if plannable {
|
|
self.check_sql_permission(&stmt, &query_ctx).await?;
|
|
|
|
let plan = self
|
|
.query_engine
|
|
.planner()
|
|
.plan(&QueryStatement::Sql(stmt), query_ctx.clone())
|
|
.await
|
|
.context(PlanStatementSnafu)?;
|
|
self.query_engine
|
|
.describe(plan, query_ctx)
|
|
.await
|
|
.map(Some)
|
|
.context(error::DescribeStatementSnafu)
|
|
} else {
|
|
Ok(None)
|
|
}
|
|
}
|
|
|
|
async fn is_valid_schema_inner(&self, catalog: &str, schema: &str) -> Result<bool> {
|
|
self.catalog_manager
|
|
.schema_exists(catalog, schema, None)
|
|
.await
|
|
.context(error::CatalogSnafu)
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl SqlQueryHandler for Instance {
|
|
async fn do_query(
|
|
&self,
|
|
query: &str,
|
|
query_ctx: QueryContextRef,
|
|
) -> Vec<server_error::Result<Output>> {
|
|
self.do_query_inner(query, query_ctx)
|
|
.await
|
|
.into_iter()
|
|
.map(|result| result.map_err(BoxedError::new).context(ExecuteQuerySnafu))
|
|
.collect()
|
|
}
|
|
|
|
async fn do_analyze_stream_query(
|
|
&self,
|
|
query: &str,
|
|
query_ctx: QueryContextRef,
|
|
) -> server_error::Result<Output> {
|
|
self.do_analyze_stream_query_inner(query, query_ctx)
|
|
.await
|
|
.map_err(BoxedError::new)
|
|
.context(ExecuteQuerySnafu)
|
|
}
|
|
|
|
async fn do_exec_plan(
|
|
&self,
|
|
plan: LogicalPlan,
|
|
stmt: Option<Statement>,
|
|
query_ctx: QueryContextRef,
|
|
) -> server_error::Result<Output> {
|
|
self.do_exec_plan_inner(plan, stmt, query_ctx)
|
|
.await
|
|
.map_err(BoxedError::new)
|
|
.context(server_error::ExecutePlanSnafu)
|
|
}
|
|
|
|
async fn do_promql_query(
|
|
&self,
|
|
query: &PromQuery,
|
|
query_ctx: QueryContextRef,
|
|
) -> Vec<server_error::Result<Output>> {
|
|
self.do_promql_query_inner(query, query_ctx)
|
|
.await
|
|
.into_iter()
|
|
.map(|result| result.map_err(BoxedError::new).context(ExecuteQuerySnafu))
|
|
.collect()
|
|
}
|
|
|
|
async fn do_describe(
|
|
&self,
|
|
stmt: Statement,
|
|
query_ctx: QueryContextRef,
|
|
) -> server_error::Result<Option<DescribeResult>> {
|
|
self.do_describe_inner(stmt, query_ctx)
|
|
.await
|
|
.map_err(BoxedError::new)
|
|
.context(server_error::DescribeStatementSnafu)
|
|
}
|
|
|
|
async fn is_valid_schema(&self, catalog: &str, schema: &str) -> server_error::Result<bool> {
|
|
self.is_valid_schema_inner(catalog, schema)
|
|
.await
|
|
.map_err(BoxedError::new)
|
|
.context(server_error::CheckDatabaseValiditySnafu)
|
|
}
|
|
}
|
|
|
|
/// Expands scan-time dictionaries only when a query result leaves the frontend.
|
|
pub(crate) fn map_query_output(output: Output) -> Result<Output> {
|
|
output
|
|
.map_dictionary_to_values()
|
|
.context(CollectRecordbatchSnafu)
|
|
}
|
|
|
|
/// Attaches a timer to the output and observes it once the output is exhausted.
|
|
pub fn attach_timer(output: Output, timer: HistogramTimer) -> Output {
|
|
match output.data {
|
|
OutputData::AffectedRows(_) | OutputData::RecordBatches(_) => output,
|
|
OutputData::Stream(stream) => {
|
|
let stream = OnDone::new(stream, move || {
|
|
timer.observe_duration();
|
|
});
|
|
Output::new(OutputData::Stream(Box::pin(stream)), output.meta)
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Instance {
|
|
fn check_prom_query_privilege(&self, query_ctx: &QueryContextRef) -> server_error::Result<()> {
|
|
self.plugins
|
|
.get::<PermissionCheckerRef>()
|
|
.as_ref()
|
|
.check_permission(
|
|
query_ctx.current_user(),
|
|
PermissionReq::Action(PROMQL_QUERY),
|
|
)
|
|
.context(AuthSnafu)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn prom_expr_permission_targets(
|
|
&self,
|
|
expr: &promql_parser::parser::Expr,
|
|
query_ctx: &QueryContextRef,
|
|
) -> server_error::Result<Option<Vec<PermissionTableTarget>>> {
|
|
extract_tables_from_prom_expr_checked(expr)
|
|
.map(|tables| {
|
|
tables
|
|
.map(|name| {
|
|
table_idents_to_full_name(&name, query_ctx).map(
|
|
|(catalog, schema, table)| {
|
|
PermissionTableTarget::new(catalog, schema, table)
|
|
},
|
|
)
|
|
})
|
|
.collect::<std::result::Result<Vec<_>, _>>()
|
|
.map_err(BoxedError::new)
|
|
.context(ExecuteQuerySnafu)
|
|
})
|
|
.transpose()
|
|
}
|
|
|
|
async fn is_physical_query_permission_target(
|
|
&self,
|
|
target: &PermissionTableTarget,
|
|
query_ctx: &QueryContextRef,
|
|
) -> server_error::Result<bool> {
|
|
self.catalog_manager
|
|
.table(
|
|
&target.catalog,
|
|
&target.schema,
|
|
&target.table,
|
|
Some(query_ctx),
|
|
)
|
|
.await
|
|
.map(|table| table.is_some_and(|table| table.table_info().is_physical_table()))
|
|
.map_err(BoxedError::new)
|
|
.context(ExecuteQuerySnafu)
|
|
}
|
|
|
|
async fn resolve_query_permission_targets(
|
|
&self,
|
|
targets: PermissionTableTargets,
|
|
query_ctx: &QueryContextRef,
|
|
) -> server_error::Result<PermissionTableTargets> {
|
|
const CONCURRENCY: usize = 8;
|
|
|
|
let checker = self.plugins.get::<PermissionCheckerRef>();
|
|
if !checker.as_ref().uses_table_targets() {
|
|
return Ok(targets);
|
|
}
|
|
|
|
let PermissionTableTargets::Resolved(mut targets) = targets else {
|
|
return Ok(PermissionTableTargets::Unresolved);
|
|
};
|
|
if targets.len() > 1 {
|
|
let mut seen = HashSet::with_capacity(targets.len());
|
|
targets.retain(|target| seen.insert(target.clone()));
|
|
}
|
|
if let [target] = targets.as_slice() {
|
|
return if self
|
|
.is_physical_query_permission_target(target, query_ctx)
|
|
.await?
|
|
{
|
|
Ok(PermissionTableTargets::Unresolved)
|
|
} else {
|
|
Ok(PermissionTableTargets::resolved(targets))
|
|
};
|
|
}
|
|
|
|
// Bound catalog load and inspect results in target order to preserve serial semantics.
|
|
for chunk in targets.chunks(CONCURRENCY) {
|
|
let results = future::join_all(
|
|
chunk
|
|
.iter()
|
|
.map(|target| self.is_physical_query_permission_target(target, query_ctx)),
|
|
)
|
|
.await;
|
|
for result in results {
|
|
if result? {
|
|
return Ok(PermissionTableTargets::Unresolved);
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(PermissionTableTargets::resolved(targets))
|
|
}
|
|
|
|
fn prom_queries_permission_targets(
|
|
&self,
|
|
queries: &[ParsedPromQuery],
|
|
query_ctx: &QueryContextRef,
|
|
) -> server_error::Result<PermissionTableTargets> {
|
|
let mut targets = Vec::new();
|
|
let mut resolved = true;
|
|
|
|
for query in queries {
|
|
let QueryStatement::Promql(eval_stmt, _) = query.statement() else {
|
|
unreachable!("query is parsed from promql");
|
|
};
|
|
|
|
if let Some(query_targets) =
|
|
self.prom_expr_permission_targets(&eval_stmt.expr, query_ctx)?
|
|
{
|
|
targets.extend(query_targets);
|
|
} else {
|
|
resolved = false;
|
|
}
|
|
}
|
|
|
|
Ok(if resolved {
|
|
PermissionTableTargets::resolved(targets)
|
|
} else {
|
|
PermissionTableTargets::Unresolved
|
|
})
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl PrometheusHandler for Instance {
|
|
#[tracing::instrument(skip_all)]
|
|
async fn do_query(
|
|
&self,
|
|
query: &PromQuery,
|
|
query_ctx: QueryContextRef,
|
|
) -> server_error::Result<Output> {
|
|
let query = ParsedPromQuery::parse(query.clone(), &query_ctx)?;
|
|
self.do_query_parsed(query, query_ctx).await
|
|
}
|
|
|
|
#[tracing::instrument(skip_all)]
|
|
async fn do_query_parsed(
|
|
&self,
|
|
query: ParsedPromQuery,
|
|
query_ctx: QueryContextRef,
|
|
) -> server_error::Result<Output> {
|
|
let interceptor = self
|
|
.plugins
|
|
.get::<PromQueryInterceptorRef<server_error::Error>>();
|
|
|
|
self.check_prom_query_privilege(&query_ctx)?;
|
|
|
|
let targets =
|
|
self.prom_queries_permission_targets(std::slice::from_ref(&query), &query_ctx)?;
|
|
self.check_query_target_permission(targets, &query_ctx)
|
|
.await?;
|
|
|
|
let requires_output_ordering = query.requires_output_ordering();
|
|
let (query, stmt) = query.into_parts();
|
|
|
|
let QueryStatement::Promql(eval_stmt, _) = &stmt else {
|
|
unreachable!("query is parsed from promql");
|
|
};
|
|
|
|
let plan = self
|
|
.statement_executor
|
|
.plan(&stmt, query_ctx.clone())
|
|
.await
|
|
.map_err(BoxedError::new)
|
|
.context(ExecuteQuerySnafu)?;
|
|
|
|
let plan = if requires_output_ordering {
|
|
plan
|
|
} else {
|
|
promql::remove_output_sort(plan)
|
|
};
|
|
|
|
interceptor.pre_execute(&query, &eval_stmt.expr, Some(&plan), query_ctx.clone())?;
|
|
|
|
// Take the EvalStmt from the original QueryStatement and use it to create the CatalogQueryStatement.
|
|
let query_statement = if let QueryStatement::Promql(eval_stmt, alias) = stmt {
|
|
CatalogQueryStatement::Promql(eval_stmt, alias)
|
|
} else {
|
|
// It should not happen since the query is already parsed successfully.
|
|
return UnexpectedResultSnafu {
|
|
reason: "The query should always be promql.".to_string(),
|
|
}
|
|
.fail();
|
|
};
|
|
let raw_query = query_statement.to_string();
|
|
|
|
let slow_query_timer = self.slow_query_options.enable.then(|| {
|
|
SlowQueryTimer::new(
|
|
query_statement,
|
|
query_ctx.current_schema(),
|
|
self.slow_query_options.threshold,
|
|
self.slow_query_options.sample_ratio,
|
|
self.slow_query_options.record_type,
|
|
self.slow_query_recorder.clone(),
|
|
)
|
|
});
|
|
|
|
let ticket = self.process_manager.register_query(
|
|
query_ctx.current_catalog().to_string(),
|
|
vec![query_ctx.current_schema()],
|
|
raw_query,
|
|
query_ctx.conn_info().to_string(),
|
|
Some(query_ctx.process_id()),
|
|
slow_query_timer,
|
|
);
|
|
|
|
let query_fut = self.statement_executor.exec_plan(plan, query_ctx.clone());
|
|
|
|
let output = CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
|
|
.await
|
|
.map_err(|_| servers::error::CancelledSnafu.build())?
|
|
.map_err(BoxedError::new)
|
|
.context(ExecuteQuerySnafu)?;
|
|
let output = map_query_output(output)
|
|
.map_err(BoxedError::new)
|
|
.context(ExecuteQuerySnafu)?;
|
|
let Output { meta, data } = output;
|
|
let data = match data {
|
|
OutputData::Stream(stream) => {
|
|
OutputData::Stream(Box::pin(CancellableStreamWrapper::new(stream, ticket)))
|
|
}
|
|
other => other,
|
|
};
|
|
let output = Output { data, meta };
|
|
Ok(interceptor.post_execute(output, query_ctx)?)
|
|
}
|
|
|
|
async fn check_query_permission(
|
|
&self,
|
|
queries: &[PromQuery],
|
|
query_ctx: &QueryContextRef,
|
|
) -> server_error::Result<()> {
|
|
let queries = queries
|
|
.iter()
|
|
.cloned()
|
|
.map(|query| ParsedPromQuery::parse(query, query_ctx))
|
|
.collect::<server_error::Result<Vec<_>>>()?;
|
|
self.check_query_permission_parsed(&queries, query_ctx)
|
|
.await
|
|
}
|
|
|
|
async fn check_query_permission_parsed(
|
|
&self,
|
|
queries: &[ParsedPromQuery],
|
|
query_ctx: &QueryContextRef,
|
|
) -> server_error::Result<()> {
|
|
self.check_prom_query_privilege(query_ctx)?;
|
|
let targets = self.prom_queries_permission_targets(queries, query_ctx)?;
|
|
self.check_query_target_permission(targets, query_ctx).await
|
|
}
|
|
|
|
async fn check_query_target_permission(
|
|
&self,
|
|
targets: PermissionTableTargets,
|
|
query_ctx: &QueryContextRef,
|
|
) -> server_error::Result<()> {
|
|
let targets = self
|
|
.resolve_query_permission_targets(targets, query_ctx)
|
|
.await?;
|
|
self.check_table_permission(query_ctx, PermissionReq::Action(PROMQL_QUERY), targets)
|
|
.context(AuthSnafu)?;
|
|
Ok(())
|
|
}
|
|
|
|
async fn filter_metadata_metric_names(
|
|
&self,
|
|
metric_names: Vec<String>,
|
|
schema: &str,
|
|
query_ctx: &QueryContextRef,
|
|
) -> server_error::Result<Vec<String>> {
|
|
let checker = self.plugins.get::<PermissionCheckerRef>();
|
|
if !checker.as_ref().uses_table_targets() {
|
|
let Some(metric) = metric_names.first() else {
|
|
return Ok(metric_names);
|
|
};
|
|
let target =
|
|
PermissionTableTarget::new(query_ctx.current_catalog(), schema, metric.as_str());
|
|
let result = checker
|
|
.as_ref()
|
|
.check_permission_with_table_targets(
|
|
query_ctx.current_user(),
|
|
PermissionReq::Action(PROMQL_QUERY),
|
|
PermissionTableTargets::resolved(vec![target]),
|
|
)
|
|
.context(AuthSnafu);
|
|
return match result {
|
|
Ok(_) => Ok(metric_names),
|
|
Err(error)
|
|
if error.status_code()
|
|
== common_error::status_code::StatusCode::PermissionDenied =>
|
|
{
|
|
Ok(Vec::new())
|
|
}
|
|
Err(error) => Err(error),
|
|
};
|
|
}
|
|
|
|
let mut allowed = Vec::with_capacity(metric_names.len());
|
|
for metric in metric_names {
|
|
let target =
|
|
PermissionTableTarget::new(query_ctx.current_catalog(), schema, metric.as_str());
|
|
match checker
|
|
.as_ref()
|
|
.check_permission_with_table_targets(
|
|
query_ctx.current_user(),
|
|
PermissionReq::Action(PROMQL_QUERY),
|
|
PermissionTableTargets::resolved(vec![target]),
|
|
)
|
|
.context(AuthSnafu)
|
|
{
|
|
Ok(_) => allowed.push(metric),
|
|
Err(error)
|
|
if error.status_code()
|
|
== common_error::status_code::StatusCode::PermissionDenied => {}
|
|
Err(error) => return Err(error),
|
|
}
|
|
}
|
|
Ok(allowed)
|
|
}
|
|
|
|
async fn query_metric_names(
|
|
&self,
|
|
matchers: Vec<Matcher>,
|
|
schema: &str,
|
|
ctx: &QueryContextRef,
|
|
) -> server_error::Result<Vec<String>> {
|
|
self.handle_query_metric_names(matchers, schema, ctx)
|
|
.await
|
|
.map_err(BoxedError::new)
|
|
.context(ExecuteQuerySnafu)
|
|
}
|
|
|
|
async fn query_metric_names_by_labels(
|
|
&self,
|
|
matchers: Vec<Matcher>,
|
|
schema: &str,
|
|
start: SystemTime,
|
|
end: SystemTime,
|
|
ctx: &QueryContextRef,
|
|
) -> server_error::Result<Vec<String>> {
|
|
self.handle_query_metric_names_by_labels(matchers, schema, start, end, ctx)
|
|
.await
|
|
.map_err(BoxedError::new)
|
|
.context(ExecuteQuerySnafu)
|
|
}
|
|
|
|
async fn query_label_values(
|
|
&self,
|
|
metric: String,
|
|
label_name: String,
|
|
matchers: Vec<Matcher>,
|
|
start: SystemTime,
|
|
end: SystemTime,
|
|
ctx: &QueryContextRef,
|
|
) -> server_error::Result<Vec<String>> {
|
|
let schema =
|
|
resolve_schema_from_matchers(&matchers)?.unwrap_or_else(|| ctx.current_schema());
|
|
let target = PermissionTableTarget::new(ctx.current_catalog(), schema.as_str(), &metric);
|
|
self.check_query_target_permission(
|
|
PermissionTableTargets::resolved(vec![target.clone()]),
|
|
ctx,
|
|
)
|
|
.await?;
|
|
|
|
self.handle_query_label_values(target, label_name, matchers, start, end, ctx)
|
|
.await
|
|
.map_err(BoxedError::new)
|
|
.context(ExecuteQuerySnafu)
|
|
}
|
|
|
|
fn catalog_manager(&self) -> CatalogManagerRef {
|
|
self.catalog_manager.clone()
|
|
}
|
|
}
|
|
|
|
/// Validate `stmt.database` permission if it's presented.
|
|
macro_rules! validate_db_permission {
|
|
($stmt: expr, $query_ctx: expr) => {
|
|
if let Some(database) = &$stmt.database {
|
|
validate_catalog_and_schema($query_ctx.current_catalog(), database, $query_ctx)
|
|
.map_err(BoxedError::new)
|
|
.context(SqlExecInterceptedSnafu)?;
|
|
}
|
|
};
|
|
}
|
|
|
|
pub fn check_permission(
|
|
plugins: Plugins,
|
|
stmt: &Statement,
|
|
query_ctx: &QueryContextRef,
|
|
) -> Result<()> {
|
|
let need_validate = plugins
|
|
.get::<QueryOptions>()
|
|
.map(|opts| opts.disallow_cross_catalog_query)
|
|
.unwrap_or_default();
|
|
|
|
if !need_validate {
|
|
return Ok(());
|
|
}
|
|
|
|
match stmt {
|
|
// Will be checked in execution.
|
|
// TODO(dennis): add a hook for admin commands.
|
|
Statement::Admin(_) => {}
|
|
// These are executed by query engine, and will be checked there.
|
|
Statement::Query(_)
|
|
| Statement::Explain(_)
|
|
| Statement::Tql(_)
|
|
| Statement::Delete(_)
|
|
| Statement::DeclareCursor(_)
|
|
| Statement::Copy(sql::statements::copy::Copy::CopyQueryTo(_)) => {}
|
|
// database ops won't be checked
|
|
Statement::CreateDatabase(_)
|
|
| Statement::ShowDatabases(_)
|
|
| Statement::DropDatabase(_)
|
|
| Statement::AlterDatabase(_)
|
|
| Statement::DropFlow(_)
|
|
| Statement::Use(_) => {}
|
|
#[cfg(feature = "enterprise")]
|
|
Statement::DropTrigger(_) => {}
|
|
Statement::ShowCreateDatabase(stmt) => {
|
|
validate_database(&stmt.database_name, query_ctx)?;
|
|
}
|
|
Statement::ShowCreateTable(stmt) => {
|
|
validate_param(&stmt.table_name, query_ctx)?;
|
|
}
|
|
Statement::ShowCreateFlow(stmt) => {
|
|
validate_flow(&stmt.flow_name, query_ctx)?;
|
|
}
|
|
#[cfg(feature = "enterprise")]
|
|
Statement::ShowCreateTrigger(stmt) => {
|
|
validate_param(&stmt.trigger_name, query_ctx)?;
|
|
}
|
|
Statement::ShowCreateView(stmt) => {
|
|
validate_param(&stmt.view_name, query_ctx)?;
|
|
}
|
|
Statement::CreateExternalTable(stmt) => {
|
|
validate_param(&stmt.name, query_ctx)?;
|
|
}
|
|
Statement::CreateFlow(stmt) => {
|
|
// TODO: should also validate source table name here?
|
|
validate_param(&stmt.sink_table_name, query_ctx)?;
|
|
}
|
|
#[cfg(feature = "enterprise")]
|
|
Statement::CreateTrigger(stmt) => {
|
|
validate_param(&stmt.trigger_name, query_ctx)?;
|
|
}
|
|
Statement::CreateView(stmt) => {
|
|
validate_param(&stmt.name, query_ctx)?;
|
|
}
|
|
Statement::AlterTable(stmt) => {
|
|
validate_param(stmt.table_name(), query_ctx)?;
|
|
}
|
|
#[cfg(feature = "enterprise")]
|
|
Statement::AlterTrigger(_) => {}
|
|
// set/show variable now only alter/show variable in session
|
|
Statement::SetVariables(_) | Statement::ShowVariables(_) => {}
|
|
// show charset and show collation won't be checked
|
|
Statement::ShowCharset(_) | Statement::ShowCollation(_) => {}
|
|
|
|
Statement::Comment(comment) => match &comment.object {
|
|
CommentObject::Table(table) => validate_param(table, query_ctx)?,
|
|
CommentObject::Column { table, .. } => validate_param(table, query_ctx)?,
|
|
CommentObject::Flow(flow) => validate_flow(flow, query_ctx)?,
|
|
},
|
|
|
|
Statement::Insert(insert) => {
|
|
let name = insert.table_name().context(ParseSqlSnafu)?;
|
|
validate_param(name, query_ctx)?;
|
|
}
|
|
Statement::CreateTable(stmt) => {
|
|
validate_param(&stmt.name, query_ctx)?;
|
|
}
|
|
Statement::CreateTableLike(stmt) => {
|
|
validate_param(&stmt.table_name, query_ctx)?;
|
|
validate_param(&stmt.source_name, query_ctx)?;
|
|
}
|
|
Statement::DropTable(drop_stmt) => {
|
|
for table_name in drop_stmt.table_names() {
|
|
validate_param(table_name, query_ctx)?;
|
|
}
|
|
}
|
|
#[cfg(feature = "enterprise")]
|
|
Statement::UndropTable(stmt) => {
|
|
validate_param(stmt.table_name(), query_ctx)?;
|
|
}
|
|
Statement::DropView(stmt) => {
|
|
validate_param(&stmt.view_name, query_ctx)?;
|
|
}
|
|
Statement::ShowTables(stmt) => {
|
|
validate_db_permission!(stmt, query_ctx);
|
|
}
|
|
Statement::ShowTableStatus(stmt) => {
|
|
validate_db_permission!(stmt, query_ctx);
|
|
}
|
|
Statement::ShowColumns(stmt) => {
|
|
validate_db_permission!(stmt, query_ctx);
|
|
}
|
|
Statement::ShowIndex(stmt) => {
|
|
validate_db_permission!(stmt, query_ctx);
|
|
}
|
|
Statement::ShowRegion(stmt) => {
|
|
validate_db_permission!(stmt, query_ctx);
|
|
}
|
|
Statement::ShowViews(stmt) => {
|
|
validate_db_permission!(stmt, query_ctx);
|
|
}
|
|
Statement::ShowFlows(stmt) => {
|
|
validate_db_permission!(stmt, query_ctx);
|
|
}
|
|
Statement::ShowFlowStatus(_stmt) => {
|
|
// Flow statistics are organized based on the catalog dimension and
|
|
// filtered by the current catalog, so there is no need to check the
|
|
// permission of the database(schema).
|
|
}
|
|
#[cfg(feature = "enterprise")]
|
|
Statement::ShowTriggers(_stmt) => {
|
|
// The trigger is organized based on the catalog dimension, so there
|
|
// is no need to check the permission of the database(schema).
|
|
}
|
|
Statement::ShowStatus(_stmt) => {}
|
|
Statement::ShowSearchPath(_stmt) => {}
|
|
Statement::DescribeTable(stmt) => {
|
|
validate_param(stmt.name(), query_ctx)?;
|
|
}
|
|
Statement::Copy(sql::statements::copy::Copy::CopyTable(stmt)) => match stmt {
|
|
CopyTable::To(copy_table_to) => validate_param(©_table_to.table_name, query_ctx)?,
|
|
CopyTable::From(copy_table_from) => {
|
|
validate_param(©_table_from.table_name, query_ctx)?
|
|
}
|
|
},
|
|
Statement::Copy(sql::statements::copy::Copy::CopyDatabase(copy_database)) => {
|
|
match copy_database {
|
|
CopyDatabase::To(stmt) => validate_database(&stmt.database_name, query_ctx)?,
|
|
CopyDatabase::From(stmt) => validate_database(&stmt.database_name, query_ctx)?,
|
|
}
|
|
}
|
|
Statement::TruncateTable(stmt) => {
|
|
validate_param(stmt.table_name(), query_ctx)?;
|
|
}
|
|
// cursor operations are always allowed once it's created
|
|
Statement::FetchCursor(_) | Statement::CloseCursor(_) => {}
|
|
// User can only kill process in their own catalog.
|
|
Statement::Kill(_) => {}
|
|
// SHOW PROCESSLIST
|
|
Statement::ShowProcesslist(_) => {}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn validate_param(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
|
|
let (catalog, schema, _) = table_idents_to_full_name(name, query_ctx)
|
|
.map_err(BoxedError::new)
|
|
.context(ExternalSnafu)?;
|
|
|
|
validate_catalog_and_schema(&catalog, &schema, query_ctx)
|
|
.map_err(BoxedError::new)
|
|
.context(SqlExecInterceptedSnafu)
|
|
}
|
|
|
|
fn validate_flow(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
|
|
let catalog = match &name.0[..] {
|
|
[_flow] => query_ctx.current_catalog().to_string(),
|
|
[catalog, _flow] => catalog.to_string_unquoted(),
|
|
_ => {
|
|
return InvalidSqlSnafu {
|
|
err_msg: format!(
|
|
"expect flow name to be <catalog>.<flow_name> or <flow_name>, actual: {name}",
|
|
),
|
|
}
|
|
.fail();
|
|
}
|
|
};
|
|
|
|
let schema = query_ctx.current_schema();
|
|
|
|
validate_catalog_and_schema(&catalog, &schema, query_ctx)
|
|
.map_err(BoxedError::new)
|
|
.context(SqlExecInterceptedSnafu)
|
|
}
|
|
|
|
fn validate_database(name: &ObjectName, query_ctx: &QueryContextRef) -> Result<()> {
|
|
let (catalog, schema) = match &name.0[..] {
|
|
[schema] => (
|
|
query_ctx.current_catalog().to_string(),
|
|
schema.to_string_unquoted(),
|
|
),
|
|
[catalog, schema] => (catalog.to_string_unquoted(), schema.to_string_unquoted()),
|
|
_ => InvalidSqlSnafu {
|
|
err_msg: format!(
|
|
"expect database name to be <catalog>.<schema> or <schema>, actual: {name}",
|
|
),
|
|
}
|
|
.fail()?,
|
|
};
|
|
|
|
validate_catalog_and_schema(&catalog, &schema, query_ctx)
|
|
.map_err(BoxedError::new)
|
|
.context(SqlExecInterceptedSnafu)
|
|
}
|
|
|
|
fn is_readonly_plan(plan: &LogicalPlan) -> bool {
|
|
!matches!(plan, LogicalPlan::Dml(_) | LogicalPlan::Ddl(_))
|
|
}
|
|
|
|
fn should_track_statement_process(stmt: &Statement) -> bool {
|
|
stmt.is_readonly()
|
|
|| matches!(stmt, Statement::Insert(insert) if insert.has_non_values_query_source())
|
|
}
|
|
|
|
fn should_track_plan_process(stmt: Option<&Statement>, plan: &LogicalPlan) -> bool {
|
|
is_readonly_plan(plan)
|
|
|| matches!(stmt, Some(Statement::Insert(insert)) if insert.has_non_values_query_source())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::any::Any;
|
|
use std::collections::HashMap;
|
|
use std::future::Future;
|
|
use std::pin::Pin;
|
|
use std::sync::Arc;
|
|
use std::task::{Context, Poll};
|
|
use std::time::Duration;
|
|
|
|
use api::prom_store::remote::label_matcher::Type as PromMatcherType;
|
|
use api::prom_store::remote::{
|
|
Label, LabelMatcher, Query as RemoteQuery, ReadRequest, ReadResponse, Sample,
|
|
};
|
|
use api::v1::greptime_request::Request;
|
|
use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse};
|
|
use api::v1::query_request::Query;
|
|
use auth::{
|
|
DASHBOARD_DELETE, DASHBOARD_QUERY, DASHBOARD_SAVE, JAEGER_QUERY, PIPELINE_DELETE,
|
|
PIPELINE_INSERT, PIPELINE_QUERY, PermissionAction, PermissionResp, UserInfo, UserInfoRef,
|
|
};
|
|
use catalog::process_manager::{ProcessManager, QueryStatement, SlowQueryTimer};
|
|
use common_base::Plugins;
|
|
use common_catalog::consts::DEFAULT_PRIVATE_SCHEMA_NAME;
|
|
use common_error::ext::{BoxedError, ErrorExt, PlainError};
|
|
use common_error::status_code::StatusCode;
|
|
use common_event_recorder::{Event, EventRecorder, EventTypeFilter, EventTypeFilterRef};
|
|
use common_frontend::slow_query_event::SlowQueryEvent;
|
|
use common_meta::cache::LayeredCacheRegistryBuilder;
|
|
use common_meta::kv_backend::memory::MemoryKvBackend;
|
|
use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor};
|
|
use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse};
|
|
use common_meta::rpc::procedure::{
|
|
MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse,
|
|
};
|
|
use common_query::prelude::greptime_value;
|
|
use common_query::{Output, OutputMeta};
|
|
use common_recordbatch::{
|
|
OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream,
|
|
};
|
|
use common_telemetry::logging::SlowQueriesRecordType;
|
|
use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
|
|
use datafusion::physical_plan::empty::EmptyExec;
|
|
use datafusion_expr::dml::InsertOp;
|
|
use datafusion_expr::{LogicalPlanBuilder, LogicalTableSource};
|
|
use datatypes::prelude::ConcreteDataType;
|
|
use datatypes::schema::{ColumnSchema, Schema as GtSchema, SchemaRef as GtSchemaRef};
|
|
use datatypes::vectors::{
|
|
Float64Vector, StringVector, TimestampMillisecondVector, TimestampNanosecondVector,
|
|
VectorRef,
|
|
};
|
|
use log_query::LogQuery;
|
|
use prost::Message;
|
|
use query::query_engine::options::QueryOptions;
|
|
use servers::query_handler::{
|
|
DashboardHandler, JaegerQueryHandler, LogQueryHandler, PipelineHandler, PipelineHandlerRef,
|
|
PromStoreProtocolHandler,
|
|
};
|
|
use session::context::{Channel, ConnInfo, QueryContext, QueryContextBuilder};
|
|
use snafu::{Location, Snafu};
|
|
use sql::dialect::GreptimeDbDialect;
|
|
use store_api::data_source::DataSource;
|
|
use store_api::metric_engine_consts::{
|
|
LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY,
|
|
};
|
|
use store_api::storage::ScanRequest;
|
|
use strfmt::Format;
|
|
use table::metadata::{
|
|
FilterPushDownType, TableInfo, TableInfoBuilder, TableMetaBuilder, TableType,
|
|
};
|
|
use table::table_name::TableName;
|
|
use table::test_util::{EmptyTable, MemTable};
|
|
use table::{Table, TableRef};
|
|
use tokio::sync::{mpsc, oneshot};
|
|
use tower::ServiceExt;
|
|
|
|
use crate::frontend::FrontendOptions;
|
|
use crate::instance::builder::FrontendBuilder;
|
|
use crate::instance::*;
|
|
|
|
fn parse_test_sql(sql: &str) -> Vec<Statement> {
|
|
parse_stmt(sql, &GreptimeDbDialect {}).unwrap()
|
|
}
|
|
|
|
#[derive(Debug, Default)]
|
|
struct RecordingSlowQueryEventRecorder {
|
|
payloads: std::sync::Mutex<Vec<serde_json::Value>>,
|
|
}
|
|
|
|
impl EventRecorder for RecordingSlowQueryEventRecorder {
|
|
fn record(&self, event: Box<dyn Event>) {
|
|
let event = event
|
|
.as_any()
|
|
.downcast_ref::<SlowQueryEvent>()
|
|
.expect("expected a slow query event");
|
|
self.payloads.lock().unwrap().push(event.payload.clone());
|
|
}
|
|
|
|
fn event_type_filter(&self) -> EventTypeFilterRef {
|
|
Arc::new(EventTypeFilter::All)
|
|
}
|
|
|
|
fn close(&self) {}
|
|
}
|
|
|
|
#[test]
|
|
fn test_validate_analyze_stream_statement_strictness() {
|
|
for sql in [
|
|
"select 1",
|
|
"explain analyze select 1",
|
|
"explain analyze verbose format text select 1",
|
|
"explain analyze verbose format graphviz select 1",
|
|
"TQL ANALYZE (0, 10, '5s') physical_metric",
|
|
"TQL EXPLAIN VERBOSE (0, 10, '5s') physical_metric",
|
|
"TQL ANALYZE VERBOSE FORMAT TEXT (0, 10, '5s') physical_metric",
|
|
] {
|
|
let mut stmts = parse_test_sql(sql);
|
|
assert!(
|
|
validate_analyze_stream_statement(&mut stmts[0]).is_err(),
|
|
"{sql}"
|
|
);
|
|
}
|
|
|
|
for sql in [
|
|
"explain analyze verbose select 1",
|
|
"explain analyze verbose format json select 1",
|
|
"TQL ANALYZE VERBOSE (0, 10, '5s') physical_metric",
|
|
"TQL ANALYZE VERBOSE FORMAT JSON (0, 10, '5s') physical_metric",
|
|
] {
|
|
let mut stmts = parse_test_sql(sql);
|
|
assert!(
|
|
validate_analyze_stream_statement(&mut stmts[0]).is_ok(),
|
|
"{sql}"
|
|
);
|
|
match &stmts[0] {
|
|
Statement::Explain(explain) => assert!(explain.format.is_none()),
|
|
Statement::Tql(Tql::Analyze(analyze)) => assert!(analyze.format.is_none()),
|
|
_ => unreachable!(),
|
|
}
|
|
}
|
|
|
|
assert_eq!(
|
|
parse_test_sql("explain analyze verbose select 1; select 2").len(),
|
|
2
|
|
);
|
|
|
|
assert!(is_explain_analyze_verbose(
|
|
&parse_test_sql("explain analyze verbose select 1")[0]
|
|
));
|
|
assert!(is_explain_analyze_verbose(
|
|
&parse_test_sql("TQL ANALYZE VERBOSE (0, 10, '5s') physical_metric")[0]
|
|
));
|
|
for sql in [
|
|
"select 1",
|
|
"explain select 1",
|
|
"explain analyze select 1",
|
|
"explain verbose select 1",
|
|
"TQL ANALYZE (0, 10, '5s') physical_metric",
|
|
"TQL EXPLAIN VERBOSE (0, 10, '5s') physical_metric",
|
|
] {
|
|
assert!(
|
|
!is_explain_analyze_verbose(&parse_test_sql(sql)[0]),
|
|
"{sql}"
|
|
);
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Snafu)]
|
|
enum TestError {
|
|
#[snafu(display("Failed to build test cache registry"))]
|
|
BuildCacheRegistry {
|
|
source: cache::error::Error,
|
|
#[snafu(implicit)]
|
|
location: Location,
|
|
},
|
|
|
|
#[snafu(display("Failed to build test table meta for table: {table_name}"))]
|
|
BuildTableMeta {
|
|
table_name: String,
|
|
source: table::metadata::TableMetaBuilderError,
|
|
#[snafu(implicit)]
|
|
location: Location,
|
|
},
|
|
|
|
#[snafu(display("Failed to build test table info for table: {table_name}"))]
|
|
BuildTableInfo {
|
|
table_name: String,
|
|
source: table::metadata::TableInfoBuilderError,
|
|
#[snafu(implicit)]
|
|
location: Location,
|
|
},
|
|
|
|
#[snafu(display("Failed to register test table: {table_name}"))]
|
|
RegisterTable {
|
|
table_name: String,
|
|
source: catalog::error::Error,
|
|
#[snafu(implicit)]
|
|
location: Location,
|
|
},
|
|
|
|
#[snafu(display("Failed to build test frontend instance"))]
|
|
BuildFrontend {
|
|
source: crate::error::Error,
|
|
#[snafu(implicit)]
|
|
location: Location,
|
|
},
|
|
|
|
#[snafu(display("Expected exactly one output for SQL `{sql}`, got {actual}"))]
|
|
UnexpectedOutputCount {
|
|
sql: String,
|
|
actual: usize,
|
|
#[snafu(implicit)]
|
|
location: Location,
|
|
},
|
|
|
|
#[snafu(display("Failed to execute SQL `{sql}`"))]
|
|
ExecuteSql {
|
|
sql: String,
|
|
source: crate::error::Error,
|
|
#[snafu(implicit)]
|
|
location: Location,
|
|
},
|
|
|
|
#[snafu(display("Timed out waiting for insert-select start notification"))]
|
|
InsertStartTimeout {
|
|
source: tokio::time::error::Elapsed,
|
|
#[snafu(implicit)]
|
|
location: Location,
|
|
},
|
|
|
|
#[snafu(display("Insert-select start notification channel closed"))]
|
|
InsertStartChannelClosed {
|
|
#[snafu(implicit)]
|
|
location: Location,
|
|
},
|
|
|
|
#[snafu(display("Failed to release blocking insert-select interceptor"))]
|
|
ReleaseBlockedInsert {
|
|
#[snafu(implicit)]
|
|
location: Location,
|
|
},
|
|
|
|
#[snafu(display("Timed out waiting for insert-select source to be polled"))]
|
|
SourcePollTimeout {
|
|
source: tokio::time::error::Elapsed,
|
|
#[snafu(implicit)]
|
|
location: Location,
|
|
},
|
|
|
|
#[snafu(display("Insert-select source poll notification channel closed"))]
|
|
SourcePollChannelClosed {
|
|
source: oneshot::error::RecvError,
|
|
#[snafu(implicit)]
|
|
location: Location,
|
|
},
|
|
|
|
#[snafu(display("Timed out waiting for insert task to finish"))]
|
|
InsertTaskTimeout {
|
|
source: tokio::time::error::Elapsed,
|
|
#[snafu(implicit)]
|
|
location: Location,
|
|
},
|
|
|
|
#[snafu(display("Insert task panicked"))]
|
|
InsertTaskPanic {
|
|
source: tokio::task::JoinError,
|
|
#[snafu(implicit)]
|
|
location: Location,
|
|
},
|
|
|
|
#[snafu(display("Expected insert-select to be cancelled"))]
|
|
InsertSelectNotCancelled {
|
|
#[snafu(implicit)]
|
|
location: Location,
|
|
},
|
|
}
|
|
|
|
type TestResult<T> = std::result::Result<T, TestError>;
|
|
|
|
fn parse_one_sql(sql: &str) -> Statement {
|
|
parse_stmt(sql, &GreptimeDbDialect {}).unwrap().remove(0)
|
|
}
|
|
|
|
fn test_query_ctx(process_id: u32) -> QueryContextRef {
|
|
Arc::new(
|
|
QueryContextBuilder::default()
|
|
.channel(Channel::Mysql)
|
|
.conn_info(ConnInfo::new(None, Channel::Mysql))
|
|
.process_id(process_id)
|
|
.build(),
|
|
)
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct AdminUserInfo;
|
|
|
|
impl UserInfo for AdminUserInfo {
|
|
fn as_any(&self) -> &dyn Any {
|
|
self
|
|
}
|
|
|
|
fn username(&self) -> &str {
|
|
"admin"
|
|
}
|
|
|
|
fn is_admin(&self) -> bool {
|
|
true
|
|
}
|
|
}
|
|
|
|
struct RejectUnresolvedPermissionChecker;
|
|
|
|
impl PermissionChecker for RejectUnresolvedPermissionChecker {
|
|
fn check_permission(
|
|
&self,
|
|
_user_info: UserInfoRef,
|
|
_req: PermissionReq,
|
|
) -> auth::error::Result<PermissionResp> {
|
|
Ok(PermissionResp::Allow)
|
|
}
|
|
|
|
fn check_permission_with_table_targets(
|
|
&self,
|
|
_user_info: UserInfoRef,
|
|
_req: PermissionReq,
|
|
targets: PermissionTableTargets,
|
|
) -> auth::error::Result<PermissionResp> {
|
|
let reject = match targets {
|
|
PermissionTableTargets::Unresolved => true,
|
|
PermissionTableTargets::Resolved(targets) => {
|
|
targets.iter().any(|target| target.table == "denied")
|
|
}
|
|
};
|
|
Ok(if reject {
|
|
PermissionResp::Reject
|
|
} else {
|
|
PermissionResp::Allow
|
|
})
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, PartialEq, Eq)]
|
|
struct CheckedAction {
|
|
action: PermissionAction,
|
|
targets: Option<PermissionTableTargets>,
|
|
}
|
|
|
|
#[derive(Default)]
|
|
struct RejectEndpointPermissionChecker {
|
|
checks: std::sync::Mutex<Vec<CheckedAction>>,
|
|
}
|
|
|
|
impl RejectEndpointPermissionChecker {
|
|
fn reject(
|
|
&self,
|
|
action: PermissionAction,
|
|
targets: Option<PermissionTableTargets>,
|
|
) -> PermissionResp {
|
|
self.checks
|
|
.lock()
|
|
.unwrap()
|
|
.push(CheckedAction { action, targets });
|
|
PermissionResp::Reject
|
|
}
|
|
|
|
fn take_check(&self) -> CheckedAction {
|
|
let mut checks = self.checks.lock().unwrap();
|
|
assert_eq!(1, checks.len());
|
|
checks.pop().unwrap()
|
|
}
|
|
}
|
|
|
|
impl PermissionChecker for RejectEndpointPermissionChecker {
|
|
fn check_permission(
|
|
&self,
|
|
_user_info: UserInfoRef,
|
|
req: PermissionReq,
|
|
) -> auth::error::Result<PermissionResp> {
|
|
Ok(match req {
|
|
PermissionReq::Action(action) => self.reject(action, None),
|
|
_ => PermissionResp::Allow,
|
|
})
|
|
}
|
|
|
|
fn check_permission_with_table_targets(
|
|
&self,
|
|
_user_info: UserInfoRef,
|
|
req: PermissionReq,
|
|
targets: PermissionTableTargets,
|
|
) -> auth::error::Result<PermissionResp> {
|
|
Ok(match req {
|
|
PermissionReq::Action(action) => self.reject(action, Some(targets)),
|
|
_ => PermissionResp::Allow,
|
|
})
|
|
}
|
|
}
|
|
|
|
struct WriteOnlyPermissionChecker;
|
|
|
|
impl PermissionChecker for WriteOnlyPermissionChecker {
|
|
fn check_permission(
|
|
&self,
|
|
_user_info: UserInfoRef,
|
|
req: PermissionReq,
|
|
) -> auth::error::Result<PermissionResp> {
|
|
Ok(if req.is_readonly() {
|
|
PermissionResp::Reject
|
|
} else {
|
|
PermissionResp::Allow
|
|
})
|
|
}
|
|
|
|
fn check_permission_with_table_targets(
|
|
&self,
|
|
user_info: UserInfoRef,
|
|
req: PermissionReq,
|
|
_targets: PermissionTableTargets,
|
|
) -> auth::error::Result<PermissionResp> {
|
|
self.check_permission(user_info, req)
|
|
}
|
|
}
|
|
|
|
#[derive(Default)]
|
|
struct TargetIndependentPermissionChecker {
|
|
checks: atomic::AtomicUsize,
|
|
}
|
|
|
|
impl PermissionChecker for TargetIndependentPermissionChecker {
|
|
fn check_permission(
|
|
&self,
|
|
_user_info: UserInfoRef,
|
|
_req: PermissionReq,
|
|
) -> auth::error::Result<PermissionResp> {
|
|
self.checks.fetch_add(1, atomic::Ordering::Relaxed);
|
|
Ok(PermissionResp::Allow)
|
|
}
|
|
|
|
fn uses_table_targets(&self) -> bool {
|
|
false
|
|
}
|
|
|
|
fn check_permission_with_table_targets(
|
|
&self,
|
|
user_info: UserInfoRef,
|
|
req: PermissionReq,
|
|
_targets: PermissionTableTargets,
|
|
) -> auth::error::Result<PermissionResp> {
|
|
self.check_permission(user_info, req)
|
|
}
|
|
}
|
|
|
|
struct BlockingInsertSelectInterceptor {
|
|
started_tx: mpsc::UnboundedSender<()>,
|
|
finish_rx: std::sync::Mutex<Option<oneshot::Receiver<()>>>,
|
|
}
|
|
|
|
impl BlockingInsertSelectInterceptor {
|
|
fn new(started_tx: mpsc::UnboundedSender<()>, finish_rx: oneshot::Receiver<()>) -> Self {
|
|
Self {
|
|
started_tx,
|
|
finish_rx: std::sync::Mutex::new(Some(finish_rx)),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl SqlQueryInterceptor for BlockingInsertSelectInterceptor {
|
|
type Error = Error;
|
|
|
|
fn pre_execute(
|
|
&self,
|
|
statement: Option<&Statement>,
|
|
_plan: Option<&LogicalPlan>,
|
|
_query_ctx: QueryContextRef,
|
|
) -> Result<()> {
|
|
let Some(Statement::Insert(insert)) = statement else {
|
|
return Ok(());
|
|
};
|
|
if !insert.has_non_values_query_source() {
|
|
return Ok(());
|
|
}
|
|
|
|
let finish_rx = self.finish_rx.lock().unwrap().take().unwrap();
|
|
let _ = self.started_tx.send(());
|
|
tokio::task::block_in_place(|| {
|
|
tokio::runtime::Handle::current()
|
|
.block_on(finish_rx)
|
|
.unwrap();
|
|
});
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
struct PendingRecordBatchStream {
|
|
schema: GtSchemaRef,
|
|
polled_tx: Option<oneshot::Sender<()>>,
|
|
_finish_tx: oneshot::Sender<()>,
|
|
finish_rx: Pin<Box<oneshot::Receiver<()>>>,
|
|
}
|
|
|
|
impl RecordBatchStream for PendingRecordBatchStream {
|
|
fn schema(&self) -> GtSchemaRef {
|
|
self.schema.clone()
|
|
}
|
|
|
|
fn output_ordering(&self) -> Option<&[OrderOption]> {
|
|
None
|
|
}
|
|
|
|
fn metrics(&self) -> Option<common_recordbatch::adapter::RecordBatchMetrics> {
|
|
None
|
|
}
|
|
}
|
|
|
|
impl Stream for PendingRecordBatchStream {
|
|
type Item = common_recordbatch::error::Result<RecordBatch>;
|
|
|
|
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
|
if let Some(polled_tx) = self.polled_tx.take() {
|
|
let _ = polled_tx.send(());
|
|
}
|
|
|
|
match self.finish_rx.as_mut().poll(cx) {
|
|
Poll::Ready(_) => Poll::Ready(None),
|
|
Poll::Pending => Poll::Pending,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Unpin for PendingRecordBatchStream {}
|
|
|
|
#[test]
|
|
fn test_record_explain_analyze_timeout_uses_empty_metrics_without_plan() {
|
|
let event_recorder = Arc::new(RecordingSlowQueryEventRecorder::default());
|
|
let timer = SlowQueryTimer::new(
|
|
QueryStatement::Plan("EXPLAIN ANALYZE VERBOSE SELECT 1".to_string()),
|
|
"public".to_string(),
|
|
Duration::from_secs(3600),
|
|
0.0,
|
|
SlowQueriesRecordType::SystemTable,
|
|
event_recorder.clone(),
|
|
);
|
|
let timeout_recorder = timer.recorder();
|
|
|
|
record_explain_analyze_timeout(Some(&timeout_recorder), None);
|
|
drop(timer);
|
|
|
|
let payloads = event_recorder.payloads.lock().unwrap();
|
|
assert_eq!(payloads.len(), 1);
|
|
assert_eq!(payloads[0]["timed_out"], true);
|
|
assert_eq!(payloads[0]["metrics"], serde_json::json!([]));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_attach_timeout_records_explain_analyze_metrics() {
|
|
let event_recorder = Arc::new(RecordingSlowQueryEventRecorder::default());
|
|
let timer = SlowQueryTimer::new(
|
|
QueryStatement::Plan("EXPLAIN ANALYZE VERBOSE SELECT 1".to_string()),
|
|
"public".to_string(),
|
|
Duration::from_secs(3600),
|
|
0.0,
|
|
SlowQueriesRecordType::SystemTable,
|
|
event_recorder.clone(),
|
|
);
|
|
let timeout_recorder = timer.recorder();
|
|
let plan: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(Arc::new(Schema::empty())));
|
|
let (finish_tx, finish_rx) = oneshot::channel();
|
|
let stream = PendingRecordBatchStream {
|
|
schema: Arc::new(GtSchema::new(vec![])),
|
|
polled_tx: None,
|
|
_finish_tx: finish_tx,
|
|
finish_rx: Box::pin(finish_rx),
|
|
};
|
|
let output = Output::new(
|
|
OutputData::Stream(Box::pin(stream)),
|
|
OutputMeta::new_with_plan(plan),
|
|
);
|
|
let output =
|
|
attach_timeout(output, Duration::from_millis(10), Some(timeout_recorder)).unwrap();
|
|
let OutputData::Stream(mut stream) = output.data else {
|
|
unreachable!();
|
|
};
|
|
|
|
let err = stream.next().await.unwrap().unwrap_err();
|
|
assert_eq!(err.to_string(), "Stream timeout");
|
|
drop(stream);
|
|
drop(timer);
|
|
|
|
let payloads = event_recorder.payloads.lock().unwrap();
|
|
assert_eq!(payloads.len(), 1);
|
|
assert_eq!(payloads[0]["timed_out"], true);
|
|
assert!(
|
|
payloads[0]["metrics"]
|
|
.as_array()
|
|
.is_some_and(|metrics| !metrics.is_empty())
|
|
);
|
|
}
|
|
|
|
struct PendingDataSource {
|
|
schema: GtSchemaRef,
|
|
polled_tx: std::sync::Mutex<Option<oneshot::Sender<()>>>,
|
|
}
|
|
|
|
impl DataSource for PendingDataSource {
|
|
fn get_stream(
|
|
&self,
|
|
_request: ScanRequest,
|
|
) -> std::result::Result<SendableRecordBatchStream, BoxedError> {
|
|
let (finish_tx, finish_rx) = oneshot::channel();
|
|
let mut polled_tx = self.polled_tx.lock().map_err(|_| {
|
|
BoxedError::new(PlainError::new(
|
|
"pending data source lock poisoned".to_string(),
|
|
StatusCode::Unexpected,
|
|
))
|
|
})?;
|
|
Ok(Box::pin(PendingRecordBatchStream {
|
|
schema: self.schema.clone(),
|
|
polled_tx: polled_tx.take(),
|
|
_finish_tx: finish_tx,
|
|
finish_rx: Box::pin(finish_rx),
|
|
}))
|
|
}
|
|
}
|
|
|
|
struct NoopProcedureExecutor;
|
|
|
|
#[async_trait::async_trait]
|
|
impl ProcedureExecutor for NoopProcedureExecutor {
|
|
async fn submit_ddl_task(
|
|
&self,
|
|
_ctx: ExecutorContext,
|
|
_request: SubmitDdlTaskRequest,
|
|
) -> common_meta::error::Result<SubmitDdlTaskResponse> {
|
|
common_meta::error::UnsupportedSnafu {
|
|
operation: "submit_ddl_task",
|
|
}
|
|
.fail()
|
|
}
|
|
|
|
async fn migrate_region(
|
|
&self,
|
|
_ctx: &ExecutorContext,
|
|
_request: MigrateRegionRequest,
|
|
) -> common_meta::error::Result<MigrateRegionResponse> {
|
|
common_meta::error::UnsupportedSnafu {
|
|
operation: "migrate_region",
|
|
}
|
|
.fail()
|
|
}
|
|
|
|
async fn reconcile(
|
|
&self,
|
|
_ctx: &ExecutorContext,
|
|
_request: ReconcileRequest,
|
|
) -> common_meta::error::Result<ReconcileResponse> {
|
|
common_meta::error::UnsupportedSnafu {
|
|
operation: "reconcile",
|
|
}
|
|
.fail()
|
|
}
|
|
|
|
async fn query_procedure_state(
|
|
&self,
|
|
_ctx: &ExecutorContext,
|
|
_pid: &str,
|
|
) -> common_meta::error::Result<ProcedureStateResponse> {
|
|
common_meta::error::UnsupportedSnafu {
|
|
operation: "query_procedure_state",
|
|
}
|
|
.fail()
|
|
}
|
|
|
|
async fn list_procedures(
|
|
&self,
|
|
_ctx: &ExecutorContext,
|
|
) -> common_meta::error::Result<ProcedureDetailResponse> {
|
|
common_meta::error::UnsupportedSnafu {
|
|
operation: "list_procedures",
|
|
}
|
|
.fail()
|
|
}
|
|
}
|
|
|
|
/// A test [`ProcedureExecutor`] that completes create/drop DDL tasks against the
|
|
/// in-memory catalog, mimicking what the meta DDL procedures do in production.
|
|
/// This allows happy-path DDL requests (create/drop table/view) to be exercised
|
|
/// end to end through the gRPC ingress.
|
|
struct MockProcedureExecutor {
|
|
catalog_manager: Arc<catalog::memory::MemoryCatalogManager>,
|
|
next_table_id: std::sync::atomic::AtomicU32,
|
|
submitted: std::sync::Mutex<Vec<DdlTask>>,
|
|
}
|
|
|
|
impl MockProcedureExecutor {
|
|
fn new(catalog_manager: Arc<catalog::memory::MemoryCatalogManager>) -> Self {
|
|
Self {
|
|
catalog_manager,
|
|
next_table_id: std::sync::atomic::AtomicU32::new(1026),
|
|
submitted: std::sync::Mutex::new(Vec::new()),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[async_trait::async_trait]
|
|
impl ProcedureExecutor for MockProcedureExecutor {
|
|
async fn submit_ddl_task(
|
|
&self,
|
|
_ctx: ExecutorContext,
|
|
request: SubmitDdlTaskRequest,
|
|
) -> common_meta::error::Result<SubmitDdlTaskResponse> {
|
|
self.submitted.lock().unwrap().push(request.task.clone());
|
|
match request.task {
|
|
DdlTask::CreateTable(task) => {
|
|
let table_id = self
|
|
.next_table_id
|
|
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
|
let mut table_info = task.table_info;
|
|
table_info.ident.table_id = table_id;
|
|
self.catalog_manager
|
|
.register_table_sync(catalog::RegisterTableRequest {
|
|
catalog: table_info.catalog_name.clone(),
|
|
schema: table_info.schema_name.clone(),
|
|
table_name: table_info.name.clone(),
|
|
table_id,
|
|
table: table::dist_table::DistTable::table(Arc::new(table_info)),
|
|
})
|
|
.map_err(BoxedError::new)
|
|
.context(common_meta::error::ExternalSnafu)?;
|
|
Ok(SubmitDdlTaskResponse {
|
|
key: Vec::new(),
|
|
table_ids: vec![table_id],
|
|
})
|
|
}
|
|
DdlTask::CreateView(task) => {
|
|
let view_id = self
|
|
.next_table_id
|
|
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
|
let mut view_info = task.view_info;
|
|
view_info.ident.table_id = view_id;
|
|
self.catalog_manager
|
|
.register_table_sync(catalog::RegisterTableRequest {
|
|
catalog: task.create_view.catalog_name.clone(),
|
|
schema: task.create_view.schema_name.clone(),
|
|
table_name: task.create_view.view_name.clone(),
|
|
table_id: view_id,
|
|
table: table::dist_table::DistTable::table(Arc::new(view_info)),
|
|
})
|
|
.map_err(BoxedError::new)
|
|
.context(common_meta::error::ExternalSnafu)?;
|
|
Ok(SubmitDdlTaskResponse {
|
|
key: Vec::new(),
|
|
table_ids: vec![view_id],
|
|
})
|
|
}
|
|
DdlTask::DropView(task) => {
|
|
self.catalog_manager
|
|
.deregister_table_sync(catalog::DeregisterTableRequest {
|
|
catalog: task.catalog.clone(),
|
|
schema: task.schema.clone(),
|
|
table_name: task.view.clone(),
|
|
})
|
|
.map_err(BoxedError::new)
|
|
.context(common_meta::error::ExternalSnafu)?;
|
|
Ok(SubmitDdlTaskResponse::default())
|
|
}
|
|
other => common_meta::error::UnsupportedSnafu {
|
|
operation: format!("mock submit_ddl_task: {other:?}"),
|
|
}
|
|
.fail(),
|
|
}
|
|
}
|
|
|
|
async fn migrate_region(
|
|
&self,
|
|
_ctx: &ExecutorContext,
|
|
_request: MigrateRegionRequest,
|
|
) -> common_meta::error::Result<MigrateRegionResponse> {
|
|
common_meta::error::UnsupportedSnafu {
|
|
operation: "migrate_region",
|
|
}
|
|
.fail()
|
|
}
|
|
|
|
async fn reconcile(
|
|
&self,
|
|
_ctx: &ExecutorContext,
|
|
_request: ReconcileRequest,
|
|
) -> common_meta::error::Result<ReconcileResponse> {
|
|
common_meta::error::UnsupportedSnafu {
|
|
operation: "reconcile",
|
|
}
|
|
.fail()
|
|
}
|
|
|
|
async fn query_procedure_state(
|
|
&self,
|
|
_ctx: &ExecutorContext,
|
|
_pid: &str,
|
|
) -> common_meta::error::Result<ProcedureStateResponse> {
|
|
common_meta::error::UnsupportedSnafu {
|
|
operation: "query_procedure_state",
|
|
}
|
|
.fail()
|
|
}
|
|
|
|
async fn list_procedures(
|
|
&self,
|
|
_ctx: &ExecutorContext,
|
|
) -> common_meta::error::Result<ProcedureDetailResponse> {
|
|
common_meta::error::UnsupportedSnafu {
|
|
operation: "list_procedures",
|
|
}
|
|
.fail()
|
|
}
|
|
}
|
|
|
|
fn test_cache_registry(
|
|
kv_backend: common_meta::kv_backend::KvBackendRef,
|
|
) -> TestResult<common_meta::cache::LayeredCacheRegistryRef> {
|
|
Ok(Arc::new(
|
|
cache::with_default_composite_cache_registry(
|
|
LayeredCacheRegistryBuilder::default()
|
|
.add_cache_registry(cache::build_fundamental_cache_registry(kv_backend)),
|
|
)
|
|
.context(BuildCacheRegistrySnafu)?
|
|
.build(),
|
|
))
|
|
}
|
|
|
|
fn test_table_info(table_id: u32, table_name: &str) -> TestResult<TableInfo> {
|
|
let schema = Arc::new(GtSchema::new(vec![
|
|
ColumnSchema::new("id", ConcreteDataType::int32_datatype(), false),
|
|
ColumnSchema::new(
|
|
"ts",
|
|
ConcreteDataType::timestamp_millisecond_datatype(),
|
|
false,
|
|
)
|
|
.with_time_index(true),
|
|
]));
|
|
let table_meta = TableMetaBuilder::empty()
|
|
.schema(schema)
|
|
.primary_key_indices(vec![0])
|
|
.value_indices(vec![1])
|
|
.next_column_id(1024)
|
|
.build()
|
|
.with_context(|_| BuildTableMetaSnafu {
|
|
table_name: table_name.to_string(),
|
|
})?;
|
|
|
|
TableInfoBuilder::new(table_name, table_meta)
|
|
.table_id(table_id)
|
|
.build()
|
|
.with_context(|_| BuildTableInfoSnafu {
|
|
table_name: table_name.to_string(),
|
|
})
|
|
}
|
|
|
|
fn test_table(table_id: u32, table_name: &str) -> TestResult<table::TableRef> {
|
|
let table_info = test_table_info(table_id, table_name)?;
|
|
Ok(EmptyTable::from_table_info(&table_info))
|
|
}
|
|
|
|
fn test_physical_table(table_id: u32, table_name: &str) -> TestResult<table::TableRef> {
|
|
let mut table_info = test_table_info(table_id, table_name)?;
|
|
table_info
|
|
.meta
|
|
.options
|
|
.extra_options
|
|
.insert(PHYSICAL_TABLE_METADATA_KEY.to_string(), String::new());
|
|
Ok(EmptyTable::from_table_info(&table_info))
|
|
}
|
|
|
|
fn test_logical_table(table_id: u32, table_name: &str) -> TestResult<table::TableRef> {
|
|
let mut table_info = test_table_info(table_id, table_name)?;
|
|
table_info.meta.engine = METRIC_ENGINE_NAME.to_string();
|
|
table_info.meta.options.extra_options.insert(
|
|
LOGICAL_TABLE_METADATA_KEY.to_string(),
|
|
"physical_metric".to_string(),
|
|
);
|
|
Ok(EmptyTable::from_table_info(&table_info))
|
|
}
|
|
|
|
fn test_metric_names_table() -> TableRef {
|
|
let schema = Arc::new(GtSchema::new(vec![
|
|
ColumnSchema::new("table_catalog", ConcreteDataType::string_datatype(), false),
|
|
ColumnSchema::new("table_schema", ConcreteDataType::string_datatype(), false),
|
|
ColumnSchema::new("table_name", ConcreteDataType::string_datatype(), false),
|
|
ColumnSchema::new("engine", ConcreteDataType::string_datatype(), false),
|
|
ColumnSchema::new("create_options", ConcreteDataType::string_datatype(), false),
|
|
]));
|
|
let columns: Vec<VectorRef> = vec![
|
|
Arc::new(StringVector::from(vec!["greptime", "greptime"])),
|
|
Arc::new(StringVector::from(vec!["public", "public"])),
|
|
Arc::new(StringVector::from(vec!["denied", "target"])),
|
|
Arc::new(StringVector::from(vec!["metric", "metric"])),
|
|
Arc::new(StringVector::from(vec![
|
|
"on_physical_table=physical_metric",
|
|
"on_physical_table=physical_metric",
|
|
])),
|
|
];
|
|
let record_batch = RecordBatch::new(schema, columns).unwrap();
|
|
MemTable::new_with_catalog(
|
|
"tables",
|
|
record_batch,
|
|
2048,
|
|
"greptime".to_string(),
|
|
"information_schema".to_string(),
|
|
)
|
|
}
|
|
|
|
fn test_pipeline_table() -> TableRef {
|
|
let schema = Arc::new(GtSchema::new(vec![
|
|
ColumnSchema::new("name", ConcreteDataType::string_datatype(), false),
|
|
ColumnSchema::new("schema", ConcreteDataType::string_datatype(), false),
|
|
ColumnSchema::new("content_type", ConcreteDataType::string_datatype(), false),
|
|
ColumnSchema::new("pipeline", ConcreteDataType::string_datatype(), false),
|
|
ColumnSchema::new(
|
|
"created_at",
|
|
ConcreteDataType::timestamp_nanosecond_datatype(),
|
|
false,
|
|
)
|
|
.with_time_index(true),
|
|
]));
|
|
let columns: Vec<VectorRef> = vec![
|
|
Arc::new(StringVector::from(vec!["pipeline"])),
|
|
Arc::new(StringVector::from(vec!["public"])),
|
|
Arc::new(StringVector::from(vec!["application/yaml"])),
|
|
Arc::new(StringVector::from(vec![
|
|
"transform:\n- field: ts\n type: timestamp, ns\n index: time\n",
|
|
])),
|
|
Arc::new(TimestampNanosecondVector::from_values([1])),
|
|
];
|
|
let record_batch = RecordBatch::new(schema, columns).unwrap();
|
|
MemTable::new_with_catalog(
|
|
"pipelines",
|
|
record_batch,
|
|
2049,
|
|
"greptime".to_string(),
|
|
DEFAULT_PRIVATE_SCHEMA_NAME.to_string(),
|
|
)
|
|
}
|
|
|
|
struct CancellableExportSource {
|
|
schema: GtSchemaRef,
|
|
channels: std::sync::Mutex<Option<(oneshot::Sender<()>, oneshot::Receiver<()>)>>,
|
|
}
|
|
|
|
impl DataSource for CancellableExportSource {
|
|
fn get_stream(
|
|
&self,
|
|
_request: ScanRequest,
|
|
) -> std::result::Result<SendableRecordBatchStream, BoxedError> {
|
|
let (started, release) = self.channels.lock().unwrap().take().unwrap();
|
|
let schema = self.schema.clone();
|
|
let stream = futures::stream::once(async move {
|
|
started.send(()).unwrap();
|
|
release.await.unwrap();
|
|
Ok(RecordBatch::new_empty(schema))
|
|
});
|
|
Ok(Box::pin(RecordBatchStreamWrapper::new(
|
|
self.schema.clone(),
|
|
Box::pin(stream),
|
|
)))
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_metric_export_http_timeout_cancels_ordinary_source() {
|
|
let destination = common_test_util::temp_dir::create_temp_dir("metric_export_timeout");
|
|
let (started_tx, started_rx) = oneshot::channel();
|
|
let (mut release_tx, release_rx) = oneshot::channel();
|
|
let info = test_table_info(1024, "source").unwrap();
|
|
let source = Arc::new(Table::new(
|
|
Arc::new(info.clone()),
|
|
FilterPushDownType::Unsupported,
|
|
Arc::new(CancellableExportSource {
|
|
schema: info.meta.schema.clone(),
|
|
channels: std::sync::Mutex::new(Some((started_tx, release_rx))),
|
|
}),
|
|
));
|
|
let catalog = catalog::memory::MemoryCatalogManager::new_with_table(source);
|
|
let kv = Arc::new(MemoryKvBackend::new());
|
|
let instance = FrontendBuilder::new(
|
|
FrontendOptions {
|
|
experimental_metric_export: true,
|
|
..Default::default()
|
|
},
|
|
kv.clone(),
|
|
test_cache_registry(kv).unwrap(),
|
|
catalog,
|
|
Arc::new(client::client_manager::NodeClients::default()),
|
|
Arc::new(NoopProcedureExecutor),
|
|
Arc::new(ProcessManager::new("export-timeout".into(), None)),
|
|
)
|
|
.with_local_file_access(
|
|
common_datasource::object_store::LocalFileAccess::sandboxed(destination.path())
|
|
.unwrap(),
|
|
)
|
|
.try_build()
|
|
.await
|
|
.unwrap();
|
|
let server = servers::http::HttpServerBuilder::new(servers::http::HttpOptions {
|
|
timeout: Duration::from_secs(2),
|
|
..Default::default()
|
|
})
|
|
.with_sql_handler(Arc::new(instance))
|
|
.build();
|
|
let app = server.build(server.make_app()).unwrap();
|
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
|
let addr = listener.local_addr().unwrap();
|
|
let server_task = tokio::spawn(async move {
|
|
axum::serve(listener, app).await.unwrap();
|
|
});
|
|
let uri = reqwest::Url::from_directory_path(destination.path()).unwrap();
|
|
let sql = format!(
|
|
"COPY DATABASE greptime.public TO '{uri}' WITH (experimental_metric_export='true')"
|
|
);
|
|
let request = tokio::spawn(async move {
|
|
reqwest::Client::builder()
|
|
.no_proxy()
|
|
.build()
|
|
.unwrap()
|
|
.post(format!("http://{addr}/v1/sql"))
|
|
.form(&[("sql", sql)])
|
|
.send()
|
|
.await
|
|
.unwrap()
|
|
});
|
|
tokio::time::timeout(Duration::from_secs(5), started_rx)
|
|
.await
|
|
.unwrap()
|
|
.unwrap();
|
|
let response = request.await.unwrap();
|
|
assert_eq!(response.status(), reqwest::StatusCode::REQUEST_TIMEOUT);
|
|
tokio::time::timeout(Duration::from_secs(5), release_tx.closed())
|
|
.await
|
|
.unwrap();
|
|
assert!(!destination.path().join("source.parquet").exists());
|
|
server_task.abort();
|
|
}
|
|
|
|
fn pending_table(
|
|
table_id: u32,
|
|
table_name: &str,
|
|
polled_tx: oneshot::Sender<()>,
|
|
) -> TestResult<table::TableRef> {
|
|
let table_info = test_table_info(table_id, table_name)?;
|
|
let data_source = Arc::new(PendingDataSource {
|
|
schema: table_info.meta.schema.clone(),
|
|
polled_tx: std::sync::Mutex::new(Some(polled_tx)),
|
|
});
|
|
|
|
Ok(Arc::new(Table::new(
|
|
Arc::new(table_info),
|
|
FilterPushDownType::Unsupported,
|
|
data_source,
|
|
)))
|
|
}
|
|
|
|
async fn test_instance_with_tables(
|
|
source_table: TableRef,
|
|
target_table: TableRef,
|
|
) -> TestResult<Instance> {
|
|
test_instance_with_plugins(source_table, target_table, Plugins::new()).await
|
|
}
|
|
|
|
async fn test_instance_with_insert_select_interceptor(
|
|
interceptor: SqlQueryInterceptorRef<Error>,
|
|
) -> TestResult<Instance> {
|
|
let plugins = Plugins::new();
|
|
plugins.insert::<SqlQueryInterceptorRef<Error>>(interceptor);
|
|
|
|
test_instance_with_plugins(
|
|
test_table(1024, "source")?,
|
|
test_table(1025, "target")?,
|
|
plugins,
|
|
)
|
|
.await
|
|
}
|
|
|
|
async fn test_instance_with_plugins(
|
|
source_table: TableRef,
|
|
target_table: TableRef,
|
|
plugins: Plugins,
|
|
) -> TestResult<Instance> {
|
|
test_instance_with_plugins_and_metric_names(source_table, target_table, plugins, None).await
|
|
}
|
|
|
|
async fn test_instance_with_plugins_and_metric_names(
|
|
source_table: TableRef,
|
|
target_table: TableRef,
|
|
plugins: Plugins,
|
|
metric_names_table: Option<TableRef>,
|
|
) -> TestResult<Instance> {
|
|
let catalog_manager = catalog::memory::MemoryCatalogManager::new_with_table(source_table);
|
|
test_instance_with_catalog_manager(
|
|
catalog_manager,
|
|
target_table,
|
|
plugins,
|
|
metric_names_table,
|
|
Arc::new(NoopProcedureExecutor),
|
|
)
|
|
.await
|
|
}
|
|
|
|
/// Builds a test frontend `Instance` over the given (already source-registered)
|
|
/// catalog manager, completing DDL tasks through `procedure_executor`.
|
|
async fn test_instance_with_catalog_manager(
|
|
catalog_manager: Arc<catalog::memory::MemoryCatalogManager>,
|
|
target_table: TableRef,
|
|
plugins: Plugins,
|
|
metric_names_table: Option<TableRef>,
|
|
procedure_executor: ProcedureExecutorRef,
|
|
) -> TestResult<Instance> {
|
|
let kv_backend = Arc::new(MemoryKvBackend::new());
|
|
let process_manager = Arc::new(ProcessManager::new("test-frontend".to_string(), None));
|
|
let target_table_name = "target";
|
|
catalog_manager
|
|
.register_table_sync(catalog::RegisterTableRequest {
|
|
catalog: "greptime".to_string(),
|
|
schema: "public".to_string(),
|
|
table_name: target_table_name.to_string(),
|
|
table_id: 1025,
|
|
table: target_table,
|
|
})
|
|
.with_context(|_| RegisterTableSnafu {
|
|
table_name: target_table_name.to_string(),
|
|
})?;
|
|
if let Some(table) = metric_names_table {
|
|
catalog_manager
|
|
.deregister_table_sync(catalog::DeregisterTableRequest {
|
|
catalog: "greptime".to_string(),
|
|
schema: "information_schema".to_string(),
|
|
table_name: "tables".to_string(),
|
|
})
|
|
.unwrap();
|
|
catalog_manager
|
|
.register_table_sync(catalog::RegisterTableRequest {
|
|
catalog: "greptime".to_string(),
|
|
schema: "information_schema".to_string(),
|
|
table_name: "tables".to_string(),
|
|
table_id: 2048,
|
|
table,
|
|
})
|
|
.unwrap();
|
|
}
|
|
catalog_manager.register_process_list_table(process_manager.clone());
|
|
|
|
let cache_registry = test_cache_registry(kv_backend.clone())?;
|
|
|
|
FrontendBuilder::new(
|
|
FrontendOptions::default(),
|
|
kv_backend,
|
|
cache_registry,
|
|
catalog_manager,
|
|
Arc::new(client::client_manager::NodeClients::default()),
|
|
procedure_executor,
|
|
process_manager,
|
|
)
|
|
.with_plugin(plugins)
|
|
.try_build()
|
|
.await
|
|
.context(BuildFrontendSnafu)
|
|
}
|
|
|
|
async fn execute_one_sql(
|
|
instance: &Instance,
|
|
sql: &str,
|
|
query_ctx: QueryContextRef,
|
|
) -> TestResult<Output> {
|
|
let mut results = instance.do_query_inner(sql, query_ctx).await;
|
|
ensure!(
|
|
results.len() == 1,
|
|
UnexpectedOutputCountSnafu {
|
|
sql: sql.to_string(),
|
|
actual: results.len(),
|
|
}
|
|
);
|
|
results.remove(0).with_context(|_| ExecuteSqlSnafu {
|
|
sql: sql.to_string(),
|
|
})
|
|
}
|
|
|
|
fn assert_permission_denied<T>(result: servers::error::Result<T>) {
|
|
let err = match result {
|
|
Ok(_) => panic!("request should be rejected"),
|
|
Err(err) => err,
|
|
};
|
|
assert_eq!(StatusCode::PermissionDenied, err.status_code());
|
|
}
|
|
|
|
fn assert_action_checked(
|
|
checker: &RejectEndpointPermissionChecker,
|
|
action: PermissionAction,
|
|
targets: Option<PermissionTableTargets>,
|
|
) {
|
|
assert_eq!(CheckedAction { action, targets }, checker.take_check());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn database_export_authorizes_all_tables_before_preparation() -> TestResult<()> {
|
|
struct ExportAcl(std::sync::Mutex<Vec<PermissionTableTargets>>);
|
|
impl PermissionChecker for ExportAcl {
|
|
fn check_permission(
|
|
&self,
|
|
_: UserInfoRef,
|
|
req: PermissionReq,
|
|
) -> auth::error::Result<PermissionResp> {
|
|
assert!(matches!(
|
|
req,
|
|
PermissionReq::SqlStatement(Statement::Copy(
|
|
sql::statements::copy::Copy::CopyDatabase(
|
|
sql::statements::copy::CopyDatabase::To(_)
|
|
)
|
|
))
|
|
));
|
|
Ok(PermissionResp::Allow)
|
|
}
|
|
fn check_permission_with_table_targets(
|
|
&self,
|
|
_: UserInfoRef,
|
|
req: PermissionReq,
|
|
targets: PermissionTableTargets,
|
|
) -> auth::error::Result<PermissionResp> {
|
|
self.0.lock().unwrap().push(targets.clone());
|
|
let PermissionTableTargets::Resolved(tables) = targets else {
|
|
panic!("unresolved export")
|
|
};
|
|
if req.is_readonly() {
|
|
// The first table is readable; the later table has only write access.
|
|
Ok(if tables.iter().any(|t| t.table == "target") {
|
|
PermissionResp::Reject
|
|
} else {
|
|
PermissionResp::Allow
|
|
})
|
|
} else {
|
|
self.check_permission(QueryContext::arc().current_user(), req)
|
|
}
|
|
}
|
|
}
|
|
let checker = Arc::new(ExportAcl(Default::default()));
|
|
let plugins = Plugins::new();
|
|
plugins.insert::<PermissionCheckerRef>(checker.clone());
|
|
let instance = test_instance_with_plugins(
|
|
test_logical_table(1024, "source")?,
|
|
test_table(1025, "target")?,
|
|
plugins,
|
|
)
|
|
.await?;
|
|
let req = table::requests::CopyDatabaseRequest {
|
|
catalog_name: "greptime".into(),
|
|
schema_name: "public".into(),
|
|
location: "invalid-destination".into(),
|
|
with: Default::default(),
|
|
connection: Default::default(),
|
|
time_range: None,
|
|
};
|
|
let result = instance
|
|
.export_database_for_test(
|
|
req.clone(),
|
|
None,
|
|
&tokio_util::sync::CancellationToken::new(),
|
|
QueryContext::arc(),
|
|
)
|
|
.await;
|
|
assert!(matches!(result, Err(Error::Permission { .. })));
|
|
let expected = PermissionTableTargets::resolved(vec![
|
|
PermissionTableTarget::new("greptime", "public", "source"),
|
|
PermissionTableTarget::new("greptime", "public", "target"),
|
|
]);
|
|
assert_eq!(*checker.0.lock().unwrap(), vec![expected.clone(), expected]);
|
|
// An empty selection must still check the operation privilege.
|
|
instance
|
|
.plugins
|
|
.map_mut::<PermissionCheckerRef, _, _>(|checker| {
|
|
*checker.unwrap() = Arc::new(WriteOnlyPermissionChecker)
|
|
});
|
|
let result = instance
|
|
.export_database_for_test(
|
|
req,
|
|
Some(&[]),
|
|
&tokio_util::sync::CancellationToken::new(),
|
|
QueryContext::arc(),
|
|
)
|
|
.await;
|
|
assert!(matches!(result, Err(Error::Permission { .. })));
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_prom_remote_read_with_custom_timestamp_and_value_columns() -> TestResult<()> {
|
|
let schema = Arc::new(GtSchema::new(vec![
|
|
ColumnSchema::new(
|
|
"custom_ts",
|
|
ConcreteDataType::timestamp_millisecond_datatype(),
|
|
false,
|
|
)
|
|
.with_time_index(true),
|
|
ColumnSchema::new("custom_value", ConcreteDataType::float64_datatype(), false),
|
|
]));
|
|
let recordbatch = RecordBatch::new(
|
|
schema,
|
|
vec![
|
|
Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as VectorRef,
|
|
Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as VectorRef,
|
|
],
|
|
)
|
|
.unwrap();
|
|
let instance = test_instance_with_tables(
|
|
MemTable::table("custom_metric", recordbatch),
|
|
test_table(1025, "target")?,
|
|
)
|
|
.await?;
|
|
|
|
let response = PromStoreProtocolHandler::read(
|
|
&instance,
|
|
ReadRequest {
|
|
queries: vec![RemoteQuery {
|
|
start_timestamp_ms: 1500,
|
|
end_timestamp_ms: 2500,
|
|
matchers: vec![LabelMatcher {
|
|
r#type: PromMatcherType::Eq as i32,
|
|
name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
|
|
value: "custom_metric".to_string(),
|
|
}],
|
|
..Default::default()
|
|
}],
|
|
..Default::default()
|
|
},
|
|
test_query_ctx(1),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
let body = servers::prom_store::snappy_decompress(&response.body).unwrap();
|
|
let response = ReadResponse::decode(body.as_slice()).unwrap();
|
|
|
|
assert_eq!(1, response.results.len());
|
|
assert_eq!(1, response.results[0].timeseries.len());
|
|
let timeseries = &response.results[0].timeseries[0];
|
|
assert_eq!(
|
|
vec![Label {
|
|
name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
|
|
value: "custom_metric".to_string(),
|
|
}],
|
|
timeseries.labels
|
|
);
|
|
assert_eq!(
|
|
vec![Sample {
|
|
value: 2.0,
|
|
timestamp: 2000,
|
|
}],
|
|
timeseries.samples
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_prom_remote_read_prefers_default_value_column() -> TestResult<()> {
|
|
let schema = Arc::new(GtSchema::new(vec![
|
|
ColumnSchema::new(
|
|
"custom_ts",
|
|
ConcreteDataType::timestamp_millisecond_datatype(),
|
|
false,
|
|
)
|
|
.with_time_index(true),
|
|
ColumnSchema::new("extra_field", ConcreteDataType::float64_datatype(), false),
|
|
ColumnSchema::new(
|
|
greptime_value(),
|
|
ConcreteDataType::float64_datatype(),
|
|
false,
|
|
),
|
|
]));
|
|
let recordbatch = RecordBatch::new(
|
|
schema,
|
|
vec![
|
|
Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as VectorRef,
|
|
Arc::new(Float64Vector::from_vec(vec![99.0, 99.0, 99.0])) as VectorRef,
|
|
Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as VectorRef,
|
|
],
|
|
)
|
|
.unwrap();
|
|
let instance = test_instance_with_tables(
|
|
MemTable::table("multi_field_metric", recordbatch),
|
|
test_table(1025, "target")?,
|
|
)
|
|
.await?;
|
|
|
|
let response = PromStoreProtocolHandler::read(
|
|
&instance,
|
|
ReadRequest {
|
|
queries: vec![RemoteQuery {
|
|
start_timestamp_ms: 1500,
|
|
end_timestamp_ms: 2500,
|
|
matchers: vec![LabelMatcher {
|
|
r#type: PromMatcherType::Eq as i32,
|
|
name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
|
|
value: "multi_field_metric".to_string(),
|
|
}],
|
|
..Default::default()
|
|
}],
|
|
..Default::default()
|
|
},
|
|
test_query_ctx(1),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
let body = servers::prom_store::snappy_decompress(&response.body).unwrap();
|
|
let response = ReadResponse::decode(body.as_slice()).unwrap();
|
|
|
|
assert_eq!(1, response.results.len());
|
|
assert_eq!(1, response.results[0].timeseries.len());
|
|
let timeseries = &response.results[0].timeseries[0];
|
|
assert_eq!(
|
|
vec![
|
|
Label {
|
|
name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
|
|
value: "multi_field_metric".to_string(),
|
|
},
|
|
Label {
|
|
name: "extra_field".to_string(),
|
|
value: "99".to_string(),
|
|
},
|
|
],
|
|
timeseries.labels
|
|
);
|
|
assert_eq!(
|
|
vec![Sample {
|
|
value: 2.0,
|
|
timestamp: 2000,
|
|
}],
|
|
timeseries.samples
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_prom_remote_read_rejects_ambiguous_value_columns() -> TestResult<()> {
|
|
let schema = Arc::new(GtSchema::new(vec![
|
|
ColumnSchema::new(
|
|
"custom_ts",
|
|
ConcreteDataType::timestamp_millisecond_datatype(),
|
|
false,
|
|
)
|
|
.with_time_index(true),
|
|
ColumnSchema::new("field_a", ConcreteDataType::float64_datatype(), false),
|
|
ColumnSchema::new("field_b", ConcreteDataType::float64_datatype(), false),
|
|
]));
|
|
let recordbatch = RecordBatch::new(
|
|
schema,
|
|
vec![
|
|
Arc::new(TimestampMillisecondVector::from_vec(vec![1000])) as VectorRef,
|
|
Arc::new(Float64Vector::from_vec(vec![1.0])) as VectorRef,
|
|
Arc::new(Float64Vector::from_vec(vec![2.0])) as VectorRef,
|
|
],
|
|
)
|
|
.unwrap();
|
|
let instance = test_instance_with_tables(
|
|
MemTable::table("ambiguous_metric", recordbatch),
|
|
test_table(1025, "target")?,
|
|
)
|
|
.await?;
|
|
|
|
let err = PromStoreProtocolHandler::read(
|
|
&instance,
|
|
ReadRequest {
|
|
queries: vec![RemoteQuery {
|
|
matchers: vec![LabelMatcher {
|
|
r#type: PromMatcherType::Eq as i32,
|
|
name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
|
|
value: "ambiguous_metric".to_string(),
|
|
}],
|
|
..Default::default()
|
|
}],
|
|
..Default::default()
|
|
},
|
|
test_query_ctx(1),
|
|
)
|
|
.await
|
|
.err()
|
|
.expect("ambiguous value columns should fail remote read");
|
|
|
|
assert_eq!(StatusCode::InvalidArguments, err.status_code());
|
|
assert!(format!("{err:?}").contains("Ambiguous value column"));
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_event_recorder_is_exposed() -> TestResult<()> {
|
|
let instance =
|
|
test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
|
|
.await?;
|
|
|
|
let _event_recorder = instance.event_recorder();
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_restricted_endpoint_handlers_check_permissions() -> TestResult<()> {
|
|
let checker = Arc::new(RejectEndpointPermissionChecker::default());
|
|
let plugins = Plugins::new();
|
|
plugins.insert::<PermissionCheckerRef>(checker.clone());
|
|
let instance = test_instance_with_plugins(
|
|
test_table(1024, "denied")?,
|
|
test_table(1025, "target")?,
|
|
plugins,
|
|
)
|
|
.await?;
|
|
let mut ctx = test_query_ctx(1);
|
|
Arc::get_mut(&mut ctx).unwrap().set_extension(
|
|
servers::http::jaeger::JAEGER_QUERY_TABLE_NAME_KEY,
|
|
"denied".to_string(),
|
|
);
|
|
let jaeger_targets = Some(PermissionTableTargets::resolved(vec![
|
|
PermissionTableTarget::new("greptime", "public", "denied"),
|
|
]));
|
|
|
|
assert_permission_denied(JaegerQueryHandler::get_services(&instance, ctx.clone()).await);
|
|
assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
|
|
assert_permission_denied(
|
|
JaegerQueryHandler::get_operations(&instance, ctx.clone(), "service", None).await,
|
|
);
|
|
assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
|
|
assert_permission_denied(
|
|
JaegerQueryHandler::get_trace(&instance, ctx.clone(), "trace", None, None, None).await,
|
|
);
|
|
assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets.clone());
|
|
assert_permission_denied(
|
|
JaegerQueryHandler::find_traces(
|
|
&instance,
|
|
ctx.clone(),
|
|
servers::http::jaeger::QueryTraceParams {
|
|
service_name: "service".to_string(),
|
|
..Default::default()
|
|
},
|
|
)
|
|
.await,
|
|
);
|
|
assert_action_checked(&checker, JAEGER_QUERY, jaeger_targets);
|
|
|
|
assert_permission_denied(
|
|
PipelineHandler::get_pipeline_str(&instance, "pipeline", None, ctx.clone()).await,
|
|
);
|
|
assert_action_checked(&checker, PIPELINE_QUERY, None);
|
|
assert_permission_denied(
|
|
PipelineHandler::insert_pipeline(
|
|
&instance,
|
|
"pipeline",
|
|
"application/yaml",
|
|
"",
|
|
ctx.clone(),
|
|
)
|
|
.await,
|
|
);
|
|
assert_action_checked(&checker, PIPELINE_INSERT, None);
|
|
assert_permission_denied(
|
|
PipelineHandler::delete_pipeline(&instance, "pipeline", None, ctx.clone()).await,
|
|
);
|
|
assert_action_checked(&checker, PIPELINE_DELETE, None);
|
|
let app = axum::Router::new()
|
|
.route(
|
|
"/pipelines/_dryrun",
|
|
axum::routing::post(servers::http::event::pipeline_dryrun),
|
|
)
|
|
.with_state(servers::http::event::LogState {
|
|
log_handler: Arc::new(instance.clone()),
|
|
log_validator: None,
|
|
ingest_interceptor: None,
|
|
})
|
|
.layer(axum::Extension((*ctx).clone()));
|
|
let response = app
|
|
.oneshot(
|
|
axum::http::Request::post("/pipelines/_dryrun")
|
|
.header("content-type", "application/json")
|
|
.body(axum::body::Body::from("{}"))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(axum::http::StatusCode::FORBIDDEN, response.status());
|
|
assert_action_checked(&checker, PIPELINE_QUERY, None);
|
|
|
|
assert_permission_denied(
|
|
DashboardHandler::save(&instance, "dashboard", "{}", ctx.clone()).await,
|
|
);
|
|
assert_action_checked(&checker, DASHBOARD_SAVE, None);
|
|
assert_permission_denied(DashboardHandler::list(&instance, ctx.clone()).await);
|
|
assert_action_checked(&checker, DASHBOARD_QUERY, None);
|
|
assert_permission_denied(
|
|
DashboardHandler::delete(&instance, "dashboard", ctx.clone()).await,
|
|
);
|
|
assert_action_checked(&checker, DASHBOARD_DELETE, None);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_write_only_ingestion_loads_named_pipeline() -> TestResult<()> {
|
|
let plugins = Plugins::new();
|
|
plugins.insert::<PermissionCheckerRef>(Arc::new(WriteOnlyPermissionChecker));
|
|
let instance = test_instance_with_plugins(
|
|
test_table(1024, "source")?,
|
|
test_table(1025, "target")?,
|
|
plugins,
|
|
)
|
|
.await?;
|
|
instance
|
|
.catalog_manager()
|
|
.as_any()
|
|
.downcast_ref::<catalog::memory::MemoryCatalogManager>()
|
|
.unwrap()
|
|
.register_table_sync(catalog::RegisterTableRequest {
|
|
catalog: "greptime".to_string(),
|
|
schema: DEFAULT_PRIVATE_SCHEMA_NAME.to_string(),
|
|
table_name: "pipelines".to_string(),
|
|
table_id: 2049,
|
|
table: test_pipeline_table(),
|
|
})
|
|
.with_context(|_| RegisterTableSnafu {
|
|
table_name: "pipelines".to_string(),
|
|
})?;
|
|
let ctx = test_query_ctx(1);
|
|
let handler: PipelineHandlerRef = Arc::new(instance.clone());
|
|
|
|
handler
|
|
.get_pipeline("pipeline", None, ctx.clone())
|
|
.await
|
|
.unwrap();
|
|
assert_permission_denied(
|
|
PipelineHandler::get_pipeline_str(&instance, "pipeline", None, ctx.clone()).await,
|
|
);
|
|
|
|
let app = axum::Router::new()
|
|
.route(
|
|
"/pipelines/_dryrun",
|
|
axum::routing::post(servers::http::event::pipeline_dryrun),
|
|
)
|
|
.with_state(servers::http::event::LogState {
|
|
log_handler: handler,
|
|
log_validator: None,
|
|
ingest_interceptor: None,
|
|
})
|
|
.layer(axum::Extension((*ctx).clone()));
|
|
let response = app
|
|
.oneshot(
|
|
axum::http::Request::post("/pipelines/_dryrun")
|
|
.header("content-type", "application/json")
|
|
.body(axum::body::Body::from("{}"))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(axum::http::StatusCode::FORBIDDEN, response.status());
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_write_only_grpc_sql_is_checked_after_parsing() -> TestResult<()> {
|
|
let plugins = Plugins::new();
|
|
plugins.insert::<PermissionCheckerRef>(Arc::new(WriteOnlyPermissionChecker));
|
|
let instance = test_instance_with_plugins(
|
|
test_table(1024, "source")?,
|
|
test_table(1025, "target")?,
|
|
plugins,
|
|
)
|
|
.await?;
|
|
|
|
let insert = Request::Query(api::v1::QueryRequest {
|
|
query: Some(Query::Sql(
|
|
"INSERT INTO target SELECT * FROM source".to_string(),
|
|
)),
|
|
});
|
|
servers::query_handler::grpc::GrpcQueryHandler::do_query(
|
|
&instance,
|
|
insert,
|
|
QueryContext::arc(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
let select = Request::Query(api::v1::QueryRequest {
|
|
query: Some(Query::Sql("SELECT * FROM source".to_string())),
|
|
});
|
|
assert_permission_denied(
|
|
servers::query_handler::grpc::GrpcQueryHandler::do_query(
|
|
&instance,
|
|
select,
|
|
QueryContext::arc(),
|
|
)
|
|
.await,
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_target_independent_checker_skips_target_resolution() -> TestResult<()> {
|
|
let physical_table = "physical_metric";
|
|
let checker = Arc::new(TargetIndependentPermissionChecker::default());
|
|
let plugins = Plugins::new();
|
|
plugins.insert::<PermissionCheckerRef>(checker.clone());
|
|
let instance = test_instance_with_plugins(
|
|
test_physical_table(1024, physical_table)?,
|
|
test_table(1025, "target")?,
|
|
plugins,
|
|
)
|
|
.await?;
|
|
|
|
let ctx = test_query_ctx(1);
|
|
let physical_target = PermissionTableTarget::new("greptime", "public", physical_table);
|
|
assert_eq!(
|
|
PermissionTableTargets::Resolved(vec![physical_target.clone()]),
|
|
instance
|
|
.resolve_query_permission_targets(
|
|
PermissionTableTargets::resolved(vec![physical_target]),
|
|
&ctx,
|
|
)
|
|
.await
|
|
.unwrap()
|
|
);
|
|
assert_eq!(
|
|
vec![physical_table.to_string(), "target".to_string()],
|
|
PrometheusHandler::filter_metadata_metric_names(
|
|
&instance,
|
|
vec![physical_table.to_string(), "target".to_string()],
|
|
"public",
|
|
&ctx,
|
|
)
|
|
.await
|
|
.unwrap()
|
|
);
|
|
assert_eq!(1, checker.checks.load(atomic::Ordering::Relaxed));
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_query_permission_targets_are_deduplicated() -> TestResult<()> {
|
|
let plugins = Plugins::new();
|
|
plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
|
|
let instance = test_instance_with_plugins(
|
|
test_table(1024, "source")?,
|
|
test_table(1025, "target")?,
|
|
plugins,
|
|
)
|
|
.await?;
|
|
let ctx = test_query_ctx(1);
|
|
let target = PermissionTableTarget::new("greptime", "public", "target");
|
|
|
|
assert_eq!(
|
|
PermissionTableTargets::Resolved(vec![target.clone()]),
|
|
instance
|
|
.resolve_query_permission_targets(
|
|
PermissionTableTargets::resolved(vec![target.clone(), target]),
|
|
&ctx,
|
|
)
|
|
.await
|
|
.unwrap()
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_physical_query_targets_fail_closed() -> TestResult<()> {
|
|
let physical_table = "physical_metric";
|
|
let plugins = Plugins::new();
|
|
plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
|
|
let instance = test_instance_with_plugins(
|
|
test_physical_table(1024, physical_table)?,
|
|
test_table(1025, "target")?,
|
|
plugins,
|
|
)
|
|
.await?;
|
|
|
|
let ctx = test_query_ctx(1);
|
|
let logical_target = PermissionTableTarget::new("greptime", "public", "target");
|
|
assert_eq!(
|
|
PermissionTableTargets::Resolved(vec![logical_target.clone()]),
|
|
instance
|
|
.resolve_query_permission_targets(
|
|
PermissionTableTargets::resolved(vec![logical_target.clone()]),
|
|
&ctx,
|
|
)
|
|
.await
|
|
.unwrap()
|
|
);
|
|
let physical_target = PermissionTableTarget::new("greptime", "public", physical_table);
|
|
assert_eq!(
|
|
PermissionTableTargets::Unresolved,
|
|
instance
|
|
.resolve_query_permission_targets(
|
|
PermissionTableTargets::resolved(
|
|
vec![logical_target, physical_target.clone(),]
|
|
),
|
|
&ctx,
|
|
)
|
|
.await
|
|
.unwrap()
|
|
);
|
|
assert_eq!(
|
|
vec!["target".to_string()],
|
|
PrometheusHandler::filter_metadata_metric_names(
|
|
&instance,
|
|
vec!["target".to_string(), "denied".to_string()],
|
|
"public",
|
|
&ctx,
|
|
)
|
|
.await
|
|
.unwrap()
|
|
);
|
|
|
|
let query = PromQuery {
|
|
query: physical_table.to_string(),
|
|
..Default::default()
|
|
};
|
|
let err = PrometheusHandler::check_query_target_permission(
|
|
&instance,
|
|
PermissionTableTargets::resolved(vec![physical_target]),
|
|
&ctx,
|
|
)
|
|
.await
|
|
.unwrap_err();
|
|
assert_eq!(StatusCode::PermissionDenied, err.status_code());
|
|
let err = PrometheusHandler::check_query_permission(
|
|
&instance,
|
|
std::slice::from_ref(&query),
|
|
&ctx,
|
|
)
|
|
.await
|
|
.unwrap_err();
|
|
assert_eq!(StatusCode::PermissionDenied, err.status_code());
|
|
let err = PrometheusHandler::do_query(&instance, &query, ctx.clone())
|
|
.await
|
|
.unwrap_err();
|
|
assert_eq!(StatusCode::PermissionDenied, err.status_code());
|
|
|
|
for sql in [
|
|
"SELECT * FROM physical_metric",
|
|
"TQL EVAL (0, 10, '5s') physical_metric",
|
|
"INSERT INTO target SELECT * FROM physical_metric",
|
|
] {
|
|
let mut results = instance.do_query_inner(sql, ctx.clone()).await;
|
|
assert_eq!(1, results.len(), "{sql}");
|
|
let err = results.remove(0).unwrap_err();
|
|
assert_eq!(StatusCode::PermissionDenied, err.status_code(), "{sql}");
|
|
}
|
|
let err = LogQueryHandler::query(
|
|
&instance,
|
|
LogQuery {
|
|
table: TableName::new("greptime", "public", physical_table),
|
|
..Default::default()
|
|
},
|
|
ctx.clone(),
|
|
)
|
|
.await
|
|
.unwrap_err();
|
|
assert_eq!(StatusCode::PermissionDenied, err.status_code());
|
|
let err = instance
|
|
.do_describe_inner(parse_one_sql("SELECT * FROM physical_metric"), ctx.clone())
|
|
.await
|
|
.unwrap_err();
|
|
assert_eq!(StatusCode::PermissionDenied, err.status_code());
|
|
|
|
let request = ReadRequest {
|
|
queries: vec![RemoteQuery {
|
|
matchers: vec![LabelMatcher {
|
|
r#type: PromMatcherType::Eq as i32,
|
|
name: servers::prom_store::METRIC_NAME_LABEL.to_string(),
|
|
value: physical_table.to_string(),
|
|
}],
|
|
..Default::default()
|
|
}],
|
|
..Default::default()
|
|
};
|
|
let Err(err) = PromStoreProtocolHandler::read(&instance, request, ctx.clone()).await else {
|
|
panic!("physical remote-read target must be rejected");
|
|
};
|
|
assert_eq!(StatusCode::PermissionDenied, err.status_code());
|
|
|
|
let err = PrometheusHandler::query_label_values(
|
|
&instance,
|
|
physical_table.to_string(),
|
|
"host".to_string(),
|
|
vec![],
|
|
SystemTime::UNIX_EPOCH,
|
|
SystemTime::UNIX_EPOCH,
|
|
&ctx,
|
|
)
|
|
.await
|
|
.unwrap_err();
|
|
assert_eq!(StatusCode::PermissionDenied, err.status_code());
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_non_exact_query_discovery_keeps_denied_targets_for_batch_check() -> TestResult<()>
|
|
{
|
|
let plugins = Plugins::new();
|
|
plugins.insert::<PermissionCheckerRef>(Arc::new(RejectUnresolvedPermissionChecker));
|
|
let instance = test_instance_with_plugins_and_metric_names(
|
|
test_logical_table(1024, "denied")?,
|
|
test_logical_table(1025, "target")?,
|
|
plugins,
|
|
Some(test_metric_names_table()),
|
|
)
|
|
.await?;
|
|
let ctx = test_query_ctx(1);
|
|
|
|
let mut metric_names = PrometheusHandler::query_metric_names(
|
|
&instance,
|
|
vec![Matcher::new(
|
|
promql_parser::label::MatchOp::NotEqual,
|
|
"__name__",
|
|
"",
|
|
)],
|
|
"public",
|
|
&ctx,
|
|
)
|
|
.await
|
|
.unwrap();
|
|
metric_names.sort_unstable();
|
|
assert_eq!(
|
|
vec!["denied".to_string(), "target".to_string()],
|
|
metric_names
|
|
);
|
|
|
|
let queries = metric_names
|
|
.into_iter()
|
|
.map(|query| PromQuery {
|
|
query,
|
|
..Default::default()
|
|
})
|
|
.collect::<Vec<_>>();
|
|
let err = PrometheusHandler::check_query_permission(&instance, &queries, &ctx)
|
|
.await
|
|
.unwrap_err();
|
|
assert_eq!(StatusCode::PermissionDenied, err.status_code());
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_fast_legacy_check_is_read_only() {
|
|
let cache = DashMap::new();
|
|
cache.insert("metric1".to_string(), true);
|
|
|
|
let names = vec!["metric1".to_string(), "metric2".to_string()];
|
|
assert_eq!(Some(true), fast_legacy_check(&cache, &names).unwrap());
|
|
assert!(!cache.contains_key("metric2"));
|
|
|
|
cache_legacy_mode(&cache, &names, true).unwrap();
|
|
assert!(*cache.get("metric2").unwrap().value());
|
|
assert!(cache_legacy_mode(&cache, &names, false).is_err());
|
|
assert!(*cache.get("metric2").unwrap().value());
|
|
|
|
let cache_incompatible = DashMap::new();
|
|
cache_incompatible.insert("metric1".to_string(), true);
|
|
cache_incompatible.insert("metric2".to_string(), false);
|
|
assert!(fast_legacy_check(&cache_incompatible, &names).is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn test_should_track_statement_process() {
|
|
assert!(should_track_statement_process(&parse_one_sql(
|
|
"SELECT * FROM demo"
|
|
)));
|
|
assert!(should_track_statement_process(&parse_one_sql(
|
|
"INSERT INTO demo SELECT * FROM source"
|
|
)));
|
|
assert!(!should_track_statement_process(&parse_one_sql(
|
|
"INSERT INTO demo VALUES (1)"
|
|
)));
|
|
assert!(!should_track_statement_process(&parse_one_sql(
|
|
"INSERT INTO demo VALUES (now())"
|
|
)));
|
|
}
|
|
|
|
#[test]
|
|
fn test_should_track_plan_process() {
|
|
let select_stmt = parse_one_sql("SELECT * FROM demo");
|
|
let insert_select_stmt = parse_one_sql("INSERT INTO demo SELECT * FROM source");
|
|
let insert_values_stmt = parse_one_sql("INSERT INTO demo VALUES (now())");
|
|
|
|
let empty_plan = LogicalPlanBuilder::empty(false).build().unwrap();
|
|
assert!(should_track_plan_process(Some(&select_stmt), &empty_plan));
|
|
assert!(should_track_plan_process(
|
|
Some(&insert_select_stmt),
|
|
&insert_dml_plan()
|
|
));
|
|
assert!(!should_track_plan_process(
|
|
Some(&insert_values_stmt),
|
|
&insert_dml_plan()
|
|
));
|
|
assert!(!should_track_plan_process(None, &insert_dml_plan()));
|
|
}
|
|
|
|
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
|
async fn test_insert_select_is_visible_in_show_processlist() -> TestResult<()> {
|
|
let insert_sql = "INSERT INTO target SELECT * FROM source";
|
|
let (started_tx, mut started_rx) = mpsc::unbounded_channel();
|
|
let (finish_tx, finish_rx) = oneshot::channel();
|
|
let interceptor = Arc::new(BlockingInsertSelectInterceptor::new(started_tx, finish_rx));
|
|
let instance = Arc::new(test_instance_with_insert_select_interceptor(interceptor).await?);
|
|
|
|
let insert_task = tokio::spawn({
|
|
let instance = instance.clone();
|
|
async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
|
|
});
|
|
|
|
tokio::time::timeout(Duration::from_secs(5), started_rx.recv())
|
|
.await
|
|
.context(InsertStartTimeoutSnafu)?
|
|
.context(InsertStartChannelClosedSnafu)?;
|
|
|
|
let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
|
|
let process_list = output.data.pretty_print().await;
|
|
assert!(
|
|
process_list.contains(insert_sql),
|
|
"process list did not contain running insert:\n{process_list}"
|
|
);
|
|
|
|
finish_tx
|
|
.send(())
|
|
.map_err(|_| ReleaseBlockedInsertSnafu.build())?;
|
|
insert_task.await.context(InsertTaskPanicSnafu)??;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_show_processlist_catalog_scope() -> TestResult<()> {
|
|
let instance =
|
|
test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
|
|
.await?;
|
|
let _current_catalog = instance.process_manager().register_query(
|
|
"greptime".to_string(),
|
|
vec!["public".to_string()],
|
|
"current_catalog_query".to_string(),
|
|
String::new(),
|
|
None,
|
|
None,
|
|
);
|
|
let _other_catalog = instance.process_manager().register_query(
|
|
"other".to_string(),
|
|
vec!["public".to_string()],
|
|
"other_catalog_query".to_string(),
|
|
String::new(),
|
|
None,
|
|
None,
|
|
);
|
|
|
|
for sql in ["SHOW PROCESSLIST", "SHOW FULL PROCESSLIST"] {
|
|
let output = execute_one_sql(&instance, sql, test_query_ctx(43)).await?;
|
|
let process_list = output.data.pretty_print().await;
|
|
assert!(
|
|
process_list.contains("current_catalog_query"),
|
|
"{process_list}"
|
|
);
|
|
assert!(
|
|
!process_list.contains("other_catalog_query"),
|
|
"{process_list}"
|
|
);
|
|
|
|
let admin_ctx = test_query_ctx(44);
|
|
admin_ctx.set_current_user(Arc::new(AdminUserInfo));
|
|
let output = execute_one_sql(&instance, sql, admin_ctx).await?;
|
|
let process_list = output.data.pretty_print().await;
|
|
assert!(
|
|
process_list.contains("current_catalog_query"),
|
|
"{process_list}"
|
|
);
|
|
assert!(
|
|
process_list.contains("other_catalog_query"),
|
|
"{process_list}"
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
|
async fn test_kill_query_cancels_insert_select() -> TestResult<()> {
|
|
assert_kill_cancels_insert_select("KILL QUERY 4242").await
|
|
}
|
|
|
|
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
|
async fn test_kill_process_id_cancels_insert_select() -> TestResult<()> {
|
|
assert_kill_cancels_insert_select("KILL 'test-frontend/4242'").await
|
|
}
|
|
|
|
async fn assert_kill_cancels_insert_select(kill_sql: &str) -> TestResult<()> {
|
|
let insert_sql = "INSERT INTO target SELECT * FROM source";
|
|
let (source_polled_tx, source_polled_rx) = oneshot::channel();
|
|
let instance = Arc::new(
|
|
test_instance_with_tables(
|
|
pending_table(1024, "source", source_polled_tx)?,
|
|
test_table(1025, "target")?,
|
|
)
|
|
.await?,
|
|
);
|
|
|
|
let insert_task = tokio::spawn({
|
|
let instance = instance.clone();
|
|
async move { execute_one_sql(&instance, insert_sql, test_query_ctx(4242)).await }
|
|
});
|
|
|
|
tokio::time::timeout(Duration::from_secs(5), source_polled_rx)
|
|
.await
|
|
.context(SourcePollTimeoutSnafu)?
|
|
.context(SourcePollChannelClosedSnafu)?;
|
|
|
|
let output = execute_one_sql(&instance, kill_sql, test_query_ctx(43)).await?;
|
|
assert!(matches!(output.data, OutputData::AffectedRows(1)));
|
|
|
|
let insert_result = tokio::time::timeout(Duration::from_secs(5), insert_task)
|
|
.await
|
|
.context(InsertTaskTimeoutSnafu)?
|
|
.context(InsertTaskPanicSnafu)?;
|
|
let err = match insert_result {
|
|
Ok(_) => return InsertSelectNotCancelledSnafu.fail(),
|
|
Err(TestError::ExecuteSql { source, .. }) => source,
|
|
Err(err) => return Err(err),
|
|
};
|
|
assert_eq!(StatusCode::Cancelled, err.status_code());
|
|
|
|
let output = execute_one_sql(&instance, "SHOW PROCESSLIST", test_query_ctx(43)).await?;
|
|
let process_list = output.data.pretty_print().await;
|
|
assert!(
|
|
!process_list.contains(insert_sql),
|
|
"process list still contains killed insert:\n{process_list}"
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn insert_dml_plan() -> LogicalPlan {
|
|
let schema = SchemaRef::new(Schema::new(vec![Field::new(
|
|
"value",
|
|
DataType::Int64,
|
|
true,
|
|
)]));
|
|
let target = Arc::new(LogicalTableSource::new(schema));
|
|
let input = LogicalPlanBuilder::empty(false).build().unwrap();
|
|
|
|
LogicalPlanBuilder::insert_into(input, "demo", target, InsertOp::Append)
|
|
.unwrap()
|
|
.build()
|
|
.unwrap()
|
|
}
|
|
|
|
#[test]
|
|
fn test_exec_validation() {
|
|
let query_ctx = QueryContext::arc();
|
|
let plugins: Plugins = Plugins::new();
|
|
plugins.insert(QueryOptions {
|
|
disallow_cross_catalog_query: true,
|
|
});
|
|
|
|
let sql = r#"
|
|
SELECT * FROM demo;
|
|
EXPLAIN SELECT * FROM demo;
|
|
CREATE DATABASE test_database;
|
|
SHOW DATABASES;
|
|
"#;
|
|
let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
|
|
assert_eq!(stmts.len(), 4);
|
|
for stmt in stmts {
|
|
let re = check_permission(plugins.clone(), &stmt, &query_ctx);
|
|
re.unwrap();
|
|
}
|
|
|
|
let sql = r#"
|
|
SHOW CREATE TABLE demo;
|
|
ALTER TABLE demo ADD COLUMN new_col INT;
|
|
"#;
|
|
let stmts = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
|
|
assert_eq!(stmts.len(), 2);
|
|
for stmt in stmts {
|
|
let re = check_permission(plugins.clone(), &stmt, &query_ctx);
|
|
re.unwrap();
|
|
}
|
|
|
|
fn replace_test(template_sql: &str, plugins: Plugins, query_ctx: &QueryContextRef) {
|
|
// test right
|
|
let right = vec![("", ""), ("", "public."), ("greptime.", "public.")];
|
|
for (catalog, schema) in right {
|
|
let sql = do_fmt(template_sql, catalog, schema);
|
|
do_test(&sql, plugins.clone(), query_ctx, true);
|
|
}
|
|
|
|
let wrong = vec![
|
|
("wrongcatalog.", "public."),
|
|
("wrongcatalog.", "wrongschema."),
|
|
];
|
|
for (catalog, schema) in wrong {
|
|
let sql = do_fmt(template_sql, catalog, schema);
|
|
do_test(&sql, plugins.clone(), query_ctx, false);
|
|
}
|
|
}
|
|
|
|
fn do_fmt(template: &str, catalog: &str, schema: &str) -> String {
|
|
let vars = HashMap::from([
|
|
("catalog".to_string(), catalog),
|
|
("schema".to_string(), schema),
|
|
]);
|
|
template.format(&vars).unwrap()
|
|
}
|
|
|
|
fn do_test(sql: &str, plugins: Plugins, query_ctx: &QueryContextRef, is_ok: bool) {
|
|
let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
|
|
let re = check_permission(plugins, stmt, query_ctx);
|
|
if is_ok {
|
|
re.unwrap();
|
|
} else {
|
|
assert!(re.is_err());
|
|
}
|
|
}
|
|
|
|
// test insert
|
|
let sql = "INSERT INTO {catalog}{schema}monitor(host) VALUES ('host1');";
|
|
replace_test(sql, plugins.clone(), &query_ctx);
|
|
|
|
// test create table
|
|
let sql = r#"CREATE TABLE {catalog}{schema}demo(
|
|
host STRING,
|
|
ts TIMESTAMP,
|
|
TIME INDEX (ts),
|
|
PRIMARY KEY(host)
|
|
) engine=mito;"#;
|
|
replace_test(sql, plugins.clone(), &query_ctx);
|
|
|
|
// test drop table
|
|
let sql = "DROP TABLE {catalog}{schema}demo;";
|
|
replace_test(sql, plugins.clone(), &query_ctx);
|
|
|
|
// test undrop table
|
|
#[cfg(feature = "enterprise")]
|
|
{
|
|
let sql = "UNDROP TABLE {catalog}{schema}demo;";
|
|
replace_test(sql, plugins.clone(), &query_ctx);
|
|
}
|
|
|
|
// test show tables
|
|
let sql = "SHOW TABLES FROM public";
|
|
let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
|
|
check_permission(plugins.clone(), &stmt[0], &query_ctx).unwrap();
|
|
|
|
let sql = "SHOW TABLES FROM private";
|
|
let stmt = parse_stmt(sql, &GreptimeDbDialect {}).unwrap();
|
|
let re = check_permission(plugins.clone(), &stmt[0], &query_ctx);
|
|
assert!(re.is_ok());
|
|
|
|
// test describe table
|
|
let sql = "DESC TABLE {catalog}{schema}demo;";
|
|
replace_test(sql, plugins.clone(), &query_ctx);
|
|
|
|
let comment_flow_cases = [
|
|
("COMMENT ON FLOW my_flow IS 'comment';", true),
|
|
("COMMENT ON FLOW greptime.my_flow IS 'comment';", true),
|
|
("COMMENT ON FLOW wrongcatalog.my_flow IS 'comment';", false),
|
|
];
|
|
for (sql, is_ok) in comment_flow_cases {
|
|
let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
|
|
let result = check_permission(plugins.clone(), stmt, &query_ctx);
|
|
assert_eq!(result.is_ok(), is_ok);
|
|
}
|
|
|
|
let show_flow_cases = [
|
|
("SHOW CREATE FLOW my_flow;", true),
|
|
("SHOW CREATE FLOW greptime.my_flow;", true),
|
|
("SHOW CREATE FLOW wrongcatalog.my_flow;", false),
|
|
];
|
|
for (sql, is_ok) in show_flow_cases {
|
|
let stmt = &parse_stmt(sql, &GreptimeDbDialect {}).unwrap()[0];
|
|
let result = check_permission(plugins.clone(), stmt, &query_ctx);
|
|
assert_eq!(result.is_ok(), is_ok);
|
|
}
|
|
}
|
|
|
|
/// A `DropView` DDL sent through the direct gRPC ingress must return an error
|
|
/// (e.g. table not found) instead of panicking on `todo!()`.
|
|
#[tokio::test]
|
|
async fn qx_152_drop_view_via_grpc_ddl_returns_error_not_panic() -> TestResult<()> {
|
|
let instance =
|
|
test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
|
|
.await?;
|
|
|
|
let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
|
|
expr: Some(api::v1::ddl_request::Expr::DropView(
|
|
api::v1::DropViewExpr {
|
|
catalog_name: String::new(),
|
|
schema_name: String::new(),
|
|
view_name: "non_existent_view".to_string(),
|
|
view_id: None,
|
|
drop_if_exists: false,
|
|
},
|
|
)),
|
|
});
|
|
|
|
let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
|
|
&instance,
|
|
request,
|
|
QueryContext::arc(),
|
|
)
|
|
.await;
|
|
|
|
let err = match result {
|
|
Ok(_) => panic!("DropView DDL request must return an error instead of panicking"),
|
|
Err(err) => err,
|
|
};
|
|
assert_eq!(
|
|
err.status_code(),
|
|
StatusCode::TableNotFound,
|
|
"dropping a non-existent view without IF EXISTS must report TableNotFound, got {err}"
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
/// `DROP VIEW IF EXISTS` on a missing view through the direct gRPC ingress must
|
|
/// succeed with 0 affected rows (no error, no DDL task submitted), instead of
|
|
/// returning `TableNotFound`.
|
|
#[tokio::test]
|
|
async fn qx_152_drop_view_if_exists_missing_view_via_grpc_ddl_succeeds() -> TestResult<()> {
|
|
let catalog_manager =
|
|
catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
|
|
let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
|
|
let instance = test_instance_with_catalog_manager(
|
|
catalog_manager,
|
|
test_table(1025, "target")?,
|
|
Plugins::new(),
|
|
None,
|
|
procedure_executor.clone() as ProcedureExecutorRef,
|
|
)
|
|
.await?;
|
|
|
|
let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
|
|
expr: Some(api::v1::ddl_request::Expr::DropView(
|
|
api::v1::DropViewExpr {
|
|
catalog_name: String::new(),
|
|
schema_name: String::new(),
|
|
view_name: "non_existent_view".to_string(),
|
|
view_id: None,
|
|
drop_if_exists: true,
|
|
},
|
|
)),
|
|
});
|
|
|
|
let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
|
|
&instance,
|
|
request,
|
|
QueryContext::arc(),
|
|
)
|
|
.await;
|
|
|
|
let output = match result {
|
|
Ok(output) => output,
|
|
Err(err) => {
|
|
panic!("DROP VIEW IF EXISTS on a missing view must succeed, got error: {err}")
|
|
}
|
|
};
|
|
assert!(
|
|
matches!(output.data, OutputData::AffectedRows(0)),
|
|
"DROP VIEW IF EXISTS on a missing view must report 0 affected rows"
|
|
);
|
|
assert!(
|
|
procedure_executor.submitted.lock().unwrap().is_empty(),
|
|
"DROP VIEW IF EXISTS on a missing view must not submit a DDL task"
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
/// A `CREATE VIEW` followed by `DROP VIEW` through the direct gRPC ingress must
|
|
/// succeed end to end: the view is registered in the catalog and then removed.
|
|
#[tokio::test]
|
|
async fn qx_152_drop_existing_view_via_grpc_ddl_succeeds() -> TestResult<()> {
|
|
let catalog_manager =
|
|
catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
|
|
let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
|
|
let instance = test_instance_with_catalog_manager(
|
|
catalog_manager,
|
|
test_table(1025, "target")?,
|
|
Plugins::new(),
|
|
None,
|
|
procedure_executor.clone() as ProcedureExecutorRef,
|
|
)
|
|
.await?;
|
|
|
|
// The default "greptime.public" schema must be visible to the kv-backed table
|
|
// metadata manager for `CREATE VIEW`/`CREATE TABLE` to pass validation.
|
|
instance
|
|
.table_metadata_manager()
|
|
.schema_manager()
|
|
.create(
|
|
common_meta::key::schema_name::SchemaNameKey::new("greptime", "public"),
|
|
None,
|
|
true,
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
let create_view_request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
|
|
expr: Some(api::v1::ddl_request::Expr::CreateView(
|
|
api::v1::CreateViewExpr {
|
|
catalog_name: String::new(),
|
|
schema_name: String::new(),
|
|
view_name: "my_view".to_string(),
|
|
logical_plan: vec![1, 2, 3],
|
|
create_if_not_exists: false,
|
|
or_replace: false,
|
|
table_names: vec![],
|
|
columns: vec![],
|
|
plan_columns: vec![],
|
|
definition: "CREATE VIEW my_view AS SELECT * FROM source".to_string(),
|
|
},
|
|
)),
|
|
});
|
|
|
|
let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
|
|
&instance,
|
|
create_view_request,
|
|
QueryContext::arc(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(output) => output,
|
|
Err(err) => panic!("CREATE VIEW via gRPC DDL must succeed, got error: {err}"),
|
|
};
|
|
assert!(
|
|
matches!(output.data, OutputData::AffectedRows(0)),
|
|
"CREATE VIEW via gRPC DDL must report 0 affected rows"
|
|
);
|
|
|
|
// The view is registered in the catalog as a view.
|
|
let view = instance
|
|
.catalog_manager()
|
|
.table("greptime", "public", "my_view", None)
|
|
.await
|
|
.unwrap()
|
|
.expect("view should exist after CREATE VIEW");
|
|
assert_eq!(view.table_info().table_type, TableType::View);
|
|
|
|
let drop_view_request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
|
|
expr: Some(api::v1::ddl_request::Expr::DropView(
|
|
api::v1::DropViewExpr {
|
|
catalog_name: String::new(),
|
|
schema_name: String::new(),
|
|
view_name: "my_view".to_string(),
|
|
view_id: None,
|
|
drop_if_exists: false,
|
|
},
|
|
)),
|
|
});
|
|
|
|
let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
|
|
&instance,
|
|
drop_view_request,
|
|
QueryContext::arc(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(output) => output,
|
|
Err(err) => panic!("DROP VIEW via gRPC DDL must succeed, got error: {err}"),
|
|
};
|
|
assert!(
|
|
matches!(output.data, OutputData::AffectedRows(0)),
|
|
"DROP VIEW via gRPC DDL must report 0 affected rows"
|
|
);
|
|
|
|
// The view is gone after the drop.
|
|
assert!(
|
|
instance
|
|
.catalog_manager()
|
|
.table("greptime", "public", "my_view", None)
|
|
.await
|
|
.unwrap()
|
|
.is_none(),
|
|
"view should be removed after DROP VIEW"
|
|
);
|
|
|
|
let submitted = procedure_executor.submitted.lock().unwrap();
|
|
assert_eq!(
|
|
submitted.len(),
|
|
2,
|
|
"expected create and drop view tasks, got {submitted:?}"
|
|
);
|
|
assert!(matches!(&submitted[0], DdlTask::CreateView(_)));
|
|
assert!(matches!(&submitted[1], DdlTask::DropView(_)));
|
|
Ok(())
|
|
}
|
|
|
|
/// A direct gRPC `CreateTable` whose time index column is not a timestamp must
|
|
/// be rejected with `InvalidArguments` instead of panicking while building the schema.
|
|
#[tokio::test]
|
|
async fn qx_153_create_table_with_non_timestamp_time_index_via_grpc_returns_error()
|
|
-> TestResult<()> {
|
|
let instance =
|
|
test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?)
|
|
.await?;
|
|
|
|
let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
|
|
expr: Some(api::v1::ddl_request::Expr::CreateTable(
|
|
api::v1::CreateTableExpr {
|
|
catalog_name: String::new(),
|
|
schema_name: String::new(),
|
|
table_name: "demo".to_string(),
|
|
desc: String::new(),
|
|
column_defs: vec![api::v1::ColumnDef {
|
|
name: "host".to_string(),
|
|
data_type: api::v1::ColumnDataType::String as i32,
|
|
is_nullable: true,
|
|
default_constraint: vec![],
|
|
semantic_type: 0,
|
|
comment: String::new(),
|
|
datatype_extension: None,
|
|
options: None,
|
|
}],
|
|
time_index: "host".to_string(),
|
|
primary_keys: vec![],
|
|
create_if_not_exists: false,
|
|
table_options: HashMap::new(),
|
|
table_id: None,
|
|
engine: "mito".to_string(),
|
|
},
|
|
)),
|
|
});
|
|
|
|
let result = servers::query_handler::grpc::GrpcQueryHandler::do_query(
|
|
&instance,
|
|
request,
|
|
QueryContext::arc(),
|
|
)
|
|
.await;
|
|
|
|
let err = match result {
|
|
Ok(_) => panic!("CreateTable with a non-timestamp time index must be rejected"),
|
|
Err(err) => err,
|
|
};
|
|
assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}");
|
|
Ok(())
|
|
}
|
|
|
|
/// A valid `CREATE TABLE` (timestamp time index) through the direct gRPC ingress
|
|
/// must succeed, guarding that the `validate_create_expr` ingress check doesn't
|
|
/// accidentally reject good requests.
|
|
#[tokio::test]
|
|
async fn qx_153_create_table_with_timestamp_time_index_via_grpc_succeeds() -> TestResult<()> {
|
|
let catalog_manager =
|
|
catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?);
|
|
let procedure_executor = Arc::new(MockProcedureExecutor::new(catalog_manager.clone()));
|
|
let instance = test_instance_with_catalog_manager(
|
|
catalog_manager,
|
|
test_table(1025, "target")?,
|
|
Plugins::new(),
|
|
None,
|
|
procedure_executor.clone() as ProcedureExecutorRef,
|
|
)
|
|
.await?;
|
|
|
|
// The default "greptime.public" schema must be visible to the kv-backed table
|
|
// metadata manager for `CREATE TABLE` to pass validation.
|
|
instance
|
|
.table_metadata_manager()
|
|
.schema_manager()
|
|
.create(
|
|
common_meta::key::schema_name::SchemaNameKey::new("greptime", "public"),
|
|
None,
|
|
true,
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest {
|
|
expr: Some(api::v1::ddl_request::Expr::CreateTable(
|
|
api::v1::CreateTableExpr {
|
|
catalog_name: String::new(),
|
|
schema_name: String::new(),
|
|
table_name: "demo".to_string(),
|
|
desc: String::new(),
|
|
column_defs: vec![
|
|
api::v1::ColumnDef {
|
|
name: "host".to_string(),
|
|
data_type: api::v1::ColumnDataType::String as i32,
|
|
is_nullable: true,
|
|
default_constraint: vec![],
|
|
semantic_type: 0,
|
|
comment: String::new(),
|
|
datatype_extension: None,
|
|
options: None,
|
|
},
|
|
api::v1::ColumnDef {
|
|
name: "ts".to_string(),
|
|
data_type: api::v1::ColumnDataType::TimestampMillisecond as i32,
|
|
is_nullable: true,
|
|
default_constraint: vec![],
|
|
semantic_type: 0,
|
|
comment: String::new(),
|
|
datatype_extension: None,
|
|
options: None,
|
|
},
|
|
],
|
|
time_index: "ts".to_string(),
|
|
primary_keys: vec![],
|
|
create_if_not_exists: false,
|
|
table_options: HashMap::new(),
|
|
table_id: None,
|
|
engine: "mito".to_string(),
|
|
},
|
|
)),
|
|
});
|
|
|
|
let output = match servers::query_handler::grpc::GrpcQueryHandler::do_query(
|
|
&instance,
|
|
request,
|
|
QueryContext::arc(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(output) => output,
|
|
Err(err) => panic!("CREATE TABLE via gRPC DDL must succeed, got error: {err}"),
|
|
};
|
|
assert!(
|
|
matches!(output.data, OutputData::AffectedRows(0)),
|
|
"CREATE TABLE via gRPC DDL must report 0 affected rows"
|
|
);
|
|
|
|
// The table is registered in the catalog.
|
|
let table = instance
|
|
.catalog_manager()
|
|
.table("greptime", "public", "demo", None)
|
|
.await
|
|
.unwrap()
|
|
.expect("table should exist after CREATE TABLE");
|
|
assert_eq!(table.table_info().table_type, TableType::Base);
|
|
|
|
let submitted = procedure_executor.submitted.lock().unwrap();
|
|
assert_eq!(
|
|
submitted.len(),
|
|
1,
|
|
"expected one create table task, got {submitted:?}"
|
|
);
|
|
assert!(matches!(&submitted[0], DdlTask::CreateTable(_)));
|
|
Ok(())
|
|
}
|
|
}
|