mirror of https://github.com/apache/superset.git
fix: Leverage actual database for rendering Jinjarized SQL (#27646)
This commit is contained in:
parent
ed9e542781
commit
28cbedb82f
|
@ -73,7 +73,7 @@ class SqlTablesMixin: # pylint: disable=too-few-public-methods
|
||||||
return list(
|
return list(
|
||||||
extract_tables_from_jinja_sql(
|
extract_tables_from_jinja_sql(
|
||||||
self.sql, # type: ignore
|
self.sql, # type: ignore
|
||||||
self.database.db_engine_spec.engine, # type: ignore
|
self.database, # type: ignore
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
except SupersetSecurityException:
|
except SupersetSecurityException:
|
||||||
|
|
|
@ -1963,9 +1963,7 @@ class SupersetSecurityManager( # pylint: disable=too-many-public-methods
|
||||||
default_schema = database.get_default_schema_for_query(query)
|
default_schema = database.get_default_schema_for_query(query)
|
||||||
tables = {
|
tables = {
|
||||||
Table(table_.table, table_.schema or default_schema)
|
Table(table_.table, table_.schema or default_schema)
|
||||||
for table_ in extract_tables_from_jinja_sql(
|
for table_ in extract_tables_from_jinja_sql(query.sql, database)
|
||||||
query.sql, database.db_engine_spec.engine
|
|
||||||
)
|
|
||||||
}
|
}
|
||||||
elif table:
|
elif table:
|
||||||
tables = {table}
|
tables = {table}
|
||||||
|
|
|
@ -25,8 +25,7 @@ import re
|
||||||
import urllib.parse
|
import urllib.parse
|
||||||
from collections.abc import Iterable, Iterator
|
from collections.abc import Iterable, Iterator
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from typing import Any, cast, Generic, TypeVar
|
from typing import Any, cast, Generic, TYPE_CHECKING, TypeVar
|
||||||
from unittest.mock import Mock
|
|
||||||
|
|
||||||
import sqlglot
|
import sqlglot
|
||||||
import sqlparse
|
import sqlparse
|
||||||
|
@ -75,6 +74,9 @@ try:
|
||||||
except (ImportError, ModuleNotFoundError):
|
except (ImportError, ModuleNotFoundError):
|
||||||
sqloxide_parse = None
|
sqloxide_parse = None
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from superset.models.core import Database
|
||||||
|
|
||||||
RESULT_OPERATIONS = {"UNION", "INTERSECT", "EXCEPT", "SELECT"}
|
RESULT_OPERATIONS = {"UNION", "INTERSECT", "EXCEPT", "SELECT"}
|
||||||
ON_KEYWORD = "ON"
|
ON_KEYWORD = "ON"
|
||||||
PRECEDES_TABLE_NAME = {"FROM", "JOIN", "DESCRIBE", "WITH", "LEFT JOIN", "RIGHT JOIN"}
|
PRECEDES_TABLE_NAME = {"FROM", "JOIN", "DESCRIBE", "WITH", "LEFT JOIN", "RIGHT JOIN"}
|
||||||
|
@ -1509,7 +1511,7 @@ def extract_table_references(
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
def extract_tables_from_jinja_sql(sql: str, engine: str | None = None) -> set[Table]:
|
def extract_tables_from_jinja_sql(sql: str, database: Database) -> set[Table]:
|
||||||
"""
|
"""
|
||||||
Extract all table references in the Jinjafied SQL statement.
|
Extract all table references in the Jinjafied SQL statement.
|
||||||
|
|
||||||
|
@ -1522,7 +1524,7 @@ def extract_tables_from_jinja_sql(sql: str, engine: str | None = None) -> set[Ta
|
||||||
SQLGlot.
|
SQLGlot.
|
||||||
|
|
||||||
:param sql: The Jinjafied SQL statement
|
:param sql: The Jinjafied SQL statement
|
||||||
:param engine: The associated database engine
|
:param database: The database associated with the SQL statement
|
||||||
:returns: The set of tables referenced in the SQL statement
|
:returns: The set of tables referenced in the SQL statement
|
||||||
:raises SupersetSecurityException: If SQLGlot is unable to parse the SQL statement
|
:raises SupersetSecurityException: If SQLGlot is unable to parse the SQL statement
|
||||||
"""
|
"""
|
||||||
|
@ -1531,8 +1533,7 @@ def extract_tables_from_jinja_sql(sql: str, engine: str | None = None) -> set[Ta
|
||||||
get_template_processor,
|
get_template_processor,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Mock the required database as the processor signature is exposed publically.
|
processor = get_template_processor(database)
|
||||||
processor = get_template_processor(database=Mock(backend=engine))
|
|
||||||
template = processor.env.parse(sql)
|
template = processor.env.parse(sql)
|
||||||
|
|
||||||
tables = set()
|
tables = set()
|
||||||
|
@ -1562,6 +1563,6 @@ def extract_tables_from_jinja_sql(sql: str, engine: str | None = None) -> set[Ta
|
||||||
tables
|
tables
|
||||||
| ParsedQuery(
|
| ParsedQuery(
|
||||||
sql_statement=processor.process_template(template),
|
sql_statement=processor.process_template(template),
|
||||||
engine=engine,
|
engine=database.db_engine_spec.engine,
|
||||||
).tables
|
).tables
|
||||||
)
|
)
|
||||||
|
|
|
@ -17,6 +17,7 @@
|
||||||
# pylint: disable=invalid-name, redefined-outer-name, too-many-lines
|
# pylint: disable=invalid-name, redefined-outer-name, too-many-lines
|
||||||
|
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
|
from unittest.mock import Mock
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
import sqlparse
|
import sqlparse
|
||||||
|
@ -1959,7 +1960,10 @@ def test_extract_tables_from_jinja_sql(
|
||||||
expected: set[Table],
|
expected: set[Table],
|
||||||
) -> None:
|
) -> None:
|
||||||
assert (
|
assert (
|
||||||
extract_tables_from_jinja_sql(sql.format(engine=engine, macro=macro), engine)
|
extract_tables_from_jinja_sql(
|
||||||
|
sql=sql.format(engine=engine, macro=macro),
|
||||||
|
database=Mock(),
|
||||||
|
)
|
||||||
== expected
|
== expected
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue