use crate::casing::IdentifierStyle;
use crate::catalog::Catalog;
use crate::diagnostic::{TableLevelDiagnostic, TableLevelDiagnosticKind};
use crate::error::Error;
use crate::extractor::{classify_statement, ExtractorOptions, StatementKind};
use crate::reference::{TableRead, TableWrite};
use crate::resolver::MergeActions;
use sqlparser::ast::Statement;
use sqlparser::dialect::Dialect;
pub fn extract_table_operations(
dialect: &dyn Dialect,
sql: &str,
) -> Result<Vec<Result<TableOperation, Error>>, Error> {
TableOperationExtractor::extract(dialect, sql)
}
pub fn extract_table_operations_with_options(
dialect: &dyn Dialect,
sql: &str,
options: ExtractorOptions,
) -> Result<Vec<Result<TableOperation, Error>>, Error> {
TableOperationExtractor::extract_with_options(dialect, sql, options)
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct TableOperation {
pub statement_kind: StatementKind,
pub reads: Vec<TableRead>,
pub writes: Vec<TableWrite>,
pub lineage: Vec<TableLineageEdge>,
pub diagnostics: Vec<TableLevelDiagnostic>,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct TableLineageEdge {
pub source: TableRead,
pub target: TableWrite,
}
#[derive(Default, Debug)]
pub struct TableOperationExtractor;
impl TableOperationExtractor {
pub fn extract(
dialect: &dyn Dialect,
sql: &str,
) -> Result<Vec<Result<TableOperation, Error>>, Error> {
Self::extract_with_options(dialect, sql, ExtractorOptions::new())
}
pub fn extract_with_options(
dialect: &dyn Dialect,
sql: &str,
options: ExtractorOptions,
) -> Result<Vec<Result<TableOperation, Error>>, Error> {
crate::extractor::extract_each(dialect, sql, options, Self::extract_from_statement)
}
pub(crate) fn extract_from_statement(
statement: &Statement,
catalog: Option<&Catalog>,
style: IdentifierStyle,
) -> Result<TableOperation, Error> {
Self::extract_inner(statement, catalog, style).map(|(op, ..)| op)
}
pub(crate) fn extract_inner(
statement: &Statement,
catalog: Option<&Catalog>,
style: IdentifierStyle,
) -> Result<
(
TableOperation,
Option<MergeActions>,
bool,
Option<crate::resolver::DataModifyingCteCrud>,
),
Error,
> {
let statement_kind = classify_statement(statement);
if statement_kind == StatementKind::Unsupported {
return Ok((
unsupported_table_operation(statement_kind, statement),
None,
false,
None,
));
}
let (plan, column_diagnostics) = crate::resolver::build(statement, catalog, style);
let merge_actions = crate::resolver::merge_actions(&plan);
let insert_updates = crate::resolver::insert_updates_on_conflict(&plan);
let outer_moves_data =
writes_data(&statement_kind) && merge_actions.is_none_or(|a| a.writes_data());
let emits_lineage = outer_moves_data || crate::resolver::has_data_modifying_cte(&plan);
let lineage = if emits_lineage {
crate::resolver::table_lineage(&plan, style.casing)
} else {
Vec::new()
};
let op = TableOperation {
statement_kind,
reads: crate::resolver::table_reads(&plan),
writes: crate::resolver::table_writes(&plan),
lineage,
diagnostics: column_diagnostics
.iter()
.filter_map(|d| d.to_table_level())
.collect(),
};
let cte_crud = crate::resolver::data_modifying_cte_crud(&plan);
Ok((op, merge_actions, insert_updates, cte_crud))
}
}
fn writes_data(kind: &StatementKind) -> bool {
matches!(
kind,
StatementKind::Insert
| StatementKind::Update
| StatementKind::Merge
| StatementKind::CreateTable
| StatementKind::CreateView
| StatementKind::AlterView
)
}
fn unsupported_table_operation(
statement_kind: StatementKind,
statement: &Statement,
) -> TableOperation {
TableOperation {
statement_kind,
reads: Vec::new(),
writes: Vec::new(),
lineage: Vec::new(),
diagnostics: vec![TableLevelDiagnostic {
kind: TableLevelDiagnosticKind::UnsupportedStatement,
message: crate::extractor::unsupported_message(statement),
span: None,
}],
}
}