(superset) branch feat/async-db-perm-sync updated: Improving the approach

[vavila]

This is an automated email from the ASF dual-hosted git repository.

vavila pushed a commit to branch feat/async-db-perm-sync
in repository https://gitbox.apache.org/repos/asf/superset.git


The following commit(s) were added to refs/heads/feat/async-db-perm-sync by 
this push:
     new e06dcf3acf Improving the approach
e06dcf3acf is described below

commit e06dcf3acfce2c3254d2633cb2e51a653072738d
Author: Vitor Avila <vitor.av...@preset.io>
AuthorDate: Fri Feb 7 15:16:27 2025 -0300

    Improving the approach
---
 superset/commands/database/resync_permissions.py | 11 +++++++---
 superset/databases/api.py                        | 11 +++++-----
 superset/tasks/permissions.py                    | 27 ++++++++++++------------
 3 files changed, 27 insertions(+), 22 deletions(-)

diff --git a/superset/commands/database/resync_permissions.py 
b/superset/commands/database/resync_permissions.py
index 28d7879ec5..f72920518c 100644
--- a/superset/commands/database/resync_permissions.py
+++ b/superset/commands/database/resync_permissions.py
@@ -30,6 +30,7 @@ from superset.commands.database.exceptions import (
     DatabaseNotFoundError,
     UserNotFoundError,
 )
+from superset.daos.database import DatabaseDAO
 from superset.models.core import Database
 from superset.utils.core import timeout
 
@@ -37,9 +38,10 @@ 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()
@@ -48,12 +50,15 @@ class ResyncPermissionsCommand(BaseCommand):
         """
         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()
 
diff --git a/superset/databases/api.py b/superset/databases/api.py
index 94b85ef75e..f6c3d6c548 100644
--- a/superset/databases/api.py
+++ b/superset/databases/api.py
@@ -659,14 +659,13 @@ class DatabaseRestApi(BaseSupersetModelRestApi):
             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)
 
diff --git a/superset/tasks/permissions.py b/superset/tasks/permissions.py
index 5b0f6ada63..902cd11a59 100644
--- a/superset/tasks/permissions.py
+++ b/superset/tasks/permissions.py
@@ -24,26 +24,27 @@ from superset import security_manager
 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,
         )