Skip to content

Commit

Permalink
fix: pass string to process_template (apache#31329)
Browse files Browse the repository at this point in the history
  • Loading branch information
betodealmeida authored Dec 7, 2024
1 parent 592564b commit 9315a88
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 5 deletions.
7 changes: 5 additions & 2 deletions superset/sql_parse.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@

import sqlparse
from flask_babel import gettext as __
from jinja2 import nodes
from jinja2 import nodes, Template
from sqlalchemy import and_
from sqlparse import keywords
from sqlparse.lexer import Lexer
Expand Down Expand Up @@ -999,10 +999,13 @@ def extract_tables_from_jinja_sql(sql: str, database: Database) -> set[Table]:
node.fields = nodes.TemplateData.fields
node.data = "NULL"

# re-render template back into a string
rendered_template = Template(template).render()

return (
tables
| ParsedQuery(
sql_statement=processor.process_template(template),
sql_statement=processor.process_template(rendered_template),
engine=database.db_engine_spec.engine,
).tables
)
27 changes: 24 additions & 3 deletions tests/unit_tests/sql_parse_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
# pylint: disable=invalid-name, redefined-outer-name, too-many-lines

from typing import Optional
from unittest.mock import Mock
from unittest import mock

import pytest
import sqlparse
Expand Down Expand Up @@ -1888,12 +1888,33 @@ def test_is_select() -> None:
],
)
def test_extract_tables_from_jinja_sql(
engine: str, macro: str, expected: set[Table]
mocker: MockerFixture,
engine: str,
macro: str,
expected: set[Table],
) -> None:
assert (
extract_tables_from_jinja_sql(
sql=f"'{{{{ {engine}.{macro} }}}}'",
database=Mock(),
database=mocker.Mock(),
)
== expected
)


@mock.patch.dict(
"superset.extensions.feature_flag_manager._feature_flags",
{"ENABLE_TEMPLATE_PROCESSING": False},
clear=True,
)
def test_extract_tables_from_jinja_sql_disabled(mocker: MockerFixture) -> None:
"""
Test the function when the feature flag is disabled.
"""
database = mocker.Mock()
database.db_engine_spec.engine = "mssql"

assert extract_tables_from_jinja_sql(
sql="SELECT 1 FROM t",
database=database,
) == {Table("t")}

0 comments on commit 9315a88

Please sign in to comment.