332 lines
10 KiB
Python
332 lines
10 KiB
Python
# 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.
|
|
|
|
from pytest_mock import MockerFixture
|
|
from sqlalchemy.orm.session import Session
|
|
|
|
from superset.migrations.shared.catalogs import (
|
|
downgrade_catalog_perms,
|
|
upgrade_catalog_perms,
|
|
)
|
|
from superset.migrations.shared.security_converge import (
|
|
Permission,
|
|
PermissionView,
|
|
ViewMenu,
|
|
)
|
|
|
|
|
|
def test_upgrade_catalog_perms(mocker: MockerFixture, session: Session) -> None:
|
|
"""
|
|
Test the `upgrade_catalog_perms` function.
|
|
|
|
The function is called when catalogs are introduced into a new DB engine spec.
|
|
"""
|
|
from superset.connectors.sqla.models import SqlaTable
|
|
from superset.models.core import Database
|
|
from superset.models.slice import Slice
|
|
from superset.models.sql_lab import Query, SavedQuery, TableSchema, TabState
|
|
|
|
engine = session.get_bind()
|
|
Database.metadata.create_all(engine)
|
|
|
|
mocker.patch("superset.migrations.shared.catalogs.op")
|
|
db = mocker.patch("superset.migrations.shared.catalogs.db")
|
|
db.Session.return_value = session
|
|
|
|
mocker.patch.object(
|
|
Database,
|
|
"get_all_schema_names",
|
|
return_value=["public", "information_schema"],
|
|
)
|
|
mocker.patch.object(
|
|
Database,
|
|
"get_all_catalog_names",
|
|
return_value=["db", "other_catalog"],
|
|
)
|
|
|
|
database = Database(
|
|
database_name="my_db",
|
|
sqlalchemy_uri="postgresql://localhost/db",
|
|
)
|
|
dataset = SqlaTable(
|
|
table_name="my_table",
|
|
database=database,
|
|
catalog=None,
|
|
schema="public",
|
|
catalog_perm=None,
|
|
schema_perm="[my_db].[public]",
|
|
)
|
|
session.add(dataset)
|
|
session.commit()
|
|
|
|
chart = Slice(
|
|
slice_name="my_chart",
|
|
datasource_type="table",
|
|
datasource_id=dataset.id,
|
|
catalog_perm=None,
|
|
schema_perm="[my_db].[public]",
|
|
)
|
|
query = Query(
|
|
client_id="foo",
|
|
database=database,
|
|
catalog=None,
|
|
schema="public",
|
|
)
|
|
saved_query = SavedQuery(
|
|
database=database,
|
|
sql="SELECT * FROM public.t",
|
|
catalog=None,
|
|
schema="public",
|
|
)
|
|
tab_state = TabState(
|
|
database=database,
|
|
catalog=None,
|
|
schema="public",
|
|
)
|
|
table_schema = TableSchema(
|
|
database=database,
|
|
catalog=None,
|
|
schema="public",
|
|
)
|
|
session.add_all([chart, query, saved_query, tab_state, table_schema])
|
|
session.commit()
|
|
|
|
# before migration
|
|
assert dataset.catalog is None
|
|
assert query.catalog is None
|
|
assert saved_query.catalog is None
|
|
assert tab_state.catalog is None
|
|
assert table_schema.catalog is None
|
|
assert dataset.catalog_perm is None
|
|
assert dataset.schema_perm == "[my_db].[public]"
|
|
assert chart.catalog_perm is None
|
|
assert chart.schema_perm == "[my_db].[public]"
|
|
assert (
|
|
session.query(ViewMenu.name, Permission.name)
|
|
.join(PermissionView, ViewMenu.id == PermissionView.view_menu_id)
|
|
.join(Permission, PermissionView.permission_id == Permission.id)
|
|
.all()
|
|
) == [
|
|
("[my_db].(id:1)", "database_access"),
|
|
("[my_db].[my_table](id:1)", "datasource_access"),
|
|
("[my_db].[public]", "schema_access"),
|
|
]
|
|
|
|
upgrade_catalog_perms()
|
|
session.commit()
|
|
|
|
# add dataset/chart in new catalog
|
|
new_dataset = SqlaTable(
|
|
table_name="my_table",
|
|
database=database,
|
|
catalog="other_catalog",
|
|
schema="public",
|
|
schema_perm="[my_db].[other_catalog].[public]",
|
|
catalog_perm="[my_db].[other_catalog]",
|
|
)
|
|
session.add(new_dataset)
|
|
session.commit()
|
|
|
|
new_chart = Slice(
|
|
slice_name="my_chart",
|
|
datasource_type="table",
|
|
datasource_id=new_dataset.id,
|
|
)
|
|
session.add(new_chart)
|
|
session.commit()
|
|
|
|
# after migration
|
|
assert dataset.catalog == "db"
|
|
assert query.catalog == "db"
|
|
assert saved_query.catalog == "db"
|
|
assert tab_state.catalog == "db"
|
|
assert table_schema.catalog == "db"
|
|
assert dataset.catalog_perm == "[my_db].[db]"
|
|
assert dataset.schema_perm == "[my_db].[db].[public]"
|
|
assert chart.catalog_perm == "[my_db].[db]"
|
|
assert chart.schema_perm == "[my_db].[db].[public]"
|
|
assert (
|
|
session.query(ViewMenu.name, Permission.name)
|
|
.join(PermissionView, ViewMenu.id == PermissionView.view_menu_id)
|
|
.join(Permission, PermissionView.permission_id == Permission.id)
|
|
.all()
|
|
) == [
|
|
("[my_db].(id:1)", "database_access"),
|
|
("[my_db].[my_table](id:1)", "datasource_access"),
|
|
("[my_db].[db].[public]", "schema_access"),
|
|
("[my_db].[db]", "catalog_access"),
|
|
("[my_db].[other_catalog]", "catalog_access"),
|
|
("[my_db].[other_catalog].[public]", "schema_access"),
|
|
("[my_db].[other_catalog].[information_schema]", "schema_access"),
|
|
("[my_db].[my_table](id:2)", "datasource_access"),
|
|
]
|
|
|
|
# do a downgrade
|
|
downgrade_catalog_perms()
|
|
session.commit()
|
|
|
|
# revert
|
|
assert dataset.catalog is None
|
|
assert query.catalog is None
|
|
assert saved_query.catalog is None
|
|
assert tab_state.catalog is None
|
|
assert table_schema.catalog is None
|
|
assert dataset.catalog_perm is None
|
|
assert dataset.schema_perm == "[my_db].[public]"
|
|
assert chart.catalog_perm is None
|
|
assert chart.schema_perm == "[my_db].[public]"
|
|
assert (
|
|
session.query(ViewMenu.name, Permission.name)
|
|
.join(PermissionView, ViewMenu.id == PermissionView.view_menu_id)
|
|
.join(Permission, PermissionView.permission_id == Permission.id)
|
|
.all()
|
|
) == [
|
|
("[my_db].(id:1)", "database_access"),
|
|
("[my_db].[my_table](id:1)", "datasource_access"),
|
|
("[my_db].[public]", "schema_access"),
|
|
]
|
|
|
|
# make sure new dataset/chart were deleted
|
|
assert session.query(SqlaTable).all() == [dataset]
|
|
assert session.query(Slice).all() == [chart]
|
|
|
|
|
|
def test_upgrade_catalog_perms_graceful(
|
|
mocker: MockerFixture,
|
|
session: Session,
|
|
) -> None:
|
|
"""
|
|
Test the `upgrade_catalog_perms` function when it fails to connect to the DB.
|
|
|
|
During the migration we try to connect to the analytical database to get the list of
|
|
schemas. This should fail gracefully and not raise an exception, since the database
|
|
could be offline, and the permissions can be generated later then the admin enables
|
|
catalog browsing on the database (permissions are always synced on a DB update, see
|
|
`UpdateDatabaseCommand`).
|
|
"""
|
|
from superset.connectors.sqla.models import SqlaTable
|
|
from superset.models.core import Database
|
|
from superset.models.slice import Slice
|
|
from superset.models.sql_lab import Query, SavedQuery, TableSchema, TabState
|
|
|
|
engine = session.get_bind()
|
|
Database.metadata.create_all(engine)
|
|
|
|
mocker.patch("superset.migrations.shared.catalogs.op")
|
|
db = mocker.patch("superset.migrations.shared.catalogs.db")
|
|
db.Session.return_value = session
|
|
|
|
mocker.patch.object(
|
|
Database,
|
|
"get_all_schema_names",
|
|
side_effect=Exception("Failed to connect to the database"),
|
|
)
|
|
mocker.patch("superset.migrations.shared.catalogs.op", session)
|
|
|
|
database = Database(
|
|
database_name="my_db",
|
|
sqlalchemy_uri="postgresql://localhost/db",
|
|
)
|
|
dataset = SqlaTable(
|
|
table_name="my_table",
|
|
database=database,
|
|
catalog=None,
|
|
schema="public",
|
|
schema_perm="[my_db].[public]",
|
|
)
|
|
session.add(dataset)
|
|
session.commit()
|
|
|
|
chart = Slice(
|
|
slice_name="my_chart",
|
|
datasource_type="table",
|
|
datasource_id=dataset.id,
|
|
)
|
|
query = Query(
|
|
client_id="foo",
|
|
database=database,
|
|
catalog=None,
|
|
schema="public",
|
|
)
|
|
saved_query = SavedQuery(
|
|
database=database,
|
|
sql="SELECT * FROM public.t",
|
|
catalog=None,
|
|
schema="public",
|
|
)
|
|
tab_state = TabState(
|
|
database=database,
|
|
catalog=None,
|
|
schema="public",
|
|
)
|
|
table_schema = TableSchema(
|
|
database=database,
|
|
catalog=None,
|
|
schema="public",
|
|
)
|
|
session.add_all([chart, query, saved_query, tab_state, table_schema])
|
|
session.commit()
|
|
|
|
# before migration
|
|
assert dataset.catalog is None
|
|
assert query.catalog is None
|
|
assert saved_query.catalog is None
|
|
assert tab_state.catalog is None
|
|
assert table_schema.catalog is None
|
|
assert dataset.schema_perm == "[my_db].[public]"
|
|
assert chart.schema_perm == "[my_db].[public]"
|
|
assert session.query(ViewMenu.name).all() == [
|
|
("[my_db].(id:1)",),
|
|
("[my_db].[my_table](id:1)",),
|
|
("[my_db].[public]",),
|
|
]
|
|
|
|
upgrade_catalog_perms()
|
|
session.commit()
|
|
|
|
# after migration
|
|
assert dataset.catalog == "db"
|
|
assert query.catalog == "db"
|
|
assert saved_query.catalog == "db"
|
|
assert tab_state.catalog == "db"
|
|
assert table_schema.catalog == "db"
|
|
assert dataset.schema_perm == "[my_db].[db].[public]"
|
|
assert chart.schema_perm == "[my_db].[db].[public]"
|
|
assert session.query(ViewMenu.name).all() == [
|
|
("[my_db].(id:1)",),
|
|
("[my_db].[my_table](id:1)",),
|
|
("[my_db].[db].[public]",),
|
|
("[my_db].[db]",),
|
|
]
|
|
|
|
downgrade_catalog_perms()
|
|
session.commit()
|
|
|
|
# revert
|
|
assert dataset.catalog is None
|
|
assert query.catalog is None
|
|
assert saved_query.catalog is None
|
|
assert tab_state.catalog is None
|
|
assert table_schema.catalog is None
|
|
assert dataset.schema_perm == "[my_db].[public]"
|
|
assert chart.schema_perm == "[my_db].[public]"
|
|
assert session.query(ViewMenu.name).all() == [
|
|
("[my_db].(id:1)",),
|
|
("[my_db].[my_table](id:1)",),
|
|
("[my_db].[public]",),
|
|
]
|