polyglot_sql/optimizer/
qualify_tables.rs

1//! Table Qualification Module
2//!
3//! This module provides functionality for qualifying table references in SQL queries
4//! with their database and catalog names.
5//!
6//! Ported from sqlglot's optimizer/qualify_tables.py
7
8use crate::dialects::DialectType;
9use crate::expressions::{Expression, Identifier, Null, Select, Star, Subquery, TableRef};
10use crate::helper::name_sequence;
11use crate::optimizer::normalize_identifiers::{
12    get_normalization_strategy, normalize_identifier, NormalizationStrategy,
13};
14use crate::scope::traverse_scope;
15use std::collections::{HashMap, HashSet};
16
17/// Options for table qualification
18#[derive(Debug, Clone)]
19pub struct QualifyTablesOptions {
20    /// Default database name to add to unqualified tables
21    pub db: Option<String>,
22    /// Default catalog name to add to tables that have a db but no catalog
23    pub catalog: Option<String>,
24    /// The dialect to use for normalization
25    pub dialect: Option<DialectType>,
26    /// Whether to use canonical aliases (_0, _1, ...) instead of table names
27    pub canonicalize_table_aliases: bool,
28    /// Whether unaliased table references should receive aliases.
29    pub alias_unaliased_tables: bool,
30    /// Whether unaliased derived tables should receive aliases.
31    pub alias_unaliased_subqueries: bool,
32    /// Prefix used for generated aliases.
33    pub alias_prefix: String,
34    /// Whether to flatten parser-produced `SELECT * FROM (<set op>)` wrappers.
35    pub normalize_set_operation_subqueries: bool,
36}
37
38impl Default for QualifyTablesOptions {
39    fn default() -> Self {
40        Self {
41            db: None,
42            catalog: None,
43            dialect: None,
44            canonicalize_table_aliases: false,
45            alias_unaliased_tables: true,
46            alias_unaliased_subqueries: true,
47            alias_prefix: "_".to_string(),
48            normalize_set_operation_subqueries: true,
49        }
50    }
51}
52
53impl QualifyTablesOptions {
54    pub fn new() -> Self {
55        Self::default()
56    }
57
58    pub fn with_db(mut self, db: impl Into<String>) -> Self {
59        self.db = Some(db.into());
60        self
61    }
62
63    pub fn with_catalog(mut self, catalog: impl Into<String>) -> Self {
64        self.catalog = Some(catalog.into());
65        self
66    }
67
68    pub fn with_dialect(mut self, dialect: DialectType) -> Self {
69        self.dialect = Some(dialect);
70        self
71    }
72
73    pub fn with_canonical_aliases(mut self) -> Self {
74        self.canonicalize_table_aliases = true;
75        self
76    }
77
78    pub fn with_canonicalize_table_aliases(mut self, canonicalize: bool) -> Self {
79        self.canonicalize_table_aliases = canonicalize;
80        self
81    }
82
83    pub fn with_alias_unaliased_tables(mut self, alias: bool) -> Self {
84        self.alias_unaliased_tables = alias;
85        self
86    }
87
88    pub fn with_alias_unaliased_subqueries(mut self, alias: bool) -> Self {
89        self.alias_unaliased_subqueries = alias;
90        self
91    }
92
93    pub fn with_alias_prefix(mut self, prefix: impl Into<String>) -> Self {
94        let prefix = prefix.into();
95        self.alias_prefix = if prefix.is_empty() {
96            "_".to_string()
97        } else {
98            prefix
99        };
100        self
101    }
102
103    pub fn with_normalize_set_operation_subqueries(mut self, normalize: bool) -> Self {
104        self.normalize_set_operation_subqueries = normalize;
105        self
106    }
107}
108
109/// Rewrite SQL AST to have fully qualified tables.
110///
111/// This function:
112/// - Adds database/catalog prefixes to table references
113/// - Ensures all tables have aliases
114/// - Optionally canonicalizes aliases to _0, _1, etc.
115///
116/// # Examples
117///
118/// ```ignore
119/// // SELECT 1 FROM tbl -> SELECT 1 FROM db.tbl AS tbl
120/// let options = QualifyTablesOptions::new().with_db("db");
121/// let qualified = qualify_tables(expression, &options);
122/// ```
123///
124/// # Arguments
125/// * `expression` - The expression to qualify
126/// * `options` - Qualification options
127///
128/// # Returns
129/// The qualified expression
130pub fn qualify_tables(expression: Expression, options: &QualifyTablesOptions) -> Expression {
131    let strategy = get_normalization_strategy(options.dialect);
132    let alias_prefix = if options.alias_prefix.is_empty() {
133        "_"
134    } else {
135        &options.alias_prefix
136    };
137    let mut reserved_aliases: HashSet<String> = traverse_scope(&expression)
138        .into_iter()
139        .flat_map(|scope| scope.sources.into_keys())
140        .filter(|name| !name.is_empty())
141        .map(|name| normalize_identifier(Identifier::new(name), strategy).name)
142        .collect();
143    let mut alias_sequence = name_sequence(alias_prefix);
144    let mut next_alias = || loop {
145        let candidate = alias_sequence();
146        let normalized = normalize_identifier(Identifier::new(candidate.clone()), strategy).name;
147        if reserved_aliases.insert(normalized) {
148            return candidate;
149        }
150    };
151
152    qualify_tables_inner(expression, options, strategy, &mut next_alias)
153}
154
155fn qualify_tables_inner(
156    expression: Expression,
157    options: &QualifyTablesOptions,
158    strategy: NormalizationStrategy,
159    next_alias: &mut impl FnMut() -> String,
160) -> Expression {
161    match expression {
162        Expression::Select(select) => {
163            if options.normalize_set_operation_subqueries {
164                if let Some(set_operation) = unwrap_set_operation_from_passthrough_select(&select) {
165                    return qualify_tables_inner(set_operation, options, strategy, next_alias);
166                }
167            }
168
169            let qualified = qualify_select(*select, options, strategy, next_alias);
170            Expression::Select(Box::new(qualified))
171        }
172        Expression::Union(mut union) => {
173            let left = std::mem::replace(&mut union.left, Expression::Null(Null));
174            union.left = qualify_set_operation_operand(left, options, strategy, next_alias);
175            let right = std::mem::replace(&mut union.right, Expression::Null(Null));
176            union.right = qualify_set_operation_operand(right, options, strategy, next_alias);
177            Expression::Union(union)
178        }
179        Expression::Intersect(mut intersect) => {
180            let left = std::mem::replace(&mut intersect.left, Expression::Null(Null));
181            intersect.left = qualify_set_operation_operand(left, options, strategy, next_alias);
182            let right = std::mem::replace(&mut intersect.right, Expression::Null(Null));
183            intersect.right = qualify_set_operation_operand(right, options, strategy, next_alias);
184            Expression::Intersect(intersect)
185        }
186        Expression::Except(mut except) => {
187            let left = std::mem::replace(&mut except.left, Expression::Null(Null));
188            except.left = qualify_set_operation_operand(left, options, strategy, next_alias);
189            let right = std::mem::replace(&mut except.right, Expression::Null(Null));
190            except.right = qualify_set_operation_operand(right, options, strategy, next_alias);
191            Expression::Except(except)
192        }
193        _ => expression,
194    }
195}
196
197/// Qualify a SELECT expression
198fn qualify_select(
199    mut select: Select,
200    options: &QualifyTablesOptions,
201    strategy: NormalizationStrategy,
202    next_alias: &mut impl FnMut() -> String,
203) -> Select {
204    // Collect CTE names to avoid qualifying them
205    let cte_names: HashSet<String> = select
206        .with
207        .as_ref()
208        .map(|w| w.ctes.iter().map(|c| c.alias.name.clone()).collect())
209        .unwrap_or_default();
210
211    // Track canonical aliases if needed
212    let mut canonical_aliases: HashMap<String, String> = HashMap::new();
213
214    // Qualify CTEs first
215    if let Some(ref mut with) = select.with {
216        for cte in &mut with.ctes {
217            cte.this = qualify_tables_inner(cte.this.clone(), options, strategy, next_alias);
218        }
219    }
220
221    // Qualify tables in FROM clause
222    if let Some(ref mut from) = select.from {
223        for expr in &mut from.expressions {
224            *expr = qualify_table_expression(
225                expr.clone(),
226                options,
227                strategy,
228                &cte_names,
229                &mut canonical_aliases,
230                next_alias,
231            );
232        }
233    }
234
235    // Qualify tables in JOINs
236    for join in &mut select.joins {
237        join.this = qualify_table_expression(
238            join.this.clone(),
239            options,
240            strategy,
241            &cte_names,
242            &mut canonical_aliases,
243            next_alias,
244        );
245    }
246
247    // Update column references if using canonical aliases
248    if options.canonicalize_table_aliases && !canonical_aliases.is_empty() {
249        select = update_column_references(select, &canonical_aliases);
250    }
251
252    select
253}
254
255/// Qualify a table expression (Table, Subquery, etc.)
256fn qualify_table_expression(
257    expression: Expression,
258    options: &QualifyTablesOptions,
259    strategy: NormalizationStrategy,
260    cte_names: &HashSet<String>,
261    canonical_aliases: &mut HashMap<String, String>,
262    next_alias: &mut impl FnMut() -> String,
263) -> Expression {
264    match expression {
265        Expression::Table(mut table) => {
266            let table_name = table.name.name.clone();
267
268            // Don't qualify CTEs
269            if cte_names.contains(&table_name) {
270                // Still ensure it has an alias
271                ensure_table_alias(&mut table, strategy, canonical_aliases, next_alias, options);
272                return Expression::Table(table);
273            }
274
275            // Add db if specified and not already present
276            if let Some(ref db) = options.db {
277                if table.schema.is_none() {
278                    table.schema =
279                        Some(normalize_identifier(Identifier::new(db.clone()), strategy));
280                }
281            }
282
283            // Add catalog if specified, db is present, and catalog not already present
284            if let Some(ref catalog) = options.catalog {
285                if table.schema.is_some() && table.catalog.is_none() {
286                    table.catalog = Some(normalize_identifier(
287                        Identifier::new(catalog.clone()),
288                        strategy,
289                    ));
290                }
291            }
292
293            // Ensure the table has an alias
294            ensure_table_alias(&mut table, strategy, canonical_aliases, next_alias, options);
295
296            Expression::Table(table)
297        }
298        Expression::Subquery(mut subquery) => {
299            // Qualify the inner query
300            subquery.this = qualify_tables_inner(subquery.this, options, strategy, next_alias);
301
302            // Ensure the subquery has an alias
303            ensure_subquery_alias(
304                &mut subquery,
305                options,
306                strategy,
307                canonical_aliases,
308                next_alias,
309            );
310
311            Expression::Subquery(subquery)
312        }
313        Expression::Paren(mut paren) => {
314            paren.this = qualify_table_expression(
315                paren.this,
316                options,
317                strategy,
318                cte_names,
319                canonical_aliases,
320                next_alias,
321            );
322            Expression::Paren(paren)
323        }
324        _ => expression,
325    }
326}
327
328fn qualify_set_operation_operand(
329    expression: Expression,
330    options: &QualifyTablesOptions,
331    strategy: NormalizationStrategy,
332    next_alias: &mut impl FnMut() -> String,
333) -> Expression {
334    match expression {
335        Expression::Subquery(mut subquery)
336            if is_plain_unaliased_subquery(&subquery) && !is_set_operation(&subquery.this) =>
337        {
338            subquery.this = qualify_tables_inner(subquery.this, options, strategy, next_alias);
339            let mut canonical_aliases = HashMap::new();
340            ensure_subquery_alias(
341                &mut subquery,
342                options,
343                strategy,
344                &mut canonical_aliases,
345                next_alias,
346            );
347
348            let select = Select::new()
349                .column(plain_star_expression())
350                .from(Expression::Subquery(subquery));
351            Expression::Select(Box::new(select))
352        }
353        other => qualify_tables_inner(other, options, strategy, next_alias),
354    }
355}
356
357fn ensure_subquery_alias(
358    subquery: &mut Subquery,
359    options: &QualifyTablesOptions,
360    strategy: NormalizationStrategy,
361    canonical_aliases: &mut HashMap<String, String>,
362    next_alias: &mut impl FnMut() -> String,
363) {
364    if options.canonicalize_table_aliases
365        || (subquery.alias.is_none() && options.alias_unaliased_subqueries)
366    {
367        let alias_name = if options.canonicalize_table_aliases {
368            let new_name = next_alias();
369            if let Some(ref old_alias) = subquery.alias {
370                canonical_aliases.insert(old_alias.name.clone(), new_name.clone());
371            }
372            new_name
373        } else {
374            subquery
375                .alias
376                .as_ref()
377                .map(|a| a.name.clone())
378                .unwrap_or_else(|| next_alias())
379        };
380
381        subquery.alias = Some(normalize_identifier(Identifier::new(alias_name), strategy));
382        subquery.alias_explicit_as = true;
383    }
384}
385
386/// Ensure a table has an alias
387fn ensure_table_alias(
388    table: &mut TableRef,
389    strategy: NormalizationStrategy,
390    canonical_aliases: &mut HashMap<String, String>,
391    next_alias: &mut impl FnMut() -> String,
392    options: &QualifyTablesOptions,
393) {
394    let table_name = table.name.name.clone();
395
396    if options.canonicalize_table_aliases {
397        // Use canonical alias (_0, _1, etc.)
398        let new_alias = next_alias();
399        let old_alias = table
400            .alias
401            .as_ref()
402            .map(|a| a.name.clone())
403            .unwrap_or(table_name.clone());
404        canonical_aliases.insert(old_alias, new_alias.clone());
405        table.alias = Some(normalize_identifier(Identifier::new(new_alias), strategy));
406    } else if table.alias.is_none() && options.alias_unaliased_tables {
407        // Use table name as alias
408        table.alias = Some(normalize_identifier(Identifier::new(table_name), strategy));
409        table.alias_explicit_as = true;
410    }
411}
412
413fn unwrap_set_operation_from_passthrough_select(select: &Select) -> Option<Expression> {
414    if !is_passthrough_star_select(select) {
415        return None;
416    }
417
418    let source = select
419        .from
420        .as_ref()
421        .and_then(|from| from.expressions.first())?;
422
423    match source {
424        Expression::Subquery(subquery)
425            if is_plain_unaliased_subquery(subquery) && is_set_operation(&subquery.this) =>
426        {
427            Some(subquery.this.clone())
428        }
429        Expression::Paren(paren) => match &paren.this {
430            Expression::Subquery(subquery)
431                if is_plain_unaliased_subquery(subquery) && is_set_operation(&subquery.this) =>
432            {
433                Some(subquery.this.clone())
434            }
435            set_operation if is_set_operation(set_operation) => Some(set_operation.clone()),
436            _ => None,
437        },
438        set_operation if is_set_operation(set_operation) => Some(set_operation.clone()),
439        _ => None,
440    }
441}
442
443fn is_passthrough_star_select(select: &Select) -> bool {
444    if select.expressions.len() != 1 || !is_plain_star(&select.expressions[0]) {
445        return false;
446    }
447
448    if select
449        .from
450        .as_ref()
451        .map_or(true, |from| from.expressions.len() != 1)
452    {
453        return false;
454    }
455
456    let mut expected = Select::new();
457    expected.expressions = select.expressions.clone();
458    expected.from = select.from.clone();
459    *select == expected
460}
461
462fn is_plain_star(expression: &Expression) -> bool {
463    matches!(
464        expression,
465        Expression::Star(Star {
466            table: None,
467            except: None,
468            replace: None,
469            rename: None,
470            trailing_comments,
471            span: None,
472        }) if trailing_comments.is_empty()
473    )
474}
475
476fn plain_star_expression() -> Expression {
477    Expression::Star(Star {
478        table: None,
479        except: None,
480        replace: None,
481        rename: None,
482        trailing_comments: Vec::new(),
483        span: None,
484    })
485}
486
487fn is_plain_unaliased_subquery(subquery: &Subquery) -> bool {
488    subquery.alias.is_none()
489        && subquery.column_aliases.is_empty()
490        && !subquery.alias_explicit_as
491        && subquery.alias_keyword.is_none()
492        && subquery.order_by.is_none()
493        && subquery.limit.is_none()
494        && subquery.offset.is_none()
495        && subquery.distribute_by.is_none()
496        && subquery.sort_by.is_none()
497        && subquery.cluster_by.is_none()
498        && !subquery.lateral
499        && !subquery.modifiers_inside
500        && subquery.trailing_comments.is_empty()
501        && subquery.inferred_type.is_none()
502}
503
504fn is_set_operation(expression: &Expression) -> bool {
505    matches!(
506        expression,
507        Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_)
508    )
509}
510
511/// Update column references to use canonical aliases
512fn update_column_references(
513    mut select: Select,
514    canonical_aliases: &HashMap<String, String>,
515) -> Select {
516    // Update SELECT expressions
517    select.expressions = select
518        .expressions
519        .into_iter()
520        .map(|e| update_column_in_expression(e, canonical_aliases))
521        .collect();
522
523    // Update WHERE
524    if let Some(mut where_clause) = select.where_clause {
525        where_clause.this = update_column_in_expression(where_clause.this, canonical_aliases);
526        select.where_clause = Some(where_clause);
527    }
528
529    // Update GROUP BY
530    if let Some(mut group_by) = select.group_by {
531        group_by.expressions = group_by
532            .expressions
533            .into_iter()
534            .map(|e| update_column_in_expression(e, canonical_aliases))
535            .collect();
536        select.group_by = Some(group_by);
537    }
538
539    // Update HAVING
540    if let Some(mut having) = select.having {
541        having.this = update_column_in_expression(having.this, canonical_aliases);
542        select.having = Some(having);
543    }
544
545    // Update ORDER BY
546    if let Some(mut order_by) = select.order_by {
547        order_by.expressions = order_by
548            .expressions
549            .into_iter()
550            .map(|mut o| {
551                o.this = update_column_in_expression(o.this, canonical_aliases);
552                o
553            })
554            .collect();
555        select.order_by = Some(order_by);
556    }
557
558    // Update JOIN ON conditions
559    for join in &mut select.joins {
560        if let Some(on) = &mut join.on {
561            *on = update_column_in_expression(on.clone(), canonical_aliases);
562        }
563    }
564
565    select
566}
567
568/// Update column references in an expression
569fn update_column_in_expression(
570    expression: Expression,
571    canonical_aliases: &HashMap<String, String>,
572) -> Expression {
573    match expression {
574        Expression::Column(mut col) => {
575            if let Some(ref table) = col.table {
576                if let Some(canonical) = canonical_aliases.get(&table.name) {
577                    col.table = Some(Identifier {
578                        name: canonical.clone(),
579                        quoted: table.quoted,
580                        trailing_comments: table.trailing_comments.clone(),
581                        span: None,
582                    });
583                }
584            }
585            Expression::Column(col)
586        }
587        Expression::And(mut bin) => {
588            bin.left = update_column_in_expression(bin.left, canonical_aliases);
589            bin.right = update_column_in_expression(bin.right, canonical_aliases);
590            Expression::And(bin)
591        }
592        Expression::Or(mut bin) => {
593            bin.left = update_column_in_expression(bin.left, canonical_aliases);
594            bin.right = update_column_in_expression(bin.right, canonical_aliases);
595            Expression::Or(bin)
596        }
597        Expression::Eq(mut bin) => {
598            bin.left = update_column_in_expression(bin.left, canonical_aliases);
599            bin.right = update_column_in_expression(bin.right, canonical_aliases);
600            Expression::Eq(bin)
601        }
602        Expression::Neq(mut bin) => {
603            bin.left = update_column_in_expression(bin.left, canonical_aliases);
604            bin.right = update_column_in_expression(bin.right, canonical_aliases);
605            Expression::Neq(bin)
606        }
607        Expression::Lt(mut bin) => {
608            bin.left = update_column_in_expression(bin.left, canonical_aliases);
609            bin.right = update_column_in_expression(bin.right, canonical_aliases);
610            Expression::Lt(bin)
611        }
612        Expression::Lte(mut bin) => {
613            bin.left = update_column_in_expression(bin.left, canonical_aliases);
614            bin.right = update_column_in_expression(bin.right, canonical_aliases);
615            Expression::Lte(bin)
616        }
617        Expression::Gt(mut bin) => {
618            bin.left = update_column_in_expression(bin.left, canonical_aliases);
619            bin.right = update_column_in_expression(bin.right, canonical_aliases);
620            Expression::Gt(bin)
621        }
622        Expression::Gte(mut bin) => {
623            bin.left = update_column_in_expression(bin.left, canonical_aliases);
624            bin.right = update_column_in_expression(bin.right, canonical_aliases);
625            Expression::Gte(bin)
626        }
627        Expression::Not(mut un) => {
628            un.this = update_column_in_expression(un.this, canonical_aliases);
629            Expression::Not(un)
630        }
631        Expression::Paren(mut paren) => {
632            paren.this = update_column_in_expression(paren.this, canonical_aliases);
633            Expression::Paren(paren)
634        }
635        Expression::Alias(mut alias) => {
636            alias.this = update_column_in_expression(alias.this, canonical_aliases);
637            Expression::Alias(alias)
638        }
639        Expression::Function(mut func) => {
640            func.args = func
641                .args
642                .into_iter()
643                .map(|a| update_column_in_expression(a, canonical_aliases))
644                .collect();
645            Expression::Function(func)
646        }
647        Expression::AggregateFunction(mut agg) => {
648            agg.args = agg
649                .args
650                .into_iter()
651                .map(|a| update_column_in_expression(a, canonical_aliases))
652                .collect();
653            Expression::AggregateFunction(agg)
654        }
655        Expression::Case(mut case) => {
656            case.operand = case
657                .operand
658                .map(|o| update_column_in_expression(o, canonical_aliases));
659            case.whens = case
660                .whens
661                .into_iter()
662                .map(|(w, t)| {
663                    (
664                        update_column_in_expression(w, canonical_aliases),
665                        update_column_in_expression(t, canonical_aliases),
666                    )
667                })
668                .collect();
669            case.else_ = case
670                .else_
671                .map(|e| update_column_in_expression(e, canonical_aliases));
672            Expression::Case(case)
673        }
674        _ => expression,
675    }
676}
677
678#[cfg(test)]
679mod tests {
680    use super::*;
681    use crate::generator::Generator;
682    use crate::parser::Parser;
683
684    fn gen(expr: &Expression) -> String {
685        Generator::new().generate(expr).unwrap()
686    }
687
688    fn parse(sql: &str) -> Expression {
689        Parser::parse_sql(sql).expect("Failed to parse")[0].clone()
690    }
691
692    #[test]
693    fn test_qualify_with_db() {
694        let options = QualifyTablesOptions::new().with_db("mydb");
695        let expr = parse("SELECT * FROM users");
696        let qualified = qualify_tables(expr, &options);
697        let sql = gen(&qualified);
698        // Should contain mydb.users
699        assert!(sql.contains("mydb") && sql.contains("users"));
700    }
701
702    #[test]
703    fn test_qualify_with_db_and_catalog() {
704        let options = QualifyTablesOptions::new()
705            .with_db("mydb")
706            .with_catalog("mycatalog");
707        let expr = parse("SELECT * FROM users");
708        let qualified = qualify_tables(expr, &options);
709        let sql = gen(&qualified);
710        // Should contain mycatalog.mydb.users
711        assert!(sql.contains("mycatalog") && sql.contains("mydb") && sql.contains("users"));
712    }
713
714    #[test]
715    fn test_preserve_existing_schema() {
716        let options = QualifyTablesOptions::new().with_db("default_db");
717        let expr = parse("SELECT * FROM other_db.users");
718        let qualified = qualify_tables(expr, &options);
719        let sql = gen(&qualified);
720        // Should preserve other_db, not add default_db
721        assert!(sql.contains("other_db"));
722        assert!(!sql.contains("default_db"));
723    }
724
725    #[test]
726    fn test_ensure_table_alias() {
727        let options = QualifyTablesOptions::new();
728        let expr = parse("SELECT * FROM users");
729        let qualified = qualify_tables(expr, &options);
730        let sql = gen(&qualified);
731        // Should have alias (AS users)
732        assert!(sql.contains("AS") || sql.to_lowercase().contains(" users"));
733    }
734
735    #[test]
736    fn test_canonical_aliases() {
737        let options = QualifyTablesOptions::new().with_canonical_aliases();
738        let expr = parse("SELECT u.id FROM users u");
739        let qualified = qualify_tables(expr, &options);
740        let sql = gen(&qualified);
741        // Should use canonical alias like _0
742        assert!(sql.contains("_0"));
743    }
744
745    #[test]
746    fn test_qualify_join() {
747        let options = QualifyTablesOptions::new().with_db("mydb");
748        let expr = parse("SELECT * FROM users JOIN orders ON users.id = orders.user_id");
749        let qualified = qualify_tables(expr, &options);
750        let sql = gen(&qualified);
751        // Both tables should be qualified
752        assert!(sql.contains("mydb"));
753    }
754
755    #[test]
756    fn test_dont_qualify_cte() {
757        let options = QualifyTablesOptions::new().with_db("mydb");
758        let expr = parse("WITH cte AS (SELECT 1) SELECT * FROM cte");
759        let qualified = qualify_tables(expr, &options);
760        let sql = gen(&qualified);
761        // CTE reference should not be qualified with mydb
762        // The CTE definition might have mydb, but the SELECT FROM cte should not
763        assert!(sql.contains("cte"));
764    }
765
766    #[test]
767    fn test_qualify_subquery() {
768        let options = QualifyTablesOptions::new().with_db("mydb");
769        let expr = parse("SELECT * FROM (SELECT * FROM users) AS sub");
770        let qualified = qualify_tables(expr, &options);
771        let sql = gen(&qualified);
772        // Inner table should be qualified
773        assert!(sql.contains("mydb"));
774    }
775
776    #[test]
777    fn test_qualify_set_operation_subqueries_with_unique_aliases() {
778        let options = QualifyTablesOptions::new();
779        let expr = parse(
780            "SELECT * FROM (SELECT * FROM tab_1) UNION ALL SELECT * FROM (SELECT * FROM tab_1)",
781        );
782        let qualified = qualify_tables(expr, &options);
783        let sql = gen(&qualified);
784
785        assert_eq!(
786            sql,
787            "SELECT * FROM (SELECT * FROM tab_1 AS tab_1) AS _0 UNION ALL SELECT * FROM (SELECT * FROM tab_1 AS tab_1) AS _1"
788        );
789    }
790
791    #[test]
792    fn test_generated_subquery_alias_avoids_existing_source_alias() {
793        let options = QualifyTablesOptions::new().with_alias_unaliased_tables(false);
794        let expr =
795            parse("SELECT _0.a, b FROM first_table AS _0 CROSS JOIN (SELECT b FROM second_table)");
796        let qualified = qualify_tables(expr, &options);
797        let sql = gen(&qualified);
798
799        assert_eq!(
800            sql,
801            "SELECT _0.a, b FROM first_table AS _0 CROSS JOIN (SELECT b FROM second_table) AS _1"
802        );
803    }
804
805    #[test]
806    fn test_canonical_set_operation_subqueries_with_unique_aliases() {
807        let options = QualifyTablesOptions::new().with_canonical_aliases();
808        let expr = parse(
809            "SELECT * FROM (SELECT * FROM tab_1) UNION ALL SELECT * FROM (SELECT * FROM tab_1)",
810        );
811        let qualified = qualify_tables(expr, &options);
812        let sql = gen(&qualified);
813
814        assert_eq!(
815            sql,
816            "SELECT * FROM (SELECT * FROM tab_1 AS _0) AS _1 UNION ALL SELECT * FROM (SELECT * FROM tab_1 AS _2) AS _3"
817        );
818    }
819
820    #[test]
821    fn test_can_disable_unaliased_table_aliases() {
822        let options = QualifyTablesOptions::new().with_alias_unaliased_tables(false);
823        let expr = parse("SELECT * FROM users");
824        let qualified = qualify_tables(expr, &options);
825        let sql = gen(&qualified);
826
827        assert_eq!(sql, "SELECT * FROM users");
828    }
829
830    #[test]
831    fn test_can_disable_unaliased_subquery_aliases() {
832        let options = QualifyTablesOptions::new().with_alias_unaliased_subqueries(false);
833        let expr = parse("SELECT * FROM (SELECT * FROM users)");
834        let qualified = qualify_tables(expr, &options);
835        let sql = gen(&qualified);
836
837        assert_eq!(sql, "SELECT * FROM (SELECT * FROM users AS users)");
838    }
839
840    #[test]
841    fn test_can_disable_set_operation_wrapper_normalization() {
842        let options = QualifyTablesOptions::new().with_normalize_set_operation_subqueries(false);
843        let expr = parse(
844            "SELECT * FROM (SELECT * FROM tab_1) UNION ALL SELECT * FROM (SELECT * FROM tab_1)",
845        );
846        let qualified = qualify_tables(expr, &options);
847        let sql = gen(&qualified);
848
849        assert_eq!(
850            sql,
851            "SELECT * FROM (SELECT * FROM (SELECT * FROM tab_1 AS tab_1) AS _0 UNION ALL SELECT * FROM (SELECT * FROM tab_1 AS tab_1) AS _1) AS _2"
852        );
853    }
854}
polyglot_sql/optimizer/qualify_tables.rs

polyglot_sql/optimizer/
qualify_tables.rs
