From 570d7296f64f7b9e915b18ea74cf8b2bcc81920b Mon Sep 17 00:00:00 2001 From: Beto Dealmeida Date: Mon, 6 May 2024 10:37:20 -0400 Subject: [PATCH] Use migration to update schema perms --- superset/commands/database/update.py | 28 +---- superset/databases/api.py | 2 +- superset/db_engine_specs/base.py | 4 + superset/migrations/shared/catalogs.py | 116 ++++++++++++++++++ ...58d051681a3b_add_catalog_perm_to_tables.py | 11 +- superset/security/manager.py | 4 - .../commands/databases/update_test.py | 63 ---------- 7 files changed, 135 insertions(+), 93 deletions(-) create mode 100644 superset/migrations/shared/catalogs.py diff --git a/superset/commands/database/update.py b/superset/commands/database/update.py index c2f169b3e2e57..c59984238cfe4 100644 --- a/superset/commands/database/update.py +++ b/superset/commands/database/update.py @@ -247,30 +247,12 @@ def _refresh_schemas( "schema_access", perm, ) - new_name = security_manager.get_schema_perm( - database.database_name, - catalog, - schema, - ) if not existing_pvm: - # If the PVM doesn't exist it could be for 2 reasons: either the schema - # is new, or the database was updated to support catalogs and the schema - # exists with the old permission format of `[db].[schema]` and needs to - # be updated. - if catalog and catalog == database.get_default_catalog(): - perm = security_manager.get_schema_perm( - original_database_name, - None, - schema, - ) - existing_pvm = security_manager.find_permission_view_menu( - "schema_access", - perm, - ) - if existing_pvm: - existing_pvm.view_menu.name = new_name - continue - + new_name = security_manager.get_schema_perm( + database.database_name, + catalog, + schema, + ) security_manager.add_permission_view_menu("schema_access", new_name) def _rename_database_in_permissions( diff --git a/superset/databases/api.py b/superset/databases/api.py index c124afced6d31..116ac9f46b49a 100644 --- a/superset/databases/api.py +++ b/superset/databases/api.py @@ -651,7 +651,7 @@ def catalogs(self, pk: int, **kwargs: Any) -> FlaskResponse: 500: $ref: '#/components/responses/500' """ - database = self.datamodel.get(pk, self._base_filters) + database = DatabaseDAO.find_by_id(pk) if not database: return self.response_404() try: diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index 0a4955a19fc4d..4ee4e4dc0385b 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -404,6 +404,10 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods # Does the DB support catalogs? A catalog here is a group of schemas, and has # different names depending on the DB: BigQuery calles it a "project", Postgres calls # it a "database", Trino calls it a "catalog", etc. + # + # When this is changed to true in a DB engine spec it MUST support the + # `get_default_catalog` and `get_catalog_names` methods. In addition, you MUST write + # a database migration updating any existing schema permissions. supports_catalog = False # Can the catalog be changed on a per-query basis? diff --git a/superset/migrations/shared/catalogs.py b/superset/migrations/shared/catalogs.py new file mode 100644 index 0000000000000..5d01ecfbfbcc5 --- /dev/null +++ b/superset/migrations/shared/catalogs.py @@ -0,0 +1,116 @@ +# 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 __future__ import annotations + +import logging + +from alembic import op + +from superset import db, security_manager +from superset.daos.database import DatabaseDAO +from superset.models.core import Database + +logger = logging.getLogger(__name__) + + +def upgrade_schema_perms(engine: str | None = None) -> None: + """ + Update schema permissions to include the catalog part. + + Before SIP-95 schema permissions were stored in the format `[db].[schema]`. With the + introduction of catalogs, any existing permissions need to be renamed to include the + catalog: `[db].[catalog].[schema]`. + """ + bind = op.get_bind() + session = db.Session(bind=bind) + for database in session.query(Database).all(): + db_engine_spec = database.db_engine_spec + if ( + engine and db_engine_spec.engine != engine + ) or not db_engine_spec.supports_catalog: + continue + + catalog = database.get_default_catalog() + ssh_tunnel = DatabaseDAO.get_ssh_tunnel(database.id) + for schema in database.get_all_schema_names( + catalog=catalog, + cache=False, + ssh_tunnel=ssh_tunnel, + ): + perm = security_manager.get_schema_perm( + database.database_name, + None, + schema, + ) + existing_pvm = security_manager.find_permission_view_menu( + "schema_access", + perm, + ) + if existing_pvm: + existing_pvm.view_menu.name = security_manager.get_schema_perm( + database.database_name, + catalog, + schema, + ) + + session.commit() + + +def downgrade_schema_perms(engine: str | None = None) -> None: + """ + Update schema permissions to not have the catalog part. + + Before SIP-95 schema permissions were stored in the format `[db].[schema]`. With the + introduction of catalogs, any existing permissions need to be renamed to include the + catalog: `[db].[catalog].[schema]`. + + This helped function reverts the process. + """ + bind = op.get_bind() + session = db.Session(bind=bind) + for database in session.query(Database).all(): + db_engine_spec = database.db_engine_spec + if ( + engine and db_engine_spec.engine != engine + ) or not db_engine_spec.supports_catalog: + continue + + catalog = database.get_default_catalog() + ssh_tunnel = DatabaseDAO.get_ssh_tunnel(database.id) + for schema in database.get_all_schema_names( + catalog=catalog, + cache=False, + ssh_tunnel=ssh_tunnel, + ): + perm = security_manager.get_schema_perm( + database.database_name, + catalog, + schema, + ) + existing_pvm = security_manager.find_permission_view_menu( + "schema_access", + perm, + ) + if existing_pvm: + existing_pvm.view_menu.name = security_manager.get_schema_perm( + database.database_name, + None, + schema, + ) + + session.commit() diff --git a/superset/migrations/versions/2024-05-01_10-52_58d051681a3b_add_catalog_perm_to_tables.py b/superset/migrations/versions/2024-05-01_10-52_58d051681a3b_add_catalog_perm_to_tables.py index d00bff0ccbe4e..17b33e1d0a8f8 100644 --- a/superset/migrations/versions/2024-05-01_10-52_58d051681a3b_add_catalog_perm_to_tables.py +++ b/superset/migrations/versions/2024-05-01_10-52_58d051681a3b_add_catalog_perm_to_tables.py @@ -17,7 +17,7 @@ """Add catalog_perm to tables Revision ID: 58d051681a3b -Revises: 5f57af97bc3f +Revises: 4a33124c18ad Create Date: 2024-05-01 10:52:31.458433 """ @@ -25,9 +25,14 @@ import sqlalchemy as sa from alembic import op +from superset.migrations.shared.catalogs import ( + downgrade_schema_perms, + upgrade_schema_perms, +) + # revision identifiers, used by Alembic. revision = "58d051681a3b" -down_revision = "5f57af97bc3f" +down_revision = "4a33124c18ad" def upgrade(): @@ -39,8 +44,10 @@ def upgrade(): "slices", sa.Column("catalog_perm", sa.String(length=1000), nullable=True), ) + upgrade_schema_perms(engine="postgresql") def downgrade(): op.drop_column("slices", "catalog_perm") op.drop_column("tables", "catalog_perm") + downgrade_schema_perms(engine="postgresql") diff --git a/superset/security/manager.py b/superset/security/manager.py index 7cfd4b9a521c7..d28ed17893cd6 100644 --- a/superset/security/manager.py +++ b/superset/security/manager.py @@ -379,10 +379,6 @@ def get_schema_perm( [database].[schema] [database].[catalog].[schema] - For backwards compatibility, when processing the first format Superset should - use the default catalog when the database supports them. This way, migrating - existing permissions is not necessary. - :param database: The database name :param catalog: The database catalog name :param schema: The database schema name diff --git a/tests/unit_tests/commands/databases/update_test.py b/tests/unit_tests/commands/databases/update_test.py index fa1d58107379f..300efb62e7d3c 100644 --- a/tests/unit_tests/commands/databases/update_test.py +++ b/tests/unit_tests/commands/databases/update_test.py @@ -196,8 +196,6 @@ def test_rename_with_catalog( "[my_db].[catalog2]", # second catalog already exists "[my_db].[catalog2].[schema3]", # first schema already exists None, # second schema is new - # these are called when checking for existing perms in [db].[schema] format - None, # these are called when renaming the permissions: catalog2_pvm, # old [my_db].[catalog2] catalog2_schema3_pvm, # old [my_db].[catalog2].[schema3] @@ -272,64 +270,3 @@ def test_rename_without_catalog( ) assert schema2_pvm.view_menu.name == "[my_other_db].[schema2]" - - -def test_update_old_format_with_catalog( - mocker: MockerFixture, - database_with_catalog: MockerFixture, -) -> None: - """ - Test existing permissions in old format are updated correctly. - - Before catalogs were introduced, the format for schema permissions was - `[db].[schema]`. With catalogs, these needs to be updated to - `[db].[catalog].[schema]`, where `catalog` is the default catalog. - - In this test, the database has two catalogs with two schemas each: - - - catalog1 - - schema1 - - schema2 - - catalog2 - - schema3 - - schema4 - - When update is called, only `catalog2.schema3` has permissions associated with it, - so `catalog1.*` and `catalog2.schema4` are added. Furthermore, the `schema3` has - permissions in the old format of `[db].[schema]`, so it is updated to the new format. - """ - DatabaseDAO = mocker.patch("superset.commands.database.update.DatabaseDAO") - DatabaseDAO.find_by_id.return_value = database_with_catalog - DatabaseDAO.update.return_value = database_with_catalog - - find_permission_view_menu = mocker.patch.object( - security_manager, - "find_permission_view_menu", - ) - schema3_pvm = mocker.MagicMock() - find_permission_view_menu.side_effect = [ - None, # first catalog is new - "[my_db].[catalog2]", # second catalog already exists - None, # first schema has old format - schema3_pvm, # old format for first schema exists - None, # second schema is new - None, # second schema has not old format - ] - add_permission_view_menu = mocker.patch.object( - security_manager, - "add_permission_view_menu", - ) - - UpdateDatabaseCommand(1, {}).run() - - add_permission_view_menu.assert_has_calls( - [ - # first catalog is added with all schemas - mocker.call("catalog_access", "[my_db].[catalog1]"), - mocker.call("schema_access", "[my_db].[catalog1].[schema1]"), - mocker.call("schema_access", "[my_db].[catalog1].[schema2]"), - # second catalog already exists, only `schema4` is added - mocker.call("schema_access", "[my_db].[catalog2].[schema4]"), - ], - ) - assert schema3_pvm.view_menu.name == "[my_db].[catalog2].[schema3]"