rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-08-18 12:08:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: close database ACL gaps in permission checks (#8492)

Browse Source 
* fix: database ACL

Signed-off-by: shuiyisong <xixing.sys@gmail.com>

* chore: reduce code

Signed-off-by: shuiyisong <xixing.sys@gmail.com>

* fix: cr issue

Signed-off-by: shuiyisong <xixing.sys@gmail.com>

---------

Signed-off-by: shuiyisong <xixing.sys@gmail.com>
This commit is contained in:
shuiyisong
2026-07-14 22:16:08 +08:00
committed by GitHub
parent b30765e5e9
commit 589ae56a6f
8 changed files with 473 additions and 212 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/auth/src/permission.rs
+42 -7
View File
@@ -36,7 +36,11 @@ pub enum PermissionReq<'a> {
PromStoreRead,
Otlp,
LogWrite,
BulkInsert,
BulkInsert {
catalog: &'a str,
schema: &'a str,
table: &'a str,
},
}
impl<'a> PermissionReq<'a> {
@@ -57,7 +61,7 @@ impl<'a> PermissionReq<'a> {
| PermissionReq::PromStoreWrite
| PermissionReq::Otlp
| PermissionReq::LogWrite
| PermissionReq::BulkInsert => false,
| PermissionReq::BulkInsert { .. } => false,
}
}
@@ -79,6 +83,15 @@ pub trait PermissionChecker: Send + Sync {
user_info: UserInfoRef,
req: PermissionReq,
) -> Result<PermissionResp>;
fn check_permission_with_context(
&self,
user_info: UserInfoRef,
req: PermissionReq,
_current_schema: Option<&str>,
) -> Result<PermissionResp> {
self.check_permission(user_info, req)
}
}
impl PermissionChecker for Option<&PermissionCheckerRef> {
@@ -86,13 +99,24 @@ impl PermissionChecker for Option<&PermissionCheckerRef> {
&self,
user_info: UserInfoRef,
req: PermissionReq,
) -> Result<PermissionResp> {
self.check_permission_with_context(user_info, req, None)
}
fn check_permission_with_context(
&self,
user_info: UserInfoRef,
req: PermissionReq,
current_schema: Option<&str>,
) -> Result<PermissionResp> {
match self {
Some(checker) => match checker.check_permission(user_info, req) {
Ok(PermissionResp::Reject) => PermissionDeniedSnafu.fail(),
Ok(PermissionResp::Allow) => Ok(PermissionResp::Allow),
Err(e) => Err(e),
},
Some(checker) => {
match checker.check_permission_with_context(user_info, req, current_schema) {
Ok(PermissionResp::Reject) => PermissionDeniedSnafu.fail(),
Ok(PermissionResp::Allow) => Ok(PermissionResp::Allow),
Err(e) => Err(e),
}
}
None => Ok(PermissionResp::Allow),
}
}
@@ -209,4 +233,15 @@ mod tests {
assert!(req.is_write());
}
#[test]
fn test_bulk_insert_is_write_request() {
let req = PermissionReq::BulkInsert {
catalog: "greptime",
schema: "public",
table: "metrics",
};
assert!(req.is_write());
}
}
src/frontend/src/instance.rs
+12 -3
View File
@@ -611,7 +611,11 @@ impl Instance {
let checker_ref = self.plugins.get::<PermissionCheckerRef>();
checker_ref
.as_ref()
.check_permission(query_ctx.current_user(), PermissionReq::SqlStatement(&stmt))
.check_permission_with_context(
query_ctx.current_user(),
PermissionReq::SqlStatement(&stmt),
Some(&query_ctx.current_schema()),
)
.context(PermissionSnafu)?;
check_permission(self.plugins.clone(), &stmt, &query_ctx)?;
let catalog_name = query_ctx.current_catalog().to_string();
@@ -672,9 +676,10 @@ impl Instance {
let mut results = Vec::with_capacity(stmts.len());
for stmt in stmts {
if let Err(e) = checker
.check_permission(
.check_permission_with_context(
query_ctx.current_user(),
PermissionReq::SqlStatement(&stmt),
Some(&query_ctx.current_schema()),
)
.context(PermissionSnafu)
{
@@ -848,7 +853,11 @@ impl Instance {
self.plugins
.get::<PermissionCheckerRef>()
.as_ref()
.check_permission(query_ctx.current_user(), PermissionReq::SqlStatement(&stmt))
.check_permission_with_context(
query_ctx.current_user(),
PermissionReq::SqlStatement(&stmt),
Some(&query_ctx.current_schema()),
)
.context(PermissionSnafu)?;
let plan = self
src/frontend/src/instance/grpc.rs
+13 -2
View File
@@ -71,7 +71,11 @@ impl GrpcQueryHandler for Instance {
self.plugins
.get::<PermissionCheckerRef>()
.as_ref()
.check_permission(ctx.current_user(), PermissionReq::GrpcRequest(&request))
.check_permission_with_context(
ctx.current_user(),
PermissionReq::GrpcRequest(&request),
Some(&ctx.current_schema()),
)
.context(PermissionSnafu)?;
let output = match request {
@@ -356,7 +360,14 @@ impl Instance {
plugins
.get::<PermissionCheckerRef>()
.as_ref()
.check_permission(ctx.current_user(), PermissionReq::BulkInsert)
.check_permission(
ctx.current_user(),
PermissionReq::BulkInsert {
catalog: &table_name.catalog_name,
schema: &table_name.schema_name,
table: &table_name.table_name,
},
)
.context(PermissionSnafu)?;
// Resolve table reference
src/operator/src/statement.rs
+5 -3
View File
@@ -78,8 +78,7 @@ use table::table_name::TableName;
use table::table_reference::TableReference;
use self::set::{
set_bytea_output, set_datestyle, set_intervalstyle, set_search_path, set_timezone,
validate_client_encoding,
set_bytea_output, set_datestyle, set_intervalstyle, set_timezone, validate_client_encoding,
};
use crate::error::{
self, CatalogSnafu, ExecLogicalPlanSnafu, ExternalSnafu, InvalidSqlSnafu, NotSupportedSnafu,
@@ -528,7 +527,10 @@ impl StatementExecutor {
},
"SEARCH_PATH" => {
if query_ctx.channel() == Channel::Postgres {
set_search_path(set_var.value, query_ctx)?
let search_path = set_var.search_path().context(NotSupportedSnafu {
feat: "Unsupported search path in set variable statement",
})?;
query_ctx.set_current_schema(search_path);
} else {
return NotSupportedSnafu {
feat: format!("Unsupported set variable {}", var_name),
src/operator/src/statement/set.rs
-30
View File
@@ -129,36 +129,6 @@ pub fn set_bytea_output(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
Ok(())
}
pub fn set_search_path(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
let search_expr = exprs.first().context(NotSupportedSnafu {
feat: "No search path find in set variable statement",
})?;
match search_expr {
Expr::Value(ValueWithSpan {
value: Value::SingleQuotedString(search_path),
..
})
| Expr::Value(ValueWithSpan {
value: Value::DoubleQuotedString(search_path),
..
}) => {
ctx.set_current_schema(search_path);
Ok(())
}
Expr::Identifier(Ident { value, .. }) => {
ctx.set_current_schema(value);
Ok(())
}
expr => NotSupportedSnafu {
feat: format!(
"Unsupported search path expr {} in set variable statement",
expr
),
}
.fail(),
}
}
pub fn validate_client_encoding(set: SetVariables) -> Result<()> {
let Some((encoding, [])) = set.value.split_first() else {
return InvalidSqlSnafu {
src/sql/src/statements/set_variables.rs
+44 -1
View File
@@ -15,7 +15,7 @@
use std::fmt::Display;
use serde::Serialize;
use sqlparser::ast::{Expr, ObjectName};
use sqlparser::ast::{Expr, ObjectName, ObjectNamePart, Value, ValueWithSpan};
use sqlparser_derive::{Visit, VisitMut};
/// SET variables statement.
@@ -25,6 +25,27 @@ pub struct SetVariables {
pub value: Vec<Expr>,
}
impl SetVariables {
/// Returns the first supported `search_path` value.
pub fn search_path(&self) -> Option<&str> {
let [ObjectNamePart::Identifier(variable)] = self.variable.0.as_slice() else {
return None;
};
if variable.quote_style.is_some() || !variable.value.eq_ignore_ascii_case("search_path") {
return None;
}
match self.value.first()? {
Expr::Value(ValueWithSpan {
value: Value::SingleQuotedString(value) | Value::DoubleQuotedString(value),
..
}) => Some(value),
Expr::Identifier(value) => Some(&value.value),
_ => None,
}
}
}
impl Display for SetVariables {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let variable = &self.variable;
@@ -70,4 +91,26 @@ SET delayed_insert_timeout = 300"#,
}
}
}
#[test]
fn test_search_path() {
for (sql, expected) in [
("SET search_path TO public", Some("public")),
("SET SEARCH_PATH TO 'private'", Some("private")),
("SET search_path TO \"quoted\"", Some("quoted")),
("SET \"search_path\" TO private", None),
("SET timezone TO 'UTC'", None),
] {
let statements = ParserContext::create_with_dialect(
sql,
&GreptimeDbDialect {},
ParseOptions::default(),
)
.unwrap();
let [Statement::SetVariables(set)] = statements.as_slice() else {
unreachable!()
};
assert_eq!(expected, set.search_path());
}
}
}
src/sql/src/statements/statement.rs
+2 -1
View File
@@ -178,7 +178,6 @@ impl Statement {
| Statement::ShowSearchPath(_)
| Statement::ShowViews(_)
| Statement::DescribeTable(_)
| Statement::Explain(_)
| Statement::ShowVariables(_)
| Statement::ShowProcesslist(_)
| Statement::FetchCursor(_)
@@ -189,6 +188,8 @@ impl Statement {
#[cfg(feature = "enterprise")]
Statement::ShowTriggers(_) => true,
Statement::Explain(explain) => !explain.analyze || explain.statement.is_readonly(),
// Write operations
Statement::Insert(_)
| Statement::Delete(_)
src/sql/src/util.rs
+355 -165
View File
@@ -14,20 +14,19 @@
use std::collections::HashSet;
use std::fmt::{Display, Formatter};
use std::ops::ControlFlow;
use itertools::Itertools;
use promql_parser::label::{METRIC_NAME, MatchOp};
use promql_parser::parser::{
AggregateExpr as PromAggregateExpr, BinaryExpr as PromBinaryExpr, Call as PromCall,
Expr as PromExpr, MatrixSelector as PromMatrixSelector, ParenExpr as PromParenExpr,
SubqueryExpr as PromSubqueryExpr, UnaryExpr as PromUnaryExpr,
VectorSelector as PromVectorSelector,
Expr as PromExpr, MatrixSelector as PromMatrixSelector, VectorSelector as PromVectorSelector,
};
use promql_parser::util::{ExprVisitor, walk_expr};
use serde::Serialize;
use snafu::ensure;
use sqlparser::ast::{
Array, Expr, Ident, ObjectName, ObjectNamePart, SetExpr, SqlOption, TableFactor,
TableWithJoins, Value, ValueWithSpan,
Array, Expr, Ident, ObjectName, ObjectNamePart, SqlOption, Value, ValueWithSpan,
Visit as AstVisit, Visitor,
};
use sqlparser_derive::{Visit, VisitMut};
@@ -184,48 +183,60 @@ pub fn parse_option_string(option: SqlOption) -> Result<(String, OptionValue)> {
/// Walk through a [Query] and extract all the tables referenced in it.
pub fn extract_tables_from_query(query: &SqlOrTql) -> impl Iterator<Item = ObjectName> {
let mut names = HashSet::new();
match query {
SqlOrTql::Sql(query, _) => {
extract_tables_from_sql_query(&query.inner, &mut names);
extract_tables_from_hybrid_cte_query(query, &mut names);
}
SqlOrTql::Tql(tql, _) => extract_tables_from_tql(tql, &mut names),
}
names.into_iter()
extract_tables_from_query_inner(query).0.into_iter()
}
fn extract_tables_from_hybrid_cte_query(query: &Query, sql_names: &mut HashSet<ObjectName>) {
if let Some(hybrid_cte) = &query.hybrid_cte {
let mut cte_names: HashSet<String> = hybrid_cte
.cte_tables
.iter()
.map(|cte| ParserContext::canonicalize_identifier(cte.name.clone()).value)
.collect();
remove_cte_names(sql_names, &cte_names);
/// Walk through a [Query] and extract its referenced tables, returning `None`
/// if any table reference cannot be resolved statically.
pub fn extract_tables_from_query_checked(
query: &SqlOrTql,
) -> Option<impl Iterator<Item = ObjectName>> {
let (names, complete) = extract_tables_from_query_inner(query);
complete.then_some(names.into_iter())
}
cte_names.clear();
for cte in &hybrid_cte.cte_tables {
let cte_name = ParserContext::canonicalize_identifier(cte.name.clone()).value;
let mut cte_query_names = HashSet::new();
match &cte.content {
CteContent::Sql(cte_query) => {
extract_tables_from_sql_query(cte_query, &mut cte_query_names)
}
CteContent::Tql(tql) => extract_tables_from_tql(tql, &mut cte_query_names),
}
if hybrid_cte.recursive {
cte_names.insert(cte_name.clone());
}
remove_cte_names(&mut cte_query_names, &cte_names);
sql_names.extend(cte_query_names);
if !hybrid_cte.recursive {
cte_names.insert(cte_name);
}
fn extract_tables_from_query_inner(query: &SqlOrTql) -> (HashSet<ObjectName>, bool) {
let mut names = HashSet::new();
let complete = match query {
SqlOrTql::Sql(query, _) => {
extract_tables_from_sql_query(&query.inner, &mut names);
extract_tables_from_hybrid_cte_query(query, &mut names)
}
SqlOrTql::Tql(tql, _) => extract_tables_from_tql(tql, &mut names),
};
(names, complete)
}
fn extract_tables_from_hybrid_cte_query(
query: &Query,
sql_names: &mut HashSet<ObjectName>,
) -> bool {
let Some(hybrid_cte) = &query.hybrid_cte else {
return true;
};
let mut complete = true;
let cte_names: HashSet<String> = hybrid_cte
.cte_tables
.iter()
.map(|cte| ParserContext::canonicalize_identifier(cte.name.clone()).value)
.collect();
remove_cte_names(sql_names, &cte_names);
for cte in &hybrid_cte.cte_tables {
let mut cte_query_names = HashSet::new();
match &cte.content {
CteContent::Sql(cte_query) => {
extract_tables_from_sql_query(cte_query, &mut cte_query_names)
}
CteContent::Tql(tql) => complete &= extract_tables_from_tql(tql, &mut cte_query_names),
}
sql_names.extend(cte_query_names);
}
complete
}
fn remove_cte_names(names: &mut HashSet<ObjectName>, cte_names: &HashSet<String>) {
@@ -246,55 +257,55 @@ fn remove_cte_names(names: &mut HashSet<ObjectName>, cte_names: &HashSet<String>
});
}
fn extract_tables_from_tql(tql: &Tql, names: &mut HashSet<ObjectName>) {
fn extract_tables_from_tql(tql: &Tql, names: &mut HashSet<ObjectName>) -> bool {
let promql = match tql {
Tql::Eval(eval) => &eval.query,
Tql::Explain(explain) => &explain.query,
Tql::Analyze(analyze) => &analyze.query,
};
if let Ok(expr) = promql_parser::parser::parse(promql) {
extract_tables_from_prom_expr(&expr, names);
}
let Ok(expr) = promql_parser::parser::parse(promql) else {
return false;
};
extract_tables_from_prom_expr(&expr, names)
}
fn extract_tables_from_prom_expr(expr: &PromExpr, names: &mut HashSet<ObjectName>) {
match expr {
PromExpr::Aggregate(PromAggregateExpr { expr, .. }) => {
extract_tables_from_prom_expr(expr, names);
}
PromExpr::Unary(PromUnaryExpr { expr, .. }) => {
extract_tables_from_prom_expr(expr, names);
}
PromExpr::Binary(PromBinaryExpr { lhs, rhs, .. }) => {
extract_tables_from_prom_expr(lhs, names);
extract_tables_from_prom_expr(rhs, names);
}
PromExpr::Paren(PromParenExpr { expr }) => {
extract_tables_from_prom_expr(expr, names);
}
PromExpr::Subquery(PromSubqueryExpr { expr, .. }) => {
extract_tables_from_prom_expr(expr, names);
}
PromExpr::VectorSelector(selector) => {
extract_metric_name_from_vector_selector(selector, names);
}
PromExpr::MatrixSelector(PromMatrixSelector { vs, .. }) => {
extract_metric_name_from_vector_selector(vs, names);
}
PromExpr::Call(PromCall { args, .. }) => {
for arg in &args.args {
extract_tables_from_prom_expr(arg, names);
}
}
PromExpr::NumberLiteral(_) | PromExpr::StringLiteral(_) | PromExpr::Extension(_) => {}
fn extract_tables_from_prom_expr(expr: &PromExpr, names: &mut HashSet<ObjectName>) -> bool {
struct TableCollector<'a> {
names: &'a mut HashSet<ObjectName>,
complete: bool,
}
impl ExprVisitor for TableCollector<'_> {
type Error = ();
fn pre_visit(&mut self, expr: &PromExpr) -> std::result::Result<bool, Self::Error> {
self.complete &= match expr {
PromExpr::VectorSelector(selector) => {
extract_metric_name_from_vector_selector(selector, self.names)
}
PromExpr::MatrixSelector(PromMatrixSelector { vs, .. }) => {
extract_metric_name_from_vector_selector(vs, self.names)
}
PromExpr::Extension(_) => false,
_ => true,
};
Ok(true)
}
}
let mut collector = TableCollector {
names,
complete: true,
};
let _ = walk_expr(&mut collector, expr);
collector.complete
}
fn extract_metric_name_from_vector_selector(
selector: &PromVectorSelector,
names: &mut HashSet<ObjectName>,
) {
) -> bool {
let metric_name = selector.name.clone().or_else(|| {
let mut metric_name_matchers = selector.matchers.find_matchers(METRIC_NAME);
if metric_name_matchers.len() == 1 && metric_name_matchers[0].op == MatchOp::Equal {
@@ -304,9 +315,16 @@ fn extract_metric_name_from_vector_selector(
}
});
let Some(metric_name) = metric_name else {
return;
return false;
};
if selector.matchers.matchers.iter().any(|matcher| {
(matcher.name == SCHEMA_MATCHER || matcher.name == DATABASE_MATCHER)
&& matcher.op != MatchOp::Equal
}) {
return false;
}
let schema_matcher = selector.matchers.matchers.iter().rev().find(|matcher| {
matcher.op == MatchOp::Equal
&& (matcher.name == SCHEMA_MATCHER || matcher.name == DATABASE_MATCHER)
@@ -322,6 +340,7 @@ fn extract_metric_name_from_vector_selector(
metric_name,
))]));
}
true
}
/// translate the start location to the index in the sql string
@@ -340,96 +359,116 @@ pub fn location_to_index(sql: &str, location: &sqlparser::tokenizer::Location) -
index - 1
}
// Tracks CTE aliases as sqlparser's visitor enters their query bodies.
struct QueryScope {
cte_names: Vec<String>,
next_cte: usize,
recursive: bool,
scope_start: usize,
add_to_parent_after: Option<String>,
}
struct RelationCollector<'a> {
names: &'a mut HashSet<ObjectName>,
ctes_in_scope: Vec<String>,
query_scopes: Vec<QueryScope>,
#[cfg(test)]
query_visits: usize,
}
impl<'a> RelationCollector<'a> {
fn new(names: &'a mut HashSet<ObjectName>) -> Self {
Self {
names,
ctes_in_scope: Vec::new(),
query_scopes: Vec::new(),
#[cfg(test)]
query_visits: 0,
}
}
}
impl Visitor for RelationCollector<'_> {
type Break = ();
fn pre_visit_query(&mut self, query: &sqlparser::ast::Query) -> ControlFlow<Self::Break> {
#[cfg(test)]
{
self.query_visits += 1;
}
// Query::visit enters WITH definitions in declaration order before its body.
let parent_cte = self.query_scopes.last_mut().and_then(|scope| {
let cte_name = scope.cte_names.get(scope.next_cte)?.clone();
scope.next_cte += 1;
Some((cte_name, scope.recursive))
});
let add_to_parent_after = match parent_cte {
Some((cte_name, true)) => {
self.ctes_in_scope.push(cte_name);
None
}
Some((cte_name, false)) => Some(cte_name),
None => None,
};
let scope_start = self.ctes_in_scope.len();
let (cte_names, recursive) = if let Some(with) = &query.with {
(
with.cte_tables
.iter()
.map(|cte| ParserContext::canonicalize_identifier(cte.alias.name.clone()).value)
.collect(),
with.recursive,
)
} else {
(Vec::new(), false)
};
self.query_scopes.push(QueryScope {
cte_names,
next_cte: 0,
recursive,
scope_start,
add_to_parent_after,
});
ControlFlow::Continue(())
}
fn post_visit_query(&mut self, _query: &sqlparser::ast::Query) -> ControlFlow<Self::Break> {
let Some(scope) = self.query_scopes.pop() else {
return ControlFlow::Break(());
};
self.ctes_in_scope.truncate(scope.scope_start);
if let Some(cte_name) = scope.add_to_parent_after {
self.ctes_in_scope.push(cte_name);
}
ControlFlow::Continue(())
}
fn pre_visit_relation(&mut self, relation: &ObjectName) -> ControlFlow<Self::Break> {
let is_cte = matches!(
relation.0.as_slice(),
[part]
if part.as_ident().is_some_and(|ident| {
self.ctes_in_scope.contains(
&ParserContext::canonicalize_identifier(ident.clone()).value,
)
})
);
if !is_cte {
self.names.insert(relation.clone());
}
ControlFlow::Continue(())
}
}
/// Helper function for [extract_tables_from_query].
///
/// Handle [sqlparser::ast::Query].
fn extract_tables_from_sql_query(query: &sqlparser::ast::Query, names: &mut HashSet<ObjectName>) {
let mut cte_names = HashSet::new();
if let Some(with) = &query.with {
for cte in &with.cte_tables {
let cte_name = ParserContext::canonicalize_identifier(cte.alias.name.clone()).value;
let mut cte_query_names = HashSet::new();
extract_tables_from_sql_query(&cte.query, &mut cte_query_names);
if with.recursive {
cte_names.insert(cte_name.clone());
}
remove_cte_names(&mut cte_query_names, &cte_names);
names.extend(cte_query_names);
if !with.recursive {
cte_names.insert(cte_name);
}
}
}
let mut body_names = HashSet::new();
extract_tables_from_set_expr(&query.body, &mut body_names);
remove_cte_names(&mut body_names, &cte_names);
names.extend(body_names);
}
/// Helper function for [extract_tables_from_query].
///
/// Handle [SetExpr].
fn extract_tables_from_set_expr(set_expr: &SetExpr, names: &mut HashSet<ObjectName>) {
match set_expr {
SetExpr::Select(select) => {
for from in &select.from {
extract_tables_from_table_with_joins(from, names);
}
}
SetExpr::Query(query) => {
extract_tables_from_sql_query(query, names);
}
SetExpr::SetOperation { left, right, .. } => {
extract_tables_from_set_expr(left, names);
extract_tables_from_set_expr(right, names);
}
_ => {}
};
}
/// Helper function for [extract_tables_from_query].
///
/// Handle [TableWithJoins].
fn extract_tables_from_table_with_joins(
table_with_joins: &TableWithJoins,
names: &mut HashSet<ObjectName>,
) {
table_factor_to_object_name(&table_with_joins.relation, names);
for join in &table_with_joins.joins {
table_factor_to_object_name(&join.relation, names);
}
}
/// Helper function for [extract_tables_from_query].
///
/// Handle [TableFactor].
fn table_factor_to_object_name(table_factor: &TableFactor, names: &mut HashSet<ObjectName>) {
match table_factor {
TableFactor::Table { name, .. } => {
names.insert(name.to_owned());
}
TableFactor::Derived { subquery, .. } => {
extract_tables_from_sql_query(subquery, names);
}
TableFactor::NestedJoin {
table_with_joins, ..
} => {
extract_tables_from_table_with_joins(table_with_joins, names);
}
TableFactor::Pivot { table, .. }
| TableFactor::Unpivot { table, .. }
| TableFactor::MatchRecognize { table, .. } => {
table_factor_to_object_name(table, names);
}
TableFactor::TableFunction { .. }
| TableFactor::Function { .. }
| TableFactor::UNNEST { .. }
| TableFactor::JsonTable { .. }
| TableFactor::OpenJsonTable { .. }
| TableFactor::XmlTable { .. }
| TableFactor::SemanticView { .. } => {}
}
let _ = query.visit(&mut RelationCollector::new(names));
}
#[cfg(test)]
@@ -518,6 +557,65 @@ TQL EVAL (now() - '15s'::interval, now(), '5s') count_values("status_code", {__n
}
}
#[test]
fn test_extract_tables_from_tql_query_completeness() {
let testcases = [
("cpu", true, vec!["cpu"]),
(r#"{__name__="cpu"}"#, true, vec!["cpu"]),
(r#"cpu + {job="api"}"#, false, vec!["cpu"]),
(r#"{__name__=~"cpu.*"}"#, false, vec![]),
(r#"cpu{__schema__=~"private.*"}"#, false, vec![]),
];
for (promql, expected_complete, expected_tables) in testcases {
let sql = format!("TQL EVAL (0, 10, '5s') {promql}");
let mut stmts = ParserContext::create_with_dialect(
&sql,
&GreptimeDbDialect {},
ParseOptions::default(),
)
.unwrap();
let Statement::Tql(tql) = stmts.pop().unwrap() else {
unreachable!()
};
let query = SqlOrTql::Tql(tql, sql);
let (tables, complete) = extract_tables_from_query_inner(&query);
let mut tables = tables
.into_iter()
.map(|table| format_raw_object_name(&table))
.collect_vec();
tables.sort();
assert_eq!(expected_complete, complete, "{promql}");
assert_eq!(expected_tables, tables, "{promql}");
assert_eq!(
expected_complete,
extract_tables_from_query_checked(&query).is_some(),
"{promql}"
);
}
}
#[test]
fn test_extract_tables_from_chained_tql_ctes() {
let sql = "WITH denied AS (TQL EVAL (0, 10, '5s') allowed), \
leak AS (TQL EVAL (0, 10, '5s') denied) \
SELECT * FROM leak";
let mut stmts =
ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
.unwrap();
let Statement::Query(query) = stmts.pop().unwrap() else {
unreachable!()
};
let mut tables = extract_tables_from_query_checked(&SqlOrTql::Sql(*query, sql.to_string()))
.unwrap()
.map(|table| format_raw_object_name(&table))
.collect_vec();
tables.sort();
assert_eq!(vec!["allowed".to_string(), "denied".to_string()], tables);
}
#[test]
fn test_extract_tables_from_sql_query_with_derived_join() {
let sql = r#"
@@ -556,6 +654,37 @@ LEFT JOIN (
);
}
#[test]
fn test_extract_tables_from_sql_query_with_expression_subqueries() {
let sql = r#"
SELECT
(SELECT max(value) FROM scalar_source)
FROM outer_source
WHERE EXISTS (SELECT 1 FROM exists_source)
ORDER BY (SELECT max(value) FROM order_source);
"#;
let mut stmts =
ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
.unwrap();
let Statement::Query(query) = stmts.pop().unwrap() else {
unreachable!()
};
let mut tables = extract_tables_from_query(&SqlOrTql::Sql(*query, sql.to_string()))
.map(|table| format_raw_object_name(&table))
.collect_vec();
tables.sort();
assert_eq!(
vec![
"exists_source".to_string(),
"order_source".to_string(),
"outer_source".to_string(),
"scalar_source".to_string(),
],
tables
);
}
#[test]
fn test_extract_tables_from_sql_query_with_cte_scopes() {
let testcases = vec![
@@ -570,6 +699,15 @@ SELECT * FROM source;
),
(
r#"
WITH RECURSIVE source AS (
SELECT * FROM source
)
SELECT * FROM source;
"#,
vec![],
),
(
r#"
WITH first_cte AS (
SELECT * FROM physical_source
), second_cte AS (
@@ -579,6 +717,24 @@ SELECT * FROM second_cte;
"#,
vec!["physical_source".to_string()],
),
(
r#"
SELECT * FROM (
WITH nested_cte AS (SELECT * FROM nested_source)
SELECT * FROM nested_cte
);
"#,
vec!["nested_source".to_string()],
),
(
r#"
SELECT * FROM (
WITH nested_cte AS (SELECT * FROM nested_source)
SELECT (SELECT * FROM nested_cte)
);
"#,
vec!["nested_source".to_string()],
),
];
for (sql, expected_tables) in testcases {
@@ -603,6 +759,40 @@ SELECT * FROM second_cte;
}
}
#[test]
fn test_extract_tables_from_deeply_nested_ctes_visits_each_query_once() {
let mut sql = "SELECT * FROM physical_source".to_string();
for depth in 0..20 {
sql = format!("WITH cte_{depth} AS ({sql}) SELECT * FROM cte_{depth}");
}
let mut stmts = ParserContext::create_with_dialect(
&sql,
&GreptimeDbDialect {},
ParseOptions::default(),
)
.unwrap();
let Statement::Query(query) = stmts.pop().unwrap() else {
unreachable!()
};
let mut tables = HashSet::new();
let query_visits = {
let mut collector = RelationCollector::new(&mut tables);
let _ = query.inner.visit(&mut collector);
collector.query_visits
};
assert_eq!(21, query_visits);
assert_eq!(
vec!["physical_source"],
tables
.into_iter()
.map(|table| format_raw_object_name(&table))
.collect_vec()
);
}
#[test]
fn test_extract_tables_from_tql_query_with_schema_matcher() {
let sql = r#"