feat: Add dataset tagging to the back-end (#20892)

Co-authored-by: Ville Brofeldt <33317356+villebro@users.noreply.github.com>

superset/initialization/__init__.py

@@ -49,6 +49,7 @@
 )
 from superset.security import SupersetSecurityManager
 from superset.superset_typing import FlaskResponse
+from superset.tags.core import register_sqla_event_listeners
 from superset.utils.core import pessimistic_connection_handling
 from superset.utils.log import DBEventLogger, get_event_logger_from_cfg_value
 
@@ -426,6 +427,9 @@ def init_app_in_ctx(self) -> None:
         if flask_app_mutator:
             flask_app_mutator(self.superset_app)
 
+        if feature_flag_manager.is_feature_enabled("TAGGING_SYSTEM"):
+            register_sqla_event_listeners()
+
         self.init_views()
 
     def check_secret_key(self) -> None:

superset/migrations/versions/2018-07-26_11-10_c82ee8a39623_add_implicit_tags.py

@@ -33,7 +33,7 @@
 from sqlalchemy import Column, DateTime, Enum, ForeignKey, Integer, String
 from sqlalchemy.ext.declarative import declarative_base, declared_attr
 
-from superset.models.tags import ObjectTypes, TagTypes
+from superset.tags.models import ObjectTypes, TagTypes
 from superset.utils.core import get_user_id
 
 Base = declarative_base()

superset/models/core.py

@@ -53,13 +53,12 @@
 from sqlalchemy.schema import UniqueConstraint
 from sqlalchemy.sql import expression, Select
 
-from superset import app, db_engine_specs, is_feature_enabled
+from superset import app, db_engine_specs
 from superset.constants import PASSWORD_MASK
 from superset.databases.utils import make_url_safe
 from superset.db_engine_specs.base import MetricType, TimeGrain
 from superset.extensions import cache_manager, encrypted_field_factory, security_manager
 from superset.models.helpers import AuditMixinNullable, ImportExportMixin
-from superset.models.tags import FavStarUpdater
 from superset.result_set import SupersetResultSet
 from superset.utils import cache as cache_util, core as utils
 from superset.utils.core import get_username
@@ -809,9 +808,3 @@ class FavStar(Model):  # pylint: disable=too-few-public-methods
     class_name = Column(String(50))
     obj_id = Column(Integer)
     dttm = Column(DateTime, default=datetime.utcnow)
-
-
-# events for updating tags
-if is_feature_enabled("TAGGING_SYSTEM"):
-    sqla.event.listen(FavStar, "after_insert", FavStarUpdater.after_insert)
-    sqla.event.listen(FavStar, "after_delete", FavStarUpdater.after_delete)

superset/models/dashboard.py

@@ -53,7 +53,6 @@
 from superset.models.filter_set import FilterSet
 from superset.models.helpers import AuditMixinNullable, ImportExportMixin
 from superset.models.slice import Slice
-from superset.models.tags import DashboardUpdater
 from superset.models.user_attributes import UserAttribute
 from superset.tasks.thumbnails import cache_dashboard_thumbnail
 from superset.utils import core as utils
@@ -454,12 +453,6 @@ def id_or_slug_filter(id_or_slug: Union[int, str]) -> BinaryExpression:
 
 OnDashboardChange = Callable[[Mapper, Connection, Dashboard], Any]
 
-# events for updating tags
-if is_feature_enabled("TAGGING_SYSTEM"):
-    sqla.event.listen(Dashboard, "after_insert", DashboardUpdater.after_insert)
-    sqla.event.listen(Dashboard, "after_update", DashboardUpdater.after_update)
-    sqla.event.listen(Dashboard, "after_delete", DashboardUpdater.after_delete)
-
 if is_feature_enabled("THUMBNAILS_SQLA_LISTENERS"):
     update_thumbnail: OnDashboardChange = lambda _, __, dash: dash.update_thumbnail()
     sqla.event.listen(Dashboard, "after_insert", update_thumbnail)

superset/models/slice.py

@@ -42,7 +42,6 @@
 from superset import db, is_feature_enabled, security_manager
 from superset.legacy import update_time_range
 from superset.models.helpers import AuditMixinNullable, ImportExportMixin
-from superset.models.tags import ChartUpdater
 from superset.tasks.thumbnails import cache_chart_thumbnail
 from superset.utils import core as utils
 from superset.utils.hashing import md5_sha_from_str
@@ -367,13 +366,6 @@ def event_after_chart_changed(
 sqla.event.listen(Slice, "before_insert", set_related_perm)
 sqla.event.listen(Slice, "before_update", set_related_perm)
 
-# events for updating tags
-if is_feature_enabled("TAGGING_SYSTEM"):
-    sqla.event.listen(Slice, "after_insert", ChartUpdater.after_insert)
-    sqla.event.listen(Slice, "after_update", ChartUpdater.after_update)
-    sqla.event.listen(Slice, "after_delete", ChartUpdater.after_delete)
-
-# events for updating tags
 if is_feature_enabled("THUMBNAILS_SQLA_LISTENERS"):
     sqla.event.listen(Slice, "after_insert", event_after_chart_changed)
     sqla.event.listen(Slice, "after_update", event_after_chart_changed)

superset/models/sql_lab.py

@@ -49,7 +49,6 @@
     ExtraJSONMixin,
     ImportExportMixin,
 )
-from superset.models.tags import QueryUpdater
 from superset.sql_parse import CtasMethod, ParsedQuery, Table
 from superset.sqllab.limiting_factor import LimitingFactor
 from superset.superset_typing import ResultSetColumnType
@@ -509,9 +508,3 @@ def to_dict(self) -> Dict[str, Any]:
             "description": description,
             "expanded": self.expanded,
         }
-
-
-# events for updating tags
-sqla.event.listen(SavedQuery, "after_insert", QueryUpdater.after_insert)
-sqla.event.listen(SavedQuery, "after_update", QueryUpdater.after_update)
-sqla.event.listen(SavedQuery, "after_delete", QueryUpdater.after_delete)

superset/tags/core.py

@@ -0,0 +1,88 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+
+def register_sqla_event_listeners() -> None:
+    import sqlalchemy as sqla
+
+    from superset.connectors.sqla.models import SqlaTable
+    from superset.models.core import FavStar
+    from superset.models.dashboard import Dashboard
+    from superset.models.slice import Slice
+    from superset.models.sql_lab import SavedQuery
+    from superset.tags.models import (
+        ChartUpdater,
+        DashboardUpdater,
+        DatasetUpdater,
+        FavStarUpdater,
+        QueryUpdater,
+    )
+
+    sqla.event.listen(SqlaTable, "after_insert", DatasetUpdater.after_insert)
+    sqla.event.listen(SqlaTable, "after_update", DatasetUpdater.after_update)
+    sqla.event.listen(SqlaTable, "after_delete", DatasetUpdater.after_delete)
+
+    sqla.event.listen(Slice, "after_insert", ChartUpdater.after_insert)
+    sqla.event.listen(Slice, "after_update", ChartUpdater.after_update)
+    sqla.event.listen(Slice, "after_delete", ChartUpdater.after_delete)
+
+    sqla.event.listen(Dashboard, "after_insert", DashboardUpdater.after_insert)
+    sqla.event.listen(Dashboard, "after_update", DashboardUpdater.after_update)
+    sqla.event.listen(Dashboard, "after_delete", DashboardUpdater.after_delete)
+
+    sqla.event.listen(FavStar, "after_insert", FavStarUpdater.after_insert)
+    sqla.event.listen(FavStar, "after_delete", FavStarUpdater.after_delete)
+
+    sqla.event.listen(SavedQuery, "after_insert", QueryUpdater.after_insert)
+    sqla.event.listen(SavedQuery, "after_update", QueryUpdater.after_update)
+    sqla.event.listen(SavedQuery, "after_delete", QueryUpdater.after_delete)
+
+
+def clear_sqla_event_listeners() -> None:
+    import sqlalchemy as sqla
+
+    from superset.connectors.sqla.models import SqlaTable
+    from superset.models.core import FavStar
+    from superset.models.dashboard import Dashboard
+    from superset.models.slice import Slice
+    from superset.models.sql_lab import SavedQuery
+    from superset.tags.models import (
+        ChartUpdater,
+        DashboardUpdater,
+        DatasetUpdater,
+        FavStarUpdater,
+        QueryUpdater,
+    )
+
+    sqla.event.remove(SqlaTable, "after_insert", DatasetUpdater.after_insert)
+    sqla.event.remove(SqlaTable, "after_update", DatasetUpdater.after_update)
+    sqla.event.remove(SqlaTable, "after_delete", DatasetUpdater.after_delete)
+
+    sqla.event.remove(Slice, "after_insert", ChartUpdater.after_insert)
+    sqla.event.remove(Slice, "after_update", ChartUpdater.after_update)
+    sqla.event.remove(Slice, "after_delete", ChartUpdater.after_delete)
+
+    sqla.event.remove(Dashboard, "after_insert", DashboardUpdater.after_insert)
+    sqla.event.remove(Dashboard, "after_update", DashboardUpdater.after_update)
+    sqla.event.remove(Dashboard, "after_delete", DashboardUpdater.after_delete)
+
+    sqla.event.remove(FavStar, "after_insert", FavStarUpdater.after_insert)
+    sqla.event.remove(FavStar, "after_delete", FavStarUpdater.after_delete)
+
+    sqla.event.remove(SavedQuery, "after_insert", QueryUpdater.after_insert)
+    sqla.event.remove(SavedQuery, "after_update", QueryUpdater.after_update)
+    sqla.event.remove(SavedQuery, "after_delete", QueryUpdater.after_delete)

superset/models/tags.py → superset/tags/models.py

@@ -14,7 +14,13 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-from __future__ import absolute_import, division, print_function, unicode_literals
+from __future__ import (
+    absolute_import,
+    annotations,
+    division,
+    print_function,
+    unicode_literals,
+)
 
 import enum
 from typing import List, Optional, TYPE_CHECKING, Union
@@ -28,6 +34,7 @@
 from superset.models.helpers import AuditMixinNullable
 
 if TYPE_CHECKING:
+    from superset.connectors.sqla.models import SqlaTable
     from superset.models.core import FavStar
     from superset.models.dashboard import Dashboard
     from superset.models.slice import Slice
@@ -41,7 +48,7 @@ class TagTypes(enum.Enum):
     """
     Types for tags.
 
-    Objects (queries, charts and dashboards) will have with implicit tags based
+    Objects (queries, charts, dashboards, and datasets) will have with implicit tags based
     on metadata: types, owners and who favorited them. This way, user "alice"
     can find all their objects by querying for the tag `owner:alice`.
     """
@@ -64,11 +71,12 @@ class ObjectTypes(enum.Enum):
     query = 1
     chart = 2
     dashboard = 3
+    dataset = 4
 
 
 class Tag(Model, AuditMixinNullable):
 
-    """A tag attached to an object (query, chart or dashboard)."""
+    """A tag attached to an object (query, chart, dashboard, or dataset)."""
 
     __tablename__ = "tag"
     id = Column(Integer, primary_key=True)
@@ -103,6 +111,7 @@ def get_object_type(class_name: str) -> ObjectTypes:
         "slice": ObjectTypes.chart,
         "dashboard": ObjectTypes.dashboard,
         "query": ObjectTypes.query,
+        "dataset": ObjectTypes.dataset,
     }
     try:
         return mapping[class_name.lower()]
@@ -116,13 +125,15 @@ class ObjectUpdater:
 
     @classmethod
     def get_owners_ids(
-        cls, target: Union["Dashboard", "FavStar", "Slice"]
+        cls, target: Union[Dashboard, FavStar, Slice, Query, SqlaTable]
     ) -> List[int]:
         raise NotImplementedError("Subclass should implement `get_owners_ids`")
 
     @classmethod
     def _add_owners(
-        cls, session: Session, target: Union["Dashboard", "FavStar", "Slice"]
+        cls,
+        session: Session,
+        target: Union[Dashboard, FavStar, Slice, Query, SqlaTable],
     ) -> None:
         for owner_id in cls.get_owners_ids(target):
             name = "owner:{0}".format(owner_id)
@@ -137,7 +148,7 @@ def after_insert(
         cls,
         _mapper: Mapper,
         connection: Connection,
-        target: Union["Dashboard", "FavStar", "Slice"],
+        target: Union[Dashboard, FavStar, Slice, Query, SqlaTable],
     ) -> None:
         session = Session(bind=connection)
 
@@ -158,7 +169,7 @@ def after_update(
         cls,
         _mapper: Mapper,
         connection: Connection,
-        target: Union["Dashboard", "FavStar", "Slice"],
+        target: Union[Dashboard, FavStar, Slice, Query, SqlaTable],
     ) -> None:
         session = Session(bind=connection)
 
@@ -187,7 +198,7 @@ def after_delete(
         cls,
         _mapper: Mapper,
         connection: Connection,
-        target: Union["Dashboard", "FavStar", "Slice"],
+        target: Union[Dashboard, FavStar, Slice, Query, SqlaTable],
     ) -> None:
         session = Session(bind=connection)
 
@@ -205,7 +216,7 @@ class ChartUpdater(ObjectUpdater):
     object_type = "chart"
 
     @classmethod
-    def get_owners_ids(cls, target: "Slice") -> List[int]:
+    def get_owners_ids(cls, target: Slice) -> List[int]:
         return [owner.id for owner in target.owners]
 
 
@@ -214,7 +225,7 @@ class DashboardUpdater(ObjectUpdater):
     object_type = "dashboard"
 
     @classmethod
-    def get_owners_ids(cls, target: "Dashboard") -> List[int]:
+    def get_owners_ids(cls, target: Dashboard) -> List[int]:
         return [owner.id for owner in target.owners]
 
 
@@ -223,14 +234,23 @@ class QueryUpdater(ObjectUpdater):
     object_type = "query"
 
     @classmethod
-    def get_owners_ids(cls, target: "Query") -> List[int]:
+    def get_owners_ids(cls, target: Query) -> List[int]:
         return [target.user_id]
 
 
+class DatasetUpdater(ObjectUpdater):
+
+    object_type = "dataset"
+
+    @classmethod
+    def get_owners_ids(cls, target: SqlaTable) -> List[int]:
+        return [owner.id for owner in target.owners]
+
+
 class FavStarUpdater:
     @classmethod
     def after_insert(
-        cls, _mapper: Mapper, connection: Connection, target: "FavStar"
+        cls, _mapper: Mapper, connection: Connection, target: FavStar
     ) -> None:
         session = Session(bind=connection)
         name = "favorited_by:{0}".format(target.user_id)
@@ -246,7 +266,7 @@ def after_insert(
 
     @classmethod
     def after_delete(
-        cls, _mapper: Mapper, connection: Connection, target: "FavStar"
+        cls, _mapper: Mapper, connection: Connection, target: FavStar
     ) -> None:
         session = Session(bind=connection)
         name = "favorited_by:{0}".format(target.user_id)

superset/tasks/cache.py

@@ -28,7 +28,7 @@
 from superset.models.core import Log
 from superset.models.dashboard import Dashboard
 from superset.models.slice import Slice
-from superset.models.tags import Tag, TaggedObject
+from superset.tags.models import Tag, TaggedObject
 from superset.utils.date_parser import parse_human_datetime
 from superset.utils.machine_auth import MachineAuthProvider
 

superset/utils/url_map_converters.py

@@ -18,7 +18,7 @@
 
 from werkzeug.routing import BaseConverter, Map
 
-from superset.models.tags import ObjectTypes
+from superset.tags.models import ObjectTypes
 
 
 class RegexConverter(BaseConverter):