1use crate::dialects::DialectType;
9use crate::expressions::{
10 Alias, BinaryOp, Column, DotAccess, Expression, Identifier, Join, JoinKind, LateralView,
11 Literal, Over, Paren, Select, TableRef, VarArgFunc, With,
12};
13use crate::resolver::{Resolver, ResolverError};
14use crate::schema::{normalize_name, Schema};
15use crate::scope::{
16 build_scope_with_ctes, selected_reference_scope, traverse_scope, Scope, SourceInfo, SourceKind,
17};
18use std::cell::RefCell;
19use std::collections::{HashMap, HashSet};
20use thiserror::Error;
21
22#[derive(Debug, Error, Clone)]
24pub enum QualifyColumnsError {
25 #[error("Unknown table: {0}")]
26 UnknownTable(String),
27
28 #[error("Unknown column: {0}")]
29 UnknownColumn(String),
30
31 #[error("Ambiguous column: {0}")]
32 AmbiguousColumn(String),
33
34 #[error("Cannot automatically join: {0}")]
35 CannotAutoJoin(String),
36
37 #[error("Unknown output column: {0}")]
38 UnknownOutputColumn(String),
39
40 #[error("Column could not be resolved: {column}{for_table}")]
41 ColumnNotResolved { column: String, for_table: String },
42
43 #[error("Resolver error: {0}")]
44 ResolverError(#[from] ResolverError),
45}
46
47pub type QualifyColumnsResult<T> = Result<T, QualifyColumnsError>;
49
50#[derive(Debug, Clone, Default)]
52pub struct QualifyColumnsOptions {
53 pub expand_alias_refs: bool,
55 pub expand_stars: bool,
57 pub infer_schema: Option<bool>,
59 pub allow_partial_qualification: bool,
61 pub dialect: Option<DialectType>,
63}
64
65impl QualifyColumnsOptions {
66 pub fn new() -> Self {
68 Self {
69 expand_alias_refs: true,
70 expand_stars: true,
71 infer_schema: None,
72 allow_partial_qualification: false,
73 dialect: None,
74 }
75 }
76
77 pub fn with_expand_alias_refs(mut self, expand: bool) -> Self {
79 self.expand_alias_refs = expand;
80 self
81 }
82
83 pub fn with_expand_stars(mut self, expand: bool) -> Self {
85 self.expand_stars = expand;
86 self
87 }
88
89 pub fn with_dialect(mut self, dialect: DialectType) -> Self {
91 self.dialect = Some(dialect);
92 self
93 }
94
95 pub fn with_allow_partial(mut self, allow: bool) -> Self {
97 self.allow_partial_qualification = allow;
98 self
99 }
100}
101
102pub fn qualify_columns(
117 expression: Expression,
118 schema: &dyn Schema,
119 options: &QualifyColumnsOptions,
120) -> QualifyColumnsResult<Expression> {
121 qualify_query(expression, schema, options, &HashMap::new(), &[])
122}
123
124fn transform_in_scope<E>(
127 expression: Expression,
128 mut transform: impl FnMut(Expression) -> Result<Expression, E>,
129) -> Result<Expression, E> {
130 enum Task {
131 Visit(Expression),
132 Finish(Expression, usize),
133 }
134 let mut tasks = vec![Task::Visit(expression)];
135 let mut results = Vec::new();
136 while let Some(task) = tasks.pop() {
137 match task {
138 Task::Visit(mut expression) => {
139 if matches!(
140 expression,
141 Expression::Select(_)
142 | Expression::Union(_)
143 | Expression::Intersect(_)
144 | Expression::Except(_)
145 ) {
146 results.push(transform(expression)?);
147 continue;
148 }
149 let mut children = Vec::new();
150 crate::ast_children::for_each_child_mut(&mut expression, |child| {
151 children.push(std::mem::replace(
152 child,
153 Expression::Null(crate::expressions::Null),
154 ));
155 });
156 tasks.push(Task::Finish(expression, children.len()));
157 tasks.extend(children.into_iter().rev().map(Task::Visit));
158 }
159 Task::Finish(mut expression, count) => {
160 let mut children = results.split_off(results.len() - count).into_iter();
161 crate::ast_children::for_each_child_mut(&mut expression, |child| {
162 *child = children.next().expect("scope transform child");
163 });
164 results.push(transform(expression)?);
165 }
166 }
167 }
168 Ok(results.pop().expect("scope transform result"))
169}
170
171fn qualify_nested_queries(
172 expression: Expression,
173 schema: &dyn Schema,
174 options: &QualifyColumnsOptions,
175 ctes: &HashMap<String, SourceInfo>,
176 outer: &[Scope],
177) -> QualifyColumnsResult<Expression> {
178 transform_in_scope(expression, |node| match node {
179 Expression::Select(_)
180 | Expression::Union(_)
181 | Expression::Intersect(_)
182 | Expression::Except(_) => qualify_query(node, schema, options, ctes, outer),
183 other => Ok(other),
184 })
185}
186
187fn qualify_query(
188 expression: Expression,
189 schema: &dyn Schema,
190 options: &QualifyColumnsOptions,
191 inherited_ctes: &HashMap<String, SourceInfo>,
192 outer: &[Scope],
193) -> QualifyColumnsResult<Expression> {
194 #[cfg(feature = "stacker")]
195 {
196 let red_zone = if cfg!(debug_assertions) {
197 4 * 1024 * 1024
198 } else {
199 1024 * 1024
200 };
201 stacker::maybe_grow(red_zone, 8 * 1024 * 1024, || {
202 qualify_query_inner(expression, schema, options, inherited_ctes, outer)
203 })
204 }
205 #[cfg(not(feature = "stacker"))]
206 qualify_query_inner(expression, schema, options, inherited_ctes, outer)
207}
208
209fn is_lateral_source(expression: &Expression) -> bool {
210 match expression {
211 Expression::Lateral(_) => true,
212 Expression::Subquery(subquery) => subquery.lateral,
213 Expression::Alias(alias) => is_lateral_source(&alias.this),
214 Expression::Paren(paren) => is_lateral_source(&paren.this),
215 _ => false,
216 }
217}
218
219fn qualify_query_inner(
220 mut expression: Expression,
221 schema: &dyn Schema,
222 options: &QualifyColumnsOptions,
223 inherited_ctes: &HashMap<String, SourceInfo>,
224 outer: &[Scope],
225) -> QualifyColumnsResult<Expression> {
226 let infer_schema = options.infer_schema.unwrap_or(schema.is_empty());
227 let dialect = options.dialect.or_else(|| schema.dialect());
228 let mut ctes = inherited_ctes.clone();
229 let with = match &mut expression {
230 Expression::Select(query) => query.with.take(),
231 Expression::Union(query) => query.with.take(),
232 Expression::Intersect(query) => query.with.take(),
233 Expression::Except(query) => query.with.take(),
234 _ => return qualify_nested_queries(expression, schema, options, inherited_ctes, outer),
235 };
236 let mut with = with;
237 if let Some(with) = &mut with {
238 pushdown_cte_alias_columns_with(with);
239 for cte in &mut with.ctes {
240 let name = cte.alias.name.clone();
241 if with.recursive {
242 ctes.retain(|existing, _| {
243 normalize_name(existing, dialect, true, true)
244 != normalize_name(&name, dialect, true, true)
245 });
246 ctes.insert(
247 name.clone(),
248 SourceInfo::new(
249 Expression::Cte(Box::new(cte.clone())),
250 true,
251 SourceKind::Cte,
252 ),
253 );
254 }
255 cte.this = qualify_nested_queries(cte.this.clone(), schema, options, &ctes, outer)?;
258 ctes.retain(|existing, _| {
259 normalize_name(existing, dialect, true, true)
260 != normalize_name(&name, dialect, true, true)
261 });
262 ctes.insert(
263 name,
264 SourceInfo::new(
265 Expression::Cte(Box::new(cte.clone())),
266 true,
267 SourceKind::Cte,
268 ),
269 );
270 }
271 }
272
273 if let Expression::Select(select) = &mut expression {
274 let joins = std::mem::take(&mut select.joins);
277 if let Some(from) = &mut select.from {
278 let sources = std::mem::take(&mut from.expressions);
279 for source in sources {
280 let mut visible_outer = outer.to_vec();
281 if is_lateral_source(&source) {
282 visible_outer.push(selected_reference_scope(&build_scope_with_ctes(
283 &Expression::Select(select.clone()),
284 &ctes,
285 )));
286 }
287 let qualified =
288 qualify_nested_queries(source, schema, options, &ctes, &visible_outer)?;
289 select.from.as_mut().unwrap().expressions.push(qualified);
290 }
291 }
292 for mut join in joins {
293 let mut visible_outer = outer.to_vec();
294 if is_lateral_source(&join.this)
295 || matches!(
296 join.kind,
297 JoinKind::Lateral
298 | JoinKind::LeftLateral
299 | JoinKind::CrossApply
300 | JoinKind::OuterApply
301 )
302 {
303 visible_outer.push(selected_reference_scope(&build_scope_with_ctes(
304 &Expression::Select(select.clone()),
305 &ctes,
306 )));
307 }
308 join.this = qualify_nested_queries(join.this, schema, options, &ctes, &visible_outer)?;
309 select.joins.push(join);
310 }
311 }
312
313 if matches!(expression, Expression::Select(_)) {
314 let scope = selected_reference_scope(&build_scope_with_ctes(&expression, &ctes));
315 let mut nested_outer = outer.to_vec();
316 nested_outer.push(scope.clone());
317 let Expression::Select(select) = &mut expression else {
320 unreachable!()
321 };
322 let from = select.from.take();
323 let mut joins = std::mem::take(&mut select.joins);
324 let mut error = None;
325 crate::ast_children::for_each_child_mut(&mut expression, |child| {
326 if error.is_none() {
327 match qualify_nested_queries(child.clone(), schema, options, &ctes, &nested_outer) {
328 Ok(qualified) => *child = qualified,
329 Err(err) => error = Some(err),
330 }
331 }
332 });
333 if let Some(error) = error {
334 return Err(error);
335 }
336 for join in &mut joins {
337 if let Some(on) = &mut join.on {
338 *on = qualify_nested_queries(on.clone(), schema, options, &ctes, &nested_outer)?;
339 }
340 }
341 let Expression::Select(select) = &mut expression else {
342 unreachable!()
343 };
344 select.from = from;
345 select.joins = joins;
346 let mut resolver = Resolver::new(&scope, schema, infer_schema).with_outer_scopes(outer);
347 let column_tables = expand_using(select, &scope, &mut resolver)?;
348 if options.expand_alias_refs {
349 expand_alias_refs(select, &mut resolver, dialect)?;
350 }
351 normalize_dotted_columns_in_scope(select, &scope, &mut resolver)?;
352 qualify_columns_in_scope(
353 select,
354 &scope,
355 &mut resolver,
356 options.allow_partial_qualification,
357 )?;
358 if options.expand_stars {
359 expand_stars(select, &scope, &mut resolver, &column_tables)?;
360 }
361 qualify_outputs_select(select)?;
362 expand_group_by(select, dialect)?;
363 select.with = with;
364 } else {
365 let (left, right) = match &mut expression {
366 Expression::Union(query) => (&mut query.left, &mut query.right),
367 Expression::Intersect(query) => (&mut query.left, &mut query.right),
368 Expression::Except(query) => (&mut query.left, &mut query.right),
369 _ => unreachable!(),
370 };
371 *left = qualify_nested_queries(left.clone(), schema, options, &ctes, outer)?;
372 *right = qualify_nested_queries(right.clone(), schema, options, &ctes, outer)?;
373 match &mut expression {
374 Expression::Union(query) => query.with = with,
375 Expression::Intersect(query) => query.with = with,
376 Expression::Except(query) => query.with = with,
377 _ => unreachable!(),
378 }
379 }
380 Ok(expression)
381}
382
383pub fn validate_qualify_columns(expression: &Expression) -> QualifyColumnsResult<()> {
388 let mut all_unqualified = Vec::new();
389
390 for scope in traverse_scope(expression) {
391 if let Expression::Select(_) = &scope.expression {
392 let unqualified = get_unqualified_columns(&scope);
394
395 let external = get_external_columns(&scope);
397 if !external.is_empty() && !is_correlated_subquery(&scope) {
398 let first = &external[0];
399 let for_table = if first.table.is_some() {
400 format!(" for table: '{}'", first.table.as_ref().unwrap())
401 } else {
402 String::new()
403 };
404 return Err(QualifyColumnsError::ColumnNotResolved {
405 column: first.name.clone(),
406 for_table,
407 });
408 }
409
410 all_unqualified.extend(unqualified);
411 }
412 }
413
414 if !all_unqualified.is_empty() {
415 let first = &all_unqualified[0];
416 return Err(QualifyColumnsError::AmbiguousColumn(first.name.clone()));
417 }
418
419 Ok(())
420}
421
422fn get_source_name(expr: &Expression) -> Option<String> {
424 match expr {
425 Expression::Table(t) => Some(
426 t.alias
427 .as_ref()
428 .map(|a| a.name.clone())
429 .unwrap_or_else(|| t.name.name.clone()),
430 ),
431 Expression::Subquery(sq) => sq.alias.as_ref().map(|a| a.name.clone()),
432 Expression::Pivot(pivot) => Some(pivot_source_name(
433 &pivot.this,
434 pivot.alias.as_ref().map(|alias| alias.name.as_str()),
435 )),
436 Expression::Unpivot(unpivot) => Some(pivot_source_name(
437 &unpivot.this,
438 unpivot.alias.as_ref().map(|alias| alias.name.as_str()),
439 )),
440 _ => None,
441 }
442}
443
444fn pivot_source_name(source: &Expression, explicit_alias: Option<&str>) -> String {
445 if let Some(alias) = explicit_alias {
446 return alias.to_string();
447 }
448
449 match source {
450 Expression::Table(table) => table
451 .alias
452 .as_ref()
453 .map(|alias| alias.name.clone())
454 .unwrap_or_else(|| table.name.name.clone()),
455 Expression::Subquery(subquery) => subquery
456 .alias
457 .as_ref()
458 .map(|alias| alias.name.clone())
459 .unwrap_or_else(|| "_0".to_string()),
460 Expression::Paren(paren) => pivot_source_name(&paren.this, explicit_alias),
461 _ => "_0".to_string(),
462 }
463}
464
465fn get_ordered_source_names(select: &Select) -> Vec<String> {
468 let mut ordered = Vec::new();
469 if let Some(from) = &select.from {
470 for expr in &from.expressions {
471 if let Some(name) = get_source_name(expr) {
472 ordered.push(name);
473 }
474 }
475 }
476 for join in &select.joins {
477 if is_semi_or_anti_join_kind(join.kind) {
478 continue;
479 }
480 if let Some(name) = get_source_name(&join.this) {
481 ordered.push(name);
482 }
483 }
484 ordered
485}
486
487fn is_semi_or_anti_join_kind(kind: JoinKind) -> bool {
488 matches!(
489 kind,
490 JoinKind::Semi
491 | JoinKind::Anti
492 | JoinKind::LeftSemi
493 | JoinKind::LeftAnti
494 | JoinKind::RightSemi
495 | JoinKind::RightAnti
496 )
497}
498
499fn make_coalesce(column_name: &str, tables: &[String]) -> Expression {
501 let args: Vec<Expression> = tables
502 .iter()
503 .map(|t| Expression::qualified_column(t.as_str(), column_name))
504 .collect();
505 Expression::Coalesce(Box::new(VarArgFunc {
506 expressions: args,
507 original_name: None,
508 inferred_type: None,
509 }))
510}
511
512fn expand_using(
518 select: &mut Select,
519 _scope: &Scope,
520 resolver: &mut Resolver,
521) -> QualifyColumnsResult<HashMap<String, Vec<String>>> {
522 let mut columns: HashMap<String, String> = HashMap::new();
525 let mut column_order: Vec<String> = Vec::new();
528
529 let mut column_tables: HashMap<String, Vec<String>> = HashMap::new();
532
533 let join_names: HashSet<String> = select
535 .joins
536 .iter()
537 .filter_map(|j| get_source_name(&j.this))
538 .collect();
539
540 let all_ordered = get_ordered_source_names(select);
541 let mut ordered: Vec<String> = all_ordered
542 .iter()
543 .filter(|name| !join_names.contains(name.as_str()))
544 .cloned()
545 .collect();
546 let mut accumulated_schema_known = true;
547
548 if join_names.is_empty() {
549 return Ok(column_tables);
550 }
551
552 fn update_source_columns(
554 source_name: &str,
555 columns: &mut HashMap<String, String>,
556 column_order: &mut Vec<String>,
557 resolver: &mut Resolver,
558 ) -> bool {
559 let Ok(source_cols) = resolver.get_source_columns(source_name) else {
560 return false;
561 };
562 let schema_known = !source_cols.is_empty()
563 && !source_cols
564 .iter()
565 .any(|column| column == "*" || column.is_empty());
566 for col_name in source_cols {
567 let normalized = normalize_column_name(&col_name, resolver.dialect);
568 if let std::collections::hash_map::Entry::Vacant(entry) = columns.entry(normalized) {
569 entry.insert(source_name.to_string());
570 column_order.push(col_name);
571 }
572 }
573 schema_known
574 }
575
576 for source_name in &ordered {
578 accumulated_schema_known &=
579 update_source_columns(source_name, &mut columns, &mut column_order, resolver);
580 }
581
582 for i in 0..select.joins.len() {
583 let source_table = ordered.last().cloned().unwrap_or_default();
585 if !source_table.is_empty() {
586 accumulated_schema_known &=
587 update_source_columns(&source_table, &mut columns, &mut column_order, resolver);
588 }
589
590 let join_table = get_source_name(&select.joins[i].this).unwrap_or_default();
592 ordered.push(join_table.clone());
593
594 let join_columns: Vec<String> =
595 resolver.get_source_columns(&join_table).unwrap_or_default();
596 let star = normalize_column_name("*", resolver.dialect);
597 let right_schema_known = !join_columns.is_empty()
598 && !join_columns
599 .iter()
600 .any(|column| normalize_column_name(column, resolver.dialect) == star);
601
602 let join_kind = select.joins[i].kind;
607 if select.joins[i].using.is_empty()
608 && matches!(
609 join_kind,
610 JoinKind::Natural
611 | JoinKind::NaturalLeft
612 | JoinKind::NaturalRight
613 | JoinKind::NaturalFull
614 )
615 {
616 let left_schema_known =
617 accumulated_schema_known && !columns.is_empty() && !columns.contains_key(&star);
618
619 if left_schema_known && right_schema_known {
620 let right_columns: HashSet<String> = join_columns
621 .iter()
622 .map(|column| normalize_column_name(column, resolver.dialect))
623 .collect();
624 let implicit_using: Vec<Identifier> = column_order
625 .iter()
626 .filter(|column| {
627 right_columns.contains(&normalize_column_name(column, resolver.dialect))
628 })
629 .map(|column| Identifier::new(column))
630 .collect();
631
632 if !implicit_using.is_empty() {
633 select.joins[i].using = implicit_using;
634 select.joins[i].kind = match join_kind {
635 JoinKind::Natural => JoinKind::Inner,
636 JoinKind::NaturalLeft => JoinKind::Left,
637 JoinKind::NaturalRight => JoinKind::Right,
638 JoinKind::NaturalFull => JoinKind::Full,
639 _ => unreachable!("checked NATURAL join kind above"),
640 };
641 }
642 }
643 }
644
645 if select.joins[i].using.is_empty() {
648 accumulated_schema_known &= right_schema_known;
649 continue;
650 }
651
652 let using_identifiers: Vec<String> = select.joins[i]
653 .using
654 .iter()
655 .map(|id| id.name.clone())
656 .collect();
657
658 let using_count = using_identifiers.len();
659 let is_semi_or_anti = matches!(
660 select.joins[i].kind,
661 crate::expressions::JoinKind::Semi
662 | crate::expressions::JoinKind::Anti
663 | crate::expressions::JoinKind::LeftSemi
664 | crate::expressions::JoinKind::LeftAnti
665 | crate::expressions::JoinKind::RightSemi
666 | crate::expressions::JoinKind::RightAnti
667 );
668
669 let mut conditions: Vec<Expression> = Vec::new();
670
671 for identifier in &using_identifiers {
672 let normalized_identifier = normalize_column_name(identifier, resolver.dialect);
673 let table = columns
674 .get(&normalized_identifier)
675 .cloned()
676 .unwrap_or_else(|| source_table.clone());
677
678 let lhs = if i == 0 || using_count == 1 {
680 Expression::qualified_column(table.as_str(), identifier.as_str())
682 } else {
683 let coalesce_cols: Vec<String> = ordered[..ordered.len() - 1]
686 .iter()
687 .filter(|t| {
688 resolver
689 .get_source_columns(t)
690 .unwrap_or_default()
691 .iter()
692 .any(|column| {
693 normalize_column_name(column, resolver.dialect)
694 == normalized_identifier
695 })
696 })
697 .cloned()
698 .collect();
699
700 if coalesce_cols.len() > 1 {
701 make_coalesce(identifier, &coalesce_cols)
702 } else {
703 Expression::qualified_column(table.as_str(), identifier.as_str())
704 }
705 };
706
707 let rhs = Expression::qualified_column(join_table.as_str(), identifier.as_str());
709
710 conditions.push(Expression::Eq(Box::new(BinaryOp::new(lhs, rhs))));
711
712 if !is_semi_or_anti {
714 let tables = column_tables
715 .entry(normalized_identifier)
716 .or_insert_with(Vec::new);
717 if !tables.contains(&table) {
718 tables.push(table.clone());
719 }
720 if !tables.contains(&join_table) {
721 tables.push(join_table.clone());
722 }
723 }
724 }
725
726 let on_condition = conditions
728 .into_iter()
729 .reduce(|acc, cond| Expression::And(Box::new(BinaryOp::new(acc, cond))))
730 .expect("at least one USING column");
731
732 select.joins[i].on = Some(on_condition);
734 select.joins[i].using = vec![];
735 accumulated_schema_known &= right_schema_known;
736 }
737
738 if !column_tables.is_empty() {
740 let mut new_expressions = Vec::with_capacity(select.expressions.len());
742 for expr in &select.expressions {
743 match expr {
744 Expression::Column(col) if col.table.is_none() => {
745 let normalized = normalize_column_name(&col.name.name, resolver.dialect);
746 let Some(tables) = column_tables.get(&normalized) else {
747 new_expressions.push(expr.clone());
748 continue;
749 };
750 let coalesce = make_coalesce(&col.name.name, tables);
751 new_expressions.push(Expression::Alias(Box::new(Alias {
753 this: coalesce,
754 alias: col.name.clone(),
755 column_aliases: vec![],
756 alias_explicit_as: false,
757 alias_keyword: None,
758 pre_alias_comments: vec![],
759 trailing_comments: vec![],
760 inferred_type: None,
761 })));
762 }
763 _ => {
764 let mut rewritten = expr.clone();
765 rewrite_using_columns_in_expression(
766 &mut rewritten,
767 &column_tables,
768 resolver.dialect,
769 );
770 new_expressions.push(rewritten);
771 }
772 }
773 }
774 select.expressions = new_expressions;
775
776 if let Some(where_clause) = &mut select.where_clause {
778 rewrite_using_columns_in_expression(
779 &mut where_clause.this,
780 &column_tables,
781 resolver.dialect,
782 );
783 }
784
785 if let Some(group_by) = &mut select.group_by {
787 for expr in &mut group_by.expressions {
788 rewrite_using_columns_in_expression(expr, &column_tables, resolver.dialect);
789 }
790 }
791
792 if let Some(having) = &mut select.having {
794 rewrite_using_columns_in_expression(&mut having.this, &column_tables, resolver.dialect);
795 }
796
797 if let Some(qualify) = &mut select.qualify {
799 rewrite_using_columns_in_expression(
800 &mut qualify.this,
801 &column_tables,
802 resolver.dialect,
803 );
804 }
805
806 if let Some(order_by) = &mut select.order_by {
808 for ordered in &mut order_by.expressions {
809 rewrite_using_columns_in_expression(
810 &mut ordered.this,
811 &column_tables,
812 resolver.dialect,
813 );
814 }
815 }
816 }
817
818 Ok(column_tables)
819}
820
821fn rewrite_using_columns_in_expression(
823 expr: &mut Expression,
824 column_tables: &HashMap<String, Vec<String>>,
825 dialect: Option<DialectType>,
826) {
827 let transformed: crate::Result<_> = transform_in_scope(expr.clone(), |node| match node {
828 Expression::Column(col) if col.table.is_none() => {
829 let normalized = normalize_column_name(&col.name.name, dialect);
830 if let Some(tables) = column_tables.get(&normalized) {
831 Ok(make_coalesce(&col.name.name, tables))
832 } else {
833 Ok(Expression::Column(col))
834 }
835 }
836 other => Ok(other),
837 });
838
839 if let Ok(next) = transformed {
840 *expr = next;
841 }
842}
843
844pub(crate) fn normalize_dotted_columns_in_scope(
852 select: &mut Select,
853 scope: &Scope,
854 resolver: &mut Resolver,
855) -> QualifyColumnsResult<()> {
856 for expression in &mut select.expressions {
857 normalize_dotted_columns_in_expression(expression, scope, resolver)?;
858 }
859 if let Some(where_clause) = &mut select.where_clause {
860 normalize_dotted_columns_in_expression(&mut where_clause.this, scope, resolver)?;
861 }
862 if let Some(group_by) = &mut select.group_by {
863 for expression in &mut group_by.expressions {
864 normalize_dotted_columns_in_expression(expression, scope, resolver)?;
865 }
866 }
867 if let Some(having) = &mut select.having {
868 normalize_dotted_columns_in_expression(&mut having.this, scope, resolver)?;
869 }
870 if let Some(qualify) = &mut select.qualify {
871 normalize_dotted_columns_in_expression(&mut qualify.this, scope, resolver)?;
872 }
873 if let Some(order_by) = &mut select.order_by {
874 for ordered in &mut order_by.expressions {
875 normalize_dotted_columns_in_expression(&mut ordered.this, scope, resolver)?;
876 }
877 }
878 for join in &mut select.joins {
879 normalize_dotted_columns_in_expression(&mut join.this, scope, resolver)?;
880 if let Some(on) = &mut join.on {
881 normalize_dotted_columns_in_expression(on, scope, resolver)?;
882 }
883 }
884 Ok(())
885}
886
887pub(crate) fn normalize_dotted_columns_in_expression(
888 expression: &mut Expression,
889 scope: &Scope,
890 resolver: &mut Resolver,
891) -> QualifyColumnsResult<()> {
892 let resolver = RefCell::new(resolver);
893 let transformed = transform_in_scope(expression.clone(), |node| {
894 let Expression::Column(column) = node else {
895 return Ok(node);
896 };
897 let Some(root) = column.table.as_ref() else {
898 return Ok(Expression::Column(column));
899 };
900
901 let root_is_source = scope
902 .sources
903 .keys()
904 .any(|source| source.eq_ignore_ascii_case(&root.name));
905 if root_is_source || resolver.borrow().outer_source_exists(&root.name) {
906 return Ok(Expression::Column(column));
907 }
908
909 let Some(source_name) = resolver.borrow_mut().get_table(&root.name) else {
910 return Ok(Expression::Column(column));
911 };
912
913 let root_column = Expression::boxed_column(Column {
914 name: root.clone(),
915 table: Some(Identifier::new(source_name)),
916 join_mark: column.join_mark,
917 trailing_comments: column.trailing_comments.clone(),
918 span: column.span,
919 inferred_type: None,
920 });
921
922 Ok(Expression::Dot(Box::new(DotAccess {
923 this: root_column,
924 field: column.name.clone(),
925 inferred_type: None,
926 })))
927 })
928 .map_err(|error: crate::Error| QualifyColumnsError::CannotAutoJoin(error.to_string()))?;
929
930 *expression = transformed;
931 Ok(())
932}
933
934fn qualify_columns_in_scope(
936 select: &mut Select,
937 scope: &Scope,
938 resolver: &mut Resolver,
939 allow_partial: bool,
940) -> QualifyColumnsResult<()> {
941 for expr in &mut select.expressions {
942 qualify_columns_in_expression(expr, scope, resolver, allow_partial)?;
943 }
944 if let Some(where_clause) = &mut select.where_clause {
945 qualify_columns_in_expression(&mut where_clause.this, scope, resolver, allow_partial)?;
946 }
947 if let Some(group_by) = &mut select.group_by {
948 for expr in &mut group_by.expressions {
949 qualify_columns_in_expression(expr, scope, resolver, allow_partial)?;
950 }
951 }
952 if let Some(having) = &mut select.having {
953 qualify_columns_in_expression(&mut having.this, scope, resolver, allow_partial)?;
954 }
955 if let Some(qualify) = &mut select.qualify {
956 qualify_columns_in_expression(&mut qualify.this, scope, resolver, allow_partial)?;
957 }
958 if let Some(order_by) = &mut select.order_by {
959 for ordered in &mut order_by.expressions {
960 qualify_columns_in_expression(&mut ordered.this, scope, resolver, allow_partial)?;
961 }
962 }
963 for join in &mut select.joins {
964 qualify_columns_in_expression(&mut join.this, scope, resolver, allow_partial)?;
965 if let Some(on) = &mut join.on {
966 qualify_columns_in_expression(on, scope, resolver, allow_partial)?;
967 }
968 }
969 Ok(())
970}
971
972fn expand_alias_refs(
979 select: &mut Select,
980 _resolver: &mut Resolver,
981 _dialect: Option<DialectType>,
982) -> QualifyColumnsResult<()> {
983 let mut alias_to_expression: HashMap<String, (Expression, usize)> = HashMap::new();
984
985 for (i, expr) in select.expressions.iter_mut().enumerate() {
986 replace_alias_refs_in_expression(expr, &alias_to_expression, false);
987 if let Expression::Alias(alias) = expr {
988 alias_to_expression.insert(alias.alias.name.clone(), (alias.this.clone(), i + 1));
989 }
990 }
991
992 if let Some(where_clause) = &mut select.where_clause {
993 replace_alias_refs_in_expression(&mut where_clause.this, &alias_to_expression, false);
994 }
995 if let Some(group_by) = &mut select.group_by {
996 for expr in &mut group_by.expressions {
997 replace_alias_refs_in_expression(expr, &alias_to_expression, true);
998 }
999 }
1000 if let Some(having) = &mut select.having {
1001 replace_alias_refs_in_expression(&mut having.this, &alias_to_expression, false);
1002 }
1003 if let Some(qualify) = &mut select.qualify {
1004 replace_alias_refs_in_expression(&mut qualify.this, &alias_to_expression, false);
1005 }
1006 if let Some(order_by) = &mut select.order_by {
1007 for ordered in &mut order_by.expressions {
1008 replace_alias_refs_in_expression(&mut ordered.this, &alias_to_expression, false);
1009 }
1010 }
1011
1012 Ok(())
1013}
1014
1015fn expand_group_by(select: &mut Select, _dialect: Option<DialectType>) -> QualifyColumnsResult<()> {
1022 let projections = select.expressions.clone();
1023
1024 if let Some(group_by) = &mut select.group_by {
1025 for group_expr in &mut group_by.expressions {
1026 if let Some(index) = positional_reference(group_expr) {
1027 let replacement = select_expression_at_position(&projections, index)?;
1028 *group_expr = replacement;
1029 }
1030 }
1031 }
1032 Ok(())
1033}
1034
1035fn expand_stars(
1045 select: &mut Select,
1046 _scope: &Scope,
1047 resolver: &mut Resolver,
1048 column_tables: &HashMap<String, Vec<String>>,
1049) -> QualifyColumnsResult<()> {
1050 let mut new_selections: Vec<Expression> = Vec::new();
1051 let mut has_star = false;
1052 let mut coalesced_columns: HashSet<String> = HashSet::new();
1053
1054 let ordered_sources = get_ordered_source_names(select);
1056
1057 for expr in &select.expressions {
1058 match expr {
1059 Expression::Star(star) => {
1060 has_star = true;
1061 if let Some(table) = &star.table {
1062 let table_name = &table.name;
1063 if !ordered_sources.contains(table_name) {
1064 return Err(QualifyColumnsError::UnknownTable(table_name.clone()));
1065 }
1066 if let Ok(columns) = resolver.get_source_columns(table_name) {
1067 if columns.contains(&"*".to_string()) || columns.is_empty() {
1068 return Ok(());
1069 }
1070 for col_name in &columns {
1071 let normalized = normalize_column_name(col_name, resolver.dialect);
1072 if coalesced_columns.contains(&normalized) {
1073 continue;
1074 }
1075 if let Some(tables) = column_tables.get(&normalized) {
1076 if tables.contains(table_name) {
1077 coalesced_columns.insert(normalized);
1078 let coalesce = make_coalesce(col_name, tables);
1079 new_selections.push(Expression::Alias(Box::new(Alias {
1080 this: coalesce,
1081 alias: Identifier::new(col_name),
1082 column_aliases: vec![],
1083 alias_explicit_as: false,
1084 alias_keyword: None,
1085 pre_alias_comments: vec![],
1086 trailing_comments: vec![],
1087 inferred_type: None,
1088 })));
1089 continue;
1090 }
1091 }
1092 new_selections
1093 .push(create_qualified_column(col_name, Some(table_name)));
1094 }
1095 }
1096 } else {
1097 for source_name in &ordered_sources {
1098 if let Ok(columns) = resolver.get_source_columns(source_name) {
1099 if columns.contains(&"*".to_string()) || columns.is_empty() {
1100 return Ok(());
1101 }
1102 for col_name in &columns {
1103 let normalized = normalize_column_name(col_name, resolver.dialect);
1104 if coalesced_columns.contains(&normalized) {
1105 continue;
1107 }
1108 if let Some(tables) = column_tables.get(&normalized) {
1109 if tables.contains(source_name) {
1110 coalesced_columns.insert(normalized);
1112 let coalesce = make_coalesce(col_name, tables);
1113 new_selections.push(Expression::Alias(Box::new(Alias {
1114 this: coalesce,
1115 alias: Identifier::new(col_name),
1116 column_aliases: vec![],
1117 alias_explicit_as: false,
1118 alias_keyword: None,
1119 pre_alias_comments: vec![],
1120 trailing_comments: vec![],
1121 inferred_type: None,
1122 })));
1123 continue;
1124 }
1125 }
1126 new_selections
1127 .push(create_qualified_column(col_name, Some(source_name)));
1128 }
1129 }
1130 }
1131 }
1132 }
1133 Expression::Column(col) if is_star_column(col) => {
1134 has_star = true;
1135 if let Some(table) = &col.table {
1136 let table_name = &table.name;
1137 if !ordered_sources.contains(table_name) {
1138 return Err(QualifyColumnsError::UnknownTable(table_name.clone()));
1139 }
1140 if let Ok(columns) = resolver.get_source_columns(table_name) {
1141 if columns.contains(&"*".to_string()) || columns.is_empty() {
1142 return Ok(());
1143 }
1144 for col_name in &columns {
1145 let normalized = normalize_column_name(col_name, resolver.dialect);
1146 if coalesced_columns.contains(&normalized) {
1147 continue;
1148 }
1149 if let Some(tables) = column_tables.get(&normalized) {
1150 if tables.contains(table_name) {
1151 coalesced_columns.insert(normalized);
1152 let coalesce = make_coalesce(col_name, tables);
1153 new_selections.push(Expression::Alias(Box::new(Alias {
1154 this: coalesce,
1155 alias: Identifier::new(col_name),
1156 column_aliases: vec![],
1157 alias_explicit_as: false,
1158 alias_keyword: None,
1159 pre_alias_comments: vec![],
1160 trailing_comments: vec![],
1161 inferred_type: None,
1162 })));
1163 continue;
1164 }
1165 }
1166 new_selections
1167 .push(create_qualified_column(col_name, Some(table_name)));
1168 }
1169 }
1170 }
1171 }
1172 _ => new_selections.push(expr.clone()),
1173 }
1174 }
1175
1176 if has_star {
1177 select.expressions = new_selections;
1178 }
1179
1180 Ok(())
1181}
1182
1183pub fn qualify_outputs(scope: &Scope) -> QualifyColumnsResult<()> {
1190 if let Expression::Select(mut select) = scope.expression.clone() {
1191 qualify_outputs_select(&mut select)?;
1192 }
1193 Ok(())
1194}
1195
1196fn qualify_outputs_select(select: &mut Select) -> QualifyColumnsResult<()> {
1197 let mut new_selections: Vec<Expression> = Vec::new();
1198
1199 for (i, expr) in select.expressions.iter().enumerate() {
1200 match expr {
1201 Expression::Alias(_) => new_selections.push(expr.clone()),
1202 Expression::Column(col) => {
1203 new_selections.push(create_alias(expr.clone(), &col.name.name));
1204 }
1205 Expression::Star(_) => new_selections.push(expr.clone()),
1206 _ => {
1207 let alias_name = get_output_name(expr).unwrap_or_else(|| format!("_col_{}", i));
1208 new_selections.push(create_alias(expr.clone(), &alias_name));
1209 }
1210 }
1211 }
1212
1213 select.expressions = new_selections;
1214 Ok(())
1215}
1216
1217fn qualify_columns_in_expression(
1218 expr: &mut Expression,
1219 scope: &Scope,
1220 resolver: &mut Resolver,
1221 allow_partial: bool,
1222) -> QualifyColumnsResult<()> {
1223 let first_error: RefCell<Option<QualifyColumnsError>> = RefCell::new(None);
1224 let resolver_cell: RefCell<&mut Resolver> = RefCell::new(resolver);
1225
1226 let transformed = transform_in_scope(expr.clone(), |node| {
1227 if first_error.borrow().is_some() {
1228 return Ok(node);
1229 }
1230
1231 match node {
1232 Expression::Column(mut col) => {
1233 if let Err(err) = qualify_single_column(
1234 &mut col,
1235 scope,
1236 &mut resolver_cell.borrow_mut(),
1237 allow_partial,
1238 ) {
1239 *first_error.borrow_mut() = Some(err);
1240 }
1241 Ok(Expression::Column(col))
1242 }
1243 _ => Ok(node),
1244 }
1245 })
1246 .map_err(|err: crate::Error| QualifyColumnsError::CannotAutoJoin(err.to_string()))?;
1247
1248 if let Some(err) = first_error.into_inner() {
1249 return Err(err);
1250 }
1251
1252 *expr = transformed;
1253 Ok(())
1254}
1255
1256fn qualify_single_column(
1257 col: &mut Column,
1258 scope: &Scope,
1259 resolver: &mut Resolver,
1260 allow_partial: bool,
1261) -> QualifyColumnsResult<()> {
1262 if is_star_column(col) {
1263 return Ok(());
1264 }
1265
1266 if let Some(table) = &col.table {
1267 let table_name = &table.name;
1268 let source_name = scope.sources.keys().find(|source| {
1269 normalize_name(source, resolver.dialect, true, true)
1270 == normalize_name(table_name, resolver.dialect, true, true)
1271 });
1272 let Some(source_name) = source_name else {
1273 if resolver.outer_source_exists(table_name) {
1274 return Ok(());
1275 }
1276 return Err(QualifyColumnsError::UnknownTable(table_name.clone()));
1277 };
1278
1279 if let Ok(source_columns) = resolver.get_source_columns(source_name) {
1280 let normalized_column_name = normalize_column_name(&col.name.name, resolver.dialect);
1281 if !allow_partial
1282 && !source_columns.is_empty()
1283 && !source_columns.iter().any(|column| {
1284 normalize_column_name(column, resolver.dialect) == normalized_column_name
1285 })
1286 && !source_columns.contains(&"*".to_string())
1287 {
1288 return Err(QualifyColumnsError::UnknownColumn(col.name.name.clone()));
1289 }
1290 }
1291 return Ok(());
1292 }
1293
1294 if resolver.is_ambiguous(&col.name.name) {
1295 if allow_partial {
1296 return Ok(());
1297 }
1298 return Err(QualifyColumnsError::AmbiguousColumn(col.name.name.clone()));
1299 }
1300 if let Some(table_name) = resolver.get_table(&col.name.name) {
1301 col.table = Some(Identifier::new(table_name));
1302 return Ok(());
1303 }
1304
1305 if let Some(outer_table) = resolver.find_column_in_outer_scopes(&col.name.name) {
1306 col.table = Some(Identifier::new(outer_table));
1307 return Ok(());
1308 }
1309
1310 if !allow_partial {
1311 return Err(QualifyColumnsError::UnknownColumn(col.name.name.clone()));
1312 }
1313
1314 Ok(())
1315}
1316
1317fn normalize_column_name(name: &str, dialect: Option<DialectType>) -> String {
1318 normalize_name(name, dialect, false, true)
1319}
1320
1321fn replace_alias_refs_in_expression(
1322 expr: &mut Expression,
1323 alias_to_expression: &HashMap<String, (Expression, usize)>,
1324 literal_index: bool,
1325) {
1326 let transformed: crate::Result<_> = transform_in_scope(expr.clone(), |node| match node {
1327 Expression::Column(col) if col.table.is_none() => {
1328 if let Some((alias_expr, index)) = alias_to_expression.get(&col.name.name) {
1329 if literal_index && matches!(alias_expr, Expression::Literal(_)) {
1330 return Ok(Expression::number(*index as i64));
1331 }
1332 return Ok(Expression::Paren(Box::new(Paren {
1333 this: alias_expr.clone(),
1334 trailing_comments: vec![],
1335 })));
1336 }
1337 Ok(Expression::Column(col))
1338 }
1339 other => Ok(other),
1340 });
1341
1342 if let Ok(next) = transformed {
1343 *expr = next;
1344 }
1345}
1346
1347fn positional_reference(expr: &Expression) -> Option<usize> {
1348 match expr {
1349 Expression::Literal(lit) if matches!(lit.as_ref(), Literal::Number(_)) => {
1350 let Literal::Number(value) = lit.as_ref() else {
1351 unreachable!()
1352 };
1353 value.parse::<usize>().ok()
1354 }
1355 _ => None,
1356 }
1357}
1358
1359fn select_expression_at_position(
1360 projections: &[Expression],
1361 index: usize,
1362) -> QualifyColumnsResult<Expression> {
1363 if index == 0 || index > projections.len() {
1364 return Err(QualifyColumnsError::UnknownOutputColumn(index.to_string()));
1365 }
1366
1367 let projection = projections[index - 1].clone();
1368 Ok(match projection {
1369 Expression::Alias(alias) => alias.this.clone(),
1370 other => other,
1371 })
1372}
1373
1374fn get_reserved_words(dialect: Option<DialectType>) -> HashSet<&'static str> {
1377 let mut words: HashSet<&'static str> = [
1379 "ADD",
1381 "ALL",
1382 "ALTER",
1383 "AND",
1384 "ANY",
1385 "AS",
1386 "ASC",
1387 "BETWEEN",
1388 "BY",
1389 "CASE",
1390 "CAST",
1391 "CHECK",
1392 "COLUMN",
1393 "CONSTRAINT",
1394 "CREATE",
1395 "CROSS",
1396 "CURRENT",
1397 "CURRENT_DATE",
1398 "CURRENT_TIME",
1399 "CURRENT_TIMESTAMP",
1400 "CURRENT_USER",
1401 "DATABASE",
1402 "DEFAULT",
1403 "DELETE",
1404 "DESC",
1405 "DISTINCT",
1406 "DROP",
1407 "ELSE",
1408 "END",
1409 "ESCAPE",
1410 "EXCEPT",
1411 "EXISTS",
1412 "FALSE",
1413 "FETCH",
1414 "FOR",
1415 "FOREIGN",
1416 "FROM",
1417 "FULL",
1418 "GRANT",
1419 "GROUP",
1420 "HAVING",
1421 "IF",
1422 "IN",
1423 "INDEX",
1424 "INNER",
1425 "INSERT",
1426 "INTERSECT",
1427 "INTO",
1428 "IS",
1429 "JOIN",
1430 "KEY",
1431 "LEFT",
1432 "LIKE",
1433 "LIMIT",
1434 "NATURAL",
1435 "NOT",
1436 "NULL",
1437 "OFFSET",
1438 "ON",
1439 "OR",
1440 "ORDER",
1441 "OUTER",
1442 "PRIMARY",
1443 "REFERENCES",
1444 "REPLACE",
1445 "RETURNING",
1446 "RIGHT",
1447 "ROLLBACK",
1448 "ROW",
1449 "ROWS",
1450 "SELECT",
1451 "SESSION_USER",
1452 "SET",
1453 "SOME",
1454 "TABLE",
1455 "THEN",
1456 "TO",
1457 "TRUE",
1458 "TRUNCATE",
1459 "UNION",
1460 "UNIQUE",
1461 "UPDATE",
1462 "USING",
1463 "VALUES",
1464 "VIEW",
1465 "WHEN",
1466 "WHERE",
1467 "WINDOW",
1468 "WITH",
1469 ]
1470 .iter()
1471 .copied()
1472 .collect();
1473
1474 match dialect {
1476 Some(DialectType::MySQL) => {
1477 words.extend(
1478 [
1479 "ANALYZE",
1480 "BOTH",
1481 "CHANGE",
1482 "CONDITION",
1483 "DATABASES",
1484 "DAY_HOUR",
1485 "DAY_MICROSECOND",
1486 "DAY_MINUTE",
1487 "DAY_SECOND",
1488 "DELAYED",
1489 "DETERMINISTIC",
1490 "DIV",
1491 "DUAL",
1492 "EACH",
1493 "ELSEIF",
1494 "ENCLOSED",
1495 "EXPLAIN",
1496 "FLOAT4",
1497 "FLOAT8",
1498 "FORCE",
1499 "HOUR_MICROSECOND",
1500 "HOUR_MINUTE",
1501 "HOUR_SECOND",
1502 "IGNORE",
1503 "INFILE",
1504 "INT1",
1505 "INT2",
1506 "INT3",
1507 "INT4",
1508 "INT8",
1509 "ITERATE",
1510 "KEYS",
1511 "KILL",
1512 "LEADING",
1513 "LEAVE",
1514 "LINES",
1515 "LOAD",
1516 "LOCK",
1517 "LONG",
1518 "LONGBLOB",
1519 "LONGTEXT",
1520 "LOOP",
1521 "LOW_PRIORITY",
1522 "MATCH",
1523 "MEDIUMBLOB",
1524 "MEDIUMINT",
1525 "MEDIUMTEXT",
1526 "MINUTE_MICROSECOND",
1527 "MINUTE_SECOND",
1528 "MOD",
1529 "MODIFIES",
1530 "NO_WRITE_TO_BINLOG",
1531 "OPTIMIZE",
1532 "OPTIONALLY",
1533 "OUT",
1534 "OUTFILE",
1535 "PURGE",
1536 "READS",
1537 "REGEXP",
1538 "RELEASE",
1539 "RENAME",
1540 "REPEAT",
1541 "REQUIRE",
1542 "RESIGNAL",
1543 "RETURN",
1544 "REVOKE",
1545 "RLIKE",
1546 "SCHEMA",
1547 "SCHEMAS",
1548 "SECOND_MICROSECOND",
1549 "SENSITIVE",
1550 "SEPARATOR",
1551 "SHOW",
1552 "SIGNAL",
1553 "SPATIAL",
1554 "SQL",
1555 "SQLEXCEPTION",
1556 "SQLSTATE",
1557 "SQLWARNING",
1558 "SQL_BIG_RESULT",
1559 "SQL_CALC_FOUND_ROWS",
1560 "SQL_SMALL_RESULT",
1561 "SSL",
1562 "STARTING",
1563 "STRAIGHT_JOIN",
1564 "TERMINATED",
1565 "TINYBLOB",
1566 "TINYINT",
1567 "TINYTEXT",
1568 "TRAILING",
1569 "TRIGGER",
1570 "UNDO",
1571 "UNLOCK",
1572 "UNSIGNED",
1573 "USAGE",
1574 "UTC_DATE",
1575 "UTC_TIME",
1576 "UTC_TIMESTAMP",
1577 "VARBINARY",
1578 "VARCHARACTER",
1579 "WHILE",
1580 "WRITE",
1581 "XOR",
1582 "YEAR_MONTH",
1583 "ZEROFILL",
1584 ]
1585 .iter()
1586 .copied(),
1587 );
1588 }
1589 Some(DialectType::PostgreSQL) | Some(DialectType::CockroachDB) => {
1590 words.extend(
1591 [
1592 "ANALYSE",
1593 "ANALYZE",
1594 "ARRAY",
1595 "AUTHORIZATION",
1596 "BINARY",
1597 "BOTH",
1598 "COLLATE",
1599 "CONCURRENTLY",
1600 "DO",
1601 "FREEZE",
1602 "ILIKE",
1603 "INITIALLY",
1604 "ISNULL",
1605 "LATERAL",
1606 "LEADING",
1607 "LOCALTIME",
1608 "LOCALTIMESTAMP",
1609 "NOTNULL",
1610 "ONLY",
1611 "OVERLAPS",
1612 "PLACING",
1613 "SIMILAR",
1614 "SYMMETRIC",
1615 "TABLESAMPLE",
1616 "TRAILING",
1617 "VARIADIC",
1618 "VERBOSE",
1619 ]
1620 .iter()
1621 .copied(),
1622 );
1623 }
1624 Some(DialectType::BigQuery) => {
1625 words.extend(
1626 [
1627 "ASSERT_ROWS_MODIFIED",
1628 "COLLATE",
1629 "CONTAINS",
1630 "CUBE",
1631 "DEFINE",
1632 "ENUM",
1633 "EXTRACT",
1634 "FOLLOWING",
1635 "GROUPING",
1636 "GROUPS",
1637 "HASH",
1638 "IGNORE",
1639 "LATERAL",
1640 "LOOKUP",
1641 "MERGE",
1642 "NEW",
1643 "NO",
1644 "NULLS",
1645 "OF",
1646 "OVER",
1647 "PARTITION",
1648 "PRECEDING",
1649 "PROTO",
1650 "RANGE",
1651 "RECURSIVE",
1652 "RESPECT",
1653 "ROLLUP",
1654 "STRUCT",
1655 "TABLESAMPLE",
1656 "TREAT",
1657 "UNBOUNDED",
1658 "WITHIN",
1659 ]
1660 .iter()
1661 .copied(),
1662 );
1663 }
1664 Some(DialectType::Snowflake) => {
1665 words.extend(
1666 [
1667 "ACCOUNT",
1668 "BOTH",
1669 "CONNECT",
1670 "FOLLOWING",
1671 "ILIKE",
1672 "INCREMENT",
1673 "ISSUE",
1674 "LATERAL",
1675 "LEADING",
1676 "LOCALTIME",
1677 "LOCALTIMESTAMP",
1678 "MINUS",
1679 "QUALIFY",
1680 "REGEXP",
1681 "RLIKE",
1682 "SOME",
1683 "START",
1684 "TABLESAMPLE",
1685 "TOP",
1686 "TRAILING",
1687 "TRY_CAST",
1688 ]
1689 .iter()
1690 .copied(),
1691 );
1692 }
1693 Some(DialectType::TSQL) | Some(DialectType::Fabric) => {
1694 words.extend(
1695 [
1696 "BACKUP",
1697 "BREAK",
1698 "BROWSE",
1699 "BULK",
1700 "CASCADE",
1701 "CHECKPOINT",
1702 "CLOSE",
1703 "CLUSTERED",
1704 "COALESCE",
1705 "COMPUTE",
1706 "CONTAINS",
1707 "CONTAINSTABLE",
1708 "CONTINUE",
1709 "CONVERT",
1710 "DBCC",
1711 "DEALLOCATE",
1712 "DENY",
1713 "DISK",
1714 "DISTRIBUTED",
1715 "DUMP",
1716 "ERRLVL",
1717 "EXEC",
1718 "EXECUTE",
1719 "EXIT",
1720 "EXTERNAL",
1721 "FILE",
1722 "FILLFACTOR",
1723 "FREETEXT",
1724 "FREETEXTTABLE",
1725 "FUNCTION",
1726 "GOTO",
1727 "HOLDLOCK",
1728 "IDENTITY",
1729 "IDENTITYCOL",
1730 "IDENTITY_INSERT",
1731 "KILL",
1732 "LINENO",
1733 "MERGE",
1734 "NONCLUSTERED",
1735 "NULLIF",
1736 "OF",
1737 "OFF",
1738 "OFFSETS",
1739 "OPEN",
1740 "OPENDATASOURCE",
1741 "OPENQUERY",
1742 "OPENROWSET",
1743 "OPENXML",
1744 "OVER",
1745 "PERCENT",
1746 "PIVOT",
1747 "PLAN",
1748 "PRINT",
1749 "PROC",
1750 "PROCEDURE",
1751 "PUBLIC",
1752 "RAISERROR",
1753 "READ",
1754 "READTEXT",
1755 "RECONFIGURE",
1756 "REPLICATION",
1757 "RESTORE",
1758 "RESTRICT",
1759 "REVERT",
1760 "ROWCOUNT",
1761 "ROWGUIDCOL",
1762 "RULE",
1763 "SAVE",
1764 "SECURITYAUDIT",
1765 "SEMANTICKEYPHRASETABLE",
1766 "SEMANTICSIMILARITYDETAILSTABLE",
1767 "SEMANTICSIMILARITYTABLE",
1768 "SETUSER",
1769 "SHUTDOWN",
1770 "STATISTICS",
1771 "SYSTEM_USER",
1772 "TEXTSIZE",
1773 "TOP",
1774 "TRAN",
1775 "TRANSACTION",
1776 "TRIGGER",
1777 "TSEQUAL",
1778 "UNPIVOT",
1779 "UPDATETEXT",
1780 "WAITFOR",
1781 "WRITETEXT",
1782 ]
1783 .iter()
1784 .copied(),
1785 );
1786 }
1787 Some(DialectType::ClickHouse) => {
1788 words.extend(
1789 [
1790 "ANTI",
1791 "ARRAY",
1792 "ASOF",
1793 "FINAL",
1794 "FORMAT",
1795 "GLOBAL",
1796 "INF",
1797 "KILL",
1798 "MATERIALIZED",
1799 "NAN",
1800 "PREWHERE",
1801 "SAMPLE",
1802 "SEMI",
1803 "SETTINGS",
1804 "TOP",
1805 ]
1806 .iter()
1807 .copied(),
1808 );
1809 }
1810 Some(DialectType::DuckDB) => {
1811 words.extend(
1812 [
1813 "ANALYSE",
1814 "ANALYZE",
1815 "ARRAY",
1816 "BOTH",
1817 "LATERAL",
1818 "LEADING",
1819 "LOCALTIME",
1820 "LOCALTIMESTAMP",
1821 "PLACING",
1822 "QUALIFY",
1823 "SIMILAR",
1824 "TABLESAMPLE",
1825 "TRAILING",
1826 ]
1827 .iter()
1828 .copied(),
1829 );
1830 }
1831 Some(DialectType::Hive) | Some(DialectType::Spark) | Some(DialectType::Databricks) => {
1832 words.extend(
1833 [
1834 "BOTH",
1835 "CLUSTER",
1836 "DISTRIBUTE",
1837 "EXCHANGE",
1838 "EXTENDED",
1839 "FUNCTION",
1840 "LATERAL",
1841 "LEADING",
1842 "MACRO",
1843 "OVER",
1844 "PARTITION",
1845 "PERCENT",
1846 "RANGE",
1847 "READS",
1848 "REDUCE",
1849 "REGEXP",
1850 "REVOKE",
1851 "RLIKE",
1852 "ROLLUP",
1853 "SEMI",
1854 "SORT",
1855 "TABLESAMPLE",
1856 "TRAILING",
1857 "TRANSFORM",
1858 "UNBOUNDED",
1859 "UNIQUEJOIN",
1860 ]
1861 .iter()
1862 .copied(),
1863 );
1864 }
1865 Some(DialectType::Trino) | Some(DialectType::Presto) | Some(DialectType::Athena) => {
1866 words.extend(
1867 [
1868 "CUBE",
1869 "DEALLOCATE",
1870 "DESCRIBE",
1871 "EXECUTE",
1872 "EXTRACT",
1873 "GROUPING",
1874 "LATERAL",
1875 "LOCALTIME",
1876 "LOCALTIMESTAMP",
1877 "NORMALIZE",
1878 "PREPARE",
1879 "ROLLUP",
1880 "SOME",
1881 "TABLESAMPLE",
1882 "UESCAPE",
1883 "UNNEST",
1884 ]
1885 .iter()
1886 .copied(),
1887 );
1888 }
1889 Some(DialectType::Oracle) => {
1890 words.extend(
1891 [
1892 "ACCESS",
1893 "AUDIT",
1894 "CLUSTER",
1895 "COMMENT",
1896 "COMPRESS",
1897 "CONNECT",
1898 "EXCLUSIVE",
1899 "FILE",
1900 "IDENTIFIED",
1901 "IMMEDIATE",
1902 "INCREMENT",
1903 "INITIAL",
1904 "LEVEL",
1905 "LOCK",
1906 "LONG",
1907 "MAXEXTENTS",
1908 "MINUS",
1909 "MODE",
1910 "NOAUDIT",
1911 "NOCOMPRESS",
1912 "NOWAIT",
1913 "NUMBER",
1914 "OF",
1915 "OFFLINE",
1916 "ONLINE",
1917 "PCTFREE",
1918 "PRIOR",
1919 "RAW",
1920 "RENAME",
1921 "RESOURCE",
1922 "REVOKE",
1923 "SHARE",
1924 "SIZE",
1925 "START",
1926 "SUCCESSFUL",
1927 "SYNONYM",
1928 "SYSDATE",
1929 "TRIGGER",
1930 "UID",
1931 "VALIDATE",
1932 "VARCHAR2",
1933 "WHENEVER",
1934 ]
1935 .iter()
1936 .copied(),
1937 );
1938 }
1939 Some(DialectType::Redshift) => {
1940 words.extend(
1941 [
1942 "AZ64",
1943 "BZIP2",
1944 "DELTA",
1945 "DELTA32K",
1946 "DISTSTYLE",
1947 "ENCODE",
1948 "GZIP",
1949 "ILIKE",
1950 "LIMIT",
1951 "LUNS",
1952 "LZO",
1953 "LZOP",
1954 "MOSTLY13",
1955 "MOSTLY32",
1956 "MOSTLY8",
1957 "RAW",
1958 "SIMILAR",
1959 "SNAPSHOT",
1960 "SORTKEY",
1961 "SYSDATE",
1962 "TOP",
1963 "ZSTD",
1964 ]
1965 .iter()
1966 .copied(),
1967 );
1968 }
1969 _ => {
1970 words.extend(
1972 [
1973 "ANALYZE",
1974 "ARRAY",
1975 "BOTH",
1976 "CUBE",
1977 "GROUPING",
1978 "LATERAL",
1979 "LEADING",
1980 "LOCALTIME",
1981 "LOCALTIMESTAMP",
1982 "OVER",
1983 "PARTITION",
1984 "QUALIFY",
1985 "RANGE",
1986 "ROLLUP",
1987 "SIMILAR",
1988 "SOME",
1989 "TABLESAMPLE",
1990 "TRAILING",
1991 ]
1992 .iter()
1993 .copied(),
1994 );
1995 }
1996 }
1997
1998 words
1999}
2000
2001fn needs_quoting(name: &str, reserved_words: &HashSet<&str>) -> bool {
2009 if name.is_empty() {
2010 return false;
2011 }
2012
2013 if name.as_bytes()[0].is_ascii_digit() {
2015 return true;
2016 }
2017
2018 if !name.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'_') {
2020 return true;
2021 }
2022
2023 let upper = name.to_uppercase();
2025 reserved_words.contains(upper.as_str())
2026}
2027
2028fn maybe_quote(id: &mut Identifier, reserved_words: &HashSet<&str>) {
2030 if id.quoted || id.name.is_empty() || id.name == "*" {
2033 return;
2034 }
2035 if needs_quoting(&id.name, reserved_words) {
2036 id.quoted = true;
2037 }
2038}
2039
2040fn quote_identifiers_recursive(expr: &mut Expression, reserved_words: &HashSet<&str>) {
2042 match expr {
2043 Expression::Identifier(id) => {
2045 maybe_quote(id, reserved_words);
2046 }
2047
2048 Expression::Column(col) => {
2049 maybe_quote(&mut col.name, reserved_words);
2050 if let Some(ref mut table) = col.table {
2051 maybe_quote(table, reserved_words);
2052 }
2053 }
2054
2055 Expression::Table(table_ref) => {
2056 maybe_quote(&mut table_ref.name, reserved_words);
2057 if let Some(ref mut schema) = table_ref.schema {
2058 maybe_quote(schema, reserved_words);
2059 }
2060 if let Some(ref mut catalog) = table_ref.catalog {
2061 maybe_quote(catalog, reserved_words);
2062 }
2063 if let Some(ref mut alias) = table_ref.alias {
2064 maybe_quote(alias, reserved_words);
2065 }
2066 for ca in &mut table_ref.column_aliases {
2067 maybe_quote(ca, reserved_words);
2068 }
2069 for p in &mut table_ref.partitions {
2070 maybe_quote(p, reserved_words);
2071 }
2072 for h in &mut table_ref.hints {
2074 quote_identifiers_recursive(h, reserved_words);
2075 }
2076 if let Some(ref mut ver) = table_ref.version {
2077 quote_identifiers_recursive(&mut ver.this, reserved_words);
2078 if let Some(ref mut e) = ver.expression {
2079 quote_identifiers_recursive(e, reserved_words);
2080 }
2081 }
2082 }
2083
2084 Expression::Star(star) => {
2085 if let Some(ref mut table) = star.table {
2086 maybe_quote(table, reserved_words);
2087 }
2088 if let Some(ref mut except_ids) = star.except {
2089 for id in except_ids {
2090 maybe_quote(id, reserved_words);
2091 }
2092 }
2093 if let Some(ref mut replace_aliases) = star.replace {
2094 for alias in replace_aliases {
2095 maybe_quote(&mut alias.alias, reserved_words);
2096 quote_identifiers_recursive(&mut alias.this, reserved_words);
2097 }
2098 }
2099 if let Some(ref mut rename_pairs) = star.rename {
2100 for (from, to) in rename_pairs {
2101 maybe_quote(from, reserved_words);
2102 maybe_quote(to, reserved_words);
2103 }
2104 }
2105 }
2106
2107 Expression::Alias(alias) => {
2109 maybe_quote(&mut alias.alias, reserved_words);
2110 for ca in &mut alias.column_aliases {
2111 maybe_quote(ca, reserved_words);
2112 }
2113 quote_identifiers_recursive(&mut alias.this, reserved_words);
2114 }
2115
2116 Expression::Select(select) => {
2118 for e in &mut select.expressions {
2119 quote_identifiers_recursive(e, reserved_words);
2120 }
2121 if let Some(ref mut from) = select.from {
2122 for e in &mut from.expressions {
2123 quote_identifiers_recursive(e, reserved_words);
2124 }
2125 }
2126 for join in &mut select.joins {
2127 quote_join(join, reserved_words);
2128 }
2129 for lv in &mut select.lateral_views {
2130 quote_lateral_view(lv, reserved_words);
2131 }
2132 if let Some(ref mut prewhere) = select.prewhere {
2133 quote_identifiers_recursive(prewhere, reserved_words);
2134 }
2135 if let Some(ref mut wh) = select.where_clause {
2136 quote_identifiers_recursive(&mut wh.this, reserved_words);
2137 }
2138 if let Some(ref mut gb) = select.group_by {
2139 for e in &mut gb.expressions {
2140 quote_identifiers_recursive(e, reserved_words);
2141 }
2142 }
2143 if let Some(ref mut hv) = select.having {
2144 quote_identifiers_recursive(&mut hv.this, reserved_words);
2145 }
2146 if let Some(ref mut q) = select.qualify {
2147 quote_identifiers_recursive(&mut q.this, reserved_words);
2148 }
2149 if let Some(ref mut ob) = select.order_by {
2150 for o in &mut ob.expressions {
2151 quote_identifiers_recursive(&mut o.this, reserved_words);
2152 }
2153 }
2154 if let Some(ref mut lim) = select.limit {
2155 quote_identifiers_recursive(&mut lim.this, reserved_words);
2156 }
2157 if let Some(ref mut off) = select.offset {
2158 quote_identifiers_recursive(&mut off.this, reserved_words);
2159 }
2160 if let Some(ref mut with) = select.with {
2161 quote_with(with, reserved_words);
2162 }
2163 if let Some(ref mut windows) = select.windows {
2164 for nw in windows {
2165 maybe_quote(&mut nw.name, reserved_words);
2166 quote_over(&mut nw.spec, reserved_words);
2167 }
2168 }
2169 if let Some(ref mut distinct_on) = select.distinct_on {
2170 for e in distinct_on {
2171 quote_identifiers_recursive(e, reserved_words);
2172 }
2173 }
2174 if let Some(ref mut limit_by) = select.limit_by {
2175 for e in limit_by {
2176 quote_identifiers_recursive(e, reserved_words);
2177 }
2178 }
2179 if let Some(ref mut settings) = select.settings {
2180 for e in settings {
2181 quote_identifiers_recursive(e, reserved_words);
2182 }
2183 }
2184 if let Some(ref mut format) = select.format {
2185 quote_identifiers_recursive(format, reserved_words);
2186 }
2187 }
2188
2189 Expression::Union(u) => {
2191 quote_identifiers_recursive(&mut u.left, reserved_words);
2192 quote_identifiers_recursive(&mut u.right, reserved_words);
2193 if let Some(ref mut with) = u.with {
2194 quote_with(with, reserved_words);
2195 }
2196 if let Some(ref mut ob) = u.order_by {
2197 for o in &mut ob.expressions {
2198 quote_identifiers_recursive(&mut o.this, reserved_words);
2199 }
2200 }
2201 if let Some(ref mut lim) = u.limit {
2202 quote_identifiers_recursive(lim, reserved_words);
2203 }
2204 if let Some(ref mut off) = u.offset {
2205 quote_identifiers_recursive(off, reserved_words);
2206 }
2207 }
2208 Expression::Intersect(i) => {
2209 quote_identifiers_recursive(&mut i.left, reserved_words);
2210 quote_identifiers_recursive(&mut i.right, reserved_words);
2211 if let Some(ref mut with) = i.with {
2212 quote_with(with, reserved_words);
2213 }
2214 if let Some(ref mut ob) = i.order_by {
2215 for o in &mut ob.expressions {
2216 quote_identifiers_recursive(&mut o.this, reserved_words);
2217 }
2218 }
2219 }
2220 Expression::Except(e) => {
2221 quote_identifiers_recursive(&mut e.left, reserved_words);
2222 quote_identifiers_recursive(&mut e.right, reserved_words);
2223 if let Some(ref mut with) = e.with {
2224 quote_with(with, reserved_words);
2225 }
2226 if let Some(ref mut ob) = e.order_by {
2227 for o in &mut ob.expressions {
2228 quote_identifiers_recursive(&mut o.this, reserved_words);
2229 }
2230 }
2231 }
2232
2233 Expression::Subquery(sq) => {
2235 quote_identifiers_recursive(&mut sq.this, reserved_words);
2236 if let Some(ref mut alias) = sq.alias {
2237 maybe_quote(alias, reserved_words);
2238 }
2239 for ca in &mut sq.column_aliases {
2240 maybe_quote(ca, reserved_words);
2241 }
2242 if let Some(ref mut ob) = sq.order_by {
2243 for o in &mut ob.expressions {
2244 quote_identifiers_recursive(&mut o.this, reserved_words);
2245 }
2246 }
2247 }
2248
2249 Expression::Insert(ins) => {
2251 quote_table_ref(&mut ins.table, reserved_words);
2252 for c in &mut ins.columns {
2253 maybe_quote(c, reserved_words);
2254 }
2255 for row in &mut ins.values {
2256 for e in row {
2257 quote_identifiers_recursive(e, reserved_words);
2258 }
2259 }
2260 if let Some(ref mut q) = ins.query {
2261 quote_identifiers_recursive(q, reserved_words);
2262 }
2263 for (id, val) in &mut ins.partition {
2264 maybe_quote(id, reserved_words);
2265 if let Some(ref mut v) = val {
2266 quote_identifiers_recursive(v, reserved_words);
2267 }
2268 }
2269 for e in &mut ins.returning {
2270 quote_identifiers_recursive(e, reserved_words);
2271 }
2272 if let Some(ref mut on_conflict) = ins.on_conflict {
2273 quote_identifiers_recursive(on_conflict, reserved_words);
2274 }
2275 if let Some(ref mut with) = ins.with {
2276 quote_with(with, reserved_words);
2277 }
2278 if let Some(ref mut alias) = ins.alias {
2279 maybe_quote(alias, reserved_words);
2280 }
2281 if let Some(ref mut src_alias) = ins.source_alias {
2282 maybe_quote(src_alias, reserved_words);
2283 }
2284 }
2285
2286 Expression::Update(upd) => {
2287 quote_table_ref(&mut upd.table, reserved_words);
2288 for tr in &mut upd.extra_tables {
2289 quote_table_ref(tr, reserved_words);
2290 }
2291 for join in &mut upd.table_joins {
2292 quote_join(join, reserved_words);
2293 }
2294 for (id, val) in &mut upd.set {
2295 maybe_quote(id, reserved_words);
2296 quote_identifiers_recursive(val, reserved_words);
2297 }
2298 if let Some(ref mut from) = upd.from_clause {
2299 for e in &mut from.expressions {
2300 quote_identifiers_recursive(e, reserved_words);
2301 }
2302 }
2303 for join in &mut upd.from_joins {
2304 quote_join(join, reserved_words);
2305 }
2306 if let Some(ref mut wh) = upd.where_clause {
2307 quote_identifiers_recursive(&mut wh.this, reserved_words);
2308 }
2309 for e in &mut upd.returning {
2310 quote_identifiers_recursive(e, reserved_words);
2311 }
2312 if let Some(ref mut with) = upd.with {
2313 quote_with(with, reserved_words);
2314 }
2315 }
2316
2317 Expression::Delete(del) => {
2318 quote_table_ref(&mut del.table, reserved_words);
2319 if let Some(ref mut alias) = del.alias {
2320 maybe_quote(alias, reserved_words);
2321 }
2322 for tr in &mut del.using {
2323 quote_table_ref(tr, reserved_words);
2324 }
2325 if let Some(ref mut wh) = del.where_clause {
2326 quote_identifiers_recursive(&mut wh.this, reserved_words);
2327 }
2328 if let Some(ref mut with) = del.with {
2329 quote_with(with, reserved_words);
2330 }
2331 }
2332
2333 Expression::And(bin)
2335 | Expression::Or(bin)
2336 | Expression::Eq(bin)
2337 | Expression::Neq(bin)
2338 | Expression::Lt(bin)
2339 | Expression::Lte(bin)
2340 | Expression::Gt(bin)
2341 | Expression::Gte(bin)
2342 | Expression::Add(bin)
2343 | Expression::Sub(bin)
2344 | Expression::Mul(bin)
2345 | Expression::Div(bin)
2346 | Expression::Mod(bin)
2347 | Expression::BitwiseAnd(bin)
2348 | Expression::BitwiseOr(bin)
2349 | Expression::BitwiseXor(bin)
2350 | Expression::Concat(bin)
2351 | Expression::Adjacent(bin)
2352 | Expression::TsMatch(bin)
2353 | Expression::PropertyEQ(bin)
2354 | Expression::ArrayContainsAll(bin)
2355 | Expression::ArrayContainedBy(bin)
2356 | Expression::ArrayOverlaps(bin)
2357 | Expression::JSONBContainsAllTopKeys(bin)
2358 | Expression::JSONBContainsAnyTopKeys(bin)
2359 | Expression::JSONBDeleteAtPath(bin)
2360 | Expression::ExtendsLeft(bin)
2361 | Expression::ExtendsRight(bin)
2362 | Expression::Is(bin)
2363 | Expression::NullSafeEq(bin)
2364 | Expression::NullSafeNeq(bin)
2365 | Expression::Glob(bin)
2366 | Expression::Match(bin)
2367 | Expression::MemberOf(bin)
2368 | Expression::BitwiseLeftShift(bin)
2369 | Expression::BitwiseRightShift(bin) => {
2370 quote_identifiers_recursive(&mut bin.left, reserved_words);
2371 quote_identifiers_recursive(&mut bin.right, reserved_words);
2372 }
2373
2374 Expression::Like(like) | Expression::ILike(like) => {
2376 quote_identifiers_recursive(&mut like.left, reserved_words);
2377 quote_identifiers_recursive(&mut like.right, reserved_words);
2378 if let Some(ref mut esc) = like.escape {
2379 quote_identifiers_recursive(esc, reserved_words);
2380 }
2381 }
2382
2383 Expression::Not(un) | Expression::Neg(un) | Expression::BitwiseNot(un) => {
2385 quote_identifiers_recursive(&mut un.this, reserved_words);
2386 }
2387
2388 Expression::In(in_expr) => {
2390 quote_identifiers_recursive(&mut in_expr.this, reserved_words);
2391 for e in &mut in_expr.expressions {
2392 quote_identifiers_recursive(e, reserved_words);
2393 }
2394 if let Some(ref mut q) = in_expr.query {
2395 quote_identifiers_recursive(q, reserved_words);
2396 }
2397 if let Some(ref mut un) = in_expr.unnest {
2398 quote_identifiers_recursive(un, reserved_words);
2399 }
2400 }
2401
2402 Expression::Between(bw) => {
2403 quote_identifiers_recursive(&mut bw.this, reserved_words);
2404 quote_identifiers_recursive(&mut bw.low, reserved_words);
2405 quote_identifiers_recursive(&mut bw.high, reserved_words);
2406 }
2407
2408 Expression::IsNull(is_null) => {
2409 quote_identifiers_recursive(&mut is_null.this, reserved_words);
2410 }
2411
2412 Expression::IsTrue(is_tf) | Expression::IsFalse(is_tf) => {
2413 quote_identifiers_recursive(&mut is_tf.this, reserved_words);
2414 }
2415
2416 Expression::Exists(ex) => {
2417 quote_identifiers_recursive(&mut ex.this, reserved_words);
2418 }
2419
2420 Expression::Function(func) => {
2422 for arg in &mut func.args {
2423 quote_identifiers_recursive(arg, reserved_words);
2424 }
2425 }
2426
2427 Expression::AggregateFunction(agg) => {
2428 for arg in &mut agg.args {
2429 quote_identifiers_recursive(arg, reserved_words);
2430 }
2431 if let Some(ref mut filter) = agg.filter {
2432 quote_identifiers_recursive(filter, reserved_words);
2433 }
2434 for o in &mut agg.order_by {
2435 quote_identifiers_recursive(&mut o.this, reserved_words);
2436 }
2437 }
2438
2439 Expression::WindowFunction(wf) => {
2440 quote_identifiers_recursive(&mut wf.this, reserved_words);
2441 quote_over(&mut wf.over, reserved_words);
2442 }
2443
2444 Expression::Case(case) => {
2446 if let Some(ref mut operand) = case.operand {
2447 quote_identifiers_recursive(operand, reserved_words);
2448 }
2449 for (when, then) in &mut case.whens {
2450 quote_identifiers_recursive(when, reserved_words);
2451 quote_identifiers_recursive(then, reserved_words);
2452 }
2453 if let Some(ref mut else_) = case.else_ {
2454 quote_identifiers_recursive(else_, reserved_words);
2455 }
2456 }
2457
2458 Expression::Cast(cast) | Expression::TryCast(cast) | Expression::SafeCast(cast) => {
2460 quote_identifiers_recursive(&mut cast.this, reserved_words);
2461 if let Some(ref mut fmt) = cast.format {
2462 quote_identifiers_recursive(fmt, reserved_words);
2463 }
2464 }
2465
2466 Expression::Paren(paren) => {
2468 quote_identifiers_recursive(&mut paren.this, reserved_words);
2469 }
2470
2471 Expression::Annotated(ann) => {
2472 quote_identifiers_recursive(&mut ann.this, reserved_words);
2473 }
2474
2475 Expression::With(with) => {
2477 quote_with(with, reserved_words);
2478 }
2479
2480 Expression::Cte(cte) => {
2481 maybe_quote(&mut cte.alias, reserved_words);
2482 for c in &mut cte.columns {
2483 maybe_quote(c, reserved_words);
2484 }
2485 quote_identifiers_recursive(&mut cte.this, reserved_words);
2486 }
2487
2488 Expression::From(from) => {
2490 for e in &mut from.expressions {
2491 quote_identifiers_recursive(e, reserved_words);
2492 }
2493 }
2494
2495 Expression::Join(join) => {
2496 quote_join(join, reserved_words);
2497 }
2498
2499 Expression::JoinedTable(jt) => {
2500 quote_identifiers_recursive(&mut jt.left, reserved_words);
2501 for join in &mut jt.joins {
2502 quote_join(join, reserved_words);
2503 }
2504 if let Some(ref mut alias) = jt.alias {
2505 maybe_quote(alias, reserved_words);
2506 }
2507 }
2508
2509 Expression::Where(wh) => {
2510 quote_identifiers_recursive(&mut wh.this, reserved_words);
2511 }
2512
2513 Expression::GroupBy(gb) => {
2514 for e in &mut gb.expressions {
2515 quote_identifiers_recursive(e, reserved_words);
2516 }
2517 }
2518
2519 Expression::Having(hv) => {
2520 quote_identifiers_recursive(&mut hv.this, reserved_words);
2521 }
2522
2523 Expression::OrderBy(ob) => {
2524 for o in &mut ob.expressions {
2525 quote_identifiers_recursive(&mut o.this, reserved_words);
2526 }
2527 }
2528
2529 Expression::Ordered(ord) => {
2530 quote_identifiers_recursive(&mut ord.this, reserved_words);
2531 }
2532
2533 Expression::Limit(lim) => {
2534 quote_identifiers_recursive(&mut lim.this, reserved_words);
2535 }
2536
2537 Expression::Offset(off) => {
2538 quote_identifiers_recursive(&mut off.this, reserved_words);
2539 }
2540
2541 Expression::Qualify(q) => {
2542 quote_identifiers_recursive(&mut q.this, reserved_words);
2543 }
2544
2545 Expression::Window(ws) => {
2546 for e in &mut ws.partition_by {
2547 quote_identifiers_recursive(e, reserved_words);
2548 }
2549 for o in &mut ws.order_by {
2550 quote_identifiers_recursive(&mut o.this, reserved_words);
2551 }
2552 }
2553
2554 Expression::Over(over) => {
2555 quote_over(over, reserved_words);
2556 }
2557
2558 Expression::WithinGroup(wg) => {
2559 quote_identifiers_recursive(&mut wg.this, reserved_words);
2560 for o in &mut wg.order_by {
2561 quote_identifiers_recursive(&mut o.this, reserved_words);
2562 }
2563 }
2564
2565 Expression::Pivot(piv) => {
2567 quote_identifiers_recursive(&mut piv.this, reserved_words);
2568 for e in &mut piv.expressions {
2569 quote_identifiers_recursive(e, reserved_words);
2570 }
2571 for f in &mut piv.fields {
2572 quote_identifiers_recursive(f, reserved_words);
2573 }
2574 if let Some(ref mut alias) = piv.alias {
2575 maybe_quote(alias, reserved_words);
2576 }
2577 }
2578
2579 Expression::Unpivot(unpiv) => {
2580 quote_identifiers_recursive(&mut unpiv.this, reserved_words);
2581 maybe_quote(&mut unpiv.value_column, reserved_words);
2582 maybe_quote(&mut unpiv.name_column, reserved_words);
2583 for e in &mut unpiv.columns {
2584 quote_identifiers_recursive(e, reserved_words);
2585 }
2586 if let Some(ref mut alias) = unpiv.alias {
2587 maybe_quote(alias, reserved_words);
2588 }
2589 }
2590
2591 Expression::Values(vals) => {
2593 for tuple in &mut vals.expressions {
2594 for e in &mut tuple.expressions {
2595 quote_identifiers_recursive(e, reserved_words);
2596 }
2597 }
2598 if let Some(ref mut alias) = vals.alias {
2599 maybe_quote(alias, reserved_words);
2600 }
2601 for ca in &mut vals.column_aliases {
2602 maybe_quote(ca, reserved_words);
2603 }
2604 }
2605
2606 Expression::Array(arr) => {
2608 for e in &mut arr.expressions {
2609 quote_identifiers_recursive(e, reserved_words);
2610 }
2611 }
2612
2613 Expression::Struct(st) => {
2614 for (_name, e) in &mut st.fields {
2615 quote_identifiers_recursive(e, reserved_words);
2616 }
2617 }
2618
2619 Expression::Tuple(tup) => {
2620 for e in &mut tup.expressions {
2621 quote_identifiers_recursive(e, reserved_words);
2622 }
2623 }
2624
2625 Expression::Subscript(sub) => {
2627 quote_identifiers_recursive(&mut sub.this, reserved_words);
2628 quote_identifiers_recursive(&mut sub.index, reserved_words);
2629 }
2630
2631 Expression::Dot(dot) => {
2632 quote_identifiers_recursive(&mut dot.this, reserved_words);
2633 maybe_quote(&mut dot.field, reserved_words);
2634 }
2635
2636 Expression::ScopeResolution(sr) => {
2637 if let Some(ref mut this) = sr.this {
2638 quote_identifiers_recursive(this, reserved_words);
2639 }
2640 quote_identifiers_recursive(&mut sr.expression, reserved_words);
2641 }
2642
2643 Expression::Lateral(lat) => {
2645 quote_identifiers_recursive(&mut lat.this, reserved_words);
2646 }
2648
2649 Expression::DPipe(dpipe) => {
2651 quote_identifiers_recursive(&mut dpipe.this, reserved_words);
2652 quote_identifiers_recursive(&mut dpipe.expression, reserved_words);
2653 }
2654
2655 Expression::Merge(merge) => {
2657 quote_identifiers_recursive(&mut merge.this, reserved_words);
2658 quote_identifiers_recursive(&mut merge.using, reserved_words);
2659 if let Some(ref mut on) = merge.on {
2660 quote_identifiers_recursive(on, reserved_words);
2661 }
2662 if let Some(ref mut whens) = merge.whens {
2663 quote_identifiers_recursive(whens, reserved_words);
2664 }
2665 if let Some(ref mut with) = merge.with_ {
2666 quote_identifiers_recursive(with, reserved_words);
2667 }
2668 if let Some(ref mut ret) = merge.returning {
2669 quote_identifiers_recursive(ret, reserved_words);
2670 }
2671 }
2672
2673 Expression::LateralView(lv) => {
2675 quote_lateral_view(lv, reserved_words);
2676 }
2677
2678 Expression::Anonymous(anon) => {
2680 quote_identifiers_recursive(&mut anon.this, reserved_words);
2681 for e in &mut anon.expressions {
2682 quote_identifiers_recursive(e, reserved_words);
2683 }
2684 }
2685
2686 Expression::Filter(filter) => {
2688 quote_identifiers_recursive(&mut filter.this, reserved_words);
2689 quote_identifiers_recursive(&mut filter.expression, reserved_words);
2690 }
2691
2692 Expression::Returning(ret) => {
2694 for e in &mut ret.expressions {
2695 quote_identifiers_recursive(e, reserved_words);
2696 }
2697 }
2698
2699 Expression::BracedWildcard(inner) => {
2701 quote_identifiers_recursive(inner, reserved_words);
2702 }
2703
2704 Expression::ReturnStmt(inner) => {
2706 quote_identifiers_recursive(inner, reserved_words);
2707 }
2708
2709 Expression::Literal(_)
2711 | Expression::Boolean(_)
2712 | Expression::Null(_)
2713 | Expression::DataType(_)
2714 | Expression::Raw(_)
2715 | Expression::Placeholder(_)
2716 | Expression::CurrentDate(_)
2717 | Expression::CurrentTime(_)
2718 | Expression::CurrentTimestamp(_)
2719 | Expression::CurrentTimestampLTZ(_)
2720 | Expression::SessionUser(_)
2721 | Expression::RowNumber(_)
2722 | Expression::Rank(_)
2723 | Expression::DenseRank(_)
2724 | Expression::PercentRank(_)
2725 | Expression::CumeDist(_)
2726 | Expression::Random(_)
2727 | Expression::Pi(_)
2728 | Expression::JSONPathRoot(_) => {
2729 }
2731
2732 _ => {}
2736 }
2737}
2738
2739fn quote_join(join: &mut Join, reserved_words: &HashSet<&str>) {
2741 quote_identifiers_recursive(&mut join.this, reserved_words);
2742 if let Some(ref mut on) = join.on {
2743 quote_identifiers_recursive(on, reserved_words);
2744 }
2745 for id in &mut join.using {
2746 maybe_quote(id, reserved_words);
2747 }
2748 if let Some(ref mut mc) = join.match_condition {
2749 quote_identifiers_recursive(mc, reserved_words);
2750 }
2751 for piv in &mut join.pivots {
2752 quote_identifiers_recursive(piv, reserved_words);
2753 }
2754}
2755
2756fn quote_with(with: &mut With, reserved_words: &HashSet<&str>) {
2758 for cte in &mut with.ctes {
2759 maybe_quote(&mut cte.alias, reserved_words);
2760 for c in &mut cte.columns {
2761 maybe_quote(c, reserved_words);
2762 }
2763 for k in &mut cte.key_expressions {
2764 maybe_quote(k, reserved_words);
2765 }
2766 quote_identifiers_recursive(&mut cte.this, reserved_words);
2767 }
2768}
2769
2770fn quote_over(over: &mut Over, reserved_words: &HashSet<&str>) {
2772 if let Some(ref mut wn) = over.window_name {
2773 maybe_quote(wn, reserved_words);
2774 }
2775 for e in &mut over.partition_by {
2776 quote_identifiers_recursive(e, reserved_words);
2777 }
2778 for o in &mut over.order_by {
2779 quote_identifiers_recursive(&mut o.this, reserved_words);
2780 }
2781 if let Some(ref mut alias) = over.alias {
2782 maybe_quote(alias, reserved_words);
2783 }
2784}
2785
2786fn quote_table_ref(table_ref: &mut TableRef, reserved_words: &HashSet<&str>) {
2788 maybe_quote(&mut table_ref.name, reserved_words);
2789 if let Some(ref mut schema) = table_ref.schema {
2790 maybe_quote(schema, reserved_words);
2791 }
2792 if let Some(ref mut catalog) = table_ref.catalog {
2793 maybe_quote(catalog, reserved_words);
2794 }
2795 if let Some(ref mut alias) = table_ref.alias {
2796 maybe_quote(alias, reserved_words);
2797 }
2798 for ca in &mut table_ref.column_aliases {
2799 maybe_quote(ca, reserved_words);
2800 }
2801 for p in &mut table_ref.partitions {
2802 maybe_quote(p, reserved_words);
2803 }
2804 for h in &mut table_ref.hints {
2805 quote_identifiers_recursive(h, reserved_words);
2806 }
2807}
2808
2809fn quote_lateral_view(lv: &mut LateralView, reserved_words: &HashSet<&str>) {
2811 quote_identifiers_recursive(&mut lv.this, reserved_words);
2812 if let Some(ref mut ta) = lv.table_alias {
2813 maybe_quote(ta, reserved_words);
2814 }
2815 for ca in &mut lv.column_aliases {
2816 maybe_quote(ca, reserved_words);
2817 }
2818}
2819
2820pub fn quote_identifiers(expression: Expression, dialect: Option<DialectType>) -> Expression {
2831 let reserved_words = get_reserved_words(dialect);
2832 let mut result = expression;
2833 quote_identifiers_recursive(&mut result, &reserved_words);
2834 result
2835}
2836
2837pub fn pushdown_cte_alias_columns(_scope: &Scope) {
2842 }
2845
2846fn pushdown_cte_alias_columns_with(with: &mut With) {
2847 for cte in &mut with.ctes {
2848 if cte.columns.is_empty() {
2849 continue;
2850 }
2851
2852 if let Expression::Select(select) = &mut cte.this {
2853 let mut next_expressions = Vec::with_capacity(select.expressions.len());
2854
2855 for (i, projection) in select.expressions.iter().enumerate() {
2856 let Some(alias_name) = cte.columns.get(i) else {
2857 next_expressions.push(projection.clone());
2858 continue;
2859 };
2860
2861 match projection {
2862 Expression::Alias(existing) => {
2863 let mut aliased = existing.clone();
2864 aliased.alias = alias_name.clone();
2865 next_expressions.push(Expression::Alias(aliased));
2866 }
2867 _ => {
2868 next_expressions.push(create_alias(projection.clone(), &alias_name.name));
2869 }
2870 }
2871 }
2872
2873 select.expressions = next_expressions;
2874 }
2875 }
2876}
2877
2878fn get_scope_columns(scope: &Scope) -> Vec<ColumnRef> {
2884 let mut columns = Vec::new();
2885 collect_columns(&scope.expression, &mut columns);
2886 columns
2887}
2888
2889#[derive(Debug, Clone)]
2891struct ColumnRef {
2892 table: Option<String>,
2893 name: String,
2894}
2895
2896fn collect_columns(expr: &Expression, columns: &mut Vec<ColumnRef>) {
2898 match expr {
2899 Expression::Column(col) => {
2900 columns.push(ColumnRef {
2901 table: col.table.as_ref().map(|t| t.name.clone()),
2902 name: col.name.name.clone(),
2903 });
2904 }
2905 Expression::Select(select) => {
2906 for e in &select.expressions {
2907 collect_columns(e, columns);
2908 }
2909 if let Some(from) = &select.from {
2910 for e in &from.expressions {
2911 collect_columns(e, columns);
2912 }
2913 }
2914 if let Some(where_clause) = &select.where_clause {
2915 collect_columns(&where_clause.this, columns);
2916 }
2917 if let Some(group_by) = &select.group_by {
2918 for e in &group_by.expressions {
2919 collect_columns(e, columns);
2920 }
2921 }
2922 if let Some(having) = &select.having {
2923 collect_columns(&having.this, columns);
2924 }
2925 if let Some(order_by) = &select.order_by {
2926 for o in &order_by.expressions {
2927 collect_columns(&o.this, columns);
2928 }
2929 }
2930 for join in &select.joins {
2931 collect_columns(&join.this, columns);
2932 if let Some(on) = &join.on {
2933 collect_columns(on, columns);
2934 }
2935 }
2936 }
2937 Expression::Alias(alias) => {
2938 collect_columns(&alias.this, columns);
2939 }
2940 Expression::Function(func) => {
2941 for arg in &func.args {
2942 collect_columns(arg, columns);
2943 }
2944 }
2945 Expression::AggregateFunction(agg) => {
2946 for arg in &agg.args {
2947 collect_columns(arg, columns);
2948 }
2949 }
2950 Expression::And(bin)
2951 | Expression::Or(bin)
2952 | Expression::Eq(bin)
2953 | Expression::Neq(bin)
2954 | Expression::Lt(bin)
2955 | Expression::Lte(bin)
2956 | Expression::Gt(bin)
2957 | Expression::Gte(bin)
2958 | Expression::Add(bin)
2959 | Expression::Sub(bin)
2960 | Expression::Mul(bin)
2961 | Expression::Div(bin) => {
2962 collect_columns(&bin.left, columns);
2963 collect_columns(&bin.right, columns);
2964 }
2965 Expression::Not(unary) | Expression::Neg(unary) => {
2966 collect_columns(&unary.this, columns);
2967 }
2968 Expression::Paren(paren) => {
2969 collect_columns(&paren.this, columns);
2970 }
2971 Expression::Case(case) => {
2972 if let Some(operand) = &case.operand {
2973 collect_columns(operand, columns);
2974 }
2975 for (when, then) in &case.whens {
2976 collect_columns(when, columns);
2977 collect_columns(then, columns);
2978 }
2979 if let Some(else_) = &case.else_ {
2980 collect_columns(else_, columns);
2981 }
2982 }
2983 Expression::Cast(cast) => {
2984 collect_columns(&cast.this, columns);
2985 }
2986 Expression::In(in_expr) => {
2987 collect_columns(&in_expr.this, columns);
2988 for e in &in_expr.expressions {
2989 collect_columns(e, columns);
2990 }
2991 if let Some(query) = &in_expr.query {
2992 collect_columns(query, columns);
2993 }
2994 }
2995 Expression::Between(between) => {
2996 collect_columns(&between.this, columns);
2997 collect_columns(&between.low, columns);
2998 collect_columns(&between.high, columns);
2999 }
3000 Expression::Subquery(subquery) => {
3001 collect_columns(&subquery.this, columns);
3002 }
3003 _ => {}
3004 }
3005}
3006
3007fn get_unqualified_columns(scope: &Scope) -> Vec<ColumnRef> {
3009 get_scope_columns(scope)
3010 .into_iter()
3011 .filter(|c| c.table.is_none())
3012 .collect()
3013}
3014
3015fn get_external_columns(scope: &Scope) -> Vec<ColumnRef> {
3017 let source_names: HashSet<_> = scope.sources.keys().cloned().collect();
3018
3019 get_scope_columns(scope)
3020 .into_iter()
3021 .filter(|c| {
3022 if let Some(table) = &c.table {
3023 !source_names.contains(table)
3024 } else {
3025 false
3026 }
3027 })
3028 .collect()
3029}
3030
3031fn is_correlated_subquery(scope: &Scope) -> bool {
3033 scope.can_be_correlated && !get_external_columns(scope).is_empty()
3034}
3035
3036fn is_star_column(col: &Column) -> bool {
3038 col.name.name == "*"
3039}
3040
3041fn create_qualified_column(name: &str, table: Option<&str>) -> Expression {
3043 Expression::boxed_column(Column {
3044 name: Identifier::new(name),
3045 table: table.map(Identifier::new),
3046 join_mark: false,
3047 trailing_comments: vec![],
3048 span: None,
3049 inferred_type: None,
3050 })
3051}
3052
3053fn create_alias(expr: Expression, alias_name: &str) -> Expression {
3055 Expression::Alias(Box::new(Alias {
3056 this: expr,
3057 alias: Identifier::new(alias_name),
3058 column_aliases: vec![],
3059 alias_explicit_as: false,
3060 alias_keyword: None,
3061 pre_alias_comments: vec![],
3062 trailing_comments: vec![],
3063 inferred_type: None,
3064 }))
3065}
3066
3067fn get_output_name(expr: &Expression) -> Option<String> {
3069 match expr {
3070 Expression::Column(col) => Some(col.name.name.clone()),
3071 Expression::Alias(alias) => Some(alias.alias.name.clone()),
3072 Expression::Identifier(id) => Some(id.name.clone()),
3073 _ => None,
3074 }
3075}
3076
3077#[cfg(test)]
3078mod tests {
3079 use super::*;
3080 use crate::expressions::DataType;
3081 use crate::generator::Generator;
3082 use crate::parser::Parser;
3083 use crate::scope::build_scope;
3084 use crate::{MappingSchema, Schema};
3085
3086 fn gen(expr: &Expression) -> String {
3087 Generator::new().generate(expr).unwrap()
3088 }
3089
3090 fn parse(sql: &str) -> Expression {
3091 Parser::parse_sql(sql).expect("Failed to parse")[0].clone()
3092 }
3093
3094 #[test]
3095 fn test_qualify_columns_options() {
3096 let options = QualifyColumnsOptions::new()
3097 .with_expand_alias_refs(true)
3098 .with_expand_stars(false)
3099 .with_dialect(DialectType::PostgreSQL)
3100 .with_allow_partial(true);
3101
3102 assert!(options.expand_alias_refs);
3103 assert!(!options.expand_stars);
3104 assert_eq!(options.dialect, Some(DialectType::PostgreSQL));
3105 assert!(options.allow_partial_qualification);
3106 }
3107
3108 #[test]
3109 fn test_qualify_cte_chain_uses_lexically_selected_sources() {
3110 let mut schema = MappingSchema::new();
3111 schema
3112 .add_table("raw_orders", &[("amount".into(), DataType::Text)], None)
3113 .unwrap();
3114 let sql = "WITH transformed AS (SELECT CAST(amount AS INT) AS amount FROM raw_orders), final AS (SELECT amount FROM transformed) SELECT amount FROM final";
3115 let qualified =
3116 qualify_columns(parse(sql), &schema, &QualifyColumnsOptions::new()).unwrap();
3117 assert_eq!(gen(&qualified), "WITH transformed AS (SELECT CAST(raw_orders.amount AS INT) AS amount FROM raw_orders), final AS (SELECT transformed.amount AS amount FROM transformed) SELECT final.amount AS amount FROM final");
3118 }
3119
3120 #[test]
3121 fn test_qualify_does_not_capture_unrelated_schema_tables() {
3122 let mut schema = MappingSchema::new();
3123 schema
3124 .add_table("unrelated", &[("n".into(), DataType::Text)], None)
3125 .unwrap();
3126 for sql in ["SELECT n FROM missing", "SELECT unrelated.n FROM missing"] {
3127 assert!(
3128 qualify_columns(parse(sql), &schema, &QualifyColumnsOptions::new()).is_err(),
3129 "{sql}"
3130 );
3131 }
3132 let sql = "SELECT n FROM missing";
3133 let qualified = qualify_columns(
3134 parse(sql),
3135 &schema,
3136 &QualifyColumnsOptions::new().with_allow_partial(true),
3137 )
3138 .unwrap();
3139 assert_eq!(gen(&qualified), "SELECT n AS n FROM missing");
3140 assert!(matches!(
3141 qualify_columns(
3142 parse(
3143 "WITH a AS (SELECT 1 AS n), b AS (SELECT 2 AS n) SELECT n FROM a CROSS JOIN b"
3144 ),
3145 &schema,
3146 &QualifyColumnsOptions::new()
3147 ),
3148 Err(QualifyColumnsError::AmbiguousColumn(_))
3149 ));
3150 }
3151
3152 #[test]
3153 fn test_qualify_correlated_cte_alias_and_nested_shadowing() {
3154 let schema = MappingSchema::new();
3155 for (sql, expected) in [
3156 ("WITH a AS (SELECT 1 AS n) SELECT (SELECT n) AS result FROM a AS outer_a", "(SELECT outer_a.n AS n)"),
3157 ("WITH a AS (SELECT 1 AS n) SELECT (SELECT (SELECT n)) AS result FROM a AS outer_a", "(SELECT outer_a.n AS n)"),
3158 ("WITH a AS (SELECT 1 AS n) SELECT (WITH b AS (SELECT n) SELECT n FROM b) AS result FROM a AS outer_a", "WITH b AS (SELECT outer_a.n AS n) SELECT b.n AS n FROM b"),
3159 ("WITH a AS (SELECT 1 AS n) SELECT (WITH a AS (SELECT 2 AS m) SELECT m FROM a) AS result FROM a", "SELECT a.m AS m FROM a"),
3160 ("WITH a AS (SELECT 1 AS n) SELECT b.n FROM a JOIN LATERAL (SELECT a.n) AS b ON TRUE", "SELECT a.n AS n"),
3161 ("WITH a AS (SELECT 1 AS n) SELECT b.n FROM a, LATERAL (SELECT a.n) AS b", "SELECT a.n AS n"),
3162 ] {
3163 let qualified = qualify_columns(parse(sql), &schema, &QualifyColumnsOptions::new())
3164 .unwrap_or_else(|error| panic!("{sql}: {error}"));
3165 assert!(gen(&qualified).contains(expected), "{}", gen(&qualified));
3166 }
3167 }
3168
3169 #[test]
3170 fn test_get_scope_columns() {
3171 let expr = parse("SELECT a, b FROM t WHERE c = 1");
3172 let scope = build_scope(&expr);
3173 let columns = get_scope_columns(&scope);
3174
3175 assert!(columns.iter().any(|c| c.name == "a"));
3176 assert!(columns.iter().any(|c| c.name == "b"));
3177 assert!(columns.iter().any(|c| c.name == "c"));
3178 }
3179
3180 #[test]
3181 fn test_get_unqualified_columns() {
3182 let expr = parse("SELECT t.a, b FROM t");
3183 let scope = build_scope(&expr);
3184 let unqualified = get_unqualified_columns(&scope);
3185
3186 assert!(unqualified.iter().any(|c| c.name == "b"));
3188 assert!(!unqualified.iter().any(|c| c.name == "a"));
3189 }
3190
3191 #[test]
3192 fn test_is_star_column() {
3193 let col = Column {
3194 name: Identifier::new("*"),
3195 table: Some(Identifier::new("t")),
3196 join_mark: false,
3197 trailing_comments: vec![],
3198 span: None,
3199 inferred_type: None,
3200 };
3201 assert!(is_star_column(&col));
3202
3203 let col2 = Column {
3204 name: Identifier::new("id"),
3205 table: None,
3206 join_mark: false,
3207 trailing_comments: vec![],
3208 span: None,
3209 inferred_type: None,
3210 };
3211 assert!(!is_star_column(&col2));
3212 }
3213
3214 #[test]
3215 fn test_create_qualified_column() {
3216 let expr = create_qualified_column("id", Some("users"));
3217 let sql = gen(&expr);
3218 assert!(sql.contains("users"));
3219 assert!(sql.contains("id"));
3220 }
3221
3222 #[test]
3223 fn test_create_alias() {
3224 let col = Expression::boxed_column(Column {
3225 name: Identifier::new("value"),
3226 table: None,
3227 join_mark: false,
3228 trailing_comments: vec![],
3229 span: None,
3230 inferred_type: None,
3231 });
3232 let aliased = create_alias(col, "total");
3233 let sql = gen(&aliased);
3234 assert!(sql.contains("AS") || sql.contains("total"));
3235 }
3236
3237 #[test]
3238 fn test_qualify_columns_normalizes_struct_field_access_issue_408() {
3239 let struct_type = DataType::Struct {
3240 fields: vec![crate::expressions::StructField::new(
3241 "field_value".into(),
3242 DataType::Text,
3243 )],
3244 nested: true,
3245 };
3246 let mut schema = MappingSchema::with_dialect(DialectType::DuckDB);
3247 schema
3248 .add_table(
3249 "source_table",
3250 &[
3251 ("composite_value".into(), struct_type.clone()),
3252 (
3253 "nested_items".into(),
3254 DataType::Array {
3255 element_type: Box::new(struct_type),
3256 dimension: None,
3257 },
3258 ),
3259 ],
3260 None,
3261 )
3262 .expect("schema setup");
3263
3264 let cases = [
3265 (
3266 "SELECT composite_value.field_value AS output_value FROM source_table",
3267 "source_table",
3268 "composite_value",
3269 ),
3270 (
3271 "SELECT item.field_value AS output_value FROM source_table s \
3272 CROSS JOIN UNNEST(s.nested_items) AS expanded(item)",
3273 "expanded",
3274 "item",
3275 ),
3276 ];
3277
3278 for (sql, expected_table, expected_column) in cases {
3279 let qualified = qualify_columns(
3280 parse(sql),
3281 &schema,
3282 &QualifyColumnsOptions::new().with_dialect(DialectType::DuckDB),
3283 )
3284 .unwrap_or_else(|error| panic!("qualification failed for {sql:?}: {error}"));
3285 let Expression::Select(select) = qualified else {
3286 panic!("expected SELECT");
3287 };
3288 let Expression::Alias(alias) = &select.expressions[0] else {
3289 panic!("expected aliased projection");
3290 };
3291 let Expression::Dot(dot) = &alias.this else {
3292 panic!("expected normalized Dot, got {:?}", alias.this);
3293 };
3294 let Expression::Column(base) = &dot.this else {
3295 panic!("expected column-backed Dot");
3296 };
3297
3298 assert_eq!(
3299 base.table.as_ref().map(|table| table.name.as_str()),
3300 Some(expected_table)
3301 );
3302 assert_eq!(base.name.name, expected_column);
3303 assert_eq!(dot.field.name, "field_value");
3304 }
3305 }
3306
3307 #[test]
3308 fn test_validate_qualify_columns_success() {
3309 let expr = parse("SELECT t.a, t.b FROM t");
3311 let result = validate_qualify_columns(&expr);
3312 let _ = result;
3315 }
3316
3317 #[test]
3318 fn test_collect_columns_nested() {
3319 let expr = parse("SELECT a + b, c FROM t WHERE d > 0 GROUP BY e HAVING f = 1");
3320 let mut columns = Vec::new();
3321 collect_columns(&expr, &mut columns);
3322
3323 let names: Vec<_> = columns.iter().map(|c| c.name.as_str()).collect();
3324 assert!(names.contains(&"a"));
3325 assert!(names.contains(&"b"));
3326 assert!(names.contains(&"c"));
3327 assert!(names.contains(&"d"));
3328 assert!(names.contains(&"e"));
3329 assert!(names.contains(&"f"));
3330 }
3331
3332 #[test]
3333 fn test_collect_columns_in_case() {
3334 let expr = parse("SELECT CASE WHEN a = 1 THEN b ELSE c END FROM t");
3335 let mut columns = Vec::new();
3336 collect_columns(&expr, &mut columns);
3337
3338 let names: Vec<_> = columns.iter().map(|c| c.name.as_str()).collect();
3339 assert!(names.contains(&"a"));
3340 assert!(names.contains(&"b"));
3341 assert!(names.contains(&"c"));
3342 }
3343
3344 #[test]
3345 fn test_collect_columns_in_subquery() {
3346 let expr = parse("SELECT a FROM t WHERE b IN (SELECT c FROM s)");
3347 let mut columns = Vec::new();
3348 collect_columns(&expr, &mut columns);
3349
3350 let names: Vec<_> = columns.iter().map(|c| c.name.as_str()).collect();
3351 assert!(names.contains(&"a"));
3352 assert!(names.contains(&"b"));
3353 assert!(names.contains(&"c"));
3354 }
3355
3356 #[test]
3357 fn test_qualify_outputs_basic() {
3358 let expr = parse("SELECT a, b + c FROM t");
3359 let scope = build_scope(&expr);
3360 let result = qualify_outputs(&scope);
3361 assert!(result.is_ok());
3362 }
3363
3364 #[test]
3365 fn test_qualify_columns_expands_star_with_schema() {
3366 let expr = parse("SELECT * FROM users");
3367
3368 let mut schema = MappingSchema::new();
3369 schema
3370 .add_table(
3371 "users",
3372 &[
3373 (
3374 "id".to_string(),
3375 DataType::Int {
3376 length: None,
3377 integer_spelling: false,
3378 },
3379 ),
3380 ("name".to_string(), DataType::Text),
3381 ("email".to_string(), DataType::Text),
3382 ],
3383 None,
3384 )
3385 .expect("schema setup");
3386
3387 let result =
3388 qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3389 let sql = gen(&result);
3390
3391 assert!(!sql.contains("SELECT *"));
3392 assert!(sql.contains("users.id"));
3393 assert!(sql.contains("users.name"));
3394 assert!(sql.contains("users.email"));
3395 }
3396
3397 #[test]
3398 fn test_qualify_columns_expands_group_by_positions() {
3399 let expr = parse("SELECT a, b FROM t GROUP BY 1, 2");
3400
3401 let mut schema = MappingSchema::new();
3402 schema
3403 .add_table(
3404 "t",
3405 &[
3406 (
3407 "a".to_string(),
3408 DataType::Int {
3409 length: None,
3410 integer_spelling: false,
3411 },
3412 ),
3413 (
3414 "b".to_string(),
3415 DataType::Int {
3416 length: None,
3417 integer_spelling: false,
3418 },
3419 ),
3420 ],
3421 None,
3422 )
3423 .expect("schema setup");
3424
3425 let result =
3426 qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3427 let sql = gen(&result);
3428
3429 assert!(!sql.contains("GROUP BY 1"));
3430 assert!(!sql.contains("GROUP BY 2"));
3431 assert!(sql.contains("GROUP BY"));
3432 assert!(sql.contains("t.a"));
3433 assert!(sql.contains("t.b"));
3434 }
3435
3436 #[test]
3441 fn test_expand_using_simple() {
3442 let expr = parse("SELECT x.b FROM x JOIN y USING (b)");
3444
3445 let mut schema = MappingSchema::new();
3446 schema
3447 .add_table(
3448 "x",
3449 &[
3450 ("a".to_string(), DataType::BigInt { length: None }),
3451 ("b".to_string(), DataType::BigInt { length: None }),
3452 ],
3453 None,
3454 )
3455 .expect("schema setup");
3456 schema
3457 .add_table(
3458 "y",
3459 &[
3460 ("b".to_string(), DataType::BigInt { length: None }),
3461 ("c".to_string(), DataType::BigInt { length: None }),
3462 ],
3463 None,
3464 )
3465 .expect("schema setup");
3466
3467 let result =
3468 qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3469 let sql = gen(&result);
3470
3471 assert!(
3473 !sql.contains("USING"),
3474 "USING should be replaced with ON: {sql}"
3475 );
3476 assert!(
3477 sql.contains("ON x.b = y.b"),
3478 "ON condition should be x.b = y.b: {sql}"
3479 );
3480 assert!(sql.contains("SELECT x.b"), "SELECT should keep x.b: {sql}");
3482 }
3483
3484 #[test]
3485 fn test_expand_using_unqualified_coalesce() {
3486 let expr = parse("SELECT b FROM x JOIN y USING(b)");
3488
3489 let mut schema = MappingSchema::new();
3490 schema
3491 .add_table(
3492 "x",
3493 &[
3494 ("a".to_string(), DataType::BigInt { length: None }),
3495 ("b".to_string(), DataType::BigInt { length: None }),
3496 ],
3497 None,
3498 )
3499 .expect("schema setup");
3500 schema
3501 .add_table(
3502 "y",
3503 &[
3504 ("b".to_string(), DataType::BigInt { length: None }),
3505 ("c".to_string(), DataType::BigInt { length: None }),
3506 ],
3507 None,
3508 )
3509 .expect("schema setup");
3510
3511 let result =
3512 qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3513 let sql = gen(&result);
3514
3515 assert!(
3516 sql.contains("COALESCE(x.b, y.b)"),
3517 "Unqualified USING column should become COALESCE: {sql}"
3518 );
3519 assert!(
3520 sql.contains("AS b"),
3521 "COALESCE should be aliased as 'b': {sql}"
3522 );
3523 assert!(
3524 sql.contains("ON x.b = y.b"),
3525 "ON condition should be generated: {sql}"
3526 );
3527 }
3528
3529 #[test]
3530 fn test_expand_using_with_where() {
3531 let expr = parse("SELECT b FROM x JOIN y USING(b) WHERE b = 1");
3533
3534 let mut schema = MappingSchema::new();
3535 schema
3536 .add_table(
3537 "x",
3538 &[("b".to_string(), DataType::BigInt { length: None })],
3539 None,
3540 )
3541 .expect("schema setup");
3542 schema
3543 .add_table(
3544 "y",
3545 &[("b".to_string(), DataType::BigInt { length: None })],
3546 None,
3547 )
3548 .expect("schema setup");
3549
3550 let result =
3551 qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3552 let sql = gen(&result);
3553
3554 assert!(
3555 sql.contains("WHERE COALESCE(x.b, y.b)"),
3556 "WHERE should use COALESCE for USING column: {sql}"
3557 );
3558 }
3559
3560 #[test]
3561 fn test_expand_using_multi_join() {
3562 let expr = parse("SELECT b FROM x JOIN y USING(b) JOIN z USING(b)");
3564
3565 let mut schema = MappingSchema::new();
3566 for table in &["x", "y", "z"] {
3567 schema
3568 .add_table(
3569 table,
3570 &[("b".to_string(), DataType::BigInt { length: None })],
3571 None,
3572 )
3573 .expect("schema setup");
3574 }
3575
3576 let result =
3577 qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3578 let sql = gen(&result);
3579
3580 assert!(
3582 sql.contains("COALESCE(x.b, y.b, z.b)"),
3583 "Should have 3-table COALESCE: {sql}"
3584 );
3585 assert!(
3587 sql.contains("ON x.b = y.b"),
3588 "First join ON condition: {sql}"
3589 );
3590 }
3591
3592 #[test]
3593 fn test_expand_using_multi_column() {
3594 let expr = parse("SELECT b, c FROM y JOIN z USING(b, c)");
3596
3597 let mut schema = MappingSchema::new();
3598 schema
3599 .add_table(
3600 "y",
3601 &[
3602 ("b".to_string(), DataType::BigInt { length: None }),
3603 ("c".to_string(), DataType::BigInt { length: None }),
3604 ],
3605 None,
3606 )
3607 .expect("schema setup");
3608 schema
3609 .add_table(
3610 "z",
3611 &[
3612 ("b".to_string(), DataType::BigInt { length: None }),
3613 ("c".to_string(), DataType::BigInt { length: None }),
3614 ],
3615 None,
3616 )
3617 .expect("schema setup");
3618
3619 let result =
3620 qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3621 let sql = gen(&result);
3622
3623 assert!(
3624 sql.contains("COALESCE(y.b, z.b)"),
3625 "column 'b' should get COALESCE: {sql}"
3626 );
3627 assert!(
3628 sql.contains("COALESCE(y.c, z.c)"),
3629 "column 'c' should get COALESCE: {sql}"
3630 );
3631 assert!(
3633 sql.contains("y.b = z.b") && sql.contains("y.c = z.c"),
3634 "ON should have both equality conditions: {sql}"
3635 );
3636 }
3637
3638 #[test]
3639 fn test_expand_using_star() {
3640 let expr = parse("SELECT * FROM x JOIN y USING(b)");
3642
3643 let mut schema = MappingSchema::new();
3644 schema
3645 .add_table(
3646 "x",
3647 &[
3648 ("a".to_string(), DataType::BigInt { length: None }),
3649 ("b".to_string(), DataType::BigInt { length: None }),
3650 ],
3651 None,
3652 )
3653 .expect("schema setup");
3654 schema
3655 .add_table(
3656 "y",
3657 &[
3658 ("b".to_string(), DataType::BigInt { length: None }),
3659 ("c".to_string(), DataType::BigInt { length: None }),
3660 ],
3661 None,
3662 )
3663 .expect("schema setup");
3664
3665 let result =
3666 qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3667 let sql = gen(&result);
3668
3669 assert!(
3671 sql.contains("COALESCE(x.b, y.b) AS b"),
3672 "USING column should be COALESCE in star expansion: {sql}"
3673 );
3674 assert!(sql.contains("x.a"), "non-USING column a from x: {sql}");
3676 assert!(sql.contains("y.c"), "non-USING column c from y: {sql}");
3677 let coalesce_count = sql.matches("COALESCE").count();
3679 assert_eq!(
3680 coalesce_count, 1,
3681 "b should appear only once as COALESCE: {sql}"
3682 );
3683 }
3684
3685 #[test]
3686 fn test_expand_using_table_star() {
3687 let expr = parse("SELECT x.* FROM x JOIN y USING(b)");
3689
3690 let mut schema = MappingSchema::new();
3691 schema
3692 .add_table(
3693 "x",
3694 &[
3695 ("a".to_string(), DataType::BigInt { length: None }),
3696 ("b".to_string(), DataType::BigInt { length: None }),
3697 ],
3698 None,
3699 )
3700 .expect("schema setup");
3701 schema
3702 .add_table(
3703 "y",
3704 &[
3705 ("b".to_string(), DataType::BigInt { length: None }),
3706 ("c".to_string(), DataType::BigInt { length: None }),
3707 ],
3708 None,
3709 )
3710 .expect("schema setup");
3711
3712 let result =
3713 qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3714 let sql = gen(&result);
3715
3716 assert!(
3718 sql.contains("COALESCE(x.b, y.b)"),
3719 "USING column from x.* should become COALESCE: {sql}"
3720 );
3721 assert!(sql.contains("x.a"), "non-USING column a: {sql}");
3722 }
3723
3724 #[test]
3725 fn test_expand_natural_join_with_derived_table() {
3726 let expr = parse(
3727 "SELECT shared_key AS output_key FROM source_table \
3728 NATURAL JOIN (SELECT shared_key FROM source_table) AS derived",
3729 );
3730
3731 let mut schema = MappingSchema::new();
3732 schema
3733 .add_table(
3734 "source_table",
3735 &[(
3736 "shared_key".to_string(),
3737 DataType::VarChar {
3738 length: None,
3739 parenthesized_length: false,
3740 },
3741 )],
3742 None,
3743 )
3744 .expect("schema setup");
3745
3746 let result = qualify_columns(
3747 expr,
3748 &schema,
3749 &QualifyColumnsOptions::new().with_dialect(DialectType::DuckDB),
3750 )
3751 .expect("qualify");
3752 let sql = gen(&result);
3753
3754 assert!(!sql.contains("NATURAL"), "NATURAL should expand: {sql}");
3755 assert!(
3756 sql.contains("ON source_table.shared_key = derived.shared_key"),
3757 "common column should become an equality condition: {sql}"
3758 );
3759 assert!(
3760 sql.contains("COALESCE(source_table.shared_key, derived.shared_key) AS output_key"),
3761 "merged projection should preserve both sources: {sql}"
3762 );
3763 }
3764
3765 #[test]
3766 fn test_expand_natural_join_star_and_chained_sources() {
3767 let expr = parse("SELECT * FROM x NATURAL JOIN y NATURAL JOIN z");
3768
3769 let mut schema = MappingSchema::new();
3770 schema
3771 .add_table(
3772 "x",
3773 &[
3774 ("a".to_string(), DataType::BigInt { length: None }),
3775 ("shared".to_string(), DataType::BigInt { length: None }),
3776 ],
3777 None,
3778 )
3779 .expect("schema setup");
3780 schema
3781 .add_table(
3782 "y",
3783 &[
3784 ("shared".to_string(), DataType::BigInt { length: None }),
3785 ("b".to_string(), DataType::BigInt { length: None }),
3786 ],
3787 None,
3788 )
3789 .expect("schema setup");
3790 schema
3791 .add_table(
3792 "z",
3793 &[
3794 ("shared".to_string(), DataType::BigInt { length: None }),
3795 ("c".to_string(), DataType::BigInt { length: None }),
3796 ],
3797 None,
3798 )
3799 .expect("schema setup");
3800
3801 let result =
3802 qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3803 let sql = gen(&result);
3804
3805 assert!(!sql.contains("NATURAL"), "both joins should expand: {sql}");
3806 assert!(sql.contains("ON x.shared = y.shared"), "first join: {sql}");
3807 assert!(sql.contains("ON x.shared = z.shared"), "second join: {sql}");
3808 assert_eq!(
3809 sql.matches("COALESCE(x.shared, y.shared, z.shared) AS shared")
3810 .count(),
3811 1,
3812 "merged star column should be emitted once: {sql}"
3813 );
3814 assert!(sql.contains("x.a"), "left-only column should remain: {sql}");
3815 assert!(
3816 sql.contains("y.b"),
3817 "middle-only column should remain: {sql}"
3818 );
3819 assert!(
3820 sql.contains("z.c"),
3821 "right-only column should remain: {sql}"
3822 );
3823 }
3824
3825 #[test]
3826 fn test_expand_natural_outer_join_kinds() {
3827 let mut schema = MappingSchema::new();
3828 for table in ["x", "y"] {
3829 schema
3830 .add_table(
3831 table,
3832 &[("shared".to_string(), DataType::BigInt { length: None })],
3833 None,
3834 )
3835 .expect("schema setup");
3836 }
3837
3838 for (input_kind, output_kind) in [
3839 ("NATURAL LEFT JOIN", "LEFT JOIN"),
3840 ("NATURAL RIGHT JOIN", "RIGHT JOIN"),
3841 ("NATURAL FULL JOIN", "FULL JOIN"),
3842 ] {
3843 let expr = parse(&format!("SELECT shared FROM x {input_kind} y"));
3844 let result =
3845 qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3846 let sql = gen(&result);
3847
3848 assert!(
3849 !sql.contains("NATURAL"),
3850 "{input_kind} should expand: {sql}"
3851 );
3852 assert!(sql.contains(output_kind), "join kind should remain: {sql}");
3853 assert!(
3854 sql.contains("ON x.shared = y.shared"),
3855 "join condition should be derived: {sql}"
3856 );
3857 }
3858 }
3859
3860 #[test]
3861 fn test_preserve_natural_join_without_known_common_columns() {
3862 let mut no_common_schema = MappingSchema::new();
3863 no_common_schema
3864 .add_table(
3865 "x",
3866 &[("a".to_string(), DataType::BigInt { length: None })],
3867 None,
3868 )
3869 .expect("schema setup");
3870 no_common_schema
3871 .add_table(
3872 "y",
3873 &[("b".to_string(), DataType::BigInt { length: None })],
3874 None,
3875 )
3876 .expect("schema setup");
3877
3878 let no_common = qualify_columns(
3879 parse("SELECT * FROM x NATURAL JOIN y"),
3880 &no_common_schema,
3881 &QualifyColumnsOptions::new(),
3882 )
3883 .expect("qualify");
3884 assert!(
3885 gen(&no_common).contains("NATURAL JOIN"),
3886 "join without common columns should remain NATURAL"
3887 );
3888
3889 let mut partial_schema = MappingSchema::new();
3890 partial_schema
3891 .add_table(
3892 "x",
3893 &[("a".to_string(), DataType::BigInt { length: None })],
3894 None,
3895 )
3896 .expect("schema setup");
3897 let unknown_right = qualify_columns(
3898 parse("SELECT * FROM x NATURAL JOIN unknown_table"),
3899 &partial_schema,
3900 &QualifyColumnsOptions::new().with_allow_partial(true),
3901 )
3902 .expect("partial qualification");
3903 assert!(
3904 gen(&unknown_right).contains("NATURAL JOIN"),
3905 "join with an unknown schema should remain NATURAL"
3906 );
3907
3908 let unknown_left = qualify_columns(
3909 parse("SELECT * FROM unknown_table NATURAL JOIN x"),
3910 &partial_schema,
3911 &QualifyColumnsOptions::new().with_allow_partial(true),
3912 )
3913 .expect("partial qualification");
3914 assert!(
3915 gen(&unknown_left).contains("NATURAL JOIN"),
3916 "join with an unknown left schema should remain NATURAL"
3917 );
3918
3919 let partially_known_chain = qualify_columns(
3920 parse("SELECT * FROM x NATURAL JOIN unknown_table NATURAL JOIN x AS x2"),
3921 &partial_schema,
3922 &QualifyColumnsOptions::new().with_allow_partial(true),
3923 )
3924 .expect("partial qualification");
3925 assert_eq!(
3926 gen(&partially_known_chain).matches("NATURAL JOIN").count(),
3927 2,
3928 "an unknown earlier source should prevent later inferred join keys"
3929 );
3930 }
3931
3932 #[test]
3933 fn test_qualify_columns_qualified_table_name() {
3934 let expr = parse("SELECT a FROM raw.t1");
3935
3936 let mut schema = MappingSchema::new();
3937 schema
3938 .add_table(
3939 "raw.t1",
3940 &[("a".to_string(), DataType::BigInt { length: None })],
3941 None,
3942 )
3943 .expect("schema setup");
3944
3945 let result =
3946 qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3947 let sql = gen(&result);
3948
3949 assert!(
3950 sql.contains("t1.a"),
3951 "column should be qualified with table name: {sql}"
3952 );
3953
3954 let expr = parse("SELECT MAX(a) FROM raw.t1");
3956 let result =
3957 qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3958 let sql = gen(&result);
3959 assert!(
3960 sql.contains("t1.a"),
3961 "column in function should be qualified with table name: {sql}"
3962 );
3963
3964 let expr = parse("SELECT ABS(a) FROM raw.t1");
3966 let result =
3967 qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3968 let sql = gen(&result);
3969 assert!(
3970 sql.contains("t1.a"),
3971 "column in function should be qualified with table name: {sql}"
3972 );
3973 }
3974
3975 #[test]
3976 fn test_qualify_columns_count_star() {
3977 let expr = parse("SELECT COUNT(*) FROM t1");
3979
3980 let mut schema = MappingSchema::new();
3981 schema
3982 .add_table(
3983 "t1",
3984 &[("id".to_string(), DataType::BigInt { length: None })],
3985 None,
3986 )
3987 .expect("schema setup");
3988
3989 let result =
3990 qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
3991 let sql = gen(&result);
3992
3993 assert!(
3994 sql.contains("COUNT(*)"),
3995 "COUNT(*) should be preserved: {sql}"
3996 );
3997 }
3998
3999 #[test]
4000 fn test_qualify_columns_correlated_scalar_subquery() {
4001 let expr =
4002 parse("SELECT id, (SELECT AVG(val) FROM t2 WHERE t2.id = t1.id) AS avg_val FROM t1");
4003
4004 let mut schema = MappingSchema::new();
4005 schema
4006 .add_table(
4007 "t1",
4008 &[("id".to_string(), DataType::BigInt { length: None })],
4009 None,
4010 )
4011 .expect("schema setup");
4012 schema
4013 .add_table(
4014 "t2",
4015 &[
4016 ("id".to_string(), DataType::BigInt { length: None }),
4017 ("val".to_string(), DataType::BigInt { length: None }),
4018 ],
4019 None,
4020 )
4021 .expect("schema setup");
4022
4023 let result =
4024 qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
4025 let sql = gen(&result);
4026
4027 assert!(
4028 sql.contains("t1.id"),
4029 "outer column should be qualified: {sql}"
4030 );
4031 assert!(
4032 sql.contains("t2.id"),
4033 "inner column should be qualified: {sql}"
4034 );
4035 }
4036
4037 #[test]
4038 fn test_qualify_columns_correlated_scalar_subquery_unqualified() {
4039 let expr =
4040 parse("SELECT t1_id, (SELECT AVG(val) FROM t2 WHERE t2_id = t1_id) AS avg_val FROM t1");
4041
4042 let mut schema = MappingSchema::new();
4043 schema
4044 .add_table(
4045 "t1",
4046 &[("t1_id".to_string(), DataType::BigInt { length: None })],
4047 None,
4048 )
4049 .expect("schema setup");
4050 schema
4051 .add_table(
4052 "t2",
4053 &[
4054 ("t2_id".to_string(), DataType::BigInt { length: None }),
4055 ("val".to_string(), DataType::BigInt { length: None }),
4056 ],
4057 None,
4058 )
4059 .expect("schema setup");
4060
4061 let result =
4062 qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
4063 let sql = gen(&result);
4064
4065 assert!(
4066 sql.contains("t1.t1_id"),
4067 "outer column should be qualified: {sql}"
4068 );
4069 assert!(
4070 sql.contains("t2.t2_id"),
4071 "inner column should be qualified: {sql}"
4072 );
4073 assert!(
4075 sql.contains("= t1.t1_id"),
4076 "correlated column should be qualified: {sql}"
4077 );
4078 }
4079
4080 #[test]
4081 fn test_qualify_columns_correlated_exists_subquery() {
4082 let expr = parse(
4083 "SELECT o_orderpriority FROM orders \
4084 WHERE EXISTS (SELECT * FROM lineitem WHERE l_orderkey = o_orderkey)",
4085 );
4086
4087 let mut schema = MappingSchema::new();
4088 schema
4089 .add_table(
4090 "orders",
4091 &[
4092 ("o_orderpriority".to_string(), DataType::Text),
4093 ("o_orderkey".to_string(), DataType::BigInt { length: None }),
4094 ],
4095 None,
4096 )
4097 .expect("schema setup");
4098 schema
4099 .add_table(
4100 "lineitem",
4101 &[("l_orderkey".to_string(), DataType::BigInt { length: None })],
4102 None,
4103 )
4104 .expect("schema setup");
4105
4106 let result =
4107 qualify_columns(expr, &schema, &QualifyColumnsOptions::new()).expect("qualify");
4108 let sql = gen(&result);
4109
4110 assert!(
4111 sql.contains("orders.o_orderpriority"),
4112 "outer column should be qualified: {sql}"
4113 );
4114 assert!(
4115 sql.contains("lineitem.l_orderkey"),
4116 "inner column should be qualified: {sql}"
4117 );
4118 assert!(
4119 sql.contains("orders.o_orderkey"),
4120 "correlated outer column should be qualified: {sql}"
4121 );
4122 }
4123
4124 #[test]
4125 fn test_qualify_columns_rejects_unknown_table() {
4126 let expr = parse("SELECT id FROM t1 WHERE nonexistent.col = 1");
4127
4128 let mut schema = MappingSchema::new();
4129 schema
4130 .add_table(
4131 "t1",
4132 &[("id".to_string(), DataType::BigInt { length: None })],
4133 None,
4134 )
4135 .expect("schema setup");
4136
4137 let result = qualify_columns(expr, &schema, &QualifyColumnsOptions::new());
4138 assert!(
4139 result.is_err(),
4140 "should reject reference to table not in scope or schema"
4141 );
4142 }
4143
4144 #[test]
4149 fn test_needs_quoting_reserved_word() {
4150 let reserved = get_reserved_words(None);
4151 assert!(needs_quoting("select", &reserved));
4152 assert!(needs_quoting("SELECT", &reserved));
4153 assert!(needs_quoting("from", &reserved));
4154 assert!(needs_quoting("WHERE", &reserved));
4155 assert!(needs_quoting("join", &reserved));
4156 assert!(needs_quoting("table", &reserved));
4157 }
4158
4159 #[test]
4160 fn test_needs_quoting_normal_identifiers() {
4161 let reserved = get_reserved_words(None);
4162 assert!(!needs_quoting("foo", &reserved));
4163 assert!(!needs_quoting("my_column", &reserved));
4164 assert!(!needs_quoting("col1", &reserved));
4165 assert!(!needs_quoting("A", &reserved));
4166 assert!(!needs_quoting("_hidden", &reserved));
4167 }
4168
4169 #[test]
4170 fn test_needs_quoting_special_characters() {
4171 let reserved = get_reserved_words(None);
4172 assert!(needs_quoting("my column", &reserved)); assert!(needs_quoting("my-column", &reserved)); assert!(needs_quoting("my.column", &reserved)); assert!(needs_quoting("col@name", &reserved)); assert!(needs_quoting("col#name", &reserved)); }
4178
4179 #[test]
4180 fn test_needs_quoting_starts_with_digit() {
4181 let reserved = get_reserved_words(None);
4182 assert!(needs_quoting("1col", &reserved));
4183 assert!(needs_quoting("123", &reserved));
4184 assert!(needs_quoting("0_start", &reserved));
4185 }
4186
4187 #[test]
4188 fn test_needs_quoting_empty() {
4189 let reserved = get_reserved_words(None);
4190 assert!(!needs_quoting("", &reserved));
4191 }
4192
4193 #[test]
4194 fn test_maybe_quote_sets_quoted_flag() {
4195 let reserved = get_reserved_words(None);
4196 let mut id = Identifier::new("select");
4197 assert!(!id.quoted);
4198 maybe_quote(&mut id, &reserved);
4199 assert!(id.quoted);
4200 }
4201
4202 #[test]
4203 fn test_maybe_quote_skips_already_quoted() {
4204 let reserved = get_reserved_words(None);
4205 let mut id = Identifier::quoted("myname");
4206 assert!(id.quoted);
4207 maybe_quote(&mut id, &reserved);
4208 assert!(id.quoted); assert_eq!(id.name, "myname"); }
4211
4212 #[test]
4213 fn test_maybe_quote_skips_star() {
4214 let reserved = get_reserved_words(None);
4215 let mut id = Identifier::new("*");
4216 maybe_quote(&mut id, &reserved);
4217 assert!(!id.quoted); }
4219
4220 #[test]
4221 fn test_maybe_quote_skips_normal() {
4222 let reserved = get_reserved_words(None);
4223 let mut id = Identifier::new("normal_col");
4224 maybe_quote(&mut id, &reserved);
4225 assert!(!id.quoted);
4226 }
4227
4228 #[test]
4229 fn test_quote_identifiers_column_with_reserved_name() {
4230 let expr = Expression::boxed_column(Column {
4232 name: Identifier::new("select"),
4233 table: None,
4234 join_mark: false,
4235 trailing_comments: vec![],
4236 span: None,
4237 inferred_type: None,
4238 });
4239 let result = quote_identifiers(expr, None);
4240 if let Expression::Column(col) = &result {
4241 assert!(col.name.quoted, "Column named 'select' should be quoted");
4242 } else {
4243 panic!("Expected Column expression");
4244 }
4245 }
4246
4247 #[test]
4248 fn test_quote_identifiers_column_with_special_chars() {
4249 let expr = Expression::boxed_column(Column {
4250 name: Identifier::new("my column"),
4251 table: None,
4252 join_mark: false,
4253 trailing_comments: vec![],
4254 span: None,
4255 inferred_type: None,
4256 });
4257 let result = quote_identifiers(expr, None);
4258 if let Expression::Column(col) = &result {
4259 assert!(col.name.quoted, "Column with space should be quoted");
4260 } else {
4261 panic!("Expected Column expression");
4262 }
4263 }
4264
4265 #[test]
4266 fn test_quote_identifiers_preserves_normal_column() {
4267 let expr = Expression::boxed_column(Column {
4268 name: Identifier::new("normal_col"),
4269 table: Some(Identifier::new("my_table")),
4270 join_mark: false,
4271 trailing_comments: vec![],
4272 span: None,
4273 inferred_type: None,
4274 });
4275 let result = quote_identifiers(expr, None);
4276 if let Expression::Column(col) = &result {
4277 assert!(!col.name.quoted, "Normal column should not be quoted");
4278 assert!(
4279 !col.table.as_ref().unwrap().quoted,
4280 "Normal table should not be quoted"
4281 );
4282 } else {
4283 panic!("Expected Column expression");
4284 }
4285 }
4286
4287 #[test]
4288 fn test_quote_identifiers_table_ref_reserved() {
4289 let expr = Expression::Table(Box::new(TableRef::new("select")));
4290 let result = quote_identifiers(expr, None);
4291 if let Expression::Table(tr) = &result {
4292 assert!(tr.name.quoted, "Table named 'select' should be quoted");
4293 } else {
4294 panic!("Expected Table expression");
4295 }
4296 }
4297
4298 #[test]
4299 fn test_quote_identifiers_table_ref_schema_and_alias() {
4300 let mut tr = TableRef::new("my_table");
4301 tr.schema = Some(Identifier::new("from"));
4302 tr.alias = Some(Identifier::new("t"));
4303 let expr = Expression::Table(Box::new(tr));
4304 let result = quote_identifiers(expr, None);
4305 if let Expression::Table(tr) = &result {
4306 assert!(!tr.name.quoted, "Normal table name should not be quoted");
4307 assert!(
4308 tr.schema.as_ref().unwrap().quoted,
4309 "Schema named 'from' should be quoted"
4310 );
4311 assert!(
4312 !tr.alias.as_ref().unwrap().quoted,
4313 "Normal alias should not be quoted"
4314 );
4315 } else {
4316 panic!("Expected Table expression");
4317 }
4318 }
4319
4320 #[test]
4321 fn test_quote_identifiers_identifier_node() {
4322 let expr = Expression::Identifier(Identifier::new("order"));
4323 let result = quote_identifiers(expr, None);
4324 if let Expression::Identifier(id) = &result {
4325 assert!(id.quoted, "Identifier named 'order' should be quoted");
4326 } else {
4327 panic!("Expected Identifier expression");
4328 }
4329 }
4330
4331 #[test]
4332 fn test_quote_identifiers_alias() {
4333 let inner = Expression::boxed_column(Column {
4334 name: Identifier::new("val"),
4335 table: None,
4336 join_mark: false,
4337 trailing_comments: vec![],
4338 span: None,
4339 inferred_type: None,
4340 });
4341 let expr = Expression::Alias(Box::new(Alias {
4342 this: inner,
4343 alias: Identifier::new("select"),
4344 column_aliases: vec![Identifier::new("from")],
4345 alias_explicit_as: false,
4346 alias_keyword: None,
4347 pre_alias_comments: vec![],
4348 trailing_comments: vec![],
4349 inferred_type: None,
4350 }));
4351 let result = quote_identifiers(expr, None);
4352 if let Expression::Alias(alias) = &result {
4353 assert!(alias.alias.quoted, "Alias named 'select' should be quoted");
4354 assert!(
4355 alias.column_aliases[0].quoted,
4356 "Column alias named 'from' should be quoted"
4357 );
4358 if let Expression::Column(col) = &alias.this {
4360 assert!(!col.name.quoted);
4361 }
4362 } else {
4363 panic!("Expected Alias expression");
4364 }
4365 }
4366
4367 #[test]
4368 fn test_quote_identifiers_select_recursive() {
4369 let expr = parse("SELECT a, b FROM t WHERE c = 1");
4371 let result = quote_identifiers(expr, None);
4372 let sql = gen(&result);
4374 assert!(sql.contains("a"));
4376 assert!(sql.contains("b"));
4377 assert!(sql.contains("t"));
4378 }
4379
4380 #[test]
4381 fn test_quote_identifiers_digit_start() {
4382 let expr = Expression::boxed_column(Column {
4383 name: Identifier::new("1col"),
4384 table: None,
4385 join_mark: false,
4386 trailing_comments: vec![],
4387 span: None,
4388 inferred_type: None,
4389 });
4390 let result = quote_identifiers(expr, None);
4391 if let Expression::Column(col) = &result {
4392 assert!(
4393 col.name.quoted,
4394 "Column starting with digit should be quoted"
4395 );
4396 } else {
4397 panic!("Expected Column expression");
4398 }
4399 }
4400
4401 #[test]
4402 fn test_quote_identifiers_with_mysql_dialect() {
4403 let reserved = get_reserved_words(Some(DialectType::MySQL));
4404 assert!(needs_quoting("KILL", &reserved));
4406 assert!(needs_quoting("FORCE", &reserved));
4408 }
4409
4410 #[test]
4411 fn test_quote_identifiers_with_postgresql_dialect() {
4412 let reserved = get_reserved_words(Some(DialectType::PostgreSQL));
4413 assert!(needs_quoting("ILIKE", &reserved));
4415 assert!(needs_quoting("VERBOSE", &reserved));
4417 }
4418
4419 #[test]
4420 fn test_quote_identifiers_with_bigquery_dialect() {
4421 let reserved = get_reserved_words(Some(DialectType::BigQuery));
4422 assert!(needs_quoting("STRUCT", &reserved));
4424 assert!(needs_quoting("PROTO", &reserved));
4426 }
4427
4428 #[test]
4429 fn test_quote_identifiers_case_insensitive_reserved() {
4430 let reserved = get_reserved_words(None);
4431 assert!(needs_quoting("Select", &reserved));
4432 assert!(needs_quoting("sElEcT", &reserved));
4433 assert!(needs_quoting("FROM", &reserved));
4434 assert!(needs_quoting("from", &reserved));
4435 }
4436
4437 #[test]
4438 fn test_quote_identifiers_join_using() {
4439 let mut join = crate::expressions::Join {
4441 this: Expression::Table(Box::new(TableRef::new("other"))),
4442 on: None,
4443 using: vec![Identifier::new("key"), Identifier::new("value")],
4444 kind: crate::expressions::JoinKind::Inner,
4445 use_inner_keyword: false,
4446 use_outer_keyword: false,
4447 deferred_condition: false,
4448 join_hint: None,
4449 match_condition: None,
4450 pivots: vec![],
4451 comments: vec![],
4452 nesting_group: 0,
4453 directed: false,
4454 };
4455 let reserved = get_reserved_words(None);
4456 quote_join(&mut join, &reserved);
4457 assert!(
4459 join.using[0].quoted,
4460 "USING identifier 'key' should be quoted"
4461 );
4462 assert!(
4463 !join.using[1].quoted,
4464 "USING identifier 'value' should not be quoted"
4465 );
4466 }
4467
4468 #[test]
4469 fn test_quote_identifiers_cte() {
4470 let mut cte = crate::expressions::Cte {
4472 alias: Identifier::new("select"),
4473 this: Expression::boxed_column(Column {
4474 name: Identifier::new("x"),
4475 table: None,
4476 join_mark: false,
4477 trailing_comments: vec![],
4478 span: None,
4479 inferred_type: None,
4480 }),
4481 columns: vec![Identifier::new("from"), Identifier::new("normal")],
4482 materialized: None,
4483 key_expressions: vec![],
4484 alias_first: false,
4485 comments: Vec::new(),
4486 };
4487 let reserved = get_reserved_words(None);
4488 maybe_quote(&mut cte.alias, &reserved);
4489 for c in &mut cte.columns {
4490 maybe_quote(c, &reserved);
4491 }
4492 assert!(cte.alias.quoted, "CTE alias 'select' should be quoted");
4493 assert!(cte.columns[0].quoted, "CTE column 'from' should be quoted");
4494 assert!(
4495 !cte.columns[1].quoted,
4496 "CTE column 'normal' should not be quoted"
4497 );
4498 }
4499
4500 #[test]
4501 fn test_quote_identifiers_binary_ops_recurse() {
4502 let expr = Expression::Add(Box::new(crate::expressions::BinaryOp::new(
4505 Expression::boxed_column(Column {
4506 name: Identifier::new("select"),
4507 table: None,
4508 join_mark: false,
4509 trailing_comments: vec![],
4510 span: None,
4511 inferred_type: None,
4512 }),
4513 Expression::boxed_column(Column {
4514 name: Identifier::new("normal"),
4515 table: None,
4516 join_mark: false,
4517 trailing_comments: vec![],
4518 span: None,
4519 inferred_type: None,
4520 }),
4521 )));
4522 let result = quote_identifiers(expr, None);
4523 if let Expression::Add(bin) = &result {
4524 if let Expression::Column(left) = &bin.left {
4525 assert!(
4526 left.name.quoted,
4527 "'select' column should be quoted in binary op"
4528 );
4529 }
4530 if let Expression::Column(right) = &bin.right {
4531 assert!(!right.name.quoted, "'normal' column should not be quoted");
4532 }
4533 } else {
4534 panic!("Expected Add expression");
4535 }
4536 }
4537
4538 #[test]
4539 fn test_quote_identifiers_already_quoted_preserved() {
4540 let expr = Expression::boxed_column(Column {
4542 name: Identifier::quoted("normal_name"),
4543 table: None,
4544 join_mark: false,
4545 trailing_comments: vec![],
4546 span: None,
4547 inferred_type: None,
4548 });
4549 let result = quote_identifiers(expr, None);
4550 if let Expression::Column(col) = &result {
4551 assert!(
4552 col.name.quoted,
4553 "Already-quoted identifier should remain quoted"
4554 );
4555 } else {
4556 panic!("Expected Column expression");
4557 }
4558 }
4559
4560 #[test]
4561 fn test_quote_identifiers_full_parsed_query() {
4562 let mut select = crate::expressions::Select::new();
4565 select.expressions.push(Expression::boxed_column(Column {
4566 name: Identifier::new("order"),
4567 table: Some(Identifier::new("t")),
4568 join_mark: false,
4569 trailing_comments: vec![],
4570 span: None,
4571 inferred_type: None,
4572 }));
4573 select.from = Some(crate::expressions::From {
4574 expressions: vec![Expression::Table(Box::new(TableRef::new("t")))],
4575 });
4576 let expr = Expression::Select(Box::new(select));
4577
4578 let result = quote_identifiers(expr, None);
4579 if let Expression::Select(sel) = &result {
4580 if let Expression::Column(col) = &sel.expressions[0] {
4581 assert!(col.name.quoted, "Column named 'order' should be quoted");
4582 assert!(
4583 !col.table.as_ref().unwrap().quoted,
4584 "Table 't' should not be quoted"
4585 );
4586 } else {
4587 panic!("Expected Column in SELECT list");
4588 }
4589 } else {
4590 panic!("Expected Select expression");
4591 }
4592 }
4593
4594 #[test]
4595 fn test_get_reserved_words_all_dialects() {
4596 let dialects = [
4598 None,
4599 Some(DialectType::Generic),
4600 Some(DialectType::MySQL),
4601 Some(DialectType::PostgreSQL),
4602 Some(DialectType::BigQuery),
4603 Some(DialectType::Snowflake),
4604 Some(DialectType::TSQL),
4605 Some(DialectType::ClickHouse),
4606 Some(DialectType::DuckDB),
4607 Some(DialectType::Hive),
4608 Some(DialectType::Spark),
4609 Some(DialectType::Trino),
4610 Some(DialectType::Oracle),
4611 Some(DialectType::Redshift),
4612 ];
4613 for dialect in &dialects {
4614 let words = get_reserved_words(*dialect);
4615 assert!(
4617 words.contains("SELECT"),
4618 "All dialects should have SELECT as reserved"
4619 );
4620 assert!(
4621 words.contains("FROM"),
4622 "All dialects should have FROM as reserved"
4623 );
4624 }
4625 }
4626}