chore(dao): Condense delete/bulk-delete operations (#24466)

Co-authored-by: Michael S. Molina <[email protected]>

superset/annotation_layers/annotations/commands/bulk_delete.py

@@ -36,9 +36,10 @@ def __init__(self, model_ids: list[int]):
 
     def run(self) -> None:
         self.validate()
+        assert self._models
+
         try:
-            AnnotationDAO.bulk_delete(self._models)
-            return None
+            AnnotationDAO.delete(self._models)
         except DAODeleteFailedError as ex:
             logger.exception(ex.exception)
             raise AnnotationBulkDeleteFailedError() from ex

superset/annotation_layers/annotations/commands/delete.py

@@ -17,8 +17,6 @@
 import logging
 from typing import Optional
 
-from flask_appbuilder.models.sqla import Model
-
 from superset.annotation_layers.annotations.commands.exceptions import (
     AnnotationDeleteFailedError,
     AnnotationNotFoundError,
@@ -36,16 +34,15 @@ def __init__(self, model_id: int):
         self._model_id = model_id
         self._model: Optional[Annotation] = None
 
-    def run(self) -> Model:
+    def run(self) -> None:
         self.validate()
         assert self._model
 
         try:
-            annotation = AnnotationDAO.delete(self._model)
+            AnnotationDAO.delete(self._model)
         except DAODeleteFailedError as ex:
             logger.exception(ex.exception)
             raise AnnotationDeleteFailedError() from ex
-        return annotation
 
     def validate(self) -> None:
         # Validate/populate model exists

superset/annotation_layers/commands/bulk_delete.py

@@ -37,9 +37,10 @@ def __init__(self, model_ids: list[int]):
 
     def run(self) -> None:
         self.validate()
+        assert self._models
+
         try:
-            AnnotationLayerDAO.bulk_delete(self._models)
-            return None
+            AnnotationLayerDAO.delete(self._models)
         except DAODeleteFailedError as ex:
             logger.exception(ex.exception)
             raise AnnotationLayerBulkDeleteFailedError() from ex

superset/annotation_layers/commands/delete.py

@@ -17,8 +17,6 @@
 import logging
 from typing import Optional
 
-from flask_appbuilder.models.sqla import Model
-
 from superset.annotation_layers.commands.exceptions import (
     AnnotationLayerDeleteFailedError,
     AnnotationLayerDeleteIntegrityError,
@@ -37,16 +35,15 @@ def __init__(self, model_id: int):
         self._model_id = model_id
         self._model: Optional[AnnotationLayer] = None
 
-    def run(self) -> Model:
+    def run(self) -> None:
         self.validate()
         assert self._model
 
         try:
-            annotation_layer = AnnotationLayerDAO.delete(self._model)
+            AnnotationLayerDAO.delete(self._model)
         except DAODeleteFailedError as ex:
             logger.exception(ex.exception)
             raise AnnotationLayerDeleteFailedError() from ex
-        return annotation_layer
 
     def validate(self) -> None:
         # Validate/populate model exists

superset/charts/commands/bulk_delete.py

@@ -43,8 +43,10 @@ def __init__(self, model_ids: list[int]):
 
     def run(self) -> None:
         self.validate()
+        assert self._models
+
         try:
-            ChartDAO.bulk_delete(self._models)
+            ChartDAO.delete(self._models)
         except DeleteFailedError as ex:
             logger.exception(ex.exception)
             raise ChartBulkDeleteFailedError() from ex

superset/charts/commands/delete.py

@@ -17,7 +17,6 @@
 import logging
 from typing import cast, Optional
 
-from flask_appbuilder.models.sqla import Model
 from flask_babel import lazy_gettext as _
 
 from superset import security_manager
@@ -43,7 +42,7 @@ def __init__(self, model_id: int):
         self._model_id = model_id
         self._model: Optional[Slice] = None
 
-    def run(self) -> Model:
+    def run(self) -> None:
         self.validate()
         self._model = cast(Slice, self._model)
         try:
@@ -52,11 +51,10 @@ def run(self) -> Model:
             # table, sporadically Superset will error because the rows are not deleted.
             # Let's do it manually here to prevent the error.
             self._model.owners = []
-            chart = ChartDAO.delete(self._model)
+            ChartDAO.delete(self._model)
         except DAODeleteFailedError as ex:
             logger.exception(ex.exception)
             raise ChartDeleteFailedError() from ex
-        return chart
 
     def validate(self) -> None:
         # Validate/populate model exists

superset/css_templates/commands/bulk_delete.py

@@ -36,9 +36,10 @@ def __init__(self, model_ids: list[int]):
 
     def run(self) -> None:
         self.validate()
+        assert self._models
+
         try:
-            CssTemplateDAO.bulk_delete(self._models)
-            return None
+            CssTemplateDAO.delete(self._models)
         except DAODeleteFailedError as ex:
             logger.exception(ex.exception)
             raise CssTemplateBulkDeleteFailedError() from ex

superset/daos/annotation.py

@@ -17,30 +17,14 @@
 import logging
 from typing import Optional, Union
 
-from sqlalchemy.exc import SQLAlchemyError
-
 from superset.daos.base import BaseDAO
-from superset.daos.exceptions import DAODeleteFailedError
 from superset.extensions import db
 from superset.models.annotations import Annotation, AnnotationLayer
 
 logger = logging.getLogger(__name__)
 
 
 class AnnotationDAO(BaseDAO[Annotation]):
-    @staticmethod
-    def bulk_delete(models: Optional[list[Annotation]], commit: bool = True) -> None:
-        item_ids = [model.id for model in models] if models else []
-        try:
-            db.session.query(Annotation).filter(Annotation.id.in_(item_ids)).delete(
-                synchronize_session="fetch"
-            )
-            if commit:
-                db.session.commit()
-        except SQLAlchemyError as ex:
-            db.session.rollback()
-            raise DAODeleteFailedError() from ex
-
     @staticmethod
     def validate_update_uniqueness(
         layer_id: int, short_descr: str, annotation_id: Optional[int] = None
@@ -63,21 +47,6 @@ def validate_update_uniqueness(
 
 
 class AnnotationLayerDAO(BaseDAO[AnnotationLayer]):
-    @staticmethod
-    def bulk_delete(
-        models: Optional[list[AnnotationLayer]], commit: bool = True
-    ) -> None:
-        item_ids = [model.id for model in models] if models else []
-        try:
-            db.session.query(AnnotationLayer).filter(
-                AnnotationLayer.id.in_(item_ids)
-            ).delete(synchronize_session="fetch")
-            if commit:
-                db.session.commit()
-        except SQLAlchemyError as ex:
-            db.session.rollback()
-            raise DAODeleteFailedError() from ex
-
     @staticmethod
     def has_annotations(model_id: Union[int, list[int]]) -> bool:
         if isinstance(model_id, list):

superset/daos/base.py

@@ -14,7 +14,9 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-from typing import Any, Generic, get_args, Optional, TypeVar, Union
+from __future__ import annotations
+
+from typing import Any, Generic, get_args, TypeVar
 
 from flask_appbuilder.models.filters import BaseFilter
 from flask_appbuilder.models.sqla import Model
@@ -29,6 +31,7 @@
     DAOUpdateFailedError,
 )
 from superset.extensions import db
+from superset.utils.core import get_iterable
 
 T = TypeVar("T", bound=Model)  # pylint: disable=invalid-name
 
@@ -38,12 +41,12 @@ class BaseDAO(Generic[T]):
     Base DAO, implement base CRUD sqlalchemy operations
     """
 
-    model_cls: Optional[type[Model]] = None
+    model_cls: type[Model] | None = None
     """
     Child classes need to state the Model class so they don't need to implement basic
     create, update and delete methods
     """
-    base_filter: Optional[BaseFilter] = None
+    base_filter: BaseFilter | None = None
     """
     Child classes can register base filtering to be applied to all filter methods
     """
@@ -57,10 +60,10 @@ def __init_subclass__(cls) -> None:  # pylint: disable=arguments-differ
     @classmethod
     def find_by_id(
         cls,
-        model_id: Union[str, int],
+        model_id: str | int,
         session: Session = None,
         skip_base_filter: bool = False,
-    ) -> Optional[Model]:
+    ) -> Model | None:
         """
         Find a model by id, if defined applies `base_filter`
         """
@@ -81,7 +84,7 @@ def find_by_id(
     @classmethod
     def find_by_ids(
         cls,
-        model_ids: Union[list[str], list[int]],
+        model_ids: list[str] | list[int],
         session: Session = None,
         skip_base_filter: bool = False,
     ) -> list[T]:
@@ -114,7 +117,7 @@ def find_all(cls) -> list[T]:
         return query.all()
 
     @classmethod
-    def find_one_or_none(cls, **filter_by: Any) -> Optional[T]:
+    def find_one_or_none(cls, **filter_by: Any) -> T | None:
         """
         Get the first that fit the `base_filter`
         """
@@ -180,25 +183,28 @@ def update(cls, model: T, properties: dict[str, Any], commit: bool = True) -> T:
         return model
 
     @classmethod
-    def delete(cls, model: T, commit: bool = True) -> T:
+    def delete(cls, items: T | list[T], commit: bool = True) -> None:
         """
-        Generic delete a model
-        :raises: DAODeleteFailedError
+        Delete the specified item(s) including their associated relationships.
+
+        Note that bulk deletion via `delete` is not invoked in the base class as this
+        does not dispatch the ORM `after_delete` event which may be required to augment
+        additional records loosely defined via implicit relationships. Instead ORM
+        objects are deleted one-by-one via `Session.delete`.
+
+        Subclasses may invoke bulk deletion but are responsible for instrumenting any
+        post-deletion logic.
+
+        :param items: The item(s) to delete
+        :param commit: Whether to commit the transaction
+        :raises DAODeleteFailedError: If the deletion failed
+        :see: https://docs.sqlalchemy.org/en/latest/orm/queryguide/dml.html
         """
-        try:
-            db.session.delete(model)
-            if commit:
-                db.session.commit()
-        except SQLAlchemyError as ex:  # pragma: no cover
-            db.session.rollback()
-            raise DAODeleteFailedError(exception=ex) from ex
-        return model
 
-    @classmethod
-    def bulk_delete(cls, models: list[T], commit: bool = True) -> None:
         try:
-            for model in models:
-                cls.delete(model, False)
+            for item in get_iterable(items):
+                db.session.delete(item)
+
             if commit:
                 db.session.commit()
         except SQLAlchemyError as ex:

superset/daos/chart.py

@@ -15,9 +15,11 @@
 # specific language governing permissions and limitations
 # under the License.
 # pylint: disable=arguments-renamed
+from __future__ import annotations
+
 import logging
 from datetime import datetime
-from typing import Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING
 
 from sqlalchemy.exc import SQLAlchemyError
 
@@ -26,7 +28,7 @@
 from superset.extensions import db
 from superset.models.core import FavStar, FavStarClassName
 from superset.models.slice import Slice
-from superset.utils.core import get_user_id
+from superset.utils.core import get_iterable, get_user_id
 
 if TYPE_CHECKING:
     from superset.connectors.base.models import BaseDatasource
@@ -37,15 +39,14 @@
 class ChartDAO(BaseDAO[Slice]):
     base_filter = ChartFilter
 
-    @staticmethod
-    def bulk_delete(models: Optional[list[Slice]], commit: bool = True) -> None:
-        item_ids = [model.id for model in models] if models else []
+    @classmethod
+    def delete(cls, items: Slice | list[Slice], commit: bool = True) -> None:
+        item_ids = [item.id for item in get_iterable(items)]
         # bulk delete, first delete related data
-        if models:
-            for model in models:
-                model.owners = []
-                model.dashboards = []
-                db.session.merge(model)
+        for item in get_iterable(items):
+            item.owners = []
+            item.dashboards = []
+            db.session.merge(item)
         # bulk delete itself
         try:
             db.session.query(Slice).filter(Slice.id.in_(item_ids)).delete(

superset/daos/css.py

@@ -14,30 +14,9 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-import logging
-from typing import Optional
-
-from sqlalchemy.exc import SQLAlchemyError
-
 from superset.daos.base import BaseDAO
-from superset.daos.exceptions import DAODeleteFailedError
-from superset.extensions import db
 from superset.models.core import CssTemplate
 
-logger = logging.getLogger(__name__)
-
 
 class CssTemplateDAO(BaseDAO[CssTemplate]):
-    @staticmethod
-    def bulk_delete(models: Optional[list[CssTemplate]], commit: bool = True) -> None:
-        item_ids = [model.id for model in models] if models else []
-        try:
-            db.session.query(CssTemplate).filter(CssTemplate.id.in_(item_ids)).delete(
-                synchronize_session="fetch"
-            )
-            if commit:
-                db.session.commit()
-        except SQLAlchemyError as ex:
-            if commit:
-                db.session.rollback()
-            raise DAODeleteFailedError() from ex
+    pass