Skip to main content

polyglot_sql/
resolver.rs

1//! Column Resolver Module
2//!
3//! This module provides functionality for resolving column references to their
4//! source tables. It handles:
5//! - Finding which table a column belongs to
6//! - Resolving ambiguous column references
7//! - Handling join context for disambiguation
8//! - Supporting set operations (UNION, INTERSECT, EXCEPT)
9//!
10//! Based on the Python implementation in `sqlglot/optimizer/resolver.py`.
11
12use crate::dialects::DialectType;
13use crate::expressions::{Expression, Identifier, TableRef};
14#[cfg(feature = "generate")]
15use crate::generator::Generator;
16use crate::schema::{normalize_name, Schema};
17use crate::scope::{Scope, SourceInfo};
18use crate::traversal::ExpressionWalk;
19use std::collections::{HashMap, HashSet};
20use thiserror::Error;
21
22/// Errors that can occur during column resolution
23#[derive(Debug, Error, Clone)]
24pub enum ResolverError {
25    #[error("Unknown table: {0}")]
26    UnknownTable(String),
27
28    #[error("Ambiguous column: {column} appears in multiple sources: {sources}")]
29    AmbiguousColumn { column: String, sources: String },
30
31    #[error("Column not found: {0}")]
32    ColumnNotFound(String),
33
34    #[error("Unknown set operation: {0}")]
35    UnknownSetOperation(String),
36}
37
38/// Result type for resolver operations
39pub type ResolverResult<T> = Result<T, ResolverError>;
40
41/// Helper for resolving columns to their source tables.
42///
43/// This is a struct so we can lazily load some things and easily share
44/// them across functions.
45pub struct Resolver<'a> {
46    /// The scope being analyzed
47    pub scope: &'a Scope,
48    /// The schema for table/column information
49    schema: &'a dyn Schema,
50    /// The dialect being used
51    pub dialect: Option<DialectType>,
52    /// Whether to infer schema from context
53    infer_schema: bool,
54    /// Cached source columns: source_name -> column names
55    source_columns_cache: HashMap<String, Vec<String>>,
56    /// Cached unambiguous columns: column_name -> source_name
57    unambiguous_columns_cache: Option<HashMap<String, String>>,
58    /// Cached set of all available columns
59    all_columns_cache: Option<HashSet<String>>,
60}
61
62impl<'a> Resolver<'a> {
63    /// Create a new resolver for a scope
64    pub fn new(scope: &'a Scope, schema: &'a dyn Schema, infer_schema: bool) -> Self {
65        Self {
66            scope,
67            schema,
68            dialect: schema.dialect(),
69            infer_schema,
70            source_columns_cache: HashMap::new(),
71            unambiguous_columns_cache: None,
72            all_columns_cache: None,
73        }
74    }
75
76    /// Get the table for a column name.
77    ///
78    /// Returns the table name if it can be found/inferred.
79    pub fn get_table(&mut self, column_name: &str) -> Option<String> {
80        // Try to find table from all sources (unambiguous lookup)
81        let table_name = self.get_table_name_from_sources(column_name, None);
82
83        // If we found a table, return it
84        if table_name.is_some() {
85            return table_name;
86        }
87
88        // If schema inference is enabled and exactly one source has no schema,
89        // assume the column belongs to that source
90        if self.infer_schema {
91            let sources_without_schema: Vec<_> = self
92                .get_all_source_columns()
93                .iter()
94                .filter(|(_, columns)| columns.is_empty() || columns.contains(&"*".to_string()))
95                .map(|(name, _)| name.clone())
96                .collect();
97
98            if sources_without_schema.len() == 1 {
99                return Some(sources_without_schema[0].clone());
100            }
101        }
102
103        None
104    }
105
106    /// Get the table for a column, returning an Identifier
107    pub fn get_table_identifier(&mut self, column_name: &str) -> Option<Identifier> {
108        self.get_table(column_name).map(Identifier::new)
109    }
110
111    /// Check if a table exists in the schema (not necessarily in the current scope).
112    /// Used to detect correlated references to outer scope tables.
113    pub fn table_exists_in_schema(&self, table_name: &str) -> bool {
114        self.schema.column_names(table_name).is_ok()
115    }
116
117    /// Find the table for a column by searching all schema tables not in the current scope.
118    /// Used for correlated subquery resolution: if an unqualified column can't be resolved
119    /// in the current scope, check if it uniquely belongs to an outer-scope table.
120    /// Returns Some(table_name) if the column is found in exactly one non-local table.
121    pub fn find_column_in_outer_schema_tables(&self, column_name: &str) -> Option<String> {
122        let tables = self.schema.find_tables_for_column(column_name);
123        // Filter to tables NOT in the current scope
124        let outer_tables: Vec<String> = tables
125            .into_iter()
126            .filter(|t| !self.scope.sources.contains_key(t))
127            .collect();
128        // Only return if unambiguous (exactly one outer table has this column)
129        if outer_tables.len() == 1 {
130            Some(outer_tables.into_iter().next().unwrap())
131        } else {
132            None
133        }
134    }
135
136    /// Get all available columns across all sources in this scope
137    pub fn all_columns(&mut self) -> &HashSet<String> {
138        if self.all_columns_cache.is_none() {
139            let mut all = HashSet::new();
140            for columns in self.get_all_source_columns().values() {
141                all.extend(columns.iter().cloned());
142            }
143            self.all_columns_cache = Some(all);
144        }
145        self.all_columns_cache
146            .as_ref()
147            .expect("cache populated above")
148    }
149
150    /// Get column names for a source.
151    ///
152    /// Returns the list of column names available from the given source.
153    pub fn get_source_columns(&mut self, source_name: &str) -> ResolverResult<Vec<String>> {
154        // Check cache first
155        if let Some(columns) = self.source_columns_cache.get(source_name) {
156            return Ok(columns.clone());
157        }
158
159        // Get the source info
160        let source_info = self
161            .scope
162            .sources
163            .get(source_name)
164            .ok_or_else(|| ResolverError::UnknownTable(source_name.to_string()))?;
165
166        let columns = self.extract_columns_from_source(source_info)?;
167
168        // Cache the result
169        self.source_columns_cache
170            .insert(source_name.to_string(), columns.clone());
171
172        Ok(columns)
173    }
174
175    /// Extract column names from a source expression
176    fn extract_columns_from_source(&self, source_info: &SourceInfo) -> ResolverResult<Vec<String>> {
177        self.get_source_columns_for_expression(&source_info.expression)
178    }
179
180    fn get_source_columns_for_expression(
181        &self,
182        expression: &Expression,
183    ) -> ResolverResult<Vec<String>> {
184        let columns = match expression {
185            Expression::Table(table) => {
186                // For tables, try to get columns from schema.
187                // Build the fully qualified name (catalog.schema.table) to
188                // match how MappingSchema stores hierarchical keys.
189                let table_name = qualified_table_name(table);
190                match self.schema.column_names(&table_name) {
191                    Ok(cols) => cols,
192                    Err(_) => Vec::new(), // Schema might not have this table
193                }
194            }
195            Expression::Subquery(subquery) => {
196                // For subqueries, get named_selects from the inner query
197                self.get_named_selects(&subquery.this)
198            }
199            Expression::Select(select) => {
200                // For derived tables that are SELECT expressions
201                self.get_select_column_names(select)
202            }
203            Expression::Union(union) => {
204                // For UNION, columns come from the set operation
205                self.get_source_columns_from_set_op(&Expression::Union(union.clone()))?
206            }
207            Expression::Intersect(intersect) => {
208                self.get_source_columns_from_set_op(&Expression::Intersect(intersect.clone()))?
209            }
210            Expression::Except(except) => {
211                self.get_source_columns_from_set_op(&Expression::Except(except.clone()))?
212            }
213            Expression::Cte(cte) => {
214                if !cte.columns.is_empty() {
215                    cte.columns.iter().map(|c| c.name.clone()).collect()
216                } else {
217                    self.get_named_selects(&cte.this)
218                }
219            }
220            Expression::Pivot(pivot) => self.get_pivot_output_columns(pivot),
221            Expression::Unpivot(unpivot) => self.get_unpivot_output_columns(unpivot),
222            Expression::Alias(alias) if matches!(&alias.this, Expression::Unnest(_)) => {
223                alias_output_columns(alias)
224            }
225            Expression::Alias(alias) => {
226                let columns = self.get_source_columns_for_expression(&alias.this)?;
227                apply_alias_columns(columns, &alias.column_aliases)
228            }
229            Expression::Unnest(unnest) => unnest_output_columns(unnest),
230            Expression::Lateral(lateral) => lateral_output_columns(lateral),
231            Expression::LateralView(lateral_view) => lateral_view_output_columns(lateral_view),
232            Expression::Paren(paren) => self.get_source_columns_for_expression(&paren.this)?,
233            _ => Vec::new(),
234        };
235
236        Ok(columns)
237    }
238
239    /// Get named selects (column names) from an expression
240    fn get_named_selects(&self, expr: &Expression) -> Vec<String> {
241        match expr {
242            Expression::Select(select) => self.get_select_column_names(select),
243            Expression::Union(union) => {
244                // For unions, use the left side's columns
245                self.get_named_selects(&union.left)
246            }
247            Expression::Intersect(intersect) => self.get_named_selects(&intersect.left),
248            Expression::Except(except) => self.get_named_selects(&except.left),
249            Expression::Subquery(subquery) => self.get_named_selects(&subquery.this),
250            Expression::Alias(alias) => {
251                let columns = self.get_named_selects(&alias.this);
252                apply_alias_columns(columns, &alias.column_aliases)
253            }
254            Expression::Paren(paren) => self.get_named_selects(&paren.this),
255            _ => Vec::new(),
256        }
257    }
258
259    /// Get column names from a SELECT expression
260    fn get_select_column_names(&self, select: &crate::expressions::Select) -> Vec<String> {
261        select
262            .expressions
263            .iter()
264            .filter_map(|expr| self.get_expression_alias(expr))
265            .collect()
266    }
267
268    /// Get the alias or name for a select expression
269    fn get_expression_alias(&self, expr: &Expression) -> Option<String> {
270        match expr {
271            Expression::Alias(alias) => Some(alias.alias.name.clone()),
272            Expression::Column(col) => Some(col.name.name.clone()),
273            Expression::Star(_) => Some("*".to_string()),
274            Expression::Identifier(id) => Some(id.name.clone()),
275            _ => None,
276        }
277    }
278
279    fn get_pivot_output_columns(&self, pivot: &crate::expressions::Pivot) -> Vec<String> {
280        if pivot.unpivot {
281            return self.get_pivot_unpivot_output_columns(pivot);
282        }
283
284        let pre_columns = self.get_source_output_columns(&pivot.this);
285        if pre_columns.is_empty() || pre_columns.iter().any(|column| column == "*") {
286            return Vec::new();
287        }
288
289        let excluded = pivot_excluded_source_columns(pivot, self.dialect);
290        let generated = pivot_generated_output_columns(pivot, self.dialect);
291        if excluded.is_empty() || generated.is_empty() {
292            return Vec::new();
293        }
294
295        let mut columns: Vec<String> = pre_columns
296            .into_iter()
297            .filter(|column| !excluded.contains(&normalize_column_name(column, self.dialect)))
298            .collect();
299        columns.extend(generated);
300        apply_alias_columns(columns, &pivot.alias_columns)
301    }
302
303    fn get_pivot_unpivot_output_columns(&self, pivot: &crate::expressions::Pivot) -> Vec<String> {
304        let pre_columns = self.get_source_output_columns(&pivot.this);
305        if pre_columns.is_empty() || pre_columns.iter().any(|column| column == "*") {
306            return Vec::new();
307        }
308
309        let input_columns: HashSet<String> = pivot
310            .expressions
311            .iter()
312            .flat_map(expression_column_names)
313            .map(|column| normalize_column_name(&column, self.dialect))
314            .collect();
315        let mut columns: Vec<String> = pre_columns
316            .into_iter()
317            .filter(|column| !input_columns.contains(&normalize_column_name(column, self.dialect)))
318            .collect();
319
320        if let Some(Expression::UnpivotColumns(unpivot_columns)) = pivot.into.as_deref() {
321            if let Some(name) = expression_name(&unpivot_columns.this) {
322                columns.push(name);
323            }
324            for value_column in &unpivot_columns.expressions {
325                if let Some(name) = expression_name(value_column) {
326                    columns.push(name);
327                }
328            }
329        }
330
331        apply_alias_columns(columns, &pivot.alias_columns)
332    }
333
334    fn get_unpivot_output_columns(&self, unpivot: &crate::expressions::Unpivot) -> Vec<String> {
335        let pre_columns = self.get_source_output_columns(&unpivot.this);
336        if pre_columns.is_empty() || pre_columns.iter().any(|column| column == "*") {
337            return Vec::new();
338        }
339
340        let input_columns: HashSet<String> = unpivot
341            .columns
342            .iter()
343            .flat_map(expression_column_names)
344            .map(|column| normalize_column_name(&column, self.dialect))
345            .collect();
346        let mut columns: Vec<String> = pre_columns
347            .into_iter()
348            .filter(|column| !input_columns.contains(&normalize_column_name(column, self.dialect)))
349            .collect();
350        columns.push(unpivot.name_column.name.clone());
351        columns.push(unpivot.value_column.name.clone());
352        columns.extend(
353            unpivot
354                .extra_value_columns
355                .iter()
356                .map(|column| column.name.clone()),
357        );
358        apply_alias_columns(columns, &unpivot.alias_columns)
359    }
360
361    fn get_source_output_columns(&self, source: &Expression) -> Vec<String> {
362        match source {
363            Expression::Table(table) => {
364                if table.schema.is_none() && table.catalog.is_none() {
365                    if let Some(source) = self.scope.cte_sources.get(&table.name.name) {
366                        return self.extract_columns_from_source(source).unwrap_or_default();
367                    }
368                }
369
370                let table_name = qualified_table_name(table);
371                self.schema.column_names(&table_name).unwrap_or_default()
372            }
373            Expression::Subquery(subquery) => self.get_named_selects(&subquery.this),
374            Expression::Select(select) => self.get_select_column_names(select),
375            Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_) => self
376                .get_source_columns_from_set_op(source)
377                .unwrap_or_default(),
378            Expression::Alias(alias) if matches!(&alias.this, Expression::Unnest(_)) => {
379                alias_output_columns(alias)
380            }
381            Expression::Alias(alias) => {
382                let columns = self.get_source_output_columns(&alias.this);
383                apply_alias_columns(columns, &alias.column_aliases)
384            }
385            Expression::Unnest(unnest) => unnest_output_columns(unnest),
386            Expression::Lateral(lateral) => lateral_output_columns(lateral),
387            Expression::LateralView(lateral_view) => lateral_view_output_columns(lateral_view),
388            Expression::Cte(cte) => {
389                if cte.columns.is_empty() {
390                    self.get_named_selects(&cte.this)
391                } else {
392                    cte.columns
393                        .iter()
394                        .map(|column| column.name.clone())
395                        .collect()
396                }
397            }
398            Expression::Paren(paren) => self.get_source_output_columns(&paren.this),
399            _ => Vec::new(),
400        }
401    }
402
403    /// Get columns from a set operation (UNION, INTERSECT, EXCEPT)
404    pub fn get_source_columns_from_set_op(
405        &self,
406        expression: &Expression,
407    ) -> ResolverResult<Vec<String>> {
408        match expression {
409            Expression::Select(select) => Ok(self.get_select_column_names(select)),
410            Expression::Subquery(subquery) => {
411                if matches!(
412                    &subquery.this,
413                    Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_)
414                ) {
415                    self.get_source_columns_from_set_op(&subquery.this)
416                } else {
417                    Ok(self.get_named_selects(&subquery.this))
418                }
419            }
420            Expression::Alias(alias) => {
421                let columns = self.get_source_columns_from_set_op(&alias.this)?;
422                Ok(apply_alias_columns(columns, &alias.column_aliases))
423            }
424            Expression::Paren(paren) => self.get_source_columns_from_set_op(&paren.this),
425            Expression::Union(union) => {
426                // Standard UNION: columns come from the left side
427                self.get_source_columns_from_set_op(&union.left)
428            }
429            Expression::Intersect(intersect) => {
430                self.get_source_columns_from_set_op(&intersect.left)
431            }
432            Expression::Except(except) => self.get_source_columns_from_set_op(&except.left),
433            _ => Err(ResolverError::UnknownSetOperation(format!(
434                "{:?}",
435                expression
436            ))),
437        }
438    }
439
440    /// Get all source columns for all sources in the scope
441    fn get_all_source_columns(&mut self) -> HashMap<String, Vec<String>> {
442        let source_names: Vec<_> = self.scope.sources.keys().cloned().collect();
443
444        let mut result = HashMap::new();
445        for source_name in source_names {
446            if let Ok(columns) = self.get_source_columns(&source_name) {
447                result.insert(source_name, columns);
448            }
449        }
450        result
451    }
452
453    /// Get the table name for a column from the sources
454    fn get_table_name_from_sources(
455        &mut self,
456        column_name: &str,
457        source_columns: Option<&HashMap<String, Vec<String>>>,
458    ) -> Option<String> {
459        let normalized_column_name = normalize_column_name(column_name, self.dialect);
460        let unambiguous = match source_columns {
461            Some(cols) => self.compute_unambiguous_columns(cols),
462            None => {
463                if self.unambiguous_columns_cache.is_none() {
464                    let all_source_columns = self.get_all_source_columns();
465                    self.unambiguous_columns_cache =
466                        Some(self.compute_unambiguous_columns(&all_source_columns));
467                }
468                self.unambiguous_columns_cache
469                    .clone()
470                    .expect("cache populated above")
471            }
472        };
473
474        unambiguous.get(&normalized_column_name).cloned()
475    }
476
477    /// Compute unambiguous columns mapping
478    ///
479    /// A column is unambiguous if it appears in exactly one source.
480    fn compute_unambiguous_columns(
481        &self,
482        source_columns: &HashMap<String, Vec<String>>,
483    ) -> HashMap<String, String> {
484        if source_columns.is_empty() {
485            return HashMap::new();
486        }
487
488        let mut column_to_sources: HashMap<String, Vec<String>> = HashMap::new();
489
490        for (source_name, columns) in source_columns {
491            for column in columns {
492                column_to_sources
493                    .entry(normalize_column_name(column, self.dialect))
494                    .or_default()
495                    .push(source_name.clone());
496            }
497        }
498
499        // Keep only columns that appear in exactly one source
500        column_to_sources
501            .into_iter()
502            .filter(|(_, sources)| sources.len() == 1)
503            .map(|(column, sources)| (column, sources.into_iter().next().unwrap()))
504            .collect()
505    }
506
507    /// Check if a column is ambiguous (appears in multiple sources)
508    pub fn is_ambiguous(&mut self, column_name: &str) -> bool {
509        let normalized_column_name = normalize_column_name(column_name, self.dialect);
510        let all_source_columns = self.get_all_source_columns();
511        let sources_with_column: Vec<_> = all_source_columns
512            .iter()
513            .filter(|(_, columns)| {
514                columns.iter().any(|column| {
515                    normalize_column_name(column, self.dialect) == normalized_column_name
516                })
517            })
518            .map(|(name, _)| name.clone())
519            .collect();
520
521        sources_with_column.len() > 1
522    }
523
524    /// Get all sources that contain a given column
525    pub fn sources_for_column(&mut self, column_name: &str) -> Vec<String> {
526        let normalized_column_name = normalize_column_name(column_name, self.dialect);
527        let all_source_columns = self.get_all_source_columns();
528        all_source_columns
529            .iter()
530            .filter(|(_, columns)| {
531                columns.iter().any(|column| {
532                    normalize_column_name(column, self.dialect) == normalized_column_name
533                })
534            })
535            .map(|(name, _)| name.clone())
536            .collect()
537    }
538
539    /// Try to disambiguate a column based on join context
540    ///
541    /// In join conditions, a column can sometimes be disambiguated based on
542    /// which tables have been joined up to that point.
543    pub fn disambiguate_in_join_context(
544        &mut self,
545        column_name: &str,
546        available_sources: &[String],
547    ) -> Option<String> {
548        let normalized_column_name = normalize_column_name(column_name, self.dialect);
549        let mut matching_sources = Vec::new();
550
551        for source_name in available_sources {
552            if let Ok(columns) = self.get_source_columns(source_name) {
553                if columns.iter().any(|column| {
554                    normalize_column_name(column, self.dialect) == normalized_column_name
555                }) {
556                    matching_sources.push(source_name.clone());
557                }
558            }
559        }
560
561        if matching_sources.len() == 1 {
562            Some(matching_sources.remove(0))
563        } else {
564            None
565        }
566    }
567}
568
569fn normalize_column_name(name: &str, dialect: Option<DialectType>) -> String {
570    normalize_name(name, dialect, false, true)
571}
572
573fn apply_alias_columns(mut columns: Vec<String>, alias_columns: &[Identifier]) -> Vec<String> {
574    for (idx, alias) in alias_columns.iter().enumerate() {
575        if let Some(column) = columns.get_mut(idx) {
576            *column = alias.name.clone();
577        }
578    }
579    columns
580}
581
582fn unnest_output_columns(unnest: &crate::expressions::UnnestFunc) -> Vec<String> {
583    unnest
584        .alias
585        .iter()
586        .map(|alias| alias.name.clone())
587        .chain(unnest.offset_alias.iter().map(|alias| alias.name.clone()))
588        .collect()
589}
590
591fn alias_output_columns(alias: &crate::expressions::Alias) -> Vec<String> {
592    if alias.column_aliases.is_empty() {
593        vec![alias.alias.name.clone()]
594    } else {
595        alias
596            .column_aliases
597            .iter()
598            .map(|column| column.name.clone())
599            .collect()
600    }
601}
602
603fn lateral_output_columns(lateral: &crate::expressions::Lateral) -> Vec<String> {
604    if lateral.column_aliases.is_empty() {
605        default_virtual_output_columns(&lateral.this)
606    } else {
607        lateral.column_aliases.clone()
608    }
609}
610
611fn lateral_view_output_columns(lateral_view: &crate::expressions::LateralView) -> Vec<String> {
612    lateral_view
613        .column_aliases
614        .iter()
615        .map(|column| column.name.clone())
616        .collect()
617}
618
619fn default_virtual_output_columns(expression: &Expression) -> Vec<String> {
620    match expression {
621        Expression::Unnest(unnest) => unnest_output_columns(unnest),
622        Expression::Alias(alias) if matches!(&alias.this, Expression::Unnest(_)) => {
623            alias_output_columns(alias)
624        }
625        Expression::Function(function) if function.name.eq_ignore_ascii_case("FLATTEN") => {
626            ["seq", "key", "path", "index", "value", "this"]
627                .into_iter()
628                .map(String::from)
629                .collect()
630        }
631        _ => Vec::new(),
632    }
633}
634
635fn pivot_excluded_source_columns(
636    pivot: &crate::expressions::Pivot,
637    dialect: Option<DialectType>,
638) -> HashSet<String> {
639    pivot
640        .fields
641        .iter()
642        .chain(pivot.expressions.iter())
643        .chain(pivot.using.iter())
644        .flat_map(expression_column_names)
645        .map(|column| normalize_column_name(&column, dialect))
646        .collect()
647}
648
649fn pivot_generated_output_columns(
650    pivot: &crate::expressions::Pivot,
651    _dialect: Option<DialectType>,
652) -> Vec<String> {
653    let fields = pivot_field_output_names(pivot);
654    let aggregations = if pivot.using.is_empty() {
655        &pivot.expressions
656    } else {
657        &pivot.using
658    };
659
660    if fields.is_empty() || aggregations.is_empty() {
661        return Vec::new();
662    }
663
664    let needs_suffix = aggregations.len() > 1;
665    let mut outputs = Vec::new();
666    for field in fields {
667        for aggregation in aggregations {
668            if let Some(suffix) = pivot_aggregation_output_suffix(aggregation, needs_suffix) {
669                outputs.push(format!("{field}_{suffix}"));
670            } else {
671                outputs.push(field.clone());
672            }
673        }
674    }
675    outputs
676}
677
678fn pivot_field_output_names(pivot: &crate::expressions::Pivot) -> Vec<String> {
679    pivot
680        .fields
681        .iter()
682        .filter_map(|field| match field {
683            Expression::In(in_expr) => Some(
684                in_expr
685                    .expressions
686                    .iter()
687                    .filter_map(expression_name)
688                    .collect::<Vec<_>>(),
689            ),
690            _ => None,
691        })
692        .flatten()
693        .collect()
694}
695
696fn pivot_aggregation_output_suffix(expr: &Expression, needs_suffix: bool) -> Option<String> {
697    match expr {
698        Expression::Alias(alias) => Some(alias.alias.name.clone()),
699        _ if needs_suffix => pivot_generated_aggregation_suffix(expr),
700        _ => None,
701    }
702}
703
704#[cfg(feature = "generate")]
705fn pivot_generated_aggregation_suffix(expr: &Expression) -> Option<String> {
706    Generator::sql(expr).ok().map(|sql| sql.to_lowercase())
707}
708
709#[cfg(not(feature = "generate"))]
710fn pivot_generated_aggregation_suffix(expr: &Expression) -> Option<String> {
711    expression_name(expr).or_else(|| Some(expr.variant_name().to_string()))
712}
713
714fn expression_name(expr: &Expression) -> Option<String> {
715    match expr {
716        Expression::PivotAlias(alias) => expression_name(&alias.alias),
717        Expression::Alias(alias) => Some(alias.alias.name.clone()),
718        Expression::Identifier(identifier) => Some(identifier.name.clone()),
719        Expression::Column(column) => Some(column.name.name.clone()),
720        Expression::Literal(literal) => Some(literal.value_str().to_string()),
721        Expression::Var(var) => Some(var.this.clone()),
722        Expression::Tuple(tuple) => tuple.expressions.first().and_then(expression_name),
723        _ => None,
724    }
725}
726
727fn expression_column_names(expr: &Expression) -> Vec<String> {
728    expr.find_all(|node| matches!(node, Expression::Column(_)))
729        .into_iter()
730        .filter_map(|node| match node {
731            Expression::Column(column) => Some(column.name.name.clone()),
732            _ => None,
733        })
734        .collect()
735}
736
737/// Resolve a column to its source table.
738///
739/// This is a convenience function that creates a Resolver and calls get_table.
740pub fn resolve_column(
741    scope: &Scope,
742    schema: &dyn Schema,
743    column_name: &str,
744    infer_schema: bool,
745) -> Option<String> {
746    let mut resolver = Resolver::new(scope, schema, infer_schema);
747    resolver.get_table(column_name)
748}
749
750/// Check if a column is ambiguous in the given scope.
751pub fn is_column_ambiguous(scope: &Scope, schema: &dyn Schema, column_name: &str) -> bool {
752    let mut resolver = Resolver::new(scope, schema, true);
753    resolver.is_ambiguous(column_name)
754}
755
756/// Build the fully qualified table name (catalog.schema.table) from a TableRef.
757fn qualified_table_name(table: &TableRef) -> String {
758    let mut parts = Vec::new();
759    if let Some(catalog) = &table.catalog {
760        parts.push(catalog.name.clone());
761    }
762    if let Some(schema) = &table.schema {
763        parts.push(schema.name.clone());
764    }
765    parts.push(table.name.name.clone());
766    parts.join(".")
767}
768
769#[cfg(test)]
770mod tests {
771    use super::*;
772    use crate::dialects::Dialect;
773    use crate::expressions::DataType;
774    use crate::parser::Parser;
775    use crate::schema::MappingSchema;
776    use crate::scope::build_scope;
777
778    fn create_test_schema() -> MappingSchema {
779        let mut schema = MappingSchema::new();
780        // Add tables with columns
781        schema
782            .add_table(
783                "users",
784                &[
785                    (
786                        "id".to_string(),
787                        DataType::Int {
788                            length: None,
789                            integer_spelling: false,
790                        },
791                    ),
792                    ("name".to_string(), DataType::Text),
793                    ("email".to_string(), DataType::Text),
794                ],
795                None,
796            )
797            .unwrap();
798        schema
799            .add_table(
800                "orders",
801                &[
802                    (
803                        "id".to_string(),
804                        DataType::Int {
805                            length: None,
806                            integer_spelling: false,
807                        },
808                    ),
809                    (
810                        "user_id".to_string(),
811                        DataType::Int {
812                            length: None,
813                            integer_spelling: false,
814                        },
815                    ),
816                    (
817                        "amount".to_string(),
818                        DataType::Double {
819                            precision: None,
820                            scale: None,
821                        },
822                    ),
823                ],
824                None,
825            )
826            .unwrap();
827        schema
828    }
829
830    #[test]
831    fn test_resolver_basic() {
832        let ast = Parser::parse_sql("SELECT id, name FROM users").expect("Failed to parse");
833        let scope = build_scope(&ast[0]);
834        let schema = create_test_schema();
835        let mut resolver = Resolver::new(&scope, &schema, true);
836
837        // 'name' should resolve to 'users' since it's the only source
838        let table = resolver.get_table("name");
839        assert_eq!(table, Some("users".to_string()));
840    }
841
842    #[test]
843    fn test_resolver_ambiguous_column() {
844        let ast =
845            Parser::parse_sql("SELECT id FROM users JOIN orders ON users.id = orders.user_id")
846                .expect("Failed to parse");
847        let scope = build_scope(&ast[0]);
848        let schema = create_test_schema();
849        let mut resolver = Resolver::new(&scope, &schema, true);
850
851        // 'id' appears in both tables, so it's ambiguous
852        assert!(resolver.is_ambiguous("id"));
853
854        // 'name' only appears in users
855        assert!(!resolver.is_ambiguous("name"));
856
857        // 'amount' only appears in orders
858        assert!(!resolver.is_ambiguous("amount"));
859    }
860
861    #[test]
862    fn test_resolver_unambiguous_column() {
863        let ast = Parser::parse_sql(
864            "SELECT name, amount FROM users JOIN orders ON users.id = orders.user_id",
865        )
866        .expect("Failed to parse");
867        let scope = build_scope(&ast[0]);
868        let schema = create_test_schema();
869        let mut resolver = Resolver::new(&scope, &schema, true);
870
871        // 'name' should resolve to 'users'
872        let table = resolver.get_table("name");
873        assert_eq!(table, Some("users".to_string()));
874
875        // 'amount' should resolve to 'orders'
876        let table = resolver.get_table("amount");
877        assert_eq!(table, Some("orders".to_string()));
878    }
879
880    #[test]
881    fn test_resolver_with_alias() {
882        let ast = Parser::parse_sql("SELECT u.id FROM users AS u").expect("Failed to parse");
883        let scope = build_scope(&ast[0]);
884        let schema = create_test_schema();
885        let _resolver = Resolver::new(&scope, &schema, true);
886
887        // Source should be indexed by alias 'u'
888        assert!(scope.sources.contains_key("u"));
889    }
890
891    #[test]
892    fn test_sources_for_column() {
893        let ast = Parser::parse_sql("SELECT * FROM users JOIN orders ON users.id = orders.user_id")
894            .expect("Failed to parse");
895        let scope = build_scope(&ast[0]);
896        let schema = create_test_schema();
897        let mut resolver = Resolver::new(&scope, &schema, true);
898
899        // 'id' should be in both users and orders
900        let sources = resolver.sources_for_column("id");
901        assert!(sources.contains(&"users".to_string()));
902        assert!(sources.contains(&"orders".to_string()));
903
904        // 'email' should only be in users
905        let sources = resolver.sources_for_column("email");
906        assert_eq!(sources, vec!["users".to_string()]);
907    }
908
909    #[test]
910    fn test_all_columns() {
911        let ast = Parser::parse_sql("SELECT * FROM users").expect("Failed to parse");
912        let scope = build_scope(&ast[0]);
913        let schema = create_test_schema();
914        let mut resolver = Resolver::new(&scope, &schema, true);
915
916        let all = resolver.all_columns();
917        assert!(all.contains("id"));
918        assert!(all.contains("name"));
919        assert!(all.contains("email"));
920    }
921
922    #[test]
923    fn test_resolver_cte_projected_alias_column() {
924        let ast = Parser::parse_sql(
925            "WITH my_cte AS (SELECT id AS emp_id FROM users) SELECT emp_id FROM my_cte",
926        )
927        .expect("Failed to parse");
928        let scope = build_scope(&ast[0]);
929        let schema = create_test_schema();
930        let mut resolver = Resolver::new(&scope, &schema, true);
931
932        let table = resolver.get_table("emp_id");
933        assert_eq!(table, Some("my_cte".to_string()));
934    }
935
936    #[test]
937    fn test_resolve_column_helper() {
938        let ast = Parser::parse_sql("SELECT name FROM users").expect("Failed to parse");
939        let scope = build_scope(&ast[0]);
940        let schema = create_test_schema();
941
942        let table = resolve_column(&scope, &schema, "name", true);
943        assert_eq!(table, Some("users".to_string()));
944    }
945
946    #[test]
947    fn test_resolver_bigquery_mixed_case_column_names() {
948        let dialect = Dialect::get(DialectType::BigQuery);
949        let expr = dialect
950            .parse("SELECT Name AS name FROM teams")
951            .unwrap()
952            .into_iter()
953            .next()
954            .expect("expected one expression");
955        let scope = build_scope(&expr);
956
957        let mut schema = MappingSchema::with_dialect(DialectType::BigQuery);
958        schema
959            .add_table(
960                "teams",
961                &[("Name".into(), DataType::String { length: None })],
962                None,
963            )
964            .expect("schema setup");
965
966        let mut resolver = Resolver::new(&scope, &schema, true);
967        let table = resolver.get_table("Name");
968        assert_eq!(table, Some("teams".to_string()));
969
970        let table = resolver.get_table("name");
971        assert_eq!(table, Some("teams".to_string()));
972    }
973}