Skip to content

Commit

Permalink
Moving FAB overrides to utils
Browse files Browse the repository at this point in the history
  • Loading branch information
@Vitor-Avila
Vitor-Avila committed Feb 28, 2025
1 parent 49b3d10 commit f394d86
Show file tree
Hide file tree
Showing 4 changed files with 126 additions and 81 deletions.
81 changes: 9 additions & 72 deletions superset/commands/database/sync_permissions.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,11 +21,6 @@
from typing import Iterable

from flask import current_app, g
from flask_appbuilder.security.sqla.models import (
Permission,
PermissionView,
ViewMenu,
)

from superset import app, security_manager
from superset.commands.base import BaseCommand
Expand All @@ -35,15 +30,18 @@
DatabaseNotFoundError,
UserNotFoundInSessionError,
)
from superset.commands.database.utils import ping
from superset.commands.database.utils import (
add_pvm,
add_vm,
ping,
)
from superset.daos.database import DatabaseDAO
from superset.daos.dataset import DatasetDAO
from superset.databases.ssh_tunnel.models import SSHTunnel
from superset.db_engine_specs.base import GenericDBException
from superset.exceptions import OAuth2RedirectError
from superset.extensions import celery_app, db
from superset.models.core import Database
from superset.security.manager import SupersetSecurityManager
from superset.utils.decorators import on_error, transaction

logger = logging.getLogger(__name__)
Expand Down Expand Up @@ -164,6 +162,7 @@ def sync_database_permissions(self) -> None:
if not existing_pvm:
# new catalog
add_pvm(
db.session,
security_manager,
"catalog_access",
security_manager.get_catalog_perm(
Expand All @@ -173,6 +172,7 @@ def sync_database_permissions(self) -> None:
)
for schema in schemas:
add_pvm(
db.session,
security_manager,
"schema_access",
security_manager.get_schema_perm(
Expand Down Expand Up @@ -243,7 +243,7 @@ def _refresh_schemas(self, catalog: str | None, schemas: Iterable[str]) -> None:
catalog,
schema,
)
add_pvm(security_manager, "schema_access", new_name)
add_pvm(db.session, security_manager, "schema_access", new_name)

def _rename_database_in_permissions(
self, catalog: str | None, schemas: Iterable[str]
Expand All @@ -254,7 +254,7 @@ def _rename_database_in_permissions(
self.db_connection.name,
catalog,
)
new_catalog_vm = add_vm(security_manager, new_catalog_perm_name)
new_catalog_vm = add_vm(db.session, security_manager, new_catalog_perm_name)
perm = security_manager.get_catalog_perm(
self.old_db_connection_name,
catalog,
Expand Down Expand Up @@ -299,69 +299,6 @@ def _rename_database_in_permissions(
chart.schema_perm = new_schema_perm_name


def add_vm(
security_manager: SupersetSecurityManager,
view_menu_name: str | None,
) -> ViewMenu:
"""
Similar to security_manager.add_view_menu, but without commit.
This ensures an atomic operation.
"""
if view_menu := security_manager.find_view_menu(view_menu_name):
return view_menu

view_menu = security_manager.viewmenu_model()
view_menu.name = view_menu_name
db.session.add(view_menu)
return view_menu


def add_perm(
security_manager: SupersetSecurityManager,
permission_name: str | None,
) -> Permission:
"""
Similar to security_manager.add_permission, but without commit.
This ensures an atomic operation.
"""
if perm := security_manager.find_permission(permission_name):
return perm

perm = security_manager.permission_model()
perm.name = permission_name
db.session.add(perm)
return perm


def add_pvm(
security_manager: SupersetSecurityManager,
permission_name: str | None,
view_menu_name: str | None,
) -> PermissionView | None:
"""
Similar to security_manager.add_permission_view_menu, but without commit.
This ensures an atomic operation.
"""
if not (permission_name and view_menu_name):
return None

if pv := security_manager.find_permission_view_menu(
permission_name, view_menu_name
):
return pv

vm = add_vm(security_manager, view_menu_name)
perm = add_perm(security_manager, permission_name)
pv = security_manager.permissionview_model()
pv.view_menu, pv.permission = vm, perm
db.session.add(pv)

return pv


@celery_app.task(name="sync_database_permissions", soft_time_limit=600)
def sync_database_permissions_task(
database_id: int, username: str, old_db_connection_name: str
Expand Down
74 changes: 74 additions & 0 deletions superset/commands/database/utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,12 +21,19 @@
from contextlib import closing

from flask import current_app as app
from flask_appbuilder.security.sqla.models import (
Permission,
PermissionView,
ViewMenu,
)
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session

from superset import security_manager
from superset.databases.ssh_tunnel.models import SSHTunnel
from superset.db_engine_specs.base import GenericDBException
from superset.models.core import Database
from superset.security.manager import SupersetSecurityManager
from superset.utils.core import timeout

logger = logging.getLogger(__name__)
Expand All @@ -48,6 +55,7 @@ def add_permissions(database: Database, ssh_tunnel: SSHTunnel | None) -> None:
"""
Add DAR for catalogs and schemas.
"""
# TODO: Migrate this to use the non-commiting add_pvm helper instead
if database.db_engine_spec.supports_catalog:
catalogs = database.get_all_catalog_names(
cache=False,
Expand Down Expand Up @@ -83,3 +91,69 @@ def add_permissions(database: Database, ssh_tunnel: SSHTunnel | None) -> None:
except GenericDBException: # pylint: disable=broad-except
logger.warning("Error processing catalog '%s'", catalog)
continue


def add_vm(
session: Session,
security_manager: SupersetSecurityManager,
view_menu_name: str | None,
) -> ViewMenu:
"""
Similar to security_manager.add_view_menu, but without commit.
This ensures an atomic operation.
"""
if view_menu := security_manager.find_view_menu(view_menu_name):
return view_menu

view_menu = security_manager.viewmenu_model()
view_menu.name = view_menu_name
session.add(view_menu)
return view_menu


def add_perm(
session: Session,
security_manager: SupersetSecurityManager,
permission_name: str | None,
) -> Permission:
"""
Similar to security_manager.add_permission, but without commit.
This ensures an atomic operation.
"""
if perm := security_manager.find_permission(permission_name):
return perm

perm = security_manager.permission_model()
perm.name = permission_name
session.add(perm)
return perm


def add_pvm(
session: Session,
security_manager: SupersetSecurityManager,
permission_name: str | None,
view_menu_name: str | None,
) -> PermissionView | None:
"""
Similar to security_manager.add_permission_view_menu, but without commit.
This ensures an atomic operation.
"""
if not (permission_name and view_menu_name):
return None

if pv := security_manager.find_permission_view_menu(
permission_name, view_menu_name
):
return pv

vm = add_vm(session, security_manager, view_menu_name)
perm = add_perm(session, security_manager, permission_name)
pv = security_manager.permissionview_model()
pv.view_menu, pv.permission = vm, perm
session.add(pv)

return pv
16 changes: 13 additions & 3 deletions tests/unit_tests/commands/databases/sync_permissions_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import pytest
from pytest_mock import MockerFixture

from superset import db
from superset.commands.database.exceptions import (
DatabaseConnectionFailedError,
DatabaseNotFoundError,
Expand Down Expand Up @@ -66,12 +67,20 @@ def test_sync_permissions_command_sync_mode(
user_mock.assert_called_once_with("admin")
add_pvm_mock.assert_has_calls(
[
mocker.call(security_manager, "catalog_access", "[my_db].[catalog2]"),
mocker.call(
security_manager, "schema_access", "[my_db].[catalog2].[schema3]"
db.session, security_manager, "catalog_access", "[my_db].[catalog2]"
),
mocker.call(
security_manager, "schema_access", "[my_db].[catalog2].[schema4]"
db.session,
security_manager,
"schema_access",
"[my_db].[catalog2].[schema3]",
),
mocker.call(
db.session,
security_manager,
"schema_access",
"[my_db].[catalog2].[schema4]",
),
]
)
Expand Down Expand Up @@ -303,6 +312,7 @@ def test_resync_permissions_command_refresh_schemas(
cmmd._refresh_schemas("catalog1", ["schema1", "schema2"])

add_pvm_mock.assert_called_once_with(
db.session,
security_manager,
"schema_access",
f"[{database_with_catalog.name}].[catalog1].[schema2]",
Expand Down
36 changes: 30 additions & 6 deletions tests/unit_tests/commands/databases/update_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

from pytest_mock import MockerFixture

from superset import db
from superset.commands.database.update import UpdateDatabaseCommand
from superset.extensions import security_manager
from superset.utils import json
Expand Down Expand Up @@ -75,15 +76,24 @@ def test_update_with_catalog(
add_pvm.assert_has_calls(
[
# first catalog is added with all schemas
mocker.call(security_manager, "catalog_access", "[my_db].[catalog1]"),
mocker.call(
security_manager, "schema_access", "[my_db].[catalog1].[schema1]"
db.session, security_manager, "catalog_access", "[my_db].[catalog1]"
),
mocker.call(
security_manager, "schema_access", "[my_db].[catalog1].[schema2]"
db.session,
security_manager,
"schema_access",
"[my_db].[catalog1].[schema1]",
),
mocker.call(
db.session,
security_manager,
"schema_access",
"[my_db].[catalog1].[schema2]",
),
# second catalog already exists, only `schema4` is added
mocker.call(
db.session,
security_manager,
"schema_access",
f"[{database_with_catalog.name}].[catalog2].[schema4]",
Expand Down Expand Up @@ -157,6 +167,7 @@ def test_update_without_catalog(
UpdateDatabaseCommand(1, {}).run()

add_pvm.assert_called_with(
db.session,
security_manager,
"schema_access",
f"[{database_without_catalog.name}].[schema1]",
Expand Down Expand Up @@ -229,15 +240,27 @@ def test_rename_with_catalog(
add_pvm.assert_has_calls(
[
# first catalog is added with all schemas with the new DB name
mocker.call(security_manager, "catalog_access", "[my_other_db].[catalog1]"),
mocker.call(
security_manager, "schema_access", "[my_other_db].[catalog1].[schema1]"
db.session,
security_manager,
"catalog_access",
"[my_other_db].[catalog1]",
),
mocker.call(
security_manager, "schema_access", "[my_other_db].[catalog1].[schema2]"
db.session,
security_manager,
"schema_access",
"[my_other_db].[catalog1].[schema1]",
),
mocker.call(
db.session,
security_manager,
"schema_access",
"[my_other_db].[catalog1].[schema2]",
),
# second catalog already exists, only `schema4` is added
mocker.call(
db.session,
security_manager,
"schema_access",
f"[{database_with_catalog.name}].[catalog2].[schema4]",
Expand Down Expand Up @@ -303,6 +326,7 @@ def test_rename_without_catalog(
UpdateDatabaseCommand(1, {}).run()

add_pvm.assert_called_with(
db.session,
security_manager,
"schema_access",
f"[{database_without_catalog.name}].[schema1]",
Expand Down

0 comments on commit f394d86

Please # to comment.