Skip to main content

polyglot_sql/
query_analysis.rs

1//! Compact query analysis facts.
2//!
3//! This module intentionally builds on the existing parser, scope builder, type
4//! annotator, and lineage implementation. It is a convenience API: callers that
5//! need the full AST or full lineage graph should continue using those lower
6//! level APIs directly.
7
8use crate::ast_transforms::get_output_column_names_for_dialect;
9use crate::dialects::{Dialect, DialectType};
10use crate::expressions::{DataType, Expression, JoinKind, TableRef, With};
11use crate::lineage::{lineage_by_index_from_expression, LineageNode};
12use crate::optimizer::annotate_types::annotate_types;
13use crate::optimizer::qualify_schema_aware_expression;
14use crate::schema::{MappingSchema, Schema};
15use crate::scope::{build_scope, Scope, SourceInfo, SourceKind};
16use crate::traversal::{contains_aggregate, ExpressionWalk};
17use crate::validation::{mapping_schema_from_validation_schema_with_dialect, ValidationSchema};
18use crate::{parse_one, Error, Result};
19use serde::{Deserialize, Serialize};
20use std::collections::{HashMap, HashSet};
21
22/// Options for [`analyze_query`].
23#[derive(Debug, Clone, Serialize, Deserialize, Default)]
24#[serde(rename_all = "camelCase", default)]
25pub struct AnalyzeQueryOptions {
26    /// SQL dialect used for parsing and dialect-aware rendering.
27    pub dialect: DialectType,
28    /// Optional validation schema used for qualification and type annotation.
29    pub schema: Option<ValidationSchema>,
30}
31
32/// Compact facts about a query's output shape and data dependencies.
33#[derive(Debug, Clone, Serialize, Deserialize)]
34#[serde(rename_all = "camelCase")]
35pub struct QueryAnalysis {
36    pub shape: QueryShape,
37    pub ctes: Vec<String>,
38    pub cte_facts: Vec<CteFact>,
39    pub projections: Vec<ProjectionFact>,
40    pub relations: Vec<RelationFact>,
41    pub base_tables: Vec<RelationFact>,
42    pub star_projections: Vec<StarProjectionFact>,
43    pub set_operations: Vec<SetOperationFact>,
44}
45
46/// Top-level query shape.
47#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
48#[serde(rename_all = "snake_case")]
49pub enum QueryShape {
50    Select,
51    SetOperation,
52}
53
54/// Compact fact about one output projection.
55#[derive(Debug, Clone, Serialize, Deserialize)]
56#[serde(rename_all = "camelCase")]
57pub struct ProjectionFact {
58    pub index: usize,
59    pub name: Option<String>,
60    pub is_star: bool,
61    pub star_table: Option<String>,
62    pub transform_kind: TransformKind,
63    #[serde(skip_serializing_if = "Option::is_none")]
64    pub transform_function: Option<TransformFunctionFact>,
65    pub cast_type: Option<String>,
66    pub type_hint: Option<String>,
67    pub nullability: ProjectionNullability,
68    pub upstream: Vec<ColumnReferenceFact>,
69}
70
71/// Compact fact about a function-like projection transform.
72#[derive(Debug, Clone, Serialize, Deserialize)]
73#[serde(rename_all = "camelCase")]
74pub struct TransformFunctionFact {
75    pub name: String,
76    pub literal_args: Vec<String>,
77    pub column_args: Vec<ColumnReferenceFact>,
78}
79
80/// Compact fact about one top-level CTE definition.
81#[derive(Debug, Clone, Serialize, Deserialize)]
82#[serde(rename_all = "camelCase")]
83pub struct CteFact {
84    pub name: String,
85    pub columns: Vec<String>,
86    pub body_sql: String,
87    pub output_columns: Vec<String>,
88}
89
90/// Compact fact about one original star projection.
91#[derive(Debug, Clone, Serialize, Deserialize)]
92#[serde(rename_all = "camelCase")]
93pub struct StarProjectionFact {
94    pub index: usize,
95    pub table: Option<String>,
96    pub expanded_columns: Vec<String>,
97}
98
99/// Compact fact about an upstream column reference.
100#[derive(Debug, Clone, Serialize, Deserialize)]
101#[serde(rename_all = "camelCase")]
102pub struct ColumnReferenceFact {
103    pub source_name: Option<String>,
104    pub source_alias: Option<String>,
105    pub source_kind: SourceKind,
106    pub table: Option<String>,
107    pub column: String,
108    pub unqualified: bool,
109    pub confidence: ReferenceConfidence,
110}
111
112/// Compact fact about a relation visible in the root scope.
113#[derive(Debug, Clone, Serialize, Deserialize)]
114#[serde(rename_all = "camelCase")]
115pub struct RelationFact {
116    pub name: String,
117    pub alias: Option<String>,
118    pub kind: SourceKind,
119    pub columns: Vec<String>,
120    pub catalog: Option<String>,
121    pub schema: Option<String>,
122    pub table: Option<String>,
123}
124
125/// Compact fact about a set operation.
126#[derive(Debug, Clone, Serialize, Deserialize)]
127#[serde(rename_all = "camelCase")]
128pub struct SetOperationFact {
129    pub kind: String,
130    pub all: bool,
131    pub distinct: bool,
132    pub output_columns: Vec<String>,
133    pub branches: Vec<SetOperationBranchFact>,
134}
135
136/// Compact facts for one immediate set-operation branch.
137#[derive(Debug, Clone, Serialize, Deserialize)]
138#[serde(rename_all = "camelCase")]
139pub struct SetOperationBranchFact {
140    pub index: usize,
141    pub role: SetOperationBranchRole,
142    pub projections: Vec<ProjectionFact>,
143}
144
145/// Semantic contribution of one set-operation branch.
146#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
147#[serde(rename_all = "snake_case")]
148pub enum SetOperationBranchRole {
149    Value,
150    Filter,
151}
152
153/// High-level kind of transformation performed by a projection.
154#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
155#[serde(rename_all = "snake_case")]
156pub enum TransformKind {
157    Direct,
158    Cast,
159    Aggregation,
160    Constant,
161    Expression,
162    Star,
163}
164
165/// Confidence level for a compact upstream column reference.
166#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
167#[serde(rename_all = "snake_case")]
168pub enum ReferenceConfidence {
169    Resolved,
170    Ambiguous,
171    Unknown,
172}
173
174/// Conservative nullability classification for one output projection.
175#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
176#[serde(rename_all = "snake_case")]
177pub enum ProjectionNullability {
178    NonNull,
179    Nullable,
180    Unknown,
181}
182
183/// Analyze a single SELECT or set-operation query.
184pub fn analyze_query(sql: &str, options: AnalyzeQueryOptions) -> Result<QueryAnalysis> {
185    let mut expression = parse_one(sql, options.dialect)?;
186    expression = effective_query(expression);
187    ensure_query(&expression)?;
188    let original_expression = expression.clone();
189
190    let mapping_schema = options
191        .schema
192        .as_ref()
193        .map(|schema| analysis_mapping_schema(schema, options.dialect));
194    let schema_info = options.schema.as_ref().map(AnalysisSchemaInfo::from_schema);
195    let cte_facts = top_level_cte_facts(&original_expression, options.dialect)?;
196    let star_projections = star_projection_facts(
197        &original_expression,
198        mapping_schema.as_ref(),
199        options.dialect,
200    );
201
202    if let Some(schema) = mapping_schema.as_ref() {
203        expression = qualify_schema_aware_expression(expression, schema, Some(options.dialect))
204            .map_err(|e| Error::internal(format!("query analysis qualification failed: {e}")))?;
205    }
206
207    annotate_types(
208        &mut expression,
209        mapping_schema.as_ref().map(|schema| schema as &dyn Schema),
210        Some(options.dialect),
211    );
212    crate::lineage::expand_cte_stars(
213        &mut expression,
214        mapping_schema.as_ref().map(|schema| schema as &dyn Schema),
215    );
216
217    let scope = build_scope(&expression);
218    let nullability_context = NullabilityContext {
219        schema: schema_info.as_ref(),
220        nullable_sources: nullable_source_names(&expression),
221    };
222    let shape = if is_set_operation(&expression) {
223        QueryShape::SetOperation
224    } else {
225        QueryShape::Select
226    };
227
228    Ok(QueryAnalysis {
229        shape,
230        ctes: collect_cte_names(&expression),
231        cte_facts,
232        projections: projection_facts_for_query(
233            &expression,
234            &scope,
235            options.dialect,
236            &nullability_context,
237        ),
238        relations: relation_facts(&scope, mapping_schema.as_ref(), options.dialect),
239        base_tables: base_table_facts(&scope, mapping_schema.as_ref(), options.dialect),
240        star_projections,
241        set_operations: set_operation_facts(&expression, &scope, options.dialect),
242    })
243}
244
245fn analysis_mapping_schema(schema: &ValidationSchema, dialect: DialectType) -> MappingSchema {
246    mapping_schema_from_validation_schema_with_dialect(schema, dialect)
247}
248
249fn validation_table_names(table: &crate::validation::SchemaTable) -> Vec<String> {
250    let mut names = Vec::new();
251
252    names.push(table.name.to_ascii_lowercase());
253    if let Some(schema_name) = &table.schema {
254        names.push(format!(
255            "{}.{}",
256            schema_name.to_ascii_lowercase(),
257            table.name.to_ascii_lowercase()
258        ));
259    }
260    for alias in &table.aliases {
261        names.push(alias.to_ascii_lowercase());
262    }
263
264    names.sort();
265    names.dedup();
266    names
267}
268
269#[derive(Debug, Clone)]
270struct AnalysisColumnInfo {
271    nullable: Option<bool>,
272    primary_key: bool,
273}
274
275#[derive(Debug, Clone)]
276struct AnalysisSchemaInfo {
277    columns: HashMap<(String, String), AnalysisColumnInfo>,
278}
279
280impl AnalysisSchemaInfo {
281    fn from_schema(schema: &ValidationSchema) -> Self {
282        let mut columns = HashMap::new();
283
284        for table in &schema.tables {
285            let table_names = validation_table_names(table);
286            let primary_keys: HashSet<String> = table
287                .primary_key
288                .iter()
289                .map(|column| column.to_ascii_lowercase())
290                .collect();
291
292            for column in &table.columns {
293                let info = AnalysisColumnInfo {
294                    nullable: column.nullable,
295                    primary_key: column.primary_key
296                        || primary_keys.contains(&column.name.to_ascii_lowercase()),
297                };
298
299                for table_name in &table_names {
300                    columns.insert(
301                        (
302                            normalize_lookup_name(table_name),
303                            normalize_lookup_name(&column.name),
304                        ),
305                        info.clone(),
306                    );
307                }
308            }
309        }
310
311        Self { columns }
312    }
313
314    fn column(&self, table: &str, column: &str) -> Option<&AnalysisColumnInfo> {
315        self.columns
316            .get(&(normalize_lookup_name(table), normalize_lookup_name(column)))
317    }
318}
319
320struct NullabilityContext<'a> {
321    schema: Option<&'a AnalysisSchemaInfo>,
322    nullable_sources: HashSet<String>,
323}
324
325fn top_level_cte_facts(expression: &Expression, dialect: DialectType) -> Result<Vec<CteFact>> {
326    let Some(with_clause) = with_clause(expression) else {
327        return Ok(Vec::new());
328    };
329
330    with_clause
331        .ctes
332        .iter()
333        .map(|cte| {
334            Ok(CteFact {
335                name: cte.alias.name.clone(),
336                columns: cte
337                    .columns
338                    .iter()
339                    .map(|column| column.name.clone())
340                    .collect(),
341                body_sql: Dialect::get(dialect).generate(&cte.this)?,
342                output_columns: get_output_column_names_for_dialect(&cte.this, Some(dialect)),
343            })
344        })
345        .collect()
346}
347
348fn star_projection_facts(
349    expression: &Expression,
350    mapping_schema: Option<&MappingSchema>,
351    dialect: DialectType,
352) -> Vec<StarProjectionFact> {
353    let scope = build_scope(expression);
354    let ordered_sources = ordered_source_names_for_query(expression);
355
356    select_expressions_for_query(expression)
357        .iter()
358        .enumerate()
359        .filter_map(|(index, projection)| {
360            let inner = unwrap_projection_alias(projection);
361            if !projection_is_star(inner) {
362                return None;
363            }
364
365            let table = projection_star_table(inner);
366            let expanded_columns = expanded_star_columns(
367                table.as_deref(),
368                &scope,
369                &ordered_sources,
370                mapping_schema,
371                dialect,
372            );
373
374            Some(StarProjectionFact {
375                index,
376                table,
377                expanded_columns,
378            })
379        })
380        .collect()
381}
382
383fn expanded_star_columns(
384    star_table: Option<&str>,
385    scope: &Scope,
386    ordered_sources: &[String],
387    mapping_schema: Option<&MappingSchema>,
388    dialect: DialectType,
389) -> Vec<String> {
390    let mut columns = Vec::new();
391    let mut source_names: Vec<String> = if ordered_sources.is_empty() {
392        let mut names: Vec<_> = scope.sources.keys().cloned().collect();
393        names.sort();
394        names
395    } else {
396        ordered_sources.to_vec()
397    };
398
399    source_names.dedup();
400
401    for source_name in source_names {
402        let Some(source) = scope.sources.get(&source_name) else {
403            continue;
404        };
405
406        if let Some(star_table) = star_table {
407            let matches = source_name.eq_ignore_ascii_case(star_table)
408                || source
409                    .alias
410                    .as_deref()
411                    .is_some_and(|alias| alias.eq_ignore_ascii_case(star_table))
412                || source_table_name(source)
413                    .is_some_and(|table| table.eq_ignore_ascii_case(star_table));
414
415            if !matches {
416                continue;
417            }
418        }
419
420        columns.extend(source_columns(source, mapping_schema, dialect));
421    }
422
423    columns
424}
425
426fn ordered_source_names_for_query(expression: &Expression) -> Vec<String> {
427    match expression {
428        Expression::Select(select) => ordered_source_names_for_select(select),
429        Expression::Union(union) => ordered_source_names_for_query(&union.left),
430        Expression::Intersect(intersect) => ordered_source_names_for_query(&intersect.left),
431        Expression::Except(except) => ordered_source_names_for_query(&except.left),
432        Expression::Subquery(subquery) => ordered_source_names_for_query(&subquery.this),
433        _ => Vec::new(),
434    }
435}
436
437fn ordered_source_names_for_select(select: &crate::expressions::Select) -> Vec<String> {
438    let mut sources = Vec::new();
439
440    if let Some(from) = &select.from {
441        for expression in &from.expressions {
442            if let Some(source_name) = expression_source_name(expression) {
443                sources.push(source_name);
444            }
445        }
446    }
447
448    for join in &select.joins {
449        if let Some(source_name) = expression_source_name(&join.this) {
450            sources.push(source_name);
451        }
452    }
453
454    sources
455}
456
457fn nullable_source_names(expression: &Expression) -> HashSet<String> {
458    match expression {
459        Expression::Select(select) => nullable_source_names_for_select(select),
460        Expression::Union(union) => nullable_source_names(&union.left),
461        Expression::Intersect(intersect) => nullable_source_names(&intersect.left),
462        Expression::Except(except) => nullable_source_names(&except.left),
463        Expression::Subquery(subquery) => nullable_source_names(&subquery.this),
464        _ => HashSet::new(),
465    }
466}
467
468fn nullable_source_names_for_select(select: &crate::expressions::Select) -> HashSet<String> {
469    let mut nullable = HashSet::new();
470    let mut left_sources = Vec::new();
471
472    if let Some(from) = &select.from {
473        for expression in &from.expressions {
474            if let Some(source_name) = expression_source_name(expression) {
475                left_sources.push(source_name);
476            }
477        }
478    }
479
480    for join in &select.joins {
481        let right_source = expression_source_name(&join.this);
482
483        if join_nullable_left(join.kind) {
484            for source_name in &left_sources {
485                nullable.insert(normalize_lookup_name(source_name));
486            }
487        }
488
489        if join_nullable_right(join.kind) {
490            if let Some(source_name) = &right_source {
491                nullable.insert(normalize_lookup_name(source_name));
492            }
493        }
494
495        if let Some(source_name) = right_source {
496            left_sources.push(source_name);
497        }
498    }
499
500    nullable
501}
502
503fn join_nullable_left(kind: JoinKind) -> bool {
504    matches!(
505        kind,
506        JoinKind::Right
507            | JoinKind::NaturalRight
508            | JoinKind::AsOfRight
509            | JoinKind::Full
510            | JoinKind::NaturalFull
511            | JoinKind::Outer
512    )
513}
514
515fn join_nullable_right(kind: JoinKind) -> bool {
516    matches!(
517        kind,
518        JoinKind::Left
519            | JoinKind::NaturalLeft
520            | JoinKind::AsOfLeft
521            | JoinKind::LeftLateral
522            | JoinKind::OuterApply
523            | JoinKind::LeftArray
524            | JoinKind::Full
525            | JoinKind::NaturalFull
526            | JoinKind::Outer
527    )
528}
529
530fn expression_source_name(expression: &Expression) -> Option<String> {
531    match expression {
532        Expression::Table(table) => table
533            .alias
534            .as_ref()
535            .map(|alias| alias.name.clone())
536            .or_else(|| Some(table.name.name.clone())),
537        Expression::Subquery(subquery) => subquery.alias.as_ref().map(|alias| alias.name.clone()),
538        Expression::Alias(alias) => Some(alias.alias.name.clone()),
539        Expression::Cte(cte) => Some(cte.alias.name.clone()),
540        _ => None,
541    }
542}
543
544fn normalize_lookup_name(name: &str) -> String {
545    name.to_ascii_lowercase()
546}
547
548fn effective_query(expression: Expression) -> Expression {
549    match expression {
550        Expression::Prepare(prepare) => prepare.statement,
551        Expression::Subquery(subquery) if subquery.alias.is_none() => subquery.this,
552        other => other,
553    }
554}
555
556fn ensure_query(expression: &Expression) -> Result<()> {
557    if matches!(
558        expression,
559        Expression::Select(_)
560            | Expression::Union(_)
561            | Expression::Intersect(_)
562            | Expression::Except(_)
563    ) {
564        Ok(())
565    } else {
566        Err(Error::internal(
567            "analyze_query requires a SELECT or set operation query",
568        ))
569    }
570}
571
572fn is_set_operation(expression: &Expression) -> bool {
573    matches!(
574        expression,
575        Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_)
576    )
577}
578
579fn collect_cte_names(expression: &Expression) -> Vec<String> {
580    let mut names = Vec::new();
581    let mut seen = HashSet::new();
582    collect_cte_names_inner(expression, &mut names, &mut seen);
583    names
584}
585
586fn collect_cte_names_inner(
587    expression: &Expression,
588    names: &mut Vec<String>,
589    seen: &mut HashSet<String>,
590) {
591    if let Some(with_clause) = with_clause(expression) {
592        collect_with_names(with_clause, names, seen);
593    }
594
595    match expression {
596        Expression::Union(union) => {
597            collect_cte_names_inner(&union.left, names, seen);
598            collect_cte_names_inner(&union.right, names, seen);
599        }
600        Expression::Intersect(intersect) => {
601            collect_cte_names_inner(&intersect.left, names, seen);
602            collect_cte_names_inner(&intersect.right, names, seen);
603        }
604        Expression::Except(except) => {
605            collect_cte_names_inner(&except.left, names, seen);
606            collect_cte_names_inner(&except.right, names, seen);
607        }
608        Expression::Subquery(subquery) => collect_cte_names_inner(&subquery.this, names, seen),
609        _ => {}
610    }
611}
612
613fn collect_with_names(with_clause: &With, names: &mut Vec<String>, seen: &mut HashSet<String>) {
614    for cte in &with_clause.ctes {
615        if seen.insert(cte.alias.name.clone()) {
616            names.push(cte.alias.name.clone());
617        }
618        collect_cte_names_inner(&cte.this, names, seen);
619    }
620}
621
622fn with_clause(expression: &Expression) -> Option<&With> {
623    match expression {
624        Expression::Select(select) => select.with.as_ref(),
625        Expression::Union(union) => union.with.as_ref(),
626        Expression::Intersect(intersect) => intersect.with.as_ref(),
627        Expression::Except(except) => except.with.as_ref(),
628        _ => None,
629    }
630}
631
632fn projection_facts_for_query(
633    expression: &Expression,
634    scope: &Scope,
635    dialect: DialectType,
636    nullability_context: &NullabilityContext<'_>,
637) -> Vec<ProjectionFact> {
638    let expressions = projection_sources_for_query(expression, dialect);
639    let names = get_output_column_names_for_dialect(expression, Some(dialect));
640
641    expressions
642        .iter()
643        .enumerate()
644        .map(|(index, (projection, null_padded))| {
645            let mut fact = projection_fact(
646                index,
647                names
648                    .get(index)
649                    .cloned()
650                    .or_else(|| projection_name(projection)),
651                projection,
652                expression,
653                scope,
654                dialect,
655                nullability_context,
656            );
657            if *null_padded {
658                fact.nullability = ProjectionNullability::Nullable;
659            }
660            fact
661        })
662        .collect()
663}
664
665/// Return one representative projection for each result ordinal together with
666/// whether any immediate name-aligned branch contributes a synthetic NULL.
667fn projection_sources_for_query(
668    expression: &Expression,
669    dialect: DialectType,
670) -> Vec<(&Expression, bool)> {
671    match crate::set_operation::set_operation_layout(expression, Some(dialect)) {
672        Ok(Some(layout)) => layout
673            .outputs
674            .iter()
675            .filter_map(|output| {
676                let null_padded = output.left_ordinal.is_none() || output.right_ordinal.is_none();
677                output
678                    .left_ordinal
679                    .and_then(|ordinal| {
680                        projection_source_for_ordinal(
681                            set_operation_left(expression)?,
682                            ordinal,
683                            dialect,
684                        )
685                    })
686                    .or_else(|| {
687                        output.right_ordinal.and_then(|ordinal| {
688                            projection_source_for_ordinal(
689                                set_operation_right(expression)?,
690                                ordinal,
691                                dialect,
692                            )
693                        })
694                    })
695                    .map(|(projection, nested_null_padded)| {
696                        (projection, null_padded || nested_null_padded)
697                    })
698            })
699            .collect(),
700        _ => select_expressions_for_query(expression)
701            .into_iter()
702            .map(|projection| (projection, false))
703            .collect(),
704    }
705}
706
707fn projection_source_for_ordinal(
708    expression: &Expression,
709    ordinal: usize,
710    dialect: DialectType,
711) -> Option<(&Expression, bool)> {
712    match crate::set_operation::set_operation_layout(expression, Some(dialect)) {
713        Ok(Some(layout)) => {
714            let output = layout.outputs.get(ordinal)?;
715            let null_padded = output.left_ordinal.is_none() || output.right_ordinal.is_none();
716            output
717                .left_ordinal
718                .and_then(|child_ordinal| {
719                    projection_source_for_ordinal(
720                        set_operation_left(expression)?,
721                        child_ordinal,
722                        dialect,
723                    )
724                })
725                .or_else(|| {
726                    output.right_ordinal.and_then(|child_ordinal| {
727                        projection_source_for_ordinal(
728                            set_operation_right(expression)?,
729                            child_ordinal,
730                            dialect,
731                        )
732                    })
733                })
734                .map(|(projection, nested_null_padded)| {
735                    (projection, null_padded || nested_null_padded)
736                })
737        }
738        _ => match expression {
739            Expression::Select(select) => select
740                .expressions
741                .get(ordinal)
742                .map(|projection| (projection, false)),
743            Expression::Union(union) => {
744                projection_source_for_ordinal(&union.left, ordinal, dialect)
745            }
746            Expression::Intersect(intersect) => {
747                projection_source_for_ordinal(&intersect.left, ordinal, dialect)
748            }
749            Expression::Except(except) => {
750                projection_source_for_ordinal(&except.left, ordinal, dialect)
751            }
752            Expression::Subquery(subquery) => {
753                projection_source_for_ordinal(&subquery.this, ordinal, dialect)
754            }
755            Expression::Paren(paren) => {
756                projection_source_for_ordinal(&paren.this, ordinal, dialect)
757            }
758            _ => None,
759        },
760    }
761}
762
763fn set_operation_left(expression: &Expression) -> Option<&Expression> {
764    match expression {
765        Expression::Union(set_op) => Some(&set_op.left),
766        Expression::Intersect(set_op) => Some(&set_op.left),
767        Expression::Except(set_op) => Some(&set_op.left),
768        _ => None,
769    }
770}
771
772fn set_operation_right(expression: &Expression) -> Option<&Expression> {
773    match expression {
774        Expression::Union(set_op) => Some(&set_op.right),
775        Expression::Intersect(set_op) => Some(&set_op.right),
776        Expression::Except(set_op) => Some(&set_op.right),
777        _ => None,
778    }
779}
780
781fn select_expressions_for_query(expression: &Expression) -> Vec<&Expression> {
782    match expression {
783        Expression::Select(select) => select.expressions.iter().collect(),
784        Expression::Union(union) => select_expressions_for_query(&union.left),
785        Expression::Intersect(intersect) => select_expressions_for_query(&intersect.left),
786        Expression::Except(except) => select_expressions_for_query(&except.left),
787        Expression::Subquery(subquery) => select_expressions_for_query(&subquery.this),
788        _ => Vec::new(),
789    }
790}
791
792fn projection_fact(
793    index: usize,
794    name: Option<String>,
795    projection: &Expression,
796    query: &Expression,
797    scope: &Scope,
798    dialect: DialectType,
799    nullability_context: &NullabilityContext<'_>,
800) -> ProjectionFact {
801    let inner = unwrap_projection_alias(projection);
802    let is_star = projection_is_star(inner);
803    let upstream = lineage_by_index_from_expression(index, query, Some(dialect), false)
804        .map(|node| terminal_references_from_lineage(&node))
805        .ok()
806        .filter(|refs| !refs.is_empty())
807        .unwrap_or_else(|| fallback_column_references(inner, scope));
808
809    ProjectionFact {
810        index,
811        name,
812        is_star,
813        star_table: projection_star_table(inner),
814        transform_kind: transform_kind(inner),
815        transform_function: transform_function_fact(inner, scope, dialect),
816        cast_type: cast_type(inner, dialect),
817        type_hint: projection
818            .inferred_type()
819            .or_else(|| inner.inferred_type())
820            .and_then(|data_type| render_data_type(data_type, dialect)),
821        nullability: projection_nullability(inner, scope, nullability_context),
822        upstream,
823    }
824}
825
826fn transform_function_fact(
827    expression: &Expression,
828    scope: &Scope,
829    dialect: DialectType,
830) -> Option<TransformFunctionFact> {
831    let mut matches = expression
832        .find_all(|candidate| transform_function_fact_for_node(candidate, scope, dialect).is_some())
833        .into_iter();
834
835    let first = matches.next()?;
836    if matches.next().is_some() {
837        return None;
838    }
839
840    transform_function_fact_for_node(first, scope, dialect)
841}
842
843fn transform_function_fact_for_node(
844    expression: &Expression,
845    scope: &Scope,
846    dialect: DialectType,
847) -> Option<TransformFunctionFact> {
848    match expression {
849        Expression::Function(function) => Some(transform_function_from_args(
850            &function.name,
851            &function.args,
852            scope,
853            dialect,
854        )),
855        Expression::AggregateFunction(function) => Some(transform_function_from_args(
856            &function.name,
857            &function.args,
858            scope,
859            dialect,
860        )),
861        Expression::DateTrunc(function) => Some(transform_function_from_parts(
862            "DATE_TRUNC",
863            vec![datetime_field_name(&function.unit)],
864            vec![&function.this],
865            scope,
866            dialect,
867        )),
868        Expression::TimestampTrunc(function) => Some(transform_function_from_parts(
869            "TIMESTAMP_TRUNC",
870            vec![datetime_field_name(&function.unit)],
871            vec![&function.this],
872            scope,
873            dialect,
874        )),
875        Expression::TimeTrunc(function) => {
876            let mut args = vec![function.this.as_ref()];
877            if let Some(zone) = function.zone.as_deref() {
878                args.push(zone);
879            }
880            Some(transform_function_from_parts(
881                "TIME_TRUNC",
882                vec![function.unit.clone()],
883                args,
884                scope,
885                dialect,
886            ))
887        }
888        Expression::Extract(function) => Some(transform_function_from_parts(
889            "EXTRACT",
890            vec![datetime_field_name(&function.field)],
891            vec![&function.this],
892            scope,
893            dialect,
894        )),
895        Expression::DateAdd(function) => Some(transform_function_from_parts(
896            "DATE_ADD",
897            Vec::new(),
898            vec![&function.this, &function.interval],
899            scope,
900            dialect,
901        )),
902        Expression::DateSub(function) => Some(transform_function_from_parts(
903            "DATE_SUB",
904            Vec::new(),
905            vec![&function.this, &function.interval],
906            scope,
907            dialect,
908        )),
909        Expression::DateDiff(function) => Some(transform_function_from_parts(
910            "DATE_DIFF",
911            Vec::new(),
912            vec![&function.this, &function.expression],
913            scope,
914            dialect,
915        )),
916        _ => None,
917    }
918}
919
920fn transform_function_from_args(
921    name: &str,
922    args: &[Expression],
923    scope: &Scope,
924    dialect: DialectType,
925) -> TransformFunctionFact {
926    let literal_args = args
927        .iter()
928        .filter_map(|arg| literal_argument(arg, dialect))
929        .collect();
930    transform_function_from_parts(name, literal_args, args.iter().collect(), scope, dialect)
931}
932
933fn transform_function_from_parts(
934    name: &str,
935    literal_args: Vec<String>,
936    args: Vec<&Expression>,
937    scope: &Scope,
938    _dialect: DialectType,
939) -> TransformFunctionFact {
940    let column_args = dedupe_column_refs(
941        args.into_iter()
942            .flat_map(|arg| fallback_column_references(arg, scope))
943            .collect(),
944    );
945
946    TransformFunctionFact {
947        name: name.to_string(),
948        literal_args,
949        column_args,
950    }
951}
952
953fn literal_argument(expression: &Expression, dialect: DialectType) -> Option<String> {
954    match expression {
955        Expression::Literal(literal) => Some(literal.value_str().to_string()),
956        Expression::Boolean(boolean) => Some(boolean.value.to_string()),
957        Expression::Null(_) => Some("NULL".to_string()),
958        Expression::Identifier(identifier) => Some(identifier.name.clone()),
959        Expression::Var(var) => Some(var.this.clone()),
960        Expression::DataType(data_type) => render_data_type(data_type, dialect),
961        _ => None,
962    }
963}
964
965fn datetime_field_name(field: &crate::expressions::DateTimeField) -> String {
966    match field {
967        crate::expressions::DateTimeField::Year => "year".to_string(),
968        crate::expressions::DateTimeField::Month => "month".to_string(),
969        crate::expressions::DateTimeField::Day => "day".to_string(),
970        crate::expressions::DateTimeField::Hour => "hour".to_string(),
971        crate::expressions::DateTimeField::Minute => "minute".to_string(),
972        crate::expressions::DateTimeField::Second => "second".to_string(),
973        crate::expressions::DateTimeField::Millisecond => "millisecond".to_string(),
974        crate::expressions::DateTimeField::Microsecond => "microsecond".to_string(),
975        crate::expressions::DateTimeField::DayOfWeek => "day_of_week".to_string(),
976        crate::expressions::DateTimeField::DayOfYear => "day_of_year".to_string(),
977        crate::expressions::DateTimeField::Week => "week".to_string(),
978        crate::expressions::DateTimeField::WeekWithModifier(modifier) => {
979            format!("week({modifier})")
980        }
981        crate::expressions::DateTimeField::Quarter => "quarter".to_string(),
982        crate::expressions::DateTimeField::Epoch => "epoch".to_string(),
983        crate::expressions::DateTimeField::Timezone => "timezone".to_string(),
984        crate::expressions::DateTimeField::TimezoneHour => "timezone_hour".to_string(),
985        crate::expressions::DateTimeField::TimezoneMinute => "timezone_minute".to_string(),
986        crate::expressions::DateTimeField::Date => "date".to_string(),
987        crate::expressions::DateTimeField::Time => "time".to_string(),
988        crate::expressions::DateTimeField::Custom(name) => name.clone(),
989    }
990}
991
992fn unwrap_projection_alias(expression: &Expression) -> &Expression {
993    match expression {
994        Expression::Alias(alias) => unwrap_projection_alias(&alias.this),
995        Expression::Annotated(annotated) => unwrap_projection_alias(&annotated.this),
996        Expression::Paren(paren) => unwrap_projection_alias(&paren.this),
997        _ => expression,
998    }
999}
1000
1001fn projection_name(expression: &Expression) -> Option<String> {
1002    match expression {
1003        Expression::Alias(alias) => Some(alias.alias.name.clone()),
1004        Expression::Column(column) => Some(column.name.name.clone()),
1005        Expression::Identifier(identifier) => Some(identifier.name.clone()),
1006        Expression::Star(_) => Some("*".to_string()),
1007        Expression::Annotated(annotated) => projection_name(&annotated.this),
1008        _ => None,
1009    }
1010}
1011
1012fn projection_is_star(expression: &Expression) -> bool {
1013    matches!(expression, Expression::Star(_))
1014        || matches!(expression, Expression::Column(column) if column.name.name == "*")
1015}
1016
1017fn projection_star_table(expression: &Expression) -> Option<String> {
1018    match expression {
1019        Expression::Star(star) => star
1020            .table
1021            .as_ref()
1022            .map(|identifier| identifier.name.clone()),
1023        Expression::Column(column) if column.name.name == "*" => column
1024            .table
1025            .as_ref()
1026            .map(|identifier| identifier.name.clone()),
1027        _ => None,
1028    }
1029}
1030
1031fn transform_kind(expression: &Expression) -> TransformKind {
1032    if projection_is_star(expression) {
1033        TransformKind::Star
1034    } else if is_cast_expression(expression) {
1035        TransformKind::Cast
1036    } else if contains_aggregate(expression) {
1037        TransformKind::Aggregation
1038    } else if matches!(
1039        expression,
1040        Expression::Column(_) | Expression::Identifier(_)
1041    ) {
1042        TransformKind::Direct
1043    } else if is_simple_constant(expression) {
1044        TransformKind::Constant
1045    } else {
1046        TransformKind::Expression
1047    }
1048}
1049
1050fn is_cast_expression(expression: &Expression) -> bool {
1051    matches!(
1052        expression,
1053        Expression::Cast(_) | Expression::TryCast(_) | Expression::SafeCast(_)
1054    )
1055}
1056
1057fn cast_type(expression: &Expression, dialect: DialectType) -> Option<String> {
1058    match expression {
1059        Expression::Cast(cast) | Expression::TryCast(cast) | Expression::SafeCast(cast) => {
1060            render_data_type(&cast.to, dialect)
1061        }
1062        _ => None,
1063    }
1064}
1065
1066fn render_data_type(data_type: &DataType, dialect: DialectType) -> Option<String> {
1067    Dialect::get(dialect)
1068        .generate(&Expression::DataType(data_type.clone()))
1069        .ok()
1070}
1071
1072fn is_simple_constant(expression: &Expression) -> bool {
1073    match expression {
1074        Expression::Literal(_) | Expression::Boolean(_) | Expression::Null(_) => true,
1075        Expression::Cast(cast) | Expression::TryCast(cast) | Expression::SafeCast(cast) => {
1076            is_simple_constant(&cast.this)
1077        }
1078        Expression::Neg(unary) | Expression::BitwiseNot(unary) => is_simple_constant(&unary.this),
1079        _ => false,
1080    }
1081}
1082
1083fn projection_nullability(
1084    expression: &Expression,
1085    scope: &Scope,
1086    context: &NullabilityContext<'_>,
1087) -> ProjectionNullability {
1088    match expression {
1089        Expression::Alias(alias) => projection_nullability(&alias.this, scope, context),
1090        Expression::Annotated(annotated) => projection_nullability(&annotated.this, scope, context),
1091        Expression::Paren(paren) => projection_nullability(&paren.this, scope, context),
1092        Expression::Literal(_) | Expression::Boolean(_) => ProjectionNullability::NonNull,
1093        Expression::Null(_) => ProjectionNullability::Nullable,
1094        Expression::Count(_) | Expression::CountIf(_) => ProjectionNullability::NonNull,
1095        Expression::Cast(cast) => projection_nullability(&cast.this, scope, context),
1096        Expression::TryCast(_) | Expression::SafeCast(_) => ProjectionNullability::Unknown,
1097        Expression::Column(column) => column_nullability(
1098            &column.name.name,
1099            column.table.as_ref().map(|table| table.name.as_str()),
1100            scope,
1101            context,
1102        ),
1103        Expression::Identifier(identifier) => {
1104            column_nullability(&identifier.name, None, scope, context)
1105        }
1106        Expression::Coalesce(func) => coalesce_nullability(&func.expressions, scope, context),
1107        _ => ProjectionNullability::Unknown,
1108    }
1109}
1110
1111fn column_nullability(
1112    column_name: &str,
1113    source_name: Option<&str>,
1114    scope: &Scope,
1115    context: &NullabilityContext<'_>,
1116) -> ProjectionNullability {
1117    let resolved_source_name = source_name
1118        .map(str::to_string)
1119        .or_else(|| single_scope_source_name(scope));
1120
1121    if let Some(source_name) = &resolved_source_name {
1122        if context
1123            .nullable_sources
1124            .contains(&normalize_lookup_name(source_name))
1125        {
1126            return ProjectionNullability::Nullable;
1127        }
1128    }
1129
1130    let Some(schema) = context.schema else {
1131        return ProjectionNullability::Unknown;
1132    };
1133
1134    let table_name = resolved_source_name
1135        .as_ref()
1136        .and_then(|name| scope.sources.get(name).and_then(source_table_name))
1137        .or(resolved_source_name);
1138
1139    let Some(table_name) = table_name else {
1140        return ProjectionNullability::Unknown;
1141    };
1142
1143    match schema.column(&table_name, column_name) {
1144        Some(info) if info.primary_key || info.nullable == Some(false) => {
1145            ProjectionNullability::NonNull
1146        }
1147        Some(info) if info.nullable == Some(true) => ProjectionNullability::Nullable,
1148        Some(_) | None => ProjectionNullability::Unknown,
1149    }
1150}
1151
1152fn single_scope_source_name(scope: &Scope) -> Option<String> {
1153    if scope.sources.len() == 1 {
1154        scope.sources.keys().next().cloned()
1155    } else {
1156        None
1157    }
1158}
1159
1160fn coalesce_nullability(
1161    expressions: &[Expression],
1162    scope: &Scope,
1163    context: &NullabilityContext<'_>,
1164) -> ProjectionNullability {
1165    if expressions.is_empty() {
1166        return ProjectionNullability::Unknown;
1167    }
1168
1169    let mut all_nullable = true;
1170
1171    for expression in expressions {
1172        match projection_nullability(unwrap_projection_alias(expression), scope, context) {
1173            ProjectionNullability::NonNull => return ProjectionNullability::NonNull,
1174            ProjectionNullability::Nullable => {}
1175            ProjectionNullability::Unknown => all_nullable = false,
1176        }
1177    }
1178
1179    if all_nullable {
1180        ProjectionNullability::Nullable
1181    } else {
1182        ProjectionNullability::Unknown
1183    }
1184}
1185
1186fn terminal_references_from_lineage(node: &LineageNode) -> Vec<ColumnReferenceFact> {
1187    let mut refs = Vec::new();
1188    collect_terminal_references(node, &mut refs);
1189    dedupe_column_refs(refs)
1190}
1191
1192fn collect_terminal_references(node: &LineageNode, refs: &mut Vec<ColumnReferenceFact>) {
1193    if node.downstream.is_empty() {
1194        if let Some(reference) = column_reference_from_lineage_node(node) {
1195            refs.push(reference);
1196        }
1197        return;
1198    }
1199
1200    for child in &node.downstream {
1201        collect_terminal_references(child, refs);
1202    }
1203}
1204
1205fn column_reference_from_lineage_node(node: &LineageNode) -> Option<ColumnReferenceFact> {
1206    match &node.expression {
1207        Expression::Column(column) => {
1208            let source_name = non_empty_string(node.source_name.clone());
1209            let table =
1210                lineage_node_table(node).or_else(|| column.table.as_ref().map(|t| t.name.clone()));
1211            let confidence = if node.source_kind == SourceKind::Unknown && source_name.is_none() {
1212                ReferenceConfidence::Unknown
1213            } else {
1214                ReferenceConfidence::Resolved
1215            };
1216            Some(ColumnReferenceFact {
1217                source_name,
1218                source_alias: node.source_alias.clone(),
1219                source_kind: node.source_kind,
1220                table,
1221                column: column.name.name.clone(),
1222                unqualified: column.table.is_none(),
1223                confidence,
1224            })
1225        }
1226        Expression::Star(_) => Some(ColumnReferenceFact {
1227            source_name: non_empty_string(node.source_name.clone()),
1228            source_alias: node.source_alias.clone(),
1229            source_kind: node.source_kind,
1230            table: lineage_node_table(node),
1231            column: "*".to_string(),
1232            unqualified: true,
1233            confidence: if node.source_kind == SourceKind::Unknown {
1234                ReferenceConfidence::Unknown
1235            } else {
1236                ReferenceConfidence::Resolved
1237            },
1238        }),
1239        _ => None,
1240    }
1241}
1242
1243fn lineage_node_table(node: &LineageNode) -> Option<String> {
1244    match &node.source {
1245        Expression::Table(table) => Some(table_name(table)),
1246        _ => None,
1247    }
1248}
1249
1250fn fallback_column_references(expression: &Expression, scope: &Scope) -> Vec<ColumnReferenceFact> {
1251    let mut refs = Vec::new();
1252    let source_count = scope.sources.len();
1253    let single_source = if source_count == 1 {
1254        scope.sources.iter().next()
1255    } else {
1256        None
1257    };
1258
1259    for column_expr in expression.find_all(|candidate| matches!(candidate, Expression::Column(_))) {
1260        if let Expression::Column(column) = column_expr {
1261            if column.name.name == "*" {
1262                continue;
1263            }
1264            let source = column
1265                .table
1266                .as_ref()
1267                .and_then(|table| scope.sources.get(&table.name));
1268            let (source_name, source_alias, source_kind, table, confidence) =
1269                if let Some(table_identifier) = &column.table {
1270                    if let Some(source) = source {
1271                        (
1272                            Some(table_identifier.name.clone()),
1273                            source.alias.clone(),
1274                            source.kind,
1275                            source_table_name(source)
1276                                .or_else(|| Some(table_identifier.name.clone())),
1277                            ReferenceConfidence::Resolved,
1278                        )
1279                    } else {
1280                        (
1281                            Some(table_identifier.name.clone()),
1282                            None,
1283                            SourceKind::Unknown,
1284                            Some(table_identifier.name.clone()),
1285                            ReferenceConfidence::Unknown,
1286                        )
1287                    }
1288                } else if let Some((name, source)) = single_source {
1289                    (
1290                        Some(name.clone()),
1291                        source.alias.clone(),
1292                        source.kind,
1293                        source_table_name(source).or_else(|| Some(name.clone())),
1294                        ReferenceConfidence::Resolved,
1295                    )
1296                } else if source_count > 1 {
1297                    (
1298                        None,
1299                        None,
1300                        SourceKind::Unknown,
1301                        None,
1302                        ReferenceConfidence::Ambiguous,
1303                    )
1304                } else {
1305                    (
1306                        None,
1307                        None,
1308                        SourceKind::Unknown,
1309                        None,
1310                        ReferenceConfidence::Unknown,
1311                    )
1312                };
1313
1314            refs.push(ColumnReferenceFact {
1315                source_name,
1316                source_alias,
1317                source_kind,
1318                table,
1319                column: column.name.name.clone(),
1320                unqualified: column.table.is_none(),
1321                confidence,
1322            });
1323        }
1324    }
1325
1326    dedupe_column_refs(refs)
1327}
1328
1329fn dedupe_column_refs(refs: Vec<ColumnReferenceFact>) -> Vec<ColumnReferenceFact> {
1330    let mut seen = HashSet::new();
1331    let mut deduped = Vec::new();
1332
1333    for reference in refs {
1334        let key = (
1335            reference.source_name.clone(),
1336            reference.source_alias.clone(),
1337            reference.table.clone(),
1338            reference.column.clone(),
1339            format!("{:?}", reference.source_kind),
1340            reference.unqualified,
1341            format!("{:?}", reference.confidence),
1342        );
1343        if seen.insert(key) {
1344            deduped.push(reference);
1345        }
1346    }
1347
1348    deduped
1349}
1350
1351fn relation_facts(
1352    scope: &Scope,
1353    mapping_schema: Option<&crate::schema::MappingSchema>,
1354    dialect: DialectType,
1355) -> Vec<RelationFact> {
1356    let mut relations = Vec::new();
1357    let mut seen = HashSet::new();
1358    collect_relation_facts(scope, mapping_schema, dialect, &mut seen, &mut relations);
1359
1360    relations.sort_by(|left, right| {
1361        left.name
1362            .cmp(&right.name)
1363            .then_with(|| left.alias.cmp(&right.alias))
1364    });
1365    relations
1366}
1367
1368fn collect_relation_facts(
1369    scope: &Scope,
1370    mapping_schema: Option<&crate::schema::MappingSchema>,
1371    dialect: DialectType,
1372    seen: &mut HashSet<String>,
1373    relations: &mut Vec<RelationFact>,
1374) {
1375    for relation in scope.sources.iter().map(|(source_name, source)| {
1376        let identity = source_table_identity(source);
1377        RelationFact {
1378            name: source
1379                .lineage_name
1380                .clone()
1381                .or_else(|| identity.as_ref().map(|identity| identity.name.clone()))
1382                .unwrap_or_else(|| source_name.clone()),
1383            alias: source.alias.clone().or_else(|| source_alias(source)),
1384            kind: source.kind,
1385            columns: source_columns(source, mapping_schema, dialect),
1386            catalog: identity
1387                .as_ref()
1388                .and_then(|identity| identity.catalog.clone()),
1389            schema: identity
1390                .as_ref()
1391                .and_then(|identity| identity.schema.clone()),
1392            table: identity
1393                .as_ref()
1394                .and_then(|identity| identity.table.clone()),
1395        }
1396    }) {
1397        let key = format!("{:?}|{}|{:?}", relation.kind, relation.name, relation.alias);
1398        if seen.insert(key) {
1399            relations.push(relation);
1400        }
1401    }
1402
1403    for branch_scope in &scope.union_scopes {
1404        collect_relation_facts(branch_scope, mapping_schema, dialect, seen, relations);
1405    }
1406}
1407
1408fn base_table_facts(
1409    scope: &Scope,
1410    mapping_schema: Option<&crate::schema::MappingSchema>,
1411    dialect: DialectType,
1412) -> Vec<RelationFact> {
1413    let mut relations = Vec::new();
1414    let mut seen = HashSet::new();
1415
1416    collect_base_table_facts(scope, mapping_schema, dialect, &mut seen, &mut relations);
1417
1418    relations.sort_by(|left, right| left.name.cmp(&right.name));
1419    relations
1420}
1421
1422fn collect_base_table_facts(
1423    scope: &Scope,
1424    mapping_schema: Option<&crate::schema::MappingSchema>,
1425    dialect: DialectType,
1426    seen: &mut HashSet<String>,
1427    relations: &mut Vec<RelationFact>,
1428) {
1429    for source in scope.sources.values() {
1430        if source.kind != SourceKind::Table {
1431            continue;
1432        }
1433
1434        let Some(identity) = source_table_identity(source) else {
1435            continue;
1436        };
1437
1438        if seen.insert(identity.name.clone()) {
1439            relations.push(RelationFact {
1440                name: identity.name,
1441                alias: source.alias.clone().or_else(|| source_alias(source)),
1442                kind: SourceKind::Table,
1443                columns: source_columns(source, mapping_schema, dialect),
1444                catalog: identity.catalog,
1445                schema: identity.schema,
1446                table: identity.table,
1447            });
1448        }
1449    }
1450
1451    for child_scope in scope
1452        .cte_scopes
1453        .iter()
1454        .chain(scope.union_scopes.iter())
1455        .chain(scope.table_scopes.iter())
1456        .chain(scope.derived_table_scopes.iter())
1457        .chain(scope.subquery_scopes.iter())
1458    {
1459        collect_base_table_facts(child_scope, mapping_schema, dialect, seen, relations);
1460    }
1461}
1462
1463fn source_columns(
1464    source: &SourceInfo,
1465    mapping_schema: Option<&crate::schema::MappingSchema>,
1466    dialect: DialectType,
1467) -> Vec<String> {
1468    match &source.expression {
1469        Expression::Table(table) => mapping_schema
1470            .and_then(|schema| schema.column_names(&table_name(table)).ok())
1471            .unwrap_or_default(),
1472        Expression::Select(_)
1473        | Expression::Union(_)
1474        | Expression::Intersect(_)
1475        | Expression::Except(_) => {
1476            get_output_column_names_for_dialect(&source.expression, Some(dialect))
1477        }
1478        Expression::Subquery(subquery) => {
1479            get_output_column_names_for_dialect(&subquery.this, Some(dialect))
1480        }
1481        Expression::Cte(cte) if !cte.columns.is_empty() => cte
1482            .columns
1483            .iter()
1484            .map(|column| column.name.clone())
1485            .collect(),
1486        Expression::Cte(cte) => get_output_column_names_for_dialect(&cte.this, Some(dialect)),
1487        _ => Vec::new(),
1488    }
1489}
1490
1491fn source_table_name(source: &SourceInfo) -> Option<String> {
1492    source_table_identity(source).map(|identity| identity.name)
1493}
1494
1495fn source_alias(source: &SourceInfo) -> Option<String> {
1496    match &source.expression {
1497        Expression::Table(table) => table.alias.as_ref().map(|alias| alias.name.clone()),
1498        Expression::Subquery(subquery) => subquery.alias.as_ref().map(|alias| alias.name.clone()),
1499        _ => None,
1500    }
1501}
1502
1503fn table_name(table: &TableRef) -> String {
1504    let mut parts = Vec::new();
1505    if let Some(catalog) = &table.catalog {
1506        parts.push(catalog.name.clone());
1507    }
1508    if let Some(schema) = &table.schema {
1509        parts.push(schema.name.clone());
1510    }
1511    parts.push(table.name.name.clone());
1512    parts.join(".")
1513}
1514
1515#[derive(Debug, Clone)]
1516struct RelationIdentity {
1517    name: String,
1518    catalog: Option<String>,
1519    schema: Option<String>,
1520    table: Option<String>,
1521}
1522
1523fn source_table_identity(source: &SourceInfo) -> Option<RelationIdentity> {
1524    match &source.expression {
1525        Expression::Table(table) => Some(table_identity(table)),
1526        _ => None,
1527    }
1528}
1529
1530fn table_identity(table: &TableRef) -> RelationIdentity {
1531    RelationIdentity {
1532        name: table_name(table),
1533        catalog: table.catalog.as_ref().map(|catalog| catalog.name.clone()),
1534        schema: table.schema.as_ref().map(|schema| schema.name.clone()),
1535        table: Some(table.name.name.clone()),
1536    }
1537}
1538
1539fn set_operation_facts(
1540    expression: &Expression,
1541    scope: &Scope,
1542    dialect: DialectType,
1543) -> Vec<SetOperationFact> {
1544    let mut facts = Vec::new();
1545    collect_set_operation_facts(expression, scope, dialect, &mut facts);
1546    facts
1547}
1548
1549fn collect_set_operation_facts(
1550    expression: &Expression,
1551    scope: &Scope,
1552    dialect: DialectType,
1553    facts: &mut Vec<SetOperationFact>,
1554) {
1555    match expression {
1556        Expression::Union(union) => {
1557            facts.push(SetOperationFact {
1558                kind: "union".to_string(),
1559                all: union.all,
1560                distinct: union.distinct,
1561                output_columns: get_output_column_names_for_dialect(expression, Some(dialect)),
1562                branches: set_operation_branches(
1563                    &union.left,
1564                    &union.right,
1565                    scope,
1566                    dialect,
1567                    SetOperationBranchRole::Value,
1568                ),
1569            });
1570            collect_set_operation_facts(&union.left, scope, dialect, facts);
1571            collect_set_operation_facts(&union.right, scope, dialect, facts);
1572        }
1573        Expression::Intersect(intersect) => {
1574            facts.push(SetOperationFact {
1575                kind: "intersect".to_string(),
1576                all: intersect.all,
1577                distinct: intersect.distinct,
1578                output_columns: get_output_column_names_for_dialect(expression, Some(dialect)),
1579                branches: set_operation_branches(
1580                    &intersect.left,
1581                    &intersect.right,
1582                    scope,
1583                    dialect,
1584                    SetOperationBranchRole::Filter,
1585                ),
1586            });
1587            collect_set_operation_facts(&intersect.left, scope, dialect, facts);
1588            collect_set_operation_facts(&intersect.right, scope, dialect, facts);
1589        }
1590        Expression::Except(except) => {
1591            facts.push(SetOperationFact {
1592                kind: "except".to_string(),
1593                all: except.all,
1594                distinct: except.distinct,
1595                output_columns: get_output_column_names_for_dialect(expression, Some(dialect)),
1596                branches: set_operation_branches(
1597                    &except.left,
1598                    &except.right,
1599                    scope,
1600                    dialect,
1601                    SetOperationBranchRole::Filter,
1602                ),
1603            });
1604            collect_set_operation_facts(&except.left, scope, dialect, facts);
1605            collect_set_operation_facts(&except.right, scope, dialect, facts);
1606        }
1607        Expression::Subquery(subquery) => {
1608            collect_set_operation_facts(&subquery.this, scope, dialect, facts);
1609        }
1610        _ => {}
1611    }
1612}
1613
1614fn set_operation_branches(
1615    left: &Expression,
1616    right: &Expression,
1617    scope: &Scope,
1618    dialect: DialectType,
1619    right_role: SetOperationBranchRole,
1620) -> Vec<SetOperationBranchFact> {
1621    vec![
1622        SetOperationBranchFact {
1623            index: 0,
1624            role: SetOperationBranchRole::Value,
1625            projections: projection_facts_for_branch(left, scope, dialect),
1626        },
1627        SetOperationBranchFact {
1628            index: 1,
1629            role: right_role,
1630            projections: projection_facts_for_branch(right, scope, dialect),
1631        },
1632    ]
1633}
1634
1635fn projection_facts_for_branch(
1636    expression: &Expression,
1637    root_scope: &Scope,
1638    dialect: DialectType,
1639) -> Vec<ProjectionFact> {
1640    let branch_scope = build_scope(expression);
1641    let scope = if branch_scope.sources.is_empty() {
1642        root_scope
1643    } else {
1644        &branch_scope
1645    };
1646    let nullability_context = NullabilityContext {
1647        schema: None,
1648        nullable_sources: nullable_source_names(expression),
1649    };
1650    projection_facts_for_query(expression, scope, dialect, &nullability_context)
1651}
1652
1653fn non_empty_string(value: String) -> Option<String> {
1654    if value.is_empty() {
1655        None
1656    } else {
1657        Some(value)
1658    }
1659}