fix: Pre-query normalization with custom SQL (#30389)

This commit is contained in:
Michael S. Molina 2024-09-25 15:30:53 -03:00 committed by GitHub
parent 69d5f766a8
commit ad2998598f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 27 additions and 3 deletions

View File

@ -1688,10 +1688,10 @@ class SqlaTable(
if isinstance(value, np.generic):
value = value.item()
column_ = columns_by_name[dimension]
column_ = columns_by_name.get(dimension)
db_extra: dict[str, Any] = self.database.get_extra()
if column_.type and column_.is_temporal and isinstance(value, str):
if column_ and column_.type and column_.is_temporal and isinstance(value, str):
sql = self.db_engine_spec.convert_dttm(
column_.type, dateutil.parser.parse(value), db_extra=db_extra
)

View File

@ -15,13 +15,14 @@
# specific language governing permissions and limitations
# under the License.
import pandas as pd
import pytest
from pytest_mock import MockerFixture
from sqlalchemy import create_engine
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm.session import Session
from superset.connectors.sqla.models import SqlaTable
from superset.connectors.sqla.models import SqlaTable, TableColumn
from superset.daos.dataset import DatasetDAO
from superset.exceptions import OAuth2RedirectError
from superset.models.core import Database
@ -263,3 +264,26 @@ def test_dataset_uniqueness(session: Session) -> None:
database,
Table("table", "schema", "some_catalog"),
)
def test_normalize_prequery_result_type_custom_sql() -> None:
"""
Test that the `_normalize_prequery_result_type` can hanndle custom SQL.
"""
sqla_table = SqlaTable(
table_name="my_sqla_table",
columns=[],
metrics=[],
database=Database(database_name="my_db", sqlalchemy_uri="sqlite://"),
)
row: pd.Series = {
"custom_sql": "Car",
}
dimension: str = "custom_sql"
columns_by_name: dict[str, TableColumn] = {
"product_line": TableColumn(column_name="product_line"),
}
assert (
sqla_table._normalize_prequery_result_type(row, dimension, columns_by_name)
== "Car"
)