Skip to content

Commit

Permalink
Feat: support non-strict qualify_columns (#4243)
Browse files Browse the repository at this point in the history
* Feat: support qualify_columns non-strictly

* add new flag
  • Loading branch information
hsheth2 authored Oct 15, 2024
1 parent fcc05c9 commit 6f32e53
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 3 deletions.
3 changes: 3 additions & 0 deletions sqlglot/optimizer/qualify.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ def qualify(
infer_schema: t.Optional[bool] = None,
isolate_tables: bool = False,
qualify_columns: bool = True,
allow_partial_qualification: bool = False,
validate_qualify_columns: bool = True,
quote_identifiers: bool = True,
identify: bool = True,
Expand Down Expand Up @@ -56,6 +57,7 @@ def qualify(
infer_schema: Whether to infer the schema if missing.
isolate_tables: Whether to isolate table selects.
qualify_columns: Whether to qualify columns.
allow_partial_qualification: Whether to allow partial qualification.
validate_qualify_columns: Whether to validate columns.
quote_identifiers: Whether to run the quote_identifiers step.
This step is necessary to ensure correctness for case sensitive queries.
Expand Down Expand Up @@ -90,6 +92,7 @@ def qualify(
expand_alias_refs=expand_alias_refs,
expand_stars=expand_stars,
infer_schema=infer_schema,
allow_partial_qualification=allow_partial_qualification,
)

if quote_identifiers:
Expand Down
13 changes: 10 additions & 3 deletions sqlglot/optimizer/qualify_columns.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ def qualify_columns(
expand_alias_refs: bool = True,
expand_stars: bool = True,
infer_schema: t.Optional[bool] = None,
allow_partial_qualification: bool = False,
) -> exp.Expression:
"""
Rewrite sqlglot AST to have fully qualified columns.
Expand All @@ -41,6 +42,7 @@ def qualify_columns(
for most of the optimizer's rules to work; do not set to False unless you
know what you're doing!
infer_schema: Whether to infer the schema if missing.
allow_partial_qualification: Whether to allow partial qualification.
Returns:
The qualified expression.
Expand Down Expand Up @@ -68,7 +70,7 @@ def qualify_columns(
)

_convert_columns_to_dots(scope, resolver)
_qualify_columns(scope, resolver)
_qualify_columns(scope, resolver, allow_partial_qualification=allow_partial_qualification)

if not schema.empty and expand_alias_refs:
_expand_alias_refs(scope, resolver)
Expand Down Expand Up @@ -441,15 +443,20 @@ def _convert_columns_to_dots(scope: Scope, resolver: Resolver) -> None:
scope.clear_cache()


def _qualify_columns(scope: Scope, resolver: Resolver) -> None:
def _qualify_columns(scope: Scope, resolver: Resolver, allow_partial_qualification: bool) -> None:
"""Disambiguate columns, ensuring each column specifies a source"""
for column in scope.columns:
column_table = column.table
column_name = column.name

if column_table and column_table in scope.sources:
source_columns = resolver.get_source_columns(column_table)
if source_columns and column_name not in source_columns and "*" not in source_columns:
if (
not allow_partial_qualification
and source_columns
and column_name not in source_columns
and "*" not in source_columns
):
raise OptimizeError(f"Unknown column: {column_name}")

if not column_table:
Expand Down
6 changes: 6 additions & 0 deletions tests/fixtures/optimizer/qualify_columns.sql
Original file line number Diff line number Diff line change
Expand Up @@ -190,6 +190,12 @@ SELECT x._col_0 AS _col_0, x._col_1 AS _col_1 FROM (VALUES (1, 2)) AS x(_col_0,
SELECT SOME_UDF(data).* FROM t;
SELECT SOME_UDF(t.data).* FROM t AS t;

# execute: false
# allow_partial_qualification: true
# validate_qualify_columns: false
SELECT a + 1 AS i, missing_column FROM x;
SELECT x.a + 1 AS i, missing_column AS missing_column FROM x AS x;

--------------------------------------
-- Derived tables
--------------------------------------
Expand Down

0 comments on commit 6f32e53

Please sign in to comment.