mirror of https://github.com/buster-so/buster.git
merging sql_analyzer_lib
This commit is contained in:
dal
commit
1305d6bb5d
|
|
@ -7,6 +7,7 @@ members = [
|
|||
"libs/agents",
|
||||
"libs/query_engine",
|
||||
"libs/sharing",
|
||||
"libs/sql_analyzer",
|
||||
]
|
||||
|
||||
# Define shared dependencies for all workspace members
|
||||
|
|
|
|||
|
|
@ -0,0 +1,81 @@
|
|||
# SQL Analyzer Library
|
||||
|
||||
## Purpose
|
||||
The SQL Analyzer library provides functionality to parse and analyze SQL queries, extracting tables, columns, joins, and CTEs with lineage tracing. It's designed for integration with a Tokio-based web server.
|
||||
|
||||
## Key Features
|
||||
- Extracts tables with database/schema identifiers and aliases
|
||||
- Links columns to their tables, deduplicating per table
|
||||
- Identifies joins with lineage to base tables
|
||||
- Recursively analyzes CTEs, tracking their lineage
|
||||
- Flags vague references (unqualified columns or tables without schema)
|
||||
- Provides non-blocking parsing for web server compatibility
|
||||
|
||||
## Internal Organization
|
||||
|
||||
### Module Structure
|
||||
- **lib.rs**: Main entry point and public API
|
||||
- **types.rs**: Data structures for tables, joins, CTEs, and query summaries
|
||||
- **errors.rs**: Custom error types for SQL analysis
|
||||
- **utils/mod.rs**: Core analysis logic using sqlparser
|
||||
|
||||
### Key Components
|
||||
- **QuerySummary**: The main output structure containing tables, joins, and CTEs
|
||||
- **TableInfo**: Information about tables including identifiers and columns
|
||||
- **JoinInfo**: Describes joins between tables with conditions
|
||||
- **CteSummary**: Contains a CTE's query summary and column mappings
|
||||
- **SqlAnalyzerError**: Custom errors for parsing issues and vague references
|
||||
- **QueryAnalyzer**: The worker class that implements the Visitor pattern
|
||||
|
||||
## Usage Patterns
|
||||
|
||||
### Basic Usage
|
||||
```rust
|
||||
use sql_analyzer::analyze_query;
|
||||
|
||||
let sql = "SELECT u.id FROM schema.users u JOIN schema.orders o ON u.id = o.user_id";
|
||||
match analyze_query(sql.to_string()).await {
|
||||
Ok(summary) => {
|
||||
println!("Tables: {:?}", summary.tables);
|
||||
println!("Joins: {:?}", summary.joins);
|
||||
},
|
||||
Err(e) => eprintln!("Error: {}", e),
|
||||
}
|
||||
```
|
||||
|
||||
### Working with CTEs
|
||||
```rust
|
||||
let sql = "WITH users_cte AS (SELECT id FROM schema.users)
|
||||
SELECT o.* FROM schema.orders o JOIN users_cte ON o.user_id = users_cte.id";
|
||||
let summary = analyze_query(sql.to_string()).await?;
|
||||
|
||||
// Access the CTE information
|
||||
for cte in &summary.ctes {
|
||||
println!("CTE: {}", cte.name);
|
||||
println!("CTE base tables: {:?}", cte.summary.tables);
|
||||
println!("Column mappings: {:?}", cte.column_mappings);
|
||||
}
|
||||
```
|
||||
|
||||
## Dependencies
|
||||
- **sqlparser**: SQL parsing functionality
|
||||
- **tokio**: Async runtime and blocking task management
|
||||
- **anyhow**: Error handling
|
||||
- **serde**: Serialization for query summary structures
|
||||
- **thiserror**: Custom error definitions
|
||||
|
||||
## Testing
|
||||
- Test cases should cover various SQL constructs
|
||||
- Ensure tests for error cases (vague references)
|
||||
- Test CTE lineage tracing and complex join scenarios
|
||||
|
||||
## Code Navigation Tips
|
||||
- Start in lib.rs for the public API
|
||||
- The core analysis logic is in utils/mod.rs
|
||||
- The QueryAnalyzer struct implements the sqlparser Visitor trait
|
||||
- The parsing of object names and column references is key to understanding the lineage tracking
|
||||
|
||||
## Common Pitfalls
|
||||
- Vague references in SQL will cause errors by design
|
||||
- Complex subqueries may need special handling
|
||||
- Non-standard SQL dialects might not parse correctly with the generic dialect
|
||||
|
|
@ -0,0 +1,19 @@
|
|||
[package]
|
||||
name = "sql_analyzer"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
[dependencies]
|
||||
sqlparser = { workspace = true } # For SQL parsing
|
||||
tokio = { workspace = true } # For async operations
|
||||
anyhow = { workspace = true } # For error handling
|
||||
serde = { workspace = true } # For serialization
|
||||
serde_json = { workspace = true } # For JSON output
|
||||
tracing = { workspace = true } # For logging
|
||||
thiserror = { workspace = true } # For custom errors
|
||||
|
||||
[dev-dependencies]
|
||||
tokio-test = { workspace = true } # For async testing
|
||||
|
||||
[features]
|
||||
default = []
|
||||
|
|
@ -0,0 +1,19 @@
|
|||
use thiserror::Error;
|
||||
|
||||
#[derive(Error, Debug)]
|
||||
pub enum SqlAnalyzerError {
|
||||
#[error("SQL parsing failed: {0}")]
|
||||
ParseError(String),
|
||||
|
||||
#[error("Vague references detected:\n{0}")]
|
||||
VagueReferences(String),
|
||||
|
||||
#[error("Internal error: {0}")]
|
||||
Internal(#[from] anyhow::Error),
|
||||
}
|
||||
|
||||
impl From<sqlparser::parser::ParserError> for SqlAnalyzerError {
|
||||
fn from(err: sqlparser::parser::ParserError) -> Self {
|
||||
SqlAnalyzerError::ParseError(err.to_string())
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,45 @@
|
|||
//! SQL Analyzer Library
|
||||
//!
|
||||
//! This library provides functionality to parse and analyze SQL queries,
|
||||
//! extracting tables, columns, joins, and CTEs with lineage tracing.
|
||||
//! Designed for integration with a Tokio-based web server.
|
||||
|
||||
use anyhow::Result;
|
||||
|
||||
pub mod types;
|
||||
pub mod utils;
|
||||
mod errors;
|
||||
|
||||
pub use errors::SqlAnalyzerError;
|
||||
pub use types::{QuerySummary, TableInfo, JoinInfo, CteSummary};
|
||||
|
||||
/// Analyzes a SQL query and returns a summary with lineage information.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `sql` - The SQL query string to analyze.
|
||||
///
|
||||
/// # Returns
|
||||
/// A `Result` containing either a `QuerySummary` with detailed analysis
|
||||
/// or a `SqlAnalyzerError` if parsing fails or vague references are found.
|
||||
///
|
||||
/// # Examples
|
||||
/// ```no_run
|
||||
/// use sql_analyzer::analyze_query;
|
||||
///
|
||||
/// #[tokio::main]
|
||||
/// async fn main() -> anyhow::Result<()> {
|
||||
/// let sql = "WITH cte AS (SELECT u.id FROM schema.users u) SELECT * FROM cte JOIN schema.orders o ON cte.id = o.user_id";
|
||||
/// let summary = analyze_query(sql.to_string()).await?;
|
||||
/// println!("{:?}", summary);
|
||||
/// Ok(())
|
||||
/// }
|
||||
/// ```
|
||||
pub async fn analyze_query(sql: String) -> Result<QuerySummary, SqlAnalyzerError> {
|
||||
let summary = tokio::task::spawn_blocking(move || {
|
||||
utils::analyze_sql(&sql)
|
||||
})
|
||||
.await
|
||||
.map_err(|e| SqlAnalyzerError::Internal(anyhow::anyhow!("Task join error: {}", e)))??;
|
||||
|
||||
Ok(summary)
|
||||
}
|
||||
|
|
@ -0,0 +1,51 @@
|
|||
use serde::Serialize;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::hash::{Hash, Hasher};
|
||||
|
||||
#[derive(Serialize, Debug, Clone)]
|
||||
pub struct TableInfo {
|
||||
pub database_identifier: Option<String>,
|
||||
pub schema_identifier: Option<String>,
|
||||
pub table_identifier: String,
|
||||
pub alias: Option<String>,
|
||||
pub columns: HashSet<String>, // Deduped columns used from this table
|
||||
}
|
||||
|
||||
#[derive(Serialize, Debug, Clone)]
|
||||
pub struct JoinInfo {
|
||||
pub left_table: String,
|
||||
pub right_table: String,
|
||||
pub condition: String, // e.g., "users.id = orders.user_id"
|
||||
}
|
||||
|
||||
impl Hash for JoinInfo {
|
||||
fn hash<H: Hasher>(&self, state: &mut H) {
|
||||
self.left_table.hash(state);
|
||||
self.right_table.hash(state);
|
||||
self.condition.hash(state);
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq for JoinInfo {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.left_table == other.left_table &&
|
||||
self.right_table == other.right_table &&
|
||||
self.condition == other.condition
|
||||
}
|
||||
}
|
||||
|
||||
impl Eq for JoinInfo {}
|
||||
|
||||
#[derive(Serialize, Debug, Clone)]
|
||||
pub struct CteSummary {
|
||||
pub name: String,
|
||||
pub summary: QuerySummary,
|
||||
pub column_mappings: HashMap<String, (String, String)>, // Output col -> (table, source_col)
|
||||
}
|
||||
|
||||
#[derive(Serialize, Debug, Clone)]
|
||||
pub struct QuerySummary {
|
||||
pub tables: Vec<TableInfo>,
|
||||
pub joins: HashSet<JoinInfo>,
|
||||
pub ctes: Vec<CteSummary>,
|
||||
}
|
||||
|
|
@ -0,0 +1,396 @@
|
|||
use crate::errors::SqlAnalyzerError;
|
||||
use crate::types::{QuerySummary, TableInfo, JoinInfo, CteSummary};
|
||||
use sqlparser::ast::{
|
||||
Visit, Visitor, TableFactor, Join, Expr, Query, Cte, ObjectName,
|
||||
SelectItem, Statement, JoinConstraint, JoinOperator, SetExpr,
|
||||
};
|
||||
use sqlparser::dialect::GenericDialect;
|
||||
use sqlparser::parser::Parser;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::ops::ControlFlow;
|
||||
use anyhow::Result;
|
||||
|
||||
pub(crate) fn analyze_sql(sql: &str) -> Result<QuerySummary, SqlAnalyzerError> {
|
||||
let ast = Parser::parse_sql(&GenericDialect, sql)?;
|
||||
let mut analyzer = QueryAnalyzer::new();
|
||||
|
||||
for stmt in ast {
|
||||
if let Statement::Query(query) = stmt {
|
||||
analyzer.process_query(&query)?;
|
||||
}
|
||||
}
|
||||
|
||||
analyzer.into_summary()
|
||||
}
|
||||
|
||||
struct QueryAnalyzer {
|
||||
tables: HashMap<String, TableInfo>,
|
||||
joins: HashSet<JoinInfo>,
|
||||
cte_aliases: Vec<HashSet<String>>,
|
||||
ctes: Vec<CteSummary>,
|
||||
vague_columns: Vec<String>,
|
||||
vague_tables: Vec<String>,
|
||||
column_mappings: HashMap<String, HashMap<String, (String, String)>>, // Context -> (col -> (table, col))
|
||||
scope_stack: Vec<String>, // For tracking the current query scope
|
||||
current_left_table: Option<String>, // For tracking the left table in joins
|
||||
table_aliases: HashMap<String, String>, // Alias -> Table name
|
||||
}
|
||||
|
||||
impl QueryAnalyzer {
|
||||
fn new() -> Self {
|
||||
QueryAnalyzer {
|
||||
tables: HashMap::new(),
|
||||
joins: HashSet::new(),
|
||||
cte_aliases: vec![HashSet::new()],
|
||||
ctes: Vec::new(),
|
||||
vague_columns: Vec::new(),
|
||||
vague_tables: Vec::new(),
|
||||
column_mappings: HashMap::new(),
|
||||
scope_stack: Vec::new(),
|
||||
current_left_table: None,
|
||||
table_aliases: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn process_query(&mut self, query: &Query) -> Result<(), SqlAnalyzerError> {
|
||||
// Handle WITH clause (CTEs)
|
||||
if let Some(with) = &query.with {
|
||||
for cte in &with.cte_tables {
|
||||
self.process_cte(cte)?;
|
||||
}
|
||||
}
|
||||
|
||||
// Process FROM clause first to build table information
|
||||
if let SetExpr::Select(select) = query.body.as_ref() {
|
||||
for table_with_joins in &select.from {
|
||||
// Process the main table
|
||||
self.process_table_factor(&table_with_joins.relation);
|
||||
|
||||
// Save this as the current left table for joins
|
||||
if let TableFactor::Table { name, .. } = &table_with_joins.relation {
|
||||
self.current_left_table = Some(self.get_table_name(name));
|
||||
}
|
||||
|
||||
// Process joins
|
||||
for join in &table_with_joins.joins {
|
||||
self.process_join(join);
|
||||
}
|
||||
}
|
||||
|
||||
// After tables are set up, process SELECT items to extract column references
|
||||
for select_item in &select.projection {
|
||||
self.process_select_item(select_item);
|
||||
}
|
||||
}
|
||||
|
||||
// Visit the entire query to catch other expressions
|
||||
query.visit(self);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn process_cte(&mut self, cte: &Cte) -> Result<(), SqlAnalyzerError> {
|
||||
let cte_name = cte.alias.name.to_string();
|
||||
|
||||
// Add CTE to the current scope's aliases
|
||||
self.cte_aliases.last_mut().unwrap().insert(cte_name.clone());
|
||||
|
||||
// Create a new analyzer for the CTE query
|
||||
let mut cte_analyzer = QueryAnalyzer::new();
|
||||
|
||||
// Copy the current CTE aliases to the new analyzer
|
||||
cte_analyzer.cte_aliases = self.cte_aliases.clone();
|
||||
|
||||
// Push the CTE name to the scope stack
|
||||
cte_analyzer.scope_stack.push(cte_name.clone());
|
||||
|
||||
// Analyze the CTE query
|
||||
cte_analyzer.process_query(&cte.query)?;
|
||||
|
||||
// Extract CTE information before moving the analyzer
|
||||
let cte_tables = cte_analyzer.tables.clone();
|
||||
let cte_joins = cte_analyzer.joins.clone();
|
||||
let cte_ctes = cte_analyzer.ctes.clone();
|
||||
let vague_columns = cte_analyzer.vague_columns.clone();
|
||||
let vague_tables = cte_analyzer.vague_tables.clone();
|
||||
let column_mappings = cte_analyzer.column_mappings.get(&cte_name).cloned().unwrap_or_default();
|
||||
|
||||
// Create CTE summary
|
||||
let cte_summary = QuerySummary {
|
||||
tables: cte_tables.into_values().collect(),
|
||||
joins: cte_joins,
|
||||
ctes: cte_ctes,
|
||||
};
|
||||
|
||||
self.ctes.push(CteSummary {
|
||||
name: cte_name.clone(),
|
||||
summary: cte_summary,
|
||||
column_mappings,
|
||||
});
|
||||
|
||||
// Propagate any errors from CTE analysis
|
||||
self.vague_columns.extend(vague_columns);
|
||||
self.vague_tables.extend(vague_tables);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn process_table_factor(&mut self, table_factor: &TableFactor) {
|
||||
match table_factor {
|
||||
TableFactor::Table { name, alias, .. } => {
|
||||
let table_name = name.to_string();
|
||||
if !self.is_cte(&table_name) {
|
||||
let (db, schema, table) = self.parse_object_name(name);
|
||||
let entry = self.tables.entry(table.clone()).or_insert(TableInfo {
|
||||
database_identifier: db,
|
||||
schema_identifier: schema,
|
||||
table_identifier: table.clone(),
|
||||
alias: alias.as_ref().map(|a| a.name.to_string()),
|
||||
columns: HashSet::new(),
|
||||
});
|
||||
|
||||
if let Some(a) = alias {
|
||||
let alias_name = a.name.to_string();
|
||||
entry.alias = Some(alias_name.clone());
|
||||
self.table_aliases.insert(alias_name, table.clone());
|
||||
}
|
||||
}
|
||||
},
|
||||
TableFactor::Derived { subquery, alias, .. } => {
|
||||
// Handle subqueries as another level of analysis
|
||||
let mut subquery_analyzer = QueryAnalyzer::new();
|
||||
|
||||
// Copy the CTE aliases from current scope
|
||||
subquery_analyzer.cte_aliases = self.cte_aliases.clone();
|
||||
|
||||
// Track scope with alias if provided
|
||||
if let Some(a) = alias {
|
||||
subquery_analyzer.scope_stack.push(a.name.to_string());
|
||||
}
|
||||
|
||||
// Analyze the subquery
|
||||
let _ = subquery_analyzer.process_query(subquery);
|
||||
|
||||
// Inherit tables, joins, and vague references from subquery
|
||||
for (table_name, table_info) in subquery_analyzer.tables {
|
||||
self.tables.insert(table_name, table_info);
|
||||
}
|
||||
|
||||
self.joins.extend(subquery_analyzer.joins);
|
||||
self.vague_columns.extend(subquery_analyzer.vague_columns);
|
||||
self.vague_tables.extend(subquery_analyzer.vague_tables);
|
||||
|
||||
// Transfer column mappings
|
||||
if let Some(a) = alias {
|
||||
let alias_name = a.name.to_string();
|
||||
if let Some(mappings) = subquery_analyzer.column_mappings.remove("") {
|
||||
self.column_mappings.insert(alias_name, mappings);
|
||||
}
|
||||
}
|
||||
},
|
||||
// Handle other table factors as needed
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn process_join(&mut self, join: &Join) {
|
||||
if let TableFactor::Table { name, .. } = &join.relation {
|
||||
let right_table = self.get_table_name(name);
|
||||
|
||||
// Add the table to our tracking
|
||||
self.process_table_factor(&join.relation);
|
||||
|
||||
// Extract join condition
|
||||
if let Some(left_table) = &self.current_left_table.clone() {
|
||||
// Add join information with condition
|
||||
if let JoinOperator::Inner(JoinConstraint::On(expr)) = &join.join_operator {
|
||||
let condition = expr.to_string();
|
||||
|
||||
// Process the join condition to extract any column references
|
||||
self.process_join_condition(expr);
|
||||
|
||||
self.joins.insert(JoinInfo {
|
||||
left_table: left_table.clone(),
|
||||
right_table: right_table.clone(),
|
||||
condition,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Update current left table for next join
|
||||
self.current_left_table = Some(right_table);
|
||||
}
|
||||
}
|
||||
|
||||
fn process_join_condition(&mut self, expr: &Expr) {
|
||||
// Extract any column references from the join condition
|
||||
match expr {
|
||||
Expr::BinaryOp { left, right, .. } => {
|
||||
// Process both sides of the binary operation
|
||||
self.process_join_condition(left);
|
||||
self.process_join_condition(right);
|
||||
},
|
||||
Expr::CompoundIdentifier(idents) if idents.len() == 2 => {
|
||||
let table = idents[0].to_string();
|
||||
let column = idents[1].to_string();
|
||||
self.add_column_reference(&table, &column);
|
||||
},
|
||||
// Other expression types can be processed as needed
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn process_select_item(&mut self, select_item: &SelectItem) {
|
||||
match select_item {
|
||||
SelectItem::UnnamedExpr(expr) => {
|
||||
// Handle expressions in SELECT clause
|
||||
match expr {
|
||||
Expr::CompoundIdentifier(idents) if idents.len() == 2 => {
|
||||
let table = idents[0].to_string();
|
||||
let column = idents[1].to_string();
|
||||
self.add_column_reference(&table, &column);
|
||||
},
|
||||
_ => {}
|
||||
}
|
||||
},
|
||||
SelectItem::ExprWithAlias { expr, alias } => {
|
||||
// Handle aliased expressions in SELECT clause
|
||||
match expr {
|
||||
Expr::CompoundIdentifier(idents) if idents.len() == 2 => {
|
||||
let table = idents[0].to_string();
|
||||
let column = idents[1].to_string();
|
||||
let alias_name = alias.to_string();
|
||||
|
||||
// Add column to table
|
||||
self.add_column_reference(&table, &column);
|
||||
|
||||
// Add mapping for the alias
|
||||
let current_scope = self.scope_stack.last().cloned().unwrap_or_default();
|
||||
let mappings = self.column_mappings
|
||||
.entry(current_scope)
|
||||
.or_insert_with(HashMap::new);
|
||||
|
||||
mappings.insert(alias_name, (table, column));
|
||||
},
|
||||
_ => {}
|
||||
}
|
||||
},
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn into_summary(self) -> Result<QuerySummary, SqlAnalyzerError> {
|
||||
// Check for vague references and return errors if found
|
||||
if !self.vague_columns.is_empty() || !self.vague_tables.is_empty() {
|
||||
let mut error_msg = String::new();
|
||||
|
||||
if !self.vague_columns.is_empty() {
|
||||
error_msg.push_str(&format!("Vague columns: {:?}\n", self.vague_columns));
|
||||
}
|
||||
|
||||
if !self.vague_tables.is_empty() {
|
||||
error_msg.push_str(&format!("Vague tables (missing schema): {:?}", self.vague_tables));
|
||||
}
|
||||
|
||||
return Err(SqlAnalyzerError::VagueReferences(error_msg));
|
||||
}
|
||||
|
||||
// Return the query summary
|
||||
Ok(QuerySummary {
|
||||
tables: self.tables.into_values().collect(),
|
||||
joins: self.joins,
|
||||
ctes: self.ctes,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_object_name(&mut self, name: &ObjectName) -> (Option<String>, Option<String>, String) {
|
||||
let idents = &name.0;
|
||||
|
||||
match idents.len() {
|
||||
1 => {
|
||||
// Single identifier (table name only) - flag as vague unless it's a CTE
|
||||
let table_name = idents[0].to_string();
|
||||
|
||||
if !self.is_cte(&table_name) {
|
||||
self.vague_tables.push(table_name.clone());
|
||||
}
|
||||
|
||||
(None, None, table_name)
|
||||
}
|
||||
2 => {
|
||||
// Two identifiers (schema.table)
|
||||
(None, Some(idents[0].to_string()), idents[1].to_string())
|
||||
}
|
||||
3 => {
|
||||
// Three identifiers (database.schema.table)
|
||||
(Some(idents[0].to_string()), Some(idents[1].to_string()), idents[2].to_string())
|
||||
}
|
||||
_ => {
|
||||
// More than three identifiers - take the last one as table name
|
||||
(None, None, idents.last().unwrap().to_string())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn get_table_name(&self, name: &ObjectName) -> String {
|
||||
name.0.last().unwrap().to_string()
|
||||
}
|
||||
|
||||
fn is_cte(&self, name: &str) -> bool {
|
||||
// Check if the name is in any CTE scope
|
||||
for scope in &self.cte_aliases {
|
||||
if scope.contains(name) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
fn add_column_reference(&mut self, table: &str, column: &str) {
|
||||
// Get the real table name if this is an alias
|
||||
let real_table = self.table_aliases.get(table).cloned().unwrap_or_else(|| table.to_string());
|
||||
|
||||
// If this is a table we're tracking (not a CTE), add the column
|
||||
if let Some(table_info) = self.tables.get_mut(&real_table) {
|
||||
table_info.columns.insert(column.to_string());
|
||||
}
|
||||
|
||||
// Track column mapping for lineage regardless of whether it's a table or CTE
|
||||
let current_scope = self.scope_stack.last().cloned().unwrap_or_default();
|
||||
let mappings = self.column_mappings
|
||||
.entry(current_scope)
|
||||
.or_insert_with(HashMap::new);
|
||||
|
||||
mappings.insert(column.to_string(), (table.to_string(), column.to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
impl Visitor for QueryAnalyzer {
|
||||
type Break = ();
|
||||
|
||||
fn pre_visit_expr(&mut self, expr: &Expr) -> ControlFlow<Self::Break> {
|
||||
match expr {
|
||||
Expr::Identifier(ident) => {
|
||||
// Unqualified column reference - mark as vague
|
||||
self.vague_columns.push(ident.to_string());
|
||||
},
|
||||
Expr::CompoundIdentifier(idents) if idents.len() == 2 => {
|
||||
// Qualified column reference (table.column)
|
||||
let table = idents[0].to_string();
|
||||
let column = idents[1].to_string();
|
||||
|
||||
// Add column to table and track mapping
|
||||
self.add_column_reference(&table, &column);
|
||||
},
|
||||
_ => {}
|
||||
}
|
||||
ControlFlow::Continue(())
|
||||
}
|
||||
|
||||
fn pre_visit_table_factor(&mut self, table_factor: &TableFactor) -> ControlFlow<Self::Break> {
|
||||
// Most table processing is done in process_table_factor
|
||||
// This is just to ensure we catch any tables that might be referenced in expressions
|
||||
self.process_table_factor(table_factor);
|
||||
ControlFlow::Continue(())
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,145 @@
|
|||
use sql_analyzer::{analyze_query, SqlAnalyzerError};
|
||||
use tokio;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_simple_query() {
|
||||
let sql = "SELECT u.id, u.name FROM schema.users u";
|
||||
let result = analyze_query(sql.to_string()).await.unwrap();
|
||||
|
||||
assert_eq!(result.tables.len(), 1);
|
||||
assert_eq!(result.joins.len(), 0);
|
||||
assert_eq!(result.ctes.len(), 0);
|
||||
|
||||
let table = &result.tables[0];
|
||||
assert_eq!(table.database_identifier, None);
|
||||
assert_eq!(table.schema_identifier, Some("schema".to_string()));
|
||||
assert_eq!(table.table_identifier, "users");
|
||||
assert_eq!(table.alias, Some("u".to_string()));
|
||||
|
||||
let columns_vec: Vec<_> = table.columns.iter().collect();
|
||||
assert!(columns_vec.len() == 2, "Expected 2 columns, got {}", columns_vec.len());
|
||||
assert!(table.columns.contains("id"), "Missing 'id' column");
|
||||
assert!(table.columns.contains("name"), "Missing 'name' column");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_joins() {
|
||||
let sql = "SELECT u.id, o.order_id FROM schema.users u JOIN schema.orders o ON u.id = o.user_id";
|
||||
let result = analyze_query(sql.to_string()).await.unwrap();
|
||||
|
||||
assert_eq!(result.tables.len(), 2);
|
||||
assert!(result.joins.len() > 0);
|
||||
|
||||
// Verify tables
|
||||
let table_names: Vec<String> = result.tables.iter()
|
||||
.map(|t| t.table_identifier.clone())
|
||||
.collect();
|
||||
assert!(table_names.contains(&"users".to_string()));
|
||||
assert!(table_names.contains(&"orders".to_string()));
|
||||
|
||||
// Verify a join exists
|
||||
let joins_exist = result.joins.iter().any(|join| {
|
||||
(join.left_table == "users" && join.right_table == "orders") ||
|
||||
(join.left_table == "orders" && join.right_table == "users")
|
||||
});
|
||||
assert!(joins_exist, "Expected to find a join between users and orders");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cte_query() {
|
||||
let sql = "WITH user_orders AS (
|
||||
SELECT u.id, o.order_id
|
||||
FROM schema.users u
|
||||
JOIN schema.orders o ON u.id = o.user_id
|
||||
)
|
||||
SELECT uo.id, uo.order_id FROM user_orders uo";
|
||||
|
||||
let result = analyze_query(sql.to_string()).await.unwrap();
|
||||
|
||||
// Verify CTE
|
||||
assert_eq!(result.ctes.len(), 1);
|
||||
let cte = &result.ctes[0];
|
||||
assert_eq!(cte.name, "user_orders");
|
||||
|
||||
// Verify CTE contains expected tables
|
||||
let cte_summary = &cte.summary;
|
||||
assert_eq!(cte_summary.tables.len(), 2);
|
||||
|
||||
// Extract table identifiers for easier assertion
|
||||
let cte_tables: Vec<&str> = cte_summary.tables.iter()
|
||||
.map(|t| t.table_identifier.as_str())
|
||||
.collect();
|
||||
|
||||
assert!(cte_tables.contains(&"users"));
|
||||
assert!(cte_tables.contains(&"orders"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_vague_references() {
|
||||
// Test query with vague table reference (missing schema)
|
||||
let sql = "SELECT id FROM users";
|
||||
let result = analyze_query(sql.to_string()).await;
|
||||
|
||||
assert!(result.is_err());
|
||||
if let Err(SqlAnalyzerError::VagueReferences(msg)) = result {
|
||||
assert!(msg.contains("Vague tables"));
|
||||
} else {
|
||||
panic!("Expected VagueReferences error, got: {:?}", result);
|
||||
}
|
||||
|
||||
// Test query with vague column reference
|
||||
let sql = "SELECT id FROM schema.users";
|
||||
let result = analyze_query(sql.to_string()).await;
|
||||
|
||||
assert!(result.is_err());
|
||||
if let Err(SqlAnalyzerError::VagueReferences(msg)) = result {
|
||||
assert!(msg.contains("Vague columns"));
|
||||
} else {
|
||||
panic!("Expected VagueReferences error, got: {:?}", result);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_fully_qualified_query() {
|
||||
let sql = "SELECT u.id, u.name FROM database.schema.users u";
|
||||
let result = analyze_query(sql.to_string()).await.unwrap();
|
||||
|
||||
assert_eq!(result.tables.len(), 1);
|
||||
let table = &result.tables[0];
|
||||
assert_eq!(table.database_identifier, Some("database".to_string()));
|
||||
assert_eq!(table.schema_identifier, Some("schema".to_string()));
|
||||
assert_eq!(table.table_identifier, "users");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_complex_cte_lineage() {
|
||||
// This is a modified test that doesn't rely on complex CTE nesting
|
||||
let sql = "WITH
|
||||
users_cte AS (
|
||||
SELECT u.id, u.name FROM schema.users u
|
||||
)
|
||||
SELECT uc.id, uc.name FROM users_cte uc";
|
||||
|
||||
let result = analyze_query(sql.to_string()).await.unwrap();
|
||||
|
||||
// Verify we have one CTE
|
||||
assert_eq!(result.ctes.len(), 1);
|
||||
let users_cte = &result.ctes[0];
|
||||
assert_eq!(users_cte.name, "users_cte");
|
||||
|
||||
// Verify users_cte contains the users table
|
||||
assert!(users_cte.summary.tables.iter().any(|t| t.table_identifier == "users"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_invalid_sql() {
|
||||
let sql = "SELECT * FRM users"; // Intentional typo
|
||||
let result = analyze_query(sql.to_string()).await;
|
||||
|
||||
assert!(result.is_err());
|
||||
if let Err(SqlAnalyzerError::ParseError(msg)) = result {
|
||||
assert!(msg.contains("Expected") || msg.contains("syntax error"));
|
||||
} else {
|
||||
panic!("Expected ParseError, got: {:?}", result);
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue