diff --git a/src/auth/src/permission.rs b/src/auth/src/permission.rs index 8914635290..cff9b41477 100644 --- a/src/auth/src/permission.rs +++ b/src/auth/src/permission.rs @@ -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; + + fn check_permission_with_context( + &self, + user_info: UserInfoRef, + req: PermissionReq, + _current_schema: Option<&str>, + ) -> Result { + 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 { + 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 { 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()); + } } diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 7242d601c1..5beb8eb827 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -611,7 +611,11 @@ impl Instance { let checker_ref = self.plugins.get::(); 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::() .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 diff --git a/src/frontend/src/instance/grpc.rs b/src/frontend/src/instance/grpc.rs index 0ca1a2cf20..808a2bd6ad 100644 --- a/src/frontend/src/instance/grpc.rs +++ b/src/frontend/src/instance/grpc.rs @@ -71,7 +71,11 @@ impl GrpcQueryHandler for Instance { self.plugins .get::() .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::() .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 diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 34b2550acf..5ec7b4bc3c 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -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), diff --git a/src/operator/src/statement/set.rs b/src/operator/src/statement/set.rs index 38a16e1b05..77515e62e0 100644 --- a/src/operator/src/statement/set.rs +++ b/src/operator/src/statement/set.rs @@ -129,36 +129,6 @@ pub fn set_bytea_output(exprs: Vec, ctx: QueryContextRef) -> Result<()> { Ok(()) } -pub fn set_search_path(exprs: Vec, 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 { diff --git a/src/sql/src/statements/set_variables.rs b/src/sql/src/statements/set_variables.rs index 4c6eb16692..346f1762e6 100644 --- a/src/sql/src/statements/set_variables.rs +++ b/src/sql/src/statements/set_variables.rs @@ -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, } +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()); + } + } } diff --git a/src/sql/src/statements/statement.rs b/src/sql/src/statements/statement.rs index bac7e367cd..ec7c73573b 100644 --- a/src/sql/src/statements/statement.rs +++ b/src/sql/src/statements/statement.rs @@ -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(_) diff --git a/src/sql/src/util.rs b/src/sql/src/util.rs index 6a7abe23bb..5868be6aa0 100644 --- a/src/sql/src/util.rs +++ b/src/sql/src/util.rs @@ -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 { - 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) { - if let Some(hybrid_cte) = &query.hybrid_cte { - let mut cte_names: HashSet = 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> { + 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, 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, +) -> bool { + let Some(hybrid_cte) = &query.hybrid_cte else { + return true; + }; + + let mut complete = true; + let cte_names: HashSet = 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, cte_names: &HashSet) { @@ -246,55 +257,55 @@ fn remove_cte_names(names: &mut HashSet, cte_names: &HashSet }); } -fn extract_tables_from_tql(tql: &Tql, names: &mut HashSet) { +fn extract_tables_from_tql(tql: &Tql, names: &mut HashSet) -> 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) { - 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) -> bool { + struct TableCollector<'a> { + names: &'a mut HashSet, + complete: bool, } + + impl ExprVisitor for TableCollector<'_> { + type Error = (); + + fn pre_visit(&mut self, expr: &PromExpr) -> std::result::Result { + 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, -) { +) -> 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, + next_cte: usize, + recursive: bool, + scope_start: usize, + add_to_parent_after: Option, +} + +struct RelationCollector<'a> { + names: &'a mut HashSet, + ctes_in_scope: Vec, + query_scopes: Vec, + #[cfg(test)] + query_visits: usize, +} + +impl<'a> RelationCollector<'a> { + fn new(names: &'a mut HashSet) -> 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 { + #[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 { + 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 { + 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) { - 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) { - 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, -) { - 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) { - 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#"