diff --git a/sqlglot/dialects/dialect.py b/sqlglot/dialects/dialect.py index 65582c4cbe..3f01b2be65 100644 --- a/sqlglot/dialects/dialect.py +++ b/sqlglot/dialects/dialect.py @@ -671,6 +671,14 @@ class Dialect(metaclass=_Dialect): Whether JSON_EXTRACT_SCALAR returns null if a non-scalar value is selected. """ + DEFAULT_FUNCTIONS_COLUMN_NAMES: t.Dict[t.Type[exp.Func], t.Union[str, t.Tuple[str, ...]]] = {} + """ + Maps function expressions to their default output column name(s). + + For example, in Postgres, generate_series function outputs a column named "generate_series" by default, + so we map the ExplodingGenerateSeries expression to "generate_series" string. + """ + # --- Autofilled --- tokenizer_class = Tokenizer diff --git a/sqlglot/dialects/postgres.py b/sqlglot/dialects/postgres.py index 7436659515..86c19ccc4f 100644 --- a/sqlglot/dialects/postgres.py +++ b/sqlglot/dialects/postgres.py @@ -297,6 +297,11 @@ class Postgres(Dialect): NULL_ORDERING = "nulls_are_large" TIME_FORMAT = "'YYYY-MM-DD HH24:MI:SS'" TABLESAMPLE_SIZE_IS_PERCENT = True + TABLES_REFERENCEABLE_AS_COLUMNS = True + + DEFAULT_FUNCTIONS_COLUMN_NAMES = { + exp.ExplodingGenerateSeries: "generate_series", + } TIME_MAPPING = { "d": "%u", # 1-based day of week diff --git a/sqlglot/optimizer/qualify_columns.py b/sqlglot/optimizer/qualify_columns.py index d0037b92a9..a31daf8867 100644 --- a/sqlglot/optimizer/qualify_columns.py +++ b/sqlglot/optimizer/qualify_columns.py @@ -598,7 +598,7 @@ def _qualify_columns( and len(column.parts) == 1 and column_name in scope.selected_sources ): - # BigQuery allows tables to be referenced as columns, treating them as structs + # BigQuery and Postgres allow tables to be referenced as columns, treating them as structs/records scope.replace(column, exp.TableColumn(this=column.this)) for pivot in scope.pivots: diff --git a/sqlglot/optimizer/qualify_tables.py b/sqlglot/optimizer/qualify_tables.py index 47269f8539..ebe0fe16bc 100644 --- a/sqlglot/optimizer/qualify_tables.py +++ b/sqlglot/optimizer/qualify_tables.py @@ -4,7 +4,7 @@ from sqlglot import exp from sqlglot.dialects.dialect import Dialect, DialectType -from sqlglot.helper import name_sequence, seq_get +from sqlglot.helper import name_sequence, seq_get, ensure_list from sqlglot.optimizer.normalize_identifiers import normalize_identifiers from sqlglot.optimizer.scope import Scope, traverse_scope @@ -79,6 +79,7 @@ def _set_alias( target_alias: t.Optional[str] = None, scope: t.Optional[Scope] = None, normalize: bool = False, + columns: t.Optional[t.List[t.Union[str, exp.Identifier]]] = None, ) -> None: alias = expression.args.get("alias") or exp.TableAlias() @@ -96,6 +97,10 @@ def _set_alias( quoted = True if canonicalize_table_aliases or not target_alias else None alias.set("this", exp.to_identifier(new_alias_name, quoted=quoted)) + + if columns: + alias.set("columns", [exp.to_identifier(c) for c in columns]) + expression.set("alias", alias) if scope: @@ -132,11 +137,27 @@ def _set_alias( if pivot := seq_get(source.args.get("pivots") or [], 0): name = source.name + table_this = source.this + table_alias = source.args.get("alias") + function_columns: t.List[t.Union[str, exp.Identifier]] = [] + if isinstance(table_this, exp.Func): + if not table_alias: + function_columns = ensure_list( + dialect.DEFAULT_FUNCTIONS_COLUMN_NAMES.get(type(table_this)) + ) + elif columns := table_alias.columns: + function_columns = columns + elif type(table_this) in dialect.DEFAULT_FUNCTIONS_COLUMN_NAMES: + function_columns = ensure_list(source.alias_or_name) + source.set("alias", None) + name = None + _set_alias( source, canonical_aliases, target_alias=name or source.name or None, normalize=True, + columns=function_columns, ) source_fqn = ".".join(p.name for p in source.parts) diff --git a/tests/fixtures/optimizer/qualify_columns.sql b/tests/fixtures/optimizer/qualify_columns.sql index 9e1a8fe922..61c6522828 100644 --- a/tests/fixtures/optimizer/qualify_columns.sql +++ b/tests/fixtures/optimizer/qualify_columns.sql @@ -287,6 +287,11 @@ SELECT DATE_TRUNC(t.col1, WEEK(MONDAY)) AS _col_0, t.col2 AS col2 FROM t AS t; SELECT first, second FROM (SELECT 'val' AS col, STACK(2, 1, 2, 3) AS (first, second)) AS tbl; SELECT tbl.first AS first, tbl.second AS second FROM (SELECT 'val' AS col, STACK(2, 1, 2, 3) AS (first, second)) AS tbl; +# execute: false +# dialect: postgres +WITH t AS (SELECT 1 AS c) SELECT t FROM t; +WITH t AS (SELECT 1 AS c) SELECT t AS _col_0 FROM t AS t; + -------------------------------------- -- Derived tables -------------------------------------- diff --git a/tests/fixtures/optimizer/qualify_tables.sql b/tests/fixtures/optimizer/qualify_tables.sql index 66653574a6..82ac03e2dd 100644 --- a/tests/fixtures/optimizer/qualify_tables.sql +++ b/tests/fixtures/optimizer/qualify_tables.sql @@ -260,3 +260,24 @@ SELECT * FROM c.db.x AS "_1" WHERE "_1".a = (SELECT SUM("_0".c) AS c FROM c.db.y # canonicalize_table_aliases: true SELECT t.foo FROM t AS t, (SELECT t.bar FROM t AS t); SELECT "_2".foo FROM c.db.t AS "_2", (SELECT "_0".bar FROM c.db.t AS "_0") AS "_1"; + +# title: Qualify GENERATE_SERIES with its default column generate_series +# dialect: postgres +SELECT generate_series FROM GENERATE_SERIES(1,2); +SELECT generate_series FROM GENERATE_SERIES(1, 2) AS "_0"(generate_series); + +# title: Qualify GENERATE_SERIES with alias by wrapping it +# dialect: postgres +SELECT g FROM GENERATE_SERIES(1,2) AS g; +SELECT g FROM GENERATE_SERIES(1, 2) AS "_0"(g); + +# title: Qualify GENERATE_SERIES with alias on table and columns +# dialect: postgres +SELECT g FROM GENERATE_SERIES(1,2) AS t(g); +SELECT g FROM GENERATE_SERIES(1, 2) AS t(g); + +# title: Qualify GENERATE_SERIES with explicit column and canonicalize_table_aliases +# dialect: postgres +# canonicalize_table_aliases: true +SELECT g FROM GENERATE_SERIES(1,2) AS t(g); +SELECT g FROM GENERATE_SERIES(1, 2) AS "_0"(g);