chore: upgrade SQLAlchemy to 1.4 (#19890)

* chore: upgrade SQLAlchemy

* Convert integration test to unit test

* Fix SQLite

* Update method names/docstrings

* Skip test

* Fix SQLite
This commit is contained in:
Beto Dealmeida
2022-07-18 15:21:38 -07:00
committed by GitHub
parent 90600d1883
commit e60083b45b
32 changed files with 656 additions and 255 deletions

View File

@@ -0,0 +1,330 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=import-outside-toplevel, unused-argument, redefined-outer-name, invalid-name
from functools import partial
from typing import Any, Dict, TYPE_CHECKING
import pytest
from pytest_mock import MockerFixture
from sqlalchemy.orm.session import Session
if TYPE_CHECKING:
from superset.connectors.sqla.models import SqlaTable
FULL_DTTM_DEFAULTS_EXAMPLE = {
"main_dttm_col": "id",
"dttm_columns": {
"dttm": {
"python_date_format": "epoch_s",
"expression": "CAST(dttm as INTEGER)",
},
"id": {"python_date_format": "epoch_ms"},
"month": {
"python_date_format": "%Y-%m-%d",
"expression": (
"CASE WHEN length(month) = 7 THEN month || '-01' ELSE month END"
),
},
},
}
def apply_dttm_defaults(table: "SqlaTable", dttm_defaults: Dict[str, Any]) -> None:
"""Applies dttm defaults to the table, mutates in place."""
for dbcol in table.columns:
# Set is_dttm is column is listed in dttm_columns.
if dbcol.column_name in dttm_defaults.get("dttm_columns", {}):
dbcol.is_dttm = True
# Skip non dttm columns.
if dbcol.column_name not in dttm_defaults.get("dttm_columns", {}):
continue
# Set table main_dttm_col.
if dbcol.column_name == dttm_defaults.get("main_dttm_col"):
table.main_dttm_col = dbcol.column_name
# Apply defaults if empty.
dttm_column_defaults = dttm_defaults.get("dttm_columns", {}).get(
dbcol.column_name, {}
)
dbcol.is_dttm = True
if (
not dbcol.python_date_format
and "python_date_format" in dttm_column_defaults
):
dbcol.python_date_format = dttm_column_defaults["python_date_format"]
if not dbcol.expression and "expression" in dttm_column_defaults:
dbcol.expression = dttm_column_defaults["expression"]
@pytest.fixture
def test_table(app_context: None, session: Session) -> "SqlaTable":
"""
Fixture that generates an in-memory table.
"""
from superset.connectors.sqla.models import SqlaTable, TableColumn
from superset.models.core import Database
engine = session.get_bind()
SqlaTable.metadata.create_all(engine) # pylint: disable=no-member
columns = [
TableColumn(column_name="ds", is_dttm=1, type="TIMESTAMP"),
TableColumn(column_name="event_time", is_dttm=1, type="TIMESTAMP"),
TableColumn(column_name="id", type="INTEGER"),
TableColumn(column_name="dttm", type="INTEGER"),
TableColumn(column_name="duration_ms", type="INTEGER"),
]
return SqlaTable(
table_name="test_table",
columns=columns,
metrics=[],
main_dttm_col=None,
database=Database(database_name="my_database", sqlalchemy_uri="sqlite://"),
)
def test_main_dttm_col(mocker: MockerFixture, test_table: "SqlaTable") -> None:
"""
Test the ``SQLA_TABLE_MUTATOR`` config.
"""
dttm_defaults = {
"main_dttm_col": "event_time",
"dttm_columns": {"ds": {}, "event_time": {}},
}
mocker.patch(
"superset.connectors.sqla.models.config",
new={
"SQLA_TABLE_MUTATOR": partial(
apply_dttm_defaults,
dttm_defaults=dttm_defaults,
)
},
)
mocker.patch(
"superset.connectors.sqla.models.get_physical_table_metadata",
return_value=[
{"name": "ds", "type": "TIMESTAMP", "is_dttm": True},
{"name": "event_time", "type": "TIMESTAMP", "is_dttm": True},
{"name": "id", "type": "INTEGER", "is_dttm": False},
],
)
assert test_table.main_dttm_col is None
test_table.fetch_metadata()
assert test_table.main_dttm_col == "event_time"
def test_main_dttm_col_nonexistent(
mocker: MockerFixture,
test_table: "SqlaTable",
) -> None:
"""
Test the ``SQLA_TABLE_MUTATOR`` config when main datetime column doesn't exist.
"""
dttm_defaults = {
"main_dttm_col": "nonexistent",
}
mocker.patch(
"superset.connectors.sqla.models.config",
new={
"SQLA_TABLE_MUTATOR": partial(
apply_dttm_defaults,
dttm_defaults=dttm_defaults,
)
},
)
mocker.patch(
"superset.connectors.sqla.models.get_physical_table_metadata",
return_value=[
{"name": "ds", "type": "TIMESTAMP", "is_dttm": True},
{"name": "event_time", "type": "TIMESTAMP", "is_dttm": True},
{"name": "id", "type": "INTEGER", "is_dttm": False},
],
)
assert test_table.main_dttm_col is None
test_table.fetch_metadata()
# fall back to ds
assert test_table.main_dttm_col == "ds"
def test_main_dttm_col_nondttm(
mocker: MockerFixture,
test_table: "SqlaTable",
) -> None:
"""
Test the ``SQLA_TABLE_MUTATOR`` config when main datetime column has wrong type.
"""
dttm_defaults = {
"main_dttm_col": "id",
}
mocker.patch(
"superset.connectors.sqla.models.config",
new={
"SQLA_TABLE_MUTATOR": partial(
apply_dttm_defaults,
dttm_defaults=dttm_defaults,
)
},
)
mocker.patch(
"superset.connectors.sqla.models.get_physical_table_metadata",
return_value=[
{"name": "ds", "type": "TIMESTAMP", "is_dttm": True},
{"name": "event_time", "type": "TIMESTAMP", "is_dttm": True},
{"name": "id", "type": "INTEGER", "is_dttm": False},
],
)
assert test_table.main_dttm_col is None
test_table.fetch_metadata()
# fall back to ds
assert test_table.main_dttm_col == "ds"
def test_python_date_format_by_column_name(
mocker: MockerFixture,
test_table: "SqlaTable",
) -> None:
"""
Test the ``SQLA_TABLE_MUTATOR`` setting for "python_date_format".
"""
table_defaults = {
"dttm_columns": {
"id": {"python_date_format": "epoch_ms"},
"dttm": {"python_date_format": "epoch_s"},
"duration_ms": {"python_date_format": "invalid"},
},
}
mocker.patch(
"superset.connectors.sqla.models.config",
new={
"SQLA_TABLE_MUTATOR": partial(
apply_dttm_defaults,
dttm_defaults=table_defaults,
)
},
)
mocker.patch(
"superset.connectors.sqla.models.get_physical_table_metadata",
return_value=[
{"name": "id", "type": "INTEGER", "is_dttm": False},
{"name": "dttm", "type": "INTEGER", "is_dttm": False},
{"name": "duration_ms", "type": "INTEGER", "is_dttm": False},
],
)
test_table.fetch_metadata()
id_col = [c for c in test_table.columns if c.column_name == "id"][0]
assert id_col.is_dttm
assert id_col.python_date_format == "epoch_ms"
dttm_col = [c for c in test_table.columns if c.column_name == "dttm"][0]
assert dttm_col.is_dttm
assert dttm_col.python_date_format == "epoch_s"
duration_ms_col = [c for c in test_table.columns if c.column_name == "duration_ms"][
0
]
assert duration_ms_col.is_dttm
assert duration_ms_col.python_date_format == "invalid"
def test_expression_by_column_name(
mocker: MockerFixture,
test_table: "SqlaTable",
) -> None:
"""
Test the ``SQLA_TABLE_MUTATOR`` setting for expression.
"""
table_defaults = {
"dttm_columns": {
"dttm": {"expression": "CAST(dttm as INTEGER)"},
"duration_ms": {"expression": "CAST(duration_ms as DOUBLE)"},
},
}
mocker.patch(
"superset.connectors.sqla.models.config",
new={
"SQLA_TABLE_MUTATOR": partial(
apply_dttm_defaults,
dttm_defaults=table_defaults,
)
},
)
mocker.patch(
"superset.connectors.sqla.models.get_physical_table_metadata",
return_value=[
{"name": "dttm", "type": "INTEGER", "is_dttm": False},
{"name": "duration_ms", "type": "INTEGER", "is_dttm": False},
],
)
test_table.fetch_metadata()
dttm_col = [c for c in test_table.columns if c.column_name == "dttm"][0]
assert dttm_col.is_dttm
assert dttm_col.expression == "CAST(dttm as INTEGER)"
duration_ms_col = [c for c in test_table.columns if c.column_name == "duration_ms"][
0
]
assert duration_ms_col.is_dttm
assert duration_ms_col.expression == "CAST(duration_ms as DOUBLE)"
def test_full_setting(
mocker: MockerFixture,
test_table: "SqlaTable",
) -> None:
"""
Test the ``SQLA_TABLE_MUTATOR`` with full settings.
"""
mocker.patch(
"superset.connectors.sqla.models.config",
new={
"SQLA_TABLE_MUTATOR": partial(
apply_dttm_defaults,
dttm_defaults=FULL_DTTM_DEFAULTS_EXAMPLE,
)
},
)
mocker.patch(
"superset.connectors.sqla.models.get_physical_table_metadata",
return_value=[
{"name": "id", "type": "INTEGER", "is_dttm": False},
{"name": "dttm", "type": "INTEGER", "is_dttm": False},
{"name": "duration_ms", "type": "INTEGER", "is_dttm": False},
],
)
test_table.fetch_metadata()
id_col = [c for c in test_table.columns if c.column_name == "id"][0]
assert id_col.is_dttm
assert id_col.python_date_format == "epoch_ms"
assert id_col.expression == ""
dttm_col = [c for c in test_table.columns if c.column_name == "dttm"][0]
assert dttm_col.is_dttm
assert dttm_col.python_date_format == "epoch_s"
assert dttm_col.expression == "CAST(dttm as INTEGER)"

View File

@@ -47,10 +47,12 @@ def get_session(mocker: MockFixture) -> Callable[[], Session]:
in_memory_session.remove = lambda: None
# patch session
mocker.patch(
get_session = mocker.patch(
"superset.security.SupersetSecurityManager.get_session",
return_value=in_memory_session,
)
get_session.return_value = in_memory_session
# FAB calls get_session.get_bind() to get a handler to the engine
get_session.get_bind.return_value = engine
mocker.patch("superset.db.session", in_memory_session)
return in_memory_session

View File

@@ -124,11 +124,11 @@ def test_import_dataset(app_context: None, session: Session) -> None:
assert len(sqla_table.columns) == 1
assert sqla_table.columns[0].column_name == "profit"
assert sqla_table.columns[0].verbose_name is None
assert sqla_table.columns[0].is_dttm is None
assert sqla_table.columns[0].is_active is None
assert sqla_table.columns[0].is_dttm is False
assert sqla_table.columns[0].is_active is True
assert sqla_table.columns[0].type == "INTEGER"
assert sqla_table.columns[0].groupby is None
assert sqla_table.columns[0].filterable is None
assert sqla_table.columns[0].groupby is True
assert sqla_table.columns[0].filterable is True
assert sqla_table.columns[0].expression == "revenue-expenses"
assert sqla_table.columns[0].description is None
assert sqla_table.columns[0].python_date_format is None

View File

@@ -259,7 +259,6 @@ def test_dataset_attributes(app_context: None, session: Session) -> None:
"main_dttm_col",
"metrics",
"offset",
"owners",
"params",
"perm",
"schema",

View File

@@ -22,7 +22,7 @@ from pytest import raises
def test_odbc_impersonation(app_context: AppContext) -> None:
"""
Test ``modify_url_for_impersonation`` method when driver == odbc.
Test ``get_url_for_impersonation`` method when driver == odbc.
The method adds the parameter ``DelegationUID`` to the query string.
"""
@@ -32,13 +32,13 @@ def test_odbc_impersonation(app_context: AppContext) -> None:
url = URL("drill+odbc")
username = "DoAsUser"
DrillEngineSpec.modify_url_for_impersonation(url, True, username)
url = DrillEngineSpec.get_url_for_impersonation(url, True, username)
assert url.query["DelegationUID"] == username
def test_jdbc_impersonation(app_context: AppContext) -> None:
"""
Test ``modify_url_for_impersonation`` method when driver == jdbc.
Test ``get_url_for_impersonation`` method when driver == jdbc.
The method adds the parameter ``impersonation_target`` to the query string.
"""
@@ -48,13 +48,13 @@ def test_jdbc_impersonation(app_context: AppContext) -> None:
url = URL("drill+jdbc")
username = "DoAsUser"
DrillEngineSpec.modify_url_for_impersonation(url, True, username)
url = DrillEngineSpec.get_url_for_impersonation(url, True, username)
assert url.query["impersonation_target"] == username
def test_sadrill_impersonation(app_context: AppContext) -> None:
"""
Test ``modify_url_for_impersonation`` method when driver == sadrill.
Test ``get_url_for_impersonation`` method when driver == sadrill.
The method adds the parameter ``impersonation_target`` to the query string.
"""
@@ -64,13 +64,13 @@ def test_sadrill_impersonation(app_context: AppContext) -> None:
url = URL("drill+sadrill")
username = "DoAsUser"
DrillEngineSpec.modify_url_for_impersonation(url, True, username)
url = DrillEngineSpec.get_url_for_impersonation(url, True, username)
assert url.query["impersonation_target"] == username
def test_invalid_impersonation(app_context: AppContext) -> None:
"""
Test ``modify_url_for_impersonation`` method when driver == foobar.
Test ``get_url_for_impersonation`` method when driver == foobar.
The method raises an exception because impersonation is not supported
for drill+foobar.
@@ -84,4 +84,4 @@ def test_invalid_impersonation(app_context: AppContext) -> None:
username = "DoAsUser"
with raises(SupersetDBAPIProgrammingError):
DrillEngineSpec.modify_url_for_impersonation(url, True, username)
DrillEngineSpec.get_url_for_impersonation(url, True, username)