feat: Added datadoc tags functionality and search filtering

querybook/config/elasticsearch.yaml

@@ -154,6 +154,8 @@ datadocs:
                     type: boolean
                 readable_user_ids:
                     type: integer
+                tags:
+                    type: keyword
 tables:
     index_name: search_tables_v1
     type_name: tables # Keep this same in mappings

querybook/migrations/versions/8f194c1b59a3_data_doc_tag_item.py

@@ -0,0 +1,38 @@
+"""data_doc_tag_item
+
+Revision ID: 8f194c1b59a3
+Revises: 7f6cdb3621f7
+Create Date: 2023-06-01 20:40:07.871141
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '8f194c1b59a3'
+down_revision = '7f6cdb3621f7'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('data_doc_tag_item',
+    sa.Column('id', sa.Integer(), nullable=False),
+    sa.Column('created_at', sa.DateTime(), nullable=True),
+    sa.Column('tag_name', sa.String(length=255), nullable=False),
+    sa.Column('datadoc_id', sa.Integer(), nullable=True),
+    sa.Column('uid', sa.Integer(), nullable=True),
+    sa.ForeignKeyConstraint(['datadoc_id'], ['data_doc.id'], ondelete='CASCADE'),
+    sa.ForeignKeyConstraint(['tag_name'], ['tag.name'], ondelete='CASCADE'),
+    sa.ForeignKeyConstraint(['uid'], ['user.id'], ondelete='SET NULL'),
+    sa.PrimaryKeyConstraint('id')
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table('data_doc_tag_item')
+    # ### end Alembic commands ###

querybook/server/datasources/tag.py

@@ -5,6 +5,7 @@
 from app.auth.permission import (
     verify_data_table_permission,
     verify_data_column_permission,
+    verify_data_doc_permission,
 )
 from logic import tag as logic
 from models.tag import Tag
@@ -19,6 +20,15 @@ def get_tags_by_table_id(table_id: int):
     return logic.get_tags_by_table_id(table_id=table_id)
 
 
+@register(
+    "/datadoc/<int:datadoc_id>/tag/",
+    methods=["GET"],
+)
+def get_tags_by_datadoc_id(datadoc_id: int):
+    verify_data_doc_permission(datadoc_id)
+    return logic.get_tags_by_datadoc_id(datadoc_id=datadoc_id)
+
+
 @register(
     "/column/<int:column_id>/tag/",
     methods=["GET"],
@@ -63,6 +73,22 @@ def add_tag_to_table(table_id, tag):
         )
 
 
+@register(
+    "/datadoc/<int:datadoc_id>/tag/",
+    methods=["POST"],
+)
+def add_tag_to_datadoc(datadoc_id, tag):
+    with DBSession() as session:
+        verify_data_doc_permission(datadoc_id, session=session)
+        return logic.add_tag_to_datadoc(
+            datadoc_id=datadoc_id,
+            tag_name=tag,
+            uid=current_user.id,
+            user_is_admin=current_user.is_admin,
+            session=session,
+        )
+
+
 @register(
     "/table/<int:table_id>/tag/",
     methods=["DELETE"],
@@ -76,3 +102,18 @@ def delete_tag_from_table(table_id: int, tag_name: str):
             user_is_admin=current_user.is_admin,
             session=session,
         )
+
+
+@register(
+    "/datadoc/<int:datadoc_id>/tag/",
+    methods=["DELETE"],
+)
+def delete_tag_from_datadoc(datadoc_id: int, tag_name: str):
+    with DBSession() as session:
+        verify_data_doc_permission(datadoc_id, session=session)
+        return logic.delete_tag_from_datadoc(
+            datadoc_id=datadoc_id,
+            tag_name=tag_name,
+            user_is_admin=current_user.is_admin,
+            session=session,
+        )

querybook/server/lib/elasticsearch/search_datadoc.py

@@ -44,7 +44,7 @@ def construct_datadoc_query(
         keywords,
         search_fields=_match_data_doc_fields(fields),
     )
-    search_filter = match_filters(filters)
+    search_filter = match_filters(filters, and_filter_names=["tags"])
     search_filter.setdefault("filter", {}).setdefault("bool", {}).setdefault(
         "must", []
     ).append({"bool": {"should": _data_doc_access_terms(uid)}})
@@ -53,7 +53,7 @@ def construct_datadoc_query(
         "query": {
             "bool": combine_keyword_and_filter_query(keywords_query, search_filter)
         },
-        "_source": ["id", "title", "owner_uid", "created_at"],
+        "_source": ["id", "title", "owner_uid", "created_at", "scheduled", "tags"],
         "size": limit,
         "from": offset,
     }

querybook/server/logic/elasticsearch.py

@@ -438,6 +438,8 @@ def datadocs_to_es(datadoc, fields=None, session=None):
         "title": datadoc.title,
         "public": datadoc.public,
         "readable_user_ids": lambda: _get_datadoc_editors(datadoc, session=session),
+        "scheduled": datadoc.scheduled,
+        "tags": [tag.tag_name for tag in datadoc.tags],
     }
     return _get_dict_by_field(field_to_getter, fields=fields)
 

querybook/server/logic/tag.py

@@ -4,7 +4,8 @@
 from const.metastore import DataTag
 from lib.utils.color import find_nearest_palette_color
 from logic.metastore import update_es_tables_by_id
-from models.tag import Tag, TagItem
+from logic.datadoc import update_es_data_doc_by_id
+from models.tag import Tag, TagItem, DataDocTagItem
 
 
 @with_session
@@ -18,6 +19,17 @@ def get_tags_by_table_id(table_id, session=None):
     )
 
 
+@with_session
+def get_tags_by_datadoc_id(datadoc_id, session=None):
+    return (
+        session.query(Tag)
+        .join(DataDocTagItem)
+        .filter(DataDocTagItem.datadoc_id == datadoc_id)
+        .order_by(Tag.count.desc())
+        .all()
+    )
+
+
 @with_session
 def get_tags_by_column_id(column_id: int, session=None):
     return (
@@ -84,6 +96,27 @@ def add_tag_to_table(table_id, tag_name, uid, user_is_admin=False, session=None)
     return tag
 
 
+@with_session
+def add_tag_to_datadoc(datadoc_id, tag_name, uid, user_is_admin=False, session=None):
+    existing_tag_item = DataDocTagItem.get(
+        datadoc_id=datadoc_id, tag_name=tag_name, session=session
+    )
+
+    if existing_tag_item:
+        return
+
+    tag = create_or_update_tag(tag_name=tag_name, commit=False, session=session)
+    if (tag.meta or {}).get("admin"):
+        assert user_is_admin, f"Tag {tag_name} can only be modified by admin"
+
+    DataDocTagItem.create(
+        {"tag_name": tag_name, "datadoc_id": datadoc_id, "uid": uid}, session=session
+    )
+    update_es_data_doc_by_id(datadoc_id)
+
+    return tag
+
+
 @with_session
 def delete_tag_from_table(
     table_id, tag_name, user_is_admin=False, commit=True, session=None
@@ -105,6 +138,29 @@ def delete_tag_from_table(
         session.flush()
 
 
+@with_session
+def delete_tag_from_datadoc(
+    datadoc_id, tag_name, user_is_admin=False, commit=True, session=None
+):
+    tag_item = DataDocTagItem.get(
+        datadoc_id=datadoc_id, tag_name=tag_name, session=session
+    )
+    tag = tag_item.tag
+
+    tag.count = tag_item.tag.count - 1
+    tag.update_at = datetime.datetime.now()
+    if (tag.meta or {}).get("admin"):
+        assert user_is_admin, f"Tag {tag_name} can only be modified by admin"
+
+    session.delete(tag_item)
+
+    if commit:
+        session.commit()
+        update_es_data_doc_by_id(tag_item.datadoc_id)
+    else:
+        session.flush()
+
+
 @with_session
 def create_table_tags(
     table_id: int,
@@ -148,6 +204,49 @@ def create_table_tags(
         session.flush()
 
 
+@with_session
+def create_datadoc_tags(
+    datadoc_id: int,
+    tags: list[DataTag] = [],
+    commit=True,
+    session=None,
+):
+    """This function is used for loading datadoc tags from metastore."""
+    # delete all tags from the table
+    session.query(DataDocTagItem).filter_by(datadoc_id=datadoc_id).delete()
+
+    for tag in tags:
+        tag_color_name = (
+            find_nearest_palette_color(tag.color)["name"]
+            if tag.color is not None
+            else None
+        )
+        meta = {
+            "type": tag.type,
+            "tooltip": tag.description,
+            "color": tag_color_name,
+            "admin": True,
+        }
+        # filter out properties with none values
+        meta = {k: v for k, v in meta.items() if v is not None}
+
+        # update or create a new tag if not exist
+        create_or_update_tag(
+            tag_name=tag.name, meta=meta, commit=commit, session=session
+        )
+
+        # add a new tag_item to associate with the datadoc
+        TagItem.create(
+            {"tag_name": tag.name, "datadoc_id": datadoc_id, "uid": None},
+            session=session,
+        )
+
+    if commit:
+        session.commit()
+    else:
+        session.flush()
+
+
 @with_session
 def create_column_tags(
     column_id: int,

querybook/server/models/datadoc.py

@@ -86,6 +86,7 @@ def to_dict(self, with_cells=False):
             "updated_at": self.updated_at,
             "meta": self.meta,
             "title": self.title,
+            "tags": self.tags,
         }
 
         if with_cells:

querybook/server/models/tag.py

@@ -65,3 +65,37 @@ class TagItem(CRUDMixin, Base):
         backref=backref("tags", cascade="all, delete", passive_deletes=True),
         foreign_keys=[column_id],
     )
+
+
+class DataDocTagItem(CRUDMixin, Base):
+    __tablename__ = "data_doc_tag_item"
+
+    id = sql.Column(sql.Integer, primary_key=True)
+    created_at = sql.Column(sql.DateTime, default=now)
+
+    tag_name = sql.Column(
+        sql.String(length=name_length),
+        sql.ForeignKey("tag.name", ondelete="CASCADE"),
+        nullable=False,
+    )
+
+    datadoc_id = sql.Column(
+        sql.Integer, sql.ForeignKey("data_doc.id", ondelete="CASCADE"), nullable=True
+    )
+
+    uid = sql.Column(
+        sql.Integer, sql.ForeignKey("user.id", ondelete="SET NULL"), nullable=True
+    )
+
+    tag = relationship(
+        "Tag",
+        backref=backref(
+            "data_doc_tag_item", cascade="all, delete", passive_deletes=True
+        ),
+        foreign_keys=[tag_name],
+    )
+    datadoc = relationship(
+        "DataDoc",
+        backref=backref("tags", cascade="all, delete", passive_deletes=True),
+        foreign_keys=[datadoc_id],
+    )

querybook/tests/test_lib/test_elasticsearch/test_elasticsearch.py

@@ -245,6 +245,9 @@ class DataDocTestCase(TestCase):
     ENVIRONMENT_ID = 7
     OWNER_UID = "bob"
     DATADOC_TITLE = "Test DataDoc"
+    SCHEDULED = True
+
+    TAGS = [type("", (object,), {"tag_name": "1"})()]
 
     def _get_datadoc_cells_mock(self):
         return [
@@ -273,6 +276,8 @@ def _get_datadoc_mock(self):
             title=self.DATADOC_TITLE,
             public=False,
             cells=self._get_datadoc_cells_mock(),
+            scheduled=self.SCHEDULED,
+            tags=self.TAGS,
         )
         return mock_doc
 
@@ -304,6 +309,8 @@ def test_data_doc_to_es(self):
             "title": self.DATADOC_TITLE,
             "public": False,
             "readable_user_ids": ["alice", "charlie"],
+            "scheduled": self.SCHEDULED,
+            "tags": [tag.tag_name for tag in self.TAGS],
         }
         self.assertEqual(result, expected_result)
 

querybook/webapp/components/DataDoc/DataDoc.scss

@@ -190,3 +190,8 @@
         z-index: 40;
     }
 }
+
+.data-doc-tag {
+    margin-left: 18px;
+    margin-bottom: 8px;
+}

querybook/webapp/components/DataDoc/DataDoc.tsx

@@ -65,6 +65,8 @@ import { DataDocHeader } from './DataDocHeader';
 import { DataDocLoading } from './DataDocLoading';
 
 import './DataDoc.scss';
+import { AccentText } from '../../ui/StyledText/StyledText';
+import { DataDocTags } from '../DataDocTags/DataDocTags';
 
 interface IOwnProps {
     docId: number;
@@ -745,6 +747,17 @@ class DataDocComponent extends React.PureComponent<IProps, IState> {
                         isSaving={isSaving}
                         lastUpdated={lastUpdated}
                     />
+                    <div className="data-doc-tag flex-row">
+                        <AccentText
+                            className="header-subtitle mr20"
+                            weight="bold"
+                            color="lightest"
+
+                        >
+                            Tags
+                        </AccentText>
+                        <DataDocTags datadocId={dataDoc.id} />
+                    </div>
                     <div className={docClassName}>
                         <DataDocTemplateCell
                             dataDoc={dataDoc}