feat: add column level tag support (#1169)

* feat: add column level tag support

* comments

* set commit False for create_table

* tagsDOM

* use verify_data_column_permission

* ffix linter

querybook/migrations/versions/cf334666f8f6_add_column_tags_support.py

@@ -0,0 +1,30 @@
+"""add column tags support
+
+Revision ID: cf334666f8f6
+Revises: adfbd684e30c
+Create Date: 2023-02-16 00:25:05.116326
+
+"""
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import mysql
+
+# revision identifiers, used by Alembic.
+revision = 'cf334666f8f6'
+down_revision = 'adfbd684e30c'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('tag_item', sa.Column('column_id', sa.Integer(), nullable=True))
+    op.create_foreign_key("fk_tag_item_data_table_column", 'tag_item', 'data_table_column', ['column_id'], ['id'], ondelete='CASCADE')
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint("fk_tag_item_data_table_column", 'tag_item', type_='foreignkey')
+    op.drop_column('tag_item', 'column_id')
+    # ### end Alembic commands ###

querybook/server/datasources/tag.py

@@ -2,7 +2,10 @@
 
 from app.datasource import register, api_assert
 from app.db import DBSession
-from app.auth.permission import verify_data_table_permission
+from app.auth.permission import (
+    verify_data_table_permission,
+    verify_data_column_permission,
+)
 from logic import tag as logic
 from models.tag import Tag
 
@@ -11,10 +14,18 @@
     "/table/<int:table_id>/tag/",
     methods=["GET"],
 )
-def get_tag_by_table_id(table_id: int):
-    with DBSession() as session:
-        verify_data_table_permission(table_id, session=session)
-        return logic.get_tag_by_table_id(table_id=table_id, session=session)
+def get_tags_by_table_id(table_id: int):
+    verify_data_table_permission(table_id)
+    return logic.get_tags_by_table_id(table_id=table_id)
+
+
+@register(
+    "/column/<int:column_id>/tag/",
+    methods=["GET"],
+)
+def get_tags_by_column_id(column_id: int):
+    verify_data_column_permission(column_id)
+    return logic.get_tags_by_column_id(column_id=column_id)
 
 
 @register(

querybook/server/lib/metastore/base_metastore_loader.py

@@ -26,7 +26,7 @@
     get_schema_by_name,
     get_table_by_schema_id_and_name,
 )
-from logic.tag import create_table_tags
+from logic.tag import create_table_tags, create_column_tags
 
 from .metastore_data_types import DataTable, DataColumn
 from .utils import MetastoreTableACLChecker
@@ -312,6 +312,7 @@ def _create_table_table(
                 location=table.location,
                 column_count=len(columns),
                 schema_id=schema_id,
+                commit=False,
                 session=session,
             ).id
             create_table_information(
@@ -329,24 +330,34 @@ def _create_table_table(
             )
 
             for column in columns:
-                create_column(
+                column_id = create_column(
                     name=column.name,
                     type=column.type,
                     comment=column.comment,
                     description=column.description,
                     table_id=table_id,
                     commit=False,
                     session=session,
-                )
+                ).id
 
-            # create tags if the metastore is configured to sync tags
+                # create tags only if the metastore is configured to sync tags
+                if self.loader_config.can_load_external_metadata(MetadataType.TAG):
+                    create_column_tags(
+                        column_id=column_id,
+                        tags=column.tags,
+                        commit=False,
+                        session=session,
+                    )
+
+            # create tags only if the metastore is configured to sync tags
             if self.loader_config.can_load_external_metadata(MetadataType.TAG):
                 create_table_tags(
                     table_id=table_id,
                     tags=table.tags,
                     commit=False,
                     session=session,
                 )
+
             session.commit()
             update_table_by_id(table_id, session=session)
             return table_id

querybook/server/lib/metastore/metastore_data_types.py

@@ -58,3 +58,6 @@ class DataColumn(NamedTuple):
 
     # user edited description from metastore, expect HTML format
     description: str = None
+
+    # list of column level tags from metastore
+    tags: List[DataTag] = []

querybook/server/logic/tag.py

@@ -6,7 +6,7 @@
 
 
 @with_session
-def get_tag_by_table_id(table_id, session=None):
+def get_tags_by_table_id(table_id, session=None):
     return (
         session.query(Tag)
         .join(TagItem)
@@ -16,6 +16,17 @@ def get_tag_by_table_id(table_id, session=None):
     )
 
 
+@with_session
+def get_tags_by_column_id(column_id: int, session=None):
+    return (
+        session.query(Tag)
+        .join(TagItem)
+        .filter(TagItem.column_id == column_id)
+        .order_by(Tag.count.desc())
+        .all()
+    )
+
+
 @with_session
 def get_tags_by_keyword(keyword, limit=10, session=None):
     return (
@@ -94,12 +105,12 @@ def delete_tag_from_table(
 
 @with_session
 def create_table_tags(
-    table_id: int = None,
+    table_id: int,
     tags: list[DataTag] = [],
     commit=True,
     session=None,
 ):
-    """This function is used for loading tags from metastore."""
+    """This function is used for loading table tags from metastore."""
     # delete all tags from the table
     session.query(TagItem).filter_by(table_id=table_id).delete()
 
@@ -128,3 +139,41 @@ def create_table_tags(
         session.commit()
     else:
         session.flush()
+
+
+@with_session
+def create_column_tags(
+    column_id: int,
+    tags: list[DataTag] = [],
+    commit=True,
+    session=None,
+):
+    """This function is used for loading column tags from metastore."""
+    # delete all tags from the table
+    session.query(TagItem).filter_by(column_id=column_id).delete()
+
+    for tag in tags:
+        meta = {
+            "type": tag.type,
+            "tooltip": tag.description,
+            "color": tag.color,
+            "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 table
+        TagItem.create(
+            {"tag_name": tag.name, "column_id": column_id, "uid": None},
+            session=session,
+        )
+
+    if commit:
+        session.commit()
+    else:
+        session.flush()

querybook/server/models/tag.py

@@ -40,6 +40,11 @@ class TagItem(CRUDMixin, Base):
     table_id = sql.Column(
         sql.Integer, sql.ForeignKey("data_table.id", ondelete="CASCADE"), nullable=True
     )
+    column_id = sql.Column(
+        sql.Integer,
+        sql.ForeignKey("data_table_column.id", ondelete="CASCADE"),
+        nullable=True,
+    )
 
     uid = sql.Column(
         sql.Integer, sql.ForeignKey("user.id", ondelete="SET NULL"), nullable=True
@@ -55,3 +60,8 @@ class TagItem(CRUDMixin, Base):
         backref=backref("tags", cascade="all, delete", passive_deletes=True),
         foreign_keys=[table_id],
     )
+    column = relationship(
+        "DataTableColumn",
+        backref=backref("tags", cascade="all, delete", passive_deletes=True),
+        foreign_keys=[column_id],
+    )

querybook/webapp/components/DataTableViewColumn/DataTableColumnCard.tsx

@@ -2,10 +2,13 @@ import { ContentState } from 'draft-js';
 import React, { useMemo } from 'react';
 
 import { DataTableColumnStats } from 'components/DataTableStats/DataTableColumnStats';
+import { TableTag } from 'components/DataTableTags/DataTableTags';
 import { IDataColumn } from 'const/metastore';
+import { useResource } from 'hooks/useResource';
 import { useToggleState } from 'hooks/useToggleState';
 import { Nullable } from 'lib/typescript';
 import { parseType } from 'lib/utils/complex-types';
+import { TableColumnResource } from 'resource/table';
 import { Card } from 'ui/Card/Card';
 import { EditableTextField } from 'ui/EditableTextField/EditableTextField';
 import { Icon } from 'ui/Icon/Icon';
@@ -30,9 +33,19 @@ export const DataTableColumnCard: React.FunctionComponent<IProps> = ({
     onEditColumnDescriptionRedirect,
     updateDataColumnDescription,
 }) => {
+    const { data: columnTags } = useResource(
+        React.useCallback(
+            () => TableColumnResource.getTags(column.id),
+            [column.id]
+        )
+    );
     const [expanded, , toggleExpanded] = useToggleState(false);
     const parsedType = useMemo(() => parseType('', column.type), [column.type]);
 
+    const tagsDOM = (columnTags || []).map((tag) => (
+        <TableTag tag={tag} readonly={true} key={tag.id} mini={true} />
+    ));
+
     const userCommentsContent = (
         <EditableTextField
             value={column.description as ContentState}
@@ -69,6 +82,13 @@ export const DataTableColumnCard: React.FunctionComponent<IProps> = ({
                                 />
                             </KeyContentDisplay>
                         )}
+                        {tagsDOM.length > 0 && (
+                            <KeyContentDisplay keyString="Tags">
+                                <div className="DataTableTags flex-row">
+                                    {tagsDOM}
+                                </div>
+                            </KeyContentDisplay>
+                        )}
                         {column.comment && (
                             <KeyContentDisplay keyString="Definition">
                                 {column.comment}

querybook/webapp/resource/table.ts

@@ -177,6 +177,7 @@ export const TableColumnResource = {
 
         return ds.update<IDataColumn>(`/column/${columnId}/`, params);
     },
+    getTags: (columnId: number) => ds.fetch<ITag[]>(`/column/${columnId}/tag/`),
 };
 
 export const TableLineageResource = {