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

polyglot_sql/optimizer/
qualify_tables.rs
