Skip to content

Commit

Permalink
Improving the approach
Browse files Browse the repository at this point in the history
  • Loading branch information
Vitor-Avila committed Feb 12, 2025
1 parent 3a26d56 commit 5d79c04
Show file tree
Hide file tree
Showing 3 changed files with 27 additions and 22 deletions.
11 changes: 8 additions & 3 deletions superset/commands/database/resync_permissions.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,16 +30,18 @@
DatabaseNotFoundError,
UserNotFoundError,
)
from superset.daos.database import DatabaseDAO
from superset.models.core import Database
from superset.utils.core import timeout

logger = logging.getLogger(__name__)


class ResyncPermissionsCommand(BaseCommand):
def __init__(self, model: Database | None, username: str):
self._model = model
self._username = username
def __init__(self, model_id: int, username: str | None):
self._model_id = model_id
self._username: str | None = username
self._model: Database | None = None

def run(self) -> None:
self.validate()
Expand All @@ -48,12 +50,15 @@ def validate(self) -> None:
"""
Validates the command.
"""
self._model = DatabaseDAO.find_by_id(self._model_id)
if not self._model:
raise DatabaseNotFoundError()

# If OAuth2 connection, we need to impersonate
# the current user to trigger the resync
if self._model.is_oauth2_enabled():
if not self._username:
raise UserNotFoundError()
if not security_manager.get_user_by_username(self._username):
raise UserNotFoundError()

Expand Down
11 changes: 5 additions & 6 deletions superset/databases/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -659,14 +659,13 @@ def resync_permissions(self, pk: int, **kwargs: Any) -> FlaskResponse:
500:
$ref: '#/components/responses/500'
"""
database = DatabaseDAO.find_by_id(pk)
if not database:
return self.response_404()
try:
current_username = get_username() or ""
ResyncPermissionsCommand(database, current_username).run()
resync_database_permissions.delay(database, current_username)
current_username = get_username()
ResyncPermissionsCommand(pk, current_username).run()
resync_database_permissions.delay(pk, current_username)
return self.response(202, message="OK")
except DatabaseNotFoundError:
return self.response_404()
except SupersetException as ex:
return self.response(ex.status, message=ex.message)

Expand Down
27 changes: 14 additions & 13 deletions superset/tasks/permissions.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,26 +24,27 @@
from superset.commands.database.update import UpdateDatabaseCommand
from superset.daos.database import DatabaseDAO
from superset.extensions import celery_app
from superset.models.core import Database

logger = logging.getLogger(__name__)


@celery_app.task(name="resync_database_permissions", soft_time_limit=600)
def resync_database_permissions(
database: Database,
username: str,
) -> None:
logger.info("Resyncing database permissions for connection ID %s", database.id)
if user := security_manager.get_user_by_username(username):
def resync_database_permissions(database_id: int, username: str | None) -> None:
logger.info("Resyncing permissions for DB connection ID %s", database_id)
database = DatabaseDAO.find_by_id(database_id)
ssh_tunnel = DatabaseDAO.get_ssh_tunnel(database_id)
if not database:
logger.error("Database ID %s not found", database_id)
return
if username and (user := security_manager.get_user_by_username(username)):
g.user = user
logger.info("Impersonating user ID %s", g.user.id)
cmmd = UpdateDatabaseCommand(database.id, {})
cmmd = UpdateDatabaseCommand(database_id, {})
try:
cmmd._refresh_catalogs(
database, database.name, DatabaseDAO.get_ssh_tunnel(database.id)
)
except Exception as ex:
cmmd._refresh_catalogs(database, database.name, ssh_tunnel)
except Exception:
logger.error(
"An error occurred while resyncing database permissions", exc_info=ex
"An error occurred while resyncing permissions for DB connection ID %s",
database_id,
exc_info=True,
)

0 comments on commit 5d79c04

Please sign in to comment.