Skip to content
16 changes: 15 additions & 1 deletion sqlalchemy_bigquery/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -1377,7 +1377,12 @@ def get_view_definition(self, connection, view_name, schema=None, **kw):
return view.view_query


class unnest(sqlalchemy.sql.functions.GenericFunction):
# unnest is a reserved keyword in some dialects.
# It is defined here to avoid conflicts.
# https://github.com/googleapis/python-bigquery-sqlalchemy/issues/882
class _unnest(sqlalchemy.sql.expression.FunctionElement):
inherit_cache = True

def __init__(self, *args, **kwargs):
expr = kwargs.pop("expr", None)
if expr is not None:
Expand All @@ -1395,9 +1400,18 @@ def __init__(self, *args, **kwargs):
):
raise TypeError("The argument to unnest must have an ARRAY type.")
self.type = arg.type.item_type

super().__init__(*args, **kwargs)


@compiles(_unnest, "bigquery")
def bigquery_unnest(element, compiler, **kw):
return "UNNEST({})".format(compiler.process(element.clauses, **kw))


sqlalchemy.sql.functions._FunctionGenerator.unnest = _unnest


dialect = BigQueryDialect

try:
Expand Down
30 changes: 15 additions & 15 deletions tests/unit/test_compiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ def test_no_alias_for_known_tables(faux_conn, metadata):

expected_sql = (
"SELECT `table1`.`foo` \n"
"FROM `table1`, unnest(`table1`.`bar`) AS `anon_1` \n"
"FROM `table1`, UNNEST(`table1`.`bar`) AS `anon_1` \n"
"WHERE `anon_1` = %(param_1:INT64)s"
)
found_sql = q.compile(faux_conn).string
Expand All @@ -116,7 +116,7 @@ def test_no_alias_for_known_tables_cte(faux_conn, metadata):

expected_initial_sql = (
"SELECT `table1`.`foo`, `bar` \n"
"FROM `table1`, unnest(`table1`.`bars`) AS `bar`"
"FROM `table1`, UNNEST(`table1`.`bars`) AS `bar`"
)
found_initial_sql = q.compile(faux_conn).string
assert found_initial_sql == expected_initial_sql
Expand All @@ -127,7 +127,7 @@ def test_no_alias_for_known_tables_cte(faux_conn, metadata):
expected_cte_sql = (
"WITH `cte` AS \n"
"(SELECT `table1`.`foo` AS `foo`, `bar` \n"
"FROM `table1`, unnest(`table1`.`bars`) AS `bar`)\n"
"FROM `table1`, UNNEST(`table1`.`bars`) AS `bar`)\n"
" SELECT `cte`.`foo`, `cte`.`bar` \n"
"FROM `cte`"
)
Expand Down Expand Up @@ -196,7 +196,7 @@ def test_no_implicit_join_asterix_for_inner_unnest_before_2_0(faux_conn, metadat
q = prepare_implicit_join_base_query(faux_conn, metadata, True, False)
expected_initial_sql = (
"SELECT `table1`.`foo`, `table2`.`bar` \n"
"FROM `table2`, unnest(`table2`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`"
"FROM `table2`, UNNEST(`table2`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`"
)
found_initial_sql = q.compile(faux_conn).string
assert found_initial_sql == expected_initial_sql
Expand All @@ -207,7 +207,7 @@ def test_no_implicit_join_asterix_for_inner_unnest_before_2_0(faux_conn, metadat
expected_outer_sql = (
"SELECT * \n"
"FROM (SELECT `table1`.`foo` AS `foo`, `table2`.`bar` AS `bar` \n"
"FROM `table2`, unnest(`table2`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`) AS `anon_1`"
"FROM `table2`, UNNEST(`table2`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`) AS `anon_1`"
)
found_outer_sql = q.compile(faux_conn).string
assert found_outer_sql == expected_outer_sql
Expand All @@ -219,7 +219,7 @@ def test_no_implicit_join_asterix_for_inner_unnest(faux_conn, metadata):
q = prepare_implicit_join_base_query(faux_conn, metadata, True, False)
expected_initial_sql = (
"SELECT `table1`.`foo`, `table2`.`bar` \n"
"FROM unnest(`table2`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`, `table2`"
"FROM UNNEST(`table2`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`, `table2`"
)
found_initial_sql = q.compile(faux_conn).string
assert found_initial_sql == expected_initial_sql
Expand All @@ -230,7 +230,7 @@ def test_no_implicit_join_asterix_for_inner_unnest(faux_conn, metadata):
expected_outer_sql = (
"SELECT * \n"
"FROM (SELECT `table1`.`foo` AS `foo`, `table2`.`bar` AS `bar` \n"
"FROM unnest(`table2`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`, `table2`) AS `anon_1`"
"FROM UNNEST(`table2`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`, `table2`) AS `anon_1`"
)
found_outer_sql = q.compile(faux_conn).string
assert found_outer_sql == expected_outer_sql
Expand All @@ -242,7 +242,7 @@ def test_no_implicit_join_for_inner_unnest_before_2_0(faux_conn, metadata):
q = prepare_implicit_join_base_query(faux_conn, metadata, True, False)
expected_initial_sql = (
"SELECT `table1`.`foo`, `table2`.`bar` \n"
"FROM `table2`, unnest(`table2`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`"
"FROM `table2`, UNNEST(`table2`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`"
)
found_initial_sql = q.compile(faux_conn).string
assert found_initial_sql == expected_initial_sql
Expand All @@ -253,7 +253,7 @@ def test_no_implicit_join_for_inner_unnest_before_2_0(faux_conn, metadata):
expected_outer_sql = (
"SELECT `anon_1`.`foo` \n"
"FROM (SELECT `table1`.`foo` AS `foo`, `table2`.`bar` AS `bar` \n"
"FROM `table2`, unnest(`table2`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`) AS `anon_1`"
"FROM `table2`, UNNEST(`table2`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`) AS `anon_1`"
)
found_outer_sql = q.compile(faux_conn).string
assert found_outer_sql == expected_outer_sql
Expand All @@ -265,7 +265,7 @@ def test_no_implicit_join_for_inner_unnest(faux_conn, metadata):
q = prepare_implicit_join_base_query(faux_conn, metadata, True, False)
expected_initial_sql = (
"SELECT `table1`.`foo`, `table2`.`bar` \n"
"FROM unnest(`table2`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`, `table2`"
"FROM UNNEST(`table2`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`, `table2`"
)
found_initial_sql = q.compile(faux_conn).string
assert found_initial_sql == expected_initial_sql
Expand All @@ -276,7 +276,7 @@ def test_no_implicit_join_for_inner_unnest(faux_conn, metadata):
expected_outer_sql = (
"SELECT `anon_1`.`foo` \n"
"FROM (SELECT `table1`.`foo` AS `foo`, `table2`.`bar` AS `bar` \n"
"FROM unnest(`table2`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`, `table2`) AS `anon_1`"
"FROM UNNEST(`table2`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`, `table2`) AS `anon_1`"
)
found_outer_sql = q.compile(faux_conn).string
assert found_outer_sql == expected_outer_sql
Expand All @@ -289,7 +289,7 @@ def test_no_implicit_join_asterix_for_inner_unnest_no_table2_column(
q = prepare_implicit_join_base_query(faux_conn, metadata, False, False)
expected_initial_sql = (
"SELECT `table1`.`foo` \n"
"FROM `table2` `table2_1`, unnest(`table2_1`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`"
"FROM `table2` `table2_1`, UNNEST(`table2_1`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`"
)
found_initial_sql = q.compile(faux_conn).string
assert found_initial_sql == expected_initial_sql
Expand All @@ -300,7 +300,7 @@ def test_no_implicit_join_asterix_for_inner_unnest_no_table2_column(
expected_outer_sql = (
"SELECT * \n"
"FROM (SELECT `table1`.`foo` AS `foo` \n"
"FROM `table2` `table2_1`, unnest(`table2_1`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`) AS `anon_1`"
"FROM `table2` `table2_1`, UNNEST(`table2_1`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`) AS `anon_1`"
)
found_outer_sql = q.compile(faux_conn).string
assert found_outer_sql == expected_outer_sql
Expand All @@ -311,7 +311,7 @@ def test_no_implicit_join_for_inner_unnest_no_table2_column(faux_conn, metadata)
q = prepare_implicit_join_base_query(faux_conn, metadata, False, False)
expected_initial_sql = (
"SELECT `table1`.`foo` \n"
"FROM `table2` `table2_1`, unnest(`table2_1`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`"
"FROM `table2` `table2_1`, UNNEST(`table2_1`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`"
)
found_initial_sql = q.compile(faux_conn).string
assert found_initial_sql == expected_initial_sql
Expand All @@ -322,7 +322,7 @@ def test_no_implicit_join_for_inner_unnest_no_table2_column(faux_conn, metadata)
expected_outer_sql = (
"SELECT `anon_1`.`foo` \n"
"FROM (SELECT `table1`.`foo` AS `foo` \n"
"FROM `table2` `table2_1`, unnest(`table2_1`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`) AS `anon_1`"
"FROM `table2` `table2_1`, UNNEST(`table2_1`.`foos`) AS `unnested_foos` JOIN `table1` ON `table1`.`foo` = `unnested_foos`) AS `anon_1`"
)
found_outer_sql = q.compile(faux_conn).string
assert found_outer_sql == expected_outer_sql
Expand Down
4 changes: 2 additions & 2 deletions tests/unit/test_select.py
Original file line number Diff line number Diff line change
Expand Up @@ -419,7 +419,7 @@ def test_unnest(faux_conn, alias):
query = fcall.column_valued("foo_objects")
compiled = str(sqlalchemy.select(query).compile(faux_conn.engine))
assert " ".join(compiled.strip().split()) == (
"SELECT `foo_objects` FROM `t` `t_1`, unnest(`t_1`.`objects`) AS `foo_objects`"
"SELECT `foo_objects` FROM `t` `t_1`, UNNEST(`t_1`.`objects`) AS `foo_objects`"
)


Expand Down Expand Up @@ -450,7 +450,7 @@ def test_unnest_w_no_table_references(faux_conn, alias):
query = fcall.column_valued()
compiled = str(sqlalchemy.select(query).compile(faux_conn.engine))
assert " ".join(compiled.strip().split()) == (
"SELECT `anon_1` FROM unnest(%(unnest_1)s) AS `anon_1`"
"SELECT `anon_1` FROM UNNEST(%(param_1)s) AS `anon_1`"
)


Expand Down
5 changes: 4 additions & 1 deletion tests/unit/test_sqlalchemy_bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -226,7 +226,10 @@ def test_unnest_function(args, kw):

f = sqlalchemy.func.unnest(*args, **kw)
assert isinstance(f.type, sqlalchemy.String)
assert isinstance(sqlalchemy.select(f).subquery().c.unnest.type, sqlalchemy.String)
assert isinstance(
sqlalchemy.select(f.label("unnested_value")).subquery().c.unnested_value.type,
sqlalchemy.String,
)


@mock.patch("sqlalchemy_bigquery._helpers.create_bigquery_client")
Expand Down
Loading