use sqlparser::ast::Ident;
use super::logical_plan::{Binding, BoundColumn, Cte, Expr, LogicalPlan, NamedExpr};
use super::reads::column_read;
use crate::casing::{CaseRule, IdentifierCasing};
use crate::extractor::ColumnLineageKind;
use crate::reference::{ColumnRead, ColumnReference, ColumnWrite, ResolutionKind, TableReference};
pub(super) struct TraceContext<'a> {
pub(super) ctes: Vec<&'a Cte>,
pub(super) active: Vec<String>,
pub(super) casing: IdentifierCasing,
}
impl<'a> TraceContext<'a> {
pub(super) fn new(op: &'a LogicalPlan, casing: IdentifierCasing) -> Self {
let mut ctes = Vec::new();
let mut node = op;
while let LogicalPlan::With(w) = node {
ctes.extend(w.ctes.iter());
node = &w.body;
}
TraceContext {
ctes,
active: Vec::new(),
casing,
}
}
fn eq_column(&self, a: &Ident, b: &Ident) -> bool {
self.casing.column.normalize(a) == self.casing.column.normalize(b)
}
fn eq_alias(&self, a: &Ident, b: &Ident) -> bool {
self.casing.table_alias.normalize(a) == self.casing.table_alias.normalize(b)
}
pub(super) fn with_decls<R>(
&mut self,
ctes: &'a [Cte],
f: impl FnOnce(&mut TraceContext<'a>) -> R,
) -> R {
let added = ctes.len();
ctes.iter().for_each(|c| self.ctes.push(c));
let r = f(self);
self.ctes.truncate(self.ctes.len() - added);
r
}
pub(super) fn with_cte_refs<R>(
&mut self,
ctes: &[&'a Cte],
f: impl FnOnce(&mut TraceContext<'a>) -> R,
) -> R {
let added = ctes.len();
self.ctes.extend_from_slice(ctes);
let r = f(self);
self.ctes.truncate(self.ctes.len() - added);
r
}
pub(super) fn enter_cte<R>(
&mut self,
name: &Ident,
f: impl FnOnce(&mut TraceContext<'a>, &'a LogicalPlan) -> R,
) -> Option<R> {
if self.active.iter().any(|n| n == &name.value) {
return None; }
let alias_fold = self.casing.table_alias;
let cte = self
.ctes
.iter()
.rev()
.find(|c| alias_fold.normalize(&c.name) == alias_fold.normalize(name))
.copied()?;
self.active.push(name.value.clone());
let r = f(self, &cte.body);
self.active.pop();
Some(r)
}
}
pub(super) fn origins_of_expr<'a>(
expr: &'a Expr,
input: &'a LogicalPlan,
context: &mut TraceContext<'a>,
) -> Vec<(ColumnRead, ColumnLineageKind)> {
match expr {
Expr::Column(c) => origins_of_ref(c, input, context),
Expr::Call { args } => {
transform(args.iter().flat_map(|e| origins_of_expr(e, input, context)))
}
Expr::Case {
then, else_result, ..
} => {
let mut sources: Vec<_> = then
.iter()
.flat_map(|e| origins_of_expr(e, input, context))
.collect();
if let Some(e) = else_result {
sources.extend(origins_of_expr(e, input, context));
}
transform(sources)
}
Expr::Window { arg, .. } => transform(origins_of_expr(arg, input, context)),
Expr::Subquery { plan, output } => {
transform(subquery_output_origins(plan, *output, context))
}
Expr::Fanin(refs) => refs
.iter()
.flat_map(|c| origins_of_ref(c, input, context))
.collect(),
Expr::InSubquery { expr, .. } => transform(origins_of_expr(expr, input, context)),
Expr::Exists(_) | Expr::Filter(_) => Vec::new(),
}
}
fn origins_of_ref<'a>(
c: &'a BoundColumn,
input: &'a LogicalPlan,
context: &mut TraceContext<'a>,
) -> Vec<(ColumnRead, ColumnLineageKind)> {
match &c.binding {
Binding::Base { .. } | Binding::Unresolved | Binding::Ambiguous => column_read(c)
.map(|r| vec![(r, ColumnLineageKind::Passthrough)])
.unwrap_or_default(),
Binding::Derived => origins_into(input, c.qualifier.as_ref(), &c.name, context),
Binding::Local => Vec::new(),
}
}
fn values_backed(op: &LogicalPlan) -> bool {
match op {
LogicalPlan::Values(_) => true,
LogicalPlan::Sort(s) => values_backed(&s.input),
LogicalPlan::Filter(f) => values_backed(&f.input),
LogicalPlan::With(w) => values_backed(&w.body),
LogicalPlan::Scan(_)
| LogicalPlan::Join(_)
| LogicalPlan::Aggregate(_)
| LogicalPlan::Projection(_)
| LogicalPlan::SetOp(_)
| LogicalPlan::SubqueryAlias(_)
| LogicalPlan::TableFunction(_)
| LogicalPlan::CteRef(_)
| LogicalPlan::Empty
| LogicalPlan::Insert(_)
| LogicalPlan::Update(_)
| LogicalPlan::Delete(_)
| LogicalPlan::Merge(_)
| LogicalPlan::CreateTableAs(_)
| LogicalPlan::CreateView(_)
| LogicalPlan::AlterTable(_)
| LogicalPlan::Drop(_) => false,
}
}
fn origins_into<'a>(
op: &'a LogicalPlan,
qualifier: Option<&Ident>,
name: &Ident,
context: &mut TraceContext<'a>,
) -> Vec<(ColumnRead, ColumnLineageKind)> {
match op {
LogicalPlan::Projection(p) => match find_named(&p.exprs, name, context.casing.column) {
Some(ne) => origins_of_expr(&ne.expr, &p.input, context),
None => Vec::new(),
},
LogicalPlan::Aggregate(a) => origins_into(&a.input, qualifier, name, context),
LogicalPlan::Filter(f) => origins_into(&f.input, qualifier, name, context),
LogicalPlan::Sort(s) => origins_into(&s.input, qualifier, name, context),
LogicalPlan::Join(j) => {
let mut o = origins_into(&j.left, qualifier, name, context);
o.extend(origins_into(&j.right, qualifier, name, context));
o
}
LogicalPlan::SubqueryAlias(sa) => {
if !qualifier.is_none_or(|q| context.eq_alias(q, &sa.alias)) {
Vec::new()
} else if values_backed(&sa.input) {
vec![(
synthetic_source(&sa.alias, name),
ColumnLineageKind::Passthrough,
)]
} else {
origins_into(&sa.input, None, name, context)
}
}
LogicalPlan::TableFunction(tf) => match &tf.alias {
Some(alias) if qualifier.is_none_or(|q| context.eq_alias(q, alias)) => {
vec![(
synthetic_source(alias, name),
ColumnLineageKind::Passthrough,
)]
}
_ => Vec::new(),
},
LogicalPlan::SetOp(_) => {
let operands = output_operands(op);
let Some(i) = operands.first().and_then(|o| {
o.outputs
.iter()
.position(|ne| ne.name.as_ref().is_some_and(|n| context.eq_column(n, name)))
}) else {
return Vec::new();
};
trace_nth_output(&operands, i, context)
}
LogicalPlan::With(w) => context.with_decls(&w.ctes, |context| {
origins_into(&w.body, qualifier, name, context)
}),
LogicalPlan::CteRef(r) => {
let exposed = r.alias.as_ref().unwrap_or(&r.name);
if !qualifier.is_none_or(|q| context.eq_alias(q, exposed)) {
return Vec::new();
}
context
.enter_cte(&r.name, |context, body| {
if values_backed(body) {
vec![(
synthetic_source(&r.name, name),
ColumnLineageKind::Passthrough,
)]
} else {
origins_into(body, None, name, context)
}
})
.unwrap_or_default()
}
LogicalPlan::Scan(_) | LogicalPlan::Values(_) | LogicalPlan::Empty => Vec::new(),
LogicalPlan::Insert(_)
| LogicalPlan::Update(_)
| LogicalPlan::Delete(_)
| LogicalPlan::Merge(_)
| LogicalPlan::CreateTableAs(_)
| LogicalPlan::CreateView(_)
| LogicalPlan::AlterTable(_)
| LogicalPlan::Drop(_) => Vec::new(),
}
}
fn subquery_output_origins<'a>(
op: &'a LogicalPlan,
output: usize,
context: &mut TraceContext<'a>,
) -> Vec<(ColumnRead, ColumnLineageKind)> {
trace_nth_output(&output_operands(op), output, context)
}
fn trace_nth_output<'a>(
operands: &[Operand<'a>],
i: usize,
context: &mut TraceContext<'a>,
) -> Vec<(ColumnRead, ColumnLineageKind)> {
let mut out = Vec::new();
for operand in operands {
if let Some(ne) = operand.outputs.get(i) {
out.extend(operand.trace(context, |input, cx| origins_of_expr(&ne.expr, input, cx)));
}
}
out
}
pub(super) fn conflict_value_origins<'a>(
value: &'a Expr,
columns: &[ColumnWrite],
source: &'a LogicalPlan,
context: &mut TraceContext<'a>,
) -> Vec<(ColumnRead, ColumnLineageKind)> {
match value {
Expr::Column(c) if matches!(c.binding, Binding::Derived) => {
let Some(i) = columns
.iter()
.position(|t| context.eq_column(&t.reference.name, &c.name))
else {
return Vec::new();
};
let operands = output_operands(source);
if operands.is_empty() {
return vec![(excluded_source(c), ColumnLineageKind::Passthrough)];
}
trace_nth_output(&operands, i, context)
}
Expr::Column(_) => origins_of_expr(value, source, context),
Expr::Call { args } => transform(
args.iter()
.flat_map(|e| conflict_value_origins(e, columns, source, context)),
),
Expr::Case {
then, else_result, ..
} => {
let mut sources: Vec<_> = then
.iter()
.flat_map(|e| conflict_value_origins(e, columns, source, context))
.collect();
if let Some(e) = else_result {
sources.extend(conflict_value_origins(e, columns, source, context));
}
transform(sources)
}
Expr::Window { arg, .. } => {
transform(conflict_value_origins(arg, columns, source, context))
}
Expr::Subquery { plan, output } => {
transform(subquery_output_origins(plan, *output, context))
}
Expr::Fanin(refs) => refs
.iter()
.filter_map(|c| column_read(c).map(|r| (r, ColumnLineageKind::Passthrough)))
.collect(),
Expr::InSubquery { expr, .. } => {
transform(conflict_value_origins(expr, columns, source, context))
}
Expr::Exists(_) | Expr::Filter(_) => Vec::new(),
}
}
fn excluded_source(c: &BoundColumn) -> ColumnRead {
ColumnRead {
reference: ColumnReference {
table: c.qualifier.clone().map(|q| TableReference {
catalog: None,
schema: None,
name: q,
}),
name: c.name.clone(),
},
resolution: ResolutionKind::Inferred,
}
}
fn synthetic_source(table: &Ident, name: &Ident) -> ColumnRead {
ColumnRead {
reference: ColumnReference {
table: Some(TableReference {
catalog: None,
schema: None,
name: table.clone(),
}),
name: name.clone(),
},
resolution: ResolutionKind::Inferred,
}
}
pub(super) struct Operand<'a> {
pub(super) outputs: &'a [NamedExpr],
input: &'a LogicalPlan,
ctes: Vec<&'a Cte>,
}
impl<'a> Operand<'a> {
pub(super) fn trace<R>(
&self,
context: &mut TraceContext<'a>,
f: impl FnOnce(&'a LogicalPlan, &mut TraceContext<'a>) -> R,
) -> R {
context.with_cte_refs(&self.ctes, |context| f(self.input, context))
}
}
pub(super) fn output_operands(op: &LogicalPlan) -> Vec<Operand<'_>> {
let mut out = Vec::new();
collect_operands(op, &[], &mut out);
out
}
fn collect_operands<'a>(op: &'a LogicalPlan, ctes: &[&'a Cte], out: &mut Vec<Operand<'a>>) {
match op {
LogicalPlan::Projection(p) => out.push(Operand {
outputs: &p.exprs,
input: &p.input,
ctes: ctes.to_vec(),
}),
LogicalPlan::Sort(s) => collect_operands(&s.input, ctes, out),
LogicalPlan::Filter(f) => collect_operands(&f.input, ctes, out),
LogicalPlan::With(w) => {
let mut inner: Vec<&Cte> = ctes.to_vec();
inner.extend(w.ctes.iter());
collect_operands(&w.body, &inner, out);
}
LogicalPlan::SetOp(so) => {
collect_operands(&so.left, ctes, out);
collect_operands(&so.right, ctes, out);
}
LogicalPlan::Scan(_)
| LogicalPlan::Join(_)
| LogicalPlan::Aggregate(_)
| LogicalPlan::SubqueryAlias(_)
| LogicalPlan::TableFunction(_)
| LogicalPlan::CteRef(_)
| LogicalPlan::Values(_)
| LogicalPlan::Empty
| LogicalPlan::Insert(_)
| LogicalPlan::Update(_)
| LogicalPlan::Delete(_)
| LogicalPlan::Merge(_)
| LogicalPlan::CreateTableAs(_)
| LogicalPlan::CreateView(_)
| LogicalPlan::AlterTable(_)
| LogicalPlan::Drop(_) => {}
}
}
pub(super) fn enter_withs<'a>(
op: &'a LogicalPlan,
context: &mut TraceContext<'a>,
) -> &'a LogicalPlan {
let mut node = op;
while let LogicalPlan::With(w) = node {
w.ctes.iter().for_each(|c| context.ctes.push(c));
node = &w.body;
}
node
}
fn transform(
sources: impl IntoIterator<Item = (ColumnRead, ColumnLineageKind)>,
) -> Vec<(ColumnRead, ColumnLineageKind)> {
sources
.into_iter()
.map(|(r, _)| (r, ColumnLineageKind::Transformation))
.collect()
}
fn find_named<'a>(exprs: &'a [NamedExpr], name: &Ident, fold: CaseRule) -> Option<&'a NamedExpr> {
let target = fold.normalize(name);
exprs.iter().find(|ne| {
ne.name
.as_ref()
.is_some_and(|n| fold.normalize(n) == target)
})
}