Skip to content

Commit

Permalink
fix: Provide more inclusive error handling for saved queries
Browse files Browse the repository at this point in the history
  • Loading branch information
john-bodley committed Mar 25, 2024
1 parent 4fce940 commit a8b35e1
Show file tree
Hide file tree
Showing 3 changed files with 39 additions and 1 deletion.
3 changes: 2 additions & 1 deletion superset/models/sql_lab.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
from flask_appbuilder.models.decorators import renders
from flask_babel import gettext as __
from humanize import naturaltime
from jinja2.exceptions import TemplateError
from sqlalchemy import (
Boolean,
Column,
Expand Down Expand Up @@ -76,7 +77,7 @@ def sql_tables(self) -> list[Table]:
self.database.db_engine_spec.engine, # type: ignore
)
)
except SupersetSecurityException:
except (SupersetSecurityException, TemplateError):
return []


Expand Down
1 change: 1 addition & 0 deletions superset/sql_parse.py
Original file line number Diff line number Diff line change
Expand Up @@ -1525,6 +1525,7 @@ def extract_tables_from_jinja_sql(sql: str, engine: str | None = None) -> set[Ta
:param engine: The associated database engine
:returns: The set of tables referenced in the SQL statement
:raises SupersetSecurityException: If SQLGlot is unable to parse the SQL statement
:raises jinja2.exceptions.TemplateError: If the Jinjafied SQL could not be rendered
"""

from superset.jinja_context import ( # pylint: disable=import-outside-toplevel
Expand Down
36 changes: 36 additions & 0 deletions tests/unit_tests/models/sql_lab_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
from unittest.mock import MagicMock

import pytest
from flask_appbuilder import Model
from jinja2.exceptions import TemplateError
from pytest_mock import MockFixture

from superset.exceptions import SupersetSecurityException
from superset.models.sql_lab import Query, SavedQuery, SqlTablesMixin


@pytest.mark.parametrize(
"klass",
[
Query,
SavedQuery,
],
)
@pytest.mark.parametrize(
"exception",
[
SupersetSecurityException,
TemplateError,
]
)
def test_sql_tables_mixin_sql_tables_exception(
klass: type[Model],
exception: type[Exception],
mocker: MockFixture,
) -> None:
mocker.patch(
"superset.models.sql_lab.extract_tables_from_jinja_sql",
side_effect=TemplateError(),
)

assert klass(sql="SELECT 1", database=MagicMock()).sql_tables == []

0 comments on commit a8b35e1

Please sign in to comment.