mirror of https://github.com/apache/superset.git
fix(sql_parse): Provide more lenient logic when extracting latest[_sub]_partition (#28152)
(cherry picked from commit c5e7d870f0
)
This commit is contained in:
parent
a8f7ea0720
commit
f413ae43c5
|
@ -1099,16 +1099,19 @@ def extract_tables_from_jinja_sql(sql: str, database: Database) -> set[Table]:
|
|||
"latest_partition",
|
||||
"latest_sub_partition",
|
||||
):
|
||||
# Extract the table referenced in the macro.
|
||||
tables.add(
|
||||
Table(
|
||||
*[
|
||||
remove_quotes(part.strip())
|
||||
for part in node.args[0].as_const().split(".")[::-1]
|
||||
if len(node.args) == 1
|
||||
]
|
||||
# Try to extract the table referenced in the macro.
|
||||
try:
|
||||
tables.add(
|
||||
Table(
|
||||
*[
|
||||
remove_quotes(part.strip())
|
||||
for part in node.args[0].as_const().split(".")[::-1]
|
||||
if len(node.args) == 1
|
||||
]
|
||||
)
|
||||
)
|
||||
)
|
||||
except nodes.Impossible:
|
||||
pass
|
||||
|
||||
# Replace the potentially problematic Jinja macro with some benign SQL.
|
||||
node.__class__ = nodes.TemplateData
|
||||
|
|
|
@ -1891,36 +1891,40 @@ SELECT * FROM t"""
|
|||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"macro",
|
||||
[
|
||||
"latest_partition('foo.bar')",
|
||||
"latest_partition(' foo.bar ')", # Non-atypical user error which works
|
||||
"latest_partition('foo.%s'|format('bar'))",
|
||||
"latest_sub_partition('foo.bar', baz='qux')",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"sql,expected",
|
||||
"macro,expected",
|
||||
[
|
||||
(
|
||||
"SELECT '{{{{ {engine}.{macro} }}}}'",
|
||||
"latest_partition('foo.bar')",
|
||||
{Table(table="bar", schema="foo")},
|
||||
),
|
||||
(
|
||||
"SELECT * FROM foo.baz WHERE quux = '{{{{ {engine}.{macro} }}}}'",
|
||||
{Table(table="bar", schema="foo"), Table(table="baz", schema="foo")},
|
||||
"latest_partition(' foo.bar ')", # Non-atypical user error which works
|
||||
{Table(table="bar", schema="foo")},
|
||||
),
|
||||
(
|
||||
"latest_partition('foo.%s'|format('bar'))",
|
||||
{Table(table="bar", schema="foo")},
|
||||
),
|
||||
(
|
||||
"latest_sub_partition('foo.bar', baz='qux')",
|
||||
{Table(table="bar", schema="foo")},
|
||||
),
|
||||
(
|
||||
"latest_partition('foo.%s'|format(str('bar')))",
|
||||
set(),
|
||||
),
|
||||
(
|
||||
"latest_partition('foo.{}'.format('bar'))",
|
||||
set(),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_extract_tables_from_jinja_sql(
|
||||
engine: str,
|
||||
macro: str,
|
||||
sql: str,
|
||||
expected: set[Table],
|
||||
engine: str, macro: str, expected: set[Table]
|
||||
) -> None:
|
||||
assert (
|
||||
extract_tables_from_jinja_sql(
|
||||
sql=sql.format(engine=engine, macro=macro),
|
||||
sql=f"'{{{{ {engine}.{macro} }}}}'",
|
||||
database=Mock(),
|
||||
)
|
||||
== expected
|
||||
|
|
Loading…
Reference in New Issue