Skip to content

Give each CTE its own node when a nested WITH reuses its name #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
IL-William merged 1 commit into from
Oct 10, 2026
Merged

Give each CTE its own node when a nested WITH reuses its name #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 10 additions & 2 deletions crates/flowscope-core/src/analyzer/context.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,14 @@ pub(crate) const DBT_MODEL_SINK_METADATA_KEY: &str = "dbtModelSink";
pub(crate) struct PendingWildcard {
/// The source canonical name (table or CTE) with unknown columns
pub(crate) source_canonical: String,
/// The target entity (CTE name or derived table alias) receiving the wildcard
pub(crate) target_name: String,
/// Node ID of the source
pub(crate) source_node_id: Arc<str>,
/// Node ID of the target. Two CTEs of one name in nested `WITH` blocks
/// share `target_name` and nothing else.
pub(crate) target_node_id: Arc<str>,
/// The node whose own wildcards feed the source: the CTE definition the
/// source names, else the source itself.
pub(crate) upstream_node_id: Arc<str>,
}

/// A single table/view reference in a FROM or JOIN clause.
Expand Down Expand Up @@ -144,6 +148,9 @@ pub(crate) struct StatementContext {
/// a derived table reusing the alias of an earlier one in another scope
/// gets a node of its own.
pub(crate) derived_occurrences: HashMap<String, usize>,
/// How many CTEs of each name the statement has defined, so that a nested
/// `WITH` defining a name an earlier CTE has gets a node of its own.
pub(crate) cte_occurrences: HashMap<String, usize>,
/// How many subqueries a predicate reads the statement has opened, each
/// analysed against a node of its own.
pub(crate) predicate_subqueries: usize,
Expand Down Expand Up @@ -243,6 +250,7 @@ impl StatementContext {
table_aliases: HashMap::new(),
subquery_aliases: HashSet::new(),
derived_occurrences: HashMap::new(),
cte_occurrences: HashMap::new(),
predicate_subqueries: 0,
last_operation: None,
current_join_info: JoinInfo::default(),
Expand Down
77 changes: 32 additions & 45 deletions crates/flowscope-core/src/analyzer/query.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -717,17 +717,19 @@ impl<'a> Analyzer<'a> {
});
}

// Find the CTE/alias name from the node ID for backward inference
// Use the cte_node_to_name reverse mapping for efficient lookup
let target_alias_name =
ctx.cte_node_to_name.get(&Arc::from(target)).cloned();

// Record pending wildcard for backward column inference
if let Some(alias_name) = target_alias_name {
// Record pending wildcard for backward column inference, when
// the target is a CTE or a derived table
if ctx.cte_node_to_name.contains_key(&Arc::from(target)) {
let upstream_node_id = ctx
.cte_definitions
.get(&table_canonical)
.cloned()
.unwrap_or_else(|| source_node_id.clone());
ctx.pending_wildcards.push(PendingWildcard {
source_canonical: table_canonical.clone(),
target_name: alias_name,
source_node_id,
target_node_id: Arc::from(target),
upstream_node_id,
});
}
}
Expand Down Expand Up @@ -1494,16 +1496,18 @@ impl<'a> Analyzer<'a> {
///
/// # Algorithm Overview
///
/// 1. **Group wildcards by target**: Wildcards are grouped by their `target_name`
/// (the CTE or derived table alias that receives the `SELECT *` columns).
/// 1. **Group wildcards by target**: Wildcards are grouped by their
/// `target_node_id` (the CTE or derived table node that receives the
/// `SELECT *` columns), not by name: two CTEs of one name in nested `WITH`
/// blocks are two targets.
///
/// 2. **Build node index**: Creates an O(1) lookup map from node ID to node
/// to avoid repeated linear scans when collecting owned columns.
///
/// 3. **Propagate columns**: For each target, finds its owned columns and
/// creates corresponding columns on the source tables. The `source_canonical`
/// field in `PendingWildcard` matches the `target_name` of upstream wildcards,
/// enabling recursive chain propagation.
/// creates corresponding columns on the source tables. The `upstream_node_id`
/// field in `PendingWildcard` matches the `target_node_id` of upstream
/// wildcards, enabling recursive chain propagation.
///
/// 4. **Cycle detection**: Uses `visited_pairs` to track (target, source) pairs
/// and prevent infinite recursion on cyclic references.
Expand All @@ -1512,11 +1516,13 @@ impl<'a> Analyzer<'a> {
return;
}

// Build map: target_name -> Vec<PendingWildcard>
let mut wildcards_by_target: HashMap<String, Vec<PendingWildcard>> = HashMap::new();
// Build map: target node -> Vec<PendingWildcard>. Keyed by node rather
// than by name, so that two CTEs of one name in nested `WITH` blocks
// infer their own columns.
let mut wildcards_by_target: HashMap<Arc<str>, Vec<PendingWildcard>> = HashMap::new();
for pw in ctx.pending_wildcards.drain(..) {
wildcards_by_target
.entry(pw.target_name.clone())
.entry(pw.target_node_id.clone())
.or_default()
.push(pw);
}
Expand All @@ -1531,22 +1537,18 @@ impl<'a> Analyzer<'a> {
.collect();

// Track visited target/source pairs to prevent cycles
let mut visited_pairs: HashSet<(String, String)> = HashSet::new();

for (target_name, wildcards) in &wildcards_by_target {
let Some(cte_node_id) = self.lookup_inference_target_node(ctx, target_name) else {
continue;
};
let mut visited_pairs: HashSet<(Arc<str>, Arc<str>)> = HashSet::new();

let owned_columns = self.collect_owned_columns(ctx, &cte_node_id, &node_index);
for (cte_node_id, wildcards) in &wildcards_by_target {
let owned_columns = self.collect_owned_columns(ctx, cte_node_id, &node_index);
if owned_columns.is_empty() {
continue;
}

for wildcard in wildcards {
self.propagate_wildcard_columns(
ctx,
target_name,
cte_node_id,
wildcard,
&owned_columns,
&wildcards_by_target,
Expand All @@ -1557,21 +1559,6 @@ impl<'a> Analyzer<'a> {
}
}

/// Locates the node ID to use as the inference target for a wildcard.
fn lookup_inference_target_node(
&self,
ctx: &StatementContext,
target_name: &str,
) -> Option<Arc<str>> {
if let Some(node_id) = ctx.cte_definitions.get(target_name) {
return Some(node_id.clone());
}

ctx.cte_node_to_name
.iter()
.find_map(|(node_id, name)| (name == target_name).then(|| node_id.clone()))
}

/// Collects column information owned by a CTE node.
///
/// Uses the provided `node_index` for O(1) node lookups instead of linear scans.
Expand Down Expand Up @@ -1607,19 +1594,19 @@ impl<'a> Analyzer<'a> {
fn propagate_wildcard_columns(
&mut self,
ctx: &mut StatementContext,
target_name: &str,
target_node_id: &Arc<str>,
wildcard: &PendingWildcard,
columns: &[(String, Option<String>, Arc<str>)],
wildcards_by_target: &HashMap<String, Vec<PendingWildcard>>,
visited_pairs: &mut HashSet<(String, String)>,
wildcards_by_target: &HashMap<Arc<str>, Vec<PendingWildcard>>,
visited_pairs: &mut HashSet<(Arc<str>, Arc<str>)>,
depth: usize,
) {
// Enforce recursion depth limit to prevent stack overflow
if depth >= Self::MAX_INFERENCE_DEPTH {
return;
}

let pair = (target_name.to_string(), wildcard.source_canonical.clone());
let pair = (target_node_id.clone(), wildcard.upstream_node_id.clone());
if !visited_pairs.insert(pair.clone()) {
return;
}
Expand All @@ -1639,11 +1626,11 @@ impl<'a> Analyzer<'a> {
.collect();

// Recursively propagate if this source is itself a target of another wildcard
if let Some(upstream_wildcards) = wildcards_by_target.get(&wildcard.source_canonical) {
if let Some(upstream_wildcards) = wildcards_by_target.get(&wildcard.upstream_node_id) {
for upstream_wildcard in upstream_wildcards {
self.propagate_wildcard_columns(
ctx,
&wildcard.source_canonical,
&wildcard.upstream_node_id,
upstream_wildcard,
&source_columns,
wildcards_by_target,
Expand Down
49 changes: 41 additions & 8 deletions crates/flowscope-core/src/analyzer/visitor.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
//! and building lineage graphs. It separates traversal logic (the `Visitor` trait)
//! from analysis logic (the `LineageVisitor` implementation).

use super::context::StatementContext;
use super::context::{OutputColumn, StatementContext};
use super::expression::ExpressionAnalyzer;
use super::helpers::{
alias_visibility_warning, find_cte_body_span, find_cte_definition_span,
Expand Down Expand Up @@ -622,10 +622,28 @@ impl<'a, 'b> LineageVisitor<'a, 'b> {

impl<'a, 'b> Visitor for LineageVisitor<'a, 'b> {
fn visit_query(&mut self, query: &Query) {
// A name this WITH defines hides, until the query ends, the CTE an
// enclosing or earlier WITH gave it: its node and its columns, put back
// on the way out. Keyed by the name alone, the second definition would
// take the first's node and overwrite its columns, and every reader of
// either would get both lineages.
let mut hidden: Vec<(String, Arc<str>, Option<Vec<OutputColumn>>)> = Vec::new();
if let Some(with) = &query.with {
let mut cte_ids: Vec<(String, Arc<str>)> = Vec::new();
for cte in &with.cte_tables {
let cte_name = cte.alias.name.to_string();
let occurrence = self
.ctx
.cte_occurrences
.entry(cte_name.clone())
.or_insert(0);
*occurrence += 1;
// The first keeps the key it always had, so a statement that
// never repeats a name is analysed as before.
let key = match *occurrence {
1 => cte_name.clone(),
n => format!("{cte_name}#{n}"),
};
let cte_span = self.locate_cte_definition_span(&cte_name);
let body_span = cte_span.and_then(|span| {
let sql = if let Some(source) = &self.analyzer.current_statement_source {
Expand All @@ -636,11 +654,7 @@ impl<'a, 'b> Visitor for LineageVisitor<'a, 'b> {
find_cte_body_span(sql, span)
});
let cte_id = self.ctx.add_node(Node {
id: generate_statement_scoped_node_id(
"cte",
self.ctx.statement_index,
&cte_name,
),
id: generate_statement_scoped_node_id("cte", self.ctx.statement_index, &key),
node_type: NodeType::Cte,
label: cte_name.clone().into(),
qualified_name: Some(cte_name.clone().into()),
Expand All @@ -654,9 +668,16 @@ impl<'a, 'b> Visitor for LineageVisitor<'a, 'b> {
*cursor = (*cursor).max(span.end);
}

self.ctx
if let Some(previous) = self
.ctx
.cte_definitions
.insert(cte_name.clone(), cte_id.clone());
.insert(cte_name.clone(), cte_id.clone())
{
if previous != cte_id {
let columns = self.ctx.aliased_subquery_columns.remove(&cte_name);
hidden.push((cte_name.clone(), previous, columns));
}
}
self.ctx
.cte_node_to_name
.insert(cte_id.clone(), cte_name.clone());
Expand All @@ -680,6 +701,18 @@ impl<'a, 'b> Visitor for LineageVisitor<'a, 'b> {
if let Some(order_by) = &query.order_by {
self.analyze_order_by(order_by);
}

for (name, previous, columns) in hidden.into_iter().rev() {
self.ctx.cte_definitions.insert(name.clone(), previous);
match columns {
Some(columns) => {
self.ctx.aliased_subquery_columns.insert(name, columns);
}
None => {
self.ctx.aliased_subquery_columns.remove(&name);
}
}
}
}

fn visit_set_expr(&mut self, set_expr: &SetExpr) {
Expand Down
Loading
Loading