Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Appearance settings

Commit 4b05d21

Browse filesBrowse files
authored
fix: avoid aliasing known tables used in CTEs (googleapis#369)
Toward googleapis#368.
1 parent f6ad9c9 commit 4b05d21
Copy full SHA for 4b05d21

File tree

2 files changed

+52
-6
lines changed
Filter options

2 files changed

+52
-6
lines changed

‎sqlalchemy_bigquery/base.py

Copy file name to clipboardExpand all lines: sqlalchemy_bigquery/base.py
+16-6Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@
5151
from sqlalchemy.engine.base import Engine
5252
from sqlalchemy.sql.schema import Column
5353
from sqlalchemy.sql.schema import Table
54+
from sqlalchemy.sql.selectable import CTE
5455
from sqlalchemy.sql import elements, selectable
5556
import re
5657

@@ -254,6 +255,20 @@ def visit_table_valued_alias(self, element, **kw):
254255
ret = f"{aliases}, {ret}"
255256
return ret
256257

258+
def _known_tables(self):
259+
known_tables = set()
260+
261+
for from_ in self.compile_state.froms:
262+
if isinstance(from_, Table):
263+
known_tables.add(from_.name)
264+
elif isinstance(from_, CTE):
265+
for column in from_.original.selected_columns:
266+
table = getattr(column, "table", None)
267+
if table is not None:
268+
known_tables.add(table.name)
269+
270+
return known_tables
271+
257272
def visit_column(
258273
self,
259274
column,
@@ -290,12 +305,7 @@ def visit_column(
290305
if isinstance(tablename, elements._truncated_label):
291306
tablename = self._truncated_identifier("alias", tablename)
292307
elif TABLE_VALUED_ALIAS_ALIASES in kwargs:
293-
known_tables = set(
294-
from_.name
295-
for from_ in self.compile_state.froms
296-
if isinstance(from_, Table)
297-
)
298-
if tablename not in known_tables:
308+
if tablename not in self._known_tables():
299309
aliases = kwargs[TABLE_VALUED_ALIAS_ALIASES]
300310
if tablename not in aliases:
301311
aliases[tablename] = self.anon_map[

‎tests/unit/test_compiler.py

Copy file name to clipboardExpand all lines: tests/unit/test_compiler.py
+36Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,3 +76,39 @@ def test_no_alias_for_known_tables(faux_conn, metadata):
7676
)
7777
found_sql = q.compile(faux_conn).string
7878
assert found_sql == expected_sql
79+
80+
81+
@sqlalchemy_1_4_or_higher
82+
def test_no_alias_for_known_tables_cte(faux_conn, metadata):
83+
# See: https://github.com/googleapis/python-bigquery-sqlalchemy/issues/368
84+
table = setup_table(
85+
faux_conn,
86+
"table1",
87+
metadata,
88+
sqlalchemy.Column("foo", sqlalchemy.Integer),
89+
sqlalchemy.Column("bars", sqlalchemy.ARRAY(sqlalchemy.Integer)),
90+
)
91+
F = sqlalchemy.func
92+
93+
# Set up initiali query
94+
q = sqlalchemy.select(table.c.foo, F.unnest(table.c.bars).column_valued("bar"))
95+
96+
expected_initial_sql = (
97+
"SELECT `table1`.`foo`, `bar` \n"
98+
"FROM `table1`, unnest(`table1`.`bars`) AS `bar`"
99+
)
100+
found_initial_sql = q.compile(faux_conn).string
101+
assert found_initial_sql == expected_initial_sql
102+
103+
q = q.cte("cte")
104+
q = sqlalchemy.select(*q.columns)
105+
106+
expected_cte_sql = (
107+
"WITH `cte` AS \n"
108+
"(SELECT `table1`.`foo` AS `foo`, `bar` \n"
109+
"FROM `table1`, unnest(`table1`.`bars`) AS `bar`)\n"
110+
" SELECT `cte`.`foo`, `cte`.`bar` \n"
111+
"FROM `cte`"
112+
)
113+
found_cte_sql = q.compile(faux_conn).string
114+
assert found_cte_sql == expected_cte_sql

0 commit comments

Comments
0 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.