Skip to content

Commit

Permalink
fix: Avoid aliasing known CTEs (#839)
Browse files Browse the repository at this point in the history
* fix: Avoid aliasing known CTEs

* Add unit-test

---------

Co-authored-by: Tim Swast <swast@google.com>
Co-authored-by: Anthonios Partheniou <partheniou@google.com>
  • Loading branch information
3 people committed Jul 11, 2023
1 parent 81212cb commit 8a1f694
Show file tree
Hide file tree
Showing 2 changed files with 49 additions and 0 deletions.
1 change: 1 addition & 0 deletions sqlalchemy_bigquery/base.py
Expand Up @@ -261,6 +261,7 @@ def _known_tables(self):
if isinstance(from_, Table):
known_tables.add(from_.name)
elif isinstance(from_, CTE):
known_tables.add(from_.name)
for column in from_.original.selected_columns:
table = getattr(column, "table", None)
if table is not None:
Expand Down
48 changes: 48 additions & 0 deletions tests/system/test_sqlalchemy_bigquery.py
Expand Up @@ -776,6 +776,54 @@ def test_unnest(engine, bigquery_dataset):
assert sorted(r[0] for r in conn.execute(query)) == ["a", "b", "c", "x", "y"]


@pytest.mark.skipif(
packaging.version.parse(sqlalchemy.__version__) < packaging.version.parse("1.4"),
reason="unnest (and other table-valued-function) support required version 1.4",
)
def test_unnest_with_cte(engine, bigquery_dataset):
from sqlalchemy import select, func, String
from sqlalchemy_bigquery import ARRAY

conn = engine.connect()
metadata = MetaData()
table_name = "test_unnest_with_cte"
table = Table(
f"{bigquery_dataset}.{table_name}",
metadata,
Column("foo", String),
Column("bars", ARRAY(String)),
)
metadata.create_all(engine)
conn.execute(
table.insert(),
[dict(foo="first", bars=["a", "b", "c"]), dict(foo="second", bars=["x", "y"])],
)
selectable = select(table.c).select_from(table).cte("cte")
query = select(
[
selectable.c.foo,
func.unnest(selectable.c.bars).column_valued("unnest_bars"),
]
).select_from(selectable)
compiled = str(query.compile(engine))
assert " ".join(compiled.strip().split()) == (
f"WITH `cte` "
f"AS (SELECT `{bigquery_dataset}.{table_name}`.`foo` AS `foo`,"
f" `{bigquery_dataset}.{table_name}`.`bars` AS `bars`"
f" FROM `{bigquery_dataset}.{table_name}`) "
f"SELECT `cte`.`foo`, `unnest_bars` "
f"FROM `cte`, unnest(`cte`.`bars`) AS `unnest_bars`"
)

assert sorted(r for r in conn.execute(query)) == [
("first", "a"),
("first", "b"),
("first", "c"),
("second", "x"),
("second", "y"),
]


@pytest.mark.skipif(
packaging.version.parse(sqlalchemy.__version__) < packaging.version.parse("1.4"),
reason="regexp_match support requires version 1.4 or higher",
Expand Down

0 comments on commit 8a1f694

Please sign in to comment.