def setup_sample_data() -> Any:
# TODO(john-bodley): Determine a cleaner way of setting up the sample data without
# relying on `tests.integration_tests.test_app.app` leveraging an `app` fixture which is purposely
# scoped to the function level to ensure tests remain idempotent.
with app.app_context():
setup_presto_if_needed()
from superset.cli.test import load_test_users_run
load_test_users_run()
from superset.examples.css_templates import load_css_templates
load_css_templates()
yield
with app.app_context():
engine = get_example_database().get_sqla_engine()
# drop sqlachemy tables
db.session.commit()
from sqlalchemy.ext import declarative
sqla_base = declarative.declarative_base()
# uses sorted_tables to drop in proper order without violating foreign constrains
for table in sqla_base.metadata.sorted_tables:
table.__table__.drop()
db.session.commit()
def load_birth_names_data():
with app.app_context():
database = get_example_database()
df = _get_dataframe(database)
dtype = {
"ds": DateTime if database.backend != "presto" else String(255),
"gender": String(16),
"state": String(10),
"name": String(255),
}
df.to_sql(
BIRTH_NAMES_TBL_NAME,
database.get_sqla_engine(),
if_exists="replace",
chunksize=500,
dtype=dtype,
index=False,
method="multi",
schema=get_example_default_schema(),
)
yield
with app.app_context():
engine = get_example_database().get_sqla_engine()
engine.execute("DROP TABLE IF EXISTS birth_names")
def load_unicode_dashboard_with_position(load_unicode_data):
slice_name = "Unicode Cloud"
position = "{}"
with app.app_context():
dash = _create_unicode_dashboard(slice_name, position)
yield
_cleanup(dash, slice_name)
def test_df_to_sql_if_exists_replace_with_schema(mock_upload_to_s3, mock_g):
config = app.config.copy()
app.config["CSV_TO_HIVE_UPLOAD_DIRECTORY_FUNC"]: lambda *args: ""
mock_upload_to_s3.return_value = "mock-location"
mock_g.user = True
mock_database = mock.MagicMock()
mock_database.get_df.return_value.empty = False
mock_execute = mock.MagicMock(return_value=True)
mock_database.get_sqla_engine.return_value.execute = mock_execute
table_name = "foobar"
schema = "schema"
with app.app_context():
HiveEngineSpec.df_to_sql(
mock_database,
Table(table=table_name, schema=schema),
pd.DataFrame(),
{
"if_exists": "replace",
"header": 1,
"na_values": "mock",
"sep": "mock"
},
)
mock_execute.assert_any_call(f"DROP TABLE IF EXISTS {schema}.{table_name}")
app.config = config
def create_annotation_layers():
"""
Creates ANNOTATION_LAYERS_COUNT-1 layers with no annotations
and a final one with ANNOTATION_COUNT childs
:return:
"""
with app.app_context():
annotation_layers = []
annotations = []
for cx in range(ANNOTATION_LAYERS_COUNT - 1):
annotation_layers.append(
_insert_annotation_layer(name=f"name{cx}", descr=f"descr{cx}")
)
layer_with_annotations = _insert_annotation_layer("layer_with_annotations")
annotation_layers.append(layer_with_annotations)
for cx in range(ANNOTATIONS_COUNT):
annotations.append(
_insert_annotation(
layer_with_annotations,
short_descr=f"short_descr{cx}",
long_descr=f"long_descr{cx}",
start_dttm=get_start_dttm(cx),
end_dttm=get_end_dttm(cx),
)
)
yield annotation_layers
# rollback changes
for annotation_layer in annotation_layers:
db.session.delete(annotation_layer)
for annotation in annotations:
db.session.delete(annotation)
db.session.commit()
def _load_data():
table_name = "birth_names"
with app.app_context():
database = get_example_database()
df = _get_dataframe(database)
dtype = {
"ds": DateTime if database.backend != "presto" else String(255),
"gender": String(16),
"state": String(10),
"name": String(255),
}
table = _create_table(
df=df,
table_name=table_name,
database=database,
dtype=dtype,
fetch_values_predicate="123 = 123",
)
from superset.examples.birth_names import create_dashboard, create_slices
slices, _ = create_slices(table, admin_owner=False)
dash = create_dashboard(slices)
slices_ids_to_delete = [slice.id for slice in slices]
dash_id_to_delete = dash.id
return dash_id_to_delete, slices_ids_to_delete
def load_dashboard(self):
with app.app_context():
table = (db.session.query(SqlaTable).filter_by(
table_name="energy_usage").one())
# get a slice from the allowed table
slice = db.session.query(Slice).filter_by(
slice_name="Energy Sankey").one()
self.grant_public_access_to_table(table)
pytest.hidden_dash_slug = f"hidden_dash_{random()}"
pytest.published_dash_slug = f"published_dash_{random()}"
# Create a published and hidden dashboard and add them to the database
published_dash = Dashboard()
published_dash.dashboard_title = "Published Dashboard"
published_dash.slug = pytest.published_dash_slug
published_dash.slices = [slice]
published_dash.published = True
hidden_dash = Dashboard()
hidden_dash.dashboard_title = "Hidden Dashboard"
hidden_dash.slug = pytest.hidden_dash_slug
hidden_dash.slices = [slice]
hidden_dash.published = False
db.session.merge(published_dash)
db.session.merge(hidden_dash)
yield db.session.commit()
self.revoke_public_access_to_table(table)
db.session.delete(published_dash)
db.session.delete(hidden_dash)
db.session.commit()
def create_old_role(pvm_map: PvmMigrationMapType, external_pvms):
with app.app_context():
pvms = []
for old_pvm, new_pvms in pvm_map.items():
pvms.append(
security_manager.add_permission_view_menu(
old_pvm.permission, old_pvm.view))
for external_pvm in external_pvms:
pvms.append(
security_manager.find_permission_view_menu(
external_pvm.permission, external_pvm.view))
new_role = Role(name="Dummy Role", permissions=pvms)
db.session.add(new_role)
db.session.commit()
yield new_role
new_role = (db.session.query(Role).filter(
Role.name == "Dummy Role").one_or_none())
new_role.permissions = []
db.session.merge(new_role)
for old_pvm, new_pvms in pvm_map.items():
security_manager.del_permission_view_menu(old_pvm.permission,
old_pvm.view)
for new_pvm in new_pvms:
security_manager.del_permission_view_menu(
new_pvm.permission, new_pvm.view)
db.session.delete(new_role)
db.session.commit()
def load_energy_table_with_slice():
table_name = "energy_usage"
df = _get_dataframe()
with app.app_context():
_create_energy_table(df, table_name)
yield
_cleanup()
def load_energy_charts(self):
with app.app_context():
admin = self.get_user("admin")
energy_table = (db.session.query(SqlaTable).filter_by(
table_name="energy_usage").one_or_none())
energy_table_id = 1
if energy_table:
energy_table_id = energy_table.id
chart1 = self.insert_chart("foo_a", [admin.id],
energy_table_id,
description="ZY_bar")
chart2 = self.insert_chart("zy_foo", [admin.id],
energy_table_id,
description="desc1")
chart3 = self.insert_chart("foo_b", [admin.id],
energy_table_id,
description="desc1zy_")
chart4 = self.insert_chart("foo_c", [admin.id],
energy_table_id,
viz_type="viz_zy_")
chart5 = self.insert_chart("bar", [admin.id],
energy_table_id,
description="foo")
yield
# rollback changes
db.session.delete(chart1)
db.session.delete(chart2)
db.session.delete(chart3)
db.session.delete(chart4)
db.session.delete(chart5)
db.session.commit()
def __call__(self) -> Database:
with app.app_context():
if self._db is None:
self._db = get_example_database()
self._load_lazy_data_to_decouple_from_session()
return self._db
def setup_csv_upload():
with app.app_context():
login(test_client, username="******")
upload_db = superset.utils.database.get_or_create_db(
CSV_UPLOAD_DATABASE, app.config["SQLALCHEMY_EXAMPLES_URI"])
extra = upload_db.get_extra()
extra[
"explore_database_id"] = superset.utils.database.get_example_database(
).id
upload_db.extra = json.dumps(extra)
upload_db.allow_file_upload = True
db.session.commit()
yield
upload_db = get_upload_db()
engine = upload_db.get_sqla_engine()
engine.execute(f"DROP TABLE IF EXISTS {EXCEL_UPLOAD_TABLE}")
engine.execute(f"DROP TABLE IF EXISTS {CSV_UPLOAD_TABLE}")
engine.execute(f"DROP TABLE IF EXISTS {PARQUET_UPLOAD_TABLE}")
engine.execute(f"DROP TABLE IF EXISTS {CSV_UPLOAD_TABLE_W_SCHEMA}")
engine.execute(f"DROP TABLE IF EXISTS {CSV_UPLOAD_TABLE_W_EXPLORE}")
db.session.delete(upload_db)
db.session.commit()
def dashboard() -> Generator[Dashboard, None, None]:
dashboard: Dashboard
slice_: Slice
datasource: SqlaTable
database: Database
session: Session
try:
with app.app_context() as ctx:
dashboard_owner_user = security_manager.find_user(
DASHBOARD_OWNER_USERNAME)
database = create_database("test_database_filter_sets")
datasource = create_datasource_table(name="test_datasource",
database=database,
owners=[dashboard_owner_user])
slice_ = create_slice(datasource=datasource,
name="test_slice",
owners=[dashboard_owner_user])
dashboard = create_dashboard(
dashboard_title="test_dashboard",
published=True,
slices=[slice_],
owners=[dashboard_owner_user],
)
session = ctx.app.appbuilder.get_session
session.add(dashboard)
session.commit()
yield dashboard
except Exception as ex:
print(str(ex))
finally:
with app.app_context() as ctx:
session = ctx.app.appbuilder.get_session
try:
dashboard.owners = []
slice_.owners = []
datasource.owners = []
session.merge(dashboard)
session.merge(slice_)
session.merge(datasource)
session.commit()
session.delete(dashboard)
session.delete(slice_)
session.delete(datasource)
session.delete(database)
session.commit()
except Exception as ex:
print(str(ex))
def test_where_latest_partition_no_columns_no_values(mock_method):
mock_method.return_value = ("01-01-19", None)
db = mock.Mock()
with app.app_context():
result = HiveEngineSpec.where_latest_partition("test_table",
"test_schema", db,
select())
assert result is None
def load_unicode_data():
with app.app_context():
_get_dataframe().to_sql(
UNICODE_TBL_NAME,
get_example_database().get_sqla_engine(),
if_exists="replace",
chunksize=500,
dtype={"phrase": String(500)},
index=False,
method="multi",
schema=get_example_default_schema(),
)
yield
with app.app_context():
engine = get_example_database().get_sqla_engine()
engine.execute("DROP TABLE IF EXISTS unicode_test")
def test_upload_to_s3_no_bucket_path():
with app.app_context():
with pytest.raises(
Exception,
match=
"No upload bucket specified. You can specify one in the config file.",
):
upload_to_s3("filename", "prefix", Table("table"))
def create_dashboard_for_loaded_data():
with app.app_context():
table = create_table_metadata(WB_HEALTH_POPULATION, get_example_database())
slices = _create_world_bank_slices(table)
dash = _create_world_bank_dashboard(table, slices)
slices_ids_to_delete = [slice.id for slice in slices]
dash_id_to_delete = dash.id
return dash_id_to_delete, slices_ids_to_delete
def setUpClass(cls):
with app.app_context():
security_manager.add_role("override_me")
security_manager.add_role(TEST_ROLE_1)
security_manager.add_role(TEST_ROLE_2)
security_manager.add_role(DB_ACCESS_ROLE)
security_manager.add_role(SCHEMA_ACCESS_ROLE)
db.session.commit()
def delete_imports(cls):
with app.app_context():
# Imported data clean up
session = db.session
for table in session.query(SqlaTable):
if DBREF in table.params_dict:
session.delete(table)
session.commit()
def tearDownClass(cls):
with app.app_context():
override_me = security_manager.find_role("override_me")
db.session.delete(override_me)
db.session.delete(security_manager.find_role(TEST_ROLE_1))
db.session.delete(security_manager.find_role(TEST_ROLE_2))
db.session.delete(security_manager.find_role(DB_ACCESS_ROLE))
db.session.delete(security_manager.find_role(SCHEMA_ACCESS_ROLE))
db.session.commit()
def test_time_grain_denylist():
config = app.config.copy()
app.config["TIME_GRAIN_DENYLIST"] = ["PT1M"]
with app.app_context():
time_grain_functions = SqliteEngineSpec.get_time_grain_expressions()
assert not "PT1M" in time_grain_functions
app.config = config
def public_role_like_test_role():
with app.app_context():
app.config["PUBLIC_ROLE_LIKE"] = "TestRole"
security_manager.sync_role_definitions()
yield
security_manager.get_public_role().permissions = []
db.session.commit()
def load_energy_table_data():
with app.app_context():
database = get_example_database()
df = _get_dataframe()
df.to_sql(
ENERGY_USAGE_TBL_NAME,
database.get_sqla_engine(),
if_exists="replace",
chunksize=500,
index=False,
dtype={"source": String(255), "target": String(255), "value": Float()},
method="multi",
schema=get_example_default_schema(),
)
yield
with app.app_context():
engine = get_example_database().get_sqla_engine()
engine.execute("DROP TABLE IF EXISTS energy_usage")
def load_unicode_dashboard_with_position():
table_name = "unicode_test"
slice_name = "Unicode Cloud"
df = _get_dataframe()
position = "{}"
with app.app_context():
dash = _create_unicode_dashboard(df, table_name, slice_name, position)
yield
_cleanup(dash, slice_name)
def add_schedule_slice_and_dashboard(self):
with app.app_context():
self.common_data = dict(
active=True,
crontab="* * * * *",
recipients=self.RECIPIENTS,
deliver_as_group=True,
delivery_type=EmailDeliveryType.inline,
)
# Pick up a sample slice and dashboard
slice = db.session.query(Slice).filter_by(slice_name="Region Filter").one()
dashboard = (
db.session.query(Dashboard)
.filter_by(dashboard_title="World Bank's Data")
.one()
)
dashboard_schedule = DashboardEmailSchedule(**self.common_data)
dashboard_schedule.dashboard_id = dashboard.id
dashboard_schedule.user_id = 1
db.session.add(dashboard_schedule)
slice_schedule = SliceEmailSchedule(**self.common_data)
slice_schedule.slice_id = slice.id
slice_schedule.user_id = 1
slice_schedule.email_format = SliceEmailReportFormat.data
slice_schedule.slack_channel = "#test_channel"
db.session.add(slice_schedule)
db.session.commit()
self.slice_schedule = slice_schedule.id
self.dashboard_schedule = dashboard_schedule.id
yield
with app.app_context():
db.session.query(SliceEmailSchedule).filter_by(
id=self.slice_schedule
).delete()
db.session.query(DashboardEmailSchedule).filter_by(
id=self.dashboard_schedule
).delete()
db.session.commit()
def test_upload_to_s3_success(client):
config = app.config.copy()
app.config["CSV_TO_HIVE_UPLOAD_S3_BUCKET"] = "bucket"
client.return_value.upload_file.return_value = True
with app.app_context():
location = upload_to_s3("filename", "prefix", Table("table"))
assert f"s3a://bucket/prefix/table" == location
app.config = config
def test_cron_schedule_window(current_dttm: str, cron: str,
excepted: List[FakeDatetime]):
"""
Reports scheduler: Test cron schedule window
"""
with app.app_context():
with freeze_time(current_dttm):
datetimes = cron_schedule_window(cron)
assert list(datetimes) == excepted
def test_where_latest_partition_super_method_exception(mock_method):
mock_method.side_effect = Exception()
db = mock.Mock()
columns = [{"name": "ds"}, {"name": "hour"}]
with app.app_context():
result = HiveEngineSpec.where_latest_partition("test_table",
"test_schema", db,
select(), columns)
assert result is None
mock_method.assert_called()
def tabbed_dashboard():
position_json = {
"DASHBOARD_VERSION_KEY": "v2",
"GRID_ID": {
"children": ["TABS-IpViLohnyP"],
"id": "GRID_ID",
"parents": ["ROOT_ID"],
"type": "GRID",
},
"HEADER_ID": {
"id": "HEADER_ID",
"meta": {
"text": "tabbed dashboard"
},
"type": "HEADER",
},
"ROOT_ID": {
"children": ["GRID_ID"],
"id": "ROOT_ID",
"type": "ROOT"
},
"TAB-j53G4gtKGF": {
"children": [],
"id": "TAB-j53G4gtKGF",
"meta": {
"defaultText": "Tab title",
"placeholder": "Tab title",
"text": "Tab 1",
},
"parents": ["ROOT_ID", "GRID_ID", "TABS-IpViLohnyP"],
"type": "TAB",
},
"TAB-nerWR09Ju": {
"children": [],
"id": "TAB-nerWR09Ju",
"meta": {
"defaultText": "Tab title",
"placeholder": "Tab title",
"text": "Tab 2",
},
"parents": ["ROOT_ID", "GRID_ID", "TABS-IpViLohnyP"],
"type": "TAB",
},
"TABS-IpViLohnyP": {
"children": ["TAB-j53G4gtKGF", "TAB-nerWR09Ju"],
"id": "TABS-IpViLohnyP",
"meta": {},
"parents": ["ROOT_ID", "GRID_ID"],
"type": "TABS",
},
}
with app.app_context():
dash = create_dashboard("tabbed-dash-test", "Tabbed Dash Test",
json.dumps(position_json), [])
yield dash
def test_where_latest_partition(mock_method):
mock_method.return_value = ("01-01-19", 1)
db = mock.Mock()
db.get_indexes = mock.Mock(return_value=[{"column_names": ["ds", "hour"]}])
db.get_extra = mock.Mock(return_value={})
db.get_df = mock.Mock()
columns = [{"name": "ds"}, {"name": "hour"}]
with app.app_context():
result = HiveEngineSpec.where_latest_partition("test_table",
"test_schema", db,
select(), columns)
query_result = str(result.compile(compile_kwargs={"literal_binds": True}))
assert "SELECT \nWHERE ds = '01-01-19' AND hour = 1" == query_result
