chore(dao/command): Add transaction decorator to try to enforce "unit…

… of work" (#24969)

pyproject.toml

@@ -240,6 +240,7 @@ ignore_basepython_conflict = true
 commands =
     superset db upgrade
     superset init
+    superset load-test-users
     # use -s to be able to use break pointers.
     # no args or tests/* can be passed as an argument to run all tests
     pytest -s {posargs}

scripts/permissions_cleanup.py

@@ -17,8 +17,10 @@
 from collections import defaultdict
 
 from superset import security_manager
+from superset.utils.decorators import transaction
 
 
+@transaction()
 def cleanup_permissions() -> None:
     # 1. Clean up duplicates.
     pvms = security_manager.get_session.query(
@@ -29,7 +31,6 @@ def cleanup_permissions() -> None:
     for pvm in pvms:
         pvms_dict[(pvm.permission, pvm.view_menu)].append(pvm)
     duplicates = [v for v in pvms_dict.values() if len(v) > 1]
-    len(duplicates)
 
     for pvm_list in duplicates:
         first_prm = pvm_list[0]
@@ -38,7 +39,6 @@ def cleanup_permissions() -> None:
             roles = roles.union(pvm.role)
             security_manager.get_session.delete(pvm)
         first_prm.roles = list(roles)
-    security_manager.get_session.commit()
 
     pvms = security_manager.get_session.query(
         security_manager.permissionview_model
@@ -52,7 +52,6 @@ def cleanup_permissions() -> None:
     for pvm in pvms:
         if not (pvm.view_menu and pvm.permission):
             security_manager.get_session.delete(pvm)
-    security_manager.get_session.commit()
 
     pvms = security_manager.get_session.query(
         security_manager.permissionview_model
@@ -63,15 +62,13 @@ def cleanup_permissions() -> None:
     roles = security_manager.get_session.query(security_manager.role_model).all()
     for role in roles:
         role.permissions = [p for p in role.permissions if p]
-    security_manager.get_session.commit()
 
     # 4. Delete empty roles from permission view menus
     pvms = security_manager.get_session.query(
         security_manager.permissionview_model
     ).all()
     for pvm in pvms:
         pvm.role = [r for r in pvm.role if r]
-    security_manager.get_session.commit()
 
 
 cleanup_permissions()

scripts/python_tests.sh

@@ -29,6 +29,7 @@ echo "Superset config module: $SUPERSET_CONFIG"
 
 superset db upgrade
 superset init
+superset load-test-users
 
 echo "Running tests"
 

superset/cachekeys/api.py

@@ -113,8 +113,10 @@ def invalidate(self) -> Response:
                 delete_stmt = CacheKey.__table__.delete().where(  # pylint: disable=no-member
                     CacheKey.cache_key.in_(cache_keys)
                 )
-                db.session.execute(delete_stmt)
-                db.session.commit()
+
+                with db.session.begin_nested():
+                    db.session.execute(delete_stmt)
+
                 stats_logger_manager.instance.gauge(
                     "invalidated_cache", len(cache_keys)
                 )
@@ -125,7 +127,5 @@ def invalidate(self) -> Response:
                 )
             except SQLAlchemyError as ex:  # pragma: no cover
                 logger.error(ex, exc_info=True)
-                db.session.rollback()
                 return self.response_500(str(ex))
-            db.session.commit()
         return self.response(201)

superset/cli/examples.py

@@ -20,6 +20,7 @@
 from flask.cli import with_appcontext
 
 import superset.utils.database as database_utils
+from superset.utils.decorators import transaction
 
 logger = logging.getLogger(__name__)
 
@@ -89,6 +90,7 @@ def load_examples_run(
 
 @click.command()
 @with_appcontext
+@transaction()
 @click.option("--load-test-data", "-t", is_flag=True, help="Load additional test data")
 @click.option("--load-big-data", "-b", is_flag=True, help="Load additional big data")
 @click.option(

superset/cli/main.py

@@ -27,6 +27,7 @@
 from superset import app, appbuilder, cli, security_manager
 from superset.cli.lib import normalize_token
 from superset.extensions import db
+from superset.utils.decorators import transaction
 
 logger = logging.getLogger(__name__)
 
@@ -60,6 +61,7 @@ def make_shell_context() -> dict[str, Any]:
 
 @superset.command()
 @with_appcontext
+@transaction()
 def init() -> None:
     """Inits the Superset application"""
     appbuilder.add_permissions(update_perms=True)

superset/cli/test.py

@@ -22,28 +22,22 @@
 
 import superset.utils.database as database_utils
 from superset import app, security_manager
+from superset.utils.decorators import transaction
 
 logger = logging.getLogger(__name__)
 
 
 @click.command()
 @with_appcontext
+@transaction()
 def load_test_users() -> None:
     """
     Loads admin, alpha, and gamma user for testing purposes
 
     Syncs permissions for those users/roles
     """
     print(Fore.GREEN + "Loading a set of users for unit tests")
-    load_test_users_run()
 
-
-def load_test_users_run() -> None:
-    """
-    Loads admin, alpha, and gamma user for testing purposes
-
-    Syncs permissions for those users/roles
-    """
     if app.config["TESTING"]:
         sm = security_manager
 
@@ -84,4 +78,3 @@ def load_test_users_run() -> None:
                     sm.find_role(role),
                     password="general",
                 )
-        sm.get_session.commit()

superset/cli/update.py

@@ -30,13 +30,15 @@
 from flask_appbuilder.api.manager import resolver
 
 import superset.utils.database as database_utils
+from superset.utils.decorators import transaction
 from superset.utils.encrypt import SecretsMigrator
 
 logger = logging.getLogger(__name__)
 
 
 @click.command()
 @with_appcontext
+@transaction()
 @click.option("--database_name", "-d", help="Database name to change")
 @click.option("--uri", "-u", help="Database URI to change")
 @click.option(
@@ -53,6 +55,7 @@ def set_database_uri(database_name: str, uri: str, skip_create: bool) -> None:
 
 @click.command()
 @with_appcontext
+@transaction()
 def sync_tags() -> None:
     """Rebuilds special tags (owner, type, favorited by)."""
     # pylint: disable=no-member

superset/commands/annotation_layer/annotation/create.py

@@ -16,6 +16,7 @@
 # under the License.
 import logging
 from datetime import datetime
+from functools import partial
 from typing import Any, Optional
 
 from flask_appbuilder.models.sqla import Model
@@ -30,7 +31,7 @@
 from superset.commands.annotation_layer.exceptions import AnnotationLayerNotFoundError
 from superset.commands.base import BaseCommand
 from superset.daos.annotation_layer import AnnotationDAO, AnnotationLayerDAO
-from superset.daos.exceptions import DAOCreateFailedError
+from superset.utils.decorators import on_error, transaction
 
 logger = logging.getLogger(__name__)
 
@@ -39,13 +40,10 @@ class CreateAnnotationCommand(BaseCommand):
     def __init__(self, data: dict[str, Any]):
         self._properties = data.copy()
 
+    @transaction(on_error=partial(on_error, reraise=AnnotationCreateFailedError))
     def run(self) -> Model:
         self.validate()
-        try:
-            return AnnotationDAO.create(attributes=self._properties)
-        except DAOCreateFailedError as ex:
-            logger.exception(ex.exception)
-            raise AnnotationCreateFailedError() from ex
+        return AnnotationDAO.create(attributes=self._properties)
 
     def validate(self) -> None:
         exceptions: list[ValidationError] = []

superset/commands/annotation_layer/annotation/delete.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 import logging
+from functools import partial
 from typing import Optional
 
 from superset.commands.annotation_layer.annotation.exceptions import (
@@ -23,8 +24,8 @@
 )
 from superset.commands.base import BaseCommand
 from superset.daos.annotation_layer import AnnotationDAO
-from superset.daos.exceptions import DAODeleteFailedError
 from superset.models.annotations import Annotation
+from superset.utils.decorators import on_error, transaction
 
 logger = logging.getLogger(__name__)
 
@@ -34,15 +35,11 @@ def __init__(self, model_ids: list[int]):
         self._model_ids = model_ids
         self._models: Optional[list[Annotation]] = None
 
+    @transaction(on_error=partial(on_error, reraise=AnnotationDeleteFailedError))
     def run(self) -> None:
         self.validate()
         assert self._models
-
-        try:
-            AnnotationDAO.delete(self._models)
-        except DAODeleteFailedError as ex:
-            logger.exception(ex.exception)
-            raise AnnotationDeleteFailedError() from ex
+        AnnotationDAO.delete(self._models)
 
     def validate(self) -> None:
         # Validate/populate model exists

superset/commands/annotation_layer/annotation/update.py

@@ -16,6 +16,7 @@
 # under the License.
 import logging
 from datetime import datetime
+from functools import partial
 from typing import Any, Optional
 
 from flask_appbuilder.models.sqla import Model
@@ -31,8 +32,8 @@
 from superset.commands.annotation_layer.exceptions import AnnotationLayerNotFoundError
 from superset.commands.base import BaseCommand
 from superset.daos.annotation_layer import AnnotationDAO, AnnotationLayerDAO
-from superset.daos.exceptions import DAOUpdateFailedError
 from superset.models.annotations import Annotation
+from superset.utils.decorators import on_error, transaction
 
 logger = logging.getLogger(__name__)
 
@@ -43,16 +44,11 @@ def __init__(self, model_id: int, data: dict[str, Any]):
         self._properties = data.copy()
         self._model: Optional[Annotation] = None
 
+    @transaction(on_error=partial(on_error, reraise=AnnotationUpdateFailedError))
     def run(self) -> Model:
         self.validate()
         assert self._model
-
-        try:
-            annotation = AnnotationDAO.update(self._model, self._properties)
-        except DAOUpdateFailedError as ex:
-            logger.exception(ex.exception)
-            raise AnnotationUpdateFailedError() from ex
-        return annotation
+        return AnnotationDAO.update(self._model, self._properties)
 
     def validate(self) -> None:
         exceptions: list[ValidationError] = []

superset/commands/annotation_layer/create.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 import logging
+from functools import partial
 from typing import Any
 
 from flask_appbuilder.models.sqla import Model
@@ -27,7 +28,7 @@
 )
 from superset.commands.base import BaseCommand
 from superset.daos.annotation_layer import AnnotationLayerDAO
-from superset.daos.exceptions import DAOCreateFailedError
+from superset.utils.decorators import on_error, transaction
 
 logger = logging.getLogger(__name__)
 
@@ -36,13 +37,10 @@ class CreateAnnotationLayerCommand(BaseCommand):
     def __init__(self, data: dict[str, Any]):
         self._properties = data.copy()
 
+    @transaction(on_error=partial(on_error, reraise=AnnotationLayerCreateFailedError))
     def run(self) -> Model:
         self.validate()
-        try:
-            return AnnotationLayerDAO.create(attributes=self._properties)
-        except DAOCreateFailedError as ex:
-            logger.exception(ex.exception)
-            raise AnnotationLayerCreateFailedError() from ex
+        return AnnotationLayerDAO.create(attributes=self._properties)
 
     def validate(self) -> None:
         exceptions: list[ValidationError] = []

superset/commands/annotation_layer/delete.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 import logging
+from functools import partial
 from typing import Optional
 
 from superset.commands.annotation_layer.exceptions import (
@@ -24,8 +25,8 @@
 )
 from superset.commands.base import BaseCommand
 from superset.daos.annotation_layer import AnnotationLayerDAO
-from superset.daos.exceptions import DAODeleteFailedError
 from superset.models.annotations import AnnotationLayer
+from superset.utils.decorators import on_error, transaction
 
 logger = logging.getLogger(__name__)
 
@@ -35,15 +36,11 @@ def __init__(self, model_ids: list[int]):
         self._model_ids = model_ids
         self._models: Optional[list[AnnotationLayer]] = None
 
+    @transaction(on_error=partial(on_error, reraise=AnnotationLayerDeleteFailedError))
     def run(self) -> None:
         self.validate()
         assert self._models
-
-        try:
-            AnnotationLayerDAO.delete(self._models)
-        except DAODeleteFailedError as ex:
-            logger.exception(ex.exception)
-            raise AnnotationLayerDeleteFailedError() from ex
+        AnnotationLayerDAO.delete(self._models)
 
     def validate(self) -> None:
         # Validate/populate model exists

superset/commands/annotation_layer/update.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 import logging
+from functools import partial
 from typing import Any, Optional
 
 from flask_appbuilder.models.sqla import Model
@@ -28,8 +29,8 @@
 )
 from superset.commands.base import BaseCommand
 from superset.daos.annotation_layer import AnnotationLayerDAO
-from superset.daos.exceptions import DAOUpdateFailedError
 from superset.models.annotations import AnnotationLayer
+from superset.utils.decorators import on_error, transaction
 
 logger = logging.getLogger(__name__)
 
@@ -40,16 +41,11 @@ def __init__(self, model_id: int, data: dict[str, Any]):
         self._properties = data.copy()
         self._model: Optional[AnnotationLayer] = None
 
+    @transaction(on_error=partial(on_error, reraise=AnnotationLayerUpdateFailedError))
     def run(self) -> Model:
         self.validate()
         assert self._model
-
-        try:
-            annotation_layer = AnnotationLayerDAO.update(self._model, self._properties)
-        except DAOUpdateFailedError as ex:
-            logger.exception(ex.exception)
-            raise AnnotationLayerUpdateFailedError() from ex
-        return annotation_layer
+        return AnnotationLayerDAO.update(self._model, self._properties)
 
     def validate(self) -> None:
         exceptions: list[ValidationError] = []