fix(sql_parse): Ensure table extraction handles Jinja templating (#27470)

This commit is contained in:
John Bodley
2024-03-22 13:39:28 +13:00
committed by Michael S. Molina
parent 4ff331a66c
commit 7c14968e6d
7 changed files with 163 additions and 53 deletions

View File

@@ -32,6 +32,7 @@ from superset.exceptions import (
from superset.sql_parse import (
add_table_name,
extract_table_references,
extract_tables_from_jinja_sql,
get_rls_for_table,
has_table_query,
insert_rls_as_subquery,
@@ -1874,3 +1875,43 @@ WITH t AS (
)
SELECT * FROM t"""
).is_select()
@pytest.mark.parametrize(
"engine",
[
"hive",
"presto",
"trino",
],
)
@pytest.mark.parametrize(
"macro",
[
"latest_partition('foo.bar')",
"latest_sub_partition('foo.bar', baz='qux')",
],
)
@pytest.mark.parametrize(
"sql,expected",
[
(
"SELECT '{{{{ {engine}.{macro} }}}}'",
{Table(table="bar", schema="foo")},
),
(
"SELECT * FROM foo.baz WHERE quux = '{{{{ {engine}.{macro} }}}}'",
{Table(table="bar", schema="foo"), Table(table="baz", schema="foo")},
),
],
)
def test_extract_tables_from_jinja_sql(
engine: str,
macro: str,
sql: str,
expected: set[Table],
) -> None:
assert (
extract_tables_from_jinja_sql(sql.format(engine=engine, macro=macro), engine)
== expected
)